From 769cd3bdd4b0843850b531440e718572fd9269da Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 6 Mar 2026 14:07:06 +0100 Subject: [PATCH 01/12] wip --- .../src/test/groovy/JMS1Test.groovy | 20 ++++++++++++++++ .../config/TraceInstrumentationConfig.java | 2 ++ .../java/datadog/trace/core/CoreTracer.java | 24 ++++++++++++++++--- .../datadog/trace/api/InstrumenterConfig.java | 10 ++++++++ .../instrumentation/api/AgentTracer.java | 7 +++--- 5 files changed, 57 insertions(+), 6 deletions(-) diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy index 3165d2d01f6..1e6fbca7fce 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy @@ -1,3 +1,6 @@ +import static datadog.trace.api.config.TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED +import static org.junit.jupiter.api.Assumptions.assumeTrue + import datadog.trace.agent.test.asserts.ListWriterAssert import datadog.trace.agent.test.asserts.TraceAssert import datadog.trace.agent.test.naming.VersionedNamingTestBase @@ -61,6 +64,10 @@ abstract class JMS1Test extends VersionedNamingTestBase { abstract String operationForConsumer() + boolean testUnclosedScopeFinished() { + true + } + def setupSpec() { broker.start() final ActiveMQConnectionFactory connectionFactory = broker.createConnectionFactory() @@ -164,6 +171,7 @@ abstract class JMS1Test extends VersionedNamingTestBase { def "closing #destinationType session should close and finish any pending scopes"() { setup: + assumeTrue(testUnclosedScopeFinished()) def destination = destinationType.create(session) def localSession = connection.createSession(false, Session.AUTO_ACKNOWLEDGE) def producer = localSession.createProducer(destination) @@ -1045,6 +1053,18 @@ class JMS1V0Test extends JMS1Test { } } +class JMSContextSwapForkedTest extends JMS1V0Test { + @Override + protected void configurePreAgent() { + injectSysConfig(MESSAGING_CONTEXT_SWAP_ENABLED, "true") + } + + @Override + boolean testUnclosedScopeFinished() { + false + } +} + class JMS1V1ForkedTest extends JMS1Test { @Override 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 80c294550e6..b66bd67972f 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 @@ -201,5 +201,7 @@ public final class TraceInstrumentationConfig { public static final String TRACE_RESOURCE_RENAMING_ALWAYS_SIMPLIFIED_ENDPOINT = "trace.resource.renaming.always.simplified.endpoint"; + public static final String MESSAGING_CONTEXT_SWAP_ENABLED = "messaging.context.swap.enabled"; + private TraceInstrumentationConfig() {} } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index c58d1ff277d..5bcbfb41e27 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -22,6 +22,8 @@ import datadog.communication.ddagent.DDAgentFeaturesDiscovery; import datadog.communication.ddagent.ExternalAgentLauncher; import datadog.communication.ddagent.SharedCommunicationObjects; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.context.propagation.Propagators; import datadog.environment.ThreadSupport; import datadog.metrics.agent.AgentMeter; @@ -35,6 +37,7 @@ import datadog.trace.api.DynamicConfig; import datadog.trace.api.EndpointTracker; import datadog.trace.api.IdGenerationStrategy; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.Pair; import datadog.trace.api.TagMap; import datadog.trace.api.TraceConfig; @@ -1139,12 +1142,27 @@ public void setAsyncPropagationEnabled(boolean asyncPropagationEnabled) { @Override public void closePrevious(boolean finishSpan) { - scopeManager.closePrevious(finishSpan); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final Context previous = Context.root().swap(); + if (finishSpan) { + final AgentSpan span = AgentSpan.fromContext(previous); + if (span != null) { + span.finish(); + } + } + } else { + scopeManager.closePrevious(finishSpan); + } } @Override - public AgentScope activateNext(AgentSpan span) { - return scopeManager.activateNext(span); + public ContextScope activateNext(AgentSpan span) { + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + span.swap(); + return Context.current().attach(); + } else { + return scopeManager.activateNext(span); + } } public TagInterceptor getTagInterceptor() { 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 a7e5b5bd5a3..38696171ca8 100644 --- a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java +++ b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java @@ -58,6 +58,7 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.JDBC_POOL_WAITING_ENABLED; import static datadog.trace.api.config.TraceInstrumentationConfig.JDBC_PREPARED_STATEMENT_CLASS_NAME; import static datadog.trace.api.config.TraceInstrumentationConfig.MEASURE_METHODS; +import static datadog.trace.api.config.TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED; import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_CACHE_CONFIG; import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_CACHE_DIR; import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_NAMES_ARE_UNIQUE; @@ -212,6 +213,7 @@ public class InstrumenterConfig { private final boolean apiSecurityEndpointCollectionEnabled; private final boolean appLogsCollectionEnabled; + private final boolean messagingContextSwapEnabled; static { // Bind telemetry collector to config module before initializing ConfigProvider @@ -363,6 +365,8 @@ private InstrumenterConfig() { appLogsCollectionEnabled = configProvider.getBoolean(APP_LOGS_COLLECTION_ENABLED, DEFAULT_APP_LOGS_COLLECTION_ENABLED); + + messagingContextSwapEnabled = configProvider.getBoolean(MESSAGING_CONTEXT_SWAP_ENABLED, false); } public boolean isCodeOriginEnabled() { @@ -682,6 +686,10 @@ public boolean isAppLogsCollectionEnabled() { return appLogsCollectionEnabled; } + public boolean isMessagingContextSwapEnabled() { + return messagingContextSwapEnabled; + } + // This has to be placed after all other static fields to give them a chance to initialize private static final InstrumenterConfig INSTANCE = new InstrumenterConfig( @@ -801,6 +809,8 @@ public String toString() { + dataJobsEnabled + ", apiSecurityEndpointCollectionEnabled=" + apiSecurityEndpointCollectionEnabled + + ", messagingContextSwapEnabled=" + + messagingContextSwapEnabled + '}'; } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index df1e3ae0fcd..21d5335f07e 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -1,5 +1,6 @@ package datadog.trace.bootstrap.instrumentation.api; +import datadog.context.ContextScope; import datadog.trace.api.ConfigDefaults; import datadog.trace.api.DDTraceId; import datadog.trace.api.EndpointCheckpointer; @@ -182,7 +183,7 @@ public static void closePrevious(final boolean finishSpan) { * * @see datadog.trace.api.config.TracerConfig#SCOPE_ITERATION_KEEP_ALIVE */ - public static AgentScope activateNext(final AgentSpan span) { + public static ContextScope activateNext(final AgentSpan span) { return get().activateNext(span); } @@ -367,7 +368,7 @@ AgentSpan startSpan( void closePrevious(boolean finishSpan); - AgentScope activateNext(AgentSpan span); + ContextScope activateNext(AgentSpan span); AgentSpan activeSpan(); @@ -548,7 +549,7 @@ public void closeActive() {} public void closePrevious(final boolean finishSpan) {} @Override - public AgentScope activateNext(final AgentSpan span) { + public ContextScope activateNext(final AgentSpan span) { return NoopScope.INSTANCE; } From af74d1bb0d356e697ea40e97f8d4bc5a1361edb2 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 6 Mar 2026 14:57:39 +0100 Subject: [PATCH 02/12] Changes --- .../JMSMessageConsumerInstrumentation.java | 33 +++++++++++++++++-- .../src/test/groovy/JMS1Test.groovy | 1 + .../java/datadog/trace/core/CoreTracer.java | 23 ++++--------- .../instrumentation/api/AgentTracer.java | 7 ++-- 4 files changed, 41 insertions(+), 23 deletions(-) diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java index 8d6e3d62436..9baf17aebcc 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java @@ -9,6 +9,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.jms.JMSDecorator.BROKER_DECORATE; import static datadog.trace.instrumentation.jms.JMSDecorator.CONSUMER_DECORATE; import static datadog.trace.instrumentation.jms.JMSDecorator.JMS_CONSUME; @@ -24,6 +26,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.api.Config; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.bootstrap.CallDepthThreadLocalMap; @@ -90,7 +93,15 @@ public static MessageConsumerState beforeReceive(@Advice.This final MessageConsu } boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge(); - closePrevious(finishSpan); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + CONSUMER_DECORATE.beforeFinish(span); + span.finish(); + } + } else { + closePrevious(finishSpan); + } if (finishSpan) { consumerState.finishTimeInQueueSpan(false); } @@ -163,7 +174,15 @@ public static void afterReceive( CONSUMER_DECORATE.onError(span, throwable); - activateNext(span); // scope is left open until next message or it times out + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan previous = spanFromContext(span.swap()); + if (previous != null) { + CONSUMER_DECORATE.beforeFinish(previous); + previous.finish(); + } + } else { + activateNext(span); // scope is left open until next message or it times out + } JMSLogger.logIterationSpan(span); SessionState sessionState = consumerState.getSessionState(); @@ -187,7 +206,15 @@ public static void beforeClose(@Advice.This final MessageConsumer consumer) { .get(consumer); if (null != consumerState) { boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge(); - closePrevious(finishSpan); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + CONSUMER_DECORATE.beforeFinish(span); + span.finish(); + } + } else { + closePrevious(finishSpan); + } if (finishSpan) { consumerState.finishTimeInQueueSpan(true); } diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy index 1e6fbca7fce..33cddadf5f4 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy @@ -1061,6 +1061,7 @@ class JMSContextSwapForkedTest extends JMS1V0Test { @Override boolean testUnclosedScopeFinished() { + //TODO: This need to be removed when the Context manager will support it false } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 5bcbfb41e27..32adab8cc45 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -22,8 +22,6 @@ import datadog.communication.ddagent.DDAgentFeaturesDiscovery; import datadog.communication.ddagent.ExternalAgentLauncher; import datadog.communication.ddagent.SharedCommunicationObjects; -import datadog.context.Context; -import datadog.context.ContextScope; import datadog.context.propagation.Propagators; import datadog.environment.ThreadSupport; import datadog.metrics.agent.AgentMeter; @@ -1143,26 +1141,19 @@ public void setAsyncPropagationEnabled(boolean asyncPropagationEnabled) { @Override public void closePrevious(boolean finishSpan) { if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { - final Context previous = Context.root().swap(); - if (finishSpan) { - final AgentSpan span = AgentSpan.fromContext(previous); - if (span != null) { - span.finish(); - } - } - } else { - scopeManager.closePrevious(finishSpan); + throw new IllegalStateException( + "closePrevious must not be called when context swap based logic is enabled"); } + scopeManager.closePrevious(finishSpan); } @Override - public ContextScope activateNext(AgentSpan span) { + public AgentScope activateNext(AgentSpan span) { if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { - span.swap(); - return Context.current().attach(); - } else { - return scopeManager.activateNext(span); + throw new IllegalStateException( + "activateNext must not be called when context swap based logic is enabled"); } + return scopeManager.activateNext(span); } public TagInterceptor getTagInterceptor() { diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 21d5335f07e..df1e3ae0fcd 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -1,6 +1,5 @@ package datadog.trace.bootstrap.instrumentation.api; -import datadog.context.ContextScope; import datadog.trace.api.ConfigDefaults; import datadog.trace.api.DDTraceId; import datadog.trace.api.EndpointCheckpointer; @@ -183,7 +182,7 @@ public static void closePrevious(final boolean finishSpan) { * * @see datadog.trace.api.config.TracerConfig#SCOPE_ITERATION_KEEP_ALIVE */ - public static ContextScope activateNext(final AgentSpan span) { + public static AgentScope activateNext(final AgentSpan span) { return get().activateNext(span); } @@ -368,7 +367,7 @@ AgentSpan startSpan( void closePrevious(boolean finishSpan); - ContextScope activateNext(AgentSpan span); + AgentScope activateNext(AgentSpan span); AgentSpan activeSpan(); @@ -549,7 +548,7 @@ public void closeActive() {} public void closePrevious(final boolean finishSpan) {} @Override - public ContextScope activateNext(final AgentSpan span) { + public AgentScope activateNext(final AgentSpan span) { return NoopScope.INSTANCE; } From 5c718e948f6f8b1cc5bedb4b592abf31e7fdfc3c Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 6 Mar 2026 16:41:19 +0100 Subject: [PATCH 03/12] migrate everything else --- .../aws/v1/sqs/TracingIterator.java | 30 +++++++++++++++++-- .../aws/v1/sqs/TracingListIterator.java | 13 +++++++- .../src/test/groovy/SqsClientTest.groovy | 9 +++++- .../aws/v2/sqs/TracingIterator.java | 30 +++++++++++++++++-- .../aws/v2/sqs/TracingListIterator.java | 13 +++++++- .../src/test/groovy/SqsClientTest.groovy | 9 ++++++ .../JMSMessageConsumerInstrumentation.java | 6 ++-- .../kafka_clients/TracingIterator.java | 30 +++++++++++++++++-- .../kafka_clients/TracingListIterator.java | 13 +++++++- .../test/groovy/KafkaClientTestBase.groovy | 9 ++++++ .../kafka_clients38/TracingIterator.java | 30 +++++++++++++++++-- .../kafka_clients38/TracingListIterator.java | 13 +++++++- .../test/groovy/KafkaClientTestBase.groovy | 9 ++++++ ...ssageListenerContainerInstrumentation.java | 13 +++++++- metadata/supported-configurations.json | 8 +++++ 15 files changed, 214 insertions(+), 21 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java index 98aeca1ec84..5cac40c854b 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java @@ -7,6 +7,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName; import static datadog.trace.instrumentation.aws.v1.sqs.MessageExtractAdapter.GETTER; import static datadog.trace.instrumentation.aws.v1.sqs.SqsDecorator.BROKER_DECORATE; @@ -18,6 +20,7 @@ import com.amazonaws.services.sqs.model.Message; import datadog.trace.api.Config; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -43,7 +46,14 @@ public boolean hasNext() { boolean moreMessages = delegate.hasNext(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreMessages; } @@ -57,7 +67,14 @@ public Message next() { protected void startNewMessageSpan(Message message) { try { - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); + if (prevSpan != null) { + prevSpan.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } if (message != null) { AgentSpan queueSpan = null; if (batchContext == null) { @@ -92,7 +109,14 @@ protected void startNewMessageSpan(Message message) { CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, queueUrl); - activateNext(span); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan previous = spanFromContext(span.swap()); + if (previous != null) { + previous.finishWithEndToEnd(); + } + } else { + activateNext(span); + } if (queueSpan != null) { BROKER_DECORATE.beforeFinish(queueSpan); queueSpan.finish(); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java index 09a2e342b5f..63564dbd35e 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java @@ -1,8 +1,12 @@ package datadog.trace.instrumentation.aws.v1.sqs; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import com.amazonaws.services.sqs.model.Message; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.ListIterator; public class TracingListIterator extends TracingIterator> @@ -17,7 +21,14 @@ public boolean hasPrevious() { boolean moreMessages = delegate.hasPrevious(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreMessages; } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy index 43bff3872e3..5a463291cec 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy @@ -1,5 +1,6 @@ import static datadog.trace.agent.test.utils.TraceUtils.basicSpan import static java.nio.charset.StandardCharsets.UTF_8 +import datadog.trace.api.config.TraceInstrumentationConfig import com.amazon.sqs.javamessaging.ProviderConfiguration import com.amazon.sqs.javamessaging.SQSConnectionFactory @@ -726,4 +727,10 @@ class SqsClientV1DataStreamsForkedTest extends SqsClientTest { } } - +class SqsClientV0ContextSwapForkedTest extends SqsClientV0Test { + @Override + protected void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + } +} diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java index 3991bfc63b1..9d1a68796ac 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java @@ -7,6 +7,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.bootstrap.instrumentation.api.URIUtils.urlFileName; import static datadog.trace.instrumentation.aws.v2.sqs.MessageExtractAdapter.GETTER; import static datadog.trace.instrumentation.aws.v2.sqs.SqsDecorator.BROKER_DECORATE; @@ -17,6 +19,7 @@ import static java.util.concurrent.TimeUnit.MILLISECONDS; import datadog.trace.api.Config; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -45,7 +48,14 @@ public boolean hasNext() { boolean moreMessages = delegate.hasNext(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreMessages; } @@ -59,7 +69,14 @@ public Message next() { protected void startNewMessageSpan(Message message) { try { - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); + if (prevSpan != null) { + prevSpan.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } if (message != null) { AgentSpan queueSpan = null; if (batchContext == null) { @@ -94,7 +111,14 @@ protected void startNewMessageSpan(Message message) { CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, queueUrl, requestId); - activateNext(span); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan previous = spanFromContext(span.swap()); + if (previous != null) { + previous.finishWithEndToEnd(); + } + } else { + activateNext(span); + } if (queueSpan != null) { BROKER_DECORATE.beforeFinish(queueSpan); queueSpan.finish(); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java index 184d05815fc..5b192c4286e 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java @@ -1,7 +1,11 @@ package datadog.trace.instrumentation.aws.v2.sqs; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.ListIterator; import software.amazon.awssdk.services.sqs.model.Message; @@ -17,7 +21,14 @@ public boolean hasPrevious() { boolean moreMessages = delegate.hasPrevious(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreMessages; } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy index 067397524e7..a159691bfa0 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy @@ -1,5 +1,6 @@ import static datadog.trace.agent.test.utils.TraceUtils.basicSpan import static java.nio.charset.StandardCharsets.UTF_8 +import datadog.trace.api.config.TraceInstrumentationConfig import com.amazon.sqs.javamessaging.ProviderConfiguration import com.amazon.sqs.javamessaging.SQSConnectionFactory @@ -574,4 +575,12 @@ class SqsClientV1DataStreamsForkedTest extends SqsClientTest { } } +class SqsClientV0ContextSwapForkedTest extends SqsClientV0Test { + @Override + protected void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + } +} + diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java index 9baf17aebcc..64e8db97856 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java @@ -97,7 +97,7 @@ public static MessageConsumerState beforeReceive(@Advice.This final MessageConsu final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { CONSUMER_DECORATE.beforeFinish(span); - span.finish(); + span.finishWithEndToEnd(); } } else { closePrevious(finishSpan); @@ -178,7 +178,7 @@ public static void afterReceive( final AgentSpan previous = spanFromContext(span.swap()); if (previous != null) { CONSUMER_DECORATE.beforeFinish(previous); - previous.finish(); + previous.finishWithEndToEnd(); } } else { activateNext(span); // scope is left open until next message or it times out @@ -210,7 +210,7 @@ public static void beforeClose(@Advice.This final MessageConsumer consumer) { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { CONSUMER_DECORATE.beforeFinish(span); - span.finish(); + span.finishWithEndToEnd(); } } else { closePrevious(finishSpan); diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java index e0076ce5aed..84f85308cd1 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java @@ -9,6 +9,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.BROKER_DECORATE; import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.KAFKA_DELIVER; import static datadog.trace.instrumentation.kafka_clients.KafkaDecorator.TIME_IN_QUEUE_ENABLED; @@ -22,6 +24,7 @@ import datadog.context.propagation.Propagator; import datadog.context.propagation.Propagators; import datadog.trace.api.Config; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.api.datastreams.DataStreamsTransactionExtractor; @@ -65,7 +68,14 @@ public boolean hasNext() { boolean moreRecords = delegateIterator.hasNext(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreRecords; } @@ -79,7 +89,14 @@ public boolean hasNext() { protected void startNewRecordSpan(ConsumerRecord val) { try { - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); + if (prevSpan != null) { + prevSpan.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } AgentSpan span, queueSpan = null; if (val != null) { if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) { @@ -126,7 +143,14 @@ protected void startNewRecordSpan(ConsumerRecord val) { } decorator.afterStart(span); decorator.onConsume(span, val, group, bootstrapServers); - activateNext(span); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan previous = spanFromContext(span.swap()); + if (previous != null) { + previous.finishWithEndToEnd(); + } + } else { + activateNext(span); + } if (null != queueSpan) { queueSpan.finish(); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java index 5d5e726e806..dab015cf6f3 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java @@ -1,7 +1,11 @@ package datadog.trace.instrumentation.kafka_clients; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.ListIterator; import org.apache.kafka.clients.consumer.ConsumerRecord; @@ -26,7 +30,14 @@ public boolean hasPrevious() { boolean moreRecords = delegateIterator.hasPrevious(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreRecords; } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy index c4c5f52afee..a012a2d48d9 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy @@ -1,5 +1,6 @@ import datadog.trace.api.datastreams.DataStreamsTags import datadog.trace.api.datastreams.DataStreamsTransactionExtractor +import datadog.trace.api.config.TraceInstrumentationConfig import datadog.trace.instrumentation.kafka_common.ClusterIdHolder import static datadog.trace.agent.test.utils.TraceUtils.basicSpan @@ -1498,3 +1499,11 @@ class KafkaClientDataStreamsDisabledForkedTest extends KafkaClientTestBase { return false } } + +class KafkaClientContextSwapForkedTest extends KafkaClientV0ForkedTest { + @Override + void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + } +} diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java index ceba52acac8..9a15b7a0ea3 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java @@ -9,6 +9,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static datadog.trace.instrumentation.kafka_clients38.TextMapExtractAdapter.GETTER; import static datadog.trace.instrumentation.kafka_clients38.TextMapInjectAdapter.SETTER; import static java.util.concurrent.TimeUnit.MILLISECONDS; @@ -16,6 +18,7 @@ import datadog.context.propagation.Propagator; import datadog.context.propagation.Propagators; import datadog.trace.api.Config; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.api.datastreams.DataStreamsTags; import datadog.trace.api.datastreams.DataStreamsTransactionExtractor; @@ -61,7 +64,14 @@ public boolean hasNext() { boolean moreRecords = delegateIterator.hasNext(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreRecords; } @@ -75,7 +85,14 @@ public boolean hasNext() { protected void startNewRecordSpan(ConsumerRecord val) { try { - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); + if (prevSpan != null) { + prevSpan.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } AgentSpan span, queueSpan = null; if (val != null) { if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) { @@ -125,7 +142,14 @@ protected void startNewRecordSpan(ConsumerRecord val) { } decorator.afterStart(span); decorator.onConsume(span, val, group, bootstrapServers); - activateNext(span); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan previous = spanFromContext(span.swap()); + if (previous != null) { + previous.finishWithEndToEnd(); + } + } else { + activateNext(span); + } if (null != queueSpan) { queueSpan.finish(); } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java index 9014ff51966..e122fd45efc 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java @@ -1,7 +1,11 @@ package datadog.trace.instrumentation.kafka_clients38; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.util.ListIterator; import org.apache.kafka.clients.consumer.ConsumerRecord; @@ -26,7 +30,14 @@ public boolean hasPrevious() { boolean moreRecords = delegateIterator.hasPrevious(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - closePrevious(true); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(true); + } } return moreRecords; } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy index 8e53df883bf..5cb96efd097 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy @@ -1,6 +1,7 @@ import datadog.trace.agent.test.asserts.TraceAssert import datadog.trace.agent.test.naming.VersionedNamingTestBase import datadog.trace.api.Config +import datadog.trace.api.config.TraceInstrumentationConfig import datadog.trace.api.DDTags import datadog.trace.api.datastreams.DataStreamsTags import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags @@ -1206,3 +1207,11 @@ class KafkaClientDataStreamsDisabledForkedTest extends KafkaClientTestBase { return false } } + +class KafkaClientContextSwapForkedTest extends KafkaClientV0ForkedTest { + @Override + void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + } +} diff --git a/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java index cda0ab24d1b..e6827d6cc24 100644 --- a/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java @@ -2,13 +2,17 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getRootContext; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.jms.MessageConsumerState; import java.util.Map; import javax.jms.MessageConsumer; @@ -55,7 +59,14 @@ public static void afterExecute(@Advice.Argument(2) final MessageConsumer consum .get(consumer); if (null != consumerState) { boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge(); - closePrevious(finishSpan); + if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + final AgentSpan span = spanFromContext(getRootContext().swap()); + if (span != null) { + span.finishWithEndToEnd(); + } + } else { + closePrevious(finishSpan); + } if (finishSpan) { consumerState.finishTimeInQueueSpan(false); } diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 0d24d356044..ae9d6ab54e7 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -11160,6 +11160,14 @@ "default": "2000", "aliases": [] } + ], + "MESSAGING_CONTEXT_SWAP_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "false", + "aliases": [] + } ] }, "deprecations": {} From cd633081003483a8dd4b4c8cbe5cd4f7e8f2d815 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 6 Mar 2026 16:45:01 +0100 Subject: [PATCH 04/12] add test for jakarta jms --- .../jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy b/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy index a3f0d05ae03..bcb9f07a94b 100644 --- a/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy +++ b/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy @@ -1,5 +1,6 @@ import com.google.common.io.Files import datadog.trace.agent.test.InstrumentationSpecification +import datadog.trace.api.config.TraceInstrumentationConfig import datadog.trace.agent.test.asserts.ListWriterAssert import datadog.trace.api.DDSpanTypes import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags @@ -254,3 +255,11 @@ class JMS2Test extends InstrumentationSpecification { } } } + +class JMS2ContextSwapForkedTest extends JMS2Test { + @Override + protected void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + } +} From 0b0e51f3ef36c6c2c769f506d130bbf74dac1d3d Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 6 Mar 2026 16:58:33 +0100 Subject: [PATCH 05/12] use full config name on the json file --- metadata/supported-configurations.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index ae9d6ab54e7..18bf0442dd7 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -11161,7 +11161,7 @@ "aliases": [] } ], - "MESSAGING_CONTEXT_SWAP_ENABLED": [ + "DD_MESSAGING_CONTEXT_SWAP_ENABLED": [ { "version": "A", "type": "boolean", From 5fac95f53e4e1d18f02fb6bb4df71bd25949c65c Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Fri, 6 Mar 2026 18:22:13 +0100 Subject: [PATCH 06/12] force kafka consumer scope to close in the dsm test --- .../src/test/groovy/KafkaClientTestBase.groovy | 1 + 1 file changed, 1 insertion(+) diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy index a012a2d48d9..f52976d5239 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy @@ -1149,6 +1149,7 @@ abstract class KafkaClientTestBase extends VersionedNamingTestBase { def recs = pollResult.records(new TopicPartition(SHARED_TOPIC, kafkaPartition)).iterator() recs.hasNext() recs.next().value() == "test-dsm-consume-transaction" + !recs.hasNext() // The consume span is created by TracingIterator when iterating over records // Find the consumer span with the DSM transaction tags From 22451e4930a7d98112ec1297c7e92821c655bbdb Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Mon, 9 Mar 2026 10:12:59 +0100 Subject: [PATCH 07/12] Use DD_LEGACY_CONTEXT_MANAGER_ENABLED --- .../aws/v1/sqs/TracingIterator.java | 14 +++++++------- .../aws/v1/sqs/TracingListIterator.java | 6 +++--- .../src/test/groovy/SqsClientTest.groovy | 2 +- .../aws/v2/sqs/TracingIterator.java | 18 +++++++++--------- .../aws/v2/sqs/TracingListIterator.java | 6 +++--- .../src/test/groovy/SqsClientTest.groovy | 2 +- .../src/test/groovy/JMS2Test.groovy | 2 +- .../jms/JMSMessageConsumerInstrumentation.java | 18 +++++++++--------- .../src/test/groovy/JMS1Test.groovy | 4 ++-- .../kafka_clients/TracingIterator.java | 18 +++++++++--------- .../kafka_clients/TracingListIterator.java | 6 +++--- .../src/test/groovy/KafkaClientTestBase.groovy | 2 +- .../kafka_clients38/TracingIterator.java | 18 +++++++++--------- .../kafka_clients38/TracingListIterator.java | 6 +++--- .../src/test/groovy/KafkaClientTestBase.groovy | 2 +- ...essageListenerContainerInstrumentation.java | 6 +++--- .../api/config/TraceInstrumentationConfig.java | 2 +- .../java/datadog/trace/core/CoreTracer.java | 4 ++-- .../datadog/trace/api/InstrumenterConfig.java | 12 ++++++------ metadata/supported-configurations.json | 4 ++-- 20 files changed, 76 insertions(+), 76 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java index 5cac40c854b..ed358eccd50 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java @@ -46,7 +46,7 @@ public boolean hasNext() { boolean moreMessages = delegate.hasNext(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); @@ -67,13 +67,13 @@ public Message next() { protected void startNewMessageSpan(Message message) { try { - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); if (prevSpan != null) { prevSpan.finishWithEndToEnd(); } - } else { - closePrevious(true); } if (message != null) { AgentSpan queueSpan = null; @@ -109,13 +109,13 @@ protected void startNewMessageSpan(Message message) { CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, queueUrl); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + activateNext(span); + } else { final AgentSpan previous = spanFromContext(span.swap()); if (previous != null) { previous.finishWithEndToEnd(); } - } else { - activateNext(span); } if (queueSpan != null) { BROKER_DECORATE.beforeFinish(queueSpan); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java index 63564dbd35e..bc1df759b2c 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingListIterator.java @@ -21,13 +21,13 @@ public boolean hasPrevious() { boolean moreMessages = delegate.hasPrevious(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreMessages; diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy index 5a463291cec..0fca67d72e6 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/test/groovy/SqsClientTest.groovy @@ -731,6 +731,6 @@ class SqsClientV0ContextSwapForkedTest extends SqsClientV0Test { @Override protected void configurePreAgent() { super.configurePreAgent() - injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false") } } diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java index 9d1a68796ac..860e6008113 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java @@ -48,13 +48,13 @@ public boolean hasNext() { boolean moreMessages = delegate.hasNext(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreMessages; @@ -69,13 +69,13 @@ public Message next() { protected void startNewMessageSpan(Message message) { try { - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); if (prevSpan != null) { prevSpan.finishWithEndToEnd(); } - } else { - closePrevious(true); } if (message != null) { AgentSpan queueSpan = null; @@ -111,13 +111,13 @@ protected void startNewMessageSpan(Message message) { CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, queueUrl, requestId); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + activateNext(span); + } else { final AgentSpan previous = spanFromContext(span.swap()); if (previous != null) { previous.finishWithEndToEnd(); } - } else { - activateNext(span); } if (queueSpan != null) { BROKER_DECORATE.beforeFinish(queueSpan); diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java index 5b192c4286e..7b2e19b9be4 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingListIterator.java @@ -21,13 +21,13 @@ public boolean hasPrevious() { boolean moreMessages = delegate.hasPrevious(); if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreMessages; diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy index a159691bfa0..e9419c46134 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-2.0/src/test/groovy/SqsClientTest.groovy @@ -579,7 +579,7 @@ class SqsClientV0ContextSwapForkedTest extends SqsClientV0Test { @Override protected void configurePreAgent() { super.configurePreAgent() - injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false") } } diff --git a/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy b/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy index bcb9f07a94b..fe3229077a9 100644 --- a/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy +++ b/dd-java-agent/instrumentation/jms/jakarta-jms-3.0/src/test/groovy/JMS2Test.groovy @@ -260,6 +260,6 @@ class JMS2ContextSwapForkedTest extends JMS2Test { @Override protected void configurePreAgent() { super.configurePreAgent() - injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false") } } diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java index 64e8db97856..1ffd07f3ed5 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java @@ -93,14 +93,14 @@ public static MessageConsumerState beforeReceive(@Advice.This final MessageConsu } boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge(); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(finishSpan); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { CONSUMER_DECORATE.beforeFinish(span); span.finishWithEndToEnd(); } - } else { - closePrevious(finishSpan); } if (finishSpan) { consumerState.finishTimeInQueueSpan(false); @@ -174,14 +174,14 @@ public static void afterReceive( CONSUMER_DECORATE.onError(span, throwable); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + activateNext(span); // scope is left open until next message or it times out + } else { final AgentSpan previous = spanFromContext(span.swap()); if (previous != null) { CONSUMER_DECORATE.beforeFinish(previous); previous.finishWithEndToEnd(); } - } else { - activateNext(span); // scope is left open until next message or it times out } JMSLogger.logIterationSpan(span); @@ -206,14 +206,14 @@ public static void beforeClose(@Advice.This final MessageConsumer consumer) { .get(consumer); if (null != consumerState) { boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge(); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(finishSpan); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { CONSUMER_DECORATE.beforeFinish(span); span.finishWithEndToEnd(); } - } else { - closePrevious(finishSpan); } if (finishSpan) { consumerState.finishTimeInQueueSpan(true); diff --git a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy index 33cddadf5f4..5279ba4e2a7 100644 --- a/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy +++ b/dd-java-agent/instrumentation/jms/javax-jms-1.1/src/test/groovy/JMS1Test.groovy @@ -1,4 +1,4 @@ -import static datadog.trace.api.config.TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED +import static datadog.trace.api.config.TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED import static org.junit.jupiter.api.Assumptions.assumeTrue import datadog.trace.agent.test.asserts.ListWriterAssert @@ -1056,7 +1056,7 @@ class JMS1V0Test extends JMS1Test { class JMSContextSwapForkedTest extends JMS1V0Test { @Override protected void configurePreAgent() { - injectSysConfig(MESSAGING_CONTEXT_SWAP_ENABLED, "true") + injectSysConfig(LEGACY_CONTEXT_MANAGER_ENABLED, "false") } @Override diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java index 84f85308cd1..808e3d9b5eb 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java @@ -68,13 +68,13 @@ public boolean hasNext() { boolean moreRecords = delegateIterator.hasNext(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreRecords; @@ -89,13 +89,13 @@ public boolean hasNext() { protected void startNewRecordSpan(ConsumerRecord val) { try { - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); if (prevSpan != null) { prevSpan.finishWithEndToEnd(); } - } else { - closePrevious(true); } AgentSpan span, queueSpan = null; if (val != null) { @@ -143,13 +143,13 @@ protected void startNewRecordSpan(ConsumerRecord val) { } decorator.afterStart(span); decorator.onConsume(span, val, group, bootstrapServers); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + activateNext(span); + } else { final AgentSpan previous = spanFromContext(span.swap()); if (previous != null) { previous.finishWithEndToEnd(); } - } else { - activateNext(span); } if (null != queueSpan) { queueSpan.finish(); diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java index dab015cf6f3..455a37dce8b 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingListIterator.java @@ -30,13 +30,13 @@ public boolean hasPrevious() { boolean moreRecords = delegateIterator.hasPrevious(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreRecords; diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy index f52976d5239..cf0ff225f8c 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-0.11/src/test/groovy/KafkaClientTestBase.groovy @@ -1505,6 +1505,6 @@ class KafkaClientContextSwapForkedTest extends KafkaClientV0ForkedTest { @Override void configurePreAgent() { super.configurePreAgent() - injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false") } } diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java index 9a15b7a0ea3..cb1f4f372ec 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java @@ -64,13 +64,13 @@ public boolean hasNext() { boolean moreRecords = delegateIterator.hasNext(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreRecords; @@ -85,13 +85,13 @@ public boolean hasNext() { protected void startNewRecordSpan(ConsumerRecord val) { try { - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan prevSpan = spanFromContext(getRootContext().swap()); if (prevSpan != null) { prevSpan.finishWithEndToEnd(); } - } else { - closePrevious(true); } AgentSpan span, queueSpan = null; if (val != null) { @@ -142,13 +142,13 @@ protected void startNewRecordSpan(ConsumerRecord val) { } decorator.afterStart(span); decorator.onConsume(span, val, group, bootstrapServers); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + activateNext(span); + } else { final AgentSpan previous = spanFromContext(span.swap()); if (previous != null) { previous.finishWithEndToEnd(); } - } else { - activateNext(span); } if (null != queueSpan) { queueSpan.finish(); diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java index e122fd45efc..de39c566576 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingListIterator.java @@ -30,13 +30,13 @@ public boolean hasPrevious() { boolean moreRecords = delegateIterator.hasPrevious(); if (!moreRecords) { // no more records, use this as a signal to close the last iteration scope - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreRecords; diff --git a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy index 5cb96efd097..220ee7f7f65 100644 --- a/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy +++ b/dd-java-agent/instrumentation/kafka/kafka-clients-3.8/src/test/groovy/KafkaClientTestBase.groovy @@ -1212,6 +1212,6 @@ class KafkaClientContextSwapForkedTest extends KafkaClientV0ForkedTest { @Override void configurePreAgent() { super.configurePreAgent() - injectSysConfig(TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED, "true") + injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false") } } diff --git a/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java b/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java index e6827d6cc24..b72029f3486 100644 --- a/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java +++ b/dd-java-agent/instrumentation/spring/spring-jms-3.1/src/main/java/datadog/trace/instrumentation/springjms/AbstractPollingMessageListenerContainerInstrumentation.java @@ -59,13 +59,13 @@ public static void afterExecute(@Advice.Argument(2) final MessageConsumer consum .get(consumer); if (null != consumerState) { boolean finishSpan = consumerState.getSessionState().isAutoAcknowledge(); - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(finishSpan); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(finishSpan); } if (finishSpan) { consumerState.finishTimeInQueueSpan(false); 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 b66bd67972f..092d97cea5a 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 @@ -201,7 +201,7 @@ public final class TraceInstrumentationConfig { public static final String TRACE_RESOURCE_RENAMING_ALWAYS_SIMPLIFIED_ENDPOINT = "trace.resource.renaming.always.simplified.endpoint"; - public static final String MESSAGING_CONTEXT_SWAP_ENABLED = "messaging.context.swap.enabled"; + public static final String LEGACY_CONTEXT_MANAGER_ENABLED = "legacy.context-manager.enabled"; private TraceInstrumentationConfig() {} } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 32adab8cc45..bb54fa81ee9 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -1140,7 +1140,7 @@ public void setAsyncPropagationEnabled(boolean asyncPropagationEnabled) { @Override public void closePrevious(boolean finishSpan) { - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (!InstrumenterConfig.get().isLegacyContextManagerEnabled()) { throw new IllegalStateException( "closePrevious must not be called when context swap based logic is enabled"); } @@ -1149,7 +1149,7 @@ public void closePrevious(boolean finishSpan) { @Override public AgentScope activateNext(AgentSpan span) { - if (InstrumenterConfig.get().isMessagingContextSwapEnabled()) { + if (!InstrumenterConfig.get().isLegacyContextManagerEnabled()) { throw new IllegalStateException( "activateNext must not be called when context swap based logic is enabled"); } 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 38696171ca8..4fc31d861bd 100644 --- a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java +++ b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java @@ -57,8 +57,8 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.JDBC_CONNECTION_CLASS_NAME; import static datadog.trace.api.config.TraceInstrumentationConfig.JDBC_POOL_WAITING_ENABLED; import static datadog.trace.api.config.TraceInstrumentationConfig.JDBC_PREPARED_STATEMENT_CLASS_NAME; +import static datadog.trace.api.config.TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED; import static datadog.trace.api.config.TraceInstrumentationConfig.MEASURE_METHODS; -import static datadog.trace.api.config.TraceInstrumentationConfig.MESSAGING_CONTEXT_SWAP_ENABLED; import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_CACHE_CONFIG; import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_CACHE_DIR; import static datadog.trace.api.config.TraceInstrumentationConfig.RESOLVER_NAMES_ARE_UNIQUE; @@ -213,7 +213,7 @@ public class InstrumenterConfig { private final boolean apiSecurityEndpointCollectionEnabled; private final boolean appLogsCollectionEnabled; - private final boolean messagingContextSwapEnabled; + private final boolean legacyContextManagerEnabled; static { // Bind telemetry collector to config module before initializing ConfigProvider @@ -366,7 +366,7 @@ private InstrumenterConfig() { appLogsCollectionEnabled = configProvider.getBoolean(APP_LOGS_COLLECTION_ENABLED, DEFAULT_APP_LOGS_COLLECTION_ENABLED); - messagingContextSwapEnabled = configProvider.getBoolean(MESSAGING_CONTEXT_SWAP_ENABLED, false); + legacyContextManagerEnabled = configProvider.getBoolean(LEGACY_CONTEXT_MANAGER_ENABLED, true); } public boolean isCodeOriginEnabled() { @@ -686,8 +686,8 @@ public boolean isAppLogsCollectionEnabled() { return appLogsCollectionEnabled; } - public boolean isMessagingContextSwapEnabled() { - return messagingContextSwapEnabled; + public boolean isLegacyContextManagerEnabled() { + return legacyContextManagerEnabled; } // This has to be placed after all other static fields to give them a chance to initialize @@ -810,7 +810,7 @@ public String toString() { + ", apiSecurityEndpointCollectionEnabled=" + apiSecurityEndpointCollectionEnabled + ", messagingContextSwapEnabled=" - + messagingContextSwapEnabled + + legacyContextManagerEnabled + '}'; } } diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 18bf0442dd7..9928be865b6 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -11161,11 +11161,11 @@ "aliases": [] } ], - "DD_MESSAGING_CONTEXT_SWAP_ENABLED": [ + "DD_LEGACY_CONTEXT_MANAGER_ENABLED": [ { "version": "A", "type": "boolean", - "default": "false", + "default": "true", "aliases": [] } ] From 98a273e7f4bed633ad602790bcb209eb4bb5fb73 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Mon, 9 Mar 2026 10:17:13 +0100 Subject: [PATCH 08/12] Update config print --- .../src/main/java/datadog/trace/api/InstrumenterConfig.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 4fc31d861bd..8097cf4726b 100644 --- a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java +++ b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java @@ -809,7 +809,7 @@ public String toString() { + dataJobsEnabled + ", apiSecurityEndpointCollectionEnabled=" + apiSecurityEndpointCollectionEnabled - + ", messagingContextSwapEnabled=" + + ", legacyContextManagerEnabled=" + legacyContextManagerEnabled + '}'; } From ae97a09146cd027415663ef9248af04f404d81e0 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Mon, 9 Mar 2026 11:00:02 +0100 Subject: [PATCH 09/12] missed reverse --- .../trace/instrumentation/aws/v1/sqs/TracingIterator.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java index ed358eccd50..18b1f485dd9 100644 --- a/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java +++ b/dd-java-agent/instrumentation/aws-java/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java @@ -47,12 +47,12 @@ public boolean hasNext() { if (!moreMessages) { // no more messages, use this as a signal to close the last iteration scope if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + closePrevious(true); + } else { final AgentSpan span = spanFromContext(getRootContext().swap()); if (span != null) { span.finishWithEndToEnd(); } - } else { - closePrevious(true); } } return moreMessages; From f592844268030bc484db06e6ac67b474e38003aa Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 10 Mar 2026 09:41:53 +0100 Subject: [PATCH 10/12] Use context swap on akka and pekko --- .../akka23Test/groovy/AkkaActorTest.groovy | 9 ++++++ .../AkkaActorCellInstrumentation.java | 32 +++++++++++++------ .../AkkaMailboxInstrumentation.java | 24 ++++++++++---- .../PekkoActorCellInstrumentation.java | 32 +++++++++++++------ .../PekkoMailboxInstrumentation.java | 24 ++++++++++---- .../src/test/groovy/PekkoActorTest.groovy | 8 +++++ 6 files changed, 99 insertions(+), 30 deletions(-) diff --git a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/akka23Test/groovy/AkkaActorTest.groovy b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/akka23Test/groovy/AkkaActorTest.groovy index a7c4b8a8eef..2f958f5bdbf 100644 --- a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/akka23Test/groovy/AkkaActorTest.groovy +++ b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/akka23Test/groovy/AkkaActorTest.groovy @@ -1,4 +1,5 @@ import datadog.trace.agent.test.InstrumentationSpecification +import datadog.trace.api.config.TraceInstrumentationConfig import datadog.trace.bootstrap.instrumentation.api.Tags import spock.lang.Shared @@ -82,3 +83,11 @@ class AkkaActorTest extends InstrumentationSpecification { } } } + +class AkkaActorContextSwapTest extends AkkaActorTest { + @Override + void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false") + } +} diff --git a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java index 03b2ccc024a..895a837a86a 100644 --- a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java @@ -3,13 +3,16 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import akka.dispatch.Envelope; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils; @@ -54,24 +57,35 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) { + public static void enter( + @Advice.Argument(value = 0) Envelope envelope, + @Advice.Local("taskScope") AgentScope taskScope, + @Advice.Local("checkpointContext") Context checkpointContext) { // do this before checkpointing, as the envelope's task scope may already be active - AgentScope taskScope = + taskScope = AdviceUtils.startTaskScope( InstrumentationContext.get(Envelope.class, State.class), envelope); - // remember the currently active scope so we can roll back to this point - checkpointActiveForRollback(); - - return taskScope; + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + // remember the currently active scope so we can roll back to this point + checkpointActiveForRollback(); + } else { + checkpointContext = getCurrentContext().swap(); + } } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Enter AgentScope taskScope) { + public static void exit( + @Advice.Local("taskScope") AgentScope taskScope, + @Advice.Local("checkpointContext") Context checkpointContext) { - // Clean up any leaking scopes from akka-streams/akka-http etc. - rollbackActiveToCheckpoint(); + if (checkpointContext == null) { + // Clean up any leaking scopes from akka-streams/akka-http etc. + rollbackActiveToCheckpoint(); + } else { + checkpointContext.swap(); + } // close envelope's task scope if we previously started it if (taskScope != null) { diff --git a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java index 2de8d97047d..6f1eae44b34 100644 --- a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java @@ -3,13 +3,16 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static java.util.Collections.singletonList; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; import java.util.Collection; import java.util.EnumMap; @@ -59,15 +62,24 @@ public void methodAdvice(MethodTransformer transformer) { */ public static final class SuppressMailboxRunAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void enter() { - // remember the currently active scope so we can roll back to this point - checkpointActiveForRollback(); + public static Context enter() { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + // remember the currently active scope so we can roll back to this point + checkpointActiveForRollback(); + return null; + } else { + return getCurrentContext().swap(); + } } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit() { - // Clean up any leaking scopes from akka-streams/akka-http etc. - rollbackActiveToCheckpoint(); + public static void exit(@Advice.Enter final Context checkpointContext) { + if (checkpointContext == null) { + // Clean up any leaking scopes from akka-streams/akka-http etc. + rollbackActiveToCheckpoint(); + } else { + checkpointContext.swap(); + } } } } diff --git a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java index 4f7dda119ed..c730a3c54c0 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java @@ -3,12 +3,15 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils; @@ -54,24 +57,35 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) { + public static void enter( + @Advice.Argument(value = 0) Envelope envelope, + @Advice.Local("taskScope") AgentScope taskScope, + @Advice.Local("checkpointContext") Context checkpointContext) { // do this before checkpointing, as the envelope's task scope may already be active - AgentScope taskScope = + taskScope = AdviceUtils.startTaskScope( InstrumentationContext.get(Envelope.class, State.class), envelope); - // remember the currently active scope so we can roll back to this point - checkpointActiveForRollback(); - - return taskScope; + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + // remember the currently active scope so we can roll back to this point + checkpointActiveForRollback(); + } else { + checkpointContext = getCurrentContext().swap(); + } } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit(@Advice.Enter AgentScope taskScope) { + public static void exit( + @Advice.Local("taskScope") AgentScope taskScope, + @Advice.Local("checkpointContext") Context checkpointContext) { - // Clean up any leaking scopes from pekko-streams/pekko-http etc. - rollbackActiveToCheckpoint(); + if (checkpointContext == null) { + // Clean up any leaking scopes from pekko-streams/pekko-http etc. + rollbackActiveToCheckpoint(); + } else { + checkpointContext.swap(); + } // close envelope's task scope if we previously started it if (taskScope != null) { diff --git a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java index aa10d92432b..83c144a8c36 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java @@ -3,13 +3,16 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint; +import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext; import static java.util.Collections.singletonList; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; import java.util.Collection; import java.util.EnumMap; @@ -59,15 +62,24 @@ public void methodAdvice(MethodTransformer transformer) { */ public static final class SuppressMailboxRunAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void enter() { - // remember the currently active scope so we can roll back to this point - checkpointActiveForRollback(); + public static Context enter() { + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + // remember the currently active scope so we can roll back to this point + checkpointActiveForRollback(); + return null; + } else { + return getCurrentContext().swap(); + } } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) - public static void exit() { - // Clean up any leaking scopes from pekko-streams/pekko-http etc. - rollbackActiveToCheckpoint(); + public static void exit(@Advice.Enter final Context checkpointContext) { + if (checkpointContext == null) { + // Clean up any leaking scopes from pekko-streams/pekko-http etc. + rollbackActiveToCheckpoint(); + } else { + checkpointContext.swap(); + } } } } diff --git a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/test/groovy/PekkoActorTest.groovy b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/test/groovy/PekkoActorTest.groovy index 5efe546cee3..7b04c19422b 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/test/groovy/PekkoActorTest.groovy +++ b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/test/groovy/PekkoActorTest.groovy @@ -121,3 +121,11 @@ class PekkoActorTest extends InstrumentationSpecification { } } } + +class PekkoActorContextSwapTest extends PekkoActorTest { + @Override + void configurePreAgent() { + super.configurePreAgent() + injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false") + } +} From be075a4ab6164a6ce37d6329c3af0f9fba5e092b Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 10 Mar 2026 09:45:37 +0100 Subject: [PATCH 11/12] add guard --- .../src/main/java/datadog/trace/core/CoreTracer.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index bb54fa81ee9..a9a34f810ff 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -1171,11 +1171,19 @@ public AgentSpan activeSpan() { @Override public void checkpointActiveForRollback() { + if (!InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + throw new IllegalStateException( + "checkpointActiveForRollback must not be called when context swap based logic is enabled"); + } this.scopeManager.checkpointActiveForRollback(); } @Override public void rollbackActiveToCheckpoint() { + if (!InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + throw new IllegalStateException( + "rollbackActiveToCheckpoint must not be called when context swap based logic is enabled"); + } this.scopeManager.rollbackActiveToCheckpoint(); } From 63aaf0567bc8a6cf444ab0f06308324f677619db Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 10 Mar 2026 12:22:24 +0100 Subject: [PATCH 12/12] suggestions --- .../akka/concurrent/AkkaActorCellInstrumentation.java | 11 +++++------ .../concurrent/PekkoActorCellInstrumentation.java | 11 +++++------ 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java index 895a837a86a..f799267fb2d 100644 --- a/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java @@ -57,10 +57,9 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void enter( + public static Context enter( @Advice.Argument(value = 0) Envelope envelope, - @Advice.Local("taskScope") AgentScope taskScope, - @Advice.Local("checkpointContext") Context checkpointContext) { + @Advice.Local("taskScope") AgentScope taskScope) { // do this before checkpointing, as the envelope's task scope may already be active taskScope = @@ -70,15 +69,15 @@ public static void enter( if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { // remember the currently active scope so we can roll back to this point checkpointActiveForRollback(); + return null; } else { - checkpointContext = getCurrentContext().swap(); + return getCurrentContext().swap(); } } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void exit( - @Advice.Local("taskScope") AgentScope taskScope, - @Advice.Local("checkpointContext") Context checkpointContext) { + @Advice.Local("taskScope") AgentScope taskScope, @Advice.Enter Context checkpointContext) { if (checkpointContext == null) { // Clean up any leaking scopes from akka-streams/akka-http etc. diff --git a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java index c730a3c54c0..e33bae813e6 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java @@ -57,10 +57,9 @@ public void methodAdvice(MethodTransformer transformer) { */ public static class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static void enter( + public static Context enter( @Advice.Argument(value = 0) Envelope envelope, - @Advice.Local("taskScope") AgentScope taskScope, - @Advice.Local("checkpointContext") Context checkpointContext) { + @Advice.Local("taskScope") AgentScope taskScope) { // do this before checkpointing, as the envelope's task scope may already be active taskScope = @@ -70,15 +69,15 @@ public static void enter( if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { // remember the currently active scope so we can roll back to this point checkpointActiveForRollback(); + return null; } else { - checkpointContext = getCurrentContext().swap(); + return getCurrentContext().swap(); } } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void exit( - @Advice.Local("taskScope") AgentScope taskScope, - @Advice.Local("checkpointContext") Context checkpointContext) { + @Advice.Local("taskScope") AgentScope taskScope, @Advice.Enter Context checkpointContext) { if (checkpointContext == null) { // Clean up any leaking scopes from pekko-streams/pekko-http etc.