Retain per-mount context propagation on JDK 21

Author: amarziali

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

@@ -1,5 +1,7 @@
 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;
@@ -15,10 +17,15 @@
  *
  * <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. So instead the context
- * is 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.
+ * 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>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.
@@ -27,7 +34,9 @@ public final class VirtualThreadState {
   // note: cws is relying on scope listener. This is disabled by default but when enabled
   // let's use the full swap logic since otherwise listeners won't be called
   private static final boolean USE_PER_MOUNT_CONTEXT =
-      !InstrumenterConfig.get().isLegacyContextManagerEnabled() || Config.get().isCwsEnabled();
+      isJavaVersion(21)
+          || !InstrumenterConfig.get().isLegacyContextManagerEnabled()
+          || Config.get().isCwsEnabled();
 
   /** The virtual thread's saved context (scope stack snapshot). */
   private Context context;

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

@@ -41,18 +41,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 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 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>{@code mount()} / {@code unmount()}: rebind and clear carrier-local profiler context. The
- *       state-backed swap path is retained when context listeners require per-mount notification.
+ *       state-backed swap path is retained on JDK 21 and when context listeners require per-mount
+ *       notification.
  *   <li>{@code afterDone()} / {@code afterTerminate()} for early VirtualThread support: cancels the
  *       help continuation, releasing the context scope to be closed.
  * </ol>
  *
- * <p>{@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.
+ * <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
  */
@@ -104,9 +107,11 @@ public Map<String, String> contextStore() {
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
-    transformer.applyAdvice(
-        isMethod().and(named("run")).and(takesArguments(Runnable.class)).and(returns(void.class)),
-        getClass().getName() + "$Run");
+    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(

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

@@ -1,10 +1,10 @@
 package testdog.trace.instrumentation.java.lang.jdk21;
 
+import static datadog.environment.JavaVirtualMachine.isJavaVersion;
 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.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 
 import datadog.trace.agent.test.AbstractInstrumentationTest;
@@ -23,10 +23,10 @@
 public class VirtualThreadLifeCycleTest extends AbstractInstrumentationTest {
   private static final Duration TIMEOUT = Duration.ofSeconds(10);
 
-  @DisplayName("test seed-once path is selected by default")
+  @DisplayName("test context path selected for runtime")
   @Test
-  void testSeedOncePathSelectedByDefault() {
-    assertFalse(VirtualThreadState.usePerMountContext());
+  void testContextPathSelectedForRuntime() {
+    assertEquals(isJavaVersion(21), VirtualThreadState.usePerMountContext());
   }
 
   @DisplayName("test context restored after virtual thread remounts")