fix: preserve TaskBlock stack references across JFR dumps

Author: kaahos

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

@@ -127,6 +127,7 @@
     "task_block_skipped_no_stack_reference")                                 \
   X(TASK_BLOCK_RECORD_FAILED, "task_block_record_failed")                    \
   X(TASK_BLOCK_QUEUE_DROPPED, "task_block_queue_dropped")                    \
+  X(TASK_BLOCK_DROPPED_ROTATION, "task_block_dropped_rotation")              \
   X(JVMTI_STACKS_DROPPED_LOCK, "jvmti_stacks_dropped_lock")                   \
   X(SAMPLES_DROPPED_REC_LOCK, "samples_dropped_rec_lock")                     \
   X(SAMPLES_DROPPED_THREAD_LOCAL, "samples_dropped_thread_local")

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

@@ -360,20 +360,31 @@ Java_com_datadoghq_profiler_JavaProfiler_parkExit0(
     return;
   }
 
-  int tid = ProfiledThread::currentTid();
-  if (tid >= 0) {
-    recordTaskBlockWithContextIfEligible(tid, start_ticks, TSC::ticks(),
-                                         park_context, (u64)blocker,
-                                         (u64)unblockingSpanId);
+  Profiler* profiler = Profiler::instance();
+  bool activity = profiler->tryEnterTaskBlockActivity();
+  if (!activity) {
+    profiler->waitForTaskBlockRotation();
   }
-
-  ThreadFilter *tf = Profiler::instance()->threadFilter();
-  if (tf->enabled()) {
-    ThreadFilter::SlotID slot_id = ThreadFilter::tokenSlotId(park_block_token);
-    if (current->filterSlotId() == slot_id) {
-      tf->exitBlockedRun(slot_id, ThreadFilter::tokenGeneration(park_block_token));
-    }
+  ThreadFilter* tf = profiler->threadFilter();
+  ThreadFilter::SlotID slot_id = ThreadFilter::tokenSlotId(park_block_token);
+  BlockRunSnapshot snapshot{};
+  bool exited = tf->enabled() && current->filterSlotId() == slot_id &&
+      tf->snapshotAndExitBlockedRun(
+          slot_id, ThreadFilter::tokenGeneration(park_block_token), &snapshot);
+  if (!activity) {
+    Counters::increment(TASK_BLOCK_DROPPED_ROTATION);
+    return;
   }
+  if (exited) {
+    int tid = ProfiledThread::currentTid();
+    OSThreadState observed_state = snapshot.sampled_state != OSThreadState::UNKNOWN
+        ? snapshot.sampled_state : snapshot.active_state;
+    recordTaskBlockWithStackReferenceIfEligible(
+        tid, start_ticks, TSC::ticks(), park_context, (u64)blocker,
+        (u64)unblockingSpanId, snapshot.call_trace_id,
+        snapshot.correlation_id, observed_state, true);
+  }
+  profiler->leaveTaskBlockActivity();
 }
 
 static bool decodeJavaBlockState(jint state, OSThreadState &decoded) {

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

@@ -104,13 +104,28 @@ void NativeBlockScope::finish(u64 end_ticks) {
   }
   _active = false;
 
-  ThreadFilter* thread_filter = Profiler::instance()->threadFilter();
+  Profiler* profiler = Profiler::instance();
+  bool activity = profiler->tryEnterTaskBlockActivity();
+  if (!activity) {
+    profiler->waitForTaskBlockRotation();
+  }
+  ThreadFilter* thread_filter = profiler->threadFilter();
   BlockRunSnapshot snapshot{};
   snapshot.active_state = _state;
   snapshot.sampled_state = OSThreadState::UNKNOWN;
   snapshot.owner = BlockRunOwner::NATIVE;
   if (!thread_filter->enabled() ||
       !thread_filter->snapshotAndExitBlockedRun(_slot_id, _generation, &snapshot)) {
+    if (activity) {
+      profiler->leaveTaskBlockActivity();
+    } else {
+      Counters::increment(TASK_BLOCK_DROPPED_ROTATION);
+    }
+    return;
+  }
+
+  if (!activity) {
+    Counters::increment(TASK_BLOCK_DROPPED_ROTATION);
     return;
   }
 
@@ -121,7 +136,8 @@ void NativeBlockScope::finish(u64 end_ticks) {
       _tid, _start_ticks, end_ticks, _context, _blocker, 0,
       snapshot.has_stack_reference ? snapshot.call_trace_id : 0,
       snapshot.has_stack_reference ? snapshot.correlation_id : 0,
-      observed_state);
+      observed_state, true);
+  profiler->leaveTaskBlockActivity();
 }
 
 #endif // __linux__

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

@@ -47,6 +47,7 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <thread>
 #include <vector>
 #include <signal.h>
 #include <stdint.h>
@@ -162,18 +163,58 @@ void Profiler::stopTaskBlockDrain() {
   drainTaskBlockQueue(true);
 }
 
-void Profiler::drainTaskBlockQueue(bool record) {
+void Profiler::drainTaskBlockQueue(bool record, bool rotation_owner) {
+  if (!rotation_owner && !tryEnterTaskBlockActivity()) {
+    return;
+  }
   QueuedTaskBlockEvent queued;
   u64 generation = _task_block_generation.load(std::memory_order_acquire);
   while (_task_block_queue.tryPop(queued)) {
     if (record && queued.generation == generation) {
-      if (recordTaskBlockLive(queued.tid, &queued.event)) {
+      if (recordTaskBlockLiveUnchecked(queued.tid, &queued.event)) {
         Counters::increment(TASK_BLOCK_EMITTED);
       } else {
         Counters::increment(TASK_BLOCK_RECORD_FAILED);
       }
     }
   }
+  if (!rotation_owner) {
+    leaveTaskBlockActivity();
+  }
+}
+
+bool Profiler::tryEnterTaskBlockActivity() {
+  if (_task_block_rotation.load(std::memory_order_acquire)) {
+    return false;
+  }
+  _task_block_inflight.fetch_add(1, std::memory_order_acq_rel);
+  if (_task_block_rotation.load(std::memory_order_acquire)) {
+    _task_block_inflight.fetch_sub(1, std::memory_order_acq_rel);
+    return false;
+  }
+  return true;
+}
+
+void Profiler::leaveTaskBlockActivity() {
+  _task_block_inflight.fetch_sub(1, std::memory_order_release);
+}
+
+void Profiler::waitForTaskBlockRotation() {
+  while (_task_block_rotation.load(std::memory_order_acquire)) {
+    std::this_thread::yield();
+  }
+}
+
+void Profiler::beginTaskBlockRotation() {
+  _task_block_rotation.store(true, std::memory_order_release);
+  while (_task_block_inflight.load(std::memory_order_acquire) != 0) {
+    std::this_thread::yield();
+  }
+  drainTaskBlockQueue(true, true);
+}
+
+void Profiler::endTaskBlockRotation() {
+  _task_block_rotation.store(false, std::memory_order_release);
 }
 
 void Profiler::onThreadStart(jvmtiEnv *jvmti, JNIEnv *jni, jthread thread) {
@@ -679,7 +720,9 @@ void Profiler::recordDeferredSample(int tid, u64 call_trace_id, jint event_type,
 
 bool Profiler::recordSample(void *ucontext, u64 counter, int tid,
                             jint event_type, u64 call_trace_id, Event *event,
-                            u64 *recorded_call_trace_id) {
+                            u64 *recorded_call_trace_id,
+                            ThreadFilter::SlotID task_block_slot_id,
+                            OSThreadState task_block_state) {
   atomicIncRelaxed(_total_samples);
 
   u32 lock_index = getLockIndex(tid);
@@ -748,14 +791,20 @@ bool Profiler::recordSample(void *ucontext, u64 counter, int tid,
   if (recorded && recorded_call_trace_id != nullptr) {
     *recorded_call_trace_id = call_trace_id;
   }
+  if (recorded && task_block_slot_id >= 0) {
+    _thread_filter.markTaskBlockSampled(task_block_slot_id, task_block_state,
+                                       call_trace_id, 0);
+  }
 
   _locks[lock_index].unlock();
   return recorded;
 }
 
 bool Profiler::recordSampleDelegated(void *ucontext, u64 weight, int tid,
                                      jint event_type, Event *event,
-                                     u64 *recorded_correlation_id) {
+                                     u64 *recorded_correlation_id,
+                                     ThreadFilter::SlotID task_block_slot_id,
+                                     OSThreadState task_block_state) {
   if (!VM::canRequestStackTrace()) {
     return false;
   }
@@ -793,6 +842,10 @@ bool Profiler::recordSampleDelegated(void *ucontext, u64 weight, int tid,
   if (recorded && recorded_correlation_id != nullptr) {
     *recorded_correlation_id = correlation_id;
   }
+  if (recorded && task_block_slot_id >= 0) {
+    _thread_filter.markTaskBlockSampled(task_block_slot_id, task_block_state,
+                                       0, correlation_id);
+  }
   _locks[lock_index].unlock();
   return recorded;
 }
@@ -830,7 +883,7 @@ void Profiler::recordQueueTime(int tid, QueueTimeEvent *event) {
   _locks[lock_index].unlock();
 }
 
-bool Profiler::recordTaskBlockLive(int tid, TaskBlockEvent *event) {
+bool Profiler::recordTaskBlockLiveUnchecked(int tid, TaskBlockEvent *event) {
 #ifdef UNIT_TEST
   record_task_block_live_calls_for_test.fetch_add(1, std::memory_order_acq_rel);
   last_task_block_tid_for_test.store(tid, std::memory_order_release);
@@ -852,6 +905,10 @@ bool Profiler::recordTaskBlockLive(int tid, TaskBlockEvent *event) {
   return recorded;
 }
 
+bool Profiler::recordTaskBlockLive(int tid, TaskBlockEvent *event) {
+  return recordTaskBlockLiveUnchecked(tid, event);
+}
+
 bool Profiler::recordTaskBlockAsync(int tid, TaskBlockEvent *event) {
   if (!_task_block_drain_running.load(std::memory_order_acquire)) {
     Counters::increment(TASK_BLOCK_QUEUE_DROPPED);
@@ -1858,10 +1915,12 @@ Error Profiler::dump(const char *path, const int length) {
     // its own writer/reader coordination; #527's classMapSharedGuard readers
     // (deferred vtable receiver resolution) are coordinated through
     // _class_map_lock.
+    beginTaskBlockRotation();
     rotateDictsAndRun([&]{
       err = _jfr.dump(path, length);
       __atomic_add_fetch(&_epoch, 1, __ATOMIC_SEQ_CST);
     });
+    endTaskBlockRotation();
 
     _thread_info.clearAll(thread_ids);
     _thread_info.reportCounters();

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

@@ -127,6 +127,10 @@ class alignas(alignof(SpinLock)) Profiler {
   pthread_t _task_block_drain_thread;
   std::atomic<bool> _task_block_drain_running;
   std::atomic<u64> _task_block_generation;
+  // Dump closes this gate, waits for active producers/consumers, then drains the
+  // queue before rotating call-trace storage.
+  std::atomic<bool> _task_block_rotation;
+  std::atomic<u32> _task_block_inflight;
 
   // dlopen() hook support
   void **_dlopen_entry;
@@ -141,7 +145,10 @@ class alignas(alignof(SpinLock)) Profiler {
   static void *taskBlockDrainLoop(void *arg);
   void startTaskBlockDrain();
   void stopTaskBlockDrain();
-  void drainTaskBlockQueue(bool record);
+  void drainTaskBlockQueue(bool record, bool rotation_owner = false);
+  bool recordTaskBlockLiveUnchecked(int tid, TaskBlockEvent *event);
+  void beginTaskBlockRotation();
+  void endTaskBlockRotation();
 
   u32 getLockIndex(int tid);
   int getNativeTrace(void *ucontext, ASGCT_CallFrame *frames, int event_type,
@@ -208,11 +215,16 @@ class alignas(alignof(SpinLock)) Profiler {
         _num_context_attributes(0), _omit_stacktraces(false),
         _remote_symbolication(false), _task_block_drain_thread(),
         _task_block_drain_running(false), _task_block_generation(0),
+        _task_block_rotation(false), _task_block_inflight(0),
         _dlopen_entry(NULL) {
 
     for (int i = 0; i < CONCURRENCY_LEVEL; i++) {
       _calltrace_buffer[i] = NULL;
     }
+    _call_trace_storage.registerLivenessChecker(
+        [this](std::unordered_set<u64>& trace_ids) {
+          _thread_filter.collectTaskBlockLiveTraceIds(trace_ids);
+        });
   }
 
   static inline Profiler *instance() {
@@ -397,15 +409,19 @@ class alignas(alignof(SpinLock)) Profiler {
                          ASGCT_CallFrame *frames, int lock_index);
   bool recordSample(void *ucontext, u64 weight, int tid, jint event_type,
                     u64 call_trace_id, Event *event,
-                    u64 *recorded_call_trace_id = nullptr);
+                    u64 *recorded_call_trace_id = nullptr,
+                    ThreadFilter::SlotID task_block_slot_id = -1,
+                    OSThreadState task_block_state = OSThreadState::UNKNOWN);
   // Delegated sample path: stack-walking is performed by the HotSpot JFR
   // RequestStackTrace extension (the JVM emits the stack trace into its own
   // JFR recording). We only emit the CPU/wall sample event with no
   // stack-trace reference, tagged by the correlation ID we passed to
   // RequestStackTrace as user_data.
   bool recordSampleDelegated(void *ucontext, u64 weight, int tid,
                              jint event_type, Event *event,
-                             u64 *recorded_correlation_id = nullptr);
+                             u64 *recorded_correlation_id = nullptr,
+                             ThreadFilter::SlotID task_block_slot_id = -1,
+                             OSThreadState task_block_state = OSThreadState::UNKNOWN);
   u64 recordJVMTISample(u64 weight, int tid, jthread thread, jint event_type, Event *event, bool deferred);
   void recordDeferredSample(int tid, u64 call_trace_id, jint event_type, Event *event);
   void recordExternalSample(u64 weight, int tid, int num_frames,
@@ -416,6 +432,11 @@ class alignas(alignof(SpinLock)) Profiler {
   void recordQueueTime(int tid, QueueTimeEvent *event);
   bool recordTaskBlockLive(int tid, TaskBlockEvent *event);
   bool recordTaskBlockAsync(int tid, TaskBlockEvent *event);
+  // TaskBlock producers must hold an activity across stack-reference snapshot
+  // and event recording/enqueue so dump cannot rotate between those operations.
+  bool tryEnterTaskBlockActivity();
+  void leaveTaskBlockActivity();
+  void waitForTaskBlockRotation();
   bool taskBlockAsyncActive() const {
     return _task_block_drain_running.load(std::memory_order_acquire);
   }
@@ -474,6 +495,14 @@ class alignas(alignof(SpinLock)) Profiler {
     drainTaskBlockQueue(record);
   }
 
+  void beginTaskBlockRotationForTest() {
+    beginTaskBlockRotation();
+  }
+
+  void endTaskBlockRotationForTest() {
+    endTaskBlockRotation();
+  }
+
   void startTaskBlockDrainForTest() {
     startTaskBlockDrain();
   }