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..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 @@ -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 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/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 187af35e63d..c800897b9f0 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 @@ -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 index 4f7dda119ed..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 @@ -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 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") + } +} 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(); }