Fix Kotlin coroutine context isolation and continuation ownership

dd-java-agent/instrumentation/kotlin-coroutines-1.3/build.gradle

@@ -35,6 +35,9 @@ kotlin {
   }
 }
 
+// Reuse the flow and context-lifecycle tests across intermediate and current releases.
+addTestSuiteForDir('coroutines16Test', 'latestDepTest')
+addTestSuiteForDir('coroutines18Test', 'latestDepTest')
 addTestSuite('latestDepTest')
 
 dependencies {
@@ -52,10 +55,15 @@ dependencies {
   testImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0"
   testImplementation testFixtures(project(':dd-java-agent:instrumentation:kotlin-coroutines-1.3'))
 
+  coroutines16TestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.6.4"
+  coroutines18TestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.8.1"
+
   latestDepTestImplementation libs.kotlin
-  latestDepTestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:1.6.+"
+  latestDepTestImplementation "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.+"
   latestDepTestImplementation testFixtures(project(':dd-java-agent:instrumentation:kotlin-coroutines-1.3'))
 
   testFixturesCompileOnly(libs.bundles.groovy)
   testFixturesCompileOnly(libs.bundles.spock)
+  testFixturesCompileOnly libs.junit.jupiter
+  testFixturesCompileOnly libs.tabletest
 }

dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java

@@ -0,0 +1,3 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+class DatadogThreadContextElementTest extends AbstractDatadogThreadContextElementTest {}

dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/kotlin/KotlinCoroutineTests.kt

@@ -12,7 +12,7 @@ class KotlinCoroutineTests(dispatcher: CoroutineDispatcher) : CoreKotlinCoroutin
 
   @Trace
   fun tracedAcrossFlows(withModifiedContext: Boolean): Int = runTest {
-    val producer = flow {
+    val producer = flow<Int> {
       repeat(3) {
         tracedChild("produce_$it")
         if (withModifiedContext) {
@@ -36,7 +36,7 @@ class KotlinCoroutineTests(dispatcher: CoroutineDispatcher) : CoreKotlinCoroutin
 
   @Trace
   fun traceAfterFlow(): Int = runTest {
-    val f = flow {
+    val f = flow<Int> {
       childSpan("inside-flow").activateAndUse {
         println("insideFlowSpan")
       }
@@ -51,6 +51,18 @@ class KotlinCoroutineTests(dispatcher: CoroutineDispatcher) : CoreKotlinCoroutin
     3
   }
 
+  @Trace
+  override fun structuredChildrenOutliveSeedSpan(): Int = super.structuredChildrenOutliveSeedSpan()
+
+  @Trace
+  override fun structuredChildrenOfScopedCoroutines(suspendScope: Boolean): Int = super.structuredChildrenOfScopedCoroutines(suspendScope)
+
+  @Trace
+  override fun childWithReplacedJob(supervisor: Boolean): Int = super.childWithReplacedJob(supervisor)
+
+  @Trace
+  override fun structuredChildOfCancelledCoroutine(): Int = super.structuredChildOfCancelledCoroutine()
+
   @Trace
   override fun tracePreventedByCancellation(): Int = super.tracePreventedByCancellation()
 

dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElement.java

@@ -2,16 +2,23 @@
 
 import datadog.context.Context;
 import datadog.context.ContextContinuation;
+import datadog.trace.api.GenericClassValue;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import kotlin.coroutines.CoroutineContext;
+import kotlin.coroutines.EmptyCoroutineContext;
 import kotlin.jvm.functions.Function2;
 import kotlinx.coroutines.AbstractCoroutine;
+import kotlinx.coroutines.Job;
 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<Context> {
+public final class DatadogThreadContextElement
+    implements ThreadContextElement<Context>,
+        Function2<CoroutineContext, CoroutineContext.Element, CoroutineContext> {
   private static final AtomicReferenceFieldUpdater<DatadogThreadContextElement, ContextContinuation>
       CONTINUATION =
           AtomicReferenceFieldUpdater.newUpdater(
@@ -20,16 +27,63 @@ public final class DatadogThreadContextElement implements ThreadContextElement<C
   private static final CoroutineContext.Key<DatadogThreadContextElement> DATADOG_KEY =
       new CoroutineContext.Key<DatadogThreadContextElement>() {};
 
+  private static final Object COMPLETED = new Object();
+
+  /** Coroutines built by coroutineScope, withContext, etc. that their caller waits for. */
+  private static final ClassValue<Boolean> SCOPED_COROUTINE =
+      GenericClassValue.of(
+          type ->
+              ScopeCoroutine.class.isAssignableFrom(type)
+                  // older timeout coroutines are scoped but don't extend ScopeCoroutine
+                  || (AbstractCoroutine.class.isAssignableFrom(type)
+                      && "kotlinx.coroutines.TimeoutCoroutine".equals(type.getName())));
+
+  /** Adds a fresh Datadog element, replacing any inherited element in place. */
   public static CoroutineContext addDatadogElement(CoroutineContext coroutineContext) {
-    if (coroutineContext.get(DATADOG_KEY) != null) {
-      return coroutineContext; // already added
+    DatadogThreadContextElement existing = coroutineContext.get(DATADOG_KEY);
+    if (existing == null) {
+      return coroutineContext.plus(new DatadogThreadContextElement(null));
     }
-    return coroutineContext.plus(new DatadogThreadContextElement());
+    DatadogThreadContextElement fresh =
+        new DatadogThreadContextElement(existing.localRootFor(coroutineContext.get(Job.Key)));
+    // keep the element's position, so elements added after it can still update scopes on top
+    return coroutineContext.fold(EmptyCoroutineContext.INSTANCE, fresh);
+  }
+
+  /** Rebuilds the context using this replacement element, avoiding a capturing fold lambda. */
+  @Override
+  public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Element element) {
+    return rebuilt.plus(element.getKey() == DATADOG_KEY ? this : element);
   }
 
-  private Context context;
+  @Nullable private AgentSpan parentLocalRootSpan;
+  private volatile Context context;
   private volatile ContextContinuation continuation;
 
+  /** The owning Job, null before initialization, or COMPLETED after releasing the owner. */
+  @Nullable private volatile Object owner;
+
+  /** The local root of the trace kept open by this coroutine or its parent. */
+  @Nullable private volatile AgentSpan localRootSpan;
+
+  private DatadogThreadContextElement(@Nullable AgentSpan parentLocalRootSpan) {
+    this.parentLocalRootSpan = parentLocalRootSpan;
+  }
+
+  /** Returns the local root kept open for structured children launched with this job. */
+  @Nullable
+  private AgentSpan localRootFor(@Nullable Job job) {
+    AgentSpan rootSpan = localRootSpan;
+    // structured children complete before their parent, which keeps its trace open until then
+    return rootSpan != null && job != null && (job == owner || isScopedBy(job)) ? rootSpan : null;
+  }
+
+  /** Scoped coroutines share their caller's element and complete before the caller does. */
+  private boolean isScopedBy(Job job) {
+    return SCOPED_COROUTINE.get(job.getClass())
+        && ((AbstractCoroutine<?>) job).getContext().get(DATADOG_KEY) == this;
+  }
+
   @Nonnull
   @Override
   public Key<?> getKey() {
@@ -38,18 +92,59 @@ public Key<?> getKey() {
 
   public static void captureDatadogContext(@Nonnull AbstractCoroutine<?> coroutine) {
     DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY);
-    if (datadog != null && datadog.context == null) {
-      // record context to use for this coroutine
-      datadog.context = Context.current();
-      // stop enclosing trace from finishing early
-      datadog.continuation = datadog.context.capture();
+    if (datadog != null) {
+      if (datadog.owner == null) {
+        // scoped coroutines built later share this element, but aren't its owner
+        datadog.owner = coroutine;
+      }
+      if (datadog.context == null) {
+        datadog.captureInitialContext();
+      } else if (datadog.owner != coroutine && datadog.isScopedBy(coroutine)) {
+        // preserve active span changes when a scoped coroutine reinstalls this element
+        datadog.context = Context.current();
+      }
+    }
+  }
+
+  private synchronized void captureInitialContext() {
+    if (context == null) {
+      Context initial = Context.current();
+      // this coroutine isn't running yet, so its context can't change under us
+      AgentSpan rootSpan = localRootOf(initial);
+      if (rootSpan != null) {
+        if (rootSpan == parentLocalRootSpan) {
+          // the parent keeps this trace open until after this child completes
+          localRootSpan = rootSpan;
+        } else {
+          // stop enclosing trace from finishing early
+          ContextContinuation captured = initial.capture();
+          continuation = captured;
+          if (captured.context() != Context.root()) {
+            // only when the capture was accepted, e.g. async propagation is enabled
+            localRootSpan = rootSpan;
+          }
+        }
+      }
+      context = initial;
+      parentLocalRootSpan = null;
     }
   }
 
+  @Nullable
+  private static AgentSpan localRootOf(Context context) {
+    AgentSpan span = AgentSpan.fromContext(context);
+    return span != null ? span.getLocalRootSpan() : null;
+  }
+
   public static void cancelDatadogContext(@Nonnull AbstractCoroutine<?> coroutine) {
     DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY);
-    ContextContinuation continuation =
-        datadog == null ? null : CONTINUATION.getAndSet(datadog, null);
+    if (datadog == null || datadog.owner != coroutine) {
+      // scoped coroutines share their caller's element, but must not release its trace hold
+      return;
+    }
+    datadog.owner = COMPLETED;
+    datadog.localRootSpan = null;
+    ContextContinuation continuation = CONTINUATION.getAndSet(datadog, null);
     if (continuation != null) {
       // release enclosing trace now the coroutine has completed
       continuation.release();
@@ -58,13 +153,16 @@ 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();
+    Context current = context;
+    if (current == null) {
+      if (owner == null) {
+        // an already-cancelled coroutine can skip constructor capture but still install context
+        owner = coroutineContext.get(Job.Key);
+      }
+      captureInitialContext();
+      current = context;
     }
-    return context.swap();
+    return current.swap();
   }
 
   @Override

dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/java/datadog/trace/instrumentation/kotlin/coroutines/DatadogThreadContextElementTest.java

@@ -0,0 +1,3 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+class DatadogThreadContextElementTest extends AbstractDatadogThreadContextElementTest {}

dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/test/kotlin/KotlinCoroutineTests.kt

@@ -32,6 +32,18 @@ class KotlinCoroutineTests(dispatcher: CoroutineDispatcher) : CoreKotlinCoroutin
     7
   }
 
+  @Trace
+  override fun structuredChildrenOutliveSeedSpan(): Int = super.structuredChildrenOutliveSeedSpan()
+
+  @Trace
+  override fun structuredChildrenOfScopedCoroutines(suspendScope: Boolean): Int = super.structuredChildrenOfScopedCoroutines(suspendScope)
+
+  @Trace
+  override fun childWithReplacedJob(supervisor: Boolean): Int = super.childWithReplacedJob(supervisor)
+
+  @Trace
+  override fun structuredChildOfCancelledCoroutine(): Int = super.structuredChildOfCancelledCoroutine()
+
   @Trace
   override fun tracePreventedByCancellation(): Int = super.tracePreventedByCancellation()