Skip to content

Commit 804e4ef

Browse files
kaahoskaahos
committed
Add LockSupport TaskBlock profiling instrumentation
1 parent a446cd6 commit 804e4ef

11 files changed

Lines changed: 1082 additions & 0 deletions

File tree

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/LockSupportHelper.java‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
package datadog.trace.bootstrap.instrumentation.java.concurrent;
2+
3+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
4+
import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
5+
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
6+
import java.util.Collections;
7+
import java.util.Map;
8+
import java.util.WeakHashMap;
9+
10+
/**
11+
* Helper for profiling {@code LockSupport.park*} intervals from bootstrap classes.
12+
*
13+
* <h3>Known limitation: unpark attribution is best effort</h3>
14+
*
15+
* <p>{@link #UNPARKING_SPAN} pairs an {@code unpark(t)} caller's span id with the next {@code
16+
* park*} return on thread {@code t}. {@code LockSupport.park*} is documented to be allowed to
17+
* return spuriously, in which case the parked thread re-parks without ever consuming the map entry.
18+
* A subsequent, unrelated {@code park*} call on the same thread will then drain the stale entry and
19+
* incorrectly attribute the unblocking span to a {@code TaskBlock} it did not cause. Repeated
20+
* {@code unpark(t)} calls before one {@code park*} return are also lossy: LockSupport's permit is
21+
* one-bit and this map stores only the latest caller span, so an earlier causal unpark can be
22+
* overwritten by a later non-causal one.
23+
*
24+
* <p>We accept this residual race because the correct fix (per-park sequence numbers carried
25+
* through {@code ProfilerContext} and matched on entry) is disproportionate to the rarity of these
26+
* edge cases on the JDKs we target, and because the worst-case impact is a single mis-attributed
27+
* {@code TaskBlock} event per occurrence.
28+
*
29+
* <h3>Virtual thread handling</h3>
30+
*
31+
* <p>Virtual threads (JDK 21+) are multiplexed on OS carrier threads. The native {@code
32+
* parkEnter0}/{@code parkExit0} JNI methods use carrier-OS-thread-local storage, so multiple
33+
* virtual threads sharing a carrier overwrite each other's park state. For virtual threads we skip
34+
* the native park entry/exit pair and instead:
35+
*
36+
* <ol>
37+
* <li>Capture the active span/root ids on the Java side at block entry.
38+
* <li>On unpark, call {@link ProfilingContextIntegration#recordTaskBlockWithContext} which passes
39+
* those ids explicitly to the native deferred-capture path, bypassing OTEP TLS. This path is
40+
* span/root-only; custom profiling context attributes are not propagated.
41+
* </ol>
42+
*/
43+
public final class LockSupportHelper {
44+
/**
45+
* Maps a target thread to the span ID of the thread that called {@code unpark()} on it. Keyed by
46+
* the {@link Thread} object (identity equality) so that entries for terminated threads are
47+
* automatically reclaimed by the GC: the JVM holds a strong reference to every live thread
48+
* internally, so a weak key becomes unreachable only after the thread has terminated and been
49+
* collected. {@link WeakHashMap} polls its internal {@link java.lang.ref.ReferenceQueue} on every
50+
* structural operation, so stale entries are cleaned up eagerly without a dedicated thread.
51+
*
52+
* <p>Thread-safety: wrapped with {@link Collections#synchronizedMap}. All access must go through
53+
* the helper methods; direct map manipulation is only permitted in tests.
54+
*/
55+
public static final Map<Thread, Long> UNPARKING_SPAN =
56+
Collections.synchronizedMap(new WeakHashMap<>());
57+
58+
private LockSupportHelper() {}
59+
60+
/** Captured state for a {@code LockSupport.park*} interval. */
61+
public static final class ParkState {
62+
public final ProfilingContextIntegration profiling;
63+
public final long blockerHash;
64+
65+
/** {@code true} when the parking thread is a virtual thread. */
66+
public final boolean isVirtual;
67+
68+
/** TSC tick captured at block entry; only meaningful when {@code isVirtual == true}. */
69+
public final long startTicks;
70+
71+
/** Span ID captured at block entry; only meaningful when {@code isVirtual == true}. */
72+
public final long spanId;
73+
74+
/** Root span ID captured at block entry; only meaningful when {@code isVirtual == true}. */
75+
public final long rootSpanId;
76+
77+
/** Constructor for platform threads: no span context captured here (OTEP TLS used instead). */
78+
public ParkState(ProfilingContextIntegration profiling, long blockerHash) {
79+
this.profiling = profiling;
80+
this.blockerHash = blockerHash;
81+
this.isVirtual = false;
82+
this.startTicks = 0L;
83+
this.spanId = 0L;
84+
this.rootSpanId = 0L;
85+
}
86+
87+
/** Constructor for virtual threads: span/root ids are captured at entry time. */
88+
public ParkState(
89+
ProfilingContextIntegration profiling,
90+
long blockerHash,
91+
long startTicks,
92+
long spanId,
93+
long rootSpanId) {
94+
this.profiling = profiling;
95+
this.blockerHash = blockerHash;
96+
this.isVirtual = true;
97+
this.startTicks = startTicks;
98+
this.spanId = spanId;
99+
this.rootSpanId = rootSpanId;
100+
}
101+
}
102+
103+
public static ParkState captureState(Object blocker) {
104+
return captureState(blocker, AgentTracer.get().getProfilingContext());
105+
}
106+
107+
public static ParkState captureState(Object blocker, ProfilingContextIntegration profiling) {
108+
if (profiling == null) {
109+
return null;
110+
}
111+
long blockerHash = blocker != null ? System.identityHashCode(blocker) : 0L;
112+
if (VirtualThreads.isCurrent()) {
113+
// Virtual thread: skip native parkEnter0 (carrier-scoped TLS is unsafe).
114+
// Capture span/root ids now so we can pass them explicitly on unpark.
115+
ProfilerContext ctx = ProfilerContexts.of(AgentTracer.activeSpan());
116+
if (ctx == null) {
117+
// No active span - nothing to record.
118+
UNPARKING_SPAN.remove(Thread.currentThread());
119+
return null;
120+
}
121+
long startTicks;
122+
try {
123+
startTicks = profiling.getCurrentTicks();
124+
} catch (Throwable ignored) {
125+
UNPARKING_SPAN.remove(Thread.currentThread());
126+
return null;
127+
}
128+
return new ParkState(
129+
profiling, blockerHash, startTicks, ctx.getSpanId(), ctx.getRootSpanId());
130+
}
131+
132+
// Platform thread: skip parkEnter() JNI when no span is active — native would discard
133+
// the interval at parkExit() anyway (zero-span eligibility check). Mirrors the guard
134+
// already present on the virtual thread path above.
135+
ProfilerContext ctx = ProfilerContexts.of(AgentTracer.activeSpan());
136+
if (ctx == null) {
137+
UNPARKING_SPAN.remove(Thread.currentThread());
138+
return null;
139+
}
140+
try {
141+
profiling.parkEnter();
142+
} catch (Throwable ignored) {
143+
// parkEnter failed (e.g. profiler not yet initialised, JNI error); do not track this park.
144+
// Drain any stale unblocking-span entry so it is not mis-attributed to the next park.
145+
UNPARKING_SPAN.remove(Thread.currentThread());
146+
return null;
147+
}
148+
return new ParkState(profiling, blockerHash);
149+
}
150+
151+
public static void finish(ParkState state) {
152+
// Always drain the map entry before any early return. If we returned first, a stale
153+
// unblocking-span ID placed by a prior unpark() would persist and be incorrectly
154+
// attributed to the next TaskBlock event emitted on this thread.
155+
Long unblockingSpanId = UNPARKING_SPAN.remove(Thread.currentThread());
156+
finish(state, unblockingSpanId != null ? unblockingSpanId : 0L);
157+
}
158+
159+
public static void finish(ParkState state, long unblockingSpanId) {
160+
if (state == null) {
161+
return;
162+
}
163+
if (state.isVirtual) {
164+
// Virtual thread: emit via explicit-context path if a span was active at entry.
165+
if (state.spanId != 0L) {
166+
try {
167+
state.profiling.recordTaskBlockWithContext(
168+
state.startTicks,
169+
state.blockerHash,
170+
unblockingSpanId,
171+
state.spanId,
172+
state.rootSpanId);
173+
} catch (Throwable ignored) {
174+
}
175+
}
176+
} else {
177+
// Platform thread: parkExit() clears native parked state and records an eligible TaskBlock
178+
// using the entry tick saved by parkEnter().
179+
state.profiling.parkExit(state.blockerHash, unblockingSpanId);
180+
}
181+
}
182+
183+
public static void recordUnpark(Thread thread) {
184+
if (thread == null) {
185+
return;
186+
}
187+
ProfilerContext ctx = ProfilerContexts.of(AgentTracer.activeSpan());
188+
if (ctx == null) {
189+
return;
190+
}
191+
UNPARKING_SPAN.put(thread, ctx.getSpanId());
192+
}
193+
}

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ProfilerContexts.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
package datadog.trace.bootstrap.instrumentation.java.concurrent;
2+
3+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
4+
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
5+
import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
6+
7+
/** Internal helpers shared by Java-level blocking instrumentation. */
8+
final class ProfilerContexts {
9+
private ProfilerContexts() {}
10+
11+
/**
12+
* Returns the {@link ProfilerContext} for the given span, or {@code null} if the span is absent
13+
* or its context does not implement {@code ProfilerContext}. Profiling-aware instrumentation
14+
* always takes this path before reading span/root-span IDs, so centralising the {@code
15+
* instanceof} check keeps the semantics identical across helpers.
16+
*/
17+
static ProfilerContext of(AgentSpan span) {
18+
if (span == null) {
19+
return null;
20+
}
21+
AgentSpanContext context = span.context();
22+
return context instanceof ProfilerContext ? (ProfilerContext) context : null;
23+
}
24+
}

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/VirtualThreads.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package datadog.trace.bootstrap.instrumentation.java.concurrent;
2+
3+
import java.lang.invoke.MethodHandle;
4+
import java.lang.invoke.MethodHandles;
5+
import java.lang.invoke.MethodType;
6+
7+
/**
8+
* Utility for detecting virtual threads (JDK 21+) in a way that is safe on JDK 8+.
9+
*
10+
* <p>{@link Thread#isVirtual()} was introduced in JDK 21. We access it through a {@link
11+
* MethodHandle} resolved once at class-load time so that on older JDKs the resolution simply fails,
12+
* {@code IS_VIRTUAL} is set to {@code null}, and {@link #isCurrent()} returns {@code false} without
13+
* any bytecode incompatibility.
14+
*/
15+
public final class VirtualThreads {
16+
private static final MethodHandle IS_VIRTUAL;
17+
18+
static {
19+
MethodHandle mh = null;
20+
try {
21+
mh =
22+
MethodHandles.lookup()
23+
.findVirtual(Thread.class, "isVirtual", MethodType.methodType(boolean.class));
24+
} catch (NoSuchMethodException | IllegalAccessException ignored) {
25+
// JDK < 21: Thread.isVirtual() does not exist; isCurrent() will always return false.
26+
}
27+
IS_VIRTUAL = mh;
28+
}
29+
30+
private VirtualThreads() {}
31+
32+
/**
33+
* Returns {@code true} if the current thread is a virtual thread (JDK 21+), {@code false}
34+
* otherwise. Never throws.
35+
*/
36+
public static boolean isCurrent() {
37+
if (IS_VIRTUAL == null) {
38+
return false;
39+
}
40+
try {
41+
return (boolean) IS_VIRTUAL.invokeExact(Thread.currentThread());
42+
} catch (Throwable ignored) {
43+
return false;
44+
}
45+
}
46+
}

‎dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
1 java.util.concurrent.ConcurrentHashMap*
7171
1 java.util.concurrent.atomic.*
7272
1 java.util.concurrent.locks.*
73+
0 java.util.concurrent.locks.LockSupport
7374
0 java.util.logging.*
7475
# allow capturing JVM shutdown
7576
0 java.lang.Shutdown

‎dd-java-agent/instrumentation/datadog/profiling/lock-support/build.gradle‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
apply from: "$rootDir/gradle/java.gradle"
2+
3+
muzzle {
4+
pass {
5+
coreJdk()
6+
}
7+
}
8+
9+
dependencies {
10+
testImplementation libs.bundles.junit5
11+
testImplementation libs.bundles.mockito
12+
}

0 commit comments

Comments
 (0)