Fix Kotlin coroutines losing scopes when resumed on another thread

dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java

@@ -73,7 +73,7 @@ public final class VirtualThreadInstrumentation extends InstrumenterModule.Conte
    */
   private static final String[] PRELOAD_CLASS_NAMES = {
     VIRTUAL_THREAD_STATE_CLASS_NAME,
-    "datadog.trace.core.scopemanager.ScopeContext",
+    "datadog.trace.core.scopemanager.SavedScopeContext",
     "datadog.trace.core.scopemanager.ScopeStack"
   };
 

dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadScopeClassLoadingForkedTest.java

@@ -0,0 +1,53 @@
+package testdog.trace.instrumentation.java.lang.jdk21;
+
+import static datadog.trace.agent.test.assertions.SpanMatcher.span;
+import static datadog.trace.agent.test.assertions.TraceMatcher.trace;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import datadog.context.ContextScope;
+import datadog.trace.agent.test.AbstractInstrumentationTest;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.time.Duration;
+import java.util.concurrent.locks.LockSupport;
+import net.bytebuddy.agent.ByteBuddyAgent;
+import org.junit.jupiter.api.Test;
+
+/** Isolates class-loading assertions from tests that may deliberately load coroutine helpers. */
+class VirtualThreadScopeClassLoadingForkedTest extends AbstractInstrumentationTest {
+  @Test
+  void virtualThreadHandoffsDoNotLoadSharedScopeStacks() throws InterruptedException {
+    AgentSpan parent = startSpan("test", "parent");
+    AgentSpan[] active = new AgentSpan[2];
+    try (ContextScope ignored = activateSpan(parent)) {
+      Thread child =
+          Thread.startVirtualThread(
+              () -> {
+                active[0] = activeSpan();
+                LockSupport.parkNanos(MILLISECONDS.toNanos(1));
+                active[1] = activeSpan();
+              });
+      assertTrue(child.join(Duration.ofSeconds(10)));
+      assertSame(parent, active[0]);
+      assertSame(parent, active[1]);
+    } finally {
+      parent.finish();
+    }
+    assertTraces(trace(span().root().operationName("parent")));
+
+    for (Class<?> loaded : ByteBuddyAgent.getInstrumentation().getAllLoadedClasses()) {
+      if (loaded != null) {
+        String name = loaded.getName();
+        assertFalse(
+            "datadog.trace.core.scopemanager.SharedScopeStack".equals(name)
+                || "datadog.trace.core.scopemanager.SharedScopeContext".equals(name),
+            name);
+      }
+    }
+  }
+}

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

@@ -41,6 +41,13 @@ addTestSuiteForDir('coroutines164Test', 'latestDepTest')
 addTestSuiteForDir('coroutines18Test', 'latestDepTest')
 addTestSuite('latestDepTest')
 
+tasks.named('coroutines16Test') {
+  filter {
+    // kotlinx only replays undispatched restore tokens from 1.6.2 (kotlinx.coroutines#2930)
+    excludeTestsMatching 'StandardSuspensionHandoffTest.undispatched coroutine can replay its restoration token'
+  }
+}
+
 dependencies {
   api project(':dd-java-agent:instrumentation:java:java-concurrent:java-concurrent-1.8')
   compileOnly libs.kotlin

dd-java-agent/instrumentation/kotlin-coroutines-1.3/src/latestDepTest/groovy/StandardSuspensionHandoffTest.groovy

@@ -0,0 +1,43 @@
+import datadog.trace.instrumentation.kotlin.coroutines.AbstractStandardSuspensionHandoffTest
+import datadog.trace.instrumentation.kotlin.coroutines.StandardSuspensionHandoffTests
+
+class StandardSuspensionHandoffTest extends AbstractStandardSuspensionHandoffTest {
+  // withContext uses the context merge hook from coroutines 1.6.1 onward.
+  def "context merge preserves an explicitly swapped span"() {
+    when:
+    new StandardSuspensionHandoffTests().runContextSwapAcrossMerge()
+    TEST_WRITER.waitForTraces(1)
+    def trace = TEST_WRITER.get(0)
+    def spans = trace.collectEntries { [(it.operationName.toString()): it] }
+
+    then:
+    trace.size() == 3
+    spans.swapped.parentId == spans.parent.spanId
+    spans.child.parentId == spans.swapped.spanId
+  }
+
+  def "undispatched coroutine can replay its restoration token"() {
+    when:
+    new StandardSuspensionHandoffTests().runReplayedRestoration()
+
+    then:
+    noExceptionThrown()
+  }
+
+  def "scope-touching thread context elements do not corrupt the handed-off stack"() {
+    when:
+    def violations = new StandardSuspensionHandoffTests().runScopeTouchingElements(1000)
+
+    then:
+    // Overlapping workers can temporarily see each other's element scopes. The coroutine must
+    // retain its base context, but restore callbacks can observe a different top during overlap.
+    def expectedDuringOverlap = [
+      'element scope leaked from previous worker',
+      'before scope not on top at restore',
+      'after scope not on top at restore',
+      'before restore did not return to previous context',
+      'after restore did not return to previous context'
+    ]
+    (violations.keySet() - expectedDuringOverlap).isEmpty()
+  }
+}

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

@@ -3,7 +3,10 @@
 import datadog.context.Context;
 import datadog.context.ContextContinuation;
 import datadog.trace.api.GenericClassValue;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -29,6 +32,9 @@ public final class DatadogThreadContextElement
 
   private static final Object COMPLETED = new Object();
 
+  private static final boolean LEGACY_CONTEXT_MANAGER =
+      InstrumenterConfig.get().isLegacyContextManagerEnabled();
+
   /** Coroutines built by coroutineScope, withContext, etc. that their caller waits for. */
   private static final ClassValue<Boolean> SCOPED_COROUTINE =
       GenericClassValue.of(
@@ -51,15 +57,23 @@ public final class DatadogThreadContextElement
   public static CoroutineContext addDatadogElement(CoroutineContext coroutineContext) {
     DatadogThreadContextElement existing = coroutineContext.get(DATADOG_KEY);
     if (existing == null) {
-      return coroutineContext.plus(new DatadogThreadContextElement(null));
+      return coroutineContext.plus(new DatadogThreadContextElement(null, null));
     }
     return replaceDatadogElement(
-        coroutineContext, existing.localRootFor(coroutineContext.get(Job.Key)));
+        coroutineContext,
+        LEGACY_CONTEXT_MANAGER ? Context.current() : null,
+        existing.localRootFor(coroutineContext.get(Job.Key)));
   }
 
   private static CoroutineContext replaceDatadogElement(
-      CoroutineContext coroutineContext, @Nullable AgentSpan parentLocalRootSpan) {
-    DatadogThreadContextElement fresh = new DatadogThreadContextElement(parentLocalRootSpan);
+      CoroutineContext coroutineContext,
+      @Nullable Context legacySeed,
+      @Nullable AgentSpan parentLocalRootSpan) {
+    // Only legacy mode inherits a live stack. Non-legacy mode captures Context.current()
+    // later on during child initialization because the parent's saved context may be stale.
+    DatadogThreadContextElement fresh =
+        new DatadogThreadContextElement(
+            LEGACY_CONTEXT_MANAGER ? legacySeed : null, parentLocalRootSpan);
     if (!Boolean.TRUE.equals(coroutineContext.fold(null, HAS_LATER_HOOK))) {
       return coroutineContext.plus(fresh);
     }
@@ -76,22 +90,26 @@ public static CoroutineContext forContextMerge(
     DatadogThreadContextElement caller = original.get(DATADOG_KEY);
     // Scoped coroutines can borrow the caller's hold even when the requested context replaces Job.
     // Flow merges may install this element without constructing a coroutine to own it.
+    // An explicit context swap can make the active context differ from the caller's shared stack.
     return replaceDatadogElement(
-        merged, caller == null ? null : caller.localRootFor(original.get(Job.Key)));
+        merged,
+        LEGACY_CONTEXT_MANAGER ? Context.current() : null,
+        caller == null ? null : caller.localRootFor(original.get(Job.Key)));
   }
 
   /** Separates upstream flow state while retaining the caller's active context for downstream. */
   public static CoroutineContext forFlowCollection(
       CoroutineContext upstream, CoroutineContext downstream) {
     DatadogThreadContextElement caller = downstream.get(DATADOG_KEY);
-    if (caller != null) {
+    // Read context before legacySharedStack to acquire the initialized stack state.
+    if (caller != null && caller.context != null && !caller.legacySharedStack) {
       // Downstream reinstalls the caller's element without constructing a coroutine.
       caller.context = Context.current();
     }
     DatadogThreadContextElement element = upstream.get(DATADOG_KEY);
     if (element == null || element == caller) {
       // Before 1.7.2, flowOn merges with plus and shares this mutable element with its caller.
-      return replaceDatadogElement(upstream, null);
+      return replaceDatadogElement(upstream, null, null);
     }
     // Newer releases already received a fresh element from the context merge advice.
     return upstream;
@@ -113,7 +131,17 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen
     return rebuilt.plus(element.getKey() == DATADOG_KEY ? this : element);
   }
 
+  // only needed until the initial context is captured; cleared so a long-lived child doesn't keep
+  // its parent's scope stack and trace alive
+  @Nullable private volatile Context legacySeed;
   @Nullable private AgentSpan parentLocalRootSpan;
+
+  @SuppressFBWarnings(
+      value = "IS2_INCONSISTENT_SYNC",
+      justification =
+          "Written once before publishing volatile context; updateThreadContext reads context before restoration")
+  private boolean legacySharedStack;
+
   private volatile Context context;
   private volatile ContextContinuation continuation;
 
@@ -123,7 +151,9 @@ public CoroutineContext invoke(CoroutineContext rebuilt, CoroutineContext.Elemen
   /** The local root of the trace kept open by this coroutine or its parent. */
   @Nullable private volatile AgentSpan localRootSpan;
 
-  private DatadogThreadContextElement(@Nullable AgentSpan parentLocalRootSpan) {
+  private DatadogThreadContextElement(
+      @Nullable Context legacySeed, @Nullable AgentSpan parentLocalRootSpan) {
+    this.legacySeed = legacySeed;
     this.parentLocalRootSpan = parentLocalRootSpan;
   }
 
@@ -147,6 +177,15 @@ public Key<?> getKey() {
     return DATADOG_KEY;
   }
 
+  /** Lazy children inherit the starting thread's context rather than their creation-time seed. */
+  public static void captureLazyDatadogContext(@Nonnull AbstractCoroutine<?> coroutine) {
+    DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY);
+    if (datadog != null) {
+      datadog.legacySeed = null;
+    }
+    captureDatadogContext(coroutine);
+  }
+
   public static void captureDatadogContext(@Nonnull AbstractCoroutine<?> coroutine) {
     DatadogThreadContextElement datadog = coroutine.getContext().get(DATADOG_KEY);
     if (datadog != null) {
@@ -156,7 +195,9 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine<?> coroutine
       }
       if (datadog.context == null) {
         datadog.captureInitialContext(true);
-      } else if (datadog.owner != coroutine && datadog.isScopedBy(coroutine)) {
+      } else if (!datadog.legacySharedStack
+          && datadog.owner != coroutine
+          && datadog.isScopedBy(coroutine)) {
         // scoped coroutines (coroutineScope, withTimeout, or withContext before 1.6.1) share their
         // caller's element, so refresh it with the active context
         datadog.context = Context.current();
@@ -166,7 +207,17 @@ public static void captureDatadogContext(@Nonnull AbstractCoroutine<?> coroutine
 
   private synchronized void captureInitialContext(boolean owned) {
     if (context == null) {
-      Context initial = Context.current();
+      Context initial = legacySeed;
+      if (initial == null) {
+        initial = Context.current();
+      }
+      if (LEGACY_CONTEXT_MANAGER) {
+        // workers share this coroutine's scope stack, rather than copying it between them
+        Context shared = AgentTracer.get().sharedContextStack(initial);
+        // the noop tracer returns initial unchanged, without creating a shared stack
+        legacySharedStack = shared != initial;
+        initial = shared;
+      }
       // Initialize before the first context installation; later swaps update the stored context.
       AgentSpan rootSpan = localRootOf(initial);
       if (rootSpan != null) {
@@ -184,6 +235,7 @@ private synchronized void captureInitialContext(boolean owned) {
         }
       }
       context = initial;
+      legacySeed = null;
       parentLocalRootSpan = null;
     }
   }
@@ -232,7 +284,13 @@ public Context updateThreadContext(@Nonnull CoroutineContext coroutineContext) {
   @Override
   public void restoreThreadContext(
       @Nonnull CoroutineContext coroutineContext, Context originalContext) {
-    context = originalContext.swap();
+    Context coroutineContextAfter = originalContext.swap();
+    // legacy workers share the coroutine's scope stack, so it never needs saving; writing it back
+    // would usually store that same stack, but after an unbalanced swap on this thread it would
+    // rebind the coroutine to the wrong stack
+    if (!legacySharedStack) {
+      context = coroutineContextAfter;
+    }
   }
 
   /** Tracing state does not change the logical coroutine context used by flow equality checks. */

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

@@ -35,7 +35,7 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class OnStartAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onStart(@Advice.This AbstractCoroutine<?> coroutine) {
-      DatadogThreadContextElement.captureDatadogContext(coroutine);
+      DatadogThreadContextElement.captureLazyDatadogContext(coroutine);
     }
   }
 }

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