From 408e4b3350e26832f04872430a31cd06b089d469 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 20 Oct 2026 09:38:12 +0200 Subject: [PATCH] Fix Kotlin coroutine context loss during thread handoffs --- .../DatadogThreadContextElement.java | 9 +- ...ThreadContextElementHandoffForkedTest.java | 99 +++++++++++++++++++ .../java/datadog/trace/core/CoreTracer.java | 5 + .../scopemanager/ContinuableScopeManager.java | 27 ++++- .../scopemanager/MovableScopeContext.java | 29 ++++++ .../trace/core/scopemanager/ScopeStack.java | 2 +- .../scopemanager/ScopeManagerForkedTest.java | 34 +++++++ .../instrumentation/api/AgentTracer.java | 9 ++ 8 files changed, 204 insertions(+), 10 deletions(-) create mode 100644 dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementHandoffForkedTest.java create mode 100644 dd-trace-core/src/main/java/datadog/trace/core/scopemanager/MovableScopeContext.java 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 25129df46e3..82e78b62b15 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 @@ -4,6 +4,7 @@ import datadog.context.ContextContinuation; import datadog.trace.api.GenericClassValue; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -99,9 +100,6 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine } if (datadog.context == null) { datadog.captureInitialContext(); - } else if (datadog.owner != coroutine && datadog.isScopedBy(coroutine)) { - // preserve active span changes when a scoped coroutine reinstalls this element - datadog.context = Context.current(); } } } @@ -125,7 +123,7 @@ private synchronized void captureInitialContext() { } } } - context = initial; + context = AgentTracer.get().forkContext(initial); parentLocalRootSpan = null; } } @@ -168,7 +166,8 @@ public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) { @Override public void restoreThreadContext( @Nonnull CoroutineContext coroutineContext, Context originalContext) { - context = originalContext.swap(); + // The next worker may already be executing against the coroutine's stack. + originalContext.swap(); } @Nonnull diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementHandoffForkedTest.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementHandoffForkedTest.java new file mode 100644 index 00000000000..862c7b6ac33 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementHandoffForkedTest.java @@ -0,0 +1,99 @@ +package datadog.trace.instrumentation.kotlin.coroutines; + +import static java.util.concurrent.Executors.newSingleThreadExecutor; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +import datadog.context.Context; +import datadog.context.ContextScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.common.writer.ListWriter; +import datadog.trace.core.CoreTracer; +import java.util.Properties; +import java.util.concurrent.ExecutorService; +import kotlin.coroutines.EmptyCoroutineContext; +import org.junit.jupiter.api.Test; + +class DatadogThreadContextElementHandoffForkedTest { + @Test + void resumptionBeforeRestorationPreservesTheEntireScopeStack() throws Exception { + AgentTracer.installLegacyContextManager(); + Properties properties = new Properties(); + properties.setProperty("dd.trace.scope.strict.mode", "true"); + CoreTracer tracer = + CoreTracer.builder() + .withProperties(properties) + .strictTraceWrites(true) + .writer(new ListWriter()) + .build(); + AgentTracer.forceRegister(tracer); + ExecutorService second = newSingleThreadExecutor(); + ExecutorService third = newSingleThreadExecutor(); + DatadogThreadContextElement element = + (DatadogThreadContextElement) + DatadogThreadContextElement.addDatadogElement(EmptyCoroutineContext.INSTANCE); + Context firstPrevious = element.updateThreadContext(element); + AgentSpan root = tracer.startSpan("test", "root"); + ContextScope rootScope = tracer.activateManualSpan(root); + AgentSpan child = tracer.startSpan("test", "child"); + ContextScope childScope = tracer.activateManualSpan(child); + AgentSpan leaf = tracer.startSpan("test", "leaf"); + ContextScope leafScope = tracer.activateManualSpan(leaf); + try { + // Kotlin permits the next update before the previous restore has even started. + Context secondPrevious = + second + .submit( + () -> { + Context previous = element.updateThreadContext(element); + assertSame(leaf, tracer.activeSpan()); + leafScope.close(); + assertSame(child, tracer.activeSpan()); + childScope.close(); + assertSame(root, tracer.activeSpan()); + return previous; + }) + .get(10, SECONDS); + element.restoreThreadContext(element, firstPrevious); + assertNull(tracer.activeSpan(), "The first worker must recover its previous context"); + + third + .submit( + () -> { + Context previous = element.updateThreadContext(element); + try { + assertSame( + root, tracer.activeSpan(), "A late restore must not overwrite progress"); + rootScope.close(); + assertNull(tracer.activeSpan()); + } finally { + element.restoreThreadContext(element, previous); + } + }) + .get(10, SECONDS); + second + .submit( + () -> { + element.restoreThreadContext(element, secondPrevious); + assertNull( + tracer.activeSpan(), "The second worker must recover its previous context"); + Context previous = element.updateThreadContext(element); + try { + assertNull(tracer.activeSpan(), "Late restores must not resurrect closed scopes"); + } finally { + element.restoreThreadContext(element, previous); + } + }) + .get(10, SECONDS); + } finally { + second.shutdownNow(); + third.shutdownNow(); + leaf.finish(); + child.finish(); + root.finish(); + tracer.close(); + } + } +} 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..8930a20236f 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 @@ -2683,6 +2683,11 @@ public Context swap(Context context) { return scopeManager.swap(context); } + @Override + public Context forkContext(Context context) { + return scopeManager.forkContext(context); + } + @Override public ContextContinuation capture(Context context) { return scopeManager.capture(context); 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..a82d854af33 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 @@ -360,13 +360,32 @@ public ContextScope attach(@Nonnull Context context) { return activate(context); } + /** Creates a stack whose scope changes are visible before the next worker resumes execution. */ + public Context forkContext(@Nonnull Context context) { + MovableScopeContext stack = new MovableScopeContext(profilingContextIntegration); + if (context != Context.root()) { + stack.top = new ContinuableScope(this, context, CONTEXT, true, createScopeState(context)); + } + return stack; + } + public Context swap(@Nonnull Context context) { ScopeStack oldStack = tlsScopeStack.get(); - ContinuableScope oldScope = oldStack.top; + if (context == oldStack) { + // Scoped coroutines can reinstall the same stack, including during a late restore. + return context; + } + boolean moving = oldStack instanceof MovableScopeContext; + // A moving stack may already be in use by the next worker. Only detach it here. + ContinuableScope oldScope = moving ? null : oldStack.top; ScopeStack newStack; ContinuableScope newScope; - if (context instanceof ScopeContext) { + if (context instanceof MovableScopeContext) { + // The execution unit owns this stack; a departing worker must not copy or mutate it. + newStack = (MovableScopeContext) context; + newScope = newStack.top; + } else if (context instanceof ScopeContext) { // restore previously swapped context stack newStack = ((ScopeContext) context).restore(profilingContextIntegration); newScope = newStack.top; @@ -382,12 +401,12 @@ public Context swap(@Nonnull Context context) { } tlsScopeStack.set(newStack); - if (oldScope != newScope && newScope != null) { + if (oldStack != newStack && oldScope != newScope && newScope != null) { newScope.beforeActivated(); newScope.afterActivated(); } - return new ScopeContext(oldStack); + return moving ? (MovableScopeContext) oldStack : new ScopeContext(oldStack); } public ContextContinuation capture(@Nonnull Context context) { diff --git a/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/MovableScopeContext.java b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/MovableScopeContext.java new file mode 100644 index 00000000000..d8572c8aa78 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/core/scopemanager/MovableScopeContext.java @@ -0,0 +1,29 @@ +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 scope stack owned by a serial execution unit rather than a worker thread. */ +final class MovableScopeContext extends ScopeStack implements Context { + MovableScopeContext(ProfilingContextIntegration profilingContextIntegration) { + super(profilingContextIntegration); + } + + private Context current() { + ContinuableScope active = active(); + return active != null ? active.context : Context.root(); + } + + @Nullable + @Override + public T get(ContextKey key) { + return current().get(key); + } + + @Override + public Context with(ContextKey key, @Nullable T value) { + return current().with(key, value); + } +} 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..6aa10adec28 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,7 +10,7 @@ * 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 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..c9b9fe904e4 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 @@ -1054,6 +1054,40 @@ void rollbackStopsAtMostRecentCheckpoint() { assertNull(scopeManager.activeSpan()); } + @Test + void forkedContextsIsolateScopeChangesAndPreserveContextValues() { + ContextKey key = ContextKey.named("forked-context-test"); + Context parent = Context.root().with(key, "parent"); + ContextScope parentScope = scopeManager.attach(parent); + Context first = scopeManager.forkContext(parent); + Context second = scopeManager.forkContext(parent); + + Context previous = scopeManager.swap(first); + ContextScope firstScope = scopeManager.attach(parent.with(key, "first")); + assertEquals("first", first.get(key)); + assertEquals("parent", second.get(key)); + assertEquals("detached", first.with(key, "detached").get(key)); + assertEquals("first", first.get(key)); + + scopeManager.swap(second); + assertEquals("parent", scopeManager.currentContext().get(key)); + ContextScope secondScope = scopeManager.attach(parent.with(key, "second")); + + scopeManager.swap(first); + assertSame(firstScope, scopeManager.active()); + firstScope.close(); + assertEquals("parent", first.get(key)); + assertEquals("second", second.get(key)); + + scopeManager.swap(second); + assertSame(secondScope, scopeManager.active()); + secondScope.close(); + scopeManager.swap(previous); + assertSame(parentScope, scopeManager.active()); + parentScope.close(); + assertNull(scopeManager.active()); + } + @Test void contextsCanBeSwappedOutAndBack() { ContextKey testKey = ContextKey.named("test"); 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..7696412c661 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 @@ -401,6 +401,15 @@ default AgentSpan startSpan(@Nonnull SpanPrototype prototype, CharSequence opera void addShutdownListener(Runnable listener); + /** + * Creates an isolated scope stack for an execution unit that moves between threads. Execution + * must be handed off serially with a happens-before edge, but restoration of a previous worker + * may overlap the next activation. The returned context is a live view of that stack. + */ + default Context forkContext(Context context) { + return context; + } + // these methods are only used for legacy context manager migration @Deprecated