Rework context capturing concept

Author: jbachorik

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

@@ -10,5 +10,4 @@
 @Target(ElementType.TYPE)
 @Retention(RetentionPolicy.RUNTIME)
 public @interface ContextAware {
-    boolean value() default true;
 }

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

@@ -1,5 +1,14 @@
 package jdk.jfr;
 
+import java.io.Closeable;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+import java.lang.invoke.MethodHandles.Lookup;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.stream.Stream;
 
 import jdk.jfr.internal.context.BaseContextType;
@@ -10,5 +19,40 @@ public static interface Registration {
         Stream<Class<? extends ContextType>> types();
     }
 
+    public static final class Captured<T extends ContextType.Capturable<T>> implements Closeable, AutoCloseable {
+        private final T parent, current;
+
+        Captured(T parent, T current) {
+            assert parent != null;
+            assert current != null;
+            this.parent = parent;
+            this.current = current;
+        }
+
+        @Override
+        public void close() {
+            parent.set();
+        }
+
+        public T get() {
+            return current;
+        }
+
+        public Captured<T> capture() {
+            return current.capture();
+        }
+    }
+
+    public static abstract class Capturable<T extends Capturable<T>> extends ContextType {
+        public Capturable() {}
+
+        @SuppressWarnings("unchecked")
+        public final Captured<T> capture() {
+            return new Captured<T>((T) this, snapshot());
+        }
+
+        abstract protected T snapshot();
+    }
+
     protected ContextType() {}
 }

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

@@ -192,15 +192,11 @@ boolean isEnabled() {
     }
 
     boolean isContextAware() {
-        Boolean result = annotationValue(classNode, ANNOTATION_CONTEXT_AWARE_DESCRIPTOR, Boolean.class, Boolean.TRUE);
-        if (result != null) {
-            return result.booleanValue();
+        if (hasAnnotation(classNode, ANNOTATION_CONTEXT_AWARE_DESCRIPTOR)) {
+            return true;
         }
         if (superClass != null) {
-            ContextAware e = superClass.getAnnotation(ContextAware.class);
-            if (e != null) {
-                return e.value();
-            }
+            return superClass.getAnnotation(ContextAware.class) != null;
         }
         return false;
     }
@@ -236,6 +232,17 @@ private static <T> T annotationValue(ClassNode classNode, String typeDescriptor,
         return null;
     }
 
+    private static boolean hasAnnotation(ClassNode classNode, String typeDescriptor) {
+        if (classNode.visibleAnnotations != null) {
+            for (AnnotationNode a : classNode.visibleAnnotations) {
+                if (typeDescriptor.equals(a.desc)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     private static List<SettingInfo> buildSettingInfos(Class<?> superClass, ClassNode classNode) {
         Set<String> methodSet = new HashSet<>();
         List<SettingInfo> settingInfos = new ArrayList<>();

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

@@ -268,7 +268,7 @@ public static synchronized Type createType(Class<?> clazz, List<AnnotationElemen
 
         if (eventType) {
             ContextAware ctxAware = clazz.getAnnotation(ContextAware.class);            
-            addImplicitFields(type, true, true, true, true , false, ctxAware != null && ctxAware.value());
+            addImplicitFields(type, true, true, true, true , false, ctxAware != null);
             addUserFields(clazz, type, dynamicFields);
             type.trimFields();
         }

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

@@ -1,13 +1,8 @@
 package jdk.jfr.internal.context;
 
-import java.nio.LongBuffer;
-import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
-
-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 {
     public final void set() {

test/micro/org/openjdk/bench/jdk/jfr/JfrContextBenchmark.java

@@ -10,6 +10,7 @@
 import jdk.jfr.ContextType;
 import jdk.jfr.FlightRecorder;
 import jdk.jfr.Recording;
+import jdk.jfr.ContextType.Captured;
 
 @BenchmarkMode(Mode.AverageTime)
 @Warmup(iterations = 10, time = 500, timeUnit = TimeUnit.MILLISECONDS)
@@ -32,7 +33,6 @@ public class JfrContextBenchmark {
         }
         if (recordJFR) {
             Recording r = new Recording();
-            // r.setDestination(Paths.get("/tmp/bench.jfr"));
             r.start();
         }
     }
@@ -48,4 +48,23 @@ public void operationWithContext(Blackhole bh) throws IOException {
         ctx.unset();
         bh.consume(ctx);
     }
+
+    @Benchmark 
+    @Threads(2)
+    public void operationWithCapturedContext(Blackhole bh) throws IOException {
+        try (Captured<TracingJfrContext> ctxCapture = new TracingJfrContext().capture()) {
+            ctxCapture.get().withValues("trace-1", "span-1").set();
+            TracedEvent.commit(10);
+            try (Captured<TracingJfrContext> ctx1Capture = ctxCapture.capture()) {
+                TracingJfrContext ctx1 = ctx1Capture.get();
+                ctx1.spanid = "span-2";
+                ctx1.set();
+                TracedEvent.commit(20);
+                bh.consume(ctx1);
+            }
+            // this should be within context [traceid=trace-1, spanid=span-1]
+            TracedEvent.commit(30);
+        }
+        
+    }
 }

test/micro/org/openjdk/bench/jdk/jfr/TracingJfrContext.java

@@ -4,10 +4,29 @@
 import jdk.jfr.Name;
 
 @Name("tracing_context")
-public class TracingJfrContext extends ContextType {
+public final class TracingJfrContext extends ContextType.Capturable<TracingJfrContext> {
     @Name("traceid")
     public String traceid;
 
     @Name("spanid")
     public String spanid;
+
+    public TracingJfrContext withValues(String traceid, String spanid) {
+        this.traceid = traceid;
+        this.spanid = spanid;
+        return this;
+    }
+
+    @Override
+    protected TracingJfrContext snapshot() {
+        return new TracingJfrContext(traceid, spanid);
+    }
+
+    public TracingJfrContext(String traceid, String spanid) {
+        this.traceid = traceid;
+        this.spanid = spanid;
+    }
+
+    public TracingJfrContext() {}
+    
 }