From a536cb18d7722c1c3d9f54cf95f40afc46d84bb1 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Aug 2026 17:31:59 +0100 Subject: [PATCH 1/7] Rename ContextStore methods to align with new ObjectStore: * putIfAbsent(value) to getOrPut * putIfAbsent(factory) to getOrCreate and default this to delegate to getOrCompute * computeIfAbsent(function) to getOrCompute --- .../trace/bootstrap/FieldBackedContextStore.java | 14 ++++---------- .../datadog/trace/bootstrap/InstanceStore.java | 2 +- .../trace/bootstrap/WeakMapContextStore.java | 9 ++------- .../java/concurrent/ConcurrentState.java | 2 +- .../concurrent/ExecutorInstrumentationUtils.java | 2 +- .../trace/bootstrap/InstanceStoreTest.groovy | 14 +++++++------- .../utils/ConcurrentHashMapContextStore.java | 9 ++------- .../FieldInjectionTestInstrumentation.java | 4 ++-- .../aws/v2/dynamodb/DynamoDbInterceptor.java | 2 +- .../v2/eventbridge/EventBridgeInterceptor.java | 2 +- .../instrumentation/aws/v2/s3/S3Interceptor.java | 2 +- .../aws/v2/AwsSdkClientDecorator.java | 2 +- .../aws/v2/TracingExecutionInterceptor.java | 2 +- .../aws/v2/sfn/SfnInterceptor.java | 2 +- .../aws/v2/sns/SnsInterceptor.java | 2 +- .../aws/v2/sqs/SqsInterceptor.java | 2 +- .../hibernate/SessionMethodUtils.java | 2 +- .../core/v3_3/SessionFactoryInstrumentation.java | 4 ++-- .../core/v4_0/SessionFactoryInstrumentation.java | 2 +- .../net/HttpUrlConnectionInstrumentation.java | 2 +- .../jms/SessionInstrumentation.java | 6 ++---- .../KafkaConsumerInfoInstrumentation.java | 4 ++-- .../KafkaProducerInstrumentation.java | 2 +- .../kafka_clients/MetadataInstrumentation.java | 4 ++-- .../kafka_clients38/ConstructorAdvice.java | 2 +- .../kafka_clients38/LegacyConstructorAdvice.java | 2 +- .../MetadataUpdate22AndAfterAdvice.java | 2 +- .../MetadataUpdateBefore22Advice.java | 2 +- .../ProducerConstructorAdvice.java | 2 +- .../karate/KarateExecutionInstrumentation.java | 2 +- .../karate2/KarateScenarioAdvice.java | 2 +- .../mongo/MongoCommandListener.java | 2 +- .../ChannelFutureListenerInstrumentation.java | 2 +- .../netty38/NettyChannelInstrumentation.java | 2 +- .../client/HttpClientRequestTracingHandler.java | 2 +- .../client/HttpClientResponseTracingHandler.java | 4 ++-- .../server/HttpServerRequestTracingHandler.java | 2 +- .../server/HttpServerResponseTracingHandler.java | 2 +- .../netty38/server/MaybeBlockResponseHandler.java | 2 +- .../ReactiveStreamsContextPropagation.java | 2 +- .../ReactiveStreamsContextPropagationTest.java | 15 ++------------- .../reactor/core/ReactorContextBridge.java | 2 +- .../resilience4j/ReactorHelper.java | 2 +- .../scalatest/ScalatestSkipInstrumentation.java | 2 +- .../tibcobw5/JobInstrumentation.java | 2 +- .../tomcat/WsHandshakeRequestInstrumentation.java | 3 +-- .../undertow/UndertowDecorator.java | 4 ++-- .../websocket/jsr256/EndpointInstrumentation.java | 2 +- .../websocket/jetty10/MethodHandleWrappers.java | 2 +- internal-api/build.gradle.kts | 1 + .../datadog/trace/bootstrap/ContextStore.java | 8 +++++--- 51 files changed, 74 insertions(+), 101 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 896342058a2..f60f890fa17 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 @@ -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 contextFactory) { - return computeIfAbsent(key, contextFactory); - } - - @Override - public Object computeIfAbsent( - Object key, KeyAwareFactory contextFactory) { + public Object getOrCompute(Object key, KeyAwareFactory 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); } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java index 765f0df6a6c..14e091774e2 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java @@ -61,7 +61,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 instanceFactory) { + public T getOrCreate(String key, Supplier instanceFactory) { return store.computeIfAbsent(key, k -> instanceFactory.get()); } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java index 1750b7f2b9c..6485966d913 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java @@ -35,7 +35,7 @@ public void put(final K key, final V context) { } @Override - public V putIfAbsent(final K key, final V context) { + public V getOrPut(final K key, final V context) { V existingContext = get(key); if (null == existingContext) { // This whole part with using synchronized is only because @@ -56,12 +56,7 @@ public V putIfAbsent(final K key, final V context) { } @Override - public V putIfAbsent(final K key, final Factory contextFactory) { - return computeIfAbsent(key, contextFactory); - } - - @Override - public V computeIfAbsent(K key, KeyAwareFactory contextFactory) { + public V getOrCompute(K key, KeyAwareFactory contextFactory) { V existingContext = get(key); if (null == existingContext) { // This whole part with using synchronized is only because diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java index 0a2477a8b44..6767c55ac77 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java @@ -36,7 +36,7 @@ private ConcurrentState() {} public static ConcurrentState captureContinuation( ContextStore 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.", diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java index 43ca4fdf7bf..cb2d682aa75 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java @@ -41,7 +41,7 @@ public static boolean shouldAttachStateToTask(final Object task, final Context c */ public static State setupState( final ContextStore 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.", diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy index eb2dbdbf1db..6b9a3fee93f 100644 --- a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy +++ b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/InstanceStoreTest.groovy @@ -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 diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java index 19a702c959f..606707bfb12 100644 --- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java +++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java @@ -19,17 +19,12 @@ public void put(K key, C context) { } @Override - public C putIfAbsent(K key, C context) { + public C getOrPut(K key, C context) { return m.computeIfAbsent(key, k -> context); } @Override - public C putIfAbsent(K key, Factory contextFactory) { - return m.computeIfAbsent(key, k -> contextFactory.create()); - } - - @Override - public C computeIfAbsent(K key, KeyAwareFactory contextFactory) { + public C getOrCompute(K key, KeyAwareFactory contextFactory) { return m.computeIfAbsent(key, contextFactory::create); } diff --git a/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java b/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java index 192ac18e374..2a8dcb46fe3 100644 --- a/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java +++ b/dd-java-agent/instrumentation-testing/src/test/java/context/FieldInjectionTestInstrumentation.java @@ -82,7 +82,7 @@ public static void methodExit( @Advice.This final KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) { final ContextStore contextStore = InstrumentationContext.get(KeyClass.class, Context.class); - final Context context = contextStore.putIfAbsent(thiz, new Context()); + final Context context = contextStore.getOrPut(thiz, new Context()); contextCount = ++context.count; } } @@ -93,7 +93,7 @@ public static void methodExit( @Advice.This final KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) { final ContextStore contextStore = InstrumentationContext.get(KeyClass.class, Context.class); - final Context context = contextStore.putIfAbsent(thiz, Context::new); + final Context context = contextStore.getOrCreate(thiz, Context::new); contextCount = ++context.count; } } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java index 3468915c7f1..a192f3d38a1 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-dynamodb-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/dynamodb/DynamoDbInterceptor.java @@ -23,7 +23,7 @@ public class DynamoDbInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws"); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java index 7cb54b62230..3672b49b88d 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java @@ -30,7 +30,7 @@ public class EventBridgeInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final String START_TIME_KEY = "x-datadog-start-time"; private static final String RESOURCE_NAME_KEY = "x-datadog-resource-name"; diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java index e03b6bb2427..4ab9a591572 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-s3-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/s3/S3Interceptor.java @@ -21,7 +21,7 @@ public class S3Interceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final boolean CAN_ADD_SPAN_POINTERS = Config.get().isAddSpanPointers("aws"); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java index f84a2dfa7e1..7c1a1a85b90 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java @@ -94,7 +94,7 @@ public class AwsSdkClientDecorator extends HttpClientDecorator KINESIS_STREAM_ARN_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("KinesisStreamArn", () -> new ExecutionAttribute<>("KinesisStreamArn")); + .getOrCreate("KinesisStreamArn", () -> new ExecutionAttribute<>("KinesisStreamArn")); // not static because this object would be ClassLoader specific if multiple SDK instances were // loaded by different loaders diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java index 980f08ca118..366c67675e3 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java @@ -39,7 +39,7 @@ public class TracingExecutionInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private static final Logger log = LoggerFactory.getLogger(TracingExecutionInterceptor.class); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java index 5ab99ecae96..e78d3b2d188 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sfn-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sfn/SfnInterceptor.java @@ -18,7 +18,7 @@ public class SfnInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); public SfnInterceptor() {} diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index 122cc55f399..6f84107b7dd 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -30,7 +30,7 @@ public class SnsInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); private SdkBytes getMessageAttributeValueToInject( ExecutionAttributes executionAttributes, String snsTopicName) { diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java index 6be07206130..ed59931810a 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java @@ -32,7 +32,7 @@ public class SqsInterceptor implements ExecutionInterceptor { public static final ExecutionAttribute CONTEXT_ATTRIBUTE = InstanceStore.of(ExecutionAttribute.class) - .putIfAbsent("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); + .getOrCreate("DatadogContext", () -> new ExecutionAttribute<>("DatadogContext")); public SqsInterceptor() {} diff --git a/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java b/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java index ddcee1bcae5..3f7fdd6bdff 100644 --- a/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java +++ b/dd-java-agent/instrumentation/hibernate/hibernate-common/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java @@ -100,6 +100,6 @@ public static void attachSpanFromStore( return; } - targetContextStore.putIfAbsent(target, state); + targetContextStore.getOrPut(target, state); } } diff --git a/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java b/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java index ade78a85eed..be7f635672c 100644 --- a/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java +++ b/dd-java-agent/instrumentation/hibernate/hibernate-core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionFactoryInstrumentation.java @@ -64,11 +64,11 @@ public static void openSession(@Advice.Return final Object session) { if (session instanceof Session) { final ContextStore contextStore = InstrumentationContext.get(Session.class, SessionState.class); - contextStore.putIfAbsent((Session) session, new SessionState(span)); + contextStore.getOrPut((Session) session, new SessionState(span)); } else if (session instanceof StatelessSession) { final ContextStore contextStore = InstrumentationContext.get(StatelessSession.class, SessionState.class); - contextStore.putIfAbsent((StatelessSession) session, new SessionState(span)); + contextStore.getOrPut((StatelessSession) session, new SessionState(span)); } } diff --git a/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java b/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java index 66b2910d592..00482ed4d67 100644 --- a/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java +++ b/dd-java-agent/instrumentation/hibernate/hibernate-core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionFactoryInstrumentation.java @@ -57,7 +57,7 @@ public static void openSession(@Advice.Return final SharedSessionContract sessio final ContextStore contextStore = InstrumentationContext.get(SharedSessionContract.class, SessionState.class); - contextStore.putIfAbsent(session, new SessionState(span)); + contextStore.getOrPut(session, new SessionState(span)); } /** diff --git a/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java b/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java index fb59f988a66..eb78b2cf27e 100644 --- a/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-net/java-net-1.8/src/main/java/datadog/trace/instrumentation/java/net/HttpUrlConnectionInstrumentation.java @@ -74,7 +74,7 @@ public static HttpUrlState methodEnter( final ContextStore contextStore = InstrumentationContext.get(HttpURLConnection.class, HttpUrlState.class); - final HttpUrlState state = contextStore.putIfAbsent(thiz, HttpUrlState.FACTORY); + final HttpUrlState state = contextStore.getOrCreate(thiz, HttpUrlState.FACTORY); synchronized (state) { final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpURLConnection.class); diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java index 8c3ffa48231..e3e657fb5d0 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/SessionInstrumentation.java @@ -118,8 +118,7 @@ public static void bindProducerState( ackMode = Session.AUTO_ACKNOWLEDGE; } sessionState = - sessionStateStore.putIfAbsent( - session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); + sessionStateStore.getOrPut(session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); } boolean isQueue = PRODUCER_DECORATE.isQueue(destination); @@ -159,8 +158,7 @@ public static void bindConsumerState( ackMode = Session.AUTO_ACKNOWLEDGE; } sessionState = - sessionStateStore.putIfAbsent( - session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); + sessionStateStore.getOrPut(session, new SessionState(ackMode, TIME_IN_QUEUE_ENABLED)); } boolean isQueue = CONSUMER_DECORATE.isQueue(destination); diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java index a2fa481491c..e428a73d94b 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInfoInstrumentation.java @@ -161,7 +161,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, @@ -214,7 +214,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java index 227d8872648..f0309fbcb84 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java @@ -267,7 +267,7 @@ public static void captureConfiguration( @Advice.Argument(0) ProducerConfig producerConfig) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); if (Config.get().isDataStreamsEnabled()) { KafkaConfigHelper.storePendingProducerConfig( state, KafkaConfigHelper.extractProducerConfig(producerConfig)); diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java index d6acfe30369..94063de4412 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/MetadataInstrumentation.java @@ -90,7 +90,7 @@ public static void onEnter( String clusterId = newCluster.clusterResource().clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } @@ -111,7 +111,7 @@ public static void onEnter( String clusterId = response.clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java index f704eecade4..085fb7ba866 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ConstructorAdvice.java @@ -59,7 +59,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, KafkaConfigHelper.extractConsumerConfig(consumerConfig)); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java index 6557c1b6821..35ebb9f8564 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/LegacyConstructorAdvice.java @@ -52,7 +52,7 @@ public static void captureGroup( if (Config.get().isDataStreamsEnabled()) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); KafkaConfigHelper.storePendingConsumerConfig( state, normalizedConsumerGroup, KafkaConfigHelper.extractConsumerConfig(consumerConfig)); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java index 307440b06fb..af2a8ce62ee 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdate22AndAfterAdvice.java @@ -16,7 +16,7 @@ public static void onEnter( String clusterId = response.clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java index a1abe7b1959..f5eab61835c 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/MetadataUpdateBefore22Advice.java @@ -16,7 +16,7 @@ public static void onEnter( String clusterId = newCluster.clusterResource().clusterId(); MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); state.clusterId = clusterId; KafkaConfigHelper.reportPendingConfig(state, clusterId); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java index a97caca83b2..ddff8190bcb 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerConstructorAdvice.java @@ -16,7 +16,7 @@ public static void captureConfiguration( @Advice.Argument(0) ProducerConfig producerConfig) { MetadataState state = InstrumentationContext.get(Metadata.class, MetadataState.class) - .putIfAbsent(metadata, MetadataState::new); + .getOrCreate(metadata, MetadataState::new); if (Config.get().isDataStreamsEnabled()) { KafkaConfigHelper.storePendingProducerConfig( state, KafkaConfigHelper.extractProducerConfig(producerConfig)); diff --git a/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java b/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java index 1a1b030fac1..84f3feaff63 100644 --- a/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java +++ b/dd-java-agent/instrumentation/karate/karate-1.0/src/main/java/datadog/trace/instrumentation/karate/KarateExecutionInstrumentation.java @@ -82,7 +82,7 @@ public static void beforeExecute(@Advice.This ScenarioRuntime scenarioRuntime) { ExecutionContext executionContext = InstrumentationContext.get(Scenario.class, ExecutionContext.class) - .computeIfAbsent(scenarioRuntime.scenario, ExecutionContext::create); + .getOrCompute(scenarioRuntime.scenario, ExecutionContext::create); // Indicate beforehand if the failures should be suppressed. This aligns the ordering with the // rest of the frameworks diff --git a/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java b/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java index 1be2b8ee422..f950a077085 100644 --- a/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java +++ b/dd-java-agent/instrumentation/karate/karate-2.0/src/main/java21/datadog/trace/instrumentation/karate2/KarateScenarioAdvice.java @@ -23,7 +23,7 @@ public static void beforeExecute(@Advice.This ScenarioRuntime scenarioRuntime) { ExecutionContext executionContext = InstrumentationContext.get(Scenario.class, ExecutionContext.class) - .computeIfAbsent(scenarioRuntime.getScenario(), ExecutionContext::create); + .getOrCompute(scenarioRuntime.getScenario(), ExecutionContext::create); executionContext.setTestStarted(false); // Indicate beforehand whether failures should be suppressed. This aligns the ordering with diff --git a/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java b/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java index d37cdd69359..c5e7f366a13 100644 --- a/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java +++ b/dd-java-agent/instrumentation/mongo/mongo-common/src/main/java/datadog/trace/instrumentation/mongo/MongoCommandListener.java @@ -121,7 +121,7 @@ public static MongoCommandListener tryRegister( @Override public void commandStarted(final CommandStartedEvent event) { if (listenerAccessor != null) { - listenerAccessor.putIfAbsent(event.getConnectionDescription(), this); + listenerAccessor.getOrPut(event.getConnectionDescription(), this); } // If DBM comment injection is enabled, the span is created on the connection instrumentation diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java index 7148cffd723..cc59edc6757 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelFutureListenerInstrumentation.java @@ -99,7 +99,7 @@ public static ContextScope activateScope(@Advice.Argument(0) final ChannelFuture final ContextContinuation continuation = contextStore - .putIfAbsent(future.getChannel(), ChannelTraceContext.Factory.INSTANCE) + .getOrCreate(future.getChannel(), ChannelTraceContext.Factory.INSTANCE) .getConnectionContinuation(); contextStore.get(future.getChannel()).setConnectionContinuation(null); if (continuation == null) { diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java index 5daa04a82f5..59605630ec6 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/NettyChannelInstrumentation.java @@ -72,7 +72,7 @@ public static void addConnectContinuation(@Advice.This final Channel channel) { InstrumentationContext.get(Channel.class, ChannelTraceContext.class); if (contextStore - .putIfAbsent(channel, ChannelTraceContext.Factory.INSTANCE) + .getOrCreate(channel, ChannelTraceContext.Factory.INSTANCE) .getConnectionContinuation() != null) { continuation.release(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java index 974ecbe5181..07fc9be8b0f 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java @@ -41,7 +41,7 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m } final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); ContextScope parentScope = null; final ContextContinuation continuation = channelTraceContext.getConnectionContinuation(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java index e25a261bb56..e340be8ee77 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientResponseTracingHandler.java @@ -28,7 +28,7 @@ public HttpClientResponseTracingHandler( public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent msg) throws Exception { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); AgentSpan parent = channelTraceContext.getClientParentSpan(); if (parent == null) { @@ -56,7 +56,7 @@ public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent @Override public void exceptionCaught(ChannelHandlerContext ctx, ExceptionEvent e) throws Exception { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); AgentSpan parent = channelTraceContext.getClientParentSpan(); if (parent == null) { diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java index ae8cfd62eb5..3a2aeef4b3f 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java @@ -27,7 +27,7 @@ public HttpServerRequestTracingHandler( @Override public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent msg) { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); if (!(msg.getMessage() instanceof HttpRequest)) { final Context storedContext = channelTraceContext.getServerContext(); diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java index cede6cd43c9..dffde997846 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerResponseTracingHandler.java @@ -27,7 +27,7 @@ public HttpServerResponseTracingHandler( @Override public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent msg) { final ChannelTraceContext channelTraceContext = - contextStore.putIfAbsent(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); + contextStore.getOrCreate(ctx.getChannel(), ChannelTraceContext.Factory.INSTANCE); final AgentSpan span = channelTraceContext.getServerSpan(); if (span == null || !(msg.getMessage() instanceof HttpResponse)) { diff --git a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java index d7451780a5b..8cd6d9d3748 100644 --- a/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java +++ b/dd-java-agent/instrumentation/netty/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/MaybeBlockResponseHandler.java @@ -36,7 +36,7 @@ public MaybeBlockResponseHandler(final ContextStore contextFactory) { - final C existing = map.get(key); - if (existing != null) { - return existing; - } - final C created = contextFactory.create(); - map.put(key, created); - return created; - } - - @Override - public C computeIfAbsent(final K key, final KeyAwareFactory contextFactory) { + public C getOrCompute(final K key, final KeyAwareFactory contextFactory) { final C existing = map.get(key); if (existing != null) { return existing; diff --git a/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java b/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java index df2f9a91a0c..1ddbd88fb7e 100644 --- a/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java +++ b/dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ReactorContextBridge.java @@ -90,7 +90,7 @@ public static void transferToOptimizedSubscriber( context = subscriberContexts.get(source); } if (context != null) { - subscriberContexts.putIfAbsent(target, context); + subscriberContexts.getOrPut(target, context); } } diff --git a/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java b/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java index bfb2bb7cdbe..4a42aeef8fa 100644 --- a/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java +++ b/dd-java-agent/instrumentation/resilience4j/resilience4j-reactor-2.0/src/main/java/datadog/trace/instrumentation/resilience4j/ReactorHelper.java @@ -32,7 +32,7 @@ public static BiConsumer, AgentSpan> putInto( public static BiConsumer, AgentSpan> putIfAbsentInto( final ContextStore store) { - return (publisher, span) -> store.putIfAbsent(publisher, HandoffContext.anyThread(span)); + return (publisher, span) -> store.getOrPut(publisher, HandoffContext.anyThread(span)); } public static Function, Publisher> wrapFunction( diff --git a/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java b/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java index 75a6e3d4764..22149f9a2a4 100644 --- a/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java +++ b/dd-java-agent/instrumentation/scalatest-3.0.8/src/main/java/datadog/trace/instrumentation/scalatest/ScalatestSkipInstrumentation.java @@ -76,7 +76,7 @@ public static void apply( int runStamp = tracker.nextOrdinal().runStamp(); RunContext context = RunContext.getOrCreate(runStamp); RunContext existingContext = - InstrumentationContext.get(Filter.class, RunContext.class).putIfAbsent(filter, context); + InstrumentationContext.get(Filter.class, RunContext.class).getOrPut(filter, context); if (existingContext != context) { // This shouldn't happen. // If it does, instrumentation isn't working as expected, or Scalatest internals changed. diff --git a/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java b/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java index e69d7902d73..a9b03588791 100644 --- a/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java +++ b/dd-java-agent/instrumentation/tibco-businessworks/tibco-businessworks-5.14/src/main/java/datadog/trace/instrumentation/tibcobw5/JobInstrumentation.java @@ -43,7 +43,7 @@ public static void after( DECORATE.onProcessStart(span, workflowName); Map map = InstrumentationContext.get(ProcessContext.class, Map.class) - .putIfAbsent(processContext, HashMap::new); + .getOrCreate(processContext, HashMap::new); map.put(wId, span); } } diff --git a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java index 41de9e9f2d5..dd310511fdf 100644 --- a/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/WsHandshakeRequestInstrumentation.java @@ -49,8 +49,7 @@ public static void captureHandshakeSpan(@Advice.This final WsHandshakeRequest se // apply jee configuration overrides if any since the servlet instrumentation won't kick in // for this span. ClassloaderConfigurationOverrides.maybeEnrichSpan(span); - InstrumentationContext.get(WsHandshakeRequest.class, AgentSpan.class) - .putIfAbsent(self, span); + InstrumentationContext.get(WsHandshakeRequest.class, AgentSpan.class).getOrPut(self, span); } } } diff --git a/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java b/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java index c34602884bd..e4876b8c670 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java +++ b/dd-java-agent/instrumentation/undertow/undertow-common/src/main/java/datadog/trace/instrumentation/undertow/UndertowDecorator.java @@ -28,12 +28,12 @@ public class UndertowDecorator @SuppressWarnings("unchecked") public static final AttachmentKey DATADOG_UNDERTOW_CONTINUATION = - attachmentStore.putIfAbsent( + attachmentStore.getOrCreate( "DD_UNDERTOW_CONTINUATION", () -> AttachmentKey.create(ContextContinuation.class)); @SuppressWarnings("unchecked") public static final AttachmentKey PARENT_CONTEXT_KEY = - attachmentStore.putIfAbsent( + attachmentStore.getOrCreate( "DD_UNDERTOW_PARENT_CONTEXT", () -> AttachmentKey.create(Context.class)); public static final UndertowDecorator DECORATE = new UndertowDecorator(); diff --git a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java index 5f475eb82cc..ef50904e530 100644 --- a/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java +++ b/dd-java-agent/instrumentation/websocket/javax-websocket-1.0/src/main/java/datadog/trace/instrumentation/websocket/jsr256/EndpointInstrumentation.java @@ -68,7 +68,7 @@ public static void onEnter(@Advice.Argument(0) final Session session) { current.forceSamplingDecision(); } InstrumentationContext.get(Session.class, HandlerContext.Sender.class) - .putIfAbsent( + .getOrPut( session, new HandlerContext.Sender(current.getLocalRootSpan(), session.getId())); } } diff --git a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java index 982330b6805..7219b6f0ce9 100644 --- a/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java +++ b/dd-java-agent/instrumentation/websocket/jetty-websocket/jetty-websocket-10.0/src/main/java11/datadog/trace/instrumentation/websocket/jetty10/MethodHandleWrappers.java @@ -96,7 +96,7 @@ public static void onOpen( if (Config.get().isWebsocketMessagesInheritSampling()) { current.forceSamplingDecision(); } - sessionStore.putIfAbsent( + sessionStore.getOrPut( session, new HandlerContext.Sender(current.getLocalRootSpan(), session.getId())); } } catch (Throwable t) { diff --git a/internal-api/build.gradle.kts b/internal-api/build.gradle.kts index cd0a0e59723..42670eb6bcb 100644 --- a/internal-api/build.gradle.kts +++ b/internal-api/build.gradle.kts @@ -77,6 +77,7 @@ extra["excludedClassesCoverage"] = listOf( "datadog.trace.api.debugger.DebuggerConfigUpdate", // Bootstrap API "datadog.trace.bootstrap.ActiveSubsystems", + "datadog.trace.bootstrap.ContextStore", "datadog.trace.bootstrap.ContextStore.Factory", "datadog.trace.bootstrap.instrumentation.api.java.lang.ProcessImplInstrumentationHelpers", "datadog.trace.bootstrap.instrumentation.api.Tags", 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 f0528cd976e..ca7690ccfef 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java @@ -65,7 +65,7 @@ interface KeyAwareFactory { * @param context new context instance to put * @return old instance if it was present, or new instance */ - C putIfAbsent(K key, C context); + C getOrPut(K key, C context); /** * Put new context instance if key is absent. Uses context factory to avoid creating objects if @@ -75,7 +75,9 @@ interface KeyAwareFactory { * @param contextFactory factory instance to produce new context object * @return old instance if it was present, or new instance */ - C putIfAbsent(K key, Factory contextFactory); + default C getOrCreate(K key, Factory contextFactory) { + return getOrCompute(key, contextFactory); + } /** * Put new context instance if key is absent. Uses context factory to avoid creating objects if @@ -85,7 +87,7 @@ interface KeyAwareFactory { * @param contextFactory factory instance to produce new context object * @return old instance if it was present, or new instance */ - C computeIfAbsent(K key, KeyAwareFactory contextFactory); + C getOrCompute(K key, KeyAwareFactory contextFactory); /** * Removes the existing value for key and return it. From 7740c46e10f1b2d21b22fa4e5f36771e3308a421 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Aug 2026 17:55:17 +0100 Subject: [PATCH 2/7] Cleanup javadoc --- .../trace/bootstrap/InstanceStore.java | 5 +- .../trace/bootstrap/WeakMapContextStore.java | 8 +-- .../datadog/trace/bootstrap/ContextStore.java | 49 ++++++++++--------- 3 files changed, 35 insertions(+), 27 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java index 14e091774e2..6f7e4f7b560 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstanceStore.java @@ -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); } @@ -71,6 +73,7 @@ public T getOrCreate(String key, Supplier 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); } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java index 6485966d913..a55adb866ec 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java @@ -40,10 +40,10 @@ public V getOrPut(final K key, final V context) { 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) { @@ -61,10 +61,10 @@ public V getOrCompute(K key, KeyAwareFactory contextFactory) { 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) { 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 ca7690ccfef..6d22bbe92ad 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,7 @@ package datadog.trace.bootstrap; +import javax.annotation.Nullable; + /** * Interface to represent context storage for instrumentations. * @@ -43,57 +45,60 @@ interface KeyAwareFactory { } /** - * Get context given the key + * Get context instance for the given key. * - * @param key the key to lookup - * @return context object + * @param key the context key + * @return context instance; {@code null} if the key had no context */ + @Nullable C get(K key); /** - * Put new context instance for given key + * Unconditionally put new context instance for the given key. * - * @param key key to use + * @param key the context key * @param context context instance to save */ void put(K key, C context); /** - * Put new context instance if key is absent + * Gets the context instance for the given key. If no context exists then associate it with the + * new context. * - * @param key key to use - * @param context new context instance to put - * @return old instance if it was present, or new instance + * @param key the context key + * @param context new context instance + * @return existing context instance if present; otherwise new instance */ C getOrPut(K key, C context); /** - * Put new context instance if key is absent. Uses context factory to avoid creating objects if - * not needed. + * Gets the context instance for the given key. If no context exists then create one using the + * given factory and associate it with the key. * - * @param key key to use - * @param contextFactory factory instance to produce new context object - * @return old instance if it was present, or new instance + * @param key the context key + * @param contextFactory factory instance to produce new context instances + * @return existing context instance if present; otherwise new instance */ default C getOrCreate(K key, Factory contextFactory) { return getOrCompute(key, contextFactory); } /** - * Put new context instance if key is absent. Uses context factory to avoid creating objects if - * not needed. + * Gets the context instance for the given key. If no context exists then create one using the + * given factory and associate it with the key. * - * @param key key to use - * @param contextFactory factory instance to produce new context object - * @return old instance if it was present, or new instance + * @param key the context key + * @param contextFactory factory instance to produce new context instances + * @return existing context instance if present; otherwise new instance */ C getOrCompute(K key, KeyAwareFactory contextFactory); /** - * Removes the existing value for key and return it. + * Removes the context instance for the given key. * - * @param key the key remove - * @return removed context object + * @param key the context key + * @return removed context instance; {@code null} if the key had no context */ + @Nullable C remove(K key); } From f8da6158d7316f431460b0870032feec4d27fa9f Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Aug 2026 18:08:09 +0100 Subject: [PATCH 3/7] Rename ConcurrentHashMapContextStore to StrongMapContextStore to make clear it uses strong references --- .../CiVisibilityInstrumentationTest.groovy | 6 +++--- .../trace/civisibility/CiVisibilitySystem.java | 6 +++--- ...textStore.java => StrongMapContextStore.java} | 16 ++++++++-------- .../events/TestEventsHandlerImplTest.java | 6 +++--- 4 files changed, 17 insertions(+), 17 deletions(-) rename dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/{ConcurrentHashMapContextStore.java => StrongMapContextStore.java} (53%) diff --git a/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy b/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy index 2618d70d520..58a7aa1d13a 100644 --- a/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy +++ b/dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy @@ -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<>()) } } diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java index d094cdf838a..af7f653e921 100644 --- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java +++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java @@ -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 TestEventsHandler 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; } diff --git a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java similarity index 53% rename from dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java rename to dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java index 606707bfb12..55fc443769c 100644 --- a/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/ConcurrentHashMapContextStore.java +++ b/dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/utils/StrongMapContextStore.java @@ -4,32 +4,32 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; -public class ConcurrentHashMapContextStore implements ContextStore { - - private final ConcurrentMap m = new ConcurrentHashMap<>(); +/** Substitute {@link ContextStore} that uses strong-references to track contexts. */ +public class StrongMapContextStore implements ContextStore { + private final ConcurrentMap map = new ConcurrentHashMap<>(); @Override public C get(K key) { - return m.get(key); + return map.get(key); } @Override public void put(K key, C context) { - m.put(key, context); + map.put(key, context); } @Override public C getOrPut(K key, C context) { - return m.computeIfAbsent(key, k -> context); + return map.computeIfAbsent(key, k -> context); } @Override public C getOrCompute(K key, KeyAwareFactory contextFactory) { - return m.computeIfAbsent(key, contextFactory::create); + return map.computeIfAbsent(key, contextFactory::create); } @Override public C remove(K key) { - return m.remove(key); + return map.remove(key); } } diff --git a/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java b/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java index dc62a3261c7..4c20bba1d2b 100644 --- a/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java +++ b/dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/events/TestEventsHandlerImplTest.java @@ -14,7 +14,7 @@ import datadog.trace.bootstrap.ContextStore; import datadog.trace.civisibility.domain.TestFrameworkModule; import datadog.trace.civisibility.domain.TestFrameworkSession; -import datadog.trace.civisibility.utils.ConcurrentHashMapContextStore; +import datadog.trace.civisibility.utils.StrongMapContextStore; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; import org.junit.jupiter.api.Test; @@ -89,8 +89,8 @@ private static TestEventsHandlerImpl handler( private static TestEventsHandlerImpl handler( Supplier testSessionSupplier, boolean eagerSessionStart) { - ContextStore suiteStore = new ConcurrentHashMapContextStore<>(); - ContextStore testStore = new ConcurrentHashMapContextStore<>(); + ContextStore suiteStore = new StrongMapContextStore<>(); + ContextStore testStore = new StrongMapContextStore<>(); return new TestEventsHandlerImpl<>( NoOpMetricCollector.INSTANCE, testSessionSupplier, From eaf9022061e6eb7381286d1d8039aca7712e6433 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Aug 2026 18:09:58 +0100 Subject: [PATCH 4/7] Simplify WeakMapContextStore - it does not need to implement ContextStore, or allow for different max sizes --- .../trace/bootstrap/WeakMapContextStore.java | 23 ++++--------------- 1 file changed, 5 insertions(+), 18 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java index a55adb866ec..e2a480bdc5c 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java @@ -1,40 +1,29 @@ package datadog.trace.bootstrap; 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. * *

This class should be created lazily because it uses weak maps with background cleanup. */ -final class WeakMapContextStore implements ContextStore { - private static final int DEFAULT_MAX_SIZE = 50_000; +final class WeakMapContextStore { + private static final int MAX_SIZE = 50_000; - private final int maxSize; private final WeakMap map = WeakMap.Supplier.newWeakMap(); - public WeakMapContextStore(int maxSize) { - this.maxSize = maxSize; - } - - public WeakMapContextStore() { - this(DEFAULT_MAX_SIZE); - } - - @Override @SuppressWarnings("unchecked") public V get(final K key) { return (V) map.get(key); } - @Override public void put(final K key, final V context) { - if (map.size() < maxSize) { + if (map.size() < MAX_SIZE) { map.put(key, context); } } - @Override public V getOrPut(final K key, final V context) { V existingContext = get(key); if (null == existingContext) { @@ -55,7 +44,6 @@ public V getOrPut(final K key, final V context) { return existingContext; } - @Override public V getOrCompute(K key, KeyAwareFactory contextFactory) { V existingContext = get(key); if (null == existingContext) { @@ -76,7 +64,6 @@ public V getOrCompute(K key, KeyAwareFactory contextFactory) { return existingContext; } - @Override @SuppressWarnings("unchecked") public V remove(final K key) { return (V) map.remove(key); From bc8e76090a5c3d09dc281c97ed93bbb63763ae18 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Aug 2026 18:31:09 +0100 Subject: [PATCH 5/7] Move weak-map fallback static helpers to WeakMapContextStore and rename WeakMapContextStore to WeakMapPerStore --- .../bootstrap/FieldBackedContextStore.java | 6 +++--- .../bootstrap/FieldBackedContextStores.java | 10 --------- ...ContextStore.java => WeakMapPerStore.java} | 16 +++++++++++++- .../context/FieldBackedContextInjector.java | 21 +++++-------------- ...veImageGeneratorRunnerInstrumentation.java | 2 +- 5 files changed, 24 insertions(+), 31 deletions(-) rename dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/{WeakMapContextStore.java => WeakMapPerStore.java} (79%) 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 f60f890fa17..7d334d5405b 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 @@ -89,14 +89,14 @@ public Object remove(Object key) { } // only create WeakMap-based fall-back when we need it - private volatile WeakMapContextStore weakStore; + private volatile WeakMapPerStore weakStore; private final Object synchronizationInstance = new Object(); - WeakMapContextStore weakStore() { + WeakMapPerStore weakStore() { if (null == weakStore) { synchronized (synchronizationInstance) { if (null == weakStore) { - weakStore = new WeakMapContextStore<>(); + weakStore = new WeakMapPerStore<>(); } } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java index 529e475ab87..8ae3a3239c5 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/FieldBackedContextStores.java @@ -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); - } } diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java similarity index 79% rename from dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java rename to dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java index e2a480bdc5c..8929735911d 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapContextStore.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMapPerStore.java @@ -1,5 +1,7 @@ package datadog.trace.bootstrap; +import static datadog.trace.bootstrap.FieldBackedContextStores.getContextStore; + import datadog.trace.api.internal.VisibleForTesting; import datadog.trace.bootstrap.ContextStore.KeyAwareFactory; @@ -8,11 +10,13 @@ * *

This class should be created lazily because it uses weak maps with background cleanup. */ -final class WeakMapContextStore { +public final class WeakMapPerStore { private static final int MAX_SIZE = 50_000; private final WeakMap map = WeakMap.Supplier.newWeakMap(); + WeakMapPerStore() {} + @SuppressWarnings("unchecked") public V get(final K key) { return (V) map.get(key); @@ -73,4 +77,14 @@ public V remove(final K key) { int size() { return map.size(); } + + /** 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); + } + + /** 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); + } } 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 1392a08fbfa..ee8ae19ca09 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 @@ -8,7 +8,7 @@ import datadog.trace.api.Pair; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.FieldBackedContextAccessor; -import datadog.trace.bootstrap.FieldBackedContextStores; +import datadog.trace.bootstrap.WeakMapPerStore; import java.io.Serializable; import java.util.Arrays; import java.util.BitSet; @@ -35,9 +35,6 @@ 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 FIELD_BACKED_CONTEXT_ACCESSOR_CLASS = getInternalName(FieldBackedContextAccessor.class.getName()); @@ -51,12 +48,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_REDIRECT_CLASS = getInternalName(WeakMapPerStore.class.getName()); + 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_DESCRIPTOR = Type.getMethodDescriptor( Type.VOID_TYPE, Type.getType(Object.class), Type.INT_TYPE, Type.getType(Object.class)); @@ -444,11 +441,7 @@ private void invokeWeakGet(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ALOAD, 0); mv.visitIntInsn(Opcodes.ILOAD, 1); mv.visitMethodInsn( - Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, - WEAK_GET_METHOD, - WEAK_GET_METHOD_DESCRIPTOR, - false); + Opcodes.INVOKESTATIC, WEAK_REDIRECT_CLASS, "get", WEAK_GET_METHOD_DESCRIPTOR, false); mv.visitInsn(Opcodes.ARETURN); } @@ -457,11 +450,7 @@ private void invokeWeakPut(final MethodVisitor mv) { mv.visitIntInsn(Opcodes.ILOAD, 1); mv.visitIntInsn(Opcodes.ALOAD, 2); mv.visitMethodInsn( - Opcodes.INVOKESTATIC, - FIELD_BACKED_CONTEXT_STORES_CLASS, - WEAK_PUT_METHOD, - WEAK_PUT_METHOD_DESCRIPTOR, - false); + Opcodes.INVOKESTATIC, WEAK_REDIRECT_CLASS, "put", WEAK_PUT_METHOD_DESCRIPTOR, false); mv.visitInsn(Opcodes.RETURN); } 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 a90d65c2472..d32db6b803f 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 @@ -148,7 +148,7 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[ + "datadog.trace.bootstrap.instrumentation.jfr.exceptions.ExceptionSampleEvent:build_time," + "datadog.trace.bootstrap.instrumentation.jfr.backpressure.BackpressureSampleEvent:build_time," + "datadog.trace.bootstrap.instrumentation.jfr.directallocation.DirectAllocationTotalEvent:build_time," - + "datadog.trace.bootstrap.WeakMapContextStore:build_time," + + "datadog.trace.bootstrap.WeakMapPerStore:build_time," + "datadog.trace.config.inversion.ConfigHelper:rerun," + "datadog.trace.config.inversion.ConfigHelper$StrictnessPolicy:rerun," + "datadog.trace.config.inversion.GeneratedSupportedConfigurations:build_time," From a4a1653a93e8edba516b2542f650f3231212c4b3 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Aug 2026 20:05:11 +0100 Subject: [PATCH 6/7] Move public weak helpers to top of class; other methods become package-private --- .../trace/bootstrap/WeakMapPerStore.java | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) 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 8929735911d..36917246b1f 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 @@ -11,6 +11,17 @@ *

This class should be created lazily because it uses weak maps with background cleanup. */ public final class WeakMapPerStore { + + /** 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); + } + + /** 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); + } + private static final int MAX_SIZE = 50_000; private final WeakMap map = WeakMap.Supplier.newWeakMap(); @@ -18,17 +29,17 @@ public final class WeakMapPerStore { WeakMapPerStore() {} @SuppressWarnings("unchecked") - public V get(final K key) { + V get(final K key) { return (V) map.get(key); } - public void put(final K key, final V context) { + void put(final K key, final V context) { if (map.size() < MAX_SIZE) { map.put(key, context); } } - public V getOrPut(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 @@ -48,7 +59,7 @@ public V getOrPut(final K key, final V context) { return existingContext; } - public V getOrCompute(K key, KeyAwareFactory contextFactory) { + V getOrCompute(K key, KeyAwareFactory contextFactory) { V existingContext = get(key); if (null == existingContext) { // This whole part with using synchronized is only because @@ -69,7 +80,7 @@ public V getOrCompute(K key, KeyAwareFactory contextFactory) { } @SuppressWarnings("unchecked") - public V remove(final K key) { + V remove(final K key) { return (V) map.remove(key); } @@ -77,14 +88,4 @@ public V remove(final K key) { int size() { return map.size(); } - - /** 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); - } - - /** 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); - } } From c1172838687f4bb824d7f85ef121aac2d445cba5 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Aug 2026 22:12:07 +0100 Subject: [PATCH 7/7] Avoid unnecessary lambda capture in StrongMapContextStore --- .../trace/civisibility/utils/StrongMapContextStore.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 55fc443769c..2ca2586ce41 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 @@ -20,7 +20,8 @@ public void put(K key, C context) { @Override public C getOrPut(K key, C context) { - return map.computeIfAbsent(key, k -> context); + C existing = map.putIfAbsent(key, context); + return existing != null ? existing : context; } @Override