Add stateless setter API

DataDog · Oct 18, 2023 · a29d052 · a29d052
1 parent 6c1eed1
commit a29d052
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 29 deletions.
diff --git a/src/jdk.jfr/share/classes/jdk/jfr/ContextType.java b/src/jdk.jfr/share/classes/jdk/jfr/ContextType.java
@@ -15,10 +15,16 @@
 import jdk.jfr.internal.context.ContextRepository;
 
 public abstract class ContextType extends BaseContextType {
-    public static interface Registration {
+    public interface Registration {
         Stream<Class<? extends ContextType>> types();
     }
 
+    public interface Setter {
+        void setAttribute(String name, String value);
+        void clearAttribute(String name);
+        void clearAll();
+    }
+
     public static final class Captured<T extends ContextType.Capturable<T>> implements Closeable, AutoCloseable {
         private final T parent, current;
 
@@ -55,4 +61,8 @@ public final Captured<T> capture() {
     }
 
     protected ContextType() {}
+
+    public static Setter setterFor(Class<? extends ContextType> type) {
+        return BaseContextType.setterFor(type);
+    }
 }
diff --git a/src/jdk.jfr/share/classes/jdk/jfr/internal/context/BaseContextType.java b/src/jdk.jfr/share/classes/jdk/jfr/internal/context/BaseContextType.java
@@ -1,5 +1,6 @@
 package jdk.jfr.internal.context;
 
+import jdk.jfr.ContextType;
 import jdk.jfr.FlightRecorder;
 import jdk.jfr.internal.JVM;
 import jdk.jfr.internal.PlatformRecorder;
@@ -24,4 +25,8 @@ public final boolean isActive() {
     private static boolean shouldCaptureState() {
         return FlightRecorder.isInitialized() && PlatformRecorder.hasRecordings() && JVM.isContextEnabled();
     }
+
+    public static ContextType.Setter setterFor(Class<? extends ContextType> type) {
+        return ContextRepository.getOrRegister(type);
+    }
 }
diff --git a/src/jdk.jfr/share/classes/jdk/jfr/internal/context/ContextDescriptor.java b/src/jdk.jfr/share/classes/jdk/jfr/internal/context/ContextDescriptor.java
@@ -2,5 +2,5 @@
 
 import java.lang.invoke.VarHandle;
 
-public record ContextDescriptor(String holderId, String name, String label, String description, VarHandle access) {
+public record ContextDescriptor(int order, String holderId, String name, String label, String description, VarHandle access) {
 }
diff --git a/src/jdk.jfr/share/classes/jdk/jfr/internal/context/ContextRepository.java b/src/jdk.jfr/share/classes/jdk/jfr/internal/context/ContextRepository.java
@@ -6,7 +6,9 @@
 import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import jdk.jfr.ContextType;
@@ -48,12 +50,13 @@ public static <T extends BaseContextType> ContextWriter getOrRegister(Class<T> c
         return contextWriters.get(contextTypeClass);
     }
 
-    private static <T extends BaseContextType> List<ContextDescriptor> descriptorsOf(Class<T> contextTypeClass) {
+    private static <T extends BaseContextType> Set<ContextDescriptor> descriptorsOf(Class<T> contextTypeClass) {
         try {
-            List<ContextDescriptor> ctxDescriptors = new ArrayList<>(8);
+            Set<ContextDescriptor> ctxDescriptors = new HashSet<>(8);
             Name typeNameAnnot = contextTypeClass.getAnnotation(Name.class);
             String id = typeNameAnnot.value();
             MethodHandles.Lookup lookup = MethodHandles.publicLookup();
+            int order = 0;
             for (Field f : contextTypeClass.getFields()) {
                 Name nameAnnot = f.getAnnotation(Name.class);
                 Label labelAnnot = f.getAnnotation(Label.class);
@@ -62,7 +65,7 @@ private static <T extends BaseContextType> List<ContextDescriptor> descriptorsOf
                     String name = nameAnnot != null ? nameAnnot.value() : f.getName();
                     String label = labelAnnot != null ? labelAnnot.value() : name;
                     String desc = descAnnot != null ? descAnnot.value() : "";
-                    ctxDescriptors.add(new ContextDescriptor(id, name, label, desc, lookup.unreflectVarHandle(f)));
+                    ctxDescriptors.add(new ContextDescriptor(order++, id, name, label, desc, lookup.unreflectVarHandle(f)));
                 }
             }
             return ctxDescriptors;
@@ -72,7 +75,7 @@ private static <T extends BaseContextType> List<ContextDescriptor> descriptorsOf
     }
 
     private static <T extends BaseContextType> ContextWriter writerFor(Class<T> type) {
-        List<ContextDescriptor> ctxDescriptors = descriptorsOf(type);
+        Set<ContextDescriptor> ctxDescriptors = descriptorsOf(type);
         int offset = register(ctxDescriptors);
         if (offset == -2) {
             Logger.log(JFR, INFO,
@@ -85,7 +88,7 @@ private static <T extends BaseContextType> ContextWriter writerFor(Class<T> type
         return offset > -1 ? new ContextWriter(offset, ctxDescriptors) : ContextWriter.NULL;
     }
 
-    private static int register(List<ContextDescriptor> descriptors) {
+    private static int register(Set<ContextDescriptor> descriptors) {
         if (!shouldCaptureState()) {
             return -1;
         }
@@ -96,12 +99,12 @@ private static int register(List<ContextDescriptor> descriptors) {
             return -3;
         }
         final int offset = slotPointer;
-        int idx = 0;
+        int top = 0;
         for (var descriptor : descriptors) {
-            allDescriptors[idx + slotPointer] = descriptor;
-            idx++;
+            allDescriptors[descriptor.order() + slotPointer] = descriptor;
+            top = Math.max(top, descriptor.order());
         }
-        slotPointer += idx;
+        slotPointer += top + 1;
         JVM.setUsedContextSize(slotPointer);
         return offset;
     }

diff --git a/src/jdk.jfr/share/classes/jdk/jfr/internal/context/ContextWriter.java b/src/jdk.jfr/share/classes/jdk/jfr/internal/context/ContextWriter.java
@@ -1,26 +1,25 @@
 package jdk.jfr.internal.context;
 
-import java.lang.invoke.MethodHandles.Lookup;
 import java.nio.LongBuffer;
-import java.lang.invoke.MethodType;
-import java.lang.invoke.VarHandle;
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 import jdk.jfr.ContextType;
-import jdk.jfr.FlightRecorder;
 import jdk.jfr.internal.JVM;
 import jdk.jfr.internal.StringPool;
 
-public final class ContextWriter {
+public final class ContextWriter implements ContextType.Setter {
     static final ContextWriter NULL = new ContextWriter(-1, null);
     private final int offset;
-    private final List<ContextDescriptor> descriptors;
+    private final Set<ContextDescriptor> descriptors;
+    private final Map<String, ContextDescriptor> attributeIndexMap;
 
-    ContextWriter(int offset, List<ContextDescriptor> descriptors) {
+    ContextWriter(int offset, Set<ContextDescriptor> descriptors) {
         this.offset = offset;
-        this.descriptors = offset > -1 ? Collections.unmodifiableList(descriptors) : null;
+        this.descriptors = offset > -1 ? Collections.unmodifiableSet(descriptors) : null;
+        this.attributeIndexMap = this.descriptors != null ? this.descriptors.stream().collect(Collectors.toUnmodifiableMap(ContextDescriptor::name, cd -> cd)) : null;
     }
 
     public boolean isActive() {
@@ -35,12 +34,10 @@ void write(BaseContextType target) {
         if (context == null) {
             return;
         }
-        int cntr = offset;
         for (ContextDescriptor cd : descriptors) {
-            String value = (String) cd.access().get(target);
-            context.put(cntr++, value != null ? StringPool.addString(value, false) : 0);
-            if (cntr >= 8) {
-                break;
+            if (cd.order() < 8) {
+                String value = (String) cd.access().get(target);
+                context.put(offset + cd.order(), value != null ? StringPool.addString(value, false) : 0);
             }
         }
     }
@@ -53,10 +50,64 @@ void clear(BaseContextType target) {
         if (context == null) {
             return;
         }
-        int cntr = 0;
         for (ContextDescriptor cd : descriptors) {
-            context.put(offset + (cntr++), 0L);
-            cd.access().set(target, null);
+            if (cd.order() < 8) {
+                context.put(offset + cd.order(), 0L);
+                cd.access().set(target, null);
+            }
+        }
+    }
+
+    @Override
+    public void clearAll() {
+        if (offset == -1 || descriptors == null) {
+            return;
+        }
+        LongBuffer context = JVM.getThreadContextBuffer();
+        if (context == null) {
+            return;
+        }
+        for (ContextDescriptor cd : descriptors) {
+            context.put(offset + cd.order(), 0L);
+        }
+    }
+
+    @Override
+    public void setAttribute(String name, String value) {
+        LongBuffer context = JVM.getThreadContextBuffer();
+        if (context == null) {
+            return;
+        }
+        int pos = getContextIndex(name, String.class);
+        if (pos < 0) {
+            System.err.println("===> set err: " + name + ": " + pos);
+            return;
+        }
+        context.put(pos, StringPool.addString(value, false));
+    }
+
+    @Override
+    public void clearAttribute(String name) {
+        LongBuffer context = JVM.getThreadContextBuffer();
+        if (context == null) {
+            return;
+        }
+        int pos = getContextIndex(name, String.class);
+        if (pos < 0) {
+            System.err.println("===> clear err: " + name + ": " + pos);
+            return;
+        }
+        context.put(pos, 0);
+    }
+
+    private int getContextIndex(String name, Class<?> type) {
+        ContextDescriptor cd = attributeIndexMap.get(name);
+        if (cd == null) {
+            return -1;
+        }
+        if (cd.access().varType() != type) {
+            return -2;
         }
+        return offset + cd.order();
     }
 }