DataDog · gh-worker-dd-mergequeue-cf854d · Sep 1, 2026 · Aug 31, 2026 · Aug 31, 2026 · Aug 31, 2026
@@ -30,7 +30,7 @@ public void put(final Object key, final Object context) {
   }
 
   @Override
-  public Object putIfAbsent(final Object key, final Object context) {
+  public Object getOrPut(final Object key, final Object context) {
     if (key instanceof FieldBackedContextAccessor) {
       final FieldBackedContextAccessor accessor = (FieldBackedContextAccessor) key;
       Object existingContext = accessor.$get$__datadogContext$(storeId);
@@ -45,18 +45,12 @@ public Object putIfAbsent(final Object key, final Object context) {
       }
       return existingContext;
     } else {
-      return weakStore().putIfAbsent(key, context);
+      return weakStore().getOrPut(key, context);
     }
   }
 
   @Override
-  public Object putIfAbsent(final Object key, final Factory<Object> contextFactory) {
-    return computeIfAbsent(key, contextFactory);
-  }
-
-  @Override
-  public Object computeIfAbsent(
-      Object key, KeyAwareFactory<? super Object, Object> contextFactory) {
+  public Object getOrCompute(Object key, KeyAwareFactory<? super Object, Object> contextFactory) {
     if (key instanceof FieldBackedContextAccessor) {
       final FieldBackedContextAccessor accessor = (FieldBackedContextAccessor) key;
       Object existingContext = accessor.$get$__datadogContext$(storeId);
@@ -71,7 +65,7 @@ public Object computeIfAbsent(
       }
       return existingContext;
     } else {
-      return weakStore().computeIfAbsent(key, contextFactory);
+      return weakStore().getOrCompute(key, contextFactory);
     }
   }
 
@@ -95,14 +89,14 @@ public Object remove(Object key) {
   }
 
   // only create WeakMap-based fall-back when we need it
-  private volatile WeakMapContextStore<Object, Object> weakStore;
+  private volatile WeakMapPerStore<Object, Object> weakStore;
   private final Object synchronizationInstance = new Object();
 
-  WeakMapContextStore<Object, Object> weakStore() {
+  WeakMapPerStore<Object, Object> weakStore() {
     if (null == weakStore) {
       synchronized (synchronizationInstance) {
         if (null == weakStore) {
-          weakStore = new WeakMapContextStore<>();
+          weakStore = new WeakMapPerStore<>();
         }
       }
     }

@@ -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);
-  }
 }
@@ -5,6 +5,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Supplier;
+import javax.annotation.Nullable;
 
 /**
  * An {@code InstanceStore} is a class global map for registering instances. This can be useful when
@@ -36,8 +37,9 @@ private InstanceStore() {}
    * Gets the instance of {@code T} currently associated with the given key.
    *
    * @param key the instance key
-   * @return the associated instance
+   * @return the associated instance; {@code null} if there was none
    */
+  @Nullable
   public T get(String key) {
     return store.get(key);
   }
@@ -61,7 +63,7 @@ public void put(String key, T instance) {
    * @param instanceFactory the factory to create instances
    * @return final associated instance
    */
-  public T putIfAbsent(String key, Supplier<T> instanceFactory) {
+  public T getOrCreate(String key, Supplier<T> instanceFactory) {
     return store.computeIfAbsent(key, k -> instanceFactory.get());
   }
 
@@ -71,6 +73,7 @@ public T putIfAbsent(String key, Supplier<T> instanceFactory) {
    * @param key the instance key
    * @return the previously associated instance; {@code null} if there was none
    */
+  @Nullable
   public T remove(String key) {
     return store.remove(key);
   }

@@ -1,49 +1,53 @@
 package datadog.trace.bootstrap;
 
+import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStore;
+
 import datadog.trace.api.internal.VisibleForTesting;
+import datadog.trace.bootstrap.ContextStore.KeyAwareFactory;
 
 /**
- * Weak {@link ContextStore} that acts as a fall-back when field-injection isn't possible.
+ * Weak "map-per-store" fall-back to track contexts 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> {
-  private static final int DEFAULT_MAX_SIZE = 50_000;
-
-  private final int maxSize;
-  private final WeakMap<Object, Object> map = WeakMap.Supplier.newWeakMap();
+public final class WeakMapPerStore<K, V> {
 
-  public WeakMapContextStore(int maxSize) {
-    this.maxSize = maxSize;
+  /** Injection helper that immediately delegates to the weak-map for the given context store. */
+  public static Object get(final Object key, final int storeId) {
+    return getContextStore(storeId).weakStore().get(key);
   }
 
-  public WeakMapContextStore() {
-    this(DEFAULT_MAX_SIZE);
+  /** Injection helper that immediately delegates to the weak-map for the given context store. */
+  public static void put(final Object key, final int storeId, final Object context) {
+    getContextStore(storeId).weakStore().put(key, context);
   }
 
-  @Override
+  private static final int MAX_SIZE = 50_000;
+
+  private final WeakMap<Object, Object> map = WeakMap.Supplier.newWeakMap();
+
+  WeakMapPerStore() {}
+
   @SuppressWarnings("unchecked")
-  public V get(final K key) {
+  V get(final K key) {
     return (V) map.get(key);
   }
 
-  @Override
-  public void put(final K key, final V context) {
-    if (map.size() < maxSize) {
+  void put(final K key, final V context) {
+    if (map.size() < MAX_SIZE) {
       map.put(key, context);
     }
   }
 
-  @Override
-  public V putIfAbsent(final K key, final V context) {
+  V getOrPut(final K key, final V context) {
     V existingContext = get(key);
     if (null == existingContext) {
       // 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.
+      // someone else is doing a getOrPut 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.
+      // from the put happening right after the getOrPut.
       synchronized (map) {
         existingContext = get(key);
         if (null == existingContext) {
@@ -55,21 +59,15 @@ public V putIfAbsent(final K key, final V context) {
     return existingContext;
   }
 
-  @Override
-  public V putIfAbsent(final K key, final Factory<V> contextFactory) {
-    return computeIfAbsent(key, contextFactory);
-  }
-
-  @Override
-  public V computeIfAbsent(K key, KeyAwareFactory<? super K, V> contextFactory) {
+  V getOrCompute(K key, KeyAwareFactory<? super K, V> contextFactory) {
     V existingContext = get(key);
     if (null == existingContext) {
       // 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.
+      // someone else is doing a getOrCompute 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.
+      // from the put happening right after the getOrCompute.
       synchronized (map) {
         existingContext = get(key);
         if (null == existingContext) {
@@ -81,9 +79,8 @@ public V computeIfAbsent(K key, KeyAwareFactory<? super K, V> contextFactory) {
     return existingContext;
   }
 
-  @Override
   @SuppressWarnings("unchecked")
-  public V remove(final K key) {
+  V remove(final K key) {
     return (V) map.remove(key);
   }
 

@@ -36,7 +36,7 @@ private ConcurrentState() {}
   public static <K> ConcurrentState captureContinuation(
       ContextStore<K, ConcurrentState> contextStore, K key, Context context) {
     if (shouldCapture(context)) {
-      final ConcurrentState state = contextStore.putIfAbsent(key, FACTORY);
+      final ConcurrentState state = contextStore.getOrCreate(key, FACTORY);
       if (!state.captureAndSetContinuation(context) && log.isDebugEnabled()) {
         log.debug(
             "continuation was already set for {} in context {}, no continuation captured.",

@@ -41,7 +41,7 @@ public static boolean shouldAttachStateToTask(final Object task, final Context c
    */
   public static <T> State setupState(
       final ContextStore<T, State> contextStore, final T task, final Context context) {
-    final State state = contextStore.putIfAbsent(task, State.FACTORY);
+    final State state = contextStore.getOrCreate(task, State.FACTORY);
     if (!state.captureAndSetContinuation(context)) {
       log.debug(
           "continuation was already set for {} in context {}, no continuation captured.",

@@ -53,13 +53,13 @@ class InstanceStoreTest extends DDSpecification {
     InstanceStore.of(Some).put(key, some1)
 
     when:
-    def current = InstanceStore.of(Some).putIfAbsent(key, Some::new)
+    def current = InstanceStore.of(Some).getOrCreate(key, Some::new)
 
     then:
     current == some1
 
     when:
-    current = InstanceStore.of(Some).putIfAbsent(key, Some::new)
+    current = InstanceStore.of(Some).getOrCreate(key, Some::new)
 
     then:
     current == some1
@@ -71,13 +71,13 @@ class InstanceStoreTest extends DDSpecification {
     def key = nextKey()
 
     when:
-    def current = someStore.putIfAbsent(key, () -> some1)
+    def current = someStore.getOrCreate(key, () -> some1)
 
     then:
     current == some1
 
     when:
-    current = someStore.putIfAbsent(key, Some::new)
+    current = someStore.getOrCreate(key, Some::new)
 
     then:
     current == some1
@@ -91,7 +91,7 @@ class InstanceStoreTest extends DDSpecification {
     someStore.put(key, some1)
 
     when:
-    def current = someStore.putIfAbsent(key, new Creator(invocations))
+    def current = someStore.getOrCreate(key, new Creator(invocations))
 
     then:
     current == some1
@@ -105,14 +105,14 @@ class InstanceStoreTest extends DDSpecification {
     def key = nextKey()
 
     when:
-    def current = someStore.putIfAbsent(key, new Creator(invocations, some1))
+    def current = someStore.getOrCreate(key, new Creator(invocations, some1))
 
     then:
     current == some1
     invocations.get() == 1
 
     when:
-    current = someStore.putIfAbsent(key, new Creator(invocations))
+    current = someStore.getOrCreate(key, new Creator(invocations))
 
     then:
     current == some1

@@ -43,7 +43,7 @@ import datadog.trace.civisibility.source.SourcePathResolver
 import datadog.trace.civisibility.source.index.RepoIndexBuilder
 import datadog.trace.civisibility.telemetry.CiVisibilityMetricCollectorImpl
 import datadog.trace.civisibility.test.ExecutionStrategy
-import datadog.trace.civisibility.utils.ConcurrentHashMapContextStore
+import datadog.trace.civisibility.utils.StrongMapContextStore
 import datadog.trace.civisibility.writer.ddintake.CiTestCovMapperV2
 import datadog.trace.civisibility.writer.ddintake.CiTestCycleMapperV1
 import datadog.trace.common.writer.ListWriter
@@ -269,8 +269,8 @@ abstract class CiVisibilityInstrumentationTest extends InstrumentationSpecificat
       { testFrameworkSessionFactory.startSession(moduleName, component, null, capabilities) },
       moduleName,
       false,
-      suiteStore != null ? suiteStore : new ConcurrentHashMapContextStore<>(),
-      testStore != null ? testStore : new ConcurrentHashMapContextStore<>())
+      suiteStore != null ? suiteStore : new StrongMapContextStore<>(),
+      testStore != null ? testStore : new StrongMapContextStore<>())
     }
   }
 

@@ -38,7 +38,7 @@
 import datadog.trace.civisibility.source.index.RepoIndex;
 import datadog.trace.civisibility.telemetry.CiVisibilityMetricCollectorImpl;
 import datadog.trace.civisibility.test.ExecutionStrategy;
-import datadog.trace.civisibility.utils.ConcurrentHashMapContextStore;
+import datadog.trace.civisibility.utils.StrongMapContextStore;
 import datadog.trace.util.throwable.FatalAgentMisconfigurationError;
 import java.lang.instrument.Instrumentation;
 import java.nio.file.Path;
@@ -198,8 +198,8 @@ public <SuiteKey, TestKey> TestEventsHandler<SuiteKey, TestKey> create(
                       repoServices.moduleName, component, null, capabilities),
               repoServices.moduleName,
               eagerSessionStart,
-              suiteStore != null ? suiteStore : new ConcurrentHashMapContextStore<>(),
-              testStore != null ? testStore : new ConcurrentHashMapContextStore<>());
+              suiteStore != null ? suiteStore : new StrongMapContextStore<>(),
+              testStore != null ? testStore : new StrongMapContextStore<>());
       handlers.add(handler);
       return handler;
     }