feat: add paired TaskBlock lifecycle API

Author: kaahos

ddprof-lib/src/main/cpp/javaApi.cpp

@@ -396,20 +396,6 @@ static bool decodeJavaBlockState(jint state, OSThreadState &decoded) {
   return false;
 }
 
-static OSThreadState decodeTaskBlockObservedState(jint state) {
-  switch (static_cast<OSThreadState>(state)) {
-    case OSThreadState::MONITOR_WAIT:
-    case OSThreadState::CONDVAR_WAIT:
-    case OSThreadState::OBJECT_WAIT:
-    case OSThreadState::SLEEPING:
-    case OSThreadState::SYSCALL:
-    case OSThreadState::IO_WAIT:
-      return static_cast<OSThreadState>(state);
-    default:
-      return OSThreadState::UNKNOWN;
-  }
-}
-
 static bool snapshotAndExitBlockedRun(jlong token, BlockRunSnapshot *snapshot) {
   if (token <= 0) {
     return false;
@@ -433,23 +419,6 @@ static bool snapshotAndExitBlockedRun(jlong token, BlockRunSnapshot *snapshot) {
   return tf->exitBlockedRun(slot_id, generation);
 }
 
-static void writeBlockRunSnapshot(JNIEnv *env, jlongArray snapshot_array,
-                                  const BlockRunSnapshot *snapshot) {
-  if (snapshot_array == nullptr || env->GetArrayLength(snapshot_array) < 3) {
-    return;
-  }
-  jlong values[3] = {0, 0, static_cast<jlong>(OSThreadState::UNKNOWN)};
-  if (snapshot != nullptr) {
-    OSThreadState observed_state = snapshot->sampled_state != OSThreadState::UNKNOWN
-                                       ? snapshot->sampled_state
-                                       : snapshot->active_state;
-    values[0] = snapshot->has_stack_reference ? static_cast<jlong>(snapshot->call_trace_id) : 0;
-    values[1] = snapshot->has_stack_reference ? static_cast<jlong>(snapshot->correlation_id) : 0;
-    values[2] = static_cast<jlong>(observed_state);
-  }
-  env->SetLongArrayRegion(snapshot_array, 0, 3, values);
-}
-
 extern "C" DLLEXPORT jlong JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_blockEnter0(
     JNIEnv *env, jclass unused, jint state) {
@@ -479,18 +448,79 @@ Java_com_datadoghq_profiler_JavaProfiler_blockExit0(
   snapshotAndExitBlockedRun(token, nullptr);
 }
 
-extern "C" DLLEXPORT void JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_blockExitWithSnapshot0(
-    JNIEnv *env, jclass unused, jlong token, jlongArray snapshotArray) {
+extern "C" DLLEXPORT jlong JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_beginTaskBlock0(
+    JNIEnv *env, jclass unused, jint state) {
+  OSThreadState decoded;
+  if (!decodeJavaBlockState(state, decoded)) {
+    return 0;
+  }
+  ProfiledThread* current = ProfiledThread::current();
+  if (current == nullptr) {
+    return 0;
+  }
+  if (VM::isHotspot() && VM::hotspot_version() >= 21) {
+    VMThread* carrier = VMThread::current();
+    if (carrier != nullptr && carrier->isCarryingVirtualThread()) {
+      return 0;
+    }
+  }
+  Profiler* profiler = Profiler::instance();
+  if (!profiler->isRunning() || !profiler->taskBlockAsyncActive()) {
+    return 0;
+  }
+  ThreadFilter* tf = profiler->threadFilter();
+  ThreadFilter::SlotID slot_id = current->filterSlotId();
+  if (!tf->enabled() || slot_id < 0) {
+    return 0;
+  }
+  u64 token = tf->enterBlockedRun(slot_id, decoded, BlockRunOwner::JAVA);
+  if (!current->taskBlockEnter(token, TSC::ticks(), ContextApi::snapshot())) {
+    tf->exitBlockedRun(slot_id, ThreadFilter::tokenGeneration(token));
+    return 0;
+  }
+  return static_cast<jlong>(token);
+}
+
+extern "C" DLLEXPORT jboolean JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_endTaskBlock0(
+    JNIEnv *env, jclass unused, jlong token, jlong blocker,
+    jlong unblockingSpanId) {
+  ProfiledThread* current = ProfiledThread::current();
+  if (current == nullptr || token <= 0) {
+    return JNI_FALSE;
+  }
+  u64 start_ticks = 0;
+  Context context{};
+  if (!current->taskBlockExit(static_cast<u64>(token), start_ticks, context)) {
+    return JNI_FALSE;
+  }
+  Profiler* profiler = Profiler::instance();
+  bool activity = profiler->tryEnterTaskBlockActivity();
+  if (!activity) {
+    profiler->waitForTaskBlockRotation();
+  }
+  ThreadFilter* tf = profiler->threadFilter();
+  ThreadFilter::SlotID slot_id = ThreadFilter::tokenSlotId(static_cast<u64>(token));
   BlockRunSnapshot snapshot{};
-  snapshot.active_state = OSThreadState::UNKNOWN;
-  snapshot.sampled_state = OSThreadState::UNKNOWN;
-  snapshot.owner = BlockRunOwner::NONE;
-  if (snapshotAndExitBlockedRun(token, &snapshot)) {
-    writeBlockRunSnapshot(env, snapshotArray, &snapshot);
-  } else {
-    writeBlockRunSnapshot(env, snapshotArray, nullptr);
+  bool exited = current->filterSlotId() == slot_id &&
+      tf->snapshotAndExitBlockedRun(
+          slot_id, ThreadFilter::tokenGeneration(static_cast<u64>(token)), &snapshot);
+  if (!activity) {
+    Counters::increment(TASK_BLOCK_DROPPED_ROTATION);
+    return JNI_FALSE;
+  }
+  bool recorded = false;
+  if (exited) {
+    OSThreadState observed_state = snapshot.sampled_state != OSThreadState::UNKNOWN
+        ? snapshot.sampled_state : snapshot.active_state;
+    recorded = recordTaskBlockWithStackReferenceIfEligible(
+        current->tid(), start_ticks, TSC::ticks(), context, (u64)blocker,
+        (u64)unblockingSpanId, snapshot.call_trace_id,
+        snapshot.correlation_id, observed_state, true);
   }
+  profiler->leaveTaskBlockActivity();
+  return recorded ? JNI_TRUE : JNI_FALSE;
 }
 
 extern "C" DLLEXPORT jboolean JNICALL
@@ -528,40 +558,6 @@ Java_com_datadoghq_profiler_JavaProfiler_recordTaskBlockWithContext0(
              : JNI_FALSE;
 }
 
-extern "C" DLLEXPORT jboolean JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_recordTaskBlockFromContext0(
-    JNIEnv *env, jclass unused, jint tid, jlong startTicks, jlong endTicks,
-    jlong blocker, jlong unblockingSpanId, jlong spanId, jlong rootSpanId) {
-  // Drain-thread path: called from background drain thread on behalf of the sleeping thread.
-  // tid is the OS tid of the sleeping thread; span context is explicit to bypass OTEP TLS.
-  if ((int)tid < 0) {
-    return JNI_FALSE;
-  }
-  Context ctx{(u64)spanId, (u64)rootSpanId};
-  return recordTaskBlockWithStackReferenceIfEligible(
-             (int)tid, (u64)startTicks, (u64)endTicks, ctx, (u64)blocker,
-             (u64)unblockingSpanId, 0, 0, OSThreadState::UNKNOWN)
-             ? JNI_TRUE
-             : JNI_FALSE;
-}
-
-extern "C" DLLEXPORT jboolean JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_recordTaskBlockFromContextWithStackReference0(
-    JNIEnv *env, jclass unused, jint tid, jlong startTicks, jlong endTicks,
-    jlong blocker, jlong unblockingSpanId, jlong spanId, jlong rootSpanId,
-    jlong callTraceId, jlong correlationId, jint observedBlockingState) {
-  if ((int)tid < 0) {
-    return JNI_FALSE;
-  }
-  Context ctx{(u64)spanId, (u64)rootSpanId};
-  return recordTaskBlockWithStackReferenceIfEligible(
-             (int)tid, (u64)startTicks, (u64)endTicks, ctx, (u64)blocker,
-             (u64)unblockingSpanId, (u64)callTraceId, (u64)correlationId,
-             decodeTaskBlockObservedState(observedBlockingState))
-             ? JNI_TRUE
-             : JNI_FALSE;
-}
-
 extern "C" DLLEXPORT jlong JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_currentTicks0(JNIEnv *env,
                                                        jclass unused) {

ddprof-lib/src/main/cpp/threadLocalData.h

@@ -84,6 +84,9 @@ class ProfiledThread : public ThreadLocalData {
   u64 _park_start_ticks;
   u64 _park_block_token;
   Context _park_context;
+  u64 _task_block_start_ticks;
+  u64 _task_block_token;
+  Context _task_block_context;
   u64 _monitor_start_ticks;
   Context _monitor_context;
   u64 _monitor_blocker;
@@ -110,6 +113,7 @@ class ProfiledThread : public ThreadLocalData {
       : ThreadLocalData(), _jmp_buf(nullptr), _pc(0), _sp(0), _span_id(0), _crash_depth(0), _tid(tid), _cpu_epoch(0),
         _wall_epoch(0), _call_trace_id(0), _recording_epoch(0), _misc_flags(0),
         _park_start_ticks(0), _park_block_token(0), _park_context{},
+        _task_block_start_ticks(0), _task_block_token(0), _task_block_context{},
         _monitor_start_ticks(0), _monitor_context{}, _monitor_blocker(0),
         _monitor_block_token(0), _monitor_block_state(OSThreadState::UNKNOWN),
         _filter_slot_id(-1), _init_window(0),
@@ -350,6 +354,26 @@ class ProfiledThread : public ThreadLocalData {
     return false;
   }
 
+  inline bool taskBlockEnter(u64 token, u64 start_ticks, const Context& context) {
+    if (token == 0 || _task_block_token != 0) {
+      return false;
+    }
+    _task_block_start_ticks = start_ticks;
+    _task_block_context = context;
+    _task_block_token = token;
+    return true;
+  }
+
+  inline bool taskBlockExit(u64 token, u64& start_ticks, Context& context) {
+    if (token == 0 || _task_block_token != token) {
+      return false;
+    }
+    start_ticks = _task_block_start_ticks;
+    context = _task_block_context;
+    _task_block_token = 0;
+    return true;
+  }
+
   inline void setParkBlockToken(u64 token) {
     _park_block_token = token;
   }

ddprof-lib/src/main/java/com/datadoghq/profiler/JavaProfiler.java

@@ -397,16 +397,6 @@ void blockExit(long token) {
         blockExit0(token);
     }
 
-    /**
-     * Clears a blocked interval and snapshots reconstruction metadata before native state is reset.
-     *
-     * @param token opaque token returned by {@link #blockEnter(int)}
-     * @param snapshot output array: [callTraceId, correlationId, observedBlockingState]
-     */
-    public void blockExit(long token, long[] snapshot) {
-        blockExitWithSnapshot0(token, snapshot);
-    }
-
     /**
      * Get the ticks for the current thread.
      * @return ticks
@@ -415,13 +405,6 @@ public long getCurrentTicks() {
         return currentTicks0();
     }
 
-    /**
-     * Returns the OS-level thread ID (tid) of the calling thread.
-     */
-    public int getCurrentThreadId() {
-        return getTid0();
-    }
-
     /**
      * Reports whether native {@code Object.wait} TaskBlock callbacks were disabled in favor of
      * Java-owned wait instrumentation. Synchronized monitor contention remains native-owned.
@@ -431,10 +414,29 @@ public boolean isMonitorEventsDelegated() {
     }
 
     /**
-     * Returns the TSC frequency in Hz (ticks per second).
+     * Begins an explicitly instrumented blocking interval on the current platform thread.
+     * The returned token is bound to the current thread and must be passed to
+     * {@link #endTaskBlock(long, long, long)}.
+     *
+     * @param state native {@code OSThreadState} value; currently only {@code SLEEPING} is accepted
+     * @return an opaque token, or {@code 0} when the interval could not be armed
+     */
+    public long beginTaskBlock(int state) {
+        return beginTaskBlock0(state);
+    }
+
+    /**
+     * Ends a blocking interval created by {@link #beginTaskBlock(int)} and records its
+     * {@code TaskBlock} event when it satisfies the profiler's eligibility rules.
+     * Lifecycle state is cleared even when no event is recorded.
+     *
+     * @param token opaque token returned by {@link #beginTaskBlock(int)}
+     * @param blocker stable identifier describing the blocking resource
+     * @param unblockingSpanId span responsible for unblocking the interval, or {@code 0}
+     * @return {@code true} when an event was recorded
      */
-    public long getTscFrequency() {
-        return tscFrequency0();
+    public boolean endTaskBlock(long token, long blocker, long unblockingSpanId) {
+        return endTaskBlock0(token, blocker, unblockingSpanId);
     }
 
     /**
@@ -460,34 +462,6 @@ public boolean recordTaskBlockWithContext(long startTicks, long endTicks, long b
         return recordTaskBlockWithContext0(startTicks, endTicks, blocker, unblockingSpanId, spanId, rootSpanId);
     }
 
-    /**
-     * Attempts to record a TaskBlock event attributed to an explicit thread ID and explicit span
-     * context. This overload has no stack-reference metadata, so a recorder thread cannot use it to
-     * emit a self-contained TaskBlock on behalf of another thread.
-     *
-     * @return {@code true} if the event was recorded; {@code false} if it was skipped by
-     *         eligibility rules or could not be recorded
-     */
-    public boolean recordTaskBlockFromContext(int tid, long startTicks, long endTicks,
-            long blocker, long unblockingSpanId, long spanId, long rootSpanId) {
-        return recordTaskBlockFromContext0(tid, startTicks, endTicks, blocker, unblockingSpanId, spanId, rootSpanId);
-    }
-
-    /**
-     * Records a TaskBlock event with explicit thread, span context, and stack-reference metadata.
-     * This is the required overload when recording from a background thread on behalf of {@code tid}.
-     *
-     * @return {@code true} if the event was recorded; {@code false} if it was skipped by
-     *         eligibility rules or could not be recorded
-     */
-    public boolean recordTaskBlockFromContext(int tid, long startTicks, long endTicks,
-            long blocker, long unblockingSpanId, long spanId, long rootSpanId,
-            long callTraceId, long correlationId, int observedBlockingState) {
-        return recordTaskBlockFromContextWithStackReference0(tid, startTicks, endTicks, blocker,
-                unblockingSpanId, spanId, rootSpanId, callTraceId, correlationId,
-                observedBlockingState);
-    }
-
     /**
      * If the profiler is built in debug mode, returns counters recorded during profile execution.
      * These are for whitebox testing and not intended for production use.
@@ -544,7 +518,9 @@ private static ThreadContext initializeThreadContext() {
 
     private static native void blockExit0(long token);
 
-    private static native void blockExitWithSnapshot0(long token, long[] snapshot);
+    private static native long beginTaskBlock0(int state);
+
+    private static native boolean endTaskBlock0(long token, long blocker, long unblockingSpanId);
 
     private static native long currentTicks0();
 
@@ -556,13 +532,6 @@ private static native boolean recordTaskBlock0(long startTicks, long endTicks,
     private static native boolean recordTaskBlockWithContext0(long startTicks, long endTicks,
             long blocker, long unblockingSpanId, long spanId, long rootSpanId);
 
-    private static native boolean recordTaskBlockFromContext0(int tid, long startTicks, long endTicks,
-            long blocker, long unblockingSpanId, long spanId, long rootSpanId);
-
-    private static native boolean recordTaskBlockFromContextWithStackReference0(int tid, long startTicks,
-            long endTicks, long blocker, long unblockingSpanId, long spanId, long rootSpanId,
-            long callTraceId, long correlationId, int observedBlockingState);
-
     private static native void mallocArenaMax0(int max);
 
     private static native String getStatus0();