gh-worker-dd-mergequeue-cf854d · Sep 23, 2026 · Jul 9, 2026 · Jul 9, 2026 · Jul 9, 2026 · Jul 10, 2026
@@ -1,15 +1,42 @@
 package datadog.trace.bootstrap.instrumentation.java.lang;
 
+import static datadog.environment.JavaVirtualMachine.isJavaVersion;
+
 import datadog.context.Context;
 import datadog.context.ContextContinuation;
+import datadog.trace.api.Config;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 
 /**
- * This class holds the saved context and scope continuation for a virtual thread.
+ * Holds the context and continuation for a virtual thread.
+ *
+ * <p>This state belongs to the Java 21 virtual-thread instrumentation and is kept in a context
+ * store because {@code java.lang.VirtualThread} is loaded before the agent can inject fields.
+ *
+ * <p>The legacy context manager's {@code swap()} wraps the current scope stack together with the
+ * context so the original stack can be restored when the context is swapped back; doing that on
+ * every mount/unmount is costly on the virtual-thread park/unpark hot path. On JDK 22 and later the
+ * context is instead seeded once when the virtual thread starts running, then follows the thread
+ * across park/unpark and carrier migration via its virtual-thread-aware {@code ThreadLocal} scope
+ * stack. The profiler context, which is keyed by carrier thread, is rebound on mount and cleared on
+ * unmount.
  *
- * <p>Used by java-lang-21.0 {@code VirtualThreadInstrumentation} to swap the entire scope stack on
- * mount/unmount.
+ * <p>JDK 21 retains the per-mount path because early update releases enter {@code run(Runnable)} on
+ * the carrier thread before the first mount. Using one implementation for all JDK 21 updates avoids
+ * relying on the internal lifecycle change introduced in later updates.
+ *
+ * <p>With the new context manager {@code swap()} is cheap and drives the profiler through its
+ * context listener, so we simply swap in on mount and out on unmount.
  */
 public final class VirtualThreadState {
+  // CWS observes scope changes through listeners, so retain context swaps on mount and unmount.
+  private static final boolean USE_PER_MOUNT_CONTEXT =
+      isJavaVersion(21)
+          || !InstrumenterConfig.get().isLegacyContextManagerEnabled()
+          || Config.get().isCwsEnabled();
+
   /** The virtual thread's saved context (scope stack snapshot). */
   private Context context;
 
@@ -24,20 +51,55 @@ public VirtualThreadState(Context context, ContextContinuation continuation) {
     this.continuation = continuation;
   }
 
-  /** Called on mount: swaps the virtual thread's context into the carrier thread. */
+  /** Whether context propagation must retain the state-backed mount/unmount path. */
+  public static boolean usePerMountContext() {
+    return USE_PER_MOUNT_CONTEXT;
+  }
+
+  /** Seeds context before the JDK invokes the virtual thread's task on JDK 22 and later. */
+  public void onRun() {
+    previousContext = context.swap();
+    context = null;
+  }
+
+  /** Restores the preceding context after the virtual thread's task exits. */
+  public void afterRun() {
+    if (previousContext != null) {
+      previousContext.swap();
+      previousContext = null;
+    }
+  }
+
+  /** Rebinds carrier-local profiler state after the virtual thread mounts. */
+  public static void onMountWithoutStore() {
+    ProfilingContextIntegration profilingContext = AgentTracer.get().getProfilingContext();
+    if (profilingContext.isThreadContextBindingRequired()) {
+      profilingContext.setContext(Context.current());
+    }
+  }
+
+  /** Clears carrier-local profiler state before the virtual thread unmounts. */
+  public static void onUnmountWithoutStore() {
+    ProfilingContextIntegration profilingContext = AgentTracer.get().getProfilingContext();
+    if (profilingContext.isThreadContextBindingRequired()) {
+      profilingContext.setContext(Context.root());
+    }
+  }
+
+  /** Activates the virtual thread's context after it mounts on the state-backed path. */
   public void onMount() {
-    this.previousContext = this.context.swap();
+    previousContext = context.swap();
   }
 
-  /** Called on unmount: restores the carrier thread's original context. */
+  /** Restores the preceding context before the virtual thread unmounts. */
   public void onUnmount() {
-    if (this.previousContext != null) {
-      this.context = this.previousContext.swap();
-      this.previousContext = null;
+    if (previousContext != null) {
+      context = previousContext.swap();
+      previousContext = null;
     }
   }
 
-  /** Called on termination: releases the trace continuation. */
+  /** Releases the retained continuation as virtual-thread termination begins. */
   public void onTerminate() {
     if (this.continuation != null) {
       this.continuation.release();

@@ -38,6 +38,8 @@ dependencies {
 
   testImplementation libs.bundles.jmc
   testImplementation libs.bundles.junit5
+  testImplementation libs.bundles.mockito
+  testImplementation group: 'org.mockito', name: 'mockito-inline', version: '4.11.0'
 }
 
 

@@ -1,5 +1,6 @@
 package com.datadog.profiling.ddprof;
 
+import datadog.context.Context;
 import datadog.trace.api.EndpointTracker;
 import datadog.trace.api.Stateful;
 import datadog.trace.api.profiling.ProfilingContextAttribute;
@@ -80,8 +81,24 @@ public String name() {
     return "ddprof";
   }
 
-  public void clearContext() {
-    DDPROF.clearTraceContext();
+  /** Rebinds ddprof's carrier-thread context to the span contained in {@code context}. */
+  @Override
+  public void setContext(Context context) {
+    AgentSpan span = AgentSpan.fromContext(context);
+    if (span != null) {
+      contextManager.activate(span.spanContext());
+    } else {
+      clearContext();
+    }
+  }
+
+  @Override
+  public boolean isThreadContextBindingRequired() {
+    return true;
+  }
+
+  private void clearContext() {
+    contextManager.close();
   }
 
   @Override

@@ -0,0 +1,51 @@
+package com.datadog.profiling.ddprof;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Answers.CALLS_REAL_METHODS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
+
+import datadog.context.Context;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
+import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+
+class DatadogProfilingIntegrationTest {
+  @Test
+  void rebindsSpanContextAndClearsRootContext() {
+    DatadogProfiler profiler = mock(DatadogProfiler.class);
+    when(profiler.operationNameOffset()).thenReturn(5);
+    when(profiler.resourceNameOffset()).thenReturn(6);
+
+    try (MockedStatic<DatadogProfiler> profilerFactory = mockStatic(DatadogProfiler.class)) {
+      profilerFactory.when(DatadogProfiler::newInstance).thenReturn(profiler);
+
+      AgentSpanContext spanContext =
+          mock(AgentSpanContext.class, withSettings().extraInterfaces(ProfilerContext.class));
+      ProfilerContext profilerContext = (ProfilerContext) spanContext;
+      when(profilerContext.getRootSpanId()).thenReturn(11L);
+      when(profilerContext.getSpanId()).thenReturn(22L);
+      when(profilerContext.getTraceIdHigh()).thenReturn(33L);
+      when(profilerContext.getTraceIdLow()).thenReturn(44L);
+      when(profilerContext.getOperationName()).thenReturn("operation");
+      when(profilerContext.getResourceName()).thenReturn("resource");
+
+      AgentSpan span = mock(AgentSpan.class, CALLS_REAL_METHODS);
+      when(span.spanContext()).thenReturn(spanContext);
+
+      DatadogProfilingIntegration integration = new DatadogProfilingIntegration();
+      assertTrue(integration.isThreadContextBindingRequired());
+
+      integration.setContext(span);
+      verify(profiler).setTraceContext(11L, 22L, 33L, 44L, 5, "operation", 6, "resource");
+
+      integration.setContext(Context.root());
+      verify(profiler).clearTraceContext();
+    }
+  }
+}
@@ -10,6 +10,7 @@
 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.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
@@ -39,21 +40,21 @@
  * <ol>
  *   <li>{@code init()}: captures the current {@link Context} and an {@link ContextContinuation} 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()} / {@code afterTerminate()} for early VirtualThread support: cancels the
- *       help continuation, releasing the context scope to be closed.
+ *   <li>On JDK 22 and later, {@code run(Runnable)} seeds the virtual thread's saved context once
+ *       and retains its state in the continuation frame across park/unpark cycles.
+ *   <li>On JDK 21, or when context listeners require per-mount notification, {@code mount()} and
+ *       {@code unmount()} retain state-backed context swaps. Otherwise they only update
+ *       carrier-local profiler context when required.
+ *   <li>{@code afterDone()} / {@code afterTerminate()} releases the retained continuation so the
+ *       enclosing context scope can complete.
  * </ol>
  *
- * <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).
+ * <p>On JDK 22 and later, {@code run(Runnable)} is the one-shot continuation body and executes
+ * after the first mount, when the current thread is the virtual thread. Its advice local is
+ * preserved with the continuation across yields. {@code runContinuation()} cannot serve this
+ * purpose because it executes once per mount and starts while the carrier is still the current
+ * thread. JDK 21 retains the per-mount path because this internal ordering differs across its
+ * update releases.
  *
  * @see VirtualThreadState
  */
@@ -93,7 +94,7 @@ public boolean isEnabled() {
 
   @Override
   public Map<ExcludeFilter.ExcludeType, ? extends Collection<String>> excludedClasses() {
-    // VirtualThread context is activated on mount/unmount, not on Runnable.run().
+    // VirtualThread context is managed directly across its internal lifecycle.
     return singletonMap(RUNNABLE, singletonList(VIRTUAL_THREAD_CLASS_NAME));
   }
 
@@ -105,6 +106,13 @@ public Map<String, String> contextStore() {
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
+    // JDK 22+ enters run(Runnable) after the first mount with the virtual thread current. JDK 21
+    // update releases differ, so they retain context swaps on every mount and unmount.
+    if (JavaVirtualMachine.isJavaVersionAtLeast(22)) {
+      transformer.applyAdvice(
+          isMethod().and(named("run")).and(takesArguments(Runnable.class)).and(returns(void.class)),
+          getClass().getName() + "$Run");
+    }
     transformer.applyAdvice(isMethod().and(named("mount")), getClass().getName() + "$Mount");
     transformer.applyAdvice(isMethod().and(named("unmount")), getClass().getName() + "$Unmount");
     transformer.applyAdvice(
@@ -129,26 +137,57 @@ public static void afterInit(@Advice.This Object virtualThread) {
     }
   }
 
+  public static final class Run {
+    @OnMethodEnter(suppress = Throwable.class)
+    public static void onRun(
+        @Advice.This Object virtualThread,
+        @Advice.Local("virtualThreadState") VirtualThreadState state) {
+      if (!VirtualThreadState.usePerMountContext()) {
+        ContextStore<Object, VirtualThreadState> store =
+            InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
+        state = store.get(virtualThread);
+        if (state != null) {
+          state.onRun();
+        }
+      }
+    }
+
+    @OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void afterRun(@Advice.Local("virtualThreadState") VirtualThreadState state) {
+      if (state != null) {
+        state.afterRun();
+      }
+    }
+  }
+
   public static final class Mount {
     @OnMethodExit(suppress = Throwable.class)
     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();
+      if (VirtualThreadState.usePerMountContext()) {
+        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();
+        }
+      } else {
+        VirtualThreadState.onMountWithoutStore();
       }
     }
   }
 
   public static final class Unmount {
     @OnMethodEnter(suppress = Throwable.class)
     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();
+      if (VirtualThreadState.usePerMountContext()) {
+        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();
+        }
+      } else {
+        VirtualThreadState.onUnmountWithoutStore();
       }
     }
   }

@@ -0,0 +1,22 @@
+package testdog.trace.instrumentation.java.lang.jdk21;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState;
+import datadog.trace.test.junit.utils.config.WithConfig;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Runs the {@link VirtualThreadApiInstrumentationTest} cases with the legacy context manager
+ * disabled, so {@code VirtualThreadState} takes the swap path instead of seed-once. Forked because
+ * the legacy-context-manager choice is captured once per JVM.
+ */
+@WithConfig(key = "legacy.context-manager.enabled", value = "false")
+class VirtualThreadApiInstrumentationForkedTest extends VirtualThreadApiInstrumentationTest {
+  @DisplayName("test per-mount path is selected for the new context manager")
+  @Test
+  void testPerMountPathSelected() {
+    assertTrue(VirtualThreadState.usePerMountContext());
+  }
+}
@@ -0,0 +1,18 @@
+package testdog.trace.instrumentation.java.lang.jdk21;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import datadog.trace.bootstrap.instrumentation.java.lang.VirtualThreadState;
+import datadog.trace.test.junit.utils.config.WithConfig;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+/** Verifies that CWS retains scope-listener notifications on every virtual-thread mount. */
+@WithConfig(key = "cws.enabled", value = "true")
+class VirtualThreadCwsInstrumentationForkedTest extends VirtualThreadApiInstrumentationTest {
+  @DisplayName("test per-mount path is selected for CWS")
+  @Test
+  void testPerMountPathSelected() {
+    assertTrue(VirtualThreadState.usePerMountContext());
+  }
+}