Skip to content

Add support for checkpoint...rollback API to clean up leaked scopes #8516

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 7, 2025
Merged

Add support for checkpoint...rollback API to clean up leaked scopes #8516

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
adf6612
Document that activeScope() is to be removed
mcculls Mar 6, 2025
dbfdb19
Add support for checkpointActiveForRollback...rollbackActiveToCheckpo…
mcculls Mar 4, 2025
3efa8ff
Use checkpoint..rollback in Akka/Pekko instrumentations to cleanup le…
mcculls Mar 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 23 additions & 24 deletions ...main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;

Expand All @@ -14,6 +15,7 @@
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
import java.util.Map;
Expand Down Expand Up @@ -56,43 +58,40 @@ public void methodAdvice(MethodTransformer transformer) {
*/
public static class InvokeAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope enter(
public static void enter(
@Advice.Argument(value = 0) Envelope envelope,
@Advice.Local(value = "localScope") AgentScope localScope) {
AgentScope activeScope = activeScope();
localScope =
@Advice.Local(value = "taskScope") AgentScope taskScope) {
checkpointActiveForRollback();
// note: task scope may be the same as the scope we want to roll back to,
// so we must remember to close it on exit to balance the activation count
taskScope =
AdviceUtils.startTaskScope(
InstrumentationContext.get(Envelope.class, State.class), envelope);
// There was a scope created from the envelop, so use that
if (localScope != null) {
return activeScope;
// There was a scope created from the envelope, so use that
if (taskScope != null) {
return;
}
AgentSpan activeSpan = activeSpan();
// If there is no active scope, we can clean all the way to the bottom
if (null == activeScope) {
return null;
if (activeSpan == null) {
return;
}
// If there is a noop span in the active scope, we can clean all the way to this scope
if (activeSpan() == noopSpan()) {
return activeScope;
if (activeSpan == noopSpan()) {
return;
}
// Create an active scope with a noop span, and clean all the way to the previous scope
localScope = activateSpan(noopSpan(), false);
return activeScope;
activateSpan(noopSpan(), false);
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void exit(
@Advice.Enter AgentScope scope, @Advice.Local(value = "localScope") AgentScope localScope) {
if (localScope != null) {
public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) {
if (taskScope != null) {
// then we have invoked an Envelope and need to mark the work complete
localScope.close();
}
// Clean up any leaking scopes from akka-streams/akka-http et.c.
AgentScope activeScope = activeScope();
while (activeScope != null && activeScope != scope) {
activeScope.close();
activeScope = activeScope();
taskScope.close();
}
// Clean up any leaking scopes from akka-streams/akka-http etc.
rollbackActiveToCheckpoint();
}
}
}
29 changes: 13 additions & 16 deletions ...c/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,18 @@

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
import static java.util.Collections.singletonList;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;

import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.ExcludeFilterProvider;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
import java.util.Collection;
import java.util.EnumMap;
Expand Down Expand Up @@ -62,29 +63,25 @@ public void methodAdvice(MethodTransformer transformer) {
*/
public static final class SuppressMailboxRunAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope enter() {
AgentScope activeScope = activeScope();
public static void enter() {
checkpointActiveForRollback();
AgentSpan activeSpan = activeSpan();
// If there is no active scope, we can clean all the way to the bottom
if (null == activeScope) {
return null;
if (activeSpan == null) {
return;
}
// If there is a noop span in the active scope, we can clean all the way to this scope
if (activeSpan() == noopSpan()) {
return activeScope;
if (activeSpan == noopSpan()) {
return;
}
// Create an active scope with a noop span, and clean all the way to the previous scope
activateSpan(noopSpan(), false);
return activeScope;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void exit(@Advice.Enter final AgentScope scope) {
// Clean up any leaking scopes from akka-streams/akka-http et.c.
AgentScope activeScope = activeScope();
while (activeScope != null && activeScope != scope) {
activeScope.close();
activeScope = activeScope();
}
public static void exit() {
// Clean up any leaking scopes from akka-streams/akka-http etc.
rollbackActiveToCheckpoint();
}
}
}
47 changes: 23 additions & 24 deletions ...in/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;

Expand All @@ -13,6 +14,7 @@
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils;
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
import java.util.Map;
Expand Down Expand Up @@ -56,43 +58,40 @@ public void methodAdvice(MethodTransformer transformer) {
*/
public static class InvokeAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope enter(
public static void enter(
@Advice.Argument(value = 0) Envelope envelope,
@Advice.Local(value = "localScope") AgentScope localScope) {
AgentScope activeScope = activeScope();
localScope =
@Advice.Local(value = "taskScope") AgentScope taskScope) {
checkpointActiveForRollback();
// note: task scope may be the same as the scope we want to roll back to,
// so we must remember to close it on exit to balance the activation count
taskScope =
AdviceUtils.startTaskScope(
InstrumentationContext.get(Envelope.class, State.class), envelope);
// There was a scope created from the envelop, so use that
if (localScope != null) {
return activeScope;
// There was a scope created from the envelope, so use that
if (taskScope != null) {
return;
}
AgentSpan activeSpan = activeSpan();
// If there is no active scope, we can clean all the way to the bottom
if (null == activeScope) {
return null;
if (activeSpan == null) {
return;
}
// If there is a noop span in the active scope, we can clean all the way to this scope
if (activeSpan() == noopSpan()) {
return activeScope;
if (activeSpan == noopSpan()) {
return;
}
// Create an active scope with a noop span, and clean all the way to the previous scope
localScope = activateSpan(noopSpan(), false);
return activeScope;
activateSpan(noopSpan(), false);
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void exit(
@Advice.Enter AgentScope scope, @Advice.Local(value = "localScope") AgentScope localScope) {
if (localScope != null) {
public static void exit(@Advice.Local(value = "taskScope") AgentScope taskScope) {
if (taskScope != null) {
// then we have invoked an Envelope and need to mark the work complete
localScope.close();
}
// Clean up any leaking scopes from pekko-streams/pekko-http et.c.
AgentScope activeScope = activeScope();
while (activeScope != null && activeScope != scope) {
activeScope.close();
activeScope = activeScope();
taskScope.close();
}
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
rollbackActiveToCheckpoint();
}
}
}
30 changes: 14 additions & 16 deletions ...main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,18 @@

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.checkpointActiveForRollback;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.rollbackActiveToCheckpoint;
import static java.util.Collections.singletonList;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;

import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.ExcludeFilterProvider;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
import java.util.Collection;
import java.util.EnumMap;
Expand Down Expand Up @@ -62,29 +63,26 @@ public void methodAdvice(MethodTransformer transformer) {
*/
public static final class SuppressMailboxRunAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope enter() {
AgentScope activeScope = activeScope();
public static void enter() {
checkpointActiveForRollback();

AgentSpan activeSpan = activeSpan();
// If there is no active scope, we can clean all the way to the bottom
if (null == activeScope) {
return null;
if (activeSpan == null) {
return;
}
// If there is a noop span in the active scope, we can clean all the way to this scope
if (activeSpan() == noopSpan()) {
return activeScope;
if (activeSpan == noopSpan()) {
return;
}
// Create an active scope with a noop span, and clean all the way to the previous scope
activateSpan(noopSpan(), false);
return activeScope;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void exit(@Advice.Enter final AgentScope scope) {
// Clean up any leaking scopes from pekko-streams/pekko-http et.c.
AgentScope activeScope = activeScope();
while (activeScope != null && activeScope != scope) {
activeScope.close();
activeScope = activeScope();
}
public static void exit() {
// Clean up any leaking scopes from pekko-streams/pekko-http etc.
rollbackActiveToCheckpoint();
}
}
}
10 changes: 10 additions & 0 deletions dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -977,6 +977,16 @@ public AgentScope activeScope() {
return scopeManager.active();
}

@Override
public void checkpointActiveForRollback() {
this.scopeManager.checkpointActiveForRollback();
}

@Override
public void rollbackActiveToCheckpoint() {
this.scopeManager.rollbackActiveToCheckpoint();
}

@Override
public void closeActive() {
AgentScope activeScope = this.scopeManager.active();
Expand Down
38 changes: 30 additions & 8 deletions dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,15 @@ class ContinuableScope implements AgentScope, AttachableWrapper {

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

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

private final byte flags;
/** Flag that we intend to roll back the scope stack to this scope in the future. */
private static final byte CHECKPOINTED = 2;

private byte flags;

private final byte source;

private short referenceCount = 1;

Expand All @@ -36,8 +41,8 @@ class ContinuableScope implements AgentScope, AttachableWrapper {
final Stateful scopeState) {
this.scopeManager = scopeManager;
this.span = span;
this.flags = source;
this.isAsyncPropagating = isAsyncPropagating;
this.source = source;
this.flags = isAsyncPropagating ? ASYNC_PROPAGATING : 0;
this.scopeState = scopeState;
}

Expand Down Expand Up @@ -116,7 +121,7 @@ final boolean alive() {

@Override
public final boolean isAsyncPropagating() {
return isAsyncPropagating;
return (flags & ASYNC_PROPAGATING) != 0;
}

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

@Override
public final void setAsyncPropagation(final boolean value) {
isAsyncPropagating = value;
if (value) {
flags |= ASYNC_PROPAGATING;
} else {
flags &= ~ASYNC_PROPAGATING;
}
}

@Override
public final String toString() {
return super.toString() + "->" + span;
}

public void checkpoint() {
flags |= CHECKPOINTED;
}

public boolean rollback() {
if ((flags & CHECKPOINTED) != 0) {
flags &= ~CHECKPOINTED;
return false;
} else {
return true;
}
}

public final void beforeActivated() {
try {
scopeState.activate(span.context());
Expand Down Expand Up @@ -164,7 +186,7 @@ public final void afterActivated() {

@Override
public byte source() {
return (byte) (flags & 0x7F);
return (byte) (source & 0x7F);
}

@Override
Expand Down
Loading