From 0fcfdb94f66301597b6d5787931c10afc7d0ea5e Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Fri, 28 Mar 2025 22:51:20 +0000 Subject: [PATCH 1/2] Optimization when span replaces span --- .../datadog/trace/bootstrap/instrumentation/api/AgentSpan.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java index e9b21f797b9..b09a48d2547 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java @@ -210,6 +210,6 @@ default T get(@Nonnull ContextKey key) { @Override default Context with(@Nonnull ContextKey key, @Nullable T value) { - return Context.root().with(SPAN_KEY, this, key, value); + return SPAN_KEY == key ? (Context) value : Context.root().with(SPAN_KEY, this, key, value); } } From 4d4a759c4b136f5f106f71a36c76e52bff8306d5 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 31 Mar 2025 11:51:03 +0100 Subject: [PATCH 2/2] Temporary bridge between contexts and spans/scopes --- .../core/scopemanager/ContinuableScope.java | 27 +++++-- .../scopemanager/ContinuableScopeManager.java | 75 +++++++++++++------ .../core/scopemanager/ContinuingScope.java | 6 +- .../trace/core/scopemanager/ScopeStack.java | 11 ++- .../instrumentation/api/AgentScope.java | 9 ++- 5 files changed, 93 insertions(+), 35 deletions(-) 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 7993f885f92..23441c9fa3b 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 @@ -1,5 +1,6 @@ package datadog.trace.core.scopemanager; +import datadog.context.Context; import datadog.trace.api.Stateful; import datadog.trace.api.scopemanager.ExtendedScopeListener; import datadog.trace.api.scopemanager.ScopeListener; @@ -15,10 +16,11 @@ class ContinuableScope implements AgentScope, AttachableWrapper { static final byte INSTRUMENTATION = 0; static final byte MANUAL = 1; static final byte ITERATION = 2; + static final byte CONTEXT = 3; private final ContinuableScopeManager scopeManager; - final AgentSpan span; // package-private so scopeManager can access it directly + final Context context; // package-private so scopeManager can access it directly /** Flag that this scope should be allowed to propagate across async boundaries. */ private static final byte ASYNC_PROPAGATING = 1; @@ -40,12 +42,12 @@ class ContinuableScope implements AgentScope, AttachableWrapper { ContinuableScope( final ContinuableScopeManager scopeManager, - final AgentSpan span, + final Context context, final byte source, final boolean isAsyncPropagating, final Stateful scopeState) { this.scopeManager = scopeManager; - this.span = span; + this.context = context; this.source = source; this.flags = isAsyncPropagating ? ASYNC_PROPAGATING : 0; this.scopeState = scopeState; @@ -65,7 +67,7 @@ public final void close() { byte source = source(); scopeManager.healthMetrics.onScopeCloseError(source == MANUAL); if (source == MANUAL && scopeManager.strictMode) { - throw new RuntimeException("Tried to close " + span + " scope when not on top"); + throw new RuntimeException("Tried to close " + context + " scope when not on top"); } } @@ -131,7 +133,12 @@ public final boolean isAsyncPropagating() { @Override public final AgentSpan span() { - return span; + return AgentSpan.fromContext(context); + } + + @Override + public Context context() { + return context; } @Override @@ -145,7 +152,7 @@ public final void setAsyncPropagation(final boolean value) { @Override public final String toString() { - return super.toString() + "->" + span; + return super.toString() + "->" + context; } public void checkpoint() { @@ -162,6 +169,10 @@ public boolean rollback() { } public final void beforeActivated() { + AgentSpan span = span(); + if (span == null) { + return; + } try { scopeState.activate(span.context()); } catch (Throwable e) { @@ -171,6 +182,10 @@ public final void beforeActivated() { } public final void afterActivated() { + AgentSpan span = span(); + if (span == null) { + return; + } for (final ScopeListener listener : scopeManager.scopeListeners) { try { listener.afterScopeActivated(); 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 5e3cafd4024..56c9a79bb5d 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 @@ -3,6 +3,7 @@ import static datadog.trace.api.ConfigDefaults.DEFAULT_ASYNC_PROPAGATING; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan; +import static datadog.trace.core.scopemanager.ContinuableScope.CONTEXT; import static datadog.trace.core.scopemanager.ContinuableScope.INSTRUMENTATION; import static datadog.trace.core.scopemanager.ContinuableScope.ITERATION; import static datadog.trace.core.scopemanager.ContinuableScope.MANUAL; @@ -10,6 +11,9 @@ import static java.util.concurrent.TimeUnit.NANOSECONDS; import static java.util.concurrent.TimeUnit.SECONDS; +import datadog.context.Context; +import datadog.context.ContextManager; +import datadog.context.ContextScope; import datadog.trace.api.Config; import datadog.trace.api.Stateful; import datadog.trace.api.scopemanager.ExtendedScopeListener; @@ -40,7 +44,7 @@ * from being reported even if all related spans are finished. It also delegates to other * ScopeInterceptors to provide additional functionality. */ -public final class ContinuableScopeManager implements ScopeStateAware { +public final class ContinuableScopeManager implements ScopeStateAware, ContextManager { static final Logger log = LoggerFactory.getLogger(ContinuableScopeManager.class); static final RatelimitedLogger ratelimitedLog = new RatelimitedLogger(log, 1, MINUTES); @@ -83,6 +87,8 @@ public ContinuableScopeManager( this.healthMetrics = healthMetrics; this.tlsScopeStack = new ScopeStackThreadLocal(profilingContextIntegration); this.profilingContextIntegration = profilingContextIntegration; + + ContextManager.register(this); } public AgentScope activateSpan(final AgentSpan span) { @@ -96,10 +102,12 @@ public AgentScope activateManualSpan(final AgentSpan span) { public AgentScope.Continuation captureActiveSpan() { ContinuableScope activeScope = scopeStack().active(); if (null != activeScope && activeScope.isAsyncPropagating()) { - return captureSpan(activeScope.span(), activeScope.source()); - } else { - return AgentTracer.noopContinuation(); + AgentSpan span = activeScope.span(); + if (span != null) { + return captureSpan(span, activeScope.source()); + } } + return AgentTracer.noopContinuation(); } public AgentScope.Continuation captureSpan(final AgentSpan span) { @@ -111,14 +119,14 @@ private AgentScope.Continuation captureSpan(final AgentSpan span, byte source) { } private AgentScope activate( - final AgentSpan span, + final Context context, final byte source, final boolean overrideAsyncPropagation, final boolean isAsyncPropagating) { ScopeStack scopeStack = scopeStack(); final ContinuableScope top = scopeStack.top; - if (top != null && top.span.equals(span)) { + if (top != null && top.context.equals(context)) { top.incrementReferences(); return top; } @@ -131,7 +139,7 @@ private AgentScope activate( return noopScope(); } - assert span != null; + assert context != null; // Inherit the async propagation from the active scope unless the value is overridden boolean asyncPropagation = @@ -140,7 +148,7 @@ private AgentScope activate( : top != null ? top.isAsyncPropagating() : DEFAULT_ASYNC_PROPAGATING; final ContinuableScope scope = - new ContinuableScope(this, span, source, asyncPropagation, createScopeState(span)); + new ContinuableScope(this, context, source, asyncPropagation, createScopeState(context)); scopeStack.push(scope); healthMetrics.onActivateScope(); @@ -153,13 +161,13 @@ private AgentScope activate( * @param continuation {@code null} if a continuation is re-used */ ContinuableScope continueSpan( - final ScopeContinuation continuation, final AgentSpan span, final byte source) { + final ScopeContinuation continuation, final Context context, final byte source) { ScopeStack scopeStack = scopeStack(); // optimization: if the top scope is already keeping the same span alive // then re-use that scope (avoids allocation) and cancel the continuation final ContinuableScope top = scopeStack.top; - if (top != null && top.span.equals(span)) { + if (top != null && top.context.equals(context)) { top.incrementReferences(); if (continuation != null) { continuation.cancelFromContinuedScopeClose(); @@ -167,12 +175,12 @@ ContinuableScope continueSpan( return top; } - Stateful scopeState = createScopeState(span); + Stateful scopeState = createScopeState(context); final ContinuableScope scope; if (continuation != null) { - scope = new ContinuingScope(this, span, source, true, continuation, scopeState); + scope = new ContinuingScope(this, context, source, true, continuation, scopeState); } else { - scope = new ContinuableScope(this, span, source, true, scopeState); + scope = new ContinuableScope(this, context, source, true, scopeState); } scopeStack.push(scope); @@ -202,8 +210,9 @@ public void closePrevious(final boolean finishSpan) { } top.close(); scopeStack.cleanup(); - if (finishSpan) { - top.span.finishWithEndToEnd(); + AgentSpan span = top.span(); + if (finishSpan && span != null) { + span.finishWithEndToEnd(); } } } @@ -261,7 +270,7 @@ public void rollbackActiveToCheckpoint() { public AgentSpan activeSpan() { final ContinuableScope active = scopeStack().active(); - return active == null ? null : active.span; + return active == null ? null : active.span(); } /** Attach a listener to scope activation events */ @@ -289,11 +298,12 @@ private void addExtendedScopeListener(final ExtendedScopeListener listener) { } } - private Stateful createScopeState(AgentSpan span) { + private Stateful createScopeState(Context context) { // currently this just manages things the profiler has to do per scope, but could be expanded // to encapsulate other scope lifecycle activities // FIXME DDSpanContext is always a ProfilerContext anyway... - if (span.context() instanceof ProfilerContext) { + AgentSpan span = AgentSpan.fromContext(context); + if (span != null && span.context() instanceof ProfilerContext) { return profilingContextIntegration.newScopeState((ProfilerContext) span.context()); } return Stateful.DEFAULT; @@ -308,6 +318,22 @@ public ScopeState newScopeState() { return new ContinuableScopeState(); } + @Override + public Context current() { + final ContinuableScope active = scopeStack().active(); + return active == null ? Context.root() : active.context; + } + + @Override + public ContextScope attach(Context context) { + return activate(context, CONTEXT, false, true); + } + + @Override + public Context swap(Context context) { + throw new UnsupportedOperationException("Not yet implemented"); + } + private class ContinuableScopeState implements ScopeState { private ScopeStack localScopeStack = tlsScopeStack.initialValue(); @@ -383,11 +409,14 @@ public void run(Map rootIterationScopes) { if (!rootScope.alive()) { // no need to track this anymore itr.remove(); - } else if (NANOSECONDS.toMillis(rootScope.span.getStartTime()) < cutOff) { - // mark scope as overdue to allow cleanup and avoid further spans being attached - scopeStack.overdueRootScope = rootScope; - rootScope.span.finishWithEndToEnd(); - itr.remove(); + } else { + AgentSpan span = rootScope.span(); + if (span != null && NANOSECONDS.toMillis(span.getStartTime()) < cutOff) { + // mark scope as overdue to allow cleanup and avoid further spans being attached + scopeStack.overdueRootScope = rootScope; + span.finishWithEndToEnd(); + itr.remove(); + } } } } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuingScope.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuingScope.java index 0ead3d4f725..c5140bc2f35 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuingScope.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuingScope.java @@ -1,7 +1,7 @@ package datadog.trace.core.scopemanager; +import datadog.context.Context; import datadog.trace.api.Stateful; -import datadog.trace.bootstrap.instrumentation.api.AgentSpan; final class ContinuingScope extends ContinuableScope { /** Continuation that created this scope. */ @@ -9,12 +9,12 @@ final class ContinuingScope extends ContinuableScope { ContinuingScope( final ContinuableScopeManager scopeManager, - final AgentSpan span, + final Context context, final byte source, final boolean isAsyncPropagating, final ScopeContinuation continuation, final Stateful scopeState) { - super(scopeManager, span, source, isAsyncPropagating, scopeState); + super(scopeManager, context, source, isAsyncPropagating, scopeState); this.continuation = continuation; } diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java index f7de0f7d101..7dcc3e2af29 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java @@ -2,6 +2,7 @@ import static datadog.trace.core.scopemanager.ContinuableScope.ITERATION; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; import java.util.ArrayDeque; @@ -84,14 +85,20 @@ final boolean checkOverdueScopes(final ContinuableScope expectedScope) { // avoid calling close() as we're already in that method, instead just clear any // remaining references so the scope gets removed in the subsequent cleanup() call top.clearReferences(); - top.span.finishWithEndToEnd(); + AgentSpan span = top.span(); + if (span != null) { + span.finishWithEndToEnd(); + } // now do the same for any previous iteration scopes ahead of the expected scope for (ContinuableScope scope : stack) { if (scope.source() != ITERATION) { return expectedScope.equals(scope); } else { scope.clearReferences(); - scope.span.finishWithEndToEnd(); + span = scope.span(); + if (span != null) { + span.finishWithEndToEnd(); + } } } return false; // we didn't find the expected scope diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentScope.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentScope.java index ae9d619a142..d17a70add2b 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentScope.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentScope.java @@ -1,11 +1,18 @@ package datadog.trace.bootstrap.instrumentation.api; +import datadog.context.Context; +import datadog.context.ContextScope; import datadog.trace.context.TraceScope; import java.io.Closeable; -public interface AgentScope extends TraceScope, Closeable { +public interface AgentScope extends ContextScope, TraceScope, Closeable { AgentSpan span(); + @Override + default Context context() { + return span(); + } + @Override void close();