Skip to content

Commit 094fdd2

Browse files
kaahoskaahos
committed
fix: fix thread.sleep taskblock emission
1 parent f34890d commit 094fdd2

18 files changed

Lines changed: 356 additions & 299 deletions

File tree

‎dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TaskBlockHelperBenchmark.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,6 @@ public void finishEligibleDeferredWithBlockToken(BenchmarkState state) {
4545
public static class BenchmarkState {
4646
private static final long START_TICKS = 42;
4747
private static final long BLOCKER = 7;
48-
private static final long SPAN_ID = 11;
49-
private static final long ROOT_SPAN_ID = 13;
5048
private static final long BLOCK_TOKEN = 17;
5149

5250
private final ProfilingContextIntegration profiling = ProfilingContextIntegration.NoOp.INSTANCE;
@@ -69,8 +67,8 @@ public void setup() {
6967
now - 2 * TaskBlockHelper.MIN_TASK_BLOCK_NANOS,
7068
BLOCKER,
7169
true,
72-
SPAN_ID,
73-
ROOT_SPAN_ID,
70+
0L,
71+
0L,
7472
BLOCK_TOKEN);
7573
}
7674
}

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

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@
55
import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
66
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
77

8-
/** Helper for Java-level instrumentation that emits {@code datadog.TaskBlock} intervals. */
8+
/**
9+
* Helper for Java-level instrumentation that emits {@code datadog.TaskBlock} intervals for untraced
10+
* blocking runs.
11+
*/
912
public final class TaskBlockHelper {
1013
static final long MIN_TASK_BLOCK_NANOS = 1_000_000L;
1114
private static final ThreadLocal<long[]> TASK_BLOCK_SUPPRESSION_SNAPSHOT =
@@ -130,26 +133,20 @@ static State capture(
130133
return null;
131134
}
132135
ProfilerContext context = ProfilerContexts.of(span);
133-
if (context == null) {
136+
if (span != null && context == null) {
137+
return null;
138+
}
139+
if (context != null && context.getSpanId() != 0L) {
134140
return null;
135141
}
136142
long startTicks = profiling.getCurrentTicks();
137143
long startNanos = System.nanoTime();
138144
if (VirtualThreads.isCurrent()) {
139-
return new State(
140-
profiling, startTicks, startNanos, blocker, context.getSpanId(), context.getRootSpanId());
145+
return new State(profiling, startTicks, startNanos, blocker, 0L, 0L);
141146
}
142147
if (deferred) {
143148
long blockToken = profiling.blockEnter(ProfilingContextIntegration.BLOCKING_STATE_SLEEPING);
144-
return new State(
145-
profiling,
146-
startTicks,
147-
startNanos,
148-
blocker,
149-
true,
150-
context.getSpanId(),
151-
context.getRootSpanId(),
152-
blockToken);
149+
return new State(profiling, startTicks, startNanos, blocker, true, 0L, 0L, blockToken);
153150
}
154151
return new State(profiling, startTicks, startNanos, blocker);
155152
}
@@ -164,7 +161,7 @@ public static void finish(final State state) {
164161
return;
165162
}
166163
if (state.deferred) {
167-
if (state.spanId == 0L) {
164+
if (state.spanId != 0L) {
168165
return;
169166
}
170167
long durationNanos = System.nanoTime() - state.startNanos;
@@ -187,8 +184,10 @@ public static void finish(final State state) {
187184
suppressionSnapshot[
188185
ProfilingContextIntegration.TASK_BLOCK_SUPPRESSION_OBSERVED_STATE]);
189186
} else if (state.isVirtual) {
190-
state.profiling.recordTaskBlockWithContext(
191-
state.startTicks, state.blocker, 0L, state.spanId, state.rootSpanId);
187+
if (state.spanId == 0L) {
188+
state.profiling.recordTaskBlockWithContext(
189+
state.startTicks, state.blocker, 0L, state.spanId, state.rootSpanId);
190+
}
192191
} else {
193192
state.profiling.recordTaskBlock(state.startTicks, state.blocker, 0L);
194193
}

‎dd-java-agent/agent-bootstrap/src/test/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TaskBlockHelperTest.java‎

Lines changed: 68 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -44,82 +44,88 @@ void capture_returnsNull_withoutProfilingContext() {
4444
}
4545

4646
@Test
47-
void capture_returnsNull_withoutActiveSpan() {
47+
void capture_recordsEntryTimingWithoutActiveSpan() {
4848
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
49+
when(profiling.getCurrentTicks()).thenReturn(START_TICKS);
50+
51+
long before = System.nanoTime();
52+
TaskBlockHelper.State state = TaskBlockHelper.capture(BLOCKER, profiling, null);
53+
long after = System.nanoTime();
4954

50-
assertNull(TaskBlockHelper.capture(BLOCKER, profiling, null));
55+
assertNotNull(state);
56+
assertEquals(profiling, state.profiling);
57+
assertEquals(START_TICKS, state.startTicks);
58+
assertTrue(state.startNanos >= before, "startNanos should be captured after `before`");
59+
assertTrue(state.startNanos <= after, "startNanos should be captured before `after`");
60+
assertEquals(BLOCKER, state.blocker);
61+
assertEquals(0L, state.spanId);
62+
assertEquals(0L, state.rootSpanId);
5163
}
5264

5365
@Test
5466
void capture_returnsNull_whenSpanContextIsNotProfilerContext() {
5567
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
5668
AgentSpan nonProfilerSpan = mock(AgentSpan.class);
5769
AgentSpanContext nonProfilerCtx = mock(AgentSpanContext.class);
58-
when(nonProfilerSpan.context()).thenReturn(nonProfilerCtx);
70+
when(nonProfilerSpan.spanContext()).thenReturn(nonProfilerCtx);
5971

6072
assertNull(TaskBlockHelper.capture(BLOCKER, profiling, nonProfilerSpan));
6173
}
6274

6375
@Test
64-
void capture_recordsEntryTimingWithoutSpanIds() {
65-
// Span identity is captured natively at the recordTaskBlock JNI boundary; the Java-side
66-
// State only retains the fields the native side cannot recompute (start tick, start nanos,
67-
// blocker). The presence of an active span on the thread is still gated by capture() so we
68-
// skip the JNI hop for span-less intervals.
76+
void capture_returnsNullWithActiveSpan() {
6977
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
7078
AgentSpan span = mock(AgentSpan.class);
7179
ProfilerSpanContext ctx = mock(ProfilerSpanContext.class);
72-
when(span.context()).thenReturn(ctx);
73-
when(profiling.getCurrentTicks()).thenReturn(START_TICKS);
80+
when(span.spanContext()).thenReturn(ctx);
81+
when(ctx.getSpanId()).thenReturn(SPAN_ID);
7482

75-
long before = System.nanoTime();
7683
TaskBlockHelper.State state = TaskBlockHelper.capture(BLOCKER, profiling, span);
77-
long after = System.nanoTime();
7884

79-
assertNotNull(state);
80-
assertEquals(profiling, state.profiling);
81-
assertEquals(START_TICKS, state.startTicks);
82-
assertTrue(state.startNanos >= before, "startNanos should be captured after `before`");
83-
assertTrue(state.startNanos <= after, "startNanos should be captured before `after`");
84-
assertEquals(BLOCKER, state.blocker);
85-
assertEquals(0L, state.spanId);
86-
assertEquals(0L, state.rootSpanId);
85+
assertNull(state);
86+
verify(profiling, never()).getCurrentTicks();
8787
}
8888

8989
@Test
90-
void capture_deferredRecordsEntryTimingAndSpanIds() {
90+
void capture_deferredReturnsNullWithActiveSpan() {
9191
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
9292
AgentSpan span = mock(AgentSpan.class);
9393
ProfilerSpanContext ctx = mock(ProfilerSpanContext.class);
94-
when(span.context()).thenReturn(ctx);
94+
when(span.spanContext()).thenReturn(ctx);
9595
when(ctx.getSpanId()).thenReturn(SPAN_ID);
9696
when(ctx.getRootSpanId()).thenReturn(ROOT_SPAN_ID);
97+
98+
TaskBlockHelper.State state = TaskBlockHelper.capture(BLOCKER, profiling, span, true);
99+
100+
assertNull(state);
101+
verify(profiling, never()).getCurrentTicks();
102+
verify(profiling, never()).blockEnter(anyInt());
103+
}
104+
105+
@Test
106+
void capture_deferredRecordsEntryTimingWithZeroContext() {
107+
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
97108
when(profiling.getCurrentTicks()).thenReturn(START_TICKS);
98109
when(profiling.blockEnter(ProfilingContextIntegration.BLOCKING_STATE_SLEEPING))
99110
.thenReturn(BLOCK_TOKEN);
100111

101-
TaskBlockHelper.State state = TaskBlockHelper.capture(BLOCKER, profiling, span, true);
112+
TaskBlockHelper.State state = TaskBlockHelper.capture(BLOCKER, profiling, null, true);
102113

103114
assertNotNull(state);
104115
assertEquals(START_TICKS, state.startTicks);
105116
assertEquals(BLOCKER, state.blocker);
106117
assertTrue(state.deferred);
107118
assertFalse(state.isVirtual);
108119
assertEquals(BLOCK_TOKEN, state.blockToken);
109-
assertEquals(SPAN_ID, state.spanId);
110-
assertEquals(ROOT_SPAN_ID, state.rootSpanId);
120+
assertEquals(0L, state.spanId);
121+
assertEquals(0L, state.rootSpanId);
111122
verify(profiling).blockEnter(ProfilingContextIntegration.BLOCKING_STATE_SLEEPING);
112123
}
113124

114125
@Test
115-
void capture_deferredSleepOnVirtualThreadUsesVirtualState() throws Exception {
126+
void capture_deferredSleepOnVirtualThreadUsesZeroContextVirtualState() throws Exception {
116127
assumeTrue(hasVirtualThreads(), "virtual threads are not available on this JVM");
117128
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
118-
AgentSpan span = mock(AgentSpan.class);
119-
ProfilerSpanContext ctx = mock(ProfilerSpanContext.class);
120-
when(span.context()).thenReturn(ctx);
121-
when(ctx.getSpanId()).thenReturn(SPAN_ID);
122-
when(ctx.getRootSpanId()).thenReturn(ROOT_SPAN_ID);
123129
when(profiling.getCurrentTicks()).thenReturn(START_TICKS);
124130

125131
AtomicReference<TaskBlockHelper.State> stateRef = new AtomicReference<>();
@@ -128,7 +134,7 @@ void capture_deferredSleepOnVirtualThreadUsesVirtualState() throws Exception {
128134
startVirtualThread(
129135
() -> {
130136
try {
131-
stateRef.set(TaskBlockHelper.capture(BLOCKER, profiling, span, true));
137+
stateRef.set(TaskBlockHelper.capture(BLOCKER, profiling, null, true));
132138
} catch (Throwable error) {
133139
errorRef.set(error);
134140
}
@@ -142,13 +148,13 @@ void capture_deferredSleepOnVirtualThreadUsesVirtualState() throws Exception {
142148
assertNotNull(state);
143149
assertTrue(state.isVirtual);
144150
assertFalse(state.deferred);
145-
assertEquals(SPAN_ID, state.spanId);
146-
assertEquals(ROOT_SPAN_ID, state.rootSpanId);
151+
assertEquals(0L, state.spanId);
152+
assertEquals(0L, state.rootSpanId);
147153

148154
waitUntilEligible(state);
149155
TaskBlockHelper.finish(state);
150156

151-
verify(profiling).recordTaskBlockWithContext(START_TICKS, BLOCKER, 0L, SPAN_ID, ROOT_SPAN_ID);
157+
verify(profiling).recordTaskBlockWithContext(START_TICKS, BLOCKER, 0L, 0L, 0L);
152158
verify(profiling, never())
153159
.enqueueTaskBlock(
154160
anyLong(), anyLong(), anyLong(), anyLong(), anyLong(), anyLong(), anyLong(), anyInt());
@@ -160,7 +166,7 @@ void captureSafely_returnsNullWhenEntryCaptureThrows() {
160166
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
161167
AgentSpan span = mock(AgentSpan.class);
162168
ProfilerSpanContext ctx = mock(ProfilerSpanContext.class);
163-
when(span.context()).thenReturn(ctx);
169+
when(span.spanContext()).thenReturn(ctx);
164170
when(profiling.getCurrentTicks()).thenThrow(new RuntimeException("boom"));
165171

166172
assertNull(TaskBlockHelper.captureSafely(BLOCKER, profiling, span));
@@ -198,8 +204,8 @@ void finish_exitsNativeBlockEvenWhenTooShort() {
198204
System.nanoTime() + 60_000_000_000L,
199205
BLOCKER,
200206
true,
201-
SPAN_ID,
202-
ROOT_SPAN_ID,
207+
0L,
208+
0L,
203209
BLOCK_TOKEN);
204210

205211
TaskBlockHelper.finish(state);
@@ -227,7 +233,24 @@ void finish_emitsTaskBlockForEligibleInterval() {
227233
}
228234

229235
@Test
230-
void finish_emitsVirtualTaskBlockWithSpanRootOnlyContext() {
236+
void finish_emitsVirtualTaskBlockWithZeroContext() {
237+
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
238+
TaskBlockHelper.State state =
239+
new TaskBlockHelper.State(
240+
profiling,
241+
START_TICKS,
242+
System.nanoTime() - 2 * TaskBlockHelper.MIN_TASK_BLOCK_NANOS,
243+
BLOCKER,
244+
0L,
245+
0L);
246+
247+
TaskBlockHelper.finish(state);
248+
249+
verify(profiling).recordTaskBlockWithContext(START_TICKS, BLOCKER, 0L, 0L, 0L);
250+
}
251+
252+
@Test
253+
void finish_skipsVirtualTaskBlockWithActiveSpanContext() {
231254
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
232255
TaskBlockHelper.State state =
233256
new TaskBlockHelper.State(
@@ -240,11 +263,11 @@ void finish_emitsVirtualTaskBlockWithSpanRootOnlyContext() {
240263

241264
TaskBlockHelper.finish(state);
242265

243-
verify(profiling).recordTaskBlockWithContext(START_TICKS, BLOCKER, 0L, SPAN_ID, ROOT_SPAN_ID);
266+
verifyNoInteractions(profiling);
244267
}
245268

246269
@Test
247-
void finish_enqueuesDeferredTaskBlockWithEntrySpanRootOnlyContext() {
270+
void finish_enqueuesDeferredTaskBlockWithZeroContext() {
248271
long anchorSampleId = 101L;
249272
long suppressedSampleCount = 3L;
250273
ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
@@ -268,8 +291,8 @@ void finish_enqueuesDeferredTaskBlockWithEntrySpanRootOnlyContext() {
268291
System.nanoTime() - 2 * TaskBlockHelper.MIN_TASK_BLOCK_NANOS,
269292
BLOCKER,
270293
true,
271-
SPAN_ID,
272-
ROOT_SPAN_ID,
294+
0L,
295+
0L,
273296
BLOCK_TOKEN);
274297

275298
TaskBlockHelper.finish(state);
@@ -279,8 +302,8 @@ void finish_enqueuesDeferredTaskBlockWithEntrySpanRootOnlyContext() {
279302
eq(START_TICKS),
280303
anyLong(),
281304
eq(BLOCKER),
282-
eq(SPAN_ID),
283-
eq(ROOT_SPAN_ID),
305+
eq(0L),
306+
eq(0L),
284307
eq(anchorSampleId),
285308
eq(suppressedSampleCount),
286309
eq(ProfilingContextIntegration.BLOCKING_STATE_SLEEPING));

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -803,6 +803,10 @@ public long getTscFrequency() {
803803
return profiler != null ? taskBlockBridge.getTscFrequency() : 1_000_000_000L;
804804
}
805805

806+
boolean hasTaskBlockEventSupport() {
807+
return profiler != null && taskBlockBridge.hasTaskBlockEventSupport();
808+
}
809+
806810
long blockEnter(int state) {
807811
if (profiler != null && recordingFlag.get()) {
808812
return taskBlockBridge.blockEnter(state);
@@ -895,6 +899,9 @@ void parkExit(long blocker, long unblockingSpanId) {
895899

896900
private static final class TaskBlockBridge {
897901
private final JavaProfiler profiler;
902+
private final Method getCurrentThreadId;
903+
private final Method getTscFrequency;
904+
private final Method recordTaskBlock;
898905
private final Method recordTaskBlockWithContext;
899906
private final Method recordTaskBlockFromContext;
900907
private final Method recordTaskBlockFromContextWithSuppression;
@@ -906,6 +913,10 @@ private static final class TaskBlockBridge {
906913

907914
private TaskBlockBridge(JavaProfiler profiler) {
908915
this.profiler = profiler;
916+
this.getCurrentThreadId = method("getCurrentThreadId");
917+
this.getTscFrequency = method("getTscFrequency");
918+
this.recordTaskBlock =
919+
method("recordTaskBlock", long.class, long.class, long.class, long.class);
909920
this.recordTaskBlockWithContext =
910921
method(
911922
"recordTaskBlockWithContext",
@@ -949,6 +960,25 @@ private boolean hasTaskBlockEventSupport() {
949960
return recordTaskBlockWithContext != null && parkEnter != null && parkExit != null;
950961
}
951962

963+
private int getCurrentThreadId() {
964+
if (getCurrentThreadId == null) {
965+
return -1;
966+
}
967+
return ((Number) invoke(getCurrentThreadId)).intValue();
968+
}
969+
970+
private long getTscFrequency() {
971+
if (getTscFrequency == null) {
972+
return 1_000_000_000L;
973+
}
974+
return ((Number) invoke(getTscFrequency)).longValue();
975+
}
976+
977+
private void recordTaskBlock(
978+
long startTicks, long endTicks, long blocker, long unblockingSpanId) {
979+
invokeIfPresent(recordTaskBlock, startTicks, endTicks, blocker, unblockingSpanId);
980+
}
981+
952982
private void recordTaskBlockWithContext(
953983
long startTicks,
954984
long endTicks,

0 commit comments

Comments
 (0)