Skip to content

Commit 60f8046

Browse files
mccullsmcculls
authored
Simplify the Akka/Pekko checkpointing advice (#8625)
1 parent 75e5963 commit 60f8046

File tree

4 files changed

+32
-91
lines changed

4 files changed

+32
-91
lines changed

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

+15-30
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
package datadog.trace.instrumentation.akka.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
64
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
7-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
85
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
96
import static java.util.Collections.singletonMap;
107
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -15,7 +12,6 @@
1512
import datadog.trace.agent.tooling.InstrumenterModule;
1613
import datadog.trace.bootstrap.InstrumentationContext;
1714
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
18-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1915
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
2016
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
2117
import java.util.Map;
@@ -58,40 +54,29 @@ public void methodAdvice(MethodTransformer transformer) {
5854
*/
5955
public static class InvokeAdvice {
6056
@Advice.OnMethodEnter(suppress = Throwable.class)
61-
public static void enter(
62-
@Advice.Argument(value = 0) Envelope envelope,
63-
@Advice.Local(value = "taskScope") AgentScope taskScope) {
64-
checkpointActiveForRollback();
65-
// note: task scope may be the same as the scope we want to roll back to,
66-
// so we must remember to close it on exit to balance the activation count
67-
taskScope =
57+
public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) {
58+
59+
// do this before checkpointing, as the envelope's task scope may already be active
60+
AgentScope taskScope =
6861
AdviceUtils.startTaskScope(
6962
InstrumentationContext.get(Envelope.class, State.class), envelope);
70-
// There was a scope created from the envelope, so use that
71-
if (taskScope != null) {
72-
return;
73-
}
74-
AgentSpan activeSpan = activeSpan();
75-
// If there is no active scope, we can clean all the way to the bottom
76-
if (activeSpan == null) {
77-
return;
78-
}
79-
// If there is a noop span in the active scope, we can clean all the way to this scope
80-
if (activeSpan == noopSpan()) {
81-
return;
82-
}
83-
// Create an active scope with a noop span, and clean all the way to the previous scope
84-
activateSpan(noopSpan());
63+
64+
// remember the currently active scope so we can roll back to this point
65+
checkpointActiveForRollback();
66+
67+
return taskScope;
8568
}
8669

8770
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
88-
public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) {
71+
public static void exit(@Advice.Enter AgentScope taskScope) {
72+
73+
// Clean up any leaking scopes from akka-streams/akka-http etc.
74+
rollbackActiveToCheckpoint();
75+
76+
// close envelope's task scope if we previously started it
8977
if (taskScope != null) {
90-
// then we have invoked an Envelope and need to mark the work complete
9178
taskScope.close();
9279
}
93-
// Clean up any leaking scopes from akka-streams/akka-http etc.
94-
rollbackActiveToCheckpoint();
9580
}
9681
}
9782
}

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

+1-15
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
package datadog.trace.instrumentation.akka.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
64
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
7-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
85
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
96
import static java.util.Collections.singletonList;
107
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -13,7 +10,6 @@
1310
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1411
import datadog.trace.agent.tooling.Instrumenter;
1512
import datadog.trace.agent.tooling.InstrumenterModule;
16-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1713
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
1814
import java.util.Collection;
1915
import java.util.EnumMap;
@@ -64,18 +60,8 @@ public void methodAdvice(MethodTransformer transformer) {
6460
public static final class SuppressMailboxRunAdvice {
6561
@Advice.OnMethodEnter(suppress = Throwable.class)
6662
public static void enter() {
63+
// remember the currently active scope so we can roll back to this point
6764
checkpointActiveForRollback();
68-
AgentSpan activeSpan = activeSpan();
69-
// If there is no active scope, we can clean all the way to the bottom
70-
if (activeSpan == null) {
71-
return;
72-
}
73-
// If there is a noop span in the active scope, we can clean all the way to this scope
74-
if (activeSpan == noopSpan()) {
75-
return;
76-
}
77-
// Create an active scope with a noop span, and clean all the way to the previous scope
78-
activateSpan(noopSpan());
7965
}
8066

8167
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

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

+15-30
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
package datadog.trace.instrumentation.pekko.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
64
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
7-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
85
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
96
import static java.util.Collections.singletonMap;
107
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -14,7 +11,6 @@
1411
import datadog.trace.agent.tooling.InstrumenterModule;
1512
import datadog.trace.bootstrap.InstrumentationContext;
1613
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1814
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
1915
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
2016
import java.util.Map;
@@ -58,40 +54,29 @@ public void methodAdvice(MethodTransformer transformer) {
5854
*/
5955
public static class InvokeAdvice {
6056
@Advice.OnMethodEnter(suppress = Throwable.class)
61-
public static void enter(
62-
@Advice.Argument(value = 0) Envelope envelope,
63-
@Advice.Local(value = "taskScope") AgentScope taskScope) {
64-
checkpointActiveForRollback();
65-
// note: task scope may be the same as the scope we want to roll back to,
66-
// so we must remember to close it on exit to balance the activation count
67-
taskScope =
57+
public static AgentScope enter(@Advice.Argument(value = 0) Envelope envelope) {
58+
59+
// do this before checkpointing, as the envelope's task scope may already be active
60+
AgentScope taskScope =
6861
AdviceUtils.startTaskScope(
6962
InstrumentationContext.get(Envelope.class, State.class), envelope);
70-
// There was a scope created from the envelope, so use that
71-
if (taskScope != null) {
72-
return;
73-
}
74-
AgentSpan activeSpan = activeSpan();
75-
// If there is no active scope, we can clean all the way to the bottom
76-
if (activeSpan == null) {
77-
return;
78-
}
79-
// If there is a noop span in the active scope, we can clean all the way to this scope
80-
if (activeSpan == noopSpan()) {
81-
return;
82-
}
83-
// Create an active scope with a noop span, and clean all the way to the previous scope
84-
activateSpan(noopSpan());
63+
64+
// remember the currently active scope so we can roll back to this point
65+
checkpointActiveForRollback();
66+
67+
return taskScope;
8568
}
8669

8770
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
88-
public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) {
71+
public static void exit(@Advice.Enter AgentScope taskScope) {
72+
73+
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
74+
rollbackActiveToCheckpoint();
75+
76+
// close envelope's task scope if we previously started it
8977
if (taskScope != null) {
90-
// then we have invoked an Envelope and need to mark the work complete
9178
taskScope.close();
9279
}
93-
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
94-
rollbackActiveToCheckpoint();
9580
}
9681
}
9782
}

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

+1-16
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
package datadog.trace.instrumentation.pekko.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
64
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
7-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
85
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
96
import static java.util.Collections.singletonList;
107
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -13,7 +10,6 @@
1310
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1411
import datadog.trace.agent.tooling.Instrumenter;
1512
import datadog.trace.agent.tooling.InstrumenterModule;
16-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1713
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
1814
import java.util.Collection;
1915
import java.util.EnumMap;
@@ -64,19 +60,8 @@ public void methodAdvice(MethodTransformer transformer) {
6460
public static final class SuppressMailboxRunAdvice {
6561
@Advice.OnMethodEnter(suppress = Throwable.class)
6662
public static void enter() {
63+
// remember the currently active scope so we can roll back to this point
6764
checkpointActiveForRollback();
68-
69-
AgentSpan activeSpan = activeSpan();
70-
// If there is no active scope, we can clean all the way to the bottom
71-
if (activeSpan == null) {
72-
return;
73-
}
74-
// If there is a noop span in the active scope, we can clean all the way to this scope
75-
if (activeSpan == noopSpan()) {
76-
return;
77-
}
78-
// Create an active scope with a noop span, and clean all the way to the previous scope
79-
activateSpan(noopSpan());
8065
}
8166

8267
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

0 commit comments

Comments
 (0)