feat: add core TaskBlock event recording support

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

@@ -64,6 +64,8 @@
   X(THREAD_REGISTRY_INDEX_FAILURES, "thread_registry_index_failures")          \
   X(THREAD_REGISTRY_CONTEXT_RESET_RACE_DETECTED, "thread_registry_context_reset_race_detected") \
   X(THREAD_REGISTRY_JAVACRITICAL_REREGISTRATION, "thread_registry_javacritical_reregistration") \
+  X(THREAD_REGISTRY_BLOCK_GENERATION_SATURATED, "thread_registry_block_generation_saturated") \
+  X(THREAD_REGISTRY_UNREGISTER_ACTIVE_BLOCK_RUN, "thread_registry_unregister_active_block_run") \
   X(JMETHODID_SKIPPED, "jmethodid_skipped_count")                              \
   X(CODECACHE_NATIVE_SIZE_BYTES, "codecache_native_size_bytes")                \
   X(CODECACHE_NATIVE_COUNT, "native_codecache_count")                          \
@@ -73,14 +75,23 @@
   X(AGCT_NATIVE_NO_JAVA_CONTEXT, "agct_native_no_java_context")                \
   X(AGCT_BLOCKED_IN_VM, "agct_blocked_in_vm")                                  \
   X(SKIPPED_WALLCLOCK_UNWINDS, "skipped_wallclock_unwinds")                    \
-  X(WC_SIGNAL_SUPPRESSED_SAMPLED_RUN, "wc_signals_suppressed_sampled_run")     \
   X(WC_PRECHECK_REGISTRY_LOOKUPS, "wc_precheck_registry_lookups")             \
   X(WC_PRECHECK_SLOT_ID_RECOVERED, "wc_precheck_slot_id_recovered")           \
   X(WC_PRECHECK_CANDIDATES_REJECTED, "wc_precheck_candidates_rejected")       \
   X(WC_PRECHECK_LOOKUP_BUDGET_EXHAUSTED, "wc_precheck_lookup_budget_exhausted") \
+  X(WC_SIGNAL_SUPPRESSED_OWNED_BLOCK, "wc_signals_suppressed_owned_block")     \
   X(WC_UNOWNED_BLOCKED_SUPPRESSED, "wc_unowned_blocked_suppressed")            \
   X(WC_UNOWNED_BLOCKED_RECORDED, "wc_unowned_blocked_recorded")                \
   X(WC_SIGNAL_QUEUE_FULL, "wc_signals_queue_full")                             \
+  X(TASK_BLOCK_EMITTED, "task_block_emitted")                                 \
+  X(TASK_BLOCK_SKIPPED_TRACE_CONTEXT, "task_block_skipped_trace_context")     \
+  X(TASK_BLOCK_SKIPPED_CONTEXT_WINDOW, "task_block_skipped_context_window")   \
+  X(TASK_BLOCK_SKIPPED_TOO_SHORT, "task_block_skipped_too_short")             \
+  X(TASK_BLOCK_STACK_CAPTURE_FAILED, "task_block_stack_capture_failed")      \
+  X(TASK_BLOCK_RECORD_FAILED, "task_block_record_failed")                     \
+  X(TASK_BLOCK_DROPPED_ROTATION, "task_block_dropped_rotation")               \
+  X(TASK_BLOCK_SKIPPED_THREAD_MISMATCH, "task_block_skipped_thread_mismatch") \
+  X(TASK_BLOCK_ROTATION_TIMEOUT, "task_block_rotation_timeout")               \
   X(UNWINDING_TIME_ASYNC, "unwinding_ticks_async")                             \
   X(UNWINDING_TIME_JVMTI, "unwinding_ticks_jvmti")                             \
   X(CALLTRACE_STORAGE_DROPPED, "calltrace_storage_dropped_traces")             \

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

@@ -58,11 +58,10 @@ class ExecutionEvent : public Event {
   OSThreadState _thread_state;
   ExecutionMode _execution_mode;
   u64 _weight;
-  u32 _call_trace_id;
 
   ExecutionEvent()
       : Event(), _thread_state(OSThreadState::RUNNABLE), _execution_mode(ExecutionMode::UNKNOWN),
-        _weight(1), _call_trace_id(0) {}
+        _weight(1) {}
 };
 
 class AllocEvent : public Event {
@@ -123,13 +122,13 @@ class WallClockEpochEvent {
   u32 _num_failed_samples;
   u32 _num_exited_threads;
   u32 _num_permission_denied;
-  u64 _num_suppressed_sampled_run;
+  u64 _num_suppressed_owned_block;
 
   WallClockEpochEvent(u64 start_time)
       : _dirty(false), _start_time(start_time), _duration_millis(0),
         _num_samplable_threads(0), _num_successful_samples(0),
         _num_failed_samples(0), _num_exited_threads(0),
-        _num_permission_denied(0), _num_suppressed_sampled_run(0) {}
+        _num_permission_denied(0), _num_suppressed_owned_block(0) {}
 
   bool hasChanged() { return _dirty; }
 
@@ -168,10 +167,10 @@ class WallClockEpochEvent {
     }
   }
 
-  void addNumSuppressedSampledRun(u64 n) {
+  void addNumSuppressedOwnedBlock(u64 n) {
     if (n > 0) {
       _dirty = true;
-      _num_suppressed_sampled_run += n;
+      _num_suppressed_owned_block += n;
     }
   }
 
@@ -182,7 +181,7 @@ class WallClockEpochEvent {
   void newEpoch(u64 start_time) {
     _dirty = false;
     _start_time = start_time;
-    _num_suppressed_sampled_run = 0;
+    _num_suppressed_owned_block = 0;
   }
 };
 
@@ -207,4 +206,14 @@ typedef struct QueueTimeEvent {
   u32 _queueLength;
 } QueueTimeEvent;
 
+typedef struct TaskBlockEvent {
+  u64 _start;
+  u64 _end;
+  u64 _blocker;
+  u64 _unblockingSpanId;
+  Context _ctx;
+  u64 _callTraceId;
+  OSThreadState _observedBlockingState;
+} TaskBlockEvent;
+
 #endif // _EVENT_H

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

@@ -1971,6 +1971,21 @@ void Recording::recordMethodSample(Buffer *buf, int tid, u64 call_trace_id,
   flushIfNeeded(buf);
 }
 
+void Recording::recordTaskBlock(Buffer *buf, int tid, TaskBlockEvent *event) {
+  int start = buf->skip(1);
+  buf->putVar64(T_TASK_BLOCK);
+  buf->putVar64(event->_start);
+  buf->putVar64(event->_end - event->_start);
+  buf->putVar64(tid);
+  buf->putVar64(event->_blocker);
+  buf->putVar64(event->_unblockingSpanId);
+  buf->putVar64(event->_callTraceId);
+  buf->put8(static_cast<int>(event->_observedBlockingState));
+  writeContextSnapshot(buf, event->_ctx);
+  writeEventSizePrefix(buf, start);
+  flushIfNeeded(buf);
+}
+
 void Recording::recordWallClockEpoch(Buffer *buf, WallClockEpochEvent *event) {
   int start = buf->skip(1);
   buf->putVar64(T_WALLCLOCK_SAMPLE_EPOCH);
@@ -1981,7 +1996,7 @@ void Recording::recordWallClockEpoch(Buffer *buf, WallClockEpochEvent *event) {
   buf->putVar64(event->_num_failed_samples);
   buf->putVar64(event->_num_exited_threads);
   buf->putVar64(event->_num_permission_denied);
-  buf->putVar64(event->_num_suppressed_sampled_run);
+  buf->putVar64(event->_num_suppressed_owned_block);
   writeEventSizePrefix(buf, start);
   flushIfNeeded(buf);
 }
@@ -2244,6 +2259,21 @@ void FlightRecorder::recordQueueTime(int lock_index, int tid,
   }
 }
 
+bool FlightRecorder::recordTaskBlock(int lock_index, int tid,
+                                     TaskBlockEvent *event) {
+  OptionalSharedLockGuard locker(&_rec_lock);
+  if (locker.ownsLock()) {
+    Recording* rec = _rec;
+    if (rec != nullptr) {
+      Buffer *buf = rec->buffer(lock_index);
+      rec->addThread(lock_index, tid);
+      rec->recordTaskBlock(buf, tid, event);
+      return true;
+    }
+  }
+  return false;
+}
+
 void FlightRecorder::recordDatadogSetting(int lock_index, int length,
                                           const char *name, const char *value,
                                           const char *unit) {

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

@@ -338,6 +338,7 @@ class Recording {
   void recordWallClockEpoch(Buffer *buf, WallClockEpochEvent *event);
   void recordTraceRoot(Buffer *buf, int tid, TraceRootEvent *event);
   void recordQueueTime(Buffer *buf, int tid, QueueTimeEvent *event);
+  void recordTaskBlock(Buffer *buf, int tid, TaskBlockEvent *event);
   void recordAllocation(RecordingBuffer *buf, int tid, u64 call_trace_id,
                         AllocEvent *event);
   void recordMallocSample(Buffer *buf, int tid, u64 call_trace_id,
@@ -439,6 +440,7 @@ class FlightRecorder {
   void wallClockEpoch(int lock_index, WallClockEpochEvent *event);
   void recordTraceRoot(int lock_index, int tid, TraceRootEvent *event);
   void recordQueueTime(int lock_index, int tid, QueueTimeEvent *event);
+  bool recordTaskBlock(int lock_index, int tid, TaskBlockEvent *event);
 
   bool active() const { return _rec != NULL; }
 

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

@@ -31,6 +31,7 @@
 #include "os.h"
 #include "otel_process_ctx.h"
 #include "profiler.h"
+#include "taskBlockRecorder.h"
 #include "threadLocalData.inline.h"
 #include "tsc.h"
 #include "vmEntry.h"
@@ -164,40 +165,6 @@ Java_com_datadoghq_profiler_JavaProfiler_getSamples(JNIEnv *env,
 
 // some duplication between add and remove, though we want to avoid having an extra branch in the hot path
 
-static ThreadFilter::SlotID ensureCurrentThreadFilterSlot(
-    ThreadFilter *thread_filter, ProfiledThread *current) {
-  int tid = current->tid();
-  if (unlikely(tid < 0)) {
-    return -1;
-  }
-
-  ThreadFilter::SlotID slot_id = current->filterSlotId();
-  if (likely(slot_id >= 0)) {
-    if (likely(thread_filter->activeSlotForId(slot_id, tid) != nullptr)) {
-      return slot_id;
-    }
-    current->setFilterSlotId(-1);
-  }
-
-  // Startup can register this TID centrally, but it cannot update another
-  // pthread's TLS. registerThread(tid) reuses that existing slot.
-  //
-  // This is the only place a JavaCritical fast path (filterThreadAdd0,
-  // parkEnter0, blockEnter0) can block on _registry_lock. It's bounded to at
-  // most once per thread lifetime (cold TLS) plus once per recording-epoch
-  // transition this thread observes (stale cached slot) - not a per-call cost.
-  // THREAD_REGISTRY_JAVACRITICAL_REREGISTRATION makes that bound observable;
-  // if it starts firing per-sample rather than per-thread/per-recording, the
-  // "provably rare" assumption has broken and the JavaCritical dispatch should
-  // be revisited.
-  Counters::increment(THREAD_REGISTRY_JAVACRITICAL_REREGISTRATION);
-  slot_id = thread_filter->registerThread(tid);
-  if (slot_id >= 0) {
-    current->setFilterSlotId(slot_id);
-  }
-  return slot_id;
-}
-
 // JavaCritical is faster JNI, but more restrictive - parameters and return value have to be
 // primitives or arrays of primitive types.
 // We direct corresponding JNI calls to JavaCritical to make sure the parameters/return value
@@ -219,15 +186,15 @@ JavaCritical_com_datadoghq_profiler_JavaProfiler_filterThreadAdd0() {
     return;
   }
 
-  int slot_id = ensureCurrentThreadFilterSlot(thread_filter, current);
+  int slot_id = thread_filter->ensureCurrentThreadSlot(current);
   if (unlikely(slot_id < 0)) {
     return;  // Failed to register thread
   }
   if (unlikely(!thread_filter->add(tid, slot_id))) {
     // The cached slot_id was rejected (lazy tid-index fallback failed under
     // this thread's own registry reset race, or the tid index is exhausted).
     // Clear the cache so the next filterThreadAdd0()/parkEnter0()/blockEnter0()
-    // call re-runs ensureCurrentThreadFilterSlot()'s registerThread() path
+    // call re-runs ensureCurrentThreadSlot()'s registerThread() path
     // instead of leaving this thread permanently outside the context window.
     current->setFilterSlotId(-1);
   }
@@ -421,22 +388,23 @@ Java_com_datadoghq_profiler_JavaProfiler_recordQueueEnd0(
   Profiler::instance()->recordQueueTime(tid, &event);
 }
 
-extern "C" DLLEXPORT void JNICALL
+extern "C" DLLEXPORT jboolean JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_parkEnter0(JNIEnv *env, jclass unused) {
   ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
-    return;
+    return JNI_FALSE;
   }
 
   bool first_park = current->parkEnter();
   ThreadFilter *tf = Profiler::instance()->threadFilter();
   if (first_park && tf->registryActive()) {
-    ThreadFilter::SlotID slot_id = ensureCurrentThreadFilterSlot(tf, current);
+    ThreadFilter::SlotID slot_id = tf->ensureCurrentThreadSlot(current);
     if (slot_id >= 0) {
       current->setParkBlockToken(
           tf->enterBlockedRun(slot_id, OSThreadState::CONDVAR_WAIT));
     }
   }
+  return first_park ? JNI_TRUE : JNI_FALSE;
 }
 
 extern "C" DLLEXPORT void JNICALL
@@ -470,26 +438,43 @@ static bool decodeJavaBlockState(jint state, OSThreadState &decoded) {
   return false;
 }
 
+// beginTaskBlock0/endTaskBlock0 accept an explicit jthread for stack-walking and
+// classification; verify it actually identifies the calling thread, since the two
+// could otherwise diverge (e.g. a stale or mismatched jthread handle).
+static bool isCurrentJniThread(JNIEnv* env, jthread thread) {
+  if (thread == nullptr) return false;
+  jthread current_thread = nullptr;
+  if (VM::jvmti()->GetCurrentThread(&current_thread) != JVMTI_ERROR_NONE) {
+    return false;
+  }
+  bool same = env->IsSameObject(thread, current_thread);
+  env->DeleteLocalRef(current_thread);
+  return same;
+}
+
 extern "C" DLLEXPORT jlong JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_blockEnter0(
     JNIEnv *env, jclass unused, jint state) {
   ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
     return 0;
   }
-
   OSThreadState decoded;
   if (!decodeJavaBlockState(state, decoded)) {
     return 0;
   }
-  ThreadFilter *tf = Profiler::instance()->threadFilter();
-  if (!tf->registryActive()) {
+  u64 span_id = 0, root_span_id = 0;
+  ContextApi::get(span_id, root_span_id);
+  if (span_id != 0) {
     return 0;
   }
-  ThreadFilter::SlotID slot_id = ensureCurrentThreadFilterSlot(tf, current);
-  if (slot_id < 0) {
+  Profiler *profiler = Profiler::instance();
+  ThreadFilter *tf = profiler->threadFilter();
+  if (!profiler->taskBlockEnabled() && !tf->registryActive()) {
     return 0;
   }
+  ThreadFilter::SlotID slot_id = tf->ensureCurrentThreadSlot(current);
+  if (slot_id < 0) return 0;
   return static_cast<jlong>(tf->enterBlockedRun(slot_id, decoded));
 }
 
@@ -516,6 +501,69 @@ Java_com_datadoghq_profiler_JavaProfiler_blockExit0(
   }
 }
 
+extern "C" DLLEXPORT jlong JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_beginTaskBlock0(
+    JNIEnv *env, jclass unused, jthread thread) {
+  ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
+  Profiler *profiler = Profiler::instance();
+  if (current == nullptr || !profiler->isRunning() ||
+      !profiler->taskBlockEnabled()) {
+    return 0;
+  }
+  ThreadFilter *tf = profiler->threadFilter();
+  if (!tf->unfilteredWallTrackingActive()) return 0;
+  if (!isCurrentJniThread(env, thread)) {
+    Counters::increment(TASK_BLOCK_SKIPPED_THREAD_MISMATCH);
+    return 0;
+  }
+  if (!JVMSupport::isPlatformThread(env, thread)) {
+    return 0;
+  }
+  ThreadFilter::SlotID slot_id = tf->ensureCurrentThreadSlot(current);
+  if (slot_id < 0) return 0;
+
+  Context context = ContextApi::snapshot();
+  if (context.spanId != 0) {
+    Counters::increment(TASK_BLOCK_SKIPPED_TRACE_CONTEXT);
+    return 0;
+  }
+  u64 token = tf->enterBlockedRun(
+      slot_id, OSThreadState::SLEEPING, BlockRunOwner::JAVA);
+  if (!current->taskBlockEnter(token, TSC::ticks(), context)) {
+    if (token != 0) {
+      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, jthread thread, jlong token, jlong blocker,
+    jlong unblockingSpanId) {
+  u64 block_token = static_cast<u64>(token);
+  ThreadFilter::SlotID slot_id = -1;
+  u64 generation = 0;
+  if (!ThreadFilter::decodeBlockRunToken(block_token, slot_id, generation)) {
+    return JNI_FALSE;
+  }
+  if (!isCurrentJniThread(env, thread)) {
+    Counters::increment(TASK_BLOCK_SKIPPED_THREAD_MISMATCH);
+    return JNI_FALSE;
+  }
+  if (!JVMSupport::isPlatformThread(env, thread)) {
+    return JNI_FALSE;
+  }
+  ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
+  if (current == nullptr) return JNI_FALSE;
+
+  bool recorded = recordTaskBlockAtExit(
+      current, Profiler::instance()->threadFilter(), thread, 1, block_token,
+      static_cast<u64>(blocker), static_cast<u64>(unblockingSpanId));
+  return recorded ? JNI_TRUE : JNI_FALSE;
+}
+
 extern "C" DLLEXPORT jlong JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_currentTicks0(JNIEnv *env,
                                                        jclass unused) {