Fix Kotlin coroutine scope propagation before suspension

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.*

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)

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()
+  }
+}

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<Context> {
+public final class DatadogThreadContextElement
+    implements ThreadContextElement<DatadogThreadContextElement.Exchange> {
   private static final AtomicReferenceFieldUpdater<DatadogThreadContextElement, ContextContinuation>
       CONTINUATION =
           AtomicReferenceFieldUpdater.newUpdater(
@@ -20,14 +24,27 @@ public final class DatadogThreadContextElement implements ThreadContextElement<C
   private static final CoroutineContext.Key<DatadogThreadContextElement> DATADOG_KEY =
       new CoroutineContext.Key<DatadogThreadContextElement>() {};
 
+  private static final ClassValue<Boolean> 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
     }
     return coroutineContext.plus(new DatadogThreadContextElement());
   }
 
-  private Context context;
+  private static final AtomicReferenceFieldUpdater<DatadogThreadContextElement, Exchange> 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

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<Instrumenter> typeInstrumentations() {
-    return Arrays.asList(
-        new CoroutineContextInstrumentation(),
-        new CoroutineInstrumentation(),
-        new LazyCoroutineInstrumentation());
+    List<Instrumenter> 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;
   }
 }

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<TypeDescription> 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);
+    }
+  }
+}

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 {
+}