WIP: Support using GlobalObjectStore for non-field-injected contexts

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

@@ -1,10 +1,17 @@
 package datadog.trace.bootstrap;
 
+import datadog.instrument.fieldinject.GlobalObjectStore;
+import datadog.trace.api.InstrumenterConfig;
+import java.util.function.Function;
+
 /**
  * {@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.
  */
 public final class FieldBackedContextStore implements ContextStore<Object, Object> {
+  private static final boolean MAP_PER_STORE =
+      InstrumenterConfig.get().isRuntimeContextMapPerStore();
+
   final int storeId;
 
   FieldBackedContextStore(final int storeId) {
@@ -15,17 +22,21 @@ public final class FieldBackedContextStore implements ContextStore<Object, Objec
   public Object get(final Object key) {
     if (key instanceof FieldBackedContextAccessor) {
       return ((FieldBackedContextAccessor) key).$get$__datadogContext$(storeId);
-    } else {
+    } else if (MAP_PER_STORE) {
       return weakStore().get(key);
+    } else {
+      return GlobalObjectStore.get(key, storeId);
     }
   }
 
   @Override
   public void put(final Object key, final Object context) {
     if (key instanceof FieldBackedContextAccessor) {
       ((FieldBackedContextAccessor) key).$put$__datadogContext$(storeId, context);
-    } else {
+    } else if (MAP_PER_STORE) {
       weakStore().put(key, context);
+    } else {
+      GlobalObjectStore.put(key, storeId, context);
     }
   }
 
@@ -44,28 +55,32 @@ public Object getOrPut(final Object key, final Object context) {
         }
       }
       return existingContext;
-    } else {
+    } else if (MAP_PER_STORE) {
       return weakStore().getOrPut(key, context);
+    } else {
+      return GlobalObjectStore.getOrPut(key, storeId, context);
     }
   }
 
   @Override
-  public Object getOrCompute(Object key, KeyAwareFactory<? super Object, Object> contextFactory) {
+  public Object getOrCompute(Object key, Function<? super Object, Object> contextFactory) {
     if (key instanceof FieldBackedContextAccessor) {
       final FieldBackedContextAccessor accessor = (FieldBackedContextAccessor) key;
       Object existingContext = accessor.$get$__datadogContext$(storeId);
       if (null == existingContext) {
         synchronized (accessor) {
           existingContext = accessor.$get$__datadogContext$(storeId);
           if (null == existingContext) {
-            existingContext = contextFactory.create(key);
+            existingContext = contextFactory.apply(key);
             accessor.$put$__datadogContext$(storeId, existingContext);
           }
         }
       }
       return existingContext;
-    } else {
+    } else if (MAP_PER_STORE) {
       return weakStore().getOrCompute(key, contextFactory);
+    } else {
+      return GlobalObjectStore.getOrCompute(key, storeId, contextFactory);
     }
   }
 
@@ -83,8 +98,10 @@ public Object remove(Object key) {
         }
       }
       return existingContext;
-    } else {
+    } else if (MAP_PER_STORE) {
       return weakStore().remove(key);
+    } else {
+      return GlobalObjectStore.remove(key, storeId);
     }
   }
 

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

@@ -3,7 +3,7 @@
 import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStore;
 
 import datadog.trace.api.internal.VisibleForTesting;
-import datadog.trace.bootstrap.ContextStore.KeyAwareFactory;
+import java.util.function.Function;
 
 /**
  * Weak "map-per-store" fall-back to track contexts when field-injection isn't possible.
@@ -59,7 +59,7 @@ V getOrPut(final K key, final V context) {
     return existingContext;
   }
 
-  V getOrCompute(K key, KeyAwareFactory<? super K, V> contextFactory) {
+  V getOrCompute(K key, Function<? super K, V> contextFactory) {
     V existingContext = get(key);
     if (null == existingContext) {
       // This whole part with using synchronized is only because
@@ -71,7 +71,7 @@ V getOrCompute(K key, KeyAwareFactory<? super K, V> contextFactory) {
       synchronized (map) {
         existingContext = get(key);
         if (null == existingContext) {
-          existingContext = contextFactory.create(key);
+          existingContext = contextFactory.apply(key);
           put(key, existingContext);
         }
       }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java

@@ -3,6 +3,7 @@
 import datadog.trace.bootstrap.ContextStore;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 
 /** Substitute {@link ContextStore} that uses strong-references to track contexts. */
 public class StrongMapContextStore<K, C> implements ContextStore<K, C> {
@@ -25,8 +26,8 @@ public C getOrPut(K key, C context) {
   }
 
   @Override
-  public C getOrCompute(K key, KeyAwareFactory<? super K, C> contextFactory) {
-    return map.computeIfAbsent(key, contextFactory::create);
+  public C getOrCompute(K key, Function<? super K, C> contextFactory) {
+    return map.computeIfAbsent(key, contextFactory);
   }
 
   @Override

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;
@@ -48,7 +49,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_REDIRECT_CLASS = getInternalName(WeakMapPerStore.class.getName());
+  static final String WEAK_REDIRECT_CLASS =
+      getInternalName(
+          (InstrumenterConfig.get().isRuntimeContextMapPerStore()
+                  ? WeakMapPerStore.class
+                  : GlobalObjectStore.class)
+              .getName());
 
   static final String WEAK_GET_METHOD_DESCRIPTOR =
       Type.getMethodDescriptor(

dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java

@@ -175,6 +175,8 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
               + "datadog.slf4j.helpers.SubstituteLoggerFactory:build_time,"
               + "datadog.slf4j.impl.StaticLoggerBinder:build_time,"
               + "datadog.slf4j.LoggerFactory:build_time,"
+              + "datadog.instrument.fieldinject.GlobalObjectStore:build_time,"
+              + "datadog.instrument.fieldinject.GlobalObjectStore$LookupKey:build_time,"
               + "com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap:build_time,"
               + "com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap$1:build_time,"
               + "net.bytebuddy:build_time,"

dd-java-agent/instrumentation/reactive-streams-1.0/src/test/java/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsContextPropagationTest.java

@@ -11,6 +11,7 @@
 import datadog.trace.bootstrap.instrumentation.reactivestreams.HandoffContext;
 import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.function.Function;
 import org.junit.jupiter.api.Test;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
@@ -176,23 +177,13 @@ public void put(final K key, final C context) {
 
     @Override
     public C getOrPut(final K key, final C context) {
-      final C existing = map.get(key);
-      if (existing != null) {
-        return existing;
-      }
-      map.put(key, context);
-      return context;
+      final C existing = map.putIfAbsent(key, context);
+      return existing != null ? existing : context;
     }
 
     @Override
-    public C getOrCompute(final K key, final KeyAwareFactory<? super K, C> contextFactory) {
-      final C existing = map.get(key);
-      if (existing != null) {
-        return existing;
-      }
-      final C created = contextFactory.create(key);
-      map.put(key, created);
-      return created;
+    public C getOrCompute(final K key, final Function<? super K, C> contextFactory) {
+      return map.computeIfAbsent(key, contextFactory);
     }
 
     @Override

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -60,6 +60,7 @@ public final class ConfigDefaults {
 
   static final boolean DEFAULT_RUNTIME_CONTEXT_FIELD_INJECTION = true;
   static final boolean DEFAULT_SERIALVERSIONUID_FIELD_INJECTION = true;
+  static final boolean DEFAULT_RUNTIME_CONTEXT_MAP_PER_STORE = true;
 
   static final boolean DEFAULT_EXPERIMENTATAL_JEE_SPLIT_BY_DEPLOYMENT = false;
   static final boolean DEFAULT_PRIORITY_SAMPLING_ENABLED = true;

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -99,6 +99,7 @@ public final class TraceInstrumentationConfig {
       "trace.runtime.context.field.injection";
   public static final String SERIALVERSIONUID_FIELD_INJECTION =
       "trace.serialversionuid.field.injection";
+  public static final String RUNTIME_CONTEXT_MAP_PER_STORE = "trace.runtime.context.map-per-store";
 
   public static final String LOGS_INJECTION_ENABLED = "logs.injection.enabled";
   public static final String LOGS_INJECTION = "logs.injection";

internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java

@@ -17,6 +17,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_RESOLVER_RESET_INTERVAL;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_RUM_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_RUNTIME_CONTEXT_FIELD_INJECTION;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_RUNTIME_CONTEXT_MAP_PER_STORE;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_SERIALVERSIONUID_FIELD_INJECTION;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TELEMETRY_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_TRACE_ANNOTATIONS;
@@ -76,6 +77,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_USE_LOADCLASS;
 import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_USE_URL_CACHES;
 import static datadog.trace.api.config.TraceInstrumentationConfig.RUNTIME_CONTEXT_FIELD_INJECTION;
+import static datadog.trace.api.config.TraceInstrumentationConfig.RUNTIME_CONTEXT_MAP_PER_STORE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.SERIALVERSIONUID_FIELD_INJECTION;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_ANNOTATIONS;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_ANNOTATION_ASYNC;
@@ -208,6 +210,7 @@ public class InstrumenterConfig {
 
   private final boolean runtimeContextFieldInjection;
   private final boolean serialVersionUIDFieldInjection;
+  private final boolean runtimeContextMapPerStore;
 
   private final String traceAnnotations;
   private final boolean traceAnnotationAsync;
@@ -355,6 +358,9 @@ private InstrumenterConfig() {
     serialVersionUIDFieldInjection =
         configProvider.getBoolean(
             SERIALVERSIONUID_FIELD_INJECTION, DEFAULT_SERIALVERSIONUID_FIELD_INJECTION);
+    runtimeContextMapPerStore =
+        configProvider.getBoolean(
+            RUNTIME_CONTEXT_MAP_PER_STORE, DEFAULT_RUNTIME_CONTEXT_MAP_PER_STORE);
 
     instrumentationConfigId = configProvider.getString(INSTRUMENTATION_CONFIG_ID, "");
 
@@ -657,6 +663,10 @@ public boolean isSerialVersionUIDFieldInjection() {
     return serialVersionUIDFieldInjection;
   }
 
+  public boolean isRuntimeContextMapPerStore() {
+    return runtimeContextMapPerStore;
+  }
+
   public String getTraceAnnotations() {
     return traceAnnotations;
   }
@@ -835,6 +845,8 @@ public String toString() {
         + runtimeContextFieldInjection
         + ", serialVersionUIDFieldInjection="
         + serialVersionUIDFieldInjection
+        + ", runtimeContextMapPerStore="
+        + runtimeContextMapPerStore
         + ", codeOriginEnabled="
         + codeOriginEnabled
         + ", traceAnnotations='"

internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java

@@ -1,5 +1,6 @@
 package datadog.trace.bootstrap;
 
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /**
@@ -18,32 +19,18 @@ public interface ContextStore<K, C> {
    *
    * @param <C> context type
    */
-  interface Factory<C> extends KeyAwareFactory<Object, C> {
+  interface Factory<C> extends Function<Object, C> {
 
     /**
      * @return new context instance
      */
     C create();
 
-    default C create(Object key) {
+    default C apply(Object key) {
       return create();
     }
   }
 
-  /**
-   * Factory interface to create context instances using context key instances
-   *
-   * @param <K> context key type
-   * @param <C> context value type
-   */
-  interface KeyAwareFactory<K, C> {
-
-    /**
-     * @return new context instance
-     */
-    C create(K key);
-  }
-
   /**
    * Get context instance for the given key.
    *
@@ -91,7 +78,7 @@ default C getOrCreate(K key, Factory<C> contextFactory) {
    * @param contextFactory factory instance to produce new context instances
    * @return existing context instance if present; otherwise new instance
    */
-  C getOrCompute(K key, KeyAwareFactory<? super K, C> contextFactory);
+  C getOrCompute(K key, Function<? super K, C> contextFactory);
 
   /**
    * Removes the context instance for the given key.

metadata/supported-configurations.json

@@ -9780,6 +9780,14 @@
         "aliases": []
       }
     ],
+    "DD_TRACE_RUNTIME_CONTEXT_MAP_PER_STORE": [
+      {
+        "version": "A",
+        "type": "boolean",
+        "default": "true",
+        "aliases": []
+      }
+    ],
     "DD_TRACE_RXJAVA_ENABLED": [
       {
         "version": "A",