From 63144518cf7e0cd9c79f08b7606c994ce2de46ff Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Tue, 20 Oct 2026 10:34:40 +0100 Subject: [PATCH 1/5] Add shared scope stacks for coroutine thread handoffs A coroutine can resume on another worker before its previous worker has finished restoring thread context. Add a legacy scope stack that follows the coroutine so every worker observes scope changes across that handoff. Publish the current scope for lock-free reads, serialize stack mutations and shared reference-count updates, and prevent closed or overdue scopes from being retained again. Run profiler hooks, scope listeners, and span finishing outside the stack lock, refreshing the top after profiler hooks. Expose shared stack creation through the tracer API, reuse shared handles and same-thread restore tokens, and preserve saved snapshots for ordinary cross-thread restoration. Update virtual-thread preloading for the saved context rename without eagerly loading shared stacks. Cover cross-thread closure, reference-count races, overdue scopes, context views, token reuse, reentrant callbacks, and notification ordering in core tests; add a virtual-thread class-loading regression. Co-Authored-By: Claude Opus 5.5 --- .../jdk21/VirtualThreadInstrumentation.java | 2 +- ...tualThreadScopeClassLoadingForkedTest.java | 53 + .../context/ContextManagerBenchmark.java | 8 +- .../java/datadog/trace/core/CoreTracer.java | 9 + .../core/scopemanager/ContinuableScope.java | 89 +- .../scopemanager/ContinuableScopeManager.java | 55 +- ...opeContext.java => SavedScopeContext.java} | 21 +- .../trace/core/scopemanager/ScopeStack.java | 57 +- .../core/scopemanager/ScopeStackContext.java | 31 + .../core/scopemanager/SharedScopeContext.java | 24 + .../core/scopemanager/SharedScopeStack.java | 208 +++ .../core/scopemanager/ThreadScopeContext.java | 33 + .../scopemanager/SharedScopeStackTest.java | 1151 +++++++++++++++++ .../instrumentation/api/AgentTracer.java | 13 + 14 files changed, 1702 insertions(+), 52 deletions(-) create mode 100644 dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java rename dd-trace-core/src/main/java/datadog/trace/core/scopemanager/{ScopeContext.java => SavedScopeContext.java} (66%) create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java create mode 100644 dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java index e105b42e26b..81d7fa472d5 100644 --- a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java @@ -73,7 +73,7 @@ public final class VirtualThreadInstrumentation extends InstrumenterModule.Conte */ private static final String[] PRELOAD_CLASS_NAMES = { VIRTUAL_THREAD_STATE_CLASS_NAME, - "datadog.trace.core.scopemanager.ScopeContext", + "datadog.trace.core.scopemanager.SavedScopeContext", "datadog.trace.core.scopemanager.ScopeStack" }; diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java new file mode 100644 index 00000000000..56bc4c5772a --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java @@ -0,0 +1,53 @@ +package testdog.trace.instrumentation.java.lang.jdk21; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.context.ContextScope; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.time.Duration; +import java.util.concurrent.locks.LockSupport; +import net.bytebuddy.agent.ByteBuddyAgent; +import org.junit.jupiter.api.Test; + +/** Isolates class-loading assertions from tests that may deliberately load coroutine helpers. */ +class VirtualThreadScopeClassLoadingForkedTest extends AbstractInstrumentationTest { + @Test + void virtualThreadHandoffsDoNotLoadSharedScopeStacks() throws InterruptedException { + AgentSpan parent = startSpan("test", "parent"); + AgentSpan[] active = new AgentSpan[2]; + try (ContextScope ignored = activateSpan(parent)) { + Thread child = + Thread.startVirtualThread( + () -> { + active[0] = activeSpan(); + LockSupport.parkNanos(MILLISECONDS.toNanos(1)); + active[1] = activeSpan(); + }); + assertTrue(child.join(Duration.ofSeconds(10))); + assertSame(parent, active[0]); + assertSame(parent, active[1]); + } finally { + parent.finish(); + } + assertTraces(trace(span().root().operationName("parent"))); + + for (Class loaded : ByteBuddyAgent.getInstrumentation().getAllLoadedClasses()) { + if (loaded != null) { + String name = loaded.getName(); + assertFalse( + "datadog.trace.core.scopemanager.SharedScopeStack".equals(name) + || "datadog.trace.core.scopemanager.SharedScopeContext".equals(name), + name); + } + } + } +} diff --git a/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java b/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java index 0128535a9d0..1764b66f32a 100644 --- a/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java +++ b/dd-trace-core/src/jmh/java/datadog/context/ContextManagerBenchmark.java @@ -235,10 +235,10 @@ public void nestedAttachAndClose(ThreadState thread) { * ContinuableScopeManager#swap(Context)} replaces the entire scope stack, making this a heavier * operation than in {@link ThreadLocalContextManager}. * - *

Note: GCProfiler will show allocation asymmetry here by design. {@link - * ContinuableScopeManager} swap allocates a {@code ScopeStack}, a {@code ContinuableScope}, and a - * {@code ScopeContext} per invocation; {@link ThreadLocalContextManager} swap is a plain field - * write. That asymmetry is the real cost of each manager's swap operation, not scaffolding. + *

This scenario uses ordinary contexts: entering creates a new scope stack and seed scope, and + * each swap creates a saved-stack token. The unused return token may be optimized away. + * Shared-stack swaps reuse their stack and restoration tokens and are not measured here. {@link + * ThreadLocalContextManager} swap is a plain field write. */ @Benchmark public void swapContexts(ThreadState thread) { 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 8f52b49052c..995f5a2c8fa 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 @@ -2687,4 +2687,13 @@ public Context swap(Context context) { public ContextContinuation capture(Context context) { return scopeManager.capture(context); } + + @Override + public Context sharedContextStack(Context initial) { + if (!InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + throw new IllegalStateException( + "sharedContextStack must not be called when context swap based logic is enabled"); + } + return scopeManager.sharedContextStack(initial); + } } 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 0f6feaf095c..cfacc375b47 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 @@ -6,6 +6,7 @@ import datadog.trace.api.scopemanager.ExtendedScopeListener; import datadog.trace.api.scopemanager.ScopeListener; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; class ContinuableScope implements ContextScope { @@ -27,6 +28,9 @@ class ContinuableScope implements ContextScope { private short referenceCount = 1; + // set when pushed onto a coroutine's shared scope stack, which other threads can also update + private boolean shared; + private final Stateful scopeState; ContinuableScope( @@ -50,7 +54,7 @@ public final void close() { if (!scopeStack.checkTop(this) && !scopeStack.checkOverdueScopes(this)) { if (ContinuableScopeManager.log.isDebugEnabled()) { ContinuableScopeManager.log.debug( - "Tried to close {} scope when not on top. Current top: {}", this, scopeStack.top); + "Tried to close {} scope when not on top. Current top: {}", this, scopeStack.top()); } byte source = source(); @@ -60,7 +64,7 @@ public final void close() { } } - final boolean alive = decrementReferences(); + final boolean alive = release(); scopeManager.healthMetrics.onCloseScope(); if (!alive) { scopeState.close(); @@ -105,20 +109,78 @@ final void onProperClose() { } } - final void incrementReferences() { + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Set before publishing the scope through SharedScopeStack.liveTop; the flag never " + + "changes after publication.") + final void markShared() { + shared = true; + } + + /** Adds a reference to a live scope -- returns false if the scope was already closed */ + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Only non-shared scopes use the unlocked write; shared scopes retain under their " + + "own monitor in retainShared().") + final boolean retain() { + if (shared) { + return retainShared(); + } + ++referenceCount; + return true; + } + + private synchronized boolean retainShared() { + // another thread may have closed the scope since it was read from the shared stack + if (referenceCount <= 0) { + return false; + } ++referenceCount; + return true; } - /** Decrements ref count -- returns true if the scope is still alive */ - final boolean decrementReferences() { + /** Removes a reference -- returns true if the scope is still alive */ + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Only non-shared scopes use the unlocked write; shared scopes release under their " + + "own monitor in releaseShared(), including closes from another stack.") + final boolean release() { + return shared ? releaseShared() : --referenceCount > 0; + } + + private synchronized boolean releaseShared() { + // the scope may be closed on any thread, whichever stack that thread has swapped in return --referenceCount > 0; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Only non-shared scopes use the unlocked write; shared scopes clear under their " + + "own monitor in clearShared().") final void clearReferences() { + if (shared) { + clearShared(); + } else { + referenceCount = 0; + } + } + + private synchronized void clearShared() { + // stops another thread retaining the scope once it's been cleared referenceCount = 0; } - /** Returns true if the scope is still alive (non-zero ref count) */ + /** + * Returns true if the scope's reference count is positive. + * + *

For shared scopes, this unlocked read is a best-effort cleanup check. Once closed, a shared + * scope cannot be retained again, so a stale read can delay cleanup but cannot cause a live scope + * to be removed. + */ final boolean alive() { return referenceCount > 0; } @@ -136,6 +198,11 @@ public Context context() { return context; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Changes are serialized with scope execution, including coroutine handoff between " + + "workers. Overlapping scope close/restore callbacks do not modify this flag.") public final void setAsyncPropagation(final boolean value) { asyncPropagating = value; } @@ -145,10 +212,20 @@ public final String toString() { return super.toString() + "->" + context; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Checkpoint updates are serialized with scope execution; coroutine handoff orders " + + "them between workers. Overlapping scope close/restore does not modify checkpoints.") public void checkpoint() { checkpointCount++; } + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "Checkpoint updates are serialized with scope execution; coroutine handoff orders " + + "them between workers. Overlapping scope close/restore does not modify checkpoints.") public boolean rollback() { if (checkpointCount > 0) { checkpointCount--; 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 e0ac7ea14fa..d60b1f458be 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 @@ -119,9 +119,8 @@ private ContextScope activate( ScopeStack scopeStack = scopeStack(); - final ContinuableScope top = scopeStack.top; - if (top != null && span.equals(top.span())) { - top.incrementReferences(); + final ContinuableScope top = scopeStack.top(); + if (top != null && span.equals(top.span()) && top.retain()) { return top; } @@ -154,9 +153,8 @@ private ContextScope activate( private ContextScope activate(final Context context) { ScopeStack scopeStack = scopeStack(); - final ContinuableScope top = scopeStack.top; - if (top != null && top.context.equals(context)) { - top.incrementReferences(); + final ContinuableScope top = scopeStack.top(); + if (top != null && top.context.equals(context) && top.retain()) { return top; } @@ -194,9 +192,8 @@ ContinuableScope continueSpan( // 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.context.equals(context)) { - top.incrementReferences(); + final ContinuableScope top = scopeStack.top(); + if (top != null && top.context.equals(context) && top.retain()) { if (continuation != null) { continuation.cancelFromContinuedScopeClose(); } @@ -231,7 +228,7 @@ public void closePrevious(final boolean finishSpan) { ScopeStack scopeStack = scopeStack(); // close any immediately previous iteration scope - final ContinuableScope top = scopeStack.top; + final ContinuableScope top = scopeStack.top(); if (top != null && top.source() == ITERATION) { if (iterationKeepAlive > 0) { // skip depth check because cancelling is cheap cancelRootIterationScopeCleanup(scopeStack, top); @@ -264,7 +261,7 @@ public ContextScope activateNext(final AgentSpan span) { assert span != null; - final ContinuableScope top = scopeStack.top; + final ContinuableScope top = scopeStack.top(); boolean asyncPropagation = top != null ? top.isAsyncPropagating() : DEFAULT_ASYNC_PROPAGATING; @@ -362,32 +359,30 @@ public ContextScope attach(@Nonnull Context context) { public Context swap(@Nonnull Context context) { ScopeStack oldStack = tlsScopeStack.get(); - ContinuableScope oldScope = oldStack.top; + ContinuableScope oldScope = oldStack.top(); ScopeStack newStack; - ContinuableScope newScope; - if (context instanceof ScopeContext) { + if (context instanceof ScopeStackContext) { // restore previously swapped context stack - newStack = ((ScopeContext) context).restore(profilingContextIntegration); - newScope = newStack.top; + newStack = ((ScopeStackContext) context).restore(profilingContextIntegration); } else if (context != Context.root()) { // start a new stack and record the new context as active newStack = new ScopeStack(profilingContextIntegration); - newScope = new ContinuableScope(this, context, CONTEXT, true, createScopeState(context)); + ContinuableScope newScope = + new ContinuableScope(this, context, CONTEXT, true, createScopeState(context)); newStack.top = newScope; } else { // start a new stack with no active context newStack = new ScopeStack(profilingContextIntegration); - newScope = null; } tlsScopeStack.set(newStack); - if (oldScope != newScope && newScope != null) { - newScope.beforeActivated(); - newScope.afterActivated(); + newStack.pruneClosedTop(); + if (oldScope != newStack.top()) { + newStack.activateCurrentTop(); } - return new ScopeContext(oldStack); + return newStack.replacing(oldStack); } public ContextContinuation capture(@Nonnull Context context) { @@ -410,6 +405,22 @@ public ContextContinuation capture(@Nonnull Context context) { return new ScopeContinuation(this, context, CONTEXT, traceCollector).register(); } + /** + * Returns a context whose scope stack is shared by every thread it is swapped onto. If the + * initial context is already shared, the new stack starts from its current top. + */ + public Context sharedContextStack(@Nonnull Context initial) { + if (initial instanceof ScopeStackContext) { + initial = ((ScopeStackContext) initial).currentContext(); + } + SharedScopeStack sharedStack = new SharedScopeStack(profilingContextIntegration); + if (initial != Context.root()) { + sharedStack.seed( + new ContinuableScope(this, initial, CONTEXT, true, createScopeState(initial))); + } + return sharedStack.handle; + } + static final class ScopeStackThreadLocal extends ThreadLocal { private final ProfilingContextIntegration profilingContextIntegration; diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SavedScopeContext.java similarity index 66% rename from dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java rename to dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SavedScopeContext.java index 8919fae96f4..4a51bb108bb 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SavedScopeContext.java @@ -1,29 +1,28 @@ package datadog.trace.core.scopemanager; import datadog.context.Context; -import datadog.context.ContextKey; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; -import javax.annotation.Nullable; /** Wraps a {@link ScopeStack} as a {@link Context} so it can be swapped back later. */ -final class ScopeContext implements Context { +final class SavedScopeContext extends ScopeStackContext { private final Thread originalThread = Thread.currentThread(); private final ScopeStack scopeStack; private final ContinuableScope parent; private final ContinuableScope active; private final Context context; - ScopeContext(ScopeStack scopeStack) { - this(scopeStack, scopeStack.top != null ? scopeStack.top.context : Context.root()); + SavedScopeContext(ScopeStack scopeStack) { + this(scopeStack, contextOf(scopeStack.top())); } - private ScopeContext(ScopeStack scopeStack, Context context) { + private SavedScopeContext(ScopeStack scopeStack, Context context) { this.scopeStack = scopeStack; this.parent = scopeStack.parent(); this.active = scopeStack.active(); this.context = context; } + @Override ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { // restore full stack on original thread, for other threads restore shallow copy if (Thread.currentThread() == originalThread) { @@ -33,14 +32,8 @@ ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { } } - @Nullable - @Override - public T get(ContextKey key) { - return context.get(key); - } - @Override - public Context with(ContextKey key, @Nullable T value) { - return context.with(key, value); + Context currentContext() { + return context; } } 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 84223a267d2..c1860fb3949 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 @@ -10,10 +10,10 @@ * The invariant is that the top of a non-empty stack is always active. Anytime a scope is closed, * cleanup() is called to ensure the invariant */ -final class ScopeStack { +class ScopeStack { private final ProfilingContextIntegration profilingContextIntegration; - private final ArrayDeque stack = new ArrayDeque<>(); // previous scopes + final ArrayDeque stack; // previous scopes ContinuableScope top; // current scope @@ -22,6 +22,13 @@ final class ScopeStack { ScopeStack(ProfilingContextIntegration profilingContextIntegration) { this.profilingContextIntegration = profilingContextIntegration; + this.stack = new ArrayDeque<>(); + } + + /** Creates an empty stack with room for the given number of previous scopes. */ + ScopeStack(ProfilingContextIntegration profilingContextIntegration, int numPrevious) { + this.profilingContextIntegration = profilingContextIntegration; + this.stack = new ArrayDeque<>(numPrevious); } /** Restore a shallow stack for async propagation purposes. */ @@ -36,6 +43,10 @@ final class ScopeStack { top = active; } + ContinuableScope top() { + return top; + } + ContinuableScope active() { // avoid attaching further spans to the root scope when it's been marked as overdue return top != overdueRootScope ? top : null; @@ -93,7 +104,7 @@ boolean checkTop(final ContinuableScope expectedScope) { * Slower check to see if overdue scopes ahead of the expected scope are all ITERATION scopes. * These represent iterations that are now out-of-scope and can be finished ready for cleanup. */ - final boolean checkOverdueScopes(final ContinuableScope expectedScope) { + boolean checkOverdueScopes(final ContinuableScope expectedScope) { // we already know 'top' isn't the expected scope, so just need to check its source if (top == null || top.source() != ITERATION) { return false; @@ -125,6 +136,42 @@ int depth() { return top != null ? 1 + stack.size() : 0; } + /** Returns a context that swaps this stack back in. */ + ScopeStackContext swappedOut() { + return new SavedScopeContext(this); + } + + /** Removes closed scopes after this stack is installed on the current thread. */ + void pruneClosedTop() {} + + /** Notifies profiler hooks and scope listeners about the current top. */ + void activateCurrentTop() { + ContinuableScope current = top; + if (current != null) { + current.beforeActivated(); + current.afterActivated(); + } + } + + /** Returns the context to restore after this stack replaces another stack. */ + ScopeStackContext replacing(ScopeStack oldStack) { + return oldStack.swappedOut(); + } + + private ThreadScopeContext sameThreadToken; // reused, as it takes no snapshot + + /** + * Returns a reusable context that swaps this stack back in, for callers that always swap it back + * on this thread. + */ + ScopeStackContext swappedOutForSameThread() { + ThreadScopeContext token = sameThreadToken; + if (token == null) { + sameThreadToken = token = new ThreadScopeContext(this); + } + return token; + } + // DQH - regrettably needed for pre-existing tests void clear() { stack.clear(); @@ -132,7 +179,7 @@ void clear() { } /** Notifies profiler that this thread has a context now */ - private void onBecomeNonEmpty() { + final void onBecomeNonEmpty() { try { profilingContextIntegration.onAttach(); } catch (Throwable e) { @@ -141,7 +188,7 @@ private void onBecomeNonEmpty() { } /** Notifies profiler that this thread no longer has a context */ - private void onBecomeEmpty() { + final void onBecomeEmpty() { try { profilingContextIntegration.onDetach(); } catch (Throwable e) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java new file mode 100644 index 00000000000..a285bb1cf66 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java @@ -0,0 +1,31 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import javax.annotation.Nullable; + +/** A {@link Context} that swaps a {@link ScopeStack} back in. */ +abstract class ScopeStackContext implements Context { + + /** Returns the stack to install on the current thread. */ + abstract ScopeStack restore(ProfilingContextIntegration profilingContextIntegration); + + /** The context that lookups and updates on this token are forwarded to. */ + abstract Context currentContext(); + + @Nullable + @Override + public final T get(ContextKey key) { + return currentContext().get(key); + } + + @Override + public final Context with(ContextKey key, @Nullable T value) { + return currentContext().with(key, value); + } + + static Context contextOf(@Nullable ContinuableScope scope) { + return scope != null ? scope.context : Context.root(); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java new file mode 100644 index 00000000000..7c225ce9884 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeContext.java @@ -0,0 +1,24 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; + +/** Swaps the same {@link SharedScopeStack} onto every thread, instead of a copy. */ +final class SharedScopeContext extends ScopeStackContext { + private final SharedScopeStack sharedStack; + + SharedScopeContext(SharedScopeStack sharedStack) { + this.sharedStack = sharedStack; + } + + /** The context of the shared stack's current top scope. */ + @Override + Context currentContext() { + return contextOf(sharedStack.top()); + } + + @Override + ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { + return sharedStack; + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java new file mode 100644 index 00000000000..910b326677f --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/SharedScopeStack.java @@ -0,0 +1,208 @@ +package datadog.trace.core.scopemanager; + +import static datadog.trace.core.scopemanager.ContinuableScope.ITERATION; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.util.ArrayList; +import java.util.List; + +/** + * A {@link ScopeStack} owned by a coroutine and shared by whichever threads are running it. + * + *

Threads normally hand the coroutine over in sequence, but the previous thread can still be + * unwinding while the next one starts. Stack mutations take this stack's lock; reference-count + * changes take the scope's lock. The top scope is published through a volatile field so the + * frequent active-scope reads need no lock. + * + *

Scope listeners, profiler hooks and span finishing run outside the stack lock to avoid + * deadlocks with application code. Notifications reread the live top, but legacy listeners receive + * no context argument, so another worker can still change the active scope during notification. + */ +final class SharedScopeStack extends ScopeStack { + private volatile ContinuableScope liveTop; + // written before liveTop, so it's visible to readers that read top() first + private int liveDepth; + + /** Swaps this stack in on any thread; also returned when this stack is swapped out. */ + final SharedScopeContext handle; + + SharedScopeStack(ProfilingContextIntegration profilingContextIntegration) { + super(profilingContextIntegration, 4); + this.handle = new SharedScopeContext(this); + } + + /** Seeds the stack with an initial scope, before it is shared with other threads. */ + void seed(ContinuableScope scope) { + scope.markShared(); + publishTop(scope); + } + + /** Updates the top scope, publishing it to lock-free readers before any listeners run. */ + @SuppressFBWarnings( + value = "AT_STALE_THREAD_WRITE_OF_PRIMITIVE", + justification = + "The following volatile liveTop write publishes liveDepth to readers that acquire " + + "the top. Depth is a best-effort limit check, so concurrent staleness is acceptable.") + private void publishTop(ContinuableScope scope) { + top = scope; + liveDepth = scope != null ? 1 + stack.size() : 0; + liveTop = scope; + } + + @Override + ContinuableScope top() { + return liveTop; + } + + @Override + ContinuableScope active() { + // avoid attaching further spans to the root scope when it's been marked as overdue + ContinuableScope current = liveTop; + return current != overdueRootScope ? current : null; + } + + /** Removes closed scopes after swapping in; the swap activates the surviving top. */ + @Override + void pruneClosedTop() { + ContinuableScope current = liveTop; + // only take the lock when there's something to remove + if (current != null && !current.alive()) { + cleanup(false); + } + } + + @Override + ScopeStackContext swappedOut() { + // the shared handle always reads the live stack, so there's no snapshot to take + return handle; + } + + @Override + ScopeStackContext swappedOutForSameThread() { + return handle; + } + + @Override + ScopeStackContext replacing(ScopeStack oldStack) { + // Coroutines always swap the replaced stack back on this thread, so need no snapshot. + return oldStack.swappedOutForSameThread(); + } + + @Override + void cleanup() { + cleanup(true); + } + + // only reactivates the new top when asked, as swap() activates it after swapping in + private void cleanup(boolean activateTop) { + ContinuableScope closedScope; + int closedCount = 0; + synchronized (this) { + ContinuableScope curScope = top; + closedScope = curScope; + while (curScope != null && !curScope.alive()) { + closedCount++; + curScope = stack.poll(); + } + if (curScope != null && curScope == overdueRootScope) { + // we know this scope is the last on the stack and is overdue + // prevent a worker that already read the top from retaining it after removal + curScope.clearReferences(); + closedCount++; + overdueRootScope = null; + curScope = null; + } + if (closedCount > 0) { + publishTop(curScope); + } + } + // Close callbacks carry no scope identity and use the same manager, so only a count is needed. + for (int i = 0; i < closedCount; i++) { + closedScope.onProperClose(); + } + if (activateTop && closedCount > 0) { + activateCurrentTop(); + } + // A close listener may have opened another scope while notifications were running. + if (liveTop == null) { + onBecomeEmpty(); + } + } + + @Override + void push(ContinuableScope scope) { + scope.markShared(); + boolean wasEmpty; + synchronized (this) { + wasEmpty = top == null; + if (!wasEmpty) { + stack.push(top); + } + publishTop(scope); + } + if (wasEmpty) { + onBecomeNonEmpty(); + } + activateCurrentTop(); + } + + @Override + void activateCurrentTop() { + ContinuableScope current = liveTop; + if (current != null) { + current.beforeActivated(); + // Profiler hooks can block or reenter; use the latest top for scope listeners. + current = liveTop; + if (current != null) { + current.afterActivated(); + } + } + } + + @Override + boolean checkTop(ContinuableScope expectedScope) { + // only a fast check; closing re-checks the stack under the lock during cleanup + return expectedScope.equals(liveTop); + } + + @Override + boolean checkOverdueScopes(ContinuableScope expectedScope) { + List overdueSpans; + boolean found = false; + synchronized (this) { + if (top == null || top.source() != ITERATION) { + return false; + } + // Preserve ScopeStack's iteration checks, but defer finishing until after releasing the lock. + overdueSpans = new ArrayList<>(); + top.clearReferences(); + AgentSpan span = top.span(); + if (span != null) { + overdueSpans.add(span); + } + for (ContinuableScope scope : stack) { + if (scope.source() != ITERATION) { + found = expectedScope.equals(scope); + break; + } + scope.clearReferences(); + span = scope.span(); + if (span != null) { + overdueSpans.add(span); + } + } + } + for (AgentSpan span : overdueSpans) { + span.finishWithEndToEnd(); + } + return found; + } + + @Override + int depth() { + // only used to enforce the depth limit, so a slightly stale value is fine + return liveDepth; + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java new file mode 100644 index 00000000000..d70d1fbe877 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ThreadScopeContext.java @@ -0,0 +1,33 @@ +package datadog.trace.core.scopemanager; + +import datadog.context.Context; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; + +/** + * Swaps a thread's own stack back in after a coroutine's shared stack ran on that thread. Unlike + * {@link SavedScopeContext} it takes no snapshot, so one instance can be reused for every + * coroutine. + */ +final class ThreadScopeContext extends ScopeStackContext { + private final Thread originalThread = Thread.currentThread(); + private final ScopeStack scopeStack; + + ThreadScopeContext(ScopeStack scopeStack) { + this.scopeStack = scopeStack; + } + + @Override + ScopeStack restore(ProfilingContextIntegration profilingContextIntegration) { + // coroutines restore on the same thread; there's no snapshot to copy for any other thread + if (Thread.currentThread() == originalThread) { + return scopeStack; + } else { + return new ScopeStack(profilingContextIntegration); + } + } + + @Override + Context currentContext() { + return contextOf(scopeStack.top()); + } +} diff --git a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java new file mode 100644 index 00000000000..7f6ffc679c1 --- /dev/null +++ b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/SharedScopeStackTest.java @@ -0,0 +1,1151 @@ +package datadog.trace.core.scopemanager; + +import static java.util.Arrays.asList; +import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; + +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.context.ContextScope; +import datadog.trace.api.DDTraceId; +import datadog.trace.api.Stateful; +import datadog.trace.api.interceptor.MutableSpan; +import datadog.trace.api.interceptor.TraceInterceptor; +import datadog.trace.api.scopemanager.ExtendedScopeListener; +import datadog.trace.api.scopemanager.ScopeListener; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; +import datadog.trace.common.writer.ListWriter; +import datadog.trace.core.CoreTracer; +import datadog.trace.core.DDCoreJavaSpecification; +import datadog.trace.core.ScopeManagerTestBridge; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.tabletest.junit.TableTest; + +class SharedScopeStackTest extends DDCoreJavaSpecification { + private static final ContextKey KEY = ContextKey.named("shared-scope-stack-test"); + + private final Context base = Context.root().with(KEY, "base"); + + private CoreTracer tracer; + private ContinuableScopeManager scopeManager; + + @BeforeEach + void setup() { + tracer = tracerBuilder().writer(new ListWriter()).build(); + scopeManager = ScopeManagerTestBridge.getScopeManager(tracer); + } + + @AfterEach + void cleanup() { + tracer.close(); + } + + @Test + void scopeOpenedOnOneThreadIsVisibleAndClosableOnAnother() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + Context opened = base.with(KEY, "opened"); + AtomicReference scope = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + scope.set(scopeManager.attach(opened)); + scopeManager.swap(original); + }); + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(opened, scopeManager.currentContext()); + scope.get().close(); + assertSame(base, scopeManager.currentContext()); + scopeManager.swap(original); + }); + } + + @Test + void rootSeedStartsEmpty() throws Exception { + Context shared = scopeManager.sharedContextStack(Context.root()); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertNull(scopeManager.active()); + assertEquals(0, scopeManager.scopeStack().depth()); + scopeManager.swap(original); + }); + } + + @Test + void scopeClosedElsewhereIsRemovedWhenSwappedIn() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + AtomicReference scope = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + scope.set(scopeManager.attach(base.with(KEY, "closed elsewhere"))); + scopeManager.swap(original); + }); + // closed while the shared stack isn't swapped in, leaving it closed on top of that stack + scope.get().close(); + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(base, scopeManager.currentContext()); + assertEquals(1, scopeManager.scopeStack().depth()); + scopeManager.swap(original); + }); + } + + @Test + void sharedContextReadsTheLiveTopAndDerivesIndependentContexts() throws Exception { + ContextKey extra = ContextKey.named("shared-context-extra"); + Context shared = scopeManager.sharedContextStack(base); + Context before = shared.with(extra, "before"); + AtomicReference opened = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + opened.set(scopeManager.attach(base.with(KEY, "opened"))); + scopeManager.swap(original); + }); + + assertEquals("opened", shared.get(KEY)); + assertNull(shared.get(extra)); + assertEquals("base", before.get(KEY)); + assertEquals("before", before.get(extra)); + Context after = shared.with(extra, "after"); + assertEquals("opened", after.get(KEY)); + assertEquals("after", after.get(extra)); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + opened.get().close(); + scopeManager.swap(original); + }); + + assertEquals("base", shared.get(KEY)); + assertNull(shared.get(extra)); + assertEquals("opened", after.get(KEY)); + } + + @Test + void overdueSharedRootIsRemovedBeforeAnotherWorkerRuns() throws Exception { + Context shared = scopeManager.sharedContextStack(Context.root()); + AgentSpan iteration = tracer.buildSpan("test", "iteration").start(); + AtomicInteger closes = new AtomicInteger(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeClosed() { + closes.incrementAndGet(); + } + }); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + ScopeStack stack = scopeManager.scopeStack(); + ContinuableScope root = + new ContinuableScope( + scopeManager, iteration, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + stack.push(root); + assertSame(iteration, scopeManager.activeSpan()); + // simulate the cleaner marking an expired root, without scheduling its background timer + stack.overdueRootScope = root; + iteration.finishWithEndToEnd(); + assertNull(scopeManager.active()); + stack.cleanup(); + assertNull(stack.overdueRootScope); + scopeManager.swap(original); + }); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertNull(scopeManager.scopeStack().top()); + assertNull(scopeManager.active()); + assertEquals(0, scopeManager.scopeStack().depth()); + assertSame(Context.root(), scopeManager.currentContext()); + scopeManager.swap(original); + }); + assertEquals(1, closes.get()); + } + + @Test + void overdueRootCannotBeRetainedAfterRemoval() throws Exception { + SharedScopeStack stack = new SharedScopeStack(ProfilingContextIntegration.NoOp.INSTANCE); + CountDownLatch topRead = new CountDownLatch(1); + CountDownLatch rootRemoved = new CountDownLatch(1); + AtomicBoolean pause = new AtomicBoolean(true); + Context context = + new Context() { + @Override + public T get(ContextKey key) { + return base.get(key); + } + + @Override + public Context with(ContextKey key, T value) { + return base.with(key, value); + } + + @Override + public boolean equals(Object other) { + // Activation has read the top but has not yet tried to retain it. + if (pause.compareAndSet(true, false)) { + topRead.countDown(); + try { + assertTrue(rootRemoved.await(10, SECONDS)); + } catch (InterruptedException e) { + throw new AssertionError(e); + } + } + return this == other; + } + + @Override + public int hashCode() { + return System.identityHashCode(this); + } + }; + ContinuableScope root = + new ContinuableScope( + scopeManager, context, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + stack.seed(root); + + Thread worker = + new Thread( + () -> { + Context original = scopeManager.swap(stack.handle); + try (ContextScope attached = scopeManager.attach(context)) { + assertNotSame(root, attached); + assertSame(context, scopeManager.currentContext()); + } finally { + scopeManager.swap(original); + } + }, + "activating-worker"); + Thread cleaner = + new Thread( + () -> { + try { + assertTrue(topRead.await(10, SECONDS)); + stack.overdueRootScope = root; + stack.cleanup(); + } catch (InterruptedException e) { + throw new AssertionError(e); + } finally { + rootRemoved.countDown(); + } + }, + "overdue-root-cleaner"); + + runThreads(worker, cleaner); + assertFalse(root.alive()); + } + + @TableTest({ + "Scenario | Parent present", + "Active only | false ", + "With parent | true " + }) + void savedContextRestoresAnIndependentStackOnAnotherThread(boolean parentPresent) + throws Exception { + ContextScope parent = parentPresent ? scopeManager.attach(base) : null; + Context activeContext = base.with(KEY, "active"); + ContextScope active = scopeManager.attach(activeContext); + ScopeStack source = scopeManager.scopeStack(); + Context saved = scopeManager.swap(Context.root()); + + try { + runOnNewThread( + () -> { + Context original = scopeManager.swap(saved); + try { + ScopeStack restored = scopeManager.scopeStack(); + assertNotSame(source, restored); + assertSame(active, restored.top()); + assertSame(parent, restored.parent()); + assertEquals(parentPresent ? 2 : 1, restored.depth()); + assertSame(activeContext, scopeManager.currentContext()); + try (ContextScope child = scopeManager.attach(base.with(KEY, "worker"))) { + assertEquals("worker", scopeManager.currentContext().get(KEY)); + assertSame(active, source.top()); + assertEquals(parentPresent ? 2 : 1, source.depth()); + } + assertSame(activeContext, scopeManager.currentContext()); + } finally { + scopeManager.swap(original); + } + }); + } finally { + scopeManager.swap(saved); + active.close(); + if (parent != null) { + parent.close(); + } + } + } + + @Test + void threadsCanUseTheSharedStackAtTheSameTime() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + int iterations = 10_000; + CyclicBarrier start = new CyclicBarrier(2); + Runnable attachAndClose = + () -> { + Context original = scopeManager.swap(shared); + try { + start.await(10, SECONDS); + for (int i = 0; i < iterations; i++) { + scopeManager.attach(base.with(KEY, Thread.currentThread().getName() + i)).close(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + scopeManager.swap(original); + } + }; + Thread first = new Thread(attachAndClose, "first"); + Thread second = new Thread(attachAndClose, "second"); + runThreads(first, second); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(base, scopeManager.currentContext()); + assertEquals(1, scopeManager.scopeStack().depth()); + scopeManager.swap(original); + }); + } + + @Test + void threadsCanReuseTheSameScopeAtTheSameTime() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + int iterations = 100_000; + CyclicBarrier start = new CyclicBarrier(2); + Runnable reattachBase = + () -> { + Context original = scopeManager.swap(shared); + try { + start.await(10, SECONDS); + for (int i = 0; i < iterations; i++) { + // reuses the seed scope, adding and removing a reference each time + scopeManager.attach(base).close(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + scopeManager.swap(original); + } + }; + Thread first = new Thread(reattachBase, "first"); + Thread second = new Thread(reattachBase, "second"); + runThreads(first, second); + + assertSeedHoldsOnlyItsOriginalReference(shared); + } + + @Test + void scopeCanBeClosedFromAnotherStackWhileItsReused() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + int references = 20_000; + List extraReferences = new ArrayList<>(); + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + for (int i = 0; i < references; i++) { + extraReferences.add(scopeManager.attach(base)); + } + scopeManager.swap(original); + }); + CyclicBarrier start = new CyclicBarrier(2); + Thread reuser = + new Thread( + () -> { + Context original = scopeManager.swap(shared); + try { + start.await(10, SECONDS); + for (int i = 0; i < references; i++) { + scopeManager.attach(base).close(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } finally { + scopeManager.swap(original); + } + }); + // closes the extra references on its own stack, while the other thread reuses the same scope + Thread closer = + new Thread( + () -> { + try { + start.await(10, SECONDS); + } catch (Exception e) { + throw new RuntimeException(e); + } + extraReferences.forEach(ContextScope::close); + }); + runThreads(reuser, closer); + + assertSeedHoldsOnlyItsOriginalReference(shared); + } + + @Test + void clearedIterationScopeIsNotReusedByAnotherThread() throws Exception { + int rounds = 2_000; + CyclicBarrier phase = new CyclicBarrier(2); + AtomicReference shared = new AtomicReference<>(); + AtomicReference iterationSpan = new AtomicReference<>(); + AtomicReference iterationScope = new AtomicReference<>(); + Runnable reuser = + () -> { + try { + for (int round = 0; round < rounds; round++) { + phase.await(10, SECONDS); // iteration scope is on top + Context original = scopeManager.swap(shared.get()); + for (int i = 0; i < 200; i++) { + // reuses the iteration scope until it's cleared, then pushes a new scope instead + scopeManager.activateSpan(iterationSpan.get()).close(); + } + scopeManager.swap(original); + phase.await(10, SECONDS); // round finished + } + } catch (Exception e) { + throw new RuntimeException(e); + } + }; + Runnable closer = + () -> { + try { + for (int round = 0; round < rounds; round++) { + shared.set(scopeManager.sharedContextStack(base)); + Context original = scopeManager.swap(shared.get()); + ContextScope outer = scopeManager.attach(base.with(KEY, "outer")); + iterationSpan.set(tracer.buildSpan("test", "iteration").start()); + iterationScope.set((ContinuableScope) scopeManager.activateNext(iterationSpan.get())); + phase.await(10, SECONDS); + // closing the outer scope clears the overdue iteration scope above it + outer.close(); + scopeManager.swap(original); + phase.await(10, SECONDS); + assertFalse(iterationScope.get().alive(), "cleared scope retained in round " + round); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + }; + Thread first = new Thread(reuser, "reuser"); + Thread second = new Thread(closer, "closer"); + runThreads(first, second); + } + + @Test + void threadGetsItsOwnStackBackAfterNestedSwaps() { + Context outer = scopeManager.sharedContextStack(base); + Context inner = scopeManager.sharedContextStack(base.with(KEY, "inner")); + ScopeStack threadStack = scopeManager.scopeStack(); + + Context original = scopeManager.swap(outer); + ScopeStack outerStack = scopeManager.scopeStack(); + Context outerToken = scopeManager.swap(inner); + scopeManager.swap(outerToken); + assertSame(outerStack, scopeManager.scopeStack()); + scopeManager.swap(original); + + assertSame(threadStack, scopeManager.scopeStack()); + } + + @Test + void threadReusesOneTokenForItsOwnStack() throws Exception { + Context shared = scopeManager.sharedContextStack(base); + Context opened = base.with(KEY, "thread"); + AtomicReference token = new AtomicReference<>(); + + runOnNewThread( + () -> { + ContextScope scope = scopeManager.attach(opened); + Context first = scopeManager.swap(shared); + scopeManager.swap(first); + Context second = scopeManager.swap(shared); + assertSame(first, second); + scopeManager.swap(second); + assertSame(opened, scopeManager.currentContext()); + token.set(second); + scope.close(); + }); + runOnNewThread( + () -> { + // the token only restores its own thread's stack; other threads get an empty stack + scopeManager.swap(token.get()); + assertNull(scopeManager.active()); + }); + } + + @Test + void reusedThreadTokenReadsItsOwnLiveStack() { + ContextKey extra = ContextKey.named("thread-token-extra"); + Context shared = scopeManager.sharedContextStack(base.with(KEY, "coroutine")); + ContextScope parentScope = scopeManager.attach(base); + Context token = scopeManager.swap(shared); + + assertEquals("coroutine", scopeManager.currentContext().get(KEY)); + assertEquals("base", token.get(KEY)); + Context captured = token.with(extra, "captured"); + assertEquals("base", captured.get(KEY)); + assertEquals("captured", captured.get(extra)); + assertNull(token.get(extra)); + + scopeManager.swap(token); + ContextScope childScope = scopeManager.attach(base.with(KEY, "child")); + Context reused = scopeManager.swap(shared); + assertSame(token, reused); + assertEquals("child", token.get(KEY)); + assertEquals("base", captured.get(KEY)); + + scopeManager.swap(reused); + childScope.close(); + parentScope.close(); + assertNull(token.get(KEY)); + } + + @Test + void seedingFromSharedContextStartsFromItsCurrentTop() throws Exception { + Context parent = scopeManager.sharedContextStack(base); + Context opened = base.with(KEY, "opened"); + AtomicReference child = new AtomicReference<>(); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(parent); + scopeManager.attach(opened); + child.set(scopeManager.sharedContextStack(parent)); + scopeManager.swap(original); + }); + runOnNewThread( + () -> { + Context original = scopeManager.swap(child.get()); + assertSame(opened, scopeManager.currentContext()); + // the child has its own stack + scopeManager.attach(base.with(KEY, "child")); + scopeManager.swap(original); + scopeManager.swap(parent); + assertSame(opened, scopeManager.currentContext()); + scopeManager.swap(original); + }); + } + + @TableTest({ + "Scenario | Same thread token | Seed present", + "Saved context | false | true ", + "Thread context | true | true ", + "Empty saved token | false | false ", + "Empty thread token | true | false " + }) + void seedingFromRestoreTokenUsesItsContext(boolean sameThreadToken, boolean seedPresent) { + ContextScope parent = seedPresent ? scopeManager.attach(base) : null; + ScopeStack parentStack = scopeManager.scopeStack(); + Context token = + scopeManager.swap( + sameThreadToken ? scopeManager.sharedContextStack(Context.root()) : Context.root()); + try { + Context child = scopeManager.sharedContextStack(token); + Context original = scopeManager.swap(child); + try { + assertSame(seedPresent ? base : Context.root(), scopeManager.currentContext()); + assertNotSame(parentStack, scopeManager.scopeStack()); + if (!seedPresent) { + assertNull(scopeManager.active()); + } + try (ContextScope scope = scopeManager.attach(base.with(KEY, "child"))) { + assertEquals("child", scopeManager.currentContext().get(KEY)); + assertSame(parent, parentStack.top()); + } + } finally { + scopeManager.swap(original); + } + } finally { + scopeManager.swap(token); + if (parent != null) { + parent.close(); + } + } + } + + @Test + void listenersSeeTheScopeBeingActivated() throws Exception { + Context shared = scopeManager.sharedContextStack(Context.root()); + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + List activeSpans = new ArrayList<>(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + activeSpans.add(scopeManager.activeSpan()); + } + + @Override + public void afterScopeClosed() {} + }); + + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + ContextScope parentScope = scopeManager.activateSpan(parent); + scopeManager.activateSpan(child).close(); + parentScope.close(); + scopeManager.swap(original); + }); + + assertEquals(3, activeSpans.size()); + assertSame(parent, activeSpans.get(0)); + assertSame(child, activeSpans.get(1)); + assertSame(parent, activeSpans.get(2), "parent reactivated"); + } + + @TableTest({ + "Scenario | Operation", + "New scope | push ", + "Scope close | cleanup ", + "Deferred close | prune " + }) + void listenerCanWaitForAnotherWorkerUsingTheSameStack(String operation) throws Exception { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + Context shared = scopeManager.sharedContextStack(parent); + CountDownLatch callbackStarted = new CountDownLatch(1); + CountDownLatch peerFinished = new CountDownLatch(1); + AtomicBoolean armed = new AtomicBoolean(); + AtomicBoolean peerFinishedDuringCallback = new AtomicBoolean(); + Runnable waitForPeer = + () -> { + if (armed.compareAndSet(true, false)) { + callbackStarted.countDown(); + try { + peerFinishedDuringCallback.set(peerFinished.await(5, SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + }; + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + if ("push".equals(operation)) { + waitForPeer.run(); + } + } + + @Override + public void afterScopeClosed() { + if (!"push".equals(operation)) { + waitForPeer.run(); + } + } + }); + Thread notifyingWorker = + new Thread( + () -> { + Context original = scopeManager.swap(shared); + try { + if ("push".equals(operation)) { + armed.set(true); + scopeManager.activateSpan(child).close(); + } else { + ContextScope scope = scopeManager.activateSpan(child); + if ("prune".equals(operation)) { + scopeManager.swap(original); + scope.close(); + armed.set(true); + original = scopeManager.swap(shared); + } else { + armed.set(true); + scope.close(); + } + } + assertSame(parent, scopeManager.activeSpan()); + } finally { + scopeManager.swap(original); + } + }, + "notifying-worker"); + Thread peer = + new Thread( + () -> { + try { + assertTrue(callbackStarted.await(10, SECONDS)); + Context original = scopeManager.swap(shared); + try (ContextScope ignored = scopeManager.attach(base.with(KEY, "peer"))) { + assertEquals("peer", scopeManager.currentContext().get(KEY)); + } finally { + scopeManager.swap(original); + } + } catch (InterruptedException e) { + throw new AssertionError(e); + } finally { + peerFinished.countDown(); + } + }, + "peer-worker"); + try { + runThreads(notifyingWorker, peer); + assertTrue(peerFinishedDuringCallback.get(), "The listener must not block peer mutations"); + } finally { + child.finish(); + parent.finish(); + } + } + + @Test + void scopeListenersAndProfilerHooksRunOutsideTheStackLock() { + ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class); + SharedScopeStack shared = new SharedScopeStack(profiling); + List callbacks = new ArrayList<>(); + List callbacksUnderLock = new ArrayList<>(); + Consumer record = + callback -> { + callbacks.add(callback); + if (Thread.holdsLock(shared)) { + callbacksUnderLock.add(callback); + } + }; + doAnswer( + invocation -> { + record.accept("attach"); + return null; + }) + .when(profiling) + .onAttach(); + doAnswer( + invocation -> { + record.accept("detach"); + return null; + }) + .when(profiling) + .onDetach(); + Stateful state = + new Stateful() { + @Override + public void activate(Object context) { + record.accept("profiler activate"); + } + + @Override + public void close() {} + }; + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + record.accept("scope activate"); + } + + @Override + public void afterScopeClosed() { + record.accept("scope close"); + } + }); + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + Context original = scopeManager.swap(shared.handle); + try { + ContinuableScope parentScope = + new ContinuableScope(scopeManager, parent, ContinuableScope.CONTEXT, true, state); + ContinuableScope childScope = + new ContinuableScope(scopeManager, child, ContinuableScope.CONTEXT, true, state); + shared.push(parentScope); + shared.push(childScope); + childScope.release(); + shared.cleanup(); + assertSame(parent, scopeManager.activeSpan()); + parentScope.release(); + shared.pruneClosedTop(); + assertNull(scopeManager.activeSpan()); + assertEquals(emptyList(), callbacksUnderLock); + assertEquals( + asList( + "attach", + "profiler activate", + "scope activate", + "profiler activate", + "scope activate", + "scope close", + "profiler activate", + "scope activate", + "scope close", + "detach"), + callbacks); + } finally { + scopeManager.swap(original); + child.finish(); + parent.finish(); + } + } + + @TableTest({ + "Scenario | Parent present", + "Restore parent | true ", + "Empty stack | false " + }) + void swapListenersUseTheTopAfterTheProfilerHook(boolean parentPresent) throws Exception { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + SharedScopeStack shared = new SharedScopeStack(ProfilingContextIntegration.NoOp.INSTANCE); + CountDownLatch profilerStarted = new CountDownLatch(1); + CountDownLatch peerFinished = new CountDownLatch(1); + AtomicBoolean armed = new AtomicBoolean(); + AtomicBoolean peerFinishedDuringHook = new AtomicBoolean(); + Stateful state = + new Stateful() { + @Override + public void activate(Object context) { + if (armed.compareAndSet(true, false)) { + profilerStarted.countDown(); + try { + peerFinishedDuringHook.set(peerFinished.await(10, SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + } + + @Override + public void close() {} + }; + if (parentPresent) { + shared.seed( + new ContinuableScope( + scopeManager, parent, ContinuableScope.CONTEXT, true, Stateful.DEFAULT)); + } + ContinuableScope childScope = + new ContinuableScope(scopeManager, child, ContinuableScope.CONTEXT, true, state); + shared.push(childScope); + + List activatedSpans = new ArrayList<>(); + Thread notifyingWorker = + new Thread( + () -> { + Context original = scopeManager.swap(shared.handle); + try { + assertTrue(peerFinishedDuringHook.get()); + assertSame(parentPresent ? parent : null, scopeManager.activeSpan()); + assertEquals( + parentPresent ? singletonList(parent.getSpanId()) : emptyList(), + activatedSpans); + } finally { + scopeManager.swap(original); + } + }, + "notifying-worker"); + scopeManager.addScopeListener( + new ExtendedScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeActivated(DDTraceId traceId, long spanId) { + if (Thread.currentThread() == notifyingWorker) { + activatedSpans.add(spanId); + } + } + + @Override + public void afterScopeClosed() {} + }); + Thread peer = + new Thread( + () -> { + try { + assertTrue(profilerStarted.await(10, SECONDS)); + Context original = scopeManager.swap(shared.handle); + try { + childScope.close(); + } finally { + scopeManager.swap(original); + } + } catch (InterruptedException e) { + throw new AssertionError(e); + } finally { + peerFinished.countDown(); + } + }, + "peer-worker"); + armed.set(true); + try { + runThreads(notifyingWorker, peer); + } finally { + child.finish(); + parent.finish(); + } + } + + @Test + void cleanupKeepsScopeOpenedByCloseListener() { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + AgentSpan replacement = tracer.buildSpan("test", "replacement").start(); + Context shared = scopeManager.sharedContextStack(parent); + AtomicReference replacementScope = new AtomicReference<>(); + AtomicBoolean openReplacement = new AtomicBoolean(true); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeClosed() { + if (openReplacement.compareAndSet(true, false)) { + replacementScope.set(scopeManager.activateSpan(replacement)); + } + } + }); + Context original = scopeManager.swap(shared); + try { + scopeManager.activateSpan(child).close(); + assertSame(replacement, scopeManager.activeSpan()); + replacementScope.get().close(); + assertSame(parent, scopeManager.activeSpan()); + } finally { + scopeManager.swap(original); + replacement.finish(); + child.finish(); + parent.finish(); + } + } + + @Test + void cleanupNotifiesOnceForEachRemovedScope() { + AgentSpan parent = tracer.buildSpan("test", "parent").start(); + AgentSpan child = tracer.buildSpan("test", "child").start(); + AgentSpan grandchild = tracer.buildSpan("test", "grandchild").start(); + Context shared = scopeManager.sharedContextStack(parent); + AtomicInteger closes = new AtomicInteger(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() {} + + @Override + public void afterScopeClosed() { + closes.incrementAndGet(); + } + }); + Context original = scopeManager.swap(shared); + try { + ContextScope childScope = scopeManager.activateSpan(child); + ContextScope grandchildScope = scopeManager.activateSpan(grandchild); + childScope.close(); + assertEquals(0, closes.get()); + grandchildScope.close(); + assertEquals(2, closes.get()); + assertSame(parent, scopeManager.activeSpan()); + } finally { + scopeManager.swap(original); + grandchild.finish(); + child.finish(); + parent.finish(); + } + } + + @Test + void finishingOverdueScopesInvokesTraceInterceptorsOutsideTheStackLock() { + SharedScopeStack shared = new SharedScopeStack(ProfilingContextIntegration.NoOp.INSTANCE); + AtomicInteger interceptedTraces = new AtomicInteger(); + AtomicBoolean interceptedUnderLock = new AtomicBoolean(); + tracer.addTraceInterceptor( + new TraceInterceptor() { + @Override + public Collection onTraceComplete( + Collection trace) { + interceptedTraces.incrementAndGet(); + if (Thread.holdsLock(shared)) { + interceptedUnderLock.set(true); + } + return trace; + } + + @Override + public int priority() { + return 0; + } + }); + AgentSpan first = tracer.buildSpan("test", "first-iteration").start(); + AgentSpan second = tracer.buildSpan("test", "second-iteration").start(); + ContinuableScope parent = + new ContinuableScope(scopeManager, base, ContinuableScope.CONTEXT, true, Stateful.DEFAULT); + ContinuableScope firstIteration = + new ContinuableScope( + scopeManager, first, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + ContinuableScope secondIteration = + new ContinuableScope( + scopeManager, second, ContinuableScope.ITERATION, true, Stateful.DEFAULT); + shared.push(parent); + shared.push(firstIteration); + shared.push(secondIteration); + try { + assertTrue(shared.checkOverdueScopes(parent)); + assertFalse(firstIteration.alive()); + assertFalse(secondIteration.alive()); + assertEquals(2, interceptedTraces.get()); + assertFalse(interceptedUnderLock.get()); + shared.cleanup(); + assertSame(parent, shared.top()); + } finally { + second.finish(); + first.finish(); + } + } + + @TableTest({ + "Scenario | Seed parent | Close child", + "Restore parent | true | true ", + "Empty stack | false | true ", + "No closed scopes | true | false " + }) + void deferredCleanupListenersSeeTheRestoredStack(boolean seedParent, boolean closeChild) { + AgentSpan parent = tracer.buildSpan("test", "coroutine-parent").start(); + AgentSpan child = tracer.buildSpan("test", "coroutine-child").start(); + AgentSpan worker = tracer.buildSpan("test", "worker-unrelated").start(); + List activeSpans = new ArrayList<>(); + List closedStacks = new ArrayList<>(); + scopeManager.addScopeListener( + new ScopeListener() { + @Override + public void afterScopeActivated() { + activeSpans.add(scopeManager.activeSpan()); + } + + @Override + public void afterScopeClosed() { + closedStacks.add(scopeManager.scopeStack()); + } + }); + + try { + Context shared = scopeManager.sharedContextStack(seedParent ? parent : Context.root()); + ContextScope childScope; + ScopeStack sharedStack; + Context original = scopeManager.swap(shared); + try { + sharedStack = scopeManager.scopeStack(); + childScope = scopeManager.activateSpan(child); + } finally { + scopeManager.swap(original); + } + + try (ContextScope ignored = scopeManager.activateSpan(worker)) { + if (closeChild) { + // The shared stack retains the closed child until it is restored. + childScope.close(); + } + activeSpans.clear(); + closedStacks.clear(); + + Context workerContext = scopeManager.swap(shared); + try { + AgentSpan expected = closeChild ? (seedParent ? parent : null) : child; + assertSame(expected, scopeManager.activeSpan()); + assertEquals(expected == null ? emptyList() : singletonList(expected), activeSpans); + assertEquals(closeChild ? singletonList(sharedStack) : emptyList(), closedStacks); + } finally { + if (!closeChild) { + childScope.close(); + } + scopeManager.swap(workerContext); + } + } + } finally { + child.finish(); + parent.finish(); + worker.finish(); + } + } + + private void assertSeedHoldsOnlyItsOriginalReference(Context shared) throws Exception { + runOnNewThread( + () -> { + Context original = scopeManager.swap(shared); + assertSame(base, scopeManager.currentContext()); + assertEquals(1, scopeManager.scopeStack().depth()); + ContinuableScope seed = scopeManager.scopeStack().top(); + assertTrue(seed.alive()); + // releasing the original reference closes it, so there are no extra references left + assertFalse(seed.release()); + scopeManager.swap(original); + }); + } + + private static void runOnNewThread(Runnable task) throws Exception { + runThreads(new Thread(task, "worker")); + } + + /** Stops peers on failure and bounds shutdown, including when a worker cannot be interrupted. */ + private static void runThreads(Thread... threads) throws Exception { + AtomicReference failure = new AtomicReference<>(); + for (Thread thread : threads) { + // a monitor deadlock cannot be interrupted; it must not prevent the test JVM exiting + thread.setDaemon(true); + thread.setUncaughtExceptionHandler( + (failed, error) -> { + failure.compareAndSet(null, error); + for (Thread peer : threads) { + if (peer != failed) { + peer.interrupt(); + } + } + }); + } + try { + for (Thread thread : threads) { + thread.start(); + } + for (Thread thread : threads) { + thread.join(10_000); + } + } finally { + for (Thread thread : threads) { + thread.interrupt(); + } + for (Thread thread : threads) { + thread.join(10_000); + } + } + if (failure.get() instanceof Error) { + throw (Error) failure.get(); + } else if (failure.get() != null) { + throw new AssertionError(failure.get()); + } + for (Thread thread : threads) { + assertFalse(thread.isAlive(), thread.getName() + " did not terminate"); + } + } +} 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 71767ea7fae..3a267edb556 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 @@ -414,6 +414,14 @@ default AgentSpan startSpan(@Nonnull SpanPrototype prototype, CharSequence opera @Deprecated ContextContinuation capture(Context context); + + /** + * Legacy scope manager only: returns a context whose scope stack is shared by every thread it + * is swapped onto. If {@code initial} is already shared, the new stack starts from its top. + * Returns a new context when supported, or {@code initial} unchanged when unsupported. + */ + @Deprecated + Context sharedContextStack(Context initial); } public interface SpanBuilder { @@ -691,6 +699,11 @@ public Context swap(Context context) { public ContextContinuation capture(Context context) { return NoopContinuation.INSTANCE; } + + @Override + public Context sharedContextStack(Context initial) { + return initial; + } } public static class NoopAgentTraceCollector implements AgentTraceCollector { From 33d9f9edfb576c2997f8650e91f30f74b810c543 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Tue, 20 Oct 2026 10:34:42 +0100 Subject: [PATCH 2/5] Preserve Kotlin coroutine scopes across suspension and resumption Use a shared legacy scope stack for each coroutine so scopes opened before suspension remain visible when execution resumes on another worker. Seed child elements from the inherited element's current context, then release the seed after capture to avoid retaining the parent's stack. Restore each worker's original stack without rebinding the coroutine's shared stack, including overlapping callbacks and replayed restore tokens. Retain context refresh and swap-based saving for the non-legacy manager, and preserve context saving when the tracer cannot create shared stacks. Keep the existing continuation ownership and local-root tracking. Add regressions for suspension APIs, dispatcher changes, structured child waits, child isolation, nested contexts, cancellation, overlapping context elements, and restore-token replay. Bound worker shutdown and exclude the replay test on coroutines 1.6.1, which lacks the required replay behavior. Co-Authored-By: Claude Opus 5.5 --- .../kotlin-coroutines-1.3/build.gradle | 7 + .../StandardSuspensionHandoffTest.groovy | 43 ++ .../DatadogThreadContextElement.java | 65 +- .../StandardSuspensionHandoffTest.groovy | 4 + ...stractStandardSuspensionHandoffTest.groovy | 111 ++++ ...stractDatadogThreadContextElementTest.java | 31 + .../StandardSuspensionHandoffTests.kt | 612 ++++++++++++++++++ 7 files changed, 862 insertions(+), 11 deletions(-) create mode 100644 dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy create mode 100644 dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/groovy/StandardSuspensionHandoffTest.groovy create mode 100644 dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy create mode 100644 dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle index e0e53c2db59..4cef874e9fc 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle @@ -41,6 +41,13 @@ addTestSuiteForDir('coroutines164Test', 'latestDepTest') addTestSuiteForDir('coroutines18Test', 'latestDepTest') addTestSuite('latestDepTest') +tasks.named('coroutines16Test') { + filter { + // kotlinx only replays undispatched restore tokens from 1.6.2 (kotlinx.coroutines#2930) + excludeTestsMatching 'StandardSuspensionHandoffTest.undispatched coroutine can replay its restoration token' + } +} + dependencies { api project(':dd-java-agent:instrumentation:java:java-concurrent:java-concurrent-1.8') compileOnly libs.kotlin diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy new file mode 100644 index 00000000000..a348f202500 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy @@ -0,0 +1,43 @@ +import datadog.trace.instrumentation.kotlin.coroutines.AbstractStandardSuspensionHandoffTest +import datadog.trace.instrumentation.kotlin.coroutines.StandardSuspensionHandoffTests + +class StandardSuspensionHandoffTest extends AbstractStandardSuspensionHandoffTest { + // withContext uses the context merge hook from coroutines 1.6.1 onward. + def "context merge preserves an explicitly swapped span"() { + when: + new StandardSuspensionHandoffTests().runContextSwapAcrossMerge() + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 3 + spans.swapped.parentId == spans.parent.spanId + spans.child.parentId == spans.swapped.spanId + } + + def "undispatched coroutine can replay its restoration token"() { + when: + new StandardSuspensionHandoffTests().runReplayedRestoration() + + then: + noExceptionThrown() + } + + def "scope-touching thread context elements do not corrupt the handed-off stack"() { + when: + def violations = new StandardSuspensionHandoffTests().runScopeTouchingElements(1000) + + then: + // Overlapping workers can temporarily see each other's element scopes. The coroutine must + // retain its base context, but restore callbacks can observe a different top during overlap. + def expectedDuringOverlap = [ + 'element scope leaked from previous worker', + 'before scope not on top at restore', + 'after scope not on top at restore', + 'before restore did not return to previous context', + 'after restore did not return to previous context' + ] + (violations.keySet() - expectedDuringOverlap).isEmpty() + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java index 99899787e84..4753b821357 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java @@ -3,7 +3,10 @@ import datadog.context.Context; import datadog.context.ContextContinuation; import datadog.trace.api.GenericClassValue; +import datadog.trace.api.InstrumenterConfig; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -29,6 +32,9 @@ public final class DatadogThreadContextElement private static final Object COMPLETED = new Object(); + private static final boolean LEGACY_CONTEXT_MANAGER = + InstrumenterConfig.get().isLegacyContextManagerEnabled(); + /** Coroutines built by coroutineScope, withContext, etc. that their caller waits for. */ private static final ClassValue SCOPED_COROUTINE = GenericClassValue.of( @@ -51,15 +57,20 @@ public final class DatadogThreadContextElement public static CoroutineContext addDatadogElement(CoroutineContext coroutineContext) { DatadogThreadContextElement existing = coroutineContext.get(DATADOG_KEY); if (existing == null) { - return coroutineContext.plus(new DatadogThreadContextElement(null)); + return coroutineContext.plus(new DatadogThreadContextElement(null, null)); } return replaceDatadogElement( - coroutineContext, existing.localRootFor(coroutineContext.get(Job.Key))); + coroutineContext, existing.context, existing.localRootFor(coroutineContext.get(Job.Key))); } private static CoroutineContext replaceDatadogElement( - CoroutineContext coroutineContext, @Nullable AgentSpan parentLocalRootSpan) { - DatadogThreadContextElement fresh = new DatadogThreadContextElement(parentLocalRootSpan); + CoroutineContext coroutineContext, + @Nullable Context seed, + @Nullable AgentSpan parentLocalRootSpan) { + // Only legacy mode inherits a live stack. Non-legacy mode captures Context.current() + // later on during child initialization because the parent's saved context may be stale. + DatadogThreadContextElement fresh = + new DatadogThreadContextElement(LEGACY_CONTEXT_MANAGER ? seed : null, parentLocalRootSpan); if (!Boolean.TRUE.equals(coroutineContext.fold(null, HAS_LATER_HOOK))) { return coroutineContext.plus(fresh); } @@ -76,22 +87,26 @@ public static CoroutineContext forContextMerge( DatadogThreadContextElement caller = original.get(DATADOG_KEY); // Scoped coroutines can borrow the caller's hold even when the requested context replaces Job. // Flow merges may install this element without constructing a coroutine to own it. + // An explicit context swap can make the active context differ from the caller's shared stack. return replaceDatadogElement( - merged, caller == null ? null : caller.localRootFor(original.get(Job.Key))); + merged, + LEGACY_CONTEXT_MANAGER ? Context.current() : null, + caller == null ? null : caller.localRootFor(original.get(Job.Key))); } /** Separates upstream flow state while retaining the caller's active context for downstream. */ public static CoroutineContext forFlowCollection( CoroutineContext upstream, CoroutineContext downstream) { DatadogThreadContextElement caller = downstream.get(DATADOG_KEY); - if (caller != null) { + // legacy callers share a live scope stack, which a snapshot would replace + if (!LEGACY_CONTEXT_MANAGER && caller != null) { // Downstream reinstalls the caller's element without constructing a coroutine. caller.context = Context.current(); } DatadogThreadContextElement element = upstream.get(DATADOG_KEY); if (element == null || element == caller) { // Before 1.7.2, flowOn merges with plus and shares this mutable element with its caller. - return replaceDatadogElement(upstream, null); + return replaceDatadogElement(upstream, null, null); } // Newer releases already received a fresh element from the context merge advice. return upstream; @@ -113,7 +128,17 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen return rebuilt.plus(element.getKey() == DATADOG_KEY ? this : element); } + // only needed until the initial context is captured; cleared so a long-lived child doesn't keep + // its parent's scope stack and trace alive + @Nullable private Context seed; @Nullable private AgentSpan parentLocalRootSpan; + + @SuppressFBWarnings( + value = "IS2_INCONSISTENT_SYNC", + justification = + "Written once before publishing volatile context; updateThreadContext reads context before restoration") + private boolean legacySharedStack; + private volatile Context context; private volatile ContextContinuation continuation; @@ -123,7 +148,9 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen /** The local root of the trace kept open by this coroutine or its parent. */ @Nullable private volatile AgentSpan localRootSpan; - private DatadogThreadContextElement(@Nullable AgentSpan parentLocalRootSpan) { + private DatadogThreadContextElement( + @Nullable Context seed, @Nullable AgentSpan parentLocalRootSpan) { + this.seed = seed; this.parentLocalRootSpan = parentLocalRootSpan; } @@ -156,7 +183,9 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine } if (datadog.context == null) { datadog.captureInitialContext(true); - } else if (datadog.owner != coroutine && datadog.isScopedBy(coroutine)) { + } else if (!LEGACY_CONTEXT_MANAGER + && datadog.owner != coroutine + && datadog.isScopedBy(coroutine)) { // scoped coroutines (coroutineScope, withTimeout, or withContext before 1.6.1) share their // caller's element, so refresh it with the active context datadog.context = Context.current(); @@ -166,7 +195,14 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine private synchronized void captureInitialContext(boolean owned) { if (context == null) { - Context initial = Context.current(); + Context initial = seed != null ? seed : Context.current(); + if (LEGACY_CONTEXT_MANAGER) { + // workers share this coroutine's scope stack, rather than copying it between them + Context shared = AgentTracer.get().sharedContextStack(initial); + // the noop tracer returns initial unchanged, without creating a shared stack + legacySharedStack = shared != initial; + initial = shared; + } // Initialize before the first context installation; later swaps update the stored context. AgentSpan rootSpan = localRootOf(initial); if (rootSpan != null) { @@ -184,6 +220,7 @@ private synchronized void captureInitialContext(boolean owned) { } } context = initial; + seed = null; parentLocalRootSpan = null; } } @@ -232,7 +269,13 @@ public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) { @Override public void restoreThreadContext( @Nonnull CoroutineContext coroutineContext, Context originalContext) { - context = originalContext.swap(); + Context coroutineContextAfter = originalContext.swap(); + // legacy workers share the coroutine's scope stack, so it never needs saving; writing it back + // would usually store that same stack, but after an unbalanced swap on this thread it would + // rebind the coroutine to the wrong stack + if (!legacySharedStack) { + context = coroutineContextAfter; + } } /** Tracing state does not change the logical coroutine context used by flow equality checks. */ diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/groovy/StandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/groovy/StandardSuspensionHandoffTest.groovy new file mode 100644 index 00000000000..0f40bb4a9bd --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/groovy/StandardSuspensionHandoffTest.groovy @@ -0,0 +1,4 @@ +import datadog.trace.instrumentation.kotlin.coroutines.AbstractStandardSuspensionHandoffTest + +class StandardSuspensionHandoffTest extends AbstractStandardSuspensionHandoffTest { +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy new file mode 100644 index 00000000000..55a13fa6cc5 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy @@ -0,0 +1,111 @@ +package datadog.trace.instrumentation.kotlin.coroutines + +import datadog.trace.agent.test.InstrumentationSpecification + +abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { + def "#kind child wait refreshes scope closed inside block #closeInside on #outcome"() { + when: + new StandardSuspensionHandoffTests().runScopedChildWait(kind, closeInside, outcome) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 3 + spans.closed.parentId == spans.parent.spanId + spans['after-child'].parentId == spans.parent.spanId + + where: + kind | closeInside | outcome + 'scope' | false | 'return' + 'scope' | true | 'return' + 'supervisor' | false | 'return' + 'supervisor' | true | 'return' + 'context' | false | 'return' + 'context' | true | 'return' + 'timeout' | false | 'return' + 'timeout' | true | 'return' + 'scope' | true | 'throw' + 'scope' | true | 'cancel' + } + + def "early dispatcher #outcome preserves scope opened afterward #openAfter with delayed restore #delayedRestore"() { + when: + new StandardSuspensionHandoffTests().runEarlyDispatcherChange(openAfter, outcome, delayedRestore) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 2 + spans.child.parentId == spans.manual.spanId + + where: + openAfter | outcome | delayedRestore + false | 'return' | false + true | 'return' | false + false | 'throw' | false + true | 'throw' | false + false | 'cancel' | false + true | 'cancel' | false + true | 'return' | true + } + + def "manual scope survives child completion after nested withContext"() { + when: + new StandardSuspensionHandoffTests().runNestedContextChange() + + then: + noExceptionThrown() + } + + def "manual scope survives suspension after flowOn collection"() { + when: + new StandardSuspensionHandoffTests().runFlowOnCollectionThenSuspend() + + then: + noExceptionThrown() + } + + def "child scope open across suspension does not leak into parent resumption"() { + when: + new StandardSuspensionHandoffTests().runChildScopeIsolation() + + then: + noExceptionThrown() + } + + def "dispatcher switch publishes manual scope before source restoration"() { + when: + new StandardSuspensionHandoffTests().runDispatcherChange() + + then: + noExceptionThrown() + } + + def "standard #kind suspension preserves manual scope with cancellation #cancelled and nesting #nested"() { + setup: + new StandardSuspensionHandoffTests().run(kind, cancelled, nested, { + assert TEST_WRITER.empty : "The finished trace must remain held until coroutine completion" + } as Runnable) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + expect: + TEST_WRITER.size() == 1 + trace.size() == (cancelled ? 2 : 3) + spans.manual.parentId == spans.parent.spanId + cancelled || spans.child.parentId == spans.manual.spanId + + where: + kind | cancelled | nested + 'cancellable' | false | false + 'cancellable' | true | false + 'safe' | false | false + 'yield' | false | false + 'delay' | false | false + 'select' | false | false + 'cancellable' | false | true + } +} diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java index 6247ba507a0..64536171995 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java @@ -8,6 +8,9 @@ import static org.junit.jupiter.api.Assertions.assertNotSame; import static org.junit.jupiter.api.Assertions.assertSame; +import datadog.context.Context; +import datadog.context.ContextKey; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import java.lang.reflect.Field; import java.util.ArrayList; import java.util.List; @@ -25,6 +28,34 @@ /** Context replacement checks shared by the oldest and latest dependency suites. */ public abstract class AbstractDatadogThreadContextElementTest { + @Test + void contextCapturedBeforeTracerRegistrationPreservesChangesAcrossSuspension() { + DatadogThreadContextElement element = + (DatadogThreadContextElement) addDatadogElement(EmptyCoroutineContext.INSTANCE); + AgentTracer.TracerAPI tracer = AgentTracer.get(); + Context original; + try { + AgentTracer.forceRegister(AgentTracer.NOOP_TRACER); + original = element.updateThreadContext(element); + } finally { + AgentTracer.forceRegister(tracer); + } + + ContextKey key = ContextKey.named("startup-coroutine-context"); + try { + Context.current().with(key, "opened before suspension").swap(); + } finally { + element.restoreThreadContext(element, original); + } + + original = element.updateThreadContext(element); + try { + assertEquals("opened before suspension", Context.current().get(key)); + } finally { + element.restoreThreadContext(element, original); + } + } + @Test void completionReleasesOwnerWithoutAllowingAnotherCoroutineToClaimElement() throws Exception { DatadogThreadContextElement element = diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt new file mode 100644 index 00000000000..08c32d72a8a --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt @@ -0,0 +1,612 @@ +package datadog.trace.instrumentation.kotlin.coroutines + +import datadog.context.Context +import datadog.context.ContextKey +import datadog.context.ContextScope +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.get +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CoroutineDispatcher +import kotlinx.coroutines.CoroutineExceptionHandler +import kotlinx.coroutines.CoroutineName +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.Job +import kotlinx.coroutines.ThreadContextElement +import kotlinx.coroutines.asCoroutineDispatcher +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.delay +import kotlinx.coroutines.flow.collect +import kotlinx.coroutines.flow.flow +import kotlinx.coroutines.flow.flowOn +import kotlinx.coroutines.launch +import kotlinx.coroutines.newCoroutineContext +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.selects.select +import kotlinx.coroutines.supervisorScope +import kotlinx.coroutines.suspendCancellableCoroutine +import kotlinx.coroutines.withContext +import kotlinx.coroutines.withTimeout +import kotlinx.coroutines.yield +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit.SECONDS +import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference +import kotlin.coroutines.AbstractCoroutineContextElement +import kotlin.coroutines.Continuation +import kotlin.coroutines.ContinuationInterceptor +import kotlin.coroutines.CoroutineContext +import kotlin.coroutines.resume +import kotlin.coroutines.suspendCoroutine + +private val TOUCH_KEY = ContextKey.named("scope-touching") +private val BASE_KEY = ContextKey.named("scope-touching-base") + +@OptIn(ExperimentalCoroutinesApi::class) +class StandardSuspensionHandoffTests { + fun runContextSwapAcrossMerge() { + runBlocking { + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + val swapped = get().buildSpan("kotlin_coroutine", "swapped").start() + val previous = swapped.swap() + try { + withContext(CoroutineName("nested")) { + check(activeSpan() === swapped) { "Context merge discarded the explicitly swapped span" } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } + check(activeSpan() === swapped) { "Context merge did not restore the swapped span" } + } finally { + previous.swap() + swapped.finish() + } + check(activeSpan() === parent) { "Balanced swap did not restore the parent" } + } finally { + parentScope.close() + parent.finish() + } + } + } + + fun runScopedChildWait(kind: String, closeInside: Boolean, outcome: String) { + val workers = Executors.newFixedThreadPool(2) + val childWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val childDispatcher = childWorkers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val childStarted = CountDownLatch(1) + val expected = when (outcome) { + "return" -> null + "throw" -> IllegalStateException("scoped body failed") + "cancel" -> CancellationException("scoped body cancelled") + else -> error("Unknown completion: $outcome") + } + val before = RestoreGate(restoring, released) + val after = RestoreGate(restoring, released) + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + runBlocking(context) { + suspend fun closeCapturedScope() { + val transient = get().buildSpan("kotlin_coroutine", "closed").start() + val scope = get().activateManualSpan(transient) + try { + suspendCoroutine { it.resume(Unit) } + } finally { + scope.close() + transient.finish() + } + check(activeSpan() === parent) { "Inline completion did not restore parent" } + } + if (!closeInside) closeCapturedScope() + val block: suspend CoroutineScope.() -> Unit = { + if (closeInside) closeCapturedScope() + launch(childDispatcher) { + childStarted.countDown() + check(restoring.await(10, SECONDS)) { "Source did not enter restoration" } + } + check(childStarted.await(10, SECONDS)) { "Child did not start" } + if (expected != null) throw expected + Unit + } + try { + try { + when (kind) { + "scope" -> coroutineScope(block) + "supervisor" -> supervisorScope(block) + "context" -> withContext(CoroutineName("nested"), block) + "timeout" -> withTimeout(10000, block) + else -> error("Unknown scoped builder: $kind") + } + check(expected == null) { "Scoped failure was swallowed" } + } catch (failure: Exception) { + check(expected != null && (failure === expected || failure.cause === expected)) { + "Unexpected scoped failure: $failure" + } + } + check(restoring.count == 0L && released.count == 1L) { "Source restoration did not overlap resumption" } + check(before.updates.get() >= 2 && after.updates.get() >= 2) { "Parent did not resume" } + check(activeSpan() === parent) { "Child wait resumed ${activeSpan()?.operationName} instead of parent" } + get().buildSpan("kotlin_coroutine", "after-child").start().finish() + } finally { + released.countDown() + } + } + } finally { + released.countDown() + restoring.countDown() + parentScope.close() + parent.finish() + dispatcher.close() + childDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(childWorkers.awaitTermination(10, SECONDS)) + } + } + + fun runReplayedRestoration() { + val interceptor = object : AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor { + override fun interceptContinuation(continuation: Continuation): Continuation = continuation + } + var saved: Continuation? = null + var failure: Throwable? = null + var completed = false + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + val handler = CoroutineExceptionHandler { _, error -> failure = error } + try { + val job = CoroutineScope(interceptor + handler).launch { + withContext(CoroutineName("nested")) { + suspendCoroutine { saved = it } + } + check(activeSpan() === parent) { "Caller continuation lost its context" } + completed = true + } + check(saved != null) { "The test did not suspend" } + // Resume without a dispatcher installing context. Kotlin must replay its saved restore token. + val original = Context.root().swap() + try { + saved!!.resume(Unit) + check(failure == null) { "Coroutine failed: $failure" } + check(completed && job.isCompleted) { "Caller continuation did not complete" } + check(activeSpan() === parent) { "Replayed restoration lost the caller thread context" } + } finally { + original.swap() + } + } finally { + parentScope.close() + parent.finish() + } + } + + fun runEarlyDispatcherChange(openAfter: Boolean, outcome: String, delayedRestore: Boolean) { + val sourceWorkers = Executors.newSingleThreadExecutor() + val destinationWorkers = Executors.newSingleThreadExecutor() + val source = sourceWorkers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val destination = object : CoroutineDispatcher() { + override fun dispatch(context: CoroutineContext, block: Runnable) { + // Run the real destination task, including restoration, before source getResult can execute. + val task = destinationWorkers.submit(block) + if (delayedRestore) { + check(restoring.await(10, SECONDS)) { "Destination did not reach restoration" } + } else { + task.get(10, SECONDS) + } + } + } + val context: CoroutineContext = if (delayedRestore) { + // Gate both restoration orders: Kotlin 1.3 restores forwards, newer versions backwards. + val before = DestinationRestoreGate(destination, restoring, released) + val after = DestinationRestoreGate(destination, restoring, released) + CoroutineScope(source).newCoroutineContext(before) + after + } else { + source + } + val expected = when (outcome) { + "return" -> null + "throw" -> IllegalStateException("destination failure") + "cancel" -> CancellationException("destination cancellation") + else -> error("Unknown completion: $outcome") + } + try { + runBlocking(context) { + var span = if (openAfter) null else get().buildSpan("kotlin_coroutine", "manual").start() + var scope = span?.let { get().activateManualSpan(it) } + try { + try { + withContext(destination) { + if (!openAfter) { + check(activeSpan() === span) { "Destination lost the caller scope before getResult" } + } + if (expected != null) throw expected + } + check(expected == null) { "Destination failure was swallowed" } + } catch (failure: Exception) { + check(failure === expected) { "Unexpected destination failure: $failure" } + } + if (openAfter) { + span = get().buildSpan("kotlin_coroutine", "manual").start() + scope = get().activateManualSpan(span) + } + check(activeSpan() === span) { "Caller scope was lost on synchronous completion" } + suspendCancellableCoroutine { continuation -> + // This queue entry cannot run until source restoration returns its worker. + sourceWorkers.execute { continuation.resume(Unit) } + } + check(activeSpan() === span) { "Caller scope was lost after the next suspension" } + if (delayedRestore) { + check(restoring.count == 0L && released.count == 1L) { "Destination cleanup was not delayed" } + } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } finally { + scope?.close() + span?.finish() + released.countDown() + } + } + } finally { + released.countDown() + source.close() + destinationWorkers.shutdown() + check(sourceWorkers.awaitTermination(10, SECONDS)) + check(destinationWorkers.awaitTermination(10, SECONDS)) + } + } + + fun runFlowOnCollectionThenSuspend() { + val workers = Executors.newFixedThreadPool(2) + val dispatcher = workers.asCoroutineDispatcher() + try { + runBlocking(dispatcher) { + // Keeping the dispatcher makes flowOn collect upstream undispatched on this coroutine. + flow { emit(1) }.flowOn(CoroutineName("producer")).collect { } + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + repeat(10) { delay(1) } + check(activeSpan() === manual) { "Suspension after flowOn lost the manual scope" } + } finally { + scope.close() + manual.finish() + } + } + } finally { + dispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + } + } + + fun runNestedContextChange() { + val workers = Executors.newSingleThreadExecutor() + val otherWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val otherDispatcher = otherWorkers.asCoroutineDispatcher() + val childRelease = CountDownLatch(1) + try { + runBlocking(dispatcher) { + withContext(CoroutineName("nested")) { Unit } + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + coroutineScope { + launch(otherDispatcher, start = CoroutineStart.ATOMIC) { + check(childRelease.await(10, SECONDS)) + } + // The child can finish only after this coroutine returns its worker. + workers.execute { childRelease.countDown() } + } + check(activeSpan() === manual) { "Waiting for the child lost the manual scope" } + } finally { + scope.close() + manual.finish() + } + } + } finally { + childRelease.countDown() + dispatcher.close() + otherDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(otherWorkers.awaitTermination(10, SECONDS)) + } + } + + fun runChildScopeIsolation() { + val workers = Executors.newSingleThreadExecutor() + val childWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val childDispatcher = childWorkers.asCoroutineDispatcher() + val parentSuspended = CountDownLatch(1) + val childContinuation = AtomicReference>() + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + runBlocking(dispatcher + SuspendGate(parentSuspended)) { + try { + suspendCancellableCoroutine { parentContinuation -> + // The child inherits the parent's Datadog element; override the gate so only the parent signals. + launch(childDispatcher + SuspendGate(CountDownLatch(0))) { + val child = get().buildSpan("kotlin_coroutine", "child").start() + val childScope = get().activateManualSpan(child) + try { + // Suspend with the child scope still open. + suspendCancellableCoroutine { continuation -> + childContinuation.set(continuation) + // Queued behind this task, so it runs only after the child has restored its worker. + childWorkers.execute { + check(parentSuspended.await(10, SECONDS)) { "Parent did not suspend" } + parentContinuation.resume(Unit) + } + } + } finally { + childScope.close() + child.finish() + } + } + } + check(activeSpan() === parent) { "Parent resumed with ${activeSpan()?.operationName} instead of parent" } + } finally { + childContinuation.get()?.resume(Unit) + } + } + } finally { + parentScope.close() + parent.finish() + dispatcher.close() + childDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(childWorkers.awaitTermination(10, SECONDS)) + } + } + + /** + * Runs a coroutine between two elements that attach and close Datadog scopes in their own + * update/restore. With newer kotlinx versions the element after Datadog closes its scope on the + * coroutine stack after the suspension has been published, racing with the next worker reading + * that stack. Returns counts of each kind of violation seen, empty if none. + */ + fun runScopeTouchingElements(iterations: Int): Map { + val workers = Executors.newFixedThreadPool(4) + val resumers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val violations = ConcurrentHashMap() + val record: (String) -> Unit = { kind -> violations.computeIfAbsent(kind) { AtomicInteger() }.incrementAndGet() } + val before = ScopeTouchingElement("before", record) + val after = ScopeTouchingElement("after", record) + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + val base = Context.root().with(BASE_KEY, "base") + val baseScope = base.attach() + try { + runBlocking(context) { + repeat(iterations) { + // Suspend with no scope open in the body, so the element's restore pops the coroutine + // stack itself. Resume straight away so the next worker races that restore. + suspendCancellableCoroutine { continuation -> + resumers.execute { continuation.resume(Unit) } + } + // What Datadog installed for this activation, as seen by the element updated after it. + val installed = after.below.get() + when { + installed === base -> Unit + installed.get(TOUCH_KEY) != null -> record("element scope leaked from previous worker") + installed.get(BASE_KEY) == null -> record("coroutine base context lost") + else -> record("unexpected coroutine context") + } + } + } + } catch (failure: Throwable) { + record("coroutine failed: ${failure.javaClass.name}") + } finally { + baseScope.close() + dispatcher.close() + resumers.shutdown() + check(workers.awaitTermination(10, SECONDS)) + check(resumers.awaitTermination(10, SECONDS)) + } + return violations.mapValues { it.value.get() } + } + + fun runDispatcherChange() { + val workers = Executors.newSingleThreadExecutor() + val otherWorkers = Executors.newSingleThreadExecutor() + val dispatcher = workers.asCoroutineDispatcher() + val otherDispatcher = otherWorkers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val before = RestoreGate(restoring, released) + val after = RestoreGate(restoring, released) + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + // withContext dispatches before getResult. Delay the destination until getResult has run, + // while the source's Datadog restoration is still blocked by one of the gates. + val destinationGate = otherWorkers.submit { check(restoring.await(10, SECONDS)) } + try { + runBlocking(context) { + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + withContext(otherDispatcher) { + try { + check(restoring.count == 0L) { "The source did not suspend" } + check(released.count == 1L) { "The source restoration was not delayed" } + check(activeSpan() === manual) { "Dispatcher switch lost the manual scope" } + } finally { + released.countDown() + } + } + } finally { + scope.close() + manual.finish() + } + } + destinationGate.get(10, SECONDS) + check(before.updates.get() >= 2) + check(after.updates.get() >= 2) + } finally { + released.countDown() + restoring.countDown() + dispatcher.close() + otherDispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + check(otherWorkers.awaitTermination(10, SECONDS)) + } + } + + fun run(kind: String, cancelled: Boolean, nested: Boolean, beforeCompletion: Runnable) { + val workers = Executors.newFixedThreadPool(2) + val dispatcher = workers.asCoroutineDispatcher() + val restoring = CountDownLatch(1) + val released = CountDownLatch(1) + val continuation = AtomicReference>() + val before = RestoreGate(restoring, released) + val after = RestoreGate(restoring, released) + // Place a gate on either side of the auto-installed Datadog element. Kotlin 1.3 restores + // elements forwards and 1.6 restores backwards, so either version delays Datadog restoration. + val context = CoroutineScope(dispatcher).newCoroutineContext(before) + after + val resumer = Thread { + check(restoring.await(10, SECONDS)) + continuation.get()?.let { + if (cancelled) it.context[Job]?.cancel() + it.resume(Unit) + } + } + resumer.start() + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + try { + runBlocking(context) { + val suspendWithManualScope: suspend () -> Unit = { + parent.finish() + val manual = get().buildSpan("kotlin_coroutine", "manual").start() + val scope = get().activateManualSpan(manual) + try { + when (kind) { + "cancellable" -> suspendCancellableCoroutine { continuation.set(it) } + "safe" -> suspendCoroutine { continuation.set(it) } + "yield" -> yield() + "delay" -> delay(20) + "select" -> select { onTimeout(20) { Unit } } + else -> error("Unknown suspension API: $kind") + } + check(activeSpan() === manual) { "Manual scope was lost across $kind" } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } finally { + try { + try { + check(activeSpan() === manual) { "Cleanup lost the manual scope across $kind" } + } finally { + try { + scope.close() + } finally { + manual.finish() + } + } + resumer.join(10000) + check(!resumer.isAlive) + beforeCompletion.run() + } finally { + released.countDown() + } + } + } + if (nested) { + withContext(CoroutineName("nested")) { suspendWithManualScope() } + } else { + suspendWithManualScope() + } + } + check(!cancelled) + } catch (expected: CancellationException) { + check(cancelled) + } + check(before.updates.get() >= 2) { "The test did not suspend" } + check(after.updates.get() >= 2) { "The test did not suspend" } + } finally { + released.countDown() + restoring.countDown() + resumer.join(10000) + parentScope.close() + dispatcher.close() + check(workers.awaitTermination(10, SECONDS)) + } + } + + private class Touch(val scope: ContextScope, val below: Context, val token: Any) + + /** Attaches a Datadog scope in update and closes it in restore, checking it was on top. */ + private class ScopeTouchingElement( + private val name: String, + private val record: (String) -> Unit + ) : AbstractCoroutineContextElement(object : CoroutineContext.Key {}), + ThreadContextElement { + val below = ThreadLocal() + + override fun updateThreadContext(context: CoroutineContext): Touch { + val current = Context.current() + below.set(current) + val token = Any() + return Touch(current.with(TOUCH_KEY, token).attach(), current, token) + } + + override fun restoreThreadContext(context: CoroutineContext, oldState: Touch) { + if (Context.current().get(TOUCH_KEY) !== oldState.token) { + record("$name scope not on top at restore") + } + oldState.scope.close() + if (Context.current() !== oldState.below) { + record("$name restore did not return to previous context") + } + } + } + + private class SuspendGate( + private val suspended: CountDownLatch + ) : AbstractCoroutineContextElement(SuspendGate), + ThreadContextElement { + companion object Key : CoroutineContext.Key + + override fun updateThreadContext(context: CoroutineContext) = Unit + + override fun restoreThreadContext(context: CoroutineContext, oldState: Unit) { + suspended.countDown() + } + } + + private class DestinationRestoreGate( + private val destination: CoroutineDispatcher, + private val restoring: CountDownLatch, + private val released: CountDownLatch + ) : AbstractCoroutineContextElement(object : CoroutineContext.Key {}), + ThreadContextElement { + override fun updateThreadContext(context: CoroutineContext): Boolean = context[ContinuationInterceptor] === destination + + override fun restoreThreadContext(context: CoroutineContext, oldState: Boolean) { + if (oldState) { + restoring.countDown() + check(released.await(10, SECONDS)) { "Caller did not finish before destination restoration" } + } + } + } + + private class RestoreGate( + private val restoring: CountDownLatch, + private val released: CountDownLatch + ) : AbstractCoroutineContextElement(object : CoroutineContext.Key {}), + ThreadContextElement { + val updates = AtomicInteger() + override fun updateThreadContext(context: CoroutineContext): Boolean = updates.incrementAndGet() == 1 + override fun restoreThreadContext(context: CoroutineContext, oldState: Boolean) { + if (oldState) { + restoring.countDown() + check(released.await(10, SECONDS)) { "The new worker did not finish before the old restoration" } + } + } + } +} From 448ded4509d83eca4eb07d76ed462f6fa645513a Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Tue, 20 Oct 2026 23:26:35 +0100 Subject: [PATCH 3/5] Fix recursive context lookups when reusing scope stack tokens --- .../core/scopemanager/ScopeStackContext.java | 35 ++++- .../scopemanager/ScopeManagerForkedTest.java | 132 ++++++++++++++++++ 2 files changed, 166 insertions(+), 1 deletion(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java index a285bb1cf66..9e3ffd4c910 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStackContext.java @@ -1,11 +1,19 @@ package datadog.trace.core.scopemanager; import datadog.context.Context; +import datadog.context.ContextContinuation; import datadog.context.ContextKey; +import datadog.context.ContextScope; +import datadog.context.ImplicitContextKeyed; import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration; import javax.annotation.Nullable; -/** A {@link Context} that swaps a {@link ScopeStack} back in. */ +/** + * A {@link Context} that swaps a {@link ScopeStack} back in. + * + *

Value access, updates and propagation use the underlying context, keeping live stack tokens + * out of scopes, continuations and carriers. + */ abstract class ScopeStackContext implements Context { /** Returns the stack to install on the current thread. */ @@ -25,6 +33,31 @@ public final Context with(ContextKey key, @Nullable T value) { return currentContext().with(key, value); } + @Override + public final Context with(@Nullable ImplicitContextKeyed value) { + return currentContext().with(value); + } + + @Override + public final ContextScope attach() { + return currentContext().attach(); + } + + @Override + public final void attachTo(Object carrier) { + currentContext().attachTo(carrier); + } + + @Override + public final ContextContinuation capture() { + return currentContext().capture(); + } + + @Override + public final ContextScope asScope() { + return currentContext().asScope(); + } + static Context contextOf(@Nullable ContinuableScope scope) { return scope != null ? scope.context : Context.root(); } diff --git a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java index 4a1eda54ac2..442b1ace94a 100644 --- a/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java +++ b/dd-trace-core/src/test/java/datadog/trace/core/scopemanager/ScopeManagerForkedTest.java @@ -27,6 +27,7 @@ import datadog.context.ContextContinuation; import datadog.context.ContextKey; import datadog.context.ContextScope; +import datadog.context.ImplicitContextKeyed; import datadog.trace.api.DDTraceId; import datadog.trace.api.Stateful; import datadog.trace.api.interceptor.MutableSpan; @@ -112,6 +113,137 @@ void cleanup() { tracer.close(); } + @TableTest({ + "Scenario | Token | Capture", + "Attach shared handle | 'shared' | false ", + "Resume shared handle | 'shared' | true ", + "Attach worker token | 'worker' | false ", + "Resume worker token | 'worker' | true ", + "Attach saved snapshot | 'saved' | false ", + "Resume saved snapshot | 'saved' | true " + }) + void stackTokensCanBeAttachedAndCaptured(String tokenKind, boolean capture) { + ContextKey key = ContextKey.named("stack-token"); + Context base = Context.root().with(key, "base"); + Context original = base.swap(); + Context sharedBase = base.with(key, "coroutine"); + Context token; + boolean sharedToken = "shared".equals(tokenKind); + if ("saved".equals(tokenKind)) { + token = Context.root().swap(); + } else { + Context worker = scopeManager.swap(scopeManager.sharedContextStack(sharedBase)); + Context shared = worker.swap(); + token = sharedToken ? shared : worker; + } + Context expected = sharedToken ? sharedBase : base; + token.swap(); + ContextContinuation continuation = capture ? token.capture() : null; + try { + ContextScope activated = capture ? continuation.resume() : token.attach(); + try { + assertSame(expected, Context.current()); + assertEquals(sharedToken ? "coroutine" : "base", Context.current().get(key)); + } finally { + activated.close(); + } + assertSame(expected, Context.current()); + } finally { + if (continuation != null) { + continuation.release(); + } + original.swap(); + } + } + + @Test + void capturingLiveHandleKeepsContextAtCaptureTime() { + ContextKey key = ContextKey.named("live-token-capture"); + Context base = Context.root().with(key, "base"); + Context original = scopeManager.sharedContextStack(base).swap(); + Context shared = Context.root().swap(); + shared.swap(); + ContextContinuation captured = shared.capture(); + Context later = base.with(key, "later"); + try (ContextScope changed = later.attach()) { + try (ContextScope resumed = captured.resume()) { + assertSame(base, Context.current()); + } + assertSame(later, Context.current()); + } finally { + captured.release(); + original.swap(); + } + } + + @ParameterizedTest + @ValueSource(strings = {"shared", "worker", "saved"}) + void tokenDerivedContextsKeepTheirValuesAfterStackChanges(String tokenKind) { + ContextKey key = ContextKey.named("derived-token-context"); + Context base = Context.root().with(key, "base"); + Context original = base.swap(); + Context token; + if ("saved".equals(tokenKind)) { + token = Context.root().swap(); + } else { + Context worker = scopeManager.sharedContextStack(base).swap(); + Context shared = worker.swap(); + token = "shared".equals(tokenKind) ? shared : worker; + } + token.swap(); + Object carrier = new Object(); + try { + token.attachTo(carrier); + ContextScope passive = token.asScope(); + Context unchanged = token.with((ImplicitContextKeyed) null); + AtomicReference supplied = new AtomicReference<>(); + Context stored = + token.with( + context -> { + supplied.set(context); + return context; + }); + Context later = base.with(key, "later"); + try (ContextScope changed = later.attach()) { + assertSame(base, Context.from(carrier)); + assertSame(base, passive.context()); + assertSame(base, unchanged); + assertSame(base, supplied.get()); + assertSame(base, stored); + passive.close(); + assertSame(later, Context.current()); + } + assertSame(base, Context.detachFrom(carrier)); + } finally { + Context.detachFrom(carrier); + original.swap(); + } + } + + @Test + void savedTokenStillAttachesAndCapturesItsSnapshotAfterStackChanges() { + ContextKey key = ContextKey.named("saved-token-snapshot"); + Context snapshot = Context.root().with(key, "snapshot"); + Context original = snapshot.swap(); + Context saved = Context.root().swap(); + saved.swap(); + Context later = snapshot.with(key, "later"); + try (ContextScope changed = later.attach()) { + try (ContextScope attached = saved.attach()) { + assertSame(snapshot, Context.current()); + } + ContextContinuation captured = saved.capture(); + try (ContextScope resumed = captured.resume()) { + assertSame(snapshot, Context.current()); + } finally { + captured.release(); + } + assertSame(later, Context.current()); + } finally { + original.swap(); + } + } + @Test void nonDdspanActivationResultsInAContinuableScope() { ContextScope scope = scopeManager.activateSpan(noopSpan()); From ea1c93e9deb12adf2045ca09a5a89d2fdb4babc6 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 21 Oct 2026 00:12:50 +0100 Subject: [PATCH 4/5] Preserve start-time parenting for lazy Kotlin coroutines --- .../DatadogThreadContextElement.java | 27 ++++++++--- .../LazyCoroutineInstrumentation.java | 2 +- ...stractStandardSuspensionHandoffTest.groovy | 19 ++++++++ .../StandardSuspensionHandoffTests.kt | 48 +++++++++++++++++++ 4 files changed, 88 insertions(+), 8 deletions(-) diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java index 4753b821357..ec1ff75b806 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java @@ -65,12 +65,13 @@ public static CoroutineContext addDatadogElement(CoroutineContext coroutineConte private static CoroutineContext replaceDatadogElement( CoroutineContext coroutineContext, - @Nullable Context seed, + @Nullable Context legacySeed, @Nullable AgentSpan parentLocalRootSpan) { // Only legacy mode inherits a live stack. Non-legacy mode captures Context.current() // later on during child initialization because the parent's saved context may be stale. DatadogThreadContextElement fresh = - new DatadogThreadContextElement(LEGACY_CONTEXT_MANAGER ? seed : null, parentLocalRootSpan); + new DatadogThreadContextElement( + LEGACY_CONTEXT_MANAGER ? legacySeed : null, parentLocalRootSpan); if (!Boolean.TRUE.equals(coroutineContext.fold(null, HAS_LATER_HOOK))) { return coroutineContext.plus(fresh); } @@ -130,7 +131,7 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen // only needed until the initial context is captured; cleared so a long-lived child doesn't keep // its parent's scope stack and trace alive - @Nullable private Context seed; + @Nullable private volatile Context legacySeed; @Nullable private AgentSpan parentLocalRootSpan; @SuppressFBWarnings( @@ -149,8 +150,8 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen @Nullable private volatile AgentSpan localRootSpan; private DatadogThreadContextElement( - @Nullable Context seed, @Nullable AgentSpan parentLocalRootSpan) { - this.seed = seed; + @Nullable Context legacySeed, @Nullable AgentSpan parentLocalRootSpan) { + this.legacySeed = legacySeed; this.parentLocalRootSpan = parentLocalRootSpan; } @@ -174,6 +175,15 @@ public Key getKey() { return DATADOG_KEY; } + /** Lazy children inherit the starting thread's context rather than their creation-time seed. */ + public static void captureLazyDatadogContext(@Nonnull AbstractCoroutine coroutine) { + DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY); + if (datadog != null) { + datadog.legacySeed = null; + } + captureDatadogContext(coroutine); + } + public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine) { DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY); if (datadog != null) { @@ -195,7 +205,10 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine private synchronized void captureInitialContext(boolean owned) { if (context == null) { - Context initial = seed != null ? seed : Context.current(); + Context initial = legacySeed; + if (initial == null) { + initial = Context.current(); + } if (LEGACY_CONTEXT_MANAGER) { // workers share this coroutine's scope stack, rather than copying it between them Context shared = AgentTracer.get().sharedContextStack(initial); @@ -220,7 +233,7 @@ private synchronized void captureInitialContext(boolean owned) { } } context = initial; - seed = null; + legacySeed = null; parentLocalRootSpan = null; } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java index a629a3fdcb4..db42c63fc5d 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/LazyCoroutineInstrumentation.java @@ -35,7 +35,7 @@ public void methodAdvice(MethodTransformer transformer) { public static class OnStartAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void onStart(@Advice.This AbstractCoroutine coroutine) { - DatadogThreadContextElement.captureDatadogContext(coroutine); + DatadogThreadContextElement.captureLazyDatadogContext(coroutine); } } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy index 55a13fa6cc5..24cc45e72c6 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy @@ -3,6 +3,25 @@ package datadog.trace.instrumentation.kotlin.coroutines import datadog.trace.agent.test.InstrumentationSpecification abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { + def "lazy #builder inherits starting context with automatic start #automatic and active span #tracedStarter"() { + when: + new StandardSuspensionHandoffTests().runLazyChildWithInheritedContext(builder, automatic, tracedStarter) + + then: + noExceptionThrown() + + where: + builder | automatic | tracedStarter + 'async' | false | false + 'async' | false | true + 'async' | true | false + 'async' | true | true + 'launch' | false | false + 'launch' | false | true + 'launch' | true | false + 'launch' | true | true + } + def "#kind child wait refreshes scope closed inside block #closeInside on #outcome"() { when: new StandardSuspensionHandoffTests().runScopedChildWait(kind, closeInside, outcome) diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt index 08c32d72a8a..ffa7d68d456 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt @@ -11,10 +11,13 @@ import kotlinx.coroutines.CoroutineExceptionHandler import kotlinx.coroutines.CoroutineName import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.CoroutineStart +import kotlinx.coroutines.Deferred +import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.Job import kotlinx.coroutines.ThreadContextElement import kotlinx.coroutines.asCoroutineDispatcher +import kotlinx.coroutines.async import kotlinx.coroutines.coroutineScope import kotlinx.coroutines.delay import kotlinx.coroutines.flow.collect @@ -47,6 +50,51 @@ private val BASE_KEY = ContextKey.named("scope-touching-base") @OptIn(ExperimentalCoroutinesApi::class) class StandardSuspensionHandoffTests { + fun runLazyChildWithInheritedContext(builder: String, automatic: Boolean, tracedStarter: Boolean) = runBlocking { + val creator = get().buildSpan("kotlin_coroutine", "creator").start() + val creatorScope = get().activateManualSpan(creator) + val starter = if (tracedStarter) get().buildSpan("kotlin_coroutine", "starter").ignoreActiveSpan().start() else null + val worker = Executors.newSingleThreadExecutor() + try { + val block: suspend CoroutineScope.() -> Unit = { + check(activeSpan() === starter) { "Lazy child did not inherit the starting thread's context" } + delay(1) + check(activeSpan() === starter) { "Lazy child lost its starting context after suspension" } + } + val child = when (builder) { + "async" -> async(coroutineContext, CoroutineStart.LAZY, block) + "launch" -> launch(coroutineContext, CoroutineStart.LAZY, block) + else -> error("Unknown lazy builder: $builder") + } + val started = worker.submit { + // Isolate the starter from executor instrumentation's capture of the creator's context. + val original = Context.root().swap() + val scope = starter?.let { get().activateManualSpan(it) } + try { + if (automatic) { + runBlocking { + if (child is Deferred<*>) child.await() else child.join() + } + } else { + child.start() + } + } finally { + scope?.close() + original.swap() + } + } + // Let the creator's dispatcher run its child while the starter waits for it. + withContext(Dispatchers.IO) { started.get(10, SECONDS) } + child.join() + } finally { + worker.shutdownNow() + check(worker.awaitTermination(5, SECONDS)) { "Lazy starter did not terminate" } + creatorScope.close() + creator.finish() + starter?.finish() + } + } + fun runContextSwapAcrossMerge() { runBlocking { val parent = get().buildSpan("kotlin_coroutine", "parent").start() From a80b481a613534f0d60a9b662ee1d998a5fc23fc Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Wed, 21 Oct 2026 10:14:11 +0100 Subject: [PATCH 5/5] Preserve active context when initializing Kotlin coroutine children --- .../DatadogThreadContextElement.java | 10 +++-- ...stractStandardSuspensionHandoffTest.groovy | 20 ++++++++++ ...stractDatadogThreadContextElementTest.java | 18 ++++++++- .../StandardSuspensionHandoffTests.kt | 37 +++++++++++++++++++ 4 files changed, 79 insertions(+), 6 deletions(-) diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java index ec1ff75b806..bb9d34f305d 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java @@ -60,7 +60,9 @@ public static CoroutineContext addDatadogElement(CoroutineContext coroutineConte return coroutineContext.plus(new DatadogThreadContextElement(null, null)); } return replaceDatadogElement( - coroutineContext, existing.context, existing.localRootFor(coroutineContext.get(Job.Key))); + coroutineContext, + LEGACY_CONTEXT_MANAGER ? Context.current() : null, + existing.localRootFor(coroutineContext.get(Job.Key))); } private static CoroutineContext replaceDatadogElement( @@ -99,8 +101,8 @@ public static CoroutineContext forContextMerge( public static CoroutineContext forFlowCollection( CoroutineContext upstream, CoroutineContext downstream) { DatadogThreadContextElement caller = downstream.get(DATADOG_KEY); - // legacy callers share a live scope stack, which a snapshot would replace - if (!LEGACY_CONTEXT_MANAGER && caller != null) { + // Read context before legacySharedStack to acquire the initialized stack state. + if (caller != null && caller.context != null && !caller.legacySharedStack) { // Downstream reinstalls the caller's element without constructing a coroutine. caller.context = Context.current(); } @@ -193,7 +195,7 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine } if (datadog.context == null) { datadog.captureInitialContext(true); - } else if (!LEGACY_CONTEXT_MANAGER + } else if (!datadog.legacySharedStack && datadog.owner != coroutine && datadog.isScopedBy(coroutine)) { // scoped coroutines (coroutineScope, withTimeout, or withContext before 1.6.1) share their diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy index 24cc45e72c6..009680c405d 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy @@ -3,6 +3,26 @@ package datadog.trace.instrumentation.kotlin.coroutines import datadog.trace.agent.test.InstrumentationSpecification abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { + def "#builder with an enclosing context captures active span with swap #swapped"() { + when: + new StandardSuspensionHandoffTests().runChildWithEnclosingContext(builder, swapped) + TEST_WRITER.waitForTraces(1) + def trace = TEST_WRITER.get(0) + def spans = trace.collectEntries { [(it.operationName.toString()): it] } + + then: + trace.size() == 3 + spans.active.parentId == spans.parent.spanId + spans.child.parentId == spans.active.spanId + + where: + builder | swapped + 'async' | false + 'async' | true + 'launch' | false + 'launch' | true + } + def "lazy #builder inherits starting context with automatic start #automatic and active span #tracedStarter"() { when: new StandardSuspensionHandoffTests().runLazyChildWithInheritedContext(builder, automatic, tracedStarter) diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java index 64536171995..05bbed4db97 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractDatadogThreadContextElementTest.java @@ -3,6 +3,7 @@ import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.addDatadogElement; import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.cancelDatadogContext; import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.captureDatadogContext; +import static datadog.trace.instrumentation.kotlin.coroutines.DatadogThreadContextElement.forFlowCollection; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNotSame; @@ -28,8 +29,12 @@ /** Context replacement checks shared by the oldest and latest dependency suites. */ public abstract class AbstractDatadogThreadContextElementTest { - @Test - void contextCapturedBeforeTracerRegistrationPreservesChangesAcrossSuspension() { + @TableTest({ + "Scenario | Flow collection", + "Suspension | false ", + "Flow collection | true " + }) + void contextCapturedBeforeTracerRegistrationPreservesChanges(boolean flowCollection) { DatadogThreadContextElement element = (DatadogThreadContextElement) addDatadogElement(EmptyCoroutineContext.INSTANCE); AgentTracer.TracerAPI tracer = AgentTracer.get(); @@ -44,6 +49,15 @@ void contextCapturedBeforeTracerRegistrationPreservesChangesAcrossSuspension() { ContextKey key = ContextKey.named("startup-coroutine-context"); try { Context.current().with(key, "opened before suspension").swap(); + if (flowCollection) { + forFlowCollection(element, element); + Context beforeCollection = element.updateThreadContext(element); + try { + assertEquals("opened before suspension", Context.current().get(key)); + } finally { + element.restoreThreadContext(element, beforeCollection); + } + } } finally { element.restoreThreadContext(element, original); } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt index ffa7d68d456..d34ec270479 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt @@ -50,6 +50,43 @@ private val BASE_KEY = ContextKey.named("scope-touching-base") @OptIn(ExperimentalCoroutinesApi::class) class StandardSuspensionHandoffTests { + fun runChildWithEnclosingContext(builder: String, swapped: Boolean) = runBlocking { + val enclosingContext = coroutineContext + val parent = get().buildSpan("kotlin_coroutine", "parent").start() + val parentScope = get().activateManualSpan(parent) + try { + withContext(CoroutineName("nested")) { + val active = get().buildSpan("kotlin_coroutine", "active").start() + val previous = if (swapped) active.swap() else null + val scope = if (!swapped) get().activateManualSpan(active) else null + try { + val block: suspend CoroutineScope.() -> Unit = { + check(activeSpan() === active) { "Child inherited the enclosing context's older span" } + delay(1) + check(activeSpan() === active) { "Child lost its active parent after suspension" } + get().buildSpan("kotlin_coroutine", "child").start().finish() + } + val child = try { + when (builder) { + "async" -> async(enclosingContext, block = block) + "launch" -> launch(enclosingContext, block = block) + else -> error("Unknown child builder: $builder") + } + } finally { + scope?.close() + previous?.swap() + } + child.join() + } finally { + active.finish() + } + } + } finally { + parentScope.close() + parent.finish() + } + } + fun runLazyChildWithInheritedContext(builder: String, automatic: Boolean, tracedStarter: Boolean) = runBlocking { val creator = get().buildSpan("kotlin_coroutine", "creator").start() val creatorScope = get().activateManualSpan(creator)