DataDog · mcculls · Oct 14, 2025 · Oct 15, 2025 · AlexeyKuznetsov-DD · Oct 14, 2025
@@ -0,0 +1,94 @@
+package datadog.trace.bootstrap;
+
+import static java.util.concurrent.TimeUnit.MICROSECONDS;
+import static org.openjdk.jmh.annotations.Mode.AverageTime;
+
+import datadog.trace.bootstrap.weakmap.WeakMapContextStore;
+import datadog.trace.bootstrap.weakmap.WeakMaps;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiFunction;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.infra.Blackhole;
+
+@State(Scope.Benchmark)
+@BenchmarkMode(AverageTime)
+@OutputTimeUnit(MICROSECONDS)
+@SuppressWarnings({"unused", "rawtypes", "unchecked"})
+public class WeakContextStoreBenchmark {
+  private static final int NUM_STORES = 3;
+
+  @Param({"false", "true"})
+  public boolean global;
+
+  private BiFunction[] stores;
+
+  private Map<String, String> kvs;
+
+  private AtomicInteger threadNumber;
+
+  @Setup(Level.Trial)
+  public void setup() {
+    WeakMaps.registerAsSupplier();
+
+    stores = new BiFunction[NUM_STORES];
+    for (int i = 0; i < NUM_STORES; i++) {
+      stores[i] = global ? globalWeakStorePutIfAbsent(i) : weakMapStorePutIfAbsent(i);
+    }
+
+    kvs = new HashMap<>();
+    for (int i = 0; i < 1_000; i++) {
+      kvs.put("key_" + i, "value_" + i);
+    }
+
+    threadNumber = new AtomicInteger();
+  }
+
+  @Benchmark
+  @Fork(value = 1)
+  @Threads(value = 1)
+  public void singleThreaded(Blackhole blackhole) {
+    test(blackhole);
+  }
+
+  @Benchmark
+  @Fork(value = 1)
+  @Threads(value = 10)
+  public void multiThreaded10(Blackhole blackhole) {
+    test(blackhole);
+  }
+
+  @Benchmark
+  @Fork(value = 1)
+  @Threads(value = 100)
+  public void multiThreaded100(Blackhole blackhole) {
+    test(blackhole);
+  }
+
+  private void test(Blackhole blackhole) {
+    // assign each benchmark thread a single store to operate on during the benchmark
+    // the number of concurrent requests to a store goes up as more threads are added
+    BiFunction store = stores[threadNumber.getAndIncrement() % NUM_STORES];
+    for (Map.Entry e : kvs.entrySet()) {
+      blackhole.consume(store.apply(e.getKey(), e.getValue()));
+    }
+  }
+
+  private static BiFunction globalWeakStorePutIfAbsent(int storeId) {
+    return (k, v) -> GlobalWeakContextStore.weakPutIfAbsent(k, storeId, v);
+  }
+
+  private static BiFunction weakMapStorePutIfAbsent(int storeId) {
+    return new WeakMapContextStore<>(storeId)::putIfAbsent;
+  }
+}
@@ -1,4 +1,4 @@
-package datadog.trace.bootstrap;
+package datadog.trace.bootstrap.weakmap;
 
 import java.util.function.Function;
 

@@ -1,11 +1,13 @@
-package datadog.trace.bootstrap;
+package datadog.trace.bootstrap.weakmap;
+
+import datadog.trace.bootstrap.ContextStore;
 
 /**
  * Weak {@link ContextStore} that acts as a fall-back when field-injection isn't possible.
  *
  * <p>This class should be created lazily because it uses weak maps with background cleanup.
  */
-final class WeakMapContextStore<K, V> implements ContextStore<K, V> {
+public final class WeakMapContextStore<K, V> implements ContextStore<K, V> {
   private static final int DEFAULT_MAX_SIZE = 50_000;
 
   private final int maxSize;

@@ -1,8 +1,7 @@
-package datadog.trace.agent.tooling;
+package datadog.trace.bootstrap.weakmap;
 
 import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import datadog.trace.api.Platform;
-import datadog.trace.bootstrap.WeakMap;
 import datadog.trace.util.AgentTaskScheduler;
 import datadog.trace.util.AgentTaskScheduler.Task;
 import java.util.concurrent.TimeUnit;

@@ -2,7 +2,7 @@
 
 /**
  * {@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.
+ * fields. Delegates to the global weak map for keys that don't have a field for this store.
  */
 public final class FieldBackedContextStore implements ContextStore<Object, Object> {
   final int storeId;
@@ -16,7 +16,7 @@ public Object get(final Object key) {
     if (key instanceof FieldBackedContextAccessor) {
       return ((FieldBackedContextAccessor) key).$get$__datadogContext$(storeId);
     } else {
-      return weakStore().get(key);
+      return GlobalWeakContextStore.weakGet(key, storeId);
     }
   }
 
@@ -25,7 +25,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);
+      GlobalWeakContextStore.weakPut(key, storeId, context);
     }
   }
 
@@ -45,7 +45,7 @@ public Object putIfAbsent(final Object key, final Object context) {
       }
       return existingContext;
     } else {
-      return weakStore().putIfAbsent(key, context);
+      return GlobalWeakContextStore.weakPutIfAbsent(key, storeId, context);
     }
   }
 
@@ -71,7 +71,7 @@ public Object computeIfAbsent(
       }
       return existingContext;
     } else {
-      return weakStore().computeIfAbsent(key, contextFactory);
+      return GlobalWeakContextStore.weakComputeIfAbsent(key, storeId, contextFactory);
     }
   }
 
@@ -90,22 +90,7 @@ public Object remove(Object key) {
       }
       return existingContext;
     } else {
-      return weakStore().remove(key);
+      return GlobalWeakContextStore.weakRemove(key, storeId);
     }
   }
-
-  // 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 weakStore;
-  }
 }
@@ -134,14 +134,4 @@ private static FieldBackedContextStore createStore(final int storeId) {
     }
     return store;
   }
-
-  /** 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);
-  }
-
-  /** 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);
-  }
 }
@@ -0,0 +1,172 @@
+package datadog.trace.bootstrap;
+
+import datadog.trace.api.Platform;
+import datadog.trace.bootstrap.ContextStore.KeyAwareFactory;
+import datadog.trace.util.AgentTaskScheduler;
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+import java.lang.ref.Reference;
+import java.lang.ref.ReferenceQueue;
+import java.lang.ref.WeakReference;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Global weak {@link ContextStore} that acts as a fall-back when field-injection isn't possible.
+ */
+public final class GlobalWeakContextStore {
+
+  // global map of weak (key + store-id) wrappers mapped to context values
+  private static final Map<Object, Object> globalMap = new ConcurrentHashMap<>();
+
+  // stale key wrappers that are now eligible for collection
+  private static final ReferenceQueue<Object> staleKeys = new ReferenceQueue<>();
+
+  private static final long CLEAN_FREQUENCY_SECONDS = 1;
+
+  private static final int MAX_KEYS_CLEANED_PER_CYCLE = 1_000;
+
+  static {
+    if (!Platform.isNativeImageBuilder()) {
+      AgentTaskScheduler.get()
+          .scheduleAtFixedRate(
+              GlobalWeakContextStore::cleanStaleKeys,
+              CLEAN_FREQUENCY_SECONDS,
+              CLEAN_FREQUENCY_SECONDS,
+              TimeUnit.SECONDS);
+    }
+  }
+
+  /** Checks for stale key wrappers and removes them from the global map. */
+  static void cleanStaleKeys() {
+    int count = 0;
+    Reference<?> ref;
+    while ((ref = staleKeys.poll()) != null) {
+      globalMap.remove(ref);
+      if (++count >= MAX_KEYS_CLEANED_PER_CYCLE) {
+        break; // limit work done per call
+      }
+    }
+  }
+
+  private GlobalWeakContextStore() {}
+
+  public static Object weakGet(Object key, int storeId) {
+    return globalMap.get(new LookupKey(key, storeId));
+  }
+
+  public static void weakPut(Object key, int storeId, Object context) {
+    if (context != null) {
+      globalMap.put(new StoreKey(key, storeId), context);
+    } else {
+      globalMap.remove(new LookupKey(key, storeId));
+    }
+  }
+
+  public static Object weakPutIfAbsent(Object key, int storeId, Object context) {
+    LookupKey lookupKey = new LookupKey(key, storeId);
+    Object existing;
+    if (null == (existing = globalMap.get(lookupKey))) {
+      // This whole part with using synchronized is only because
+      // we want to avoid prematurely calling the factory if
+      // someone else is doing a putIfAbsent at the same time.
+      // There is still the possibility that there is a concurrent
+      // call to put that will win, but that is indistinguishable
+      // from the put happening right after the putIfAbsent.
+      synchronized (key) {
+        if (null == (existing = globalMap.get(lookupKey))) {
+          weakPut(key, storeId, existing = context);
+        }
+      }
+    }
+    return existing;
+  }
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  public static Object weakComputeIfAbsent(
+      Object key, int storeId, KeyAwareFactory contextFactory) {
+    LookupKey lookupKey = new LookupKey(key, storeId);
+    Object existing;
+    if (null == (existing = globalMap.get(lookupKey))) {
+      // This whole part with using synchronized is only because
+      // we want to avoid prematurely calling the factory if
+      // someone else is doing a putIfAbsent at the same time.
+      // There is still the possibility that there is a concurrent
+      // call to put that will win, but that is indistinguishable
+      // from the put happening right after the putIfAbsent.
+      synchronized (key) {
+        if (null == (existing = globalMap.get(lookupKey))) {
+          weakPut(key, storeId, existing = contextFactory.create(key));
+        }
+      }
+    }
+    return existing;
+  }
+
+  public static Object weakRemove(Object key, int storeId) {
+    return globalMap.remove(new LookupKey(key, storeId));
+  }
+
+  /** Reference key used to weakly associate a key and store-id with a context value. */
+  static final class StoreKey extends WeakReference<Object> {
+    final int hash;
+    final int storeId;
+
+    StoreKey(Object key, int storeId) {
+      super(key, staleKeys);
+      this.hash = (31 * storeId) + System.identityHashCode(key);
+      this.storeId = storeId;
+    }
+
+    @Override
+    public int hashCode() {
+      return hash;
+    }
+
+    @Override
+    @SuppressFBWarnings("Eq") // symmetric because it mirrors LookupKey.equals
+    public boolean equals(Object o) {
+      if (o instanceof LookupKey) {
+        LookupKey lookupKey = (LookupKey) o;
+        return storeId == lookupKey.storeId && get() == lookupKey.key;
+      } else if (o instanceof StoreKey) {
+        StoreKey storeKey = (StoreKey) o;
+        return storeId == storeKey.storeId && get() == storeKey.get();
+      } else {
+        return false;
+      }
+    }
+  }
+
+  /** Temporary key used for lookup purposes without the reference tracking overhead. */
+  static final class LookupKey {
+    final Object key;
+    final int hash;
+    final int storeId;
+
+    LookupKey(Object key, int storeId) {
+      this.key = key;
+      this.hash = (31 * storeId) + System.identityHashCode(key);
+      this.storeId = storeId;
+    }
+
+    @Override
+    public int hashCode() {
+      return hash;
+    }
+
+    @Override
+    @SuppressFBWarnings("Eq") // symmetric because it mirrors StoreKey.equals
+    public boolean equals(Object o) {
+      if (o instanceof StoreKey) {
+        StoreKey storeKey = (StoreKey) o;
+        return storeId == storeKey.storeId && key == storeKey.get();
+      } else if (o instanceof LookupKey) {
+        LookupKey lookupKey = (LookupKey) o;
+        return storeId == lookupKey.storeId && key == lookupKey.key;
+      } else {
+        return false;
+      }
+    }
+  }
+}
@@ -56,8 +56,6 @@ public class AgentInstaller {
   static {
     addByteBuddyRawSetting();
     disableByteBuddyNexus();
-    // register weak map supplier as early as possible
-    WeakMaps.registerAsSupplier();
     circularityErrorWorkaround();
   }