Reduce virtual-thread context-propagation overhead on park/unpark

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

@@ -1,15 +1,26 @@
 package datadog.trace.bootstrap.instrumentation.java.lang;
 
 import datadog.context.Context;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope.Continuation;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 
 /**
- * This class holds the saved context and scope continuation for a virtual thread.
+ * Holds the context and continuation for a virtual thread.
  *
- * <p>Used by java-lang-21.0 {@code VirtualThreadInstrumentation} to swap the entire scope stack on
- * mount/unmount.
+ * <p>With the legacy context manager, {@code swap()} rebuilds the whole scope stack on every
+ * mount/unmount, which is costly on the virtual-thread park/unpark hot path. There the context is
+ * seeded once on the first mount — it then follows the thread across park/unpark and carrier
+ * migration via its virtual-thread-aware {@code ThreadLocal} scope stack — and the profiler context
+ * (which is keyed by carrier thread) is re-applied on each subsequent mount and cleared on unmount.
+ *
+ * <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 {
+  private static final boolean LEGACY_CONTEXT_MANAGER =
+      InstrumenterConfig.get().isLegacyContextManagerEnabled();
+
   /** The virtual thread's saved context (scope stack snapshot). */
   private Context context;
 
@@ -19,21 +30,35 @@ public final class VirtualThreadState {
   /** The carrier thread's saved context, set between mount and unmount. */
   private Context previousContext;
 
+  /** Determines whenever the context has been already set for the first mount */
+  private boolean seeded;
+
   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();
+    if (LEGACY_CONTEXT_MANAGER) {
+      if (!seeded) {
+        // First mount also applies the profiler context to the carrier.
+        context.swap();
+        seeded = true;
+        context = null;
+      } else {
+        AgentTracer.get().getProfilingContext().setContext(Context.current());
+      }
+    } else {
+      previousContext = 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;
+    if (LEGACY_CONTEXT_MANAGER) {
+      AgentTracer.get().getProfilingContext().clearContext();
+    } else if (previousContext != null) {
+      context = previousContext.swap();
+      previousContext = null;
     }
   }
 

dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java

@@ -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;
@@ -78,6 +79,15 @@ public String name() {
     return "ddprof";
   }
 
+  @Override
+  public void setContext(Context context) {
+    AgentSpan span = AgentSpan.fromContext(context);
+    if (span != null) {
+      contextManager.activate(span.spanContext());
+    }
+  }
+
+  @Override
   public void clearContext() {
     DDPROF.clearSpanContext();
     DDPROF.clearContextValue(SPAN_NAME_INDEX);

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

@@ -40,10 +40,13 @@
  * <ol>
  *   <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>{@code mount()}: with the legacy context manager, seeds the virtual thread's scope stack
+ *       with the captured context on the first mount only (it then follows the thread via its
+ *       virtual-thread-aware thread-local) and re-applies the profiler context on later mounts;
+ *       with the new context manager, swaps the context in on every mount (cheap, and drives the
+ *       profiler through its context listener).
+ *   <li>{@code unmount()}: with the legacy context manager, clears the carrier's profiler context;
+ *       with the new context manager, swaps the carrier's context back.
  *   <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.

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

@@ -0,0 +1,11 @@
+package testdog.trace.instrumentation.java.lang.jdk21;
+
+import datadog.trace.junit.utils.config.WithConfig;
+
+/**
+ * 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 {}

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

@@ -212,6 +212,59 @@ public void run() {
             span().childOfPrevious().operationName("child")));
   }
 
+  @DisplayName("test context preserved across carrier migration")
+  @Test
+  void testContextPreservedAcrossCarrierMigration() {
+    // Constrain the carrier pool so parked VTs resume on different carriers.
+    String previousParallelism = System.getProperty("jdk.virtualThreadScheduler.parallelism");
+    System.setProperty("jdk.virtualThreadScheduler.parallelism", "2");
+    try {
+      int threadCount = 32;
+      String[] parentSpanId = new String[1];
+      String[] childParentSpanIds = new String[threadCount];
+
+      new Runnable() {
+        @Override
+        @Trace(operationName = "parent")
+        public void run() {
+          parentSpanId[0] = GlobalTracer.get().getSpanId();
+          List<Thread> threads = new ArrayList<>();
+          for (int i = 0; i < threadCount; i++) {
+            int index = i;
+            threads.add(
+                Thread.startVirtualThread(
+                    () -> {
+                      // Multiple park/unpark cycles to provoke carrier migration.
+                      tryUnmount();
+                      childParentSpanIds[index] = GlobalTracer.get().getSpanId();
+                    }));
+          }
+          for (Thread thread : threads) {
+            try {
+              thread.join(TIMEOUT);
+            } catch (InterruptedException e) {
+              throw new RuntimeException(e);
+            }
+          }
+        }
+      }.run();
+
+      for (int i = 0; i < threadCount; i++) {
+        assertEquals(
+            parentSpanId[0],
+            childParentSpanIds[i],
+            "context must survive park/unpark and carrier migration for VT #" + i);
+      }
+      assertTraces(trace(span().root().operationName("parent")));
+    } finally {
+      if (previousParallelism == null) {
+        System.clearProperty("jdk.virtualThreadScheduler.parallelism");
+      } else {
+        System.setProperty("jdk.virtualThreadScheduler.parallelism", previousParallelism);
+      }
+    }
+  }
+
   @Trace(operationName = "child")
   private static void childWork(String[] beforeUnmount, String[] afterRemount) {
     beforeUnmount[0] = GlobalTracer.get().getSpanId();

dd-trace-core/src/jmh/java/datadog/trace/core/VirtualThreadContextBenchmark.java

@@ -0,0 +1,185 @@
+package datadog.trace.core;
+
+import static java.util.concurrent.TimeUnit.MICROSECONDS;
+
+import datadog.context.Context;
+import datadog.trace.api.EndpointTracker;
+import datadog.trace.api.Stateful;
+import datadog.trace.api.profiling.ProfilingContextAttribute;
+import datadog.trace.api.profiling.ProfilingScope;
+import datadog.trace.api.profiling.Timer.TimerType;
+import datadog.trace.api.profiling.Timing;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
+import datadog.trace.core.monitor.HealthMetrics;
+import datadog.trace.core.scopemanager.ContinuableScopeManager;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+
+/**
+ * Per park/unpark cost of virtual-thread context propagation: the current full {@link
+ * ContinuableScopeManager#swap(Context)} on every mount/unmount versus the proposed seed-once path
+ * (nothing when profiling is off, a profiler rebind/unbind when it is on). ddprof's native {@code
+ * setContext} is modelled by a stub {@link ProfilingContextIntegration} whose {@link Stateful}
+ * writes four volatile longs and allocates nothing.
+ *
+ * <pre>
+ * {@code ./gradlew :dd-trace-core:jmh -Pjmh.includes=VirtualThreadContextBenchmark -PtestJvm=21 -Pjmh.profilers=gc}
+ * </pre>
+ *
+ * <p>Sample run (JDK 21.0.9, {@code @Threads(8)}; run-to-run variance is high, the alloc/op figures
+ * are stable):
+ *
+ * <pre>{@code
+ * Benchmark                                                   Mode  Cnt     Score     Error   Units
+ * currentCycle_profilingOff                                  thrpt    5   463.841 ± 251.726  ops/us
+ * currentCycle_profilingOff:gc.alloc.rate.norm               thrpt    5   176.002 ±   0.016    B/op
+ * currentCycle_profilingOn                                   thrpt    5   272.253 ±  21.041  ops/us
+ * currentCycle_profilingOn:gc.alloc.rate.norm                thrpt    5   288.003 ±   0.022    B/op
+ * proposedSteady_profilingOff                                thrpt    5  5010.376 ± 354.716  ops/us
+ * proposedSteady_profilingOff:gc.alloc.rate.norm             thrpt    5    ~0                   B/op
+ * proposedRebindUnbind_profilingOn                           thrpt    5  1755.815 ± 473.954  ops/us
+ * proposedRebindUnbind_profilingOn:gc.alloc.rate.norm        thrpt    5    ~0                   B/op
+ * }</pre>
+ */
+@State(Scope.Benchmark)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@BenchmarkMode(Mode.Throughput)
+@Threads(8)
+@OutputTimeUnit(MICROSECONDS)
+@Fork(value = 1)
+public class VirtualThreadContextBenchmark {
+
+  static final CoreTracer TRACER = CoreTracer.builder().build();
+
+  ContinuableScopeManager plainManager; // profiling off
+  ContinuableScopeManager profiledManager; // profiling on
+
+  @Setup
+  public void setup() {
+    plainManager = new ContinuableScopeManager(0, false);
+    profiledManager =
+        new ContinuableScopeManager(0, false, new StubProfiling(), HealthMetrics.NO_OP);
+  }
+
+  @State(Scope.Thread)
+  public static class ThreadState {
+    AgentSpan span;
+    Context spanContext;
+    Stateful profilerState;
+
+    @Setup(Level.Trial)
+    public void setup(VirtualThreadContextBenchmark bench) {
+      span = TRACER.startSpan("benchmark", "vt");
+      spanContext = span;
+      profilerState = new StubProfiling().newScopeState((ProfilerContext) span.spanContext());
+    }
+
+    @TearDown(Level.Trial)
+    public void tearDown() {
+      span.finish();
+    }
+  }
+
+  @Benchmark
+  public void currentCycle_profilingOff(ThreadState t) {
+    Context previous = plainManager.swap(t.spanContext);
+    plainManager.swap(previous);
+  }
+
+  @Benchmark
+  public void currentCycle_profilingOn(ThreadState t) {
+    Context previous = profiledManager.swap(t.spanContext);
+    profiledManager.swap(previous);
+  }
+
+  // current() is the faithful upper bound for the seed-once steady state (a scope-stack read).
+  @Benchmark
+  public Context proposedSteady_profilingOff(ThreadState t) {
+    return plainManager.current();
+  }
+
+  @Benchmark
+  public Context proposedRebindUnbind_profilingOn(ThreadState t) {
+    Context active = profiledManager.current();
+    t.profilerState.activate(t.span.spanContext());
+    t.profilerState.close();
+    return active;
+  }
+
+  static final class StubProfiling implements ProfilingContextIntegration {
+    @Override
+    public Stateful newScopeState(ProfilerContext profilerContext) {
+      // Per-scope storage mirrors ddprof's per-OS-thread native slots, avoiding the cross-thread
+      // cache contention a single shared instance would introduce.
+      return new StubState();
+    }
+
+    @Override
+    public String name() {
+      return "stub";
+    }
+
+    @Override
+    public ProfilingContextAttribute createContextAttribute(String attribute) {
+      return ProfilingContextAttribute.NoOp.INSTANCE;
+    }
+
+    @Override
+    public ProfilingScope newScope() {
+      return ProfilingScope.NO_OP;
+    }
+
+    @Override
+    public void onRootSpanFinished(AgentSpan rootSpan, EndpointTracker tracker) {}
+
+    @Override
+    public EndpointTracker onRootSpanStarted(AgentSpan rootSpan) {
+      return EndpointTracker.NO_OP;
+    }
+
+    @Override
+    public Timing start(TimerType type) {
+      return Timing.NoOp.INSTANCE;
+    }
+  }
+
+  static final class StubState implements Stateful {
+    volatile long rootSpanId;
+    volatile long spanId;
+    volatile long traceHigh;
+    volatile long traceLow;
+
+    @Override
+    public void activate(Object context) {
+      if (context instanceof ProfilerContext) {
+        ProfilerContext c = (ProfilerContext) context;
+        rootSpanId = c.getRootSpanId();
+        spanId = c.getSpanId();
+        traceHigh = c.getTraceIdHigh();
+        traceLow = c.getTraceIdLow();
+      }
+    }
+
+    @Override
+    public void close() {
+      rootSpanId = 0;
+      spanId = 0;
+      traceHigh = 0;
+      traceLow = 0;
+    }
+  }
+}