[WIP] Temporary patch to trial GlobalObjectStore in the SDK

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java

@@ -1,5 +1,7 @@
 package datadog.trace.bootstrap;
 
+import datadog.instrument.fieldinject.GlobalObjectStore;
+
 /**
  * {@link ContextStore} that attempts to store context in its keys by using bytecode-injected
  * fields. Delegates to a lazy {@link WeakMap} for keys that don't have a field for this store.
@@ -16,7 +18,7 @@ public Object get(final Object key) {
     if (key instanceof FieldBackedContextAccessor) {
       return ((FieldBackedContextAccessor) key).$get$__datadogContext$(storeId);
     } else {
-      return weakStore().get(key);
+      return GlobalObjectStore.get(key, storeId);
     }
   }
 
@@ -25,7 +27,7 @@ public void put(final Object key, final Object context) {
     if (key instanceof FieldBackedContextAccessor) {
       ((FieldBackedContextAccessor) key).$put$__datadogContext$(storeId, context);
     } else {
-      weakStore().put(key, context);
+      GlobalObjectStore.put(key, storeId, context);
     }
   }
 
@@ -45,7 +47,7 @@ public Object putIfAbsent(final Object key, final Object context) {
       }
       return existingContext;
     } else {
-      return weakStore().putIfAbsent(key, context);
+      return GlobalObjectStore.getOrPut(key, storeId, context);
     }
   }
 
@@ -71,7 +73,7 @@ public Object computeIfAbsent(
       }
       return existingContext;
     } else {
-      return weakStore().computeIfAbsent(key, contextFactory);
+      return GlobalObjectStore.getOrCompute(key, storeId, contextFactory::create);
     }
   }
 
@@ -90,22 +92,7 @@ public Object remove(Object key) {
       }
       return existingContext;
     } else {
-      return weakStore().remove(key);
-    }
-  }
-
-  // only create WeakMap-based fall-back when we need it
-  private volatile WeakMapContextStore<Object, Object> weakStore;
-  private final Object synchronizationInstance = new Object();
-
-  WeakMapContextStore<Object, Object> weakStore() {
-    if (null == weakStore) {
-      synchronized (synchronizationInstance) {
-        if (null == weakStore) {
-          weakStore = new WeakMapContextStore<>();
-        }
-      }
+      return GlobalObjectStore.remove(key, storeId);
     }
-    return weakStore;
   }
 }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java

@@ -1,5 +1,6 @@
 package datadog.trace.bootstrap;
 
+import datadog.instrument.fieldinject.GlobalObjectStore;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.Arrays;
 import java.util.concurrent.ConcurrentHashMap;
@@ -137,11 +138,11 @@ private static FieldBackedContextStore createStore(final int storeId) {
 
   /** Injection helper that immediately delegates to the weak-map for the given context store. */
   public static Object weakGet(final Object key, final int storeId) {
-    return getContextStore(storeId).weakStore().get(key);
+    return GlobalObjectStore.get(key, storeId);
   }
 
   /** Injection helper that immediately delegates to the weak-map for the given context store. */
   public static void weakPut(final Object key, final int storeId, final Object context) {
-    getContextStore(storeId).weakStore().put(key, context);
+    GlobalObjectStore.put(key, storeId, context);
   }
 }

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -6,6 +6,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer;
 
 import datadog.environment.SystemProperties;
+import datadog.instrument.fieldinject.GlobalObjectStore;
 import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
 import datadog.trace.agent.tooling.bytebuddy.iast.TaintableRedefinitionStrategyListener;
 import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers;
@@ -59,7 +60,8 @@ public class AgentInstaller {
     enableByteBuddyRawTypes();
     disableByteBuddyNexus();
     // register weak map supplier as early as possible
-    WeakMaps.registerAsSupplier();
+    // WeakMaps.registerAsSupplier();
+    AgentTaskScheduler.get().scheduleAtFixedRate(GlobalObjectStore::removeStaleEntries, 1, 1, TimeUnit.SECONDS);
     circularityErrorWorkaround();
   }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java

@@ -3,6 +3,7 @@
 import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStoreId;
 import static datadog.trace.util.Strings.getInternalName;
 
+import datadog.instrument.fieldinject.GlobalObjectStore;
 import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers;
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.api.Pair;
@@ -35,8 +36,8 @@ public final class FieldBackedContextInjector implements AsmVisitorWrapper {
 
   private static final Logger log = LoggerFactory.getLogger(FieldBackedContextInjector.class);
 
-  static final String FIELD_BACKED_CONTEXT_STORES_CLASS =
-      getInternalName(FieldBackedContextStores.class.getName());
+  static final String GLOBAL_OBJECT_STORE_CLASS =
+      getInternalName(GlobalObjectStore.class.getName());
 
   static final String FIELD_BACKED_CONTEXT_ACCESSOR_CLASS =
       getInternalName(FieldBackedContextAccessor.class.getName());
@@ -51,12 +52,12 @@ public final class FieldBackedContextInjector implements AsmVisitorWrapper {
   static final String PUTTER_METHOD_DESCRIPTOR =
       Type.getMethodDescriptor(Type.VOID_TYPE, Type.INT_TYPE, Type.getType(Object.class));
 
-  static final String WEAK_GET_METHOD = "weakGet";
+  static final String WEAK_GET_METHOD = "get";
   static final String WEAK_GET_METHOD_DESCRIPTOR =
       Type.getMethodDescriptor(
           Type.getType(Object.class), Type.getType(Object.class), Type.INT_TYPE);
 
-  static final String WEAK_PUT_METHOD = "weakPut";
+  static final String WEAK_PUT_METHOD = "put";
   static final String WEAK_PUT_METHOD_DESCRIPTOR =
       Type.getMethodDescriptor(
           Type.VOID_TYPE, Type.getType(Object.class), Type.INT_TYPE, Type.getType(Object.class));
@@ -445,7 +446,7 @@ private void invokeWeakGet(final MethodVisitor mv) {
         mv.visitIntInsn(Opcodes.ILOAD, 1);
         mv.visitMethodInsn(
             Opcodes.INVOKESTATIC,
-            FIELD_BACKED_CONTEXT_STORES_CLASS,
+            GLOBAL_OBJECT_STORE_CLASS,
             WEAK_GET_METHOD,
             WEAK_GET_METHOD_DESCRIPTOR,
             false);
@@ -458,7 +459,7 @@ private void invokeWeakPut(final MethodVisitor mv) {
         mv.visitIntInsn(Opcodes.ALOAD, 2);
         mv.visitMethodInsn(
             Opcodes.INVOKESTATIC,
-            FIELD_BACKED_CONTEXT_STORES_CLASS,
+            GLOBAL_OBJECT_STORE_CLASS,
             WEAK_PUT_METHOD,
             WEAK_PUT_METHOD_DESCRIPTOR,
             false);

gradle/libs.versions.toml

@@ -34,7 +34,7 @@ scala33 = "3.3.0"
 autoservice = "1.1.1"
 asm = "9.10.1"
 byte-buddy = "1.18.10"
-instrument-java = "0.0.4"
+instrument-java = "0.0.5-SNAPSHOT"
 
 # Benchmarks
 jmh = "1.37"