Optionally use context swap on akka and pekko

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")
+  }
+}

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) {

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();
+      }
     }
   }
 }

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();
         }
       }

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) {

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();
+      }
     }
   }
 }

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")
+  }
+}

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();
   }