feat(virtual-thread): Use context swapping instead of scope activation

Replaces the `ConcurrentState` approach (which activated/closed individual scopes) with `VirtualThreadState` that swaps the entire scope stack via `Context#swap()`. This correctly handles child spans created during virtual thread execution and avoids out-of-order scope closing. The new approach mirrors the ZIO FiberContext and Kotlin coroutines instrumentation patterns.

Author: PerfectSlayer

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadHelper.java

@@ -1,12 +1,13 @@
 package datadog.trace.bootstrap.instrumentation.java.lang;
 
+/** This class is a helper for the java-lang-21.0 {@code VirtualThreadInstrumentation}. */
 public final class VirtualThreadHelper {
   public static final String VIRTUAL_THREAD_CLASS_NAME = "java.lang.VirtualThread";
 
   /**
-   * {@link datadog.trace.bootstrap.instrumentation.api.AgentScope} class name as string literal.
-   * This is mandatory for {@link datadog.trace.bootstrap.ContextStore} API call.
+   * {@link VirtualThreadState} class name as string literal. This is mandatory for {@link
+   * datadog.trace.bootstrap.ContextStore} API call.
    */
-  public static final String AGENT_SCOPE_CLASS_NAME =
-      "datadog.trace.bootstrap.instrumentation.api.AgentScope";
+  public static final String VIRTUAL_THREAD_STATE_CLASS_NAME =
+      "datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState";
 }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java

@@ -0,0 +1,46 @@
+package datadog.trace.bootstrap.instrumentation.java.lang;
+
+import datadog.context.Context;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope.Continuation;
+
+/**
+ * This class holds the saved context and scope continuation for a virtual thread.
+ *
+ * <p>Used by java-lang-21.0 {@code VirtualThreadInstrumentation} to swap the entire scope stack on
+ * mount/unmount.
+ */
+public final class VirtualThreadState {
+  /** The virtual thread's saved context (scope stack snapshot). */
+  private Context context;
+
+  /** Prevents the enclosing context scope from completing before the virtual thread finishes. */
+  private final Continuation continuation;
+
+  /** The carrier thread's saved context, set between mount and unmount. */
+  private Context previousContext;
+
+  public VirtualThreadState(Context context, Continuation continuation) {
+    this.context = context;
+    this.continuation = continuation;
+  }
+
+  /** Called on mount: swaps the virtual thread's context into the carrier thread. */
+  public void onMount() {
+    this.previousContext = this.context.swap();
+  }
+
+  /** Called on unmount: restores the carrier thread's original context. */
+  public void onUnmount() {
+    if (this.previousContext != null) {
+      this.context = this.previousContext.swap();
+      this.previousContext = null;
+    }
+  }
+
+  /** Called on termination: releases the trace continuation. */
+  public void onTerminate() {
+    if (this.continuation != null) {
+      this.continuation.cancel();
+    }
+  }
+}

dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/AbstractInstrumentationTest.java

@@ -42,6 +42,13 @@
  */
 @ExtendWith(TestClassShadowingExtension.class)
 public abstract class AbstractInstrumentationTest {
+  static {
+    // Allow re-registration of ContextManagers so each test can use a fresh tracer.
+    // This mirrors DDSpecification.allowContextTesting() for the JUnit 5 test framework.
+    datadog.context.ContextManager.allowTesting();
+    datadog.context.ContextBinder.allowTesting();
+  }
+
   static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.getInstrumentation();
 
   static final long TIMEOUT_MILLIS = TimeUnit.SECONDS.toMillis(20);

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

@@ -1,69 +1,62 @@
 package datadog.trace.instrumentation.java.lang.jdk21;
 
+import static datadog.context.Context.current;
+import static datadog.context.Context.root;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
-import static datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState.activateAndContinueContinuation;
-import static datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState.captureContinuation;
-import static datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState.closeScope;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
-import static datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadHelper.AGENT_SCOPE_CLASS_NAME;
 import static datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadHelper.VIRTUAL_THREAD_CLASS_NAME;
+import static datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadHelper.VIRTUAL_THREAD_STATE_CLASS_NAME;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
+import datadog.context.Context;
 import datadog.environment.JavaVirtualMachine;
 import datadog.trace.agent.tooling.ExcludeFilterProvider;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
-import datadog.trace.bootstrap.instrumentation.api.AgentScope;
-import datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope.Continuation;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
+import datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.asm.Advice.OnMethodEnter;
 import net.bytebuddy.asm.Advice.OnMethodExit;
 
 /**
- * Instruments {@code VirtualThread} to capture active state at creation, activate it on mount,
- * close the scope on unmount, and cancel the continuation on thread termination.
+ * Instruments {@code VirtualThread} to propagate the context across mount/unmount cycles using
+ * {@link Context#swap()}, and {@link Continuation} to prevent context scope to complete before the
+ * thread finishes.
  *
  * <p>The lifecycle is as follows:
  *
  * <ol>
- *   <li>{@code init()}: captures and holds a continuation from the active context (span due to
- *       legacy API).
- *   <li>{@code mount()}: activates the held continuation, restoring the context on the current
- *       carrier thread.
- *   <li>{@code unmount()}: closes the scope. The continuation survives as still hold.
+ *   <li>{@code init()}: captures the current {@link Context} and an {@link Continuation} to prevent
+ *       the enclosing context scope from completing early.
+ *   <li>{@code mount()}: swaps the virtual thread's saved context into the carrier thread, saving
+ *       the carrier thread's context.
+ *   <li>{@code unmount()}: swaps the carrier thread's original context back, saving the virtual
+ *       thread's (possibly modified) context for the next mount.
  *   <li>Steps 2-3 repeat on each park/unpark cycle, potentially on different carrier threads.
- *   <li>{@code afterDone}: cancels the held continuation to let the context scope to be closed.
+ *   <li>{@code afterDone()}: cancels the help continuation, releasing the context scope to be
+ *       closed.
  * </ol>
  *
- * <p>The instrumentation uses two context stores. The first from {@link Runnable} (as {@code
- * VirtualThread} inherits from {@link Runnable}) to store the captured {@link ConcurrentState} to
- * restore later. It additionally stores the {@link AgentScope} to be able to close it later as
- * activation / close is not done around the same method (so passing the scope from {@link
- * OnMethodEnter} / {@link OnMethodExit} using advice return value is not possible).
- *
- * <p>{@link ConcurrentState} is used instead of {@code State} because virtual threads can mount and
- * unmount multiple times across different carrier threads. The held continuation in {@link
- * ConcurrentState} survives multiple activate/close cycles without being consumed, and is
- * explicitly canceled on thread termination.
- *
  * <p>Instrumenting the internal {@code VirtualThread.runContinuation()} method does not work as the
  * current thread is still the carrier thread and not a virtual thread. Activating the state when on
  * the carrier thread (ie a platform thread) would store the active context into ThreadLocal using
  * the platform thread as key, making the tracer unable to retrieve the stored context from the
  * current virtual thread (ThreadLocal will not return the value associated to the underlying
  * platform thread as they are considered to be different).
+ *
+ * @see VirtualThreadState
  */
 @SuppressWarnings("unused")
 @AutoService(InstrumenterModule.class)
@@ -95,10 +88,7 @@ public boolean isEnabled() {
 
   @Override
   public Map<String, String> contextStore() {
-    Map<String, String> contextStore = new HashMap<>();
-    contextStore.put(Runnable.class.getName(), ConcurrentState.class.getName());
-    contextStore.put(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
-    return contextStore;
+    return singletonMap(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
   }
 
   @Override
@@ -113,47 +103,51 @@ public void methodAdvice(MethodTransformer transformer) {
 
   public static final class Construct {
     @OnMethodExit(suppress = Throwable.class)
-    public static void capture(@Advice.This Object virtualThread) {
-      captureContinuation(
-          InstrumentationContext.get(Runnable.class, ConcurrentState.class),
-          (Runnable) virtualThread,
-          activeSpan());
+    public static void afterInit(@Advice.This Object virtualThread) {
+      Context context = current();
+      if (context == root()) {
+        return; // no active context to propagate
+      }
+      VirtualThreadState state = new VirtualThreadState(context, captureActiveSpan());
+      ContextStore<Object, Object> store =
+          InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
+      store.put(virtualThread, state);
     }
   }
 
   public static final class Mount {
     @OnMethodExit(suppress = Throwable.class)
-    public static void activate(@Advice.This Object virtualThread) {
-      AgentScope scope =
-          activateAndContinueContinuation(
-              InstrumentationContext.get(Runnable.class, ConcurrentState.class),
-              (Runnable) virtualThread);
-      ContextStore<Object, AgentScope> scopeStore =
-          InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
-      scopeStore.put(virtualThread, scope);
+    public static void onMount(@Advice.This Object virtualThread) {
+      ContextStore<Object, VirtualThreadState> store =
+          InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
+      VirtualThreadState state = store.get(virtualThread);
+      if (state != null) {
+        state.onMount();
+      }
     }
   }
 
   public static final class Unmount {
     @OnMethodEnter(suppress = Throwable.class)
-    public static void close(@Advice.This Object virtualThread) {
-      ContextStore<Object, AgentScope> scopeStore =
-          InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, AGENT_SCOPE_CLASS_NAME);
-      AgentScope scope = scopeStore.remove(virtualThread);
-      closeScope(
-          InstrumentationContext.get(Runnable.class, ConcurrentState.class),
-          (Runnable) virtualThread,
-          scope,
-          null);
+    public static void onUnmount(@Advice.This Object virtualThread) {
+      ContextStore<Object, VirtualThreadState> store =
+          InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
+      VirtualThreadState state = store.get(virtualThread);
+      if (state != null) {
+        state.onUnmount();
+      }
     }
   }
 
   public static final class AfterDone {
     @OnMethodEnter(suppress = Throwable.class)
-    public static void clear(@Advice.This Object virtualThread) {
-      ConcurrentState.cancelAndClearContinuation(
-          InstrumentationContext.get(Runnable.class, ConcurrentState.class),
-          (Runnable) virtualThread);
+    public static void onTerminate(@Advice.This Object virtualThread) {
+      ContextStore<Object, VirtualThreadState> store =
+          InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
+      VirtualThreadState state = store.remove(virtualThread);
+      if (state != null) {
+        state.onTerminate();
+      }
     }
   }
 }

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

@@ -13,7 +13,6 @@
 
 /** Test the {@code VirtualThread} and {@code Thread.Builder} API. */
 public class VirtualThreadApiInstrumentationTest extends AbstractInstrumentationTest {
-
   @DisplayName("test Thread.Builder.OfVirtual.start()")
   @Test
   void testBuilderOfVirtualStart() throws InterruptedException, TimeoutException {
@@ -145,4 +144,16 @@ void assertConnectedTrace() {
             span().root().operationName("parent"),
             span().childOfPrevious().operationName("asyncChild")));
   }
+
+  private static void tryUnmount() {
+    try {
+      // Multiple sleeps to attempt triggering repeated park/unpark cycles.
+      // This is not guaranteed to work, but there is no API to force mount/unmount.
+      for (int i = 0; i < 5; i++) {
+        Thread.sleep(10);
+      }
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
+  }
 }

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

@@ -1,8 +1,10 @@
 package testdog.trace.instrumentation.java.lang.jdk21;
 
 import static datadog.trace.agent.test.assertions.SpanMatcher.span;
+import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME;
 import static datadog.trace.agent.test.assertions.TraceMatcher.trace;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
 
 import datadog.trace.agent.test.AbstractInstrumentationTest;
 import datadog.trace.api.CorrelationIdentifier;
@@ -160,6 +162,63 @@ void testNoContextVirtualThreadRemount() throws InterruptedException {
     assertEquals("0", spanIdAfterRemount.get(), "there should be no active context after remount");
   }
 
+  @DisplayName("test context ordering with child span across unmount/remount")
+  @Test
+  void testContextOrderingWithChildSpanAcrossRemount() throws InterruptedException {
+    String[] parentSpanId = new String[1];
+    String[] beforeChild = new String[1];
+    String[] insideChildBeforeUnmount = new String[1];
+    String[] insideChildAfterRemount = new String[1];
+    String[] afterChild = new String[1];
+
+    new Runnable() {
+      @Override
+      @Trace(operationName = "parent")
+      public void run() {
+        parentSpanId[0] = GlobalTracer.get().getSpanId();
+
+        Thread thread =
+            Thread.startVirtualThread(
+                () -> {
+                  beforeChild[0] = GlobalTracer.get().getSpanId();
+                  childWork(insideChildBeforeUnmount, insideChildAfterRemount);
+                  afterChild[0] = GlobalTracer.get().getSpanId();
+                });
+        try {
+          thread.join();
+        } catch (InterruptedException e) {
+          throw new RuntimeException(e);
+        }
+        blockUntilChildSpansFinished(1);
+      }
+    }.run();
+
+    // Verify context ordering at each checkpoint
+    assertEquals(parentSpanId[0], beforeChild[0], "parent should be active before child span");
+    assertNotEquals("0", insideChildBeforeUnmount[0], "child should be active before unmount");
+    assertNotEquals(
+        parentSpanId[0], insideChildBeforeUnmount[0], "active span should be child, not parent");
+    assertEquals(
+        insideChildBeforeUnmount[0],
+        insideChildAfterRemount[0],
+        "child should still be active after remount (no out-of-order scope close)");
+    assertEquals(parentSpanId[0], afterChild[0], "parent should be active after child span closes");
+
+    // Verify trace structure
+    assertTraces(
+        trace(
+            SORT_BY_START_TIME,
+            span().root().operationName("parent"),
+            span().childOfPrevious().operationName("child")));
+  }
+
+  @Trace(operationName = "child")
+  private static void childWork(String[] beforeUnmount, String[] afterRemount) {
+    beforeUnmount[0] = GlobalTracer.get().getSpanId();
+    tryUnmount();
+    afterRemount[0] = GlobalTracer.get().getSpanId();
+  }
+
   private static void tryUnmount() {
     try {
       // Multiple sleeps to expect triggering repeated park/unpark cycles.