Used by java-lang-21.0 {@code VirtualThreadInstrumentation} to swap the entire scope stack on - * mount/unmount. + *
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 on the first mount, then follows the thread across park/unpark and carrier + * migration via its virtual-thread-aware {@code ThreadLocal} scope stack, while the profiler + * context (which is keyed by carrier thread) is re-applied on each subsequent mount and restored on + * unmount. + * + *
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;
@@ -24,20 +37,31 @@ public VirtualThreadState(Context context, Continuation continuation) {
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 (context != null) {
+ // First mount also applies the profiler context to the carrier.
+ previousContext = context.swap();
+ 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 (previousContext != null) {
+ if (LEGACY_CONTEXT_MANAGER) {
+ AgentTracer.get().getProfilingContext().setContext(previousContext);
+ } else {
+ context = previousContext.swap();
+ previousContext = null;
+ }
}
}
- /** Called on termination: releases the trace continuation. */
public void onTerminate() {
if (this.continuation != null) {
this.continuation.cancel();
diff --git a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java b/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java
index 65462a46bce..0d9afc3f7df 100644
--- a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java
+++ b/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,14 @@ public String name() {
return "ddprof";
}
+ @Override
+ public void setContext(Context context) {
+ AgentSpan span = AgentSpan.fromContext(context);
+ if (span != null) {
+ contextManager.activate(span.spanContext());
+ }
+ }
+
public void clearContext() {
DDPROF.clearSpanContext();
DDPROF.clearContextValue(SPAN_NAME_INDEX);
diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadApiInstrumentationForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadApiInstrumentationForkedTest.java
new file mode 100644
index 00000000000..bedb14f10b3
--- /dev/null
+++ b/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 {}
diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadLifeCycleTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadLifeCycleTest.java
index 54690e14f73..22e2be954ec 100644
--- a/dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadLifeCycleTest.java
+++ b/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 Sample run (JDK 21.0.9, {@code @Threads(8)}; run-to-run variance is high, the alloc/op figures
+ * are stable):
+ *
+ *
+ * {@code ./gradlew :dd-trace-core:jmh -Pjmh.includes=VirtualThreadContextBenchmark -PtestJvm=21 -Pjmh.profilers=gc}
+ *
+ *
+ * {@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
+ * }
+ */
+@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;
+ }
+ }
+}
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ProfilingContextIntegration.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ProfilingContextIntegration.java
index 4accced983a..f2a74fee8ce 100644
--- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ProfilingContextIntegration.java
+++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ProfilingContextIntegration.java
@@ -1,5 +1,6 @@
package datadog.trace.bootstrap.instrumentation.api;
+import datadog.context.Context;
import datadog.trace.api.EndpointCheckpointer;
import datadog.trace.api.EndpointTracker;
import datadog.trace.api.Stateful;
@@ -18,6 +19,14 @@ default void onAttach() {}
/** Invoked when a thread exits */
default void onDetach() {}
+ /**
+ * Applies {@code context} to the current thread's profiler context. Default is a no-op: only
+ * integrations that key profiler context by the running (carrier) thread need this, and only when
+ * driven by the legacy context manager, where the virtual-thread instrumentation seeds the scope
+ * stack once and calls this on each subsequent mount rather than swapping.
+ */
+ default void setContext(Context context) {}
+
default Stateful newScopeState(ProfilerContext profilerContext) {
return Stateful.DEFAULT;
}