feat: emit TaskBlock events for JVM blocking states

build-logic/conventions/src/main/kotlin/com/datadoghq/native/config/ConfigurationPresets.kt

@@ -1,4 +1,3 @@
-
 package com.datadoghq.native.config
 
 import com.datadoghq.native.model.Architecture

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

@@ -92,6 +92,7 @@
   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(TASK_BLOCK_VIRTUAL_THREAD_DETECTION_FAILED, "task_block_virtual_thread_detection_failed") \
   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/javaApi.cpp

@@ -70,7 +70,8 @@ class JniString {
 };
 
 extern "C" DLLEXPORT jboolean JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_init0(JNIEnv *env, jclass unused) {
+Java_com_datadoghq_profiler_JavaProfiler_init0(
+    JNIEnv *env, jclass unused, jboolean delegateMonitorWaitEvents) {
   Error error = Profiler::instance()->init();
   if (error) {
     throwNew(env, "java/lang/IllegalStateException", error.message());
@@ -79,13 +80,25 @@ Java_com_datadoghq_profiler_JavaProfiler_init0(JNIEnv *env, jclass unused) {
 
 
   // JavaVM* has already been stored when the native library was loaded so we can pass nullptr here
-  if (VM::initProfilerBridge(nullptr, true)) {
-    // Attach ProfiledThread
-    ProfiledThread::initCurrentThreadSignalSafe();
-    return JNI_TRUE;
-  } else {
+  ProfilerBridgeInitResult result =
+      VM::initProfilerBridge(nullptr, true, delegateMonitorWaitEvents);
+  if (result == ProfilerBridgeInitResult::MONITOR_EVENTS_DELEGATION_CONFLICT) {
+    // Keep this message identical to the one JavaProfiler.getInstance(String, String,
+    // boolean) throws for the analogous Java-singleton conflict; there is no shared
+    // constant across the JNI boundary, so both call sites must be updated together.
+    throwNew(env, "java/lang/IllegalStateException",
+             "Monitor-event ownership conflicts with the profiler's "
+             "process-wide initialization");
     return JNI_FALSE;
   }
+  if (result != ProfilerBridgeInitResult::SUCCESS) {
+    throwNew(env, "java/lang/IllegalStateException",
+             "Failed to initialize the profiler bridge");
+    return JNI_FALSE;
+  }
+  // Attach ProfiledThread
+  ProfiledThread::initCurrentThreadSignalSafe();
+  return JNI_TRUE;
 }
 
 extern "C" DLLEXPORT void JNICALL
@@ -110,6 +123,12 @@ Java_com_datadoghq_profiler_JavaProfiler_getTid0(JNIEnv *env, jclass unused) {
   return OS::threadId();
 }
 
+extern "C" DLLEXPORT jboolean JNICALL
+Java_com_datadoghq_profiler_JavaProfiler_monitorWaitEventsDelegated0(
+    JNIEnv *env, jclass unused) {
+  return VM::monitorWaitEventsDelegated();
+}
+
 extern "C" DLLEXPORT jstring JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_execute0(JNIEnv *env, jobject unused,
                                                   jstring command) {
@@ -389,51 +408,71 @@ Java_com_datadoghq_profiler_JavaProfiler_recordQueueEnd0(
 }
 
 extern "C" DLLEXPORT jboolean JNICALL
-Java_com_datadoghq_profiler_JavaProfiler_parkEnter0(JNIEnv *env, jclass unused) {
+Java_com_datadoghq_profiler_JavaProfiler_parkEnter0(
+    JNIEnv *env, jclass unused, jthread thread, jboolean isVirtual) {
+  // Virtuality is resolved once on the Java side; re-deriving it here would cost a
+  // GetVersion() plus an IsVirtualThread() JNI round-trip on every park.
+  if (isVirtual != JNI_FALSE) {
+    return JNI_FALSE;
+  }
   ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
     return JNI_FALSE;
   }
+  Profiler *profiler = Profiler::instance();
+  ThreadFilter *tf = profiler->threadFilter();
+  if (!tf->registryActive() ||
+      !(profiler->taskBlockEnabled() || tf->enabled())) {
+    return JNI_FALSE;
+  }
+  Context context = ContextApi::snapshot();
+  if (!current->parkEnter(TSC::ticks(), context)) {
+    return JNI_FALSE;
+  }
 
-  bool first_park = current->parkEnter();
-  ThreadFilter *tf = Profiler::instance()->threadFilter();
-  if (first_park && tf->registryActive()) {
+  if (context.spanId == 0) {
     ThreadFilter::SlotID slot_id = tf->ensureCurrentThreadSlot(current);
     if (slot_id >= 0) {
-      current->setParkBlockToken(
-          tf->enterBlockedRun(slot_id, OSThreadState::CONDVAR_WAIT));
+      current->setParkBlockToken(tf->enterBlockedRun(
+          slot_id, OSThreadState::CONDVAR_WAIT, BlockRunOwner::JAVA));
     }
   }
-  return first_park ? JNI_TRUE : JNI_FALSE;
+  return JNI_TRUE;
 }
 
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_parkExit0(
-    JNIEnv *env, jclass unused, jlong blocker, jlong unblockingSpanId) {
+    JNIEnv *env, jclass unused, jthread thread, jboolean isVirtual,
+    jlong blocker, jlong unblockingSpanId) {
+  if (isVirtual != JNI_FALSE) {
+    return;
+  }
   ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
     return;
   }
-
+  u64 start_ticks = 0;
   u64 park_block_token = 0;
-  if (!current->parkExit(park_block_token) || park_block_token == 0) {
+  Context context{};
+  if (!current->parkExit(start_ticks, context, park_block_token) ||
+      park_block_token == 0) {
     return;
   }
-  ThreadFilter *tf = Profiler::instance()->threadFilter();
-  if (tf->registryActive()) {
-    ThreadFilter::SlotID slot_id = ThreadFilter::tokenSlotId(park_block_token);
-    if (tf->activeSlotForId(current->filterSlotId(), current->tid()) != nullptr &&
-        current->filterSlotId() == slot_id) {
-      tf->exitBlockedRun(slot_id, ThreadFilter::tokenGeneration(park_block_token));
-    }
-  }
+  finishTaskBlockAtExit(current, Profiler::instance()->threadFilter(), thread,
+                        1, park_block_token, start_ticks, context,
+                        static_cast<u64>(blocker),
+                        static_cast<u64>(unblockingSpanId));
 }
 
 static bool decodeJavaBlockState(jint state, OSThreadState &decoded) {
   if (state == static_cast<jint>(OSThreadState::SLEEPING)) {
     decoded = OSThreadState::SLEEPING;
     return true;
   }
+  if (state == static_cast<jint>(OSThreadState::OBJECT_WAIT)) {
+    decoded = OSThreadState::OBJECT_WAIT;
+    return true;
+  }
   decoded = OSThreadState::UNKNOWN;
   return false;
 }
@@ -454,35 +493,75 @@ static bool isCurrentJniThread(JNIEnv* env, jthread thread) {
 
 extern "C" DLLEXPORT jlong JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_blockEnter0(
-    JNIEnv *env, jclass unused, jint state) {
-  ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
-  if (current == nullptr) {
+    JNIEnv *env, jclass unused, jboolean isVirtual, jint state) {
+  OSThreadState decoded;
+  if (!decodeJavaBlockState(state, decoded) || isVirtual != JNI_FALSE) {
     return 0;
   }
-  OSThreadState decoded;
-  if (!decodeJavaBlockState(state, decoded)) {
+  ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
+  if (current == nullptr) {
     return 0;
   }
-  u64 span_id = 0, root_span_id = 0;
-  ContextApi::get(span_id, root_span_id);
-  if (span_id != 0) {
+  Context context = ContextApi::snapshot();
+  if (context.spanId != 0) {
     return 0;
   }
   Profiler *profiler = Profiler::instance();
   ThreadFilter *tf = profiler->threadFilter();
-  if (!profiler->taskBlockEnabled() && !tf->registryActive()) {
+  bool task_block_armable = profiler->taskBlockEnabled();
+  if (!task_block_armable && !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));
+  u64 token = tf->enterBlockedRun(slot_id, decoded);
+  if (token == 0) return 0;
+  // Also arm the TaskBlock recording slot so blockExit0 can produce an event
+  // with blocker identity, matching the beginTaskBlock0/endTaskBlock0 pair.
+  // Skipped when TaskBlock recording is disabled so this stays purely a
+  // wall-clock suppression marker for the precheck tests that exercise it
+  // with task-block recording off.
+  if (task_block_armable &&
+      !current->taskBlockEnter(token, TSC::ticks(), context)) {
+    // FLAG_TASK_BLOCKED is sticky across recordings (unlike the ThreadFilter
+    // blocked-run, which is generation-scoped), so a worker whose previous
+    // interval was abandoned mid-run (e.g. profiler stop/restart without a
+    // matching blockExit0) can still be holding it. Mirror monitorBlockEnter's
+    // stale-slot check: only refuse the new interval if the old token still
+    // maps to a genuinely active, same-generation blocked run.
+    u64 stale_token = current->taskBlockToken();
+    bool current_owner = false;
+    if (stale_token != 0) {
+      ThreadFilter::SlotID stale_slot_id = ThreadFilter::tokenSlotId(stale_token);
+      ThreadFilter::Slot *stale_slot = current->filterSlotId() == stale_slot_id
+          ? tf->activeSlotForId(stale_slot_id, current->tid())
+          : nullptr;
+      if (stale_slot != nullptr) {
+        BlockRunSnapshot snapshot = stale_slot->snapshotBlockRun();
+        current_owner = snapshot.isActive() &&
+            snapshot.owner == BlockRunOwner::JAVA &&
+            snapshot.generation == ThreadFilter::tokenGeneration(stale_token);
+      }
+    }
+    if (current_owner) {
+      tf->exitBlockedRun(slot_id, ThreadFilter::tokenGeneration(token));
+      return 0;
+    }
+    current->clearTaskBlock();
+    if (!current->taskBlockEnter(token, TSC::ticks(), context)) {
+      tf->exitBlockedRun(slot_id, ThreadFilter::tokenGeneration(token));
+      return 0;
+    }
+  }
+  return static_cast<jlong>(token);
 }
 
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_blockExit0(
-    JNIEnv *env, jclass unused, jlong token) {
+    JNIEnv *env, jclass unused, jboolean isVirtual, jlong token, jlong blocker,
+    jlong unblockingSpanId) {
   u64 block_token = static_cast<u64>(token);
-  if (block_token == 0) {
+  if (block_token == 0 || isVirtual != JNI_FALSE) {
     return;
   }
 
@@ -496,6 +575,19 @@ Java_com_datadoghq_profiler_JavaProfiler_blockExit0(
       tf->activeSlotForId(slot_id, current->tid()) == nullptr) {
     return;
   }
+
+  u64 start_ticks = 0;
+  Context context{};
+  if (current->taskBlockExit(block_token, start_ticks, context)) {
+    // thread=nullptr: GetStackTrace treats NULL as "the calling thread", which
+    // is always correct here since blockExit0 only ever runs on the blocked
+    // thread itself.
+    finishTaskBlockAtExit(current, tf, /*thread=*/nullptr, 1, block_token,
+                          start_ticks, context, static_cast<u64>(blocker),
+                          static_cast<u64>(unblockingSpanId));
+    return;
+  }
+
   if (tf->registryActive()) {
     tf->exitBlockedRun(slot_id, ThreadFilter::tokenGeneration(block_token));
   }

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

@@ -7,6 +7,7 @@
 
 #include "asyncSampleMutex.h"
 #include "common.h"
+#include "counters.h"
 #include "frames.h"
 #include "os.h"
 #include "profiler.h"
@@ -18,6 +19,8 @@
 
 #include <jni.h>
 
+#include <atomic>
+
 using JniFunction = void (JNICALL*)();
 using IsVirtualThreadFunction = jboolean (JNICALL*)(JNIEnv*, jobject);
 
@@ -44,11 +47,23 @@ bool JVMSupport::isPlatformThread(JNIEnv* jni, jthread thread) {
 
     const JniFunction* functions =
         reinterpret_cast<const JniFunction*>(jni->functions);
+    if (functions == nullptr) return false;
     IsVirtualThreadFunction is_virtual_thread =
         reinterpret_cast<IsVirtualThreadFunction>(
             functions[IS_VIRTUAL_THREAD_INDEX]);
-    return is_virtual_thread != nullptr &&
-        is_virtual_thread(jni, thread) == JNI_FALSE;
+    if (is_virtual_thread == nullptr) {
+        Counters::increment(TASK_BLOCK_VIRTUAL_THREAD_DETECTION_FAILED);
+        static std::atomic<bool> warning_emitted{false};
+        bool expected = false;
+        if (warning_emitted.compare_exchange_strong(expected, true,
+                                                     std::memory_order_relaxed)) {
+            LOG_WARN("Failed to resolve IsVirtualThread from the JNI function table "
+                     "even though JNI version is 19 or later; JVM producer callbacks "
+                     "will be ignored");
+        }
+        return false;
+    }
+    return is_virtual_thread(jni, thread) == JNI_FALSE;
 }
 
 bool JVMSupport::initialize() {

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

@@ -1541,6 +1541,27 @@ Error Profiler::init() {
   return Error::OK;
 }
 
+void Profiler::setTaskBlockEnabled(bool enabled) {
+  if (enabled) {
+    // Keep callback admission closed until native setup has either completed
+    // or rolled back, so partial event enablement cannot create paired state.
+    bool monitor_events_enabled = setNativeMonitorTaskBlockEventsEnabled(true);
+    _task_block_monitor_events_enabled.store(monitor_events_enabled,
+                                             std::memory_order_release);
+    _task_block_enabled.store(true, std::memory_order_release);
+    return;
+  }
+
+  _task_block_enabled.store(false, std::memory_order_release);
+  // Clear the admission flag first so no consumer can observe enabled events, then
+  // always attempt teardown. A previous enable whose setup AND rollback both failed
+  // left the flag false while JVMTI events stayed on; retrying unconditionally is the
+  // only way that leak is ever reclaimed. setNativeMonitorTaskBlockEventsEnabled(false) is
+  // documented as a no-op when the capability was never enabled.
+  _task_block_monitor_events_enabled.exchange(false, std::memory_order_acq_rel);
+  setNativeMonitorTaskBlockEventsEnabled(false);
+}
+
 Error Profiler::start(Arguments &args, bool reset) {
   MutexLocker ml(_state_lock);
   Error error = checkState();
@@ -1876,9 +1897,8 @@ Error Profiler::start(Arguments &args, bool reset) {
     // Paired with drainInflight() on the stop side.
     _cpu_engine->enableEvents(true);
 
-    _task_block_enabled.store(
-        (activated & EM_WALL) && args._wall_precheck && track_unfiltered_wall,
-        std::memory_order_release);
+    setTaskBlockEnabled(
+        (activated & EM_WALL) && args._wall_precheck && track_unfiltered_wall);
     _state.store(RUNNING, std::memory_order_release);
     _start_time = time(NULL);
     __atomic_add_fetch(&_epoch, 1, __ATOMIC_RELAXED);
@@ -1903,7 +1923,7 @@ Error Profiler::stop() {
   if (state() != RUNNING) {
     return Error("Profiler is not active");
   }
-  _task_block_enabled.store(false, std::memory_order_release);
+  setTaskBlockEnabled(false);
 
   // Order matters: disable engines first so the _enabled check inside signal
   // handlers will fail for any new signal delivered from now on. drain() then
@@ -2089,8 +2109,10 @@ Error Profiler::dump(const char *path, const int length) {
     Error err = Error::OK;
     // rotateDictsAndRun rotates the dictionaries, takes lockAll() around the
     // dump (fences ASGCT/JNI writers to CallTraceStorage), then clearStandby()s
-    // the rotated buffers.  StringDictionary's RefCountGuard protocol handles
-    // its own writer/reader coordination.
+    // the rotated buffers. StringDictionary's rotate()/standby()/clearStandby()
+    // protocol handles the dump's own writer/reader coordination for each
+    // dictionary, including _class_map (see flightRecorder.cpp's writeCpool:
+    // classMap()->standby() stays stable for the dump's lifetime).
     if (beginTaskBlockRotation()) {
       rotateDictsAndRun([&]{
         err = _jfr.dump(path, length);

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

@@ -133,6 +133,7 @@ class alignas(alignof(SpinLock)) Profiler {
   alignas(DEFAULT_CACHE_LINE_SIZE) u64 _failures[ASGCT_FAILURE_TYPES];
   bool _wall_precheck = false;
   std::atomic<bool> _task_block_enabled{false};
+  std::atomic<bool> _task_block_monitor_events_enabled{false};
   std::atomic<bool> _task_block_rotation{false};
   std::atomic<u64> _task_block_inflight{0};
 
@@ -185,6 +186,7 @@ class alignas(alignof(SpinLock)) Profiler {
 
   void lockAll();
   void unlockAll();
+  void setTaskBlockEnabled(bool enabled);
   bool beginTaskBlockRotation();
   void endTaskBlockRotation();
 
@@ -494,6 +496,9 @@ class alignas(alignof(SpinLock)) Profiler {
   bool taskBlockEnabled() const {
     return _task_block_enabled.load(std::memory_order_acquire);
   }
+  bool nativeMonitorTaskBlockEnabled() const {
+    return _task_block_monitor_events_enabled.load(std::memory_order_acquire);
+  }
   void writeLog(LogLevel level, const char *message);
   void writeLog(LogLevel level, const char *message, size_t len);
   void writeDatadogProfilerSetting(int tid, int length, const char *name,