Use context swap for Akka and Pekko HTTP flows

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/WithHttpServer.groovy

@@ -33,6 +33,10 @@ abstract class WithHttpServer<SERVER> extends VersionedNamingTestBase {
     return new DefaultHttpServer()
   }
 
+  boolean recreateServerForEachTest() {
+    false
+  }
+
   private class DefaultHttpServer implements HttpServer {
     final ServerSocket socket = PortUtils.randomOpenSocket()
     final int port = socket.localPort
@@ -64,6 +68,30 @@ abstract class WithHttpServer<SERVER> extends VersionedNamingTestBase {
   }
 
   void setupSpec() {
+    if (!recreateServerForEachTest()) {
+      startHttpServer()
+    }
+  }
+
+  void setup() {
+    if (recreateServerForEachTest()) {
+      startHttpServer()
+    }
+  }
+
+  void cleanup() {
+    if (recreateServerForEachTest()) {
+      stopHttpServer()
+    }
+  }
+
+  void cleanupSpec() {
+    if (!recreateServerForEachTest()) {
+      stopHttpServer()
+    }
+  }
+
+  private void startHttpServer() {
     server = server()
     server.start()
     address = server.address()
@@ -72,7 +100,7 @@ abstract class WithHttpServer<SERVER> extends VersionedNamingTestBase {
     println "$server started at: $address"
   }
 
-  void cleanupSpec() {
+  private void stopHttpServer() {
     server.stop()
     println "$server stopped at: $address"
   }

dd-java-agent/instrumentation/akka/akka-actor-2.5/src/akka23Test/groovy/AkkaActorTest.groovy

@@ -4,7 +4,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags
 import spock.lang.AutoCleanup
 import spock.lang.Shared
 
-class AkkaActorTest extends InstrumentationSpecification {
+abstract class AbstractAkkaActorTest extends InstrumentationSpecification {
   @Shared
   @AutoCleanup
   def akkaTester = new AkkaActors()
@@ -56,8 +56,10 @@ class AkkaActorTest extends InstrumentationSpecification {
     "forward" | "Akka"  | "Hello"          | 10
     "route"   | "Rakka" | "How you doin'"  | 10
   }
+}
 
-  def "actor message handling should close leaked scopes"() {
+class AkkaActorTest extends AbstractAkkaActorTest {
+  def "legacy actor message handling should close leaked scopes"() {
     when:
     akkaTester.leak("Leaker", "drip")
 
@@ -86,7 +88,7 @@ class AkkaActorTest extends InstrumentationSpecification {
   }
 }
 
-class AkkaActorContextSwapForkedTest extends AkkaActorTest {
+class AkkaActorContextSwapForkedTest extends AbstractAkkaActorTest {
   @Override
   void configurePreAgent() {
     super.configurePreAgent()

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy

@@ -1,6 +1,7 @@
 import datadog.trace.agent.test.base.HttpServer
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions
+import datadog.trace.api.config.TraceInstrumentationConfig
 import datadog.trace.test.util.ThreadUtils
 import datadog.trace.instrumentation.akkahttp.AkkaHttpServerDecorator
 import okhttp3.HttpUrl
@@ -319,6 +320,20 @@ class AkkaHttpServerInstrumentationBindAndHandleTest extends AkkaHttpServerInstr
   }
 }
 
+class AkkaHttpServerInstrumentationBindAndHandleContextSwapForkedTest extends AkkaHttpServerInstrumentationBindAndHandleTest {
+  @Override
+  boolean recreateServerForEachTest() {
+    // This forked suite changes a process-wide setting; do not let its actor system outlive it.
+    true
+  }
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
+  }
+}
+
 class AkkaHttpServerInstrumentationBindAndHandleAsyncWithRouteAsyncHandlerTest extends AkkaHttpServerInstrumentationTest {
   String akkaHttpVersion
 

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/java/datadog/trace/instrumentation/akkahttp/SwappedContextScopeTest.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.akkahttp;
+
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+import datadog.context.Context;
+import datadog.context.ContextKey;
+import datadog.context.ContextScope;
+import org.junit.jupiter.api.Test;
+
+class SwappedContextScopeTest {
+  private static final ContextKey<String> KEY = ContextKey.named("akka-http-swap-test");
+
+  @Test
+  void restoresPreviousContext() {
+    Context previous = Context.root().with(KEY, "previous");
+    Context request = Context.root().with(KEY, "request");
+
+    try (ContextScope previousScope = previous.attach()) {
+      try (ContextScope requestScope = new DatadogWrapperHelper.SwappedContextScope(request)) {
+        assertSame(request, Context.current());
+      }
+      assertSame(previous, Context.current());
+    }
+  }
+
+  @Test
+  void waitsUntilRequestContextIsCurrent() {
+    Context previous = Context.root().with(KEY, "previous");
+    Context request = Context.root().with(KEY, "request");
+    Context other = Context.root().with(KEY, "other");
+
+    try (ContextScope previousScope = previous.attach()) {
+      ContextScope requestScope = new DatadogWrapperHelper.SwappedContextScope(request);
+      try (ContextScope otherScope = other.attach()) {
+        requestScope.close();
+        assertSame(other, Context.current());
+      }
+      assertSame(request, Context.current());
+      requestScope.close();
+      assertSame(previous, Context.current());
+    }
+  }
+}

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/scala/AkkaHttpTestWebServer.scala

@@ -58,9 +58,12 @@ class AkkaHttpTestWebServer(binder: Binder) extends HttpServer {
 
   override def stop(): Unit = {
     import materializer.executionContext
-    portBinding
-      .flatMap(_.unbind())
-      .onComplete(_ => system.terminate())
+    Await.ready(
+      portBinding
+        .flatMap(_.unbind())
+        .flatMap(_ => system.terminate()),
+      10 seconds
+    )
   }
 
   override def address(): URI = {

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttp2ServerInstrumentation.java

@@ -36,6 +36,7 @@ public String[] knownMatchingTypes() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".DatadogWrapperHelper",
+      packageName + ".DatadogWrapperHelper$SwappedContextScope",
       packageName + ".DatadogAsyncHandlerWrapper",
       packageName + ".DatadogAsyncHandlerWrapper$1",
       packageName + ".DatadogAsyncHandlerWrapper$2",

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpServerInstrumentation.java

@@ -65,6 +65,7 @@ public String instrumentedType() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".DatadogWrapperHelper",
+      packageName + ".DatadogWrapperHelper$SwappedContextScope",
       packageName + ".DatadogServerRequestResponseFlowWrapper",
       packageName + ".DatadogServerRequestResponseFlowWrapper$1",
       packageName + ".DatadogServerRequestResponseFlowWrapper$1$1",

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java

@@ -67,7 +67,7 @@ public GraphStageLogic createLogic(final Attributes inheritedAttributes) throws
               @Override
               public void onPush() throws Exception {
                 final HttpRequest request = grab(requestInlet);
-                final ContextScope scope = DatadogWrapperHelper.createSpan(request);
+                final ContextScope scope = DatadogWrapperHelper.createSpanForFlow(request);
                 final AgentSpan span = fromContext(scope.context());
                 RequestContext requestContext = span.getRequestContext();
                 if (requestContext != null) {
@@ -87,11 +87,8 @@ public void onPush() throws Exception {
 
                 scopes.add(scope);
                 push(requestOutlet, request);
-                // Since we haven't instrumented the akka stream state machine, we can't rely
-                // on spans and scopes being propagated during the push and pull of the
-                // element. Instead we let the scope leak intentionally here and clean it
-                // up when the user response comes back, or in the actor message processing
-                // instrumentation that drives this state machine.
+                // Legacy mode leaves the scope open so the surrounding actor can clean it up.
+                // Context-manager mode swaps the context and the actor restores it on exit.
               }
 
               @Override
@@ -143,9 +140,7 @@ public void onPush() throws Exception {
                     response = newResponse;
                   }
                   DatadogWrapperHelper.finishSpan(scope.context(), response);
-                  // Check if the active span matches the scope from when the request came in,
-                  // and close it. If it's not, then it will be cleaned up actor message
-                  // processing instrumentation that drives this state machine
+                  // Legacy mode may still own the scope when the response arrives.
                   AgentSpan activeSpan = activeSpan();
                   if (activeSpan == span) {
                     scope.close();

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java

@@ -7,17 +7,59 @@
 import akka.http.scaladsl.model.HttpResponse;
 import datadog.context.Context;
 import datadog.context.ContextScope;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 
 public class DatadogWrapperHelper {
+  private static final boolean LEGACY_CONTEXT_MANAGER_ENABLED =
+      InstrumenterConfig.get().isLegacyContextManagerEnabled();
+
   public static ContextScope createSpan(final HttpRequest request) {
+    return startSpan(request).attach();
+  }
+
+  public static ContextScope createSpanForFlow(final HttpRequest request) {
+    final Context context = startSpan(request);
+    if (LEGACY_CONTEXT_MANAGER_ENABLED) {
+      return context.attach();
+    }
+    return new SwappedContextScope(context);
+  }
+
+  private static Context startSpan(final HttpRequest request) {
     final Context parentContext = DECORATE.extract(request);
     final Context context = DECORATE.startSpan(request, parentContext);
     final AgentSpan span = fromContext(context);
     DECORATE.afterStart(span);
     DECORATE.onRequest(span, request, request, parentContext);
 
-    return context.attach();
+    return context;
+  }
+
+  static final class SwappedContextScope implements ContextScope {
+    private final Context context;
+    private final Context previousContext;
+    private final Thread ownerThread;
+    private boolean closed;
+
+    SwappedContextScope(final Context context) {
+      this.context = context;
+      this.ownerThread = Thread.currentThread();
+      this.previousContext = context.swap();
+    }
+
+    @Override
+    public Context context() {
+      return context;
+    }
+
+    @Override
+    public void close() {
+      if (!closed && ownerThread == Thread.currentThread() && context == Context.current()) {
+        closed = true;
+        previousContext.swap();
+      }
+    }
   }
 
   public static void finishSpan(final Context context, final HttpResponse response) {

dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/test/groovy/PekkoActorTest.groovy

@@ -4,7 +4,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags
 import spock.lang.AutoCleanup
 import spock.lang.Shared
 
-class PekkoActorTest extends InstrumentationSpecification {
+abstract class AbstractPekkoActorTest extends InstrumentationSpecification {
 
   @Shared
   @AutoCleanup
@@ -94,8 +94,10 @@ class PekkoActorTest extends InstrumentationSpecification {
       }
     }
   }
+}
 
-  def "actor message handling should close leaked scopes"() {
+class PekkoActorTest extends AbstractPekkoActorTest {
+  def "legacy actor message handling should close leaked scopes"() {
     when:
     pekkoTester.leak("Leaker", "drip")
 
@@ -124,7 +126,7 @@ class PekkoActorTest extends InstrumentationSpecification {
   }
 }
 
-class PekkoActorContextSwapForkedTest extends PekkoActorTest {
+class PekkoActorContextSwapForkedTest extends AbstractPekkoActorTest {
   @Override
   void configurePreAgent() {
     super.configurePreAgent()

dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy

@@ -1,6 +1,7 @@
 import datadog.trace.agent.test.base.HttpServer
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.agent.test.naming.TestingGenericHttpNamingConventions
+import datadog.trace.api.config.TraceInstrumentationConfig
 import datadog.trace.test.util.ThreadUtils
 import datadog.trace.instrumentation.pekkohttp.PekkoHttpServerDecorator
 import okhttp3.Request
@@ -126,6 +127,20 @@ class PekkoHttpServerInstrumentationBindAndHandleTest extends PekkoHttpServerIns
   }
 }
 
+class PekkoHttpServerInstrumentationBindAndHandleContextSwapForkedTest extends PekkoHttpServerInstrumentationBindAndHandleTest {
+  @Override
+  boolean recreateServerForEachTest() {
+    // This forked suite changes a process-wide setting; do not let its actor system outlive it.
+    true
+  }
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
+  }
+}
+
 class PekkoHttpServerInstrumentationBindAndHandleAsyncWithRouteAsyncHandlerTest extends PekkoHttpServerInstrumentationTest {
   @Override
   HttpServer server() {

dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/java/datadog/trace/instrumentation/pekkohttp/SwappedContextScopeTest.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.pekkohttp;
+
+import static org.junit.jupiter.api.Assertions.assertSame;
+
+import datadog.context.Context;
+import datadog.context.ContextKey;
+import datadog.context.ContextScope;
+import org.junit.jupiter.api.Test;
+
+class SwappedContextScopeTest {
+  private static final ContextKey<String> KEY = ContextKey.named("pekko-http-swap-test");
+
+  @Test
+  void restoresPreviousContext() {
+    Context previous = Context.root().with(KEY, "previous");
+    Context request = Context.root().with(KEY, "request");
+
+    try (ContextScope previousScope = previous.attach()) {
+      try (ContextScope requestScope = new DatadogWrapperHelper.SwappedContextScope(request)) {
+        assertSame(request, Context.current());
+      }
+      assertSame(previous, Context.current());
+    }
+  }
+
+  @Test
+  void waitsUntilRequestContextIsCurrent() {
+    Context previous = Context.root().with(KEY, "previous");
+    Context request = Context.root().with(KEY, "request");
+    Context other = Context.root().with(KEY, "other");
+
+    try (ContextScope previousScope = previous.attach()) {
+      ContextScope requestScope = new DatadogWrapperHelper.SwappedContextScope(request);
+      try (ContextScope otherScope = other.attach()) {
+        requestScope.close();
+        assertSame(other, Context.current());
+      }
+      assertSame(request, Context.current());
+      requestScope.close();
+      assertSame(previous, Context.current());
+    }
+  }
+}

dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala

@@ -43,9 +43,12 @@ class PekkoHttpTestWebServer(binder: Binder) extends HttpServer {
 
   override def stop(): Unit = {
     import materializer.executionContext
-    portBinding
-      .flatMap(_.unbind())
-      .onComplete(_ => system.terminate())
+    Await.ready(
+      portBinding
+        .flatMap(_.unbind())
+        .flatMap(_ => system.terminate()),
+      10 seconds
+    )
   }
 
   override def address(): URI = {