-
Notifications
You must be signed in to change notification settings - Fork 344
Reduce virtual-thread context-propagation overhead on park/unpark #11893
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
e22b0c6
b05f030
acd37e6
b48e027
480efcf
498ed03
539cb2a
430567e
05f68b5
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,15 +1,31 @@ | ||
| package datadog.trace.bootstrap.instrumentation.java.lang; | ||
|
|
||
| import datadog.context.Context; | ||
| import datadog.trace.api.Config; | ||
| 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>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. | ||
| * | ||
| * <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 { | ||
| // 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_SIMPLE_SWAP = | ||
| !InstrumenterConfig.get().isLegacyContextManagerEnabled() || Config.get().isCwsEnabled(); | ||
|
|
||
| /** The virtual thread's saved context (scope stack snapshot). */ | ||
| private Context context; | ||
|
|
||
|
|
@@ -24,20 +40,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 (USE_SIMPLE_SWAP) { | ||
| previousContext = context.swap(); | ||
| } else { | ||
| 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()); | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When a parked virtual thread remounts on a different carrier after its first mount, this fast path only updates Useful? React with 👍 / 👎. |
||
| } | ||
| } | ||
| } | ||
|
|
||
| /** 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) { | ||
|
jbachorik marked this conversation as resolved.
|
||
| if (USE_SIMPLE_SWAP) { | ||
| context = previousContext.swap(); | ||
| previousContext = null; | ||
| } else { | ||
| AgentTracer.get().getProfilingContext().setContext(previousContext); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /** Called on termination: releases the trace continuation. */ | ||
| public void onTerminate() { | ||
| if (this.continuation != null) { | ||
| this.continuation.cancel(); | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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()); | ||
| } | ||
|
Comment on lines
+83
to
+87
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When the legacy virtual-thread path unmounts a VT whose previous context is root/no-span (the normal case), Useful? React with 👍 / 👎. |
||
| } | ||
|
|
||
| public void clearContext() { | ||
| DDPROF.clearSpanContext(); | ||
| DDPROF.clearContextValue(SPAN_NAME_INDEX); | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 {} |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -212,6 +212,59 @@ public void run() { | |
| span().childOfPrevious().operationName("child"))); | ||
| } | ||
|
|
||
| @DisplayName("test context preserved across carrier migration") | ||
| @Test | ||
| void testContextPreservedAcrossCarrierMigration() { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ❔ question: How is this test different from |
||
| // 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(); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. # |
||
| })); | ||
| } | ||
| 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(); | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🔍 nitpick: That was useful to me to keep things organized compared to the other
State🥲