feat: add core TaskBlock event recording support

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

@@ -71,13 +71,19 @@
   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_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_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(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,7 +58,7 @@ class ExecutionEvent : public Event {
   OSThreadState _thread_state;
   ExecutionMode _execution_mode;
   u64 _weight;
-  u32 _call_trace_id;
+  u64 _call_trace_id;
 
   ExecutionEvent()
       : Event(), _thread_state(OSThreadState::RUNNABLE), _execution_mode(ExecutionMode::UNKNOWN),
@@ -123,13 +123,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 +168,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 +182,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 +207,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

@@ -1945,6 +1945,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);
@@ -1955,7 +1970,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);
 }
@@ -2218,6 +2233,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

@@ -319,6 +319,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,
@@ -427,6 +428,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

@@ -32,6 +32,7 @@
 #include "otel_process_ctx.h"
 #include "profiler.h"
 #include "threadLocalData.h"
+#include "taskBlockRecorder.h"
 #include "tsc.h"
 #include "vmEntry.h"
 #include <errno.h>
@@ -417,14 +418,16 @@ Java_com_datadoghq_profiler_JavaProfiler_blockEnter0(
   if (current == nullptr) {
     return 0;
   }
-  ThreadFilter *tf = Profiler::instance()->threadFilter();
-  if (!tf->registryActive()) {
+  if (ContextApi::snapshot().spanId != 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->enabled()) {
     return 0;
   }
+  ThreadFilter::SlotID slot_id = ensureCurrentThreadFilterSlot(tf, current);
+  if (slot_id < 0) return 0;
   return static_cast<jlong>(tf->enterBlockedRun(slot_id, decoded));
 }
 
@@ -450,6 +453,60 @@ Java_com_datadoghq_profiler_JavaProfiler_blockExit0(
   }
 }
 
+extern "C" DLLEXPORT jlong JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_beginTaskBlock0(
+    JNIEnv *env, jclass unused, jthread thread) {
+  if (!JVMSupport::isPlatformThread(env, thread)) {
+    return 0;
+  }
+  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;
+  ThreadFilter::SlotID slot_id = ensureCurrentThreadFilterSlot(tf, 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) ||
+      !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,
+      slot_id, generation, 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) {

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

@@ -156,8 +156,8 @@ void JfrMetadata::initialize(
                        "Number of Exited Threads Before Handling Signal")
               << field("numPermissionDenied", T_INT,
                        "Number of Permission Denied Errors")
-              << field("numSuppressedSampledRun", T_LONG,
-                       "Signals suppressed by the wall-clock once-per-run filter"))
+              << field("numSuppressedOwnedBlock", T_LONG,
+                       "Signals suppressed for lifecycle-owned blocked intervals"))
 
           << (type("datadog.ObjectSample", T_ALLOC, "Allocation sample")
                   << category("Datadog", "Profiling")
@@ -209,6 +209,20 @@ void JfrMetadata::initialize(
                   << field("localRootSpanId", T_LONG, "Local Root Span ID") ||
               contextAttributes)
 
+          << (type("datadog.TaskBlock", T_TASK_BLOCK, "Task Block")
+                  << category("Datadog")
+                  << field("startTime", T_LONG, "Start Time", F_TIME_TICKS)
+                  << field("duration", T_LONG, "Duration", F_DURATION_TICKS)
+                  << field("eventThread", T_THREAD, "Event Thread", F_CPOOL)
+                  << field("blocker", T_LONG, "Blocker Identity Hash")
+                  << field("unblockingSpanId", T_LONG, "Unblocking Span ID")
+                  << field("stackTrace", T_STACK_TRACE, "Stack Trace", F_CPOOL)
+                  << field("observedBlockingState", T_THREAD_STATE,
+                           "Observed Blocking State", F_CPOOL)
+                  << field("spanId", T_LONG, "Span ID")
+                  << field("localRootSpanId", T_LONG, "Local Root Span ID") ||
+              contextAttributes)
+
           << (type("datadog.HeapUsage", T_HEAP_USAGE, "JVM Heap Usage")
               << category("Datadog")
               << field("startTime", T_LONG, "Start Time", F_TIME_TICKS)

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

@@ -81,6 +81,7 @@ enum JfrType {
   T_UNWIND_FAILURE = 126,
   T_MALLOC = 127,
   T_NATIVE_SOCKET = 128,
+  T_TASK_BLOCK = 129,
   T_ANNOTATION = 200,
   T_LABEL = 201,
   T_CATEGORY = 202,

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

@@ -16,16 +16,38 @@
 
 #include <jni.h>
 
+using JniFunction = void (JNICALL*)();
+using IsVirtualThreadFunction = jboolean (JNICALL*)(JNIEnv*, jobject);
+
+static constexpr jint JNI_VERSION_19_VALUE = 0x00130000;
+static constexpr int IS_VIRTUAL_THREAD_INDEX = 234;
+
+static_assert(sizeof(JniFunction) == sizeof(void*),
+              "JNI function table entries must be pointer-sized");
 
 volatile JVMSupport::JMethodIDLoadStats JVMSupport::jmethodID_load_state = JVMSupport::No_loaded;
 Mutex JVMSupport::_initialization_lock;
 
-
 // This method must be called after JVM has been properly initialized, e.g. after JVMTI::VMinit()
 // callback.
 // Currently, there are two paths lead to this call
 // - JVMTI::VMInit() callback (vmEntry.cpp)
 // - JavaProfiler.getInstance() via JNI down call - JVM must have been initialized
+bool JVMSupport::isPlatformThread(JNIEnv* jni, jthread thread) {
+    if (jni == nullptr || thread == nullptr) return false;
+    jint jni_version = jni->GetVersion();
+    if (jni_version <= 0) return false;
+    if (jni_version < JNI_VERSION_19_VALUE) return true;
+
+    const JniFunction* functions =
+        reinterpret_cast<const JniFunction*>(jni->functions);
+    IsVirtualThreadFunction is_virtual_thread =
+        reinterpret_cast<IsVirtualThreadFunction>(
+            functions[IS_VIRTUAL_THREAD_INDEX]);
+    return is_virtual_thread != nullptr &&
+        is_virtual_thread(jni, thread) == JNI_FALSE;
+}
+
 bool JVMSupport::initialize() {
     MutexLocker locker(_initialization_lock);
 

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

@@ -44,6 +44,10 @@ class JVMSupport {
 
     static bool isInitialized();
 public:
+    // Java-owned profiler state is carrier-local and may only be used by platform threads.
+    // IsVirtualThread was added to the JNI function table in JDK 19.
+    static bool isPlatformThread(JNIEnv* jni, jthread thread);
+
     // Initialize JVM support - check JVM related resources are available.
     // Return false if any critical resource is not available, which should
     // result in disabling profiling.