Skip to content

Commit 0666be6

Browse files
mccullsmcculls
authored
Add support for checkpoint...rollback API to clean up leaked scopes (#8516)
* Document that activeScope() is to be removed * Add support for checkpointActiveForRollback...rollbackActiveToCheckpoint to clean up leaked scopes * Use checkpoint..rollback in Akka/Pekko instrumentations to cleanup leaky scopes
1 parent 2575c71 commit 0666be6

File tree

8 files changed

+169
-89
lines changed

8 files changed

+169
-89
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
}

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -977,6 +977,16 @@ public AgentScope activeScope() {
977977
return scopeManager.active();
978978
}
979979

980+
@Override
981+
public void checkpointActiveForRollback() {
982+
this.scopeManager.checkpointActiveForRollback();
983+
}
984+
985+
@Override
986+
public void rollbackActiveToCheckpoint() {
987+
this.scopeManager.rollbackActiveToCheckpoint();
988+
}
989+
980990
@Override
981991
public void closeActive() {
982992
AgentScope activeScope = this.scopeManager.active();

‎dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,15 @@ class ContinuableScope implements AgentScope, AttachableWrapper {
1515

1616
final AgentSpan span; // package-private so scopeManager can access it directly
1717

18-
/** Flag to propagate this scope across async boundaries. */
19-
private boolean isAsyncPropagating;
18+
/** Flag that this scope should be allowed to propagate across async boundaries. */
19+
private static final byte ASYNC_PROPAGATING = 1;
2020

21-
private final byte flags;
21+
/** Flag that we intend to roll back the scope stack to this scope in the future. */
22+
private static final byte CHECKPOINTED = 2;
23+
24+
private byte flags;
25+
26+
private final byte source;
2227

2328
private short referenceCount = 1;
2429

@@ -36,8 +41,8 @@ class ContinuableScope implements AgentScope, AttachableWrapper {
3641
final Stateful scopeState) {
3742
this.scopeManager = scopeManager;
3843
this.span = span;
39-
this.flags = source;
40-
this.isAsyncPropagating = isAsyncPropagating;
44+
this.source = source;
45+
this.flags = isAsyncPropagating ? ASYNC_PROPAGATING : 0;
4146
this.scopeState = scopeState;
4247
}
4348

@@ -116,7 +121,7 @@ final boolean alive() {
116121

117122
@Override
118123
public final boolean isAsyncPropagating() {
119-
return isAsyncPropagating;
124+
return (flags & ASYNC_PROPAGATING) != 0;
120125
}
121126

122127
@Override
@@ -126,14 +131,31 @@ public final AgentSpan span() {
126131

127132
@Override
128133
public final void setAsyncPropagation(final boolean value) {
129-
isAsyncPropagating = value;
134+
if (value) {
135+
flags |= ASYNC_PROPAGATING;
136+
} else {
137+
flags &= ~ASYNC_PROPAGATING;
138+
}
130139
}
131140

132141
@Override
133142
public final String toString() {
134143
return super.toString() + "->" + span;
135144
}
136145

146+
public void checkpoint() {
147+
flags |= CHECKPOINTED;
148+
}
149+
150+
public boolean rollback() {
151+
if ((flags & CHECKPOINTED) != 0) {
152+
flags &= ~CHECKPOINTED;
153+
return false;
154+
} else {
155+
return true;
156+
}
157+
}
158+
137159
public final void beforeActivated() {
138160
try {
139161
scopeState.activate(span.context());
@@ -164,7 +186,7 @@ public final void afterActivated() {
164186

165187
@Override
166188
public byte source() {
167-
return (byte) (flags & 0x7F);
189+
return (byte) (source & 0x7F);
168190
}
169191

170192
@Override

0 commit comments

Comments
 (0)