Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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

Copy link
Copy Markdown
Contributor

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 🥲

* 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;

Expand All @@ -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());

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Rebind native TLS listeners on virtual-thread remount

When a parked virtual thread remounts on a different carrier after its first mount, this fast path only updates ProfilingContextIntegration and skips the scope-listener activation that context.swap() used to perform. I checked the CWS TlsScopeListener: it is an ExtendedScopeListener that writes the active span into native TLS keyed by getTID(), so CWS correlation on the new carrier remains unset or stale until another scope activation happens.

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) {
Comment thread
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();
Expand Down
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;
Expand Down Expand Up @@ -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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clear profiler context when restoring a root context

When the legacy virtual-thread path unmounts a VT whose previous context is root/no-span (the normal case), VirtualThreadState.onUnmount() now calls this method, but setContext does nothing for non-span contexts. Since ddprof's span context is stored on the carrier OS thread, the virtual thread's span remains installed after unmount and later work on that carrier can be profiled under the stale VT span until another span overwrites it. Please clear the ddprof context when AgentSpan.fromContext(context) is null.

Useful? React with 👍 / 👎.

}

public void clearContext() {
DDPROF.clearSpanContext();
DDPROF.clearContextValue(SPAN_NAME_INDEX);
Expand Down
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
Expand Up @@ -212,6 +212,59 @@ public void run() {
span().childOfPrevious().operationName("child")));
}

@DisplayName("test context preserved across carrier migration")
@Test
void testContextPreservedAcrossCarrierMigration() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❔ question: ‏How is this test different from testContextRestoredAfterVirtualThreadRemount and testConcurrentVirtualThreadsWithRemount?

// 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();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#
💭 thought: ‏Why is there child / parent references here? There is no parent / span children, right?

}));
}
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();
Expand Down
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;
}
}
}
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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) {}
Comment thread
amarziali marked this conversation as resolved.

default Stateful newScopeState(ProfilerContext profilerContext) {
return Stateful.DEFAULT;
}
Expand Down