From 8c6eb3678fba93b46558d25901fd4f6ecf12eb98 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 2 Sep 2026 12:03:14 +0100 Subject: [PATCH 1/3] Replace KeyAwareFactory with standard Function type --- .../bootstrap/FieldBackedContextStore.java | 6 ++++-- .../trace/bootstrap/WeakMapPerStore.java | 6 +++--- .../utils/StrongMapContextStore.java | 5 +++-- ...ReactiveStreamsContextPropagationTest.java | 19 +++++------------ .../datadog/trace/bootstrap/ContextStore.java | 21 ++++--------------- 5 files changed, 19 insertions(+), 38 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java index 7d334d5405b..37e4210f100 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java @@ -1,5 +1,7 @@ package datadog.trace.bootstrap; +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. @@ -50,7 +52,7 @@ public Object getOrPut(final Object key, final Object context) { } @Override - public Object getOrCompute(Object key, KeyAwareFactory contextFactory) { + public Object getOrCompute(Object key, Function contextFactory) { if (key instanceof FieldBackedContextAccessor) { final FieldBackedContextAccessor accessor = (FieldBackedContextAccessor) key; Object existingContext = accessor.$get$__datadogContext$(storeId); @@ -58,7 +60,7 @@ public Object getOrCompute(Object key, KeyAwareFactory c synchronized (accessor) { existingContext = accessor.$get$__datadogContext$(storeId); if (null == existingContext) { - existingContext = contextFactory.create(key); + existingContext = contextFactory.apply(key); accessor.$put$__datadogContext$(storeId, existingContext); } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java index 36917246b1f..da184317865 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java +++ b/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 contextFactory) { + V getOrCompute(K key, Function 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 contextFactory) { synchronized (map) { existingContext = get(key); if (null == existingContext) { - existingContext = contextFactory.create(key); + existingContext = contextFactory.apply(key); put(key, existingContext); } } diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java index 2ca2586ce41..4cffe0ffb36 100644 --- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java +++ b/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 implements ContextStore { @@ -25,8 +26,8 @@ public C getOrPut(K key, C context) { } @Override - public C getOrCompute(K key, KeyAwareFactory contextFactory) { - return map.computeIfAbsent(key, contextFactory::create); + public C getOrCompute(K key, Function contextFactory) { + return map.computeIfAbsent(key, contextFactory); } @Override diff --git a/dd-java-agent/instrumentation/reactive-streams-1.0/src/test/java/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsContextPropagationTest.java b/dd-java-agent/instrumentation/reactive-streams-1.0/src/test/java/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsContextPropagationTest.java index 139464eaa2d..ab99ae9c3ae 100644 --- a/dd-java-agent/instrumentation/reactive-streams-1.0/src/test/java/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsContextPropagationTest.java +++ b/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 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 contextFactory) { + return map.computeIfAbsent(key, contextFactory); } @Override diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java b/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java index 6d22bbe92ad..18dd8344a5c 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java +++ b/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 { * * @param context type */ - interface Factory extends KeyAwareFactory { + interface Factory extends Function { /** * @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 context key type - * @param context value type - */ - interface KeyAwareFactory { - - /** - * @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 contextFactory) { * @param contextFactory factory instance to produce new context instances * @return existing context instance if present; otherwise new instance */ - C getOrCompute(K key, KeyAwareFactory contextFactory); + C getOrCompute(K key, Function contextFactory); /** * Removes the context instance for the given key. From c474b92ad9147513e897792234f4cdc65965a98b Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 2 Sep 2026 12:17:46 +0100 Subject: [PATCH 2/3] Support using GlobalObjectStore for non-field-injected contexts --- .../bootstrap/FieldBackedContextStore.java | 25 +++++++++++++++---- .../trace/agent/tooling/AgentInstaller.java | 12 +++++++++ .../context/FieldBackedContextInjector.java | 8 +++++- ...veImageGeneratorRunnerInstrumentation.java | 2 ++ .../datadog/trace/api/ConfigDefaults.java | 1 + .../config/TraceInstrumentationConfig.java | 1 + .../datadog/trace/api/InstrumenterConfig.java | 12 +++++++++ metadata/supported-configurations.json | 8 ++++++ 8 files changed, 63 insertions(+), 6 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java index 37e4210f100..aeafbe2f2a9 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStore.java +++ b/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; +import datadog.trace.api.InstrumenterConfig; import java.util.function.Function; /** @@ -7,6 +9,9 @@ * fields. Delegates to a lazy {@link WeakMap} for keys that don't have a field for this store. */ public final class FieldBackedContextStore implements ContextStore { + private static final boolean MAP_PER_STORE = + InstrumenterConfig.get().isRuntimeContextMapPerStore(); + final int storeId; FieldBackedContextStore(final int storeId) { @@ -17,8 +22,10 @@ public final class FieldBackedContextStore implements ContextStore contextF } } return existingContext; - } else { + } else if (MAP_PER_STORE) { return weakStore().getOrCompute(key, contextFactory); + } else { + return GlobalObjectStore.getOrCompute(key, storeId, contextFactory); } } @@ -85,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); } } diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index 3a8c7065362..db69109651a 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/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; @@ -55,6 +56,8 @@ public class AgentInstaller { private static final List LOG_MANAGER_CALLBACKS = new CopyOnWriteArrayList<>(); private static final List MBEAN_SERVER_BUILDER_CALLBACKS = new CopyOnWriteArrayList<>(); + private static final long GLOBAL_OBJECT_STORE_CLEAN_FREQUENCY_SECONDS = 1; + static { enableByteBuddyRawTypes(); disableByteBuddyNexus(); @@ -251,6 +254,15 @@ public void applied(Iterable instrumentationNames) { }); } + if (!InstrumenterConfig.get().isRuntimeContextMapPerStore()) { + AgentTaskScheduler.get() + .scheduleAtFixedRate( + GlobalObjectStore::removeStaleEntries, + GLOBAL_OBJECT_STORE_CLEAN_FREQUENCY_SECONDS, + GLOBAL_OBJECT_STORE_CLEAN_FREQUENCY_SECONDS, + TimeUnit.SECONDS); + } + InstrumenterState.resetDefaultState(); try { return transformerBuilder.installOn(inst); diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java index ee8ae19ca09..9f3b0a51dce 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedContextInjector.java +++ b/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( diff --git a/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java b/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java index d32db6b803f..efde0f6e533 100644 --- a/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java +++ b/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," diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java index 7fe08a8e121..20bb7b1740a 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java +++ b/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; diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java index 855a7243caf..a056c006d61 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java +++ b/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"; diff --git a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java index 74bff640024..84b407a6836 100644 --- a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java +++ b/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='" diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 6c3fe354b68..a15f6205295 100644 --- a/metadata/supported-configurations.json +++ b/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", From edd9ecf210452f19f798ae814e1de1938861f825 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Thu, 3 Sep 2026 17:30:59 +0100 Subject: [PATCH 3/3] Remove leftover reference to KeyAwareFactory --- .../nativeimage/NativeImageGeneratorRunnerInstrumentation.java | 1 - 1 file changed, 1 deletion(-) diff --git a/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java b/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java index efde0f6e533..c4c47e02233 100644 --- a/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java +++ b/dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java @@ -124,7 +124,6 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[ + "datadog.trace.bootstrap.CallDepthThreadLocalMap:build_time," + "datadog.trace.bootstrap.CallDepthThreadLocalMap$ThreadLocalDepth:build_time," + "datadog.trace.bootstrap.ContextStore$Factory:build_time," - + "datadog.trace.bootstrap.ContextStore$KeyAwareFactory:build_time," + "datadog.trace.bootstrap.DatadogClassLoader:build_time," + "datadog.trace.bootstrap.InstrumentationClassLoader:build_time," + "datadog.trace.bootstrap.FieldBackedContextStores:build_time,"