Add Thread.sleep TaskBlock profiling instrumentation

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

@@ -0,0 +1,47 @@
+// Copyright 2026 Datadog, Inc.
+package datadog.trace.bootstrap.instrumentation.java.concurrent;
+
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+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.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@Fork(1)
+public class TaskBlockHelperBenchmark {
+
+  /** Measures the rejected-entry fast path. */
+  @Benchmark
+  public TaskBlockHelper.State captureRejected(BenchmarkState state) {
+    return TaskBlockHelper.captureForSleep(state.rejected);
+  }
+
+  /** Measures an accepted entry paired with synchronous completion. */
+  @Benchmark
+  public void captureAndFinishAccepted(BenchmarkState state) {
+    TaskBlockHelper.finish(state.acceptedState);
+  }
+
+  /** Measures the null-state completion fast path. */
+  @Benchmark
+  public void finishNull() {
+    TaskBlockHelper.finish(null);
+  }
+
+  @State(Scope.Thread)
+  public static class BenchmarkState {
+    final ProfilingContextIntegration rejected = ProfilingContextIntegration.NoOp.INSTANCE;
+    final TaskBlockHelper.State acceptedState =
+        new TaskBlockHelper.State(ProfilingContextIntegration.NoOp.INSTANCE, 17L);
+  }
+}

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

@@ -0,0 +1,88 @@
+// Copyright 2026 Datadog, Inc.
+package datadog.trace.bootstrap.instrumentation.java.concurrent;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+
+/** Helper for profiling {@code LockSupport.park*} intervals from bootstrap instrumentation. */
+public final class LockSupportHelper {
+  /**
+   * Best-effort association between a parked thread and the most recent {@code unpark} caller's
+   * active span. Weak keys avoid retaining terminated threads; the synchronized wrapper protects
+   * calls made concurrently by parked and unparking threads.
+   */
+  static final Map<Thread, Long> UNPARKING_SPAN = Collections.synchronizedMap(new WeakHashMap<>());
+
+  private LockSupportHelper() {}
+
+  /** State required to balance a park entry accepted by a profiling integration. */
+  public static final class ParkState {
+    private final ProfilingContextIntegration profiling;
+    private final long blockerHash;
+
+    ParkState(ProfilingContextIntegration profiling, long blockerHash) {
+      this.profiling = profiling;
+      this.blockerHash = blockerHash;
+    }
+  }
+
+  /** Captures a park entry through the currently installed profiling integration. */
+  public static ParkState captureState(Object blocker) {
+    return captureState(blocker, AgentTracer.get().getProfilingContext());
+  }
+
+  static ParkState captureState(Object blocker, ProfilingContextIntegration profiling) {
+    if (profiling == null) {
+      return null;
+    }
+    try {
+      if (!profiling.parkEnter()) {
+        return null;
+      }
+      return new ParkState(
+          profiling,
+          blocker == null ? 0L : Integer.toUnsignedLong(System.identityHashCode(blocker)));
+    } catch (Throwable ignored) {
+      return null;
+    }
+  }
+
+  /** Drains unpark attribution and balances an accepted park entry. */
+  public static void finish(ParkState state) {
+    Long unblockingSpanId = UNPARKING_SPAN.remove(Thread.currentThread());
+    finish(state, unblockingSpanId == null ? 0L : unblockingSpanId);
+  }
+
+  static void finish(ParkState state, long unblockingSpanId) {
+    if (state == null) {
+      return;
+    }
+    try {
+      state.profiling.parkExit(state.blockerHash, unblockingSpanId);
+    } catch (Throwable ignored) {
+    }
+  }
+
+  /**
+   * Records the latest unpark caller for {@code thread}. An untraced call explicitly clears an
+   * older traced caller so the association follows last-writer semantics.
+   */
+  public static void recordUnpark(Thread thread) {
+    if (thread == null) {
+      return;
+    }
+    AgentSpan span = AgentTracer.activeSpan();
+    AgentSpanContext context = span == null ? null : span.spanContext();
+    if (context instanceof ProfilerContext) {
+      UNPARKING_SPAN.put(thread, ((ProfilerContext) context).getSpanId());
+    } else {
+      UNPARKING_SPAN.remove(thread);
+    }
+  }
+}

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

@@ -0,0 +1,53 @@
+// Copyright 2026 Datadog, Inc.
+package datadog.trace.bootstrap.instrumentation.java.concurrent;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
+
+/** Helper for synchronously bracketing untraced {@code Thread.sleep} intervals. */
+public final class TaskBlockHelper {
+  private TaskBlockHelper() {}
+
+  /** State required to balance a native TaskBlock interval accepted at sleep entry. */
+  public static final class State {
+    final ProfilingContextIntegration profiling;
+    final long token;
+
+    State(ProfilingContextIntegration profiling, long token) {
+      this.profiling = profiling;
+      this.token = token;
+    }
+  }
+
+  /** Starts a synchronous sleeping interval through the currently installed integration. */
+  public static State captureForSleep() {
+    try {
+      return captureForSleep(AgentTracer.get().getProfilingContext());
+    } catch (Throwable ignored) {
+      return null;
+    }
+  }
+
+  static State captureForSleep(ProfilingContextIntegration profiling) {
+    try {
+      if (profiling == null) {
+        return null;
+      }
+      long token = profiling.beginTaskBlock(ProfilingContextIntegration.BLOCKING_STATE_SLEEPING);
+      return token == 0L ? null : new State(profiling, token);
+    } catch (Throwable ignored) {
+      return null;
+    }
+  }
+
+  /** Completes a sleeping interval accepted by {@link #captureForSleep()}. */
+  public static void finish(State state) {
+    if (state == null) {
+      return;
+    }
+    try {
+      state.profiling.endTaskBlock(state.token, 0L, 0L);
+    } catch (Throwable ignored) {
+    }
+  }
+}

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

@@ -0,0 +1,167 @@
+// Copyright 2026 Datadog, Inc.
+package datadog.trace.bootstrap.instrumentation.java.concurrent;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
+import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class LockSupportHelperTest {
+  private interface ProfilerSpanContext extends AgentSpanContext, ProfilerContext {}
+
+  private AgentTracer.TracerAPI previousTracer;
+
+  @BeforeEach
+  void setUp() {
+    previousTracer = AgentTracer.get();
+    AgentTracer.forceRegister(tracerWithActiveSpan(null));
+    LockSupportHelper.UNPARKING_SPAN.clear();
+  }
+
+  @AfterEach
+  void tearDown() {
+    LockSupportHelper.UNPARKING_SPAN.clear();
+    AgentTracer.forceRegister(previousTracer);
+  }
+
+  @Test
+  void nullIntegrationDoesNotAcceptEntry() {
+    assertNull(LockSupportHelper.captureState(new Object(), null));
+  }
+
+  @Test
+  void rejectedEntryDoesNotRequireExit() {
+    ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
+    when(profiling.parkEnter()).thenReturn(false);
+
+    LockSupportHelper.ParkState state = LockSupportHelper.captureState(new Object(), profiling);
+    LockSupportHelper.finish(state, 17L);
+
+    assertNull(state);
+    verify(profiling).parkEnter();
+    verify(profiling, never()).parkExit(0L, 17L);
+  }
+
+  @Test
+  void acceptedEntryIsAlwaysPairedWithExit() {
+    ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
+    when(profiling.parkEnter()).thenReturn(true);
+    Object blocker = new Object();
+
+    LockSupportHelper.ParkState state = LockSupportHelper.captureState(blocker, profiling);
+    LockSupportHelper.finish(state, 23L);
+
+    assertNotNull(state);
+    verify(profiling).parkEnter();
+    verify(profiling).parkExit(Integer.toUnsignedLong(System.identityHashCode(blocker)), 23L);
+  }
+
+  @Test
+  void nullBlockerUsesZeroIdentity() {
+    ProfilingContextIntegration profiling = mock(ProfilingContextIntegration.class);
+    when(profiling.parkEnter()).thenReturn(true);
+
+    LockSupportHelper.finish(LockSupportHelper.captureState(null, profiling), 0L);
+
+    verify(profiling).parkExit(0L, 0L);
+  }
+
+  @Test
+  void entryAndExitFailuresDoNotEscapeInstrumentation() {
+    ProfilingContextIntegration entryFailure = mock(ProfilingContextIntegration.class);
+    doThrow(new IllegalStateException("enter")).when(entryFailure).parkEnter();
+    assertNull(LockSupportHelper.captureState(null, entryFailure));
+
+    ProfilingContextIntegration exitFailure = mock(ProfilingContextIntegration.class);
+    when(exitFailure.parkEnter()).thenReturn(true);
+    doThrow(new IllegalStateException("exit")).when(exitFailure).parkExit(0L, 0L);
+    LockSupportHelper.ParkState state = LockSupportHelper.captureState(null, exitFailure);
+
+    assertDoesNotThrow(() -> LockSupportHelper.finish(state, 0L));
+  }
+
+  @Test
+  void finishAlwaysDrainsUnparkAttribution() {
+    Thread current = Thread.currentThread();
+    LockSupportHelper.UNPARKING_SPAN.put(current, 31L);
+
+    LockSupportHelper.finish(null);
+
+    assertNull(LockSupportHelper.UNPARKING_SPAN.get(current));
+  }
+
+  @Test
+  void latestTracedUnparkWins() {
+    Thread target = new Thread(() -> {});
+    installActiveSpan(41L);
+    LockSupportHelper.recordUnpark(target);
+    installActiveSpan(42L);
+
+    LockSupportHelper.recordUnpark(target);
+
+    assertEquals(42L, LockSupportHelper.UNPARKING_SPAN.get(target));
+  }
+
+  @Test
+  void laterUntracedUnparkClearsOlderTracedCaller() {
+    Thread target = new Thread(() -> {});
+    installActiveSpan(51L);
+    LockSupportHelper.recordUnpark(target);
+    AgentTracer.forceRegister(tracerWithActiveSpan(null));
+
+    LockSupportHelper.recordUnpark(target);
+
+    assertNull(LockSupportHelper.UNPARKING_SPAN.get(target));
+  }
+
+  @Test
+  void nonProfilerSpanContextClearsOlderTracedCaller() {
+    Thread target = new Thread(() -> {});
+    installActiveSpan(61L);
+    LockSupportHelper.recordUnpark(target);
+    AgentSpan span = mock(AgentSpan.class);
+    when(span.spanContext()).thenReturn(mock(AgentSpanContext.class));
+    AgentTracer.forceRegister(tracerWithActiveSpan(span));
+
+    LockSupportHelper.recordUnpark(target);
+
+    assertNull(LockSupportHelper.UNPARKING_SPAN.get(target));
+  }
+
+  @Test
+  void nullUnparkTargetIsNoOp() {
+    installActiveSpan(71L);
+
+    LockSupportHelper.recordUnpark(null);
+
+    assertEquals(0, LockSupportHelper.UNPARKING_SPAN.size());
+  }
+
+  private static void installActiveSpan(long spanId) {
+    AgentSpan span = mock(AgentSpan.class);
+    ProfilerSpanContext context = mock(ProfilerSpanContext.class);
+    when(span.spanContext()).thenReturn(context);
+    when(context.getSpanId()).thenReturn(spanId);
+    AgentTracer.forceRegister(tracerWithActiveSpan(span));
+  }
+
+  private static AgentTracer.TracerAPI tracerWithActiveSpan(AgentSpan span) {
+    AgentTracer.TracerAPI tracer = mock(AgentTracer.TracerAPI.class);
+    when(tracer.activeSpan()).thenReturn(span);
+    return tracer;
+  }
+}