Optionally use context swap on akka and pekko (#10778)

amarziali · devflow.devflow-routing-intake · web-flow · commit 28a1a34033e7 · 2026-03-11T00:00:55.000Z
wip

Changes

migrate everything else

add test for jakarta jms

use full config name on the json file

force kafka consumer scope to close in the dsm test

Merge branch 'master' into andrea.marziali/swap-messaging

Use DD_LEGACY_CONTEXT_MANAGER_ENABLED

Update config print

missed reverse

Use context swap on akka and pekko

add guard

suggestions

Merge remote-tracking branch 'origin/master' into andrea.marziali/akka-swap

Merge branch 'master' into andrea.marziali/akka-swap

Co-authored-by: devflow.devflow-routing-intake &lt;devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io&gt;
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
@@ -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
@@ -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,34 @@ 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 Context enter(
+        @Advice.Argument(value = 0) Envelope envelope,
+        @Advice.Local("taskScope") AgentScope taskScope) {
 
       // 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();
+        return null;
+      } else {
+        return 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.Enter 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
@@ -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/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
@@ -98,6 +98,7 @@ public static MessageConsumerState beforeReceive(@Advice.This final MessageConsu
       } else {
         final AgentSpan previousSpan = spanFromContext(getRootContext().swap());
         if (previousSpan != null) {
+          CONSUMER_DECORATE.beforeFinish(previousSpan);
           previousSpan.finishWithEndToEnd();
         }
       }
@@ -178,6 +179,7 @@ public static void afterReceive(
       } else {
         final AgentSpan previousSpan = spanFromContext(span.swap());
         if (previousSpan != null) {
+          CONSUMER_DECORATE.beforeFinish(previousSpan);
           previousSpan.finishWithEndToEnd();
         }
       }
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
@@ -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,34 @@ 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 Context enter(
+        @Advice.Argument(value = 0) Envelope envelope,
+        @Advice.Local("taskScope") AgentScope taskScope) {
 
       // 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();
+        return null;
+      } else {
+        return 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.Enter 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
@@ -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
@@ -121,3 +121,11 @@ class PekkoActorTest extends InstrumentationSpecification {
     }
   }
 }
+
+class PekkoActorContextSwapTest extends PekkoActorTest {
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED,  "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
@@ -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();
   }
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`import datadog.trace.agent.test.InstrumentationSpecification`
	`2`	`+import datadog.trace.api.config.TraceInstrumentationConfig`
`2`	`3`	`import datadog.trace.bootstrap.instrumentation.api.Tags`
`3`	`4`	`import spock.lang.Shared`
`4`	`5`
`@@ -82,3 +83,11 @@ class AkkaActorTest extends InstrumentationSpecification {`
`82`	`83`	`}`
`83`	`84`	`}`
`84`	`85`	`}`
	`86`	`+`
	`87`	`+class AkkaActorContextSwapTest extends AkkaActorTest {`
	`88`	`+ @Override`
	`89`	`+ void configurePreAgent() {`
	`90`	`+ super.configurePreAgent()`
	`91`	`+ injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")`
	`92`	`+ }`
	`93`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -98,6 +98,7 @@ public static MessageConsumerState beforeReceive(@Advice.This final MessageConsu`
`98`	`98`	`} else {`
`99`	`99`	`final AgentSpan previousSpan = spanFromContext(getRootContext().swap());`
`100`	`100`	`if (previousSpan != null) {`
	`101`	`+ CONSUMER_DECORATE.beforeFinish(previousSpan);`
`101`	`102`	`previousSpan.finishWithEndToEnd();`
`102`	`103`	`}`
`103`	`104`	`}`
`@@ -178,6 +179,7 @@ public static void afterReceive(`
`178`	`179`	`} else {`
`179`	`180`	`final AgentSpan previousSpan = spanFromContext(span.swap());`
`180`	`181`	`if (previousSpan != null) {`
	`182`	`+ CONSUMER_DECORATE.beforeFinish(previousSpan);`
`181`	`183`	`previousSpan.finishWithEndToEnd();`
`182`	`184`	`}`
`183`	`185`	`}`
Original file line number	Diff line number	Diff line change
`@@ -121,3 +121,11 @@ class PekkoActorTest extends InstrumentationSpecification {`
`121`	`121`	`}`
`122`	`122`	`}`
`123`	`123`	`}`
	`124`	`+`
	`125`	`+class PekkoActorContextSwapTest extends PekkoActorTest {`
	`126`	`+ @Override`
	`127`	`+ void configurePreAgent() {`
	`128`	`+ super.configurePreAgent()`
	`129`	`+ injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")`
	`130`	`+ }`
	`131`	`+}`