Skip to content

Commit f592844

Browse files
amarzialiamarziali
committed
Use context swap on akka and pekko
1 parent ae97a09 commit f592844

6 files changed

Lines changed: 99 additions & 30 deletions

File tree

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import datadog.trace.agent.test.InstrumentationSpecification
2+
import datadog.trace.api.config.TraceInstrumentationConfig
23
import datadog.trace.bootstrap.instrumentation.api.Tags
34
import spock.lang.Shared
45

@@ -82,3 +83,11 @@ class AkkaActorTest extends InstrumentationSpecification {
8283
}
8384
}
8485
}
86+
87+
class AkkaActorContextSwapTest extends AkkaActorTest {
88+
@Override
89+
void configurePreAgent() {
90+
super.configurePreAgent()
91+
injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
92+
}
93+
}

‎dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,16 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
67
import static java.util.Collections.singletonMap;
78
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
89

910
import akka.dispatch.Envelope;
1011
import com.google.auto.service.AutoService;
12+
import datadog.context.Context;
1113
import datadog.trace.agent.tooling.Instrumenter;
1214
import datadog.trace.agent.tooling.InstrumenterModule;
15+
import datadog.trace.api.InstrumenterConfig;
1316
import datadog.trace.bootstrap.InstrumentationContext;
1417
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1518
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
@@ -54,24 +57,35 @@ public void methodAdvice(MethodTransformer transformer) {
5457
*/
5558
public static class InvokeAdvice {
5659
@Advice.OnMethodEnter(suppress = Throwable.class)
57-
public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) {
60+
public static void enter(
61+
@Advice.Argument(value = 0) Envelope envelope,
62+
@Advice.Local("taskScope") AgentScope taskScope,
63+
@Advice.Local("checkpointContext") Context checkpointContext) {
5864

5965
// do this before checkpointing, as the envelope's task scope may already be active
60-
AgentScope taskScope =
66+
taskScope =
6167
AdviceUtils.startTaskScope(
6268
InstrumentationContext.get(Envelope.class, State.class), envelope);
6369

64-
// remember the currently active scope so we can roll back to this point
65-
checkpointActiveForRollback();
66-
67-
return taskScope;
70+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
71+
// remember the currently active scope so we can roll back to this point
72+
checkpointActiveForRollback();
73+
} else {
74+
checkpointContext = getCurrentContext().swap();
75+
}
6876
}
6977

7078
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
71-
public static void exit(@Advice.Enter AgentScope taskScope) {
79+
public static void exit(
80+
@Advice.Local("taskScope") AgentScope taskScope,
81+
@Advice.Local("checkpointContext") Context checkpointContext) {
7282

73-
// Clean up any leaking scopes from akka-streams/akka-http etc.
74-
rollbackActiveToCheckpoint();
83+
if (checkpointContext == null) {
84+
// Clean up any leaking scopes from akka-streams/akka-http etc.
85+
rollbackActiveToCheckpoint();
86+
} else {
87+
checkpointContext.swap();
88+
}
7589

7690
// close envelope's task scope if we previously started it
7791
if (taskScope != null) {

‎dd-java-agent/instrumentation/akka/akka-actor-2.5/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,16 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
67
import static java.util.Collections.singletonList;
78
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
89

910
import com.google.auto.service.AutoService;
11+
import datadog.context.Context;
1012
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1113
import datadog.trace.agent.tooling.Instrumenter;
1214
import datadog.trace.agent.tooling.InstrumenterModule;
15+
import datadog.trace.api.InstrumenterConfig;
1316
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
1417
import java.util.Collection;
1518
import java.util.EnumMap;
@@ -59,15 +62,24 @@ public void methodAdvice(MethodTransformer transformer) {
5962
*/
6063
public static final class SuppressMailboxRunAdvice {
6164
@Advice.OnMethodEnter(suppress = Throwable.class)
62-
public static void enter() {
63-
// remember the currently active scope so we can roll back to this point
64-
checkpointActiveForRollback();
65+
public static Context enter() {
66+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
67+
// remember the currently active scope so we can roll back to this point
68+
checkpointActiveForRollback();
69+
return null;
70+
} else {
71+
return getCurrentContext().swap();
72+
}
6573
}
6674

6775
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
68-
public static void exit() {
69-
// Clean up any leaking scopes from akka-streams/akka-http etc.
70-
rollbackActiveToCheckpoint();
76+
public static void exit(@Advice.Enter final Context checkpointContext) {
77+
if (checkpointContext == null) {
78+
// Clean up any leaking scopes from akka-streams/akka-http etc.
79+
rollbackActiveToCheckpoint();
80+
} else {
81+
checkpointContext.swap();
82+
}
7183
}
7284
}
7385
}

‎dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,15 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
67
import static java.util.Collections.singletonMap;
78
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
89

910
import com.google.auto.service.AutoService;
11+
import datadog.context.Context;
1012
import datadog.trace.agent.tooling.Instrumenter;
1113
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import datadog.trace.api.InstrumenterConfig;
1215
import datadog.trace.bootstrap.InstrumentationContext;
1316
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1417
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
@@ -54,24 +57,35 @@ public void methodAdvice(MethodTransformer transformer) {
5457
*/
5558
public static class InvokeAdvice {
5659
@Advice.OnMethodEnter(suppress = Throwable.class)
57-
public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) {
60+
public static void enter(
61+
@Advice.Argument(value = 0) Envelope envelope,
62+
@Advice.Local("taskScope") AgentScope taskScope,
63+
@Advice.Local("checkpointContext") Context checkpointContext) {
5864

5965
// do this before checkpointing, as the envelope's task scope may already be active
60-
AgentScope taskScope =
66+
taskScope =
6167
AdviceUtils.startTaskScope(
6268
InstrumentationContext.get(Envelope.class, State.class), envelope);
6369

64-
// remember the currently active scope so we can roll back to this point
65-
checkpointActiveForRollback();
66-
67-
return taskScope;
70+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
71+
// remember the currently active scope so we can roll back to this point
72+
checkpointActiveForRollback();
73+
} else {
74+
checkpointContext = getCurrentContext().swap();
75+
}
6876
}
6977

7078
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
71-
public static void exit(@Advice.Enter AgentScope taskScope) {
79+
public static void exit(
80+
@Advice.Local("taskScope") AgentScope taskScope,
81+
@Advice.Local("checkpointContext") Context checkpointContext) {
7282

73-
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
74-
rollbackActiveToCheckpoint();
83+
if (checkpointContext == null) {
84+
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
85+
rollbackActiveToCheckpoint();
86+
} else {
87+
checkpointContext.swap();
88+
}
7589

7690
// close envelope's task scope if we previously started it
7791
if (taskScope != null) {

‎dd-java-agent/instrumentation/pekko/pekko-concurrent-1.0/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,16 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
6+
import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.getCurrentContext;
67
import static java.util.Collections.singletonList;
78
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
89

910
import com.google.auto.service.AutoService;
11+
import datadog.context.Context;
1012
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1113
import datadog.trace.agent.tooling.Instrumenter;
1214
import datadog.trace.agent.tooling.InstrumenterModule;
15+
import datadog.trace.api.InstrumenterConfig;
1316
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
1417
import java.util.Collection;
1518
import java.util.EnumMap;
@@ -59,15 +62,24 @@ public void methodAdvice(MethodTransformer transformer) {
5962
*/
6063
public static final class SuppressMailboxRunAdvice {
6164
@Advice.OnMethodEnter(suppress = Throwable.class)
62-
public static void enter() {
63-
// remember the currently active scope so we can roll back to this point
64-
checkpointActiveForRollback();
65+
public static Context enter() {
66+
if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) {
67+
// remember the currently active scope so we can roll back to this point
68+
checkpointActiveForRollback();
69+
return null;
70+
} else {
71+
return getCurrentContext().swap();
72+
}
6573
}
6674

6775
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
68-
public static void exit() {
69-
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
70-
rollbackActiveToCheckpoint();
76+
public static void exit(@Advice.Enter final Context checkpointContext) {
77+
if (checkpointContext == null) {
78+
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
79+
rollbackActiveToCheckpoint();
80+
} else {
81+
checkpointContext.swap();
82+
}
7183
}
7284
}
7385
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,3 +121,11 @@ class PekkoActorTest extends InstrumentationSpecification {
121121
}
122122
}
123123
}
124+
125+
class PekkoActorContextSwapTest extends PekkoActorTest {
126+
@Override
127+
void configurePreAgent() {
128+
super.configurePreAgent()
129+
injectSysConfig(TraceInstrumentationConfig.LEGACY_CONTEXT_MANAGER_ENABLED, "false")
130+
}
131+
}

0 commit comments

Comments
 (0)