Add Thread.sleep TaskBlock profiling instrumentation

.gitlab-ci.yml

@@ -64,6 +64,9 @@ variables:
   NON_DEFAULT_JVMS:
     description: "Enable tests on JVMs that are not the default"
     value: "false"
+  DDPROF_VERSION:
+    description: "ddprof version embedded in the agent artifact used by downstream tests"
+    value: "1.47.0-paul.fournillon_taskblock-native-io-SNAPSHOT"
   RUN_FLAKY_TESTS:
     description: "Enable flaky tests"
     value: "false"
@@ -343,14 +346,22 @@ build:
   script:
     - if [ $CI_PIPELINE_SOURCE == "schedule" ] ; then ./gradlew resolveAndLockAll --write-locks $GRADLE_ARGS; fi
     - ./gradlew --version
-    - ./gradlew clean :dd-java-agent:shadowJar :dd-java-agent:check :dd-trace-api:jar :dd-trace-ot:shadowJar -PskipTests -x spotlessCheck $GRADLE_ARGS
-    - echo UPSTREAM_TRACER_VERSION=$(java -jar workspace/dd-java-agent/build/libs/*.jar) >> upstream.env
+    - ./gradlew clean :dd-java-agent:shadowJar :dd-java-agent:check :dd-trace-api:jar :dd-trace-ot:shadowJar -Pddprof.version="$DDPROF_VERSION" -PskipTests -x spotlessCheck $GRADLE_ARGS
+    - AGENT_JARS=$(find workspace/dd-java-agent/build/libs -maxdepth 1 -type f -name 'dd-java-agent-*.jar' ! -name '*-javadoc.jar' ! -name '*-sources.jar' -print)
+    - test "$(printf '%s\n' "$AGENT_JARS" | grep -c .)" -eq 1
+    - AGENT_JAR="$AGENT_JARS"
+    - AGENT_SHA256=$(sha256sum "$AGENT_JAR" | cut -d' ' -f1)
+    - echo "$AGENT_SHA256  $(basename "$AGENT_JAR")" > "$AGENT_JAR.sha256"
+    - echo UPSTREAM_TRACER_VERSION=$(java -jar "$AGENT_JAR") >> upstream.env
     - echo "BUILD_JOB_NAME=$CI_JOB_NAME" >> build.env
     - echo "BUILD_JOB_ID=$CI_JOB_ID" >> build.env
+    - echo "SMOKETEST_AGENT_JAR=$AGENT_JAR" >> build.env
+    - echo "SMOKETEST_AGENT_SHA256=$AGENT_SHA256" >> build.env
   artifacts:
     when: always
     paths:
       - 'workspace/dd-java-agent/build/libs/*.jar'
+      - 'workspace/dd-java-agent/build/libs/*.sha256'
       - 'workspace/dd-trace-api/build/libs/*.jar'
       - 'workspace/dd-trace-ot/build/libs/*.jar'
       - 'upstream.env'
@@ -575,6 +586,8 @@ muzzle:
   # Keep matrix vars exact and in build_tests declaration order:
   # https://docs.gitlab.com/ci/yaml/#needsparallelmatrix
   needs: &needs_build_tests_inst
+    - job: build
+      artifacts: true
     - job: build_tests
       parallel:
         matrix:
@@ -660,7 +673,11 @@ muzzle-dep-report:
   extends: .gradle_build
   image: ${BUILDER_IMAGE_REPO}:${BUILDER_IMAGE_VERSION_PREFIX}$testJvm
   tags: [ "docker-in-docker:amd64" ] # use docker-in-docker runner for testcontainers
-  needs: [ build_tests ]
+  needs:
+    - job: build
+      artifacts: true
+    - job: build_tests
+      artifacts: false
   stage: tests
   variables:
     <<: *tier_m_variables
@@ -695,6 +712,8 @@ muzzle-dep-report:
     - *prepare_test_env
     - export GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xms$GRADLE_MEMORY_MIN -Xmx$GRADLE_MEMORY_MAX $PROFILER_COMMAND -XX:ErrorFile=/tmp/hs_err_pid%p.log -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp -Djava.util.prefs.userRoot=/tmp/.java/.userPrefs-${CI_JOB_ID}' -Ddatadog.forkedMinHeapSize=128M -Ddatadog.forkedMaxHeapSize=1024M"
     - ./gradlew --version
+    - test -n "$SMOKETEST_AGENT_JAR" && test -n "$SMOKETEST_AGENT_SHA256" && test -f "$SMOKETEST_AGENT_JAR"
+    - echo "$SMOKETEST_AGENT_SHA256  $SMOKETEST_AGENT_JAR" | sha256sum -c -
     - ./gradlew $GRADLE_TARGET $GRADLE_PARAMS -PtestJvm=$testJvm -Pslot=$CI_NODE_INDEX/$CI_NODE_TOTAL $GRADLE_ARGS --continue || $CONTINUE_ON_FAILURE
   after_script:
     - *restore_pretest_env
@@ -742,7 +761,9 @@ muzzle-dep-report:
   image: ${BUILDER_IMAGE_REPO}:${BUILDER_IMAGE_VERSION_PREFIX}$testJvm
   tags: [ "docker-in-docker:arm64" ]
   stage: tests-arm64
-  needs: []
+  needs:
+    - job: build
+      artifacts: true
   variables:
     <<: *tier_m_variables
     DEFAULT_TEST_JVMS: /^(8|11|17|21|25|27|tip)$/ # Java 27 TODO: remove 27 after GA (tip will move to 27)
@@ -821,6 +842,8 @@ muzzle-dep-report:
     - export GRADLE_OPTS="-Dorg.gradle.jvmargs='-Xshare:off -Xms$GRADLE_MEMORY_MIN -Xmx$GRADLE_MEMORY_MAX -XX:ErrorFile=/tmp/hs_err_pid%p.log -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/tmp -Djava.util.prefs.userRoot=/tmp/.java/.userPrefs-${CI_JOB_ID}' -Ddatadog.forkedMinHeapSize=128M -Ddatadog.forkedMaxHeapSize=1024M"
     - export GRADLE_ARGS="--build-cache --stacktrace --no-daemon --parallel --max-workers=$GRADLE_WORKERS"
     - ./gradlew --version
+    - test -n "$SMOKETEST_AGENT_JAR" && test -n "$SMOKETEST_AGENT_SHA256" && test -f "$SMOKETEST_AGENT_JAR"
+    - echo "$SMOKETEST_AGENT_SHA256  $SMOKETEST_AGENT_JAR" | sha256sum -c -
   script:
     - *gitlab_base_ref_params
     - ./gradlew $GRADLE_TARGET $GRADLE_PARAMS -PtestJvm=$testJvm -Pslot=$CI_NODE_INDEX/$CI_NODE_TOTAL $GRADLE_ARGS --continue
@@ -919,6 +942,8 @@ test_base:
   # Keep matrix vars exact and in build_tests declaration order:
   # https://docs.gitlab.com/ci/yaml/#needsparallelmatrix
   needs:
+    - job: build
+      artifacts: true
     - job: build_tests
       parallel:
         matrix:
@@ -970,6 +995,8 @@ test_inst_latest:
   # Keep matrix vars exact and in build_tests declaration order:
   # https://docs.gitlab.com/ci/yaml/#needsparallelmatrix
   needs:
+    - job: build
+      artifacts: true
     - job: build_tests
       parallel:
         matrix:
@@ -1038,6 +1065,8 @@ test_profiling:
   # Keep matrix vars exact and in build_tests declaration order:
   # https://docs.gitlab.com/ci/yaml/#needsparallelmatrix
   needs:
+    - job: build
+      artifacts: true
     - job: build_tests
       parallel:
         matrix:
@@ -1092,6 +1121,8 @@ test_smoke:
   # Keep matrix vars exact and in build_tests declaration order:
   # https://docs.gitlab.com/ci/yaml/#needsparallelmatrix
   needs: &needs_build_tests_smoke
+    - job: build
+      artifacts: true
     - job: build_tests
       parallel:
         matrix:

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) {
+    }
+  }
+}