diff --git a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/WithHttpServer.groovy b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/WithHttpServer.groovy index 235aa2a7933..3225b1a49b7 100644 --- a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/WithHttpServer.groovy +++ b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/WithHttpServer.groovy @@ -33,6 +33,10 @@ abstract class WithHttpServer 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 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 extends VersionedNamingTestBase { println "$server started at: $address" } - void cleanupSpec() { + private void stopHttpServer() { server.stop() println "$server stopped at: $address" } 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 e4a5d76d27a..31cb7711a6a 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 @@ -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() diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy index a1aa9510a93..2d807d5e215 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy +++ b/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 diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/java/datadog/trace/instrumentation/akkahttp/SwappedContextScopeTest.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/java/datadog/trace/instrumentation/akkahttp/SwappedContextScopeTest.java new file mode 100644 index 00000000000..a439fdd92ce --- /dev/null +++ b/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 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()); + } + } +} diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/scala/AkkaHttpTestWebServer.scala b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/scala/AkkaHttpTestWebServer.scala index 4c6d5b3f142..8fce57838e2 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/scala/AkkaHttpTestWebServer.scala +++ b/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 = { diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttp2ServerInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttp2ServerInstrumentation.java index 169cb3dbbc0..846a9c6b836 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttp2ServerInstrumentation.java +++ b/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", diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpServerInstrumentation.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpServerInstrumentation.java index 0f2649fbc3a..02d3d90a05c 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpServerInstrumentation.java +++ b/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", diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java index e35593c27e0..545aeff7600 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogServerRequestResponseFlowWrapper.java +++ b/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(); diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java index 2acf49c3ad6..cf0a5d68258 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java +++ b/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) { 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 9b70f175ce5..b3d4d6b1e95 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 @@ -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() diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy index 9f7fa23c2b1..1034cf83e26 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy +++ b/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() { diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/java/datadog/trace/instrumentation/pekkohttp/SwappedContextScopeTest.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/java/datadog/trace/instrumentation/pekkohttp/SwappedContextScopeTest.java new file mode 100644 index 00000000000..21fee39d54f --- /dev/null +++ b/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 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()); + } + } +} diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala index 300e84478ad..7e070ac16a5 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala +++ b/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 = { diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java index 52075e732f9..0de81b07a08 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogServerRequestResponseFlowWrapper.java @@ -64,14 +64,11 @@ 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); scopes.add(scope); push(requestOutlet, request); - // Since we haven't instrumented the pekko 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 @@ -115,9 +112,7 @@ public void onPush() throws Exception { final ContextScope scope = scopes.poll(); if (scope != null) { 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(); AgentSpan span = fromContext(scope.context()); if (activeSpan == span) { diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java index d4565c68bfc..3cee9b327b5 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java @@ -5,19 +5,61 @@ import datadog.context.Context; import datadog.context.ContextScope; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import org.apache.pekko.http.scaladsl.model.HttpRequest; import org.apache.pekko.http.scaladsl.model.HttpResponse; 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) { diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttp2ServerInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttp2ServerInstrumentation.java index 9899e04f6af..55f78ffab6d 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttp2ServerInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttp2ServerInstrumentation.java @@ -42,6 +42,7 @@ public String[] knownMatchingTypes() { public String[] helperClassNames() { return new String[] { packageName + ".DatadogWrapperHelper", + packageName + ".DatadogWrapperHelper$SwappedContextScope", packageName + ".DatadogAsyncHandlerWrapper", packageName + ".DatadogAsyncHandlerWrapper$1", packageName + ".DatadogAsyncHandlerWrapper$2", diff --git a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpServerInstrumentation.java b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpServerInstrumentation.java index bb74992b00c..e37c329c7a2 100644 --- a/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpServerInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpServerInstrumentation.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",