From a9bcd5ce06524c75ed90644b6685f6279b47a2c2 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 6 Oct 2026 11:24:05 +0200 Subject: [PATCH 1/5] Fix Kotlin coroutine scope propagation before suspensio --- .../bytebuddy/matcher/ignored_class_name.trie | 1 + .../StandardSuspensionHandoffTest.groovy | 4 + .../DatadogThreadContextElement.java | 44 +++++-- .../coroutines/KotlinCoroutinesModule.java | 3 +- .../coroutines/SuspensionInstrumentation.java | 45 +++++++ .../StandardSuspensionHandoffTest.groovy | 4 + ...stractStandardSuspensionHandoffTest.groovy | 30 +++++ .../StandardSuspensionHandoffTests.kt | 117 ++++++++++++++++++ 8 files changed, 240 insertions(+), 8 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/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java 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/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index e719f4684f1..c7781e57b88 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -237,6 +237,7 @@ 2 javax.xml.* # note we do not ignore kotlinx because we instrument coroutines code 2 kotlin.* +0 kotlin.coroutines.SafeContinuation 2 net.sf.cglib.* 2 org.apache.bcel.* 2 org.apache.html.* 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..0f40bb4a9bd --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/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/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 55516db89e0..004d0f6e3d7 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 @@ -5,13 +5,15 @@ import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import kotlin.coroutines.Continuation; import kotlin.coroutines.CoroutineContext; import kotlin.jvm.functions.Function2; import kotlinx.coroutines.AbstractCoroutine; import kotlinx.coroutines.ThreadContextElement; /** Manages the Datadog context for coroutines, switching contexts as coroutines switch threads. */ -public final class DatadogThreadContextElement implements ThreadContextElement { +public final class DatadogThreadContextElement + implements ThreadContextElement { private static final AtomicReferenceFieldUpdater CONTINUATION = AtomicReferenceFieldUpdater.newUpdater( @@ -27,7 +29,11 @@ public static CoroutineContext addDatadogElement(CoroutineContext coroutineConte return coroutineContext.plus(new DatadogThreadContextElement()); } - private Context context; + private static final AtomicReferenceFieldUpdater CONTEXT = + AtomicReferenceFieldUpdater.newUpdater( + DatadogThreadContextElement.class, Context.class, "context"); + + private volatile Context context; private volatile ContextContinuation continuation; @Nonnull @@ -57,20 +63,44 @@ public static void cancelDatadogContext(@Nonnull AbstractCoroutine coroutine) } @Override - public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) { + public State updateThreadContext(@Nonnull CoroutineContext coroutineContext) { if (context == null) { // record context to use for this coroutine context = Context.current(); // stop enclosing trace from finishing early continuation = context.capture(); } - return context.swap(); + Context resumed = context; + return new State(resumed.swap(), resumed); } @Override - public void restoreThreadContext( - @Nonnull CoroutineContext coroutineContext, Context originalContext) { - context = originalContext.swap(); + public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, State state) { + Context suspended = state.originalContext.swap(); + // An early suspension snapshot or a newer restore takes precedence over this worker. + CONTEXT.compareAndSet(this, state.resumedContext, suspended); + } + + /** Publishes scope changes before suspension can resume the coroutine on another worker. */ + public static void beforeSuspension(Continuation continuation) { + DatadogThreadContextElement element = continuation.getContext().get(DATADOG_KEY); + if (element != null) { + // A swap captures the complete scope stack; Context.current() only exposes the active + // context. + Context currentStack = Context.root().swap(); + currentStack.swap(); + element.context = currentStack; + } + } + + static final class State { + final Context originalContext; + final Context resumedContext; + + State(Context originalContext, Context resumedContext) { + this.originalContext = originalContext; + this.resumedContext = resumedContext; + } } @Nonnull diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java index 23d202ec40d..896af4b2db9 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java @@ -17,6 +17,7 @@ public List typeInstrumentations() { return Arrays.asList( new CoroutineContextInstrumentation(), new CoroutineInstrumentation(), - new LazyCoroutineInstrumentation()); + new LazyCoroutineInstrumentation(), + new SuspensionInstrumentation()); } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java new file mode 100644 index 00000000000..8d4699ccbb3 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java @@ -0,0 +1,45 @@ +package datadog.trace.instrumentation.kotlin.coroutines; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith; +import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import kotlin.coroutines.Continuation; +import net.bytebuddy.asm.Advice; + +/** + * Snapshots context before Kotlin publishes standard suspension or yield decisions. Custom + * low-level suspension primitives can bypass these hooks. + */ +public class SuspensionInstrumentation + implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice { + @Override + public String[] knownMatchingTypes() { + return new String[] { + "kotlinx.coroutines.CancellableContinuationImpl", + "kotlin.coroutines.SafeContinuation", + "kotlinx.coroutines.DispatchedCoroutine", + "kotlinx.coroutines.selects.SelectBuilderImpl", + "kotlinx.coroutines.DispatchedContinuation", + "kotlinx.coroutines.internal.DispatchedContinuation" + }; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("getResult").or(named("getOrThrow")).and(takesNoArguments()), + SuspensionInstrumentation.class.getName() + "$SuspensionAdvice"); + transformer.applyAdvice( + nameStartsWith("dispatchYield"), + SuspensionInstrumentation.class.getName() + "$SuspensionAdvice"); + } + + public static class SuspensionAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void beforeSuspension(@Advice.This Continuation continuation) { + DatadogThreadContextElement.beforeSuspension(continuation); + } + } +} 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..738fb6314f6 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy @@ -0,0 +1,30 @@ +package datadog.trace.instrumentation.kotlin.coroutines + +import datadog.trace.agent.test.InstrumentationSpecification + +abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { + def "standard #kind suspension preserves manual scope with cancellation #cancelled"() { + setup: + new StandardSuspensionHandoffTests().run(kind, cancelled, { + 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 + 'cancellable' | false + 'cancellable' | true + 'safe' | false + 'yield' | false + 'delay' | false + 'select' | false + } +} 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..e3e676fcf53 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt @@ -0,0 +1,117 @@ +package datadog.trace.instrumentation.kotlin.coroutines + +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.get +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.Job +import kotlinx.coroutines.ThreadContextElement +import kotlinx.coroutines.asCoroutineDispatcher +import kotlinx.coroutines.delay +import kotlinx.coroutines.newCoroutineContext +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.selects.select +import kotlinx.coroutines.suspendCancellableCoroutine +import kotlinx.coroutines.yield +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.CoroutineContext +import kotlin.coroutines.resume +import kotlin.coroutines.suspendCoroutine + +@OptIn(ExperimentalCoroutinesApi::class) +class StandardSuspensionHandoffTests { + fun run(kind: String, cancelled: 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) { + 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() + } + } + } + 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 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 c46d067580d6a969c269dec48fbe8814eb415474 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 7 Oct 2026 14:31:08 +0200 Subject: [PATCH 2/5] suggestions --- .../kotlin-coroutines-1.3/build.gradle | 2 +- .../coroutines/DatadogThreadContextElement.java | 16 ++++++++-------- .../coroutines/KotlinCoroutinesModule.java | 16 ++++++++++------ .../coroutines/SuspensionInstrumentation.java | 12 +++++++++++- 4 files changed, 30 insertions(+), 16 deletions(-) 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 86abd12eb08..b40386e2bc8 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle @@ -53,7 +53,7 @@ dependencies { testImplementation testFixtures(project(':dd-java-agent:instrumentation:kotlin-coroutines-1.3')) latestDepTestImplementation libs.kotlin - latestDepTestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.6.+" + latestDepTestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.8.+" latestDepTestImplementation testFixtures(project(':dd-java-agent:instrumentation:kotlin-coroutines-1.3')) testFixturesCompileOnly(libs.bundles.groovy) 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 004d0f6e3d7..5a2b8897480 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 @@ -13,7 +13,7 @@ /** Manages the Datadog context for coroutines, switching contexts as coroutines switch threads. */ public final class DatadogThreadContextElement - implements ThreadContextElement { + implements ThreadContextElement { private static final AtomicReferenceFieldUpdater CONTINUATION = AtomicReferenceFieldUpdater.newUpdater( @@ -63,7 +63,7 @@ public static void cancelDatadogContext(@Nonnull AbstractCoroutine coroutine) } @Override - public State updateThreadContext(@Nonnull CoroutineContext coroutineContext) { + public Exchange updateThreadContext(@Nonnull CoroutineContext coroutineContext) { if (context == null) { // record context to use for this coroutine context = Context.current(); @@ -71,14 +71,14 @@ public State updateThreadContext(@Nonnull CoroutineContext coroutineContext) { continuation = context.capture(); } Context resumed = context; - return new State(resumed.swap(), resumed); + return new Exchange(resumed.swap(), resumed); } @Override - public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, State state) { - Context suspended = state.originalContext.swap(); + public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, Exchange exchange) { + Context suspended = exchange.originalContext.swap(); // An early suspension snapshot or a newer restore takes precedence over this worker. - CONTEXT.compareAndSet(this, state.resumedContext, suspended); + CONTEXT.compareAndSet(this, exchange.resumedContext, suspended); } /** Publishes scope changes before suspension can resume the coroutine on another worker. */ @@ -93,11 +93,11 @@ public static void beforeSuspension(Continuation continuation) { } } - static final class State { + static final class Exchange { final Context originalContext; final Context resumedContext; - State(Context originalContext, Context resumedContext) { + Exchange(Context originalContext, Context resumedContext) { this.originalContext = originalContext; this.resumedContext = resumedContext; } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java index 896af4b2db9..739785057c8 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/KotlinCoroutinesModule.java @@ -3,7 +3,8 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import java.util.Arrays; +import datadog.trace.api.InstrumenterConfig; +import java.util.ArrayList; import java.util.List; @AutoService(InstrumenterModule.class) @@ -14,10 +15,13 @@ public KotlinCoroutinesModule() { @Override public List typeInstrumentations() { - return Arrays.asList( - new CoroutineContextInstrumentation(), - new CoroutineInstrumentation(), - new LazyCoroutineInstrumentation(), - new SuspensionInstrumentation()); + List instrumenters = new ArrayList<>(); + instrumenters.add(new CoroutineContextInstrumentation()); + instrumenters.add(new CoroutineInstrumentation()); + instrumenters.add(new LazyCoroutineInstrumentation()); + if (InstrumenterConfig.get().isLegacyContextManagerEnabled()) { + instrumenters.add(new SuspensionInstrumentation()); + } + return instrumenters; } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java index 8d4699ccbb3..59e41a1ca69 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.kotlin.coroutines; +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; @@ -7,13 +8,17 @@ import datadog.trace.agent.tooling.Instrumenter; import kotlin.coroutines.Continuation; import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; /** * Snapshots context before Kotlin publishes standard suspension or yield decisions. Custom * low-level suspension primitives can bypass these hooks. */ public class SuspensionInstrumentation - implements Instrumenter.ForKnownTypes, Instrumenter.HasMethodAdvice { + implements Instrumenter.ForKnownTypes, + Instrumenter.WithTypeStructure, + Instrumenter.HasMethodAdvice { @Override public String[] knownMatchingTypes() { return new String[] { @@ -26,6 +31,11 @@ public String[] knownMatchingTypes() { }; } + @Override + public ElementMatcher structureMatcher() { + return hasSuperType(named("kotlin.coroutines.Continuation")); + } + @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( From c365c5d33f28025b1e3017d65fc8df7b2f5e30a9 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 7 Oct 2026 22:12:38 +0200 Subject: [PATCH 3/5] autotest --- .../DatadogThreadContextElement.java | 77 +++++++--- .../coroutines/SuspensionInstrumentation.java | 5 +- ...stractStandardSuspensionHandoffTest.groovy | 35 +++-- .../StandardSuspensionHandoffTests.kt | 141 ++++++++++++++---- 4 files changed, 199 insertions(+), 59 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 5a2b8897480..0b6271f1c5b 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 @@ -29,11 +29,11 @@ public static CoroutineContext addDatadogElement(CoroutineContext coroutineConte return coroutineContext.plus(new DatadogThreadContextElement()); } - private static final AtomicReferenceFieldUpdater CONTEXT = + private static final AtomicReferenceFieldUpdater EXCHANGE = AtomicReferenceFieldUpdater.newUpdater( - DatadogThreadContextElement.class, Context.class, "context"); + DatadogThreadContextElement.class, Exchange.class, "exchange"); - private volatile Context context; + private volatile Exchange exchange; private volatile ContextContinuation continuation; @Nonnull @@ -44,11 +44,12 @@ public Key getKey() { public static void captureDatadogContext(@Nonnull AbstractCoroutine coroutine) { DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY); - if (datadog != null && datadog.context == null) { + if (datadog != null && datadog.exchange == null) { // record context to use for this coroutine - datadog.context = Context.current(); + Context captured = Context.current(); + datadog.exchange = new Exchange(captured); // stop enclosing trace from finishing early - datadog.continuation = datadog.context.capture(); + datadog.continuation = captured.capture(); } } @@ -64,42 +65,70 @@ public static void cancelDatadogContext(@Nonnull AbstractCoroutine coroutine) @Override public Exchange updateThreadContext(@Nonnull CoroutineContext coroutineContext) { - if (context == null) { - // record context to use for this coroutine - context = Context.current(); - // stop enclosing trace from finishing early - continuation = context.capture(); + if (exchange == null) { + Context captured = Context.current(); + continuation = captured.capture(); + exchange = new Exchange(captured); + } + Exchange current = new Exchange(Thread.currentThread(), exchange.context); + Exchange previous = EXCHANGE.getAndSet(this, current); + // Read the actual predecessor after claiming ownership, including any completed restoration. + current.context = previous.context; + if (previous.thread == current.thread && previous.active) { + current.parent = previous; } - Context resumed = context; - return new Exchange(resumed.swap(), resumed); + current.active = true; + current.originalContext = current.context.swap(); + return current; } @Override - public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, Exchange exchange) { - Context suspended = exchange.originalContext.swap(); - // An early suspension snapshot or a newer restore takes precedence over this worker. - CONTEXT.compareAndSet(this, exchange.resumedContext, suspended); + public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, Exchange restored) { + restored.active = false; + Context original = restored.originalContext; + restored.originalContext = null; + restored.context = original.swap(); + Exchange parent = restored.parent; + restored.parent = null; + if (parent != null) { + parent.context = restored.context; + // Unwind same-thread nesting only if a newer worker has not taken ownership. + EXCHANGE.compareAndSet(this, restored, parent); + } } /** Publishes scope changes before suspension can resume the coroutine on another worker. */ public static void beforeSuspension(Continuation continuation) { DatadogThreadContextElement element = continuation.getContext().get(DATADOG_KEY); - if (element != null) { + if (element == null) { + return; + } + Exchange current = element.exchange; + if (current != null && current.thread == Thread.currentThread() && current.active) { // A swap captures the complete scope stack; Context.current() only exposes the active // context. Context currentStack = Context.root().swap(); currentStack.swap(); - element.context = currentStack; + current.context = currentStack; } } static final class Exchange { - final Context originalContext; - final Context resumedContext; + final Thread thread; + volatile Context context; + // The activating thread owns the remaining fields. + Context originalContext; + Exchange parent; + boolean active; + + Exchange(Context context) { + this.thread = null; + this.context = context; + } - Exchange(Context originalContext, Context resumedContext) { - this.originalContext = originalContext; - this.resumedContext = resumedContext; + Exchange(Thread thread, Context context) { + this.thread = thread; + this.context = context; } } diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java index 59e41a1ca69..278a343030a 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java @@ -39,7 +39,10 @@ public ElementMatcher structureMatcher() { @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( - named("getResult").or(named("getOrThrow")).and(takesNoArguments()), + named("getResult") + .or(named("getResult$kotlinx_coroutines_core")) + .or(named("getOrThrow")) + .and(takesNoArguments()), SuspensionInstrumentation.class.getName() + "$SuspensionAdvice"); transformer.applyAdvice( nameStartsWith("dispatchYield"), 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 738fb6314f6..f9db38143fd 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,9 +3,25 @@ package datadog.trace.instrumentation.kotlin.coroutines import datadog.trace.agent.test.InstrumentationSpecification abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { - def "standard #kind suspension preserves manual scope with cancellation #cancelled"() { + def "manual scope survives child completion after nested withContext"() { + when: + new StandardSuspensionHandoffTests().runNestedContextChange() + + 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, { + 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) @@ -19,12 +35,13 @@ abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpec cancelled || spans.child.parentId == spans.manual.spanId where: - kind | cancelled - 'cancellable' | false - 'cancellable' | true - 'safe' | false - 'yield' | false - 'delay' | false - 'select' | false + 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/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 e3e676fcf53..815cbf83a92 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 @@ -3,16 +3,21 @@ package datadog.trace.instrumentation.kotlin.coroutines import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan import datadog.trace.bootstrap.instrumentation.api.AgentTracer.get import kotlinx.coroutines.CancellationException +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.launch import kotlinx.coroutines.newCoroutineContext import kotlinx.coroutines.runBlocking import kotlinx.coroutines.selects.select import kotlinx.coroutines.suspendCancellableCoroutine +import kotlinx.coroutines.withContext import kotlinx.coroutines.yield import java.util.concurrent.CountDownLatch import java.util.concurrent.Executors @@ -27,7 +32,86 @@ import kotlin.coroutines.suspendCoroutine @OptIn(ExperimentalCoroutinesApi::class) class StandardSuspensionHandoffTests { - fun run(kind: String, cancelled: Boolean, beforeCompletion: Runnable) { + 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 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) @@ -51,38 +135,45 @@ class StandardSuspensionHandoffTests { try { try { runBlocking(context) { - 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 { + 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 { - check(activeSpan() === manual) { "Cleanup lost the manual scope across $kind" } - } finally { try { - scope.close() + check(activeSpan() === manual) { "Cleanup lost the manual scope across $kind" } } finally { - manual.finish() + try { + scope.close() + } finally { + manual.finish() + } } + resumer.join(10000) + check(!resumer.isAlive) + beforeCompletion.run() + } finally { + released.countDown() } - resumer.join(10000) - check(!resumer.isAlive) - beforeCompletion.run() - } finally { - released.countDown() } } + if (nested) { + withContext(CoroutineName("nested")) { suspendWithManualScope() } + } else { + suspendWithManualScope() + } } check(!cancelled) } catch (expected: CancellationException) { From a40328ef547970fd05f9ec46f5ab502808584135 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 7 Oct 2026 22:58:12 +0200 Subject: [PATCH 4/5] autotest2 --- .../StandardSuspensionHandoffTest.groovy | 8 ++ .../DatadogThreadContextElement.java | 29 +++- .../coroutines/SuspensionInstrumentation.java | 35 ++++- ...stractStandardSuspensionHandoffTest.groovy | 22 +++ .../StandardSuspensionHandoffTests.kt | 131 ++++++++++++++++++ 5 files changed, 218 insertions(+), 7 deletions(-) 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 index 0f40bb4a9bd..691befe3279 100644 --- 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 @@ -1,4 +1,12 @@ import datadog.trace.instrumentation.kotlin.coroutines.AbstractStandardSuspensionHandoffTest +import datadog.trace.instrumentation.kotlin.coroutines.StandardSuspensionHandoffTests class StandardSuspensionHandoffTest extends AbstractStandardSuspensionHandoffTest { + def "undispatched coroutine can replay its restoration token"() { + when: + new StandardSuspensionHandoffTests().runReplayedRestoration() + + then: + noExceptionThrown() + } } 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 0b6271f1c5b..56461eb6f1a 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 @@ -85,9 +85,8 @@ public Exchange updateThreadContext(@Nonnull CoroutineContext coroutineContext) @Override public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, Exchange restored) { restored.active = false; - Context original = restored.originalContext; - restored.originalContext = null; - restored.context = original.swap(); + // UndispatchedCoroutine can restore the same token again when no dispatcher owns restoration. + restored.context = restored.originalContext.swap(); Exchange parent = restored.parent; restored.parent = null; if (parent != null) { @@ -99,9 +98,29 @@ public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, Exc /** Publishes scope changes before suspension can resume the coroutine on another worker. */ public static void beforeSuspension(Continuation continuation) { + captureContext(continuation); + } + + public static Object beforeContextChange(Continuation continuation) { + return captureContext(continuation); + } + + public static void afterContextChange(Continuation continuation, Object original) { + if (original == null) { + return; + } + Exchange source = (Exchange) original; + if (source.thread == Thread.currentThread() && source.active) { + // A synchronous return leaves the caller running, even if destination cleanup is still + // pending. + continuation.getContext().get(DATADOG_KEY).exchange = source; + } + } + + private static Exchange captureContext(Continuation continuation) { DatadogThreadContextElement element = continuation.getContext().get(DATADOG_KEY); if (element == null) { - return; + return null; } Exchange current = element.exchange; if (current != null && current.thread == Thread.currentThread() && current.active) { @@ -110,7 +129,9 @@ public static void beforeSuspension(Continuation continuation) { Context currentStack = Context.root().swap(); currentStack.swap(); current.context = currentStack; + return current; } + return null; } static final class Exchange { diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java index 278a343030a..b3bc2671122 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java @@ -2,7 +2,11 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static kotlin.coroutines.intrinsics.IntrinsicsKt.getCOROUTINE_SUSPENDED; +import static net.bytebuddy.matcher.ElementMatchers.isStatic; import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import datadog.trace.agent.tooling.Instrumenter; @@ -12,8 +16,8 @@ import net.bytebuddy.matcher.ElementMatcher; /** - * Snapshots context before Kotlin publishes standard suspension or yield decisions. Custom - * low-level suspension primitives can bypass these hooks. + * Snapshots context before Kotlin dispatches context changes or publishes suspension decisions. + * Custom low-level suspension primitives can bypass these hooks. */ public class SuspensionInstrumentation implements Instrumenter.ForKnownTypes, @@ -22,6 +26,7 @@ public class SuspensionInstrumentation @Override public String[] knownMatchingTypes() { return new String[] { + "kotlinx.coroutines.BuildersKt__Builders_commonKt", "kotlinx.coroutines.CancellableContinuationImpl", "kotlin.coroutines.SafeContinuation", "kotlinx.coroutines.DispatchedCoroutine", @@ -33,11 +38,18 @@ public String[] knownMatchingTypes() { @Override public ElementMatcher structureMatcher() { - return hasSuperType(named("kotlin.coroutines.Continuation")); + return named("kotlinx.coroutines.BuildersKt__Builders_commonKt") + .or(hasSuperType(named("kotlin.coroutines.Continuation"))); } @Override public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("withContext") + .and(isStatic()) + .and(takesArguments(3)) + .and(takesArgument(2, named("kotlin.coroutines.Continuation"))), + SuspensionInstrumentation.class.getName() + "$ContextChangeAdvice"); transformer.applyAdvice( named("getResult") .or(named("getResult$kotlinx_coroutines_core")) @@ -49,6 +61,23 @@ public void methodAdvice(MethodTransformer transformer) { SuspensionInstrumentation.class.getName() + "$SuspensionAdvice"); } + public static class ContextChangeAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static Object beforeContextChange(@Advice.Argument(2) Continuation continuation) { + return DatadogThreadContextElement.beforeContextChange(continuation); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void afterContextChange( + @Advice.Argument(2) Continuation continuation, + @Advice.Enter Object original, + @Advice.Return Object result) { + if (result != getCOROUTINE_SUSPENDED()) { + DatadogThreadContextElement.afterContextChange(continuation, original); + } + } + } + public static class SuspensionAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void beforeSuspension(@Advice.This Continuation continuation) { 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 f9db38143fd..fd0cdf9208e 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,28 @@ package datadog.trace.instrumentation.kotlin.coroutines import datadog.trace.agent.test.InstrumentationSpecification abstract class AbstractStandardSuspensionHandoffTest extends InstrumentationSpecification { + 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() 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 815cbf83a92..ee78e4e61d6 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 @@ -1,8 +1,11 @@ package datadog.trace.instrumentation.kotlin.coroutines +import datadog.context.Context 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 @@ -26,12 +29,124 @@ 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 @OptIn(ExperimentalCoroutinesApi::class) class StandardSuspensionHandoffTests { + 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 runNestedContextChange() { val workers = Executors.newSingleThreadExecutor() val otherWorkers = Executors.newSingleThreadExecutor() @@ -191,6 +306,22 @@ class StandardSuspensionHandoffTests { } } + 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 From 2c6da9faadda673d803518f74fb3c62c7bf7373b Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 7 Oct 2026 23:54:34 +0200 Subject: [PATCH 5/5] simplify --- .../DatadogThreadContextElement.java | 26 ++++-- .../coroutines/SuspensionInstrumentation.java | 33 ++++++-- ...stractStandardSuspensionHandoffTest.groovy | 26 ++++++ .../StandardSuspensionHandoffTests.kt | 80 +++++++++++++++++++ 4 files changed, 150 insertions(+), 15 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 56461eb6f1a..946c6c511f2 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 @@ -2,6 +2,7 @@ import datadog.context.Context; import datadog.context.ContextContinuation; +import datadog.trace.api.GenericClassValue; import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -10,6 +11,7 @@ import kotlin.jvm.functions.Function2; import kotlinx.coroutines.AbstractCoroutine; import kotlinx.coroutines.ThreadContextElement; +import kotlinx.coroutines.internal.ScopeCoroutine; /** Manages the Datadog context for coroutines, switching contexts as coroutines switch threads. */ public final class DatadogThreadContextElement @@ -22,6 +24,15 @@ public final class DatadogThreadContextElement private static final CoroutineContext.Key DATADOG_KEY = new CoroutineContext.Key() {}; + private static final ClassValue SCOPED_COROUTINE = + GenericClassValue.of( + type -> { + // Older timeout coroutines are scoped but do not extend ScopeCoroutine. + return ScopeCoroutine.class.isAssignableFrom(type) + || (AbstractCoroutine.class.isAssignableFrom(type) + && "kotlinx.coroutines.TimeoutCoroutine".equals(type.getName())); + }); + public static CoroutineContext addDatadogElement(CoroutineContext coroutineContext) { if (coroutineContext.get(DATADOG_KEY) != null) { return coroutineContext; // already added @@ -96,13 +107,11 @@ public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, Exc } } - /** Publishes scope changes before suspension can resume the coroutine on another worker. */ - public static void beforeSuspension(Continuation continuation) { - captureContext(continuation); - } - - public static Object beforeContextChange(Continuation continuation) { - return captureContext(continuation); + public static void beforeScopedCompletion(Object coroutine) { + if (SCOPED_COROUTINE.get(coroutine.getClass())) { + // The scoped body has finished changing its scopes; children can resume it during completion. + captureContext((AbstractCoroutine) coroutine); + } } public static void afterContextChange(Continuation continuation, Object original) { @@ -117,7 +126,8 @@ public static void afterContextChange(Continuation continuation, Object origi } } - private static Exchange captureContext(Continuation continuation) { + /** Publishes the active scope stack before Kotlin can hand execution to another worker. */ + public static Object captureContext(Continuation continuation) { DatadogThreadContextElement element = continuation.getContext().get(DATADOG_KEY); if (element == null) { return null; diff --git a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java index b3bc2671122..50e6c48b8b7 100644 --- a/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java @@ -5,6 +5,8 @@ import static kotlin.coroutines.intrinsics.IntrinsicsKt.getCOROUTINE_SUSPENDED; import static net.bytebuddy.matcher.ElementMatchers.isStatic; import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith; +import static net.bytebuddy.matcher.ElementMatchers.not; +import static net.bytebuddy.matcher.ElementMatchers.returns; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; @@ -27,9 +29,9 @@ public class SuspensionInstrumentation public String[] knownMatchingTypes() { return new String[] { "kotlinx.coroutines.BuildersKt__Builders_commonKt", + "kotlinx.coroutines.JobSupport", "kotlinx.coroutines.CancellableContinuationImpl", "kotlin.coroutines.SafeContinuation", - "kotlinx.coroutines.DispatchedCoroutine", "kotlinx.coroutines.selects.SelectBuilderImpl", "kotlinx.coroutines.DispatchedContinuation", "kotlinx.coroutines.internal.DispatchedContinuation" @@ -39,11 +41,24 @@ public String[] knownMatchingTypes() { @Override public ElementMatcher structureMatcher() { return named("kotlinx.coroutines.BuildersKt__Builders_commonKt") + .or(named("kotlinx.coroutines.JobSupport")) .or(hasSuperType(named("kotlin.coroutines.Continuation"))); } @Override public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + named("makeCompletingOnce$kotlinx_coroutines_core") + .and(not(isStatic())) + .and(takesArgument(0, Object.class)) + .and( + takesArguments(1) + .and(returns(Object.class)) + .or( + takesArguments(2) + .and(takesArgument(1, int.class)) + .and(returns(boolean.class)))), + SuspensionInstrumentation.class.getName() + "$ScopedCompletionAdvice"); transformer.applyAdvice( named("withContext") .and(isStatic()) @@ -51,20 +66,24 @@ public void methodAdvice(MethodTransformer transformer) { .and(takesArgument(2, named("kotlin.coroutines.Continuation"))), SuspensionInstrumentation.class.getName() + "$ContextChangeAdvice"); transformer.applyAdvice( - named("getResult") - .or(named("getResult$kotlinx_coroutines_core")) - .or(named("getOrThrow")) - .and(takesNoArguments()), + named("getResult").or(named("getOrThrow")).and(takesNoArguments()), SuspensionInstrumentation.class.getName() + "$SuspensionAdvice"); transformer.applyAdvice( nameStartsWith("dispatchYield"), SuspensionInstrumentation.class.getName() + "$SuspensionAdvice"); } + public static class ScopedCompletionAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void beforeCompletion(@Advice.This Object coroutine) { + DatadogThreadContextElement.beforeScopedCompletion(coroutine); + } + } + public static class ContextChangeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static Object beforeContextChange(@Advice.Argument(2) Continuation continuation) { - return DatadogThreadContextElement.beforeContextChange(continuation); + return DatadogThreadContextElement.captureContext(continuation); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) @@ -81,7 +100,7 @@ public static void afterContextChange( public static class SuspensionAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void beforeSuspension(@Advice.This Continuation continuation) { - DatadogThreadContextElement.beforeSuspension(continuation); + DatadogThreadContextElement.captureContext(continuation); } } } 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 fd0cdf9208e..36302f91e48 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,32 @@ 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) 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 ee78e4e61d6..e33babb317e 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 @@ -19,8 +19,10 @@ 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.CountDownLatch import java.util.concurrent.Executors @@ -36,6 +38,84 @@ import kotlin.coroutines.suspendCoroutine @OptIn(ExperimentalCoroutinesApi::class) class StandardSuspensionHandoffTests { + 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