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/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/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy new file mode 100644 index 00000000000..691befe3279 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy @@ -0,0 +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 55516db89e0..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,16 +2,20 @@ 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; +import kotlin.coroutines.Continuation; import kotlin.coroutines.CoroutineContext; 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 implements ThreadContextElement { +public final class DatadogThreadContextElement + implements ThreadContextElement { private static final AtomicReferenceFieldUpdater CONTINUATION = AtomicReferenceFieldUpdater.newUpdater( @@ -20,6 +24,15 @@ public final class DatadogThreadContextElement implements ThreadContextElement 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 @@ -27,7 +40,11 @@ public static CoroutineContext addDatadogElement(CoroutineContext coroutineConte return coroutineContext.plus(new DatadogThreadContextElement()); } - private Context context; + private static final AtomicReferenceFieldUpdater EXCHANGE = + AtomicReferenceFieldUpdater.newUpdater( + DatadogThreadContextElement.class, Exchange.class, "exchange"); + + private volatile Exchange exchange; private volatile ContextContinuation continuation; @Nonnull @@ -38,11 +55,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(); } } @@ -57,20 +75,92 @@ public static void cancelDatadogContext(@Nonnull AbstractCoroutine coroutine) } @Override - public Context 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(); + public Exchange updateThreadContext(@Nonnull CoroutineContext coroutineContext) { + 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; } - return context.swap(); + current.active = true; + current.originalContext = current.context.swap(); + return current; } @Override - public void restoreThreadContext( - @Nonnull CoroutineContext coroutineContext, Context originalContext) { - context = originalContext.swap(); + public void restoreThreadContext(@Nonnull CoroutineContext coroutineContext, Exchange restored) { + restored.active = false; + // 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) { + parent.context = restored.context; + // Unwind same-thread nesting only if a newer worker has not taken ownership. + EXCHANGE.compareAndSet(this, restored, parent); + } + } + + 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) { + 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; + } + } + + /** 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; + } + 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(); + current.context = currentStack; + return current; + } + return null; + } + + static final class Exchange { + 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(Thread thread, Context context) { + this.thread = thread; + this.context = context; + } } @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..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,9 +15,13 @@ public KotlinCoroutinesModule() { @Override public List typeInstrumentations() { - return Arrays.asList( - new CoroutineContextInstrumentation(), - new CoroutineInstrumentation(), - new LazyCoroutineInstrumentation()); + 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 new file mode 100644 index 00000000000..50e6c48b8b7 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/SuspensionInstrumentation.java @@ -0,0 +1,106 @@ +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 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; + +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 dispatches context changes or publishes suspension decisions. + * Custom low-level suspension primitives can bypass these hooks. + */ +public class SuspensionInstrumentation + implements Instrumenter.ForKnownTypes, + Instrumenter.WithTypeStructure, + Instrumenter.HasMethodAdvice { + @Override + public String[] knownMatchingTypes() { + return new String[] { + "kotlinx.coroutines.BuildersKt__Builders_commonKt", + "kotlinx.coroutines.JobSupport", + "kotlinx.coroutines.CancellableContinuationImpl", + "kotlin.coroutines.SafeContinuation", + "kotlinx.coroutines.selects.SelectBuilderImpl", + "kotlinx.coroutines.DispatchedContinuation", + "kotlinx.coroutines.internal.DispatchedContinuation" + }; + } + + @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()) + .and(takesArguments(3)) + .and(takesArgument(2, named("kotlin.coroutines.Continuation"))), + SuspensionInstrumentation.class.getName() + "$ContextChangeAdvice"); + transformer.applyAdvice( + 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.captureContext(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) { + DatadogThreadContextElement.captureContext(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..36302f91e48 --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractStandardSuspensionHandoffTest.groovy @@ -0,0 +1,95 @@ +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 "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/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..e33babb317e --- /dev/null +++ b/dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/testFixtures/kotlin/datadog/trace/instrumentation/kotlin/coroutines/StandardSuspensionHandoffTests.kt @@ -0,0 +1,419 @@ +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 +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.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 +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 + +@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 + } + 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() + 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) + 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 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" } + } + } + } +}