Skip to content

Commit 3efa8ff

Browse files
mccullsmcculls
committed
Use checkpoint..rollback in Akka/Pekko instrumentations to cleanup leaky scopes
1 parent dbfdb19 commit 3efa8ff

File tree

4 files changed

+73
-80
lines changed

4 files changed

+73
-80
lines changed

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

Lines changed: 23 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
8+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
89
import static java.util.Collections.singletonMap;
910
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1011

@@ -14,6 +15,7 @@
1415
import datadog.trace.agent.tooling.InstrumenterModule;
1516
import datadog.trace.bootstrap.InstrumentationContext;
1617
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
18+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1719
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
1820
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
1921
import java.util.Map;
@@ -56,43 +58,40 @@ public void methodAdvice(MethodTransformer transformer) {
5658
*/
5759
public static class InvokeAdvice {
5860
@Advice.OnMethodEnter(suppress = Throwable.class)
59-
public static AgentScope enter(
61+
public static void enter(
6062
@Advice.Argument(value = 0) Envelope envelope,
61-
@Advice.Local(value = "localScope") AgentScope localScope) {
62-
AgentScope activeScope = activeScope();
63-
localScope =
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 =
6468
AdviceUtils.startTaskScope(
6569
InstrumentationContext.get(Envelope.class, State.class), envelope);
66-
// There was a scope created from the envelop, so use that
67-
if (localScope != null) {
68-
return activeScope;
70+
// There was a scope created from the envelope, so use that
71+
if (taskScope != null) {
72+
return;
6973
}
74+
AgentSpan activeSpan = activeSpan();
7075
// If there is no active scope, we can clean all the way to the bottom
71-
if (null == activeScope) {
72-
return null;
76+
if (activeSpan == null) {
77+
return;
7378
}
7479
// 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 activeScope;
80+
if (activeSpan == noopSpan()) {
81+
return;
7782
}
7883
// Create an active scope with a noop span, and clean all the way to the previous scope
79-
localScope = activateSpan(noopSpan(), false);
80-
return activeScope;
84+
activateSpan(noopSpan(), false);
8185
}
8286

8387
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
84-
public static void exit(
85-
@Advice.Enter AgentScope scope, @Advice.Local(value = "localScope") AgentScope localScope) {
86-
if (localScope != null) {
88+
public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) {
89+
if (taskScope != null) {
8790
// then we have invoked an Envelope and need to mark the work complete
88-
localScope.close();
89-
}
90-
// Clean up any leaking scopes from akka-streams/akka-http et.c.
91-
AgentScope activeScope = activeScope();
92-
while (activeScope != null && activeScope != scope) {
93-
activeScope.close();
94-
activeScope = activeScope();
91+
taskScope.close();
9592
}
93+
// Clean up any leaking scopes from akka-streams/akka-http etc.
94+
rollbackActiveToCheckpoint();
9695
}
9796
}
9897
}

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

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,18 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
8+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
89
import static java.util.Collections.singletonList;
910
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1011

1112
import com.google.auto.service.AutoService;
1213
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1314
import datadog.trace.agent.tooling.Instrumenter;
1415
import datadog.trace.agent.tooling.InstrumenterModule;
15-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
16+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1617
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
1718
import java.util.Collection;
1819
import java.util.EnumMap;
@@ -62,29 +63,25 @@ public void methodAdvice(MethodTransformer transformer) {
6263
*/
6364
public static final class SuppressMailboxRunAdvice {
6465
@Advice.OnMethodEnter(suppress = Throwable.class)
65-
public static AgentScope enter() {
66-
AgentScope activeScope = activeScope();
66+
public static void enter() {
67+
checkpointActiveForRollback();
68+
AgentSpan activeSpan = activeSpan();
6769
// If there is no active scope, we can clean all the way to the bottom
68-
if (null == activeScope) {
69-
return null;
70+
if (activeSpan == null) {
71+
return;
7072
}
7173
// If there is a noop span in the active scope, we can clean all the way to this scope
72-
if (activeSpan() == noopSpan()) {
73-
return activeScope;
74+
if (activeSpan == noopSpan()) {
75+
return;
7476
}
7577
// Create an active scope with a noop span, and clean all the way to the previous scope
7678
activateSpan(noopSpan(), false);
77-
return activeScope;
7879
}
7980

8081
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
81-
public static void exit(@Advice.Enter final AgentScope scope) {
82-
// Clean up any leaking scopes from akka-streams/akka-http et.c.
83-
AgentScope activeScope = activeScope();
84-
while (activeScope != null && activeScope != scope) {
85-
activeScope.close();
86-
activeScope = activeScope();
87-
}
82+
public static void exit() {
83+
// Clean up any leaking scopes from akka-streams/akka-http etc.
84+
rollbackActiveToCheckpoint();
8885
}
8986
}
9087
}

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

Lines changed: 23 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
8+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
89
import static java.util.Collections.singletonMap;
910
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1011

@@ -13,6 +14,7 @@
1314
import datadog.trace.agent.tooling.InstrumenterModule;
1415
import datadog.trace.bootstrap.InstrumentationContext;
1516
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1618
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
1719
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
1820
import java.util.Map;
@@ -56,43 +58,40 @@ public void methodAdvice(MethodTransformer transformer) {
5658
*/
5759
public static class InvokeAdvice {
5860
@Advice.OnMethodEnter(suppress = Throwable.class)
59-
public static AgentScope enter(
61+
public static void enter(
6062
@Advice.Argument(value = 0) Envelope envelope,
61-
@Advice.Local(value = "localScope") AgentScope localScope) {
62-
AgentScope activeScope = activeScope();
63-
localScope =
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 =
6468
AdviceUtils.startTaskScope(
6569
InstrumentationContext.get(Envelope.class, State.class), envelope);
66-
// There was a scope created from the envelop, so use that
67-
if (localScope != null) {
68-
return activeScope;
70+
// There was a scope created from the envelope, so use that
71+
if (taskScope != null) {
72+
return;
6973
}
74+
AgentSpan activeSpan = activeSpan();
7075
// If there is no active scope, we can clean all the way to the bottom
71-
if (null == activeScope) {
72-
return null;
76+
if (activeSpan == null) {
77+
return;
7378
}
7479
// 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 activeScope;
80+
if (activeSpan == noopSpan()) {
81+
return;
7782
}
7883
// Create an active scope with a noop span, and clean all the way to the previous scope
79-
localScope = activateSpan(noopSpan(), false);
80-
return activeScope;
84+
activateSpan(noopSpan(), false);
8185
}
8286

8387
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
84-
public static void exit(
85-
@Advice.Enter AgentScope scope, @Advice.Local(value = "localScope") AgentScope localScope) {
86-
if (localScope != null) {
88+
public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) {
89+
if (taskScope != null) {
8790
// then we have invoked an Envelope and need to mark the work complete
88-
localScope.close();
89-
}
90-
// Clean up any leaking scopes from pekko-streams/pekko-http et.c.
91-
AgentScope activeScope = activeScope();
92-
while (activeScope != null && activeScope != scope) {
93-
activeScope.close();
94-
activeScope = activeScope();
91+
taskScope.close();
9592
}
93+
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
94+
rollbackActiveToCheckpoint();
9695
}
9796
}
9897
}

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

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,18 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
65
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
8+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
89
import static java.util.Collections.singletonList;
910
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1011

1112
import com.google.auto.service.AutoService;
1213
import datadog.trace.agent.tooling.ExcludeFilterProvider;
1314
import datadog.trace.agent.tooling.Instrumenter;
1415
import datadog.trace.agent.tooling.InstrumenterModule;
15-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
16+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1617
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
1718
import java.util.Collection;
1819
import java.util.EnumMap;
@@ -62,29 +63,26 @@ public void methodAdvice(MethodTransformer transformer) {
6263
*/
6364
public static final class SuppressMailboxRunAdvice {
6465
@Advice.OnMethodEnter(suppress = Throwable.class)
65-
public static AgentScope enter() {
66-
AgentScope activeScope = activeScope();
66+
public static void enter() {
67+
checkpointActiveForRollback();
68+
69+
AgentSpan activeSpan = activeSpan();
6770
// If there is no active scope, we can clean all the way to the bottom
68-
if (null == activeScope) {
69-
return null;
71+
if (activeSpan == null) {
72+
return;
7073
}
7174
// If there is a noop span in the active scope, we can clean all the way to this scope
72-
if (activeSpan() == noopSpan()) {
73-
return activeScope;
75+
if (activeSpan == noopSpan()) {
76+
return;
7477
}
7578
// Create an active scope with a noop span, and clean all the way to the previous scope
7679
activateSpan(noopSpan(), false);
77-
return activeScope;
7880
}
7981

8082
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
81-
public static void exit(@Advice.Enter final AgentScope scope) {
82-
// Clean up any leaking scopes from pekko-streams/pekko-http et.c.
83-
AgentScope activeScope = activeScope();
84-
while (activeScope != null && activeScope != scope) {
85-
activeScope.close();
86-
activeScope = activeScope();
87-
}
83+
public static void exit() {
84+
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
85+
rollbackActiveToCheckpoint();
8886
}
8987
}
9088
}

0 commit comments

Comments
 (0)