From adf6612bd99e124f355dcc6ed6ce5c7faf0752c7 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Thu, 6 Mar 2025 10:37:37 +0000 Subject: [PATCH 1/3] Document that activeScope() is to be removed --- .../trace/bootstrap/instrumentation/api/AgentTracer.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index 789a3e6a388..b06ba117212 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -160,6 +160,8 @@ public static AgentSpan activeSpan() { return get().activeSpan(); } + /** @deprecated To be removed, do not use. */ + @Deprecated public static AgentScope activeScope() { return get().activeScope(); } From dbfdb196a82b705f0d408e99d4ba7424591f80b3 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Tue, 4 Mar 2025 18:10:03 +0000 Subject: [PATCH 2/3] Add support for checkpointActiveForRollback...rollbackActiveToCheckpoint to clean up leaked scopes --- .../java/datadog/trace/core/CoreTracer.java | 10 +++++ .../core/scopemanager/ContinuableScope.java | 38 +++++++++++++++---- .../scopemanager/ContinuableScopeManager.java | 18 +++++++++ .../instrumentation/api/AgentTracer.java | 37 +++++++++++++++++- 4 files changed, 94 insertions(+), 9 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java index 49eb2fe5728..cd7036edb66 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java @@ -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(); diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java index a2707e9df68..f8bea7cc160 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java @@ -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; @@ -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; } @@ -116,7 +121,7 @@ final boolean alive() { @Override public final boolean isAsyncPropagating() { - return isAsyncPropagating; + return (flags & ASYNC_PROPAGATING) != 0; } @Override @@ -126,7 +131,11 @@ public final AgentSpan span() { @Override public final void setAsyncPropagation(final boolean value) { - isAsyncPropagating = value; + if (value) { + flags |= ASYNC_PROPAGATING; + } else { + flags &= ~ASYNC_PROPAGATING; + } } @Override @@ -134,6 +143,19 @@ 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()); @@ -164,7 +186,7 @@ public final void afterActivated() { @Override public byte source() { - return (byte) (flags & 0x7F); + return (byte) (source & 0x7F); } @Override diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java index 5dd4485bc46..5fcf4c6d6b3 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java @@ -218,6 +218,24 @@ public AgentScope active() { return scopeStack().active(); } + public void checkpointActiveForRollback() { + ContinuableScope active = scopeStack().active(); + if (active != null) { + active.checkpoint(); + } + } + + public void rollbackActiveToCheckpoint() { + ContinuableScope active; + while ((active = scopeStack().active()) != null) { + if (active.rollback()) { + active.close(); + } else { + break; // stop at the most recent checkpointed scope + } + } + } + public AgentSpan activeSpan() { final ContinuableScope active = scopeStack().active(); return active == null ? null : active.span; diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java index b06ba117212..7dfbc56d251 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java @@ -121,10 +121,35 @@ public static AgentScope.Continuation captureSpan(final AgentSpan span) { return get().captureSpan(span); } + /** + * Checkpoints the active scope. A subsequent call to {@link #rollbackActiveToCheckpoint()} closes + * outstanding scopes up to but not including the most recent checkpointed scope. + * + * @deprecated This should only be used when scopes might leak onto the scope stack which cannot + * be cleaned up by other means. + */ + @Deprecated + public static void checkpointActiveForRollback() { + get().checkpointActiveForRollback(); + } + + /** + * Closes outstanding scopes up to but not including the most recent scope checkpointed with + * {@link #checkpointActiveForRollback()}. Closes all scopes if none have been checkpointed. + * + * @deprecated This should only be used when scopes have leaked onto the scope stack that cannot + * be cleaned up by other means. + */ + @Deprecated + public static void rollbackActiveToCheckpoint() { + get().rollbackActiveToCheckpoint(); + } + /** * Closes the scope for the currently active span. * - * @deprecated Prefer closing the scope returned by {@link #activateSpan} when available. + * @deprecated This should only be used when an instrumentation does not have access to the + * original scope returned by {@link #activateSpan}. */ @Deprecated public static void closeActive() { @@ -322,6 +347,10 @@ AgentSpan startSpan( AgentScope.Continuation captureSpan(AgentSpan span); + void checkpointActiveForRollback(); + + void rollbackActiveToCheckpoint(); + void closeActive(); void closePrevious(boolean finishSpan); @@ -477,6 +506,12 @@ public boolean isAsyncPropagationEnabled() { @Override public void setAsyncPropagationEnabled(boolean asyncPropagationEnabled) {} + @Override + public void checkpointActiveForRollback() {} + + @Override + public void rollbackActiveToCheckpoint() {} + @Override public void closeActive() {} From 3efa8ff025246105ce628e9c1108acb81e7275d4 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Thu, 6 Mar 2025 10:39:50 +0000 Subject: [PATCH 3/3] Use checkpoint..rollback in Akka/Pekko instrumentations to cleanup leaky scopes --- .../AkkaActorCellInstrumentation.java | 47 +++++++++---------- .../AkkaMailboxInstrumentation.java | 29 +++++------- .../PekkoActorCellInstrumentation.java | 47 +++++++++---------- .../PekkoMailboxInstrumentation.java | 30 ++++++------ 4 files changed, 73 insertions(+), 80 deletions(-) diff --git a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java index 84bb970c141..db55e7ceddb 100644 --- a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaActorCellInstrumentation.java @@ -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; @@ -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; @@ -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(); } } } diff --git a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java index 64e894bb59e..4b6e6b8355a 100644 --- a/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaMailboxInstrumentation.java @@ -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.singletonList; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -12,7 +13,7 @@ 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; @@ -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(); } } } diff --git a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java index 45aae18c813..9f242c70b3b 100644 --- a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoActorCellInstrumentation.java @@ -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; @@ -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; @@ -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(); } } } diff --git a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java index 43434dce387..bf8772e5a9f 100644 --- a/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-concurrent/src/main/java/datadog/trace/instrumentation/pekko/concurrent/PekkoMailboxInstrumentation.java @@ -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.singletonList; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -12,7 +13,7 @@ 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; @@ -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(); } } }