Simplify the JFR Context API

src/hotspot/share/jfr/jni/jfrJniMethod.cpp

@@ -399,35 +399,7 @@ JVM_ENTRY_NO_ENV(void, jfr_emit_data_loss(JNIEnv* env, jclass jvm, jlong bytes))
   EventDataLoss::commit(bytes, min_jlong);
 JVM_END
 
-JVM_ENTRY_NO_ENV(jlong, jfr_get_and_set_context(JNIEnv* env, jclass jvm, jint slot, jlong value))
-  return JfrContext::get_and_set_context((uint8_t)slot, (uint64_t)value);
-JVM_END
-
-JVM_ENTRY_NO_ENV(jint, jfr_get_all_context(JNIEnv* env, jclass jvm, jlongArray data, jint length))
-  assert(length > 0 && length <= 8, "invalid length");
-  uint64_t ctx[8];
-  memset(ctx, 0, sizeof(ctx)); // zero the data structure
-  length = JfrContext::get_all_context(ctx, length);
-  if (length > -1) {
-    env->SetLongArrayRegion(data, 0, length, (jlong*)ctx);
-  }
-  return length;
-JVM_END
-
-JVM_ENTRY_NO_ENV(void, jfr_set_all_context(JNIEnv* env, jclass jvm, jlongArray context))
-  jlong* data = env->GetLongArrayElements(context, NULL);
-  JfrContext::set_all_context((uint64_t*)data, 8);
-  env->ReleaseLongArrayElements(context, data, 0);
-JVM_END
-
-JVM_ENTRY_NO_ENV(void, jfr_set_context(JNIEnv* env, jclass jvm, jint offset, jlongArray context))
-  jsize length = env->GetArrayLength(context);
-  jlong* data = env->GetLongArrayElements(context, NULL);
-  JfrContext::set_all_context((uint64_t*)data + offset, length);
-  env->ReleaseLongArrayElements(context, data, 0);
-JVM_END
-
-JVM_ENTRY_NO_ENV(void, jfr_set_context_size(JNIEnv* env, jclass jvm, jint size))
+JVM_ENTRY_NO_ENV(void, jfr_set_used_context_size(JNIEnv* env, jclass jvm, jint size))
   JfrContext::set_used_context_size(size);
 JVM_END
 

src/hotspot/share/jfr/jni/jfrJniMethod.hpp

@@ -159,10 +159,7 @@ jlong JNICALL jfr_host_total_memory(JNIEnv* env, jclass jvm);
 
 void JNICALL jfr_emit_data_loss(JNIEnv* env, jclass jvm, jlong bytes);
 
-jlong JNICALL jfr_get_and_set_context(JNIEnv* env, jclass jvm, jint slot, jlong value);
-jint JNICALL jfr_get_all_context(JNIEnv* env, jclass jvm, jlongArray data, jint length);
-void JNICALL jfr_set_all_context(JNIEnv* env, jclass jvm, jlongArray data);
-void JNICALL jfr_set_context_size(JNIEnv* env, jclass jvm, jint size);
+void JNICALL jfr_set_used_context_size(JNIEnv* env, jclass jvm, jint size);
 jboolean JNICALL jfr_is_context_enabled(JNIEnv* env, jclass jvm);
 jobject JNICALL jfr_get_thread_context_buffer(JNIEnv* env, jclass jvm);
 

src/hotspot/share/jfr/jni/jfrJniMethodRegistration.cpp

@@ -98,10 +98,7 @@ JfrJniMethodRegistration::JfrJniMethodRegistration(JNIEnv* env) {
       (char*)"isContainerized", (char*)"()Z", (void*) jfr_is_containerized,
       (char*)"hostTotalMemory", (char*)"()J", (void*) jfr_host_total_memory,
       (char*)"emitDataLoss", (char*)"(J)V", (void*)jfr_emit_data_loss,
-      (char*)"getAndSetContext", (char*)"(IJ)J", (void*)jfr_get_and_set_context,
-      (char*)"getAllContext", (char*)"([JI)I", (void*)jfr_get_all_context,
-      (char*)"setAllContext", (char*)"([J)V", (void*)jfr_set_all_context,
-      (char*)"setContextSize", (char*)"(I)V", (void*)jfr_set_context_size,
+      (char*)"setUsedContextSize", (char*)"(I)V", (void*)jfr_set_used_context_size,
       (char*)"isContextEnabled", (char*)"()Z", (void*)jfr_is_context_enabled,
       (char*)"getThreadContextBuffer0", (char*)"()Ljava/lang/Object;", (void*)jfr_get_thread_context_buffer
     };

src/jdk.jfr/share/classes/jdk/jfr/ContextType.java

@@ -1,9 +1,14 @@
 package jdk.jfr;
 
-import java.io.Closeable;
+import java.util.stream.Stream;
 
 import jdk.jfr.internal.context.BaseContextType;
+import jdk.jfr.internal.context.ContextRepository;
+
+public abstract class ContextType extends BaseContextType {
+    public static interface Registration {
+        Stream<Class<? extends ContextType>> types();
+    }
 
-public abstract class ContextType extends BaseContextType implements Closeable, AutoCloseable {
     protected ContextType() {}
 }

src/jdk.jfr/share/classes/jdk/jfr/FlightRecorder.java

@@ -37,9 +37,6 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
 
-import jdk.jfr.ContextType;
-import jdk.jfr.context.ContextRegistration;
-import jdk.jfr.context.RegisteredContext;
 import jdk.jfr.internal.JVM;
 import jdk.jfr.internal.JVMSupport;
 import static jdk.jfr.internal.LogLevel.DEBUG;
@@ -52,8 +49,7 @@
 import jdk.jfr.internal.PlatformRecording;
 import jdk.jfr.internal.Repository;
 import jdk.jfr.internal.SecuritySupport;
-import jdk.jfr.internal.context.ContextRegistry;
-import jdk.jfr.internal.context.RegisteredContextImpl;
+import jdk.jfr.internal.context.ContextRepository;
 import jdk.jfr.internal.periodic.PeriodicEvents;
 import jdk.jfr.internal.util.Utils;
 
@@ -69,7 +65,6 @@ public final class FlightRecorder {
     private static volatile FlightRecorder platformRecorder;
     private static volatile boolean initialized;
     private final PlatformRecorder internal;
-    private static final Map<String, RegisteredContextImpl<? extends ContextType>> contextMap = new ConcurrentHashMap<>(5);
 
     private FlightRecorder(PlatformRecorder internal) {
         this.internal = internal;
@@ -92,10 +87,6 @@ public List<Recording> getRecordings() {
         return Collections.unmodifiableList(recs);
     }
 
-    public static boolean hasRecordings() {
-        return PlatformRecorder.hasRecordings();
-    }
-
     /**
      * Creates a snapshot of all available recorded data.
      * <p>
@@ -183,8 +174,9 @@ public static FlightRecorder getFlightRecorder() throws IllegalStateException, S
             JVMSupport.ensureWithIllegalStateException();
             if (platformRecorder == null) {
                 try {
-                    for (ContextRegistration reg : ServiceLoader.load(ContextRegistration.class)) {
-                        reg.types().forEach(FlightRecorder::registerContext);
+                    Logger.log(JFR, DEBUG, "Automatically registering available context types");
+                    for (ContextType.Registration reg : ServiceLoader.load( ContextType.Registration.class)) {
+                        reg.types().forEach(ContextRepository::getOrRegister);
                     }
 
                     platformRecorder = new FlightRecorder(new PlatformRecorder());
@@ -265,23 +257,13 @@ public static boolean removePeriodicEvent(Runnable hook) throws SecurityExceptio
     }
 
     @SuppressWarnings("unchecked")
-    public static <T extends ContextType> RegisteredContext<T> registerContext(Class<T> ctxType) {
+    public static boolean registerContext(Class<? extends ContextType> ctxType) {
         Name typeNameAnnot = ctxType.getAnnotation(Name.class);
         if (typeNameAnnot == null) {
             throw new IllegalArgumentException("Context type must be annotated by @Name");
         }
         String id = typeNameAnnot.value();
-        return (RegisteredContext<T>)contextMap.computeIfAbsent(id, key -> {
-            synchronized (PlatformRecorder.class) {
-                if (platformRecorder != null) {
-                    if (Logger.shouldLog(JFR, INFO)) {
-                        Logger.log(JFR, INFO, "FlightRecorder already initialized. Not possible to request a new context access with id=" + id);
-                    }
-                    return null;
-                }
-                return new RegisteredContextImpl<>(ctxType);
-            }
-        });
+        return ContextRepository.getOrRegister(ctxType).isActive();
     }
 
     /**

src/jdk.jfr/share/classes/jdk/jfr/internal/JVM.java

@@ -33,7 +33,7 @@
 import jdk.internal.access.SharedSecrets;
 import jdk.internal.vm.annotation.IntrinsicCandidate;
 import jdk.jfr.Event;
-import jdk.jfr.internal.context.ContextRegistry;
+import jdk.jfr.internal.context.ContextRepository;
 import jdk.jfr.internal.event.EventConfiguration;
 import jdk.jfr.internal.event.EventWriter;
 
@@ -644,10 +644,18 @@ public void setAllContext(long[] data) {
      */
     public static native void emitDataLoss(long bytes);
 
-    public static native void setContextSize(int size);
-    public static native long getAndSetContext(int idx, long value);
-    public static native int getAllContext(long[] context, int length);
-    public static native void setAllContext(long[] context);
+    public static native void setUsedContextSize(int size);
+    public static int getAllContext(long[] context, int length) {
+        LongBuffer buffer = getThreadContextBuffer();
+        int toCopy = Math.min(buffer.limit(), length);
+        buffer.rewind().get(0, context, 0, toCopy);
+        return toCopy;
+    }
+    public static void setAllContext(long[] context) {
+        LongBuffer buffer = getThreadContextBuffer();
+        int toCopy = Math.min(buffer.limit(), context.length);
+        buffer.put(0, context, 0, toCopy).rewind();
+    }
     public static native boolean isContextEnabled();
 
     private static final ThreadLocal<LongBuffer> threadContextBufferRef = ThreadLocal.withInitial(JVM::initializeContextBuffer);

src/jdk.jfr/share/classes/jdk/jfr/internal/TypeLibrary.java

@@ -58,7 +58,7 @@
 import jdk.jfr.Timespan;
 import jdk.jfr.Timestamp;
 import jdk.jfr.ValueDescriptor;
-import jdk.jfr.internal.context.ContextRegistry;
+import jdk.jfr.internal.context.ContextRepository;
 import jdk.jfr.internal.util.Utils;
 
 public final class TypeLibrary {
@@ -353,7 +353,7 @@ private static void addFields(Type type, boolean requestable, boolean hasDuratio
         }
         if (hasContext) {
             var fldSuffix = 1;
-            for (var descriptor : ContextRegistry.getInstance().registrations()) {
+            for (var descriptor : ContextRepository.registrations()) {
                 String name = descriptor.holderId() + "_" + descriptor.name();
                 var annos = createStandardAnnotations(name, descriptor.holderId() + ":" + descriptor.label(), descriptor.description());
                 type.add(PrivateAccess.getInstance().newValueDescriptor(name, Type.STRING, annos, 0, true, name));

src/jdk.jfr/share/classes/jdk/jfr/internal/context/BaseContextType.java

@@ -6,63 +6,27 @@
 import jdk.jfr.ContextType;
 import jdk.jfr.FlightRecorder;
 import jdk.jfr.internal.JVM;
+import jdk.jfr.internal.PlatformRecorder;
 import jdk.jfr.internal.context.ContextWriter;
 
 public abstract class BaseContextType {
-    private static record State(ContextWriter writer, LongBuffer contextBuffer) {}
-
-    private State state;
-    void setup(ContextWriter writer, LongBuffer contextBuffer) {
-        this.state = new State(writer, contextBuffer);
-    }
-
-    private byte modificationMask = 0;
-    private int captured = 0;
-
     public final void apply() {
-        if (captured > 0 && shouldCaptureState()) {
-            State s = state;
-            if (s != null && s.writer() != null && s.contextBuffer() != null) {
-                s.writer().write(this, s.contextBuffer());
-            }
+        if (shouldCaptureState()) {
+            ContextRepository.getOrRegister(this.getClass()).write(this);
         }
     }
 
-    @SuppressWarnings("unchecked")
-    public final void close() {
-        if (--captured > 0 && shouldCaptureState()) {
-            State s = state;
-            if (s != null && s.writer() != null && s.contextBuffer() != null) {
-                s.writer().close(this, s.contextBuffer());
-            }
+    public final void clear() {
+        if (shouldCaptureState()) {
+            ContextRepository.getOrRegister(this.getClass()).clear(this);
         }
     }
 
-    @SuppressWarnings("unchecked")
-    public final <T extends ContextType> T capture() {
-        if (state != null && shouldCaptureState()) {
-            if (captured++ == 0) {
-                // first capture will create backup
-                state.writer().capture(this, state.contextBuffer());
-            }
-        }
-        return (T) this;
-    }
-
-    public void setModificationMask(byte mask) {
-        modificationMask |= mask;
-    }
-
-    @SuppressWarnings("lossy-conversions")
-    public void clearModificationMask(byte mask) {
-        modificationMask &= ~mask;
-    }
-
-    public byte getModificationMask() {
-        return modificationMask;
+    public final boolean isActive() {
+        return ContextRepository.getOrRegister(this.getClass()).isActive();
     }
 
     private static boolean shouldCaptureState() {
-        return FlightRecorder.isInitialized() && FlightRecorder.hasRecordings() && JVM.isContextEnabled();
+        return FlightRecorder.isInitialized() && PlatformRecorder.hasRecordings() && JVM.isContextEnabled();
     }
 }