Restore priming and fix races on acquire/release ProfiledThread

AGENTS.md

@@ -290,6 +290,13 @@ The profiler uses a sophisticated double-buffered storage system for call traces
 - **Thread-local Buffers**: Per-thread recording buffers minimize contention
 - **Atomic Operations**: Instance ID management and counter updates use atomics
 - **Memory Allocation**: Minimize malloc() in hot paths, use pre-allocated containers
+- **Thread termination**: A thread must call `blockProfilingForExit()` before releasing its `ProfiledThread` on exit — otherwise a profiling signal can race the release and allocate a new `ProfiledThread` that can never be freed (leak).
+
+### Sampler Safety
+- **Sampled thread**: The sampled thread must have a `ProfiledThread` in TLS, since it owns the `sigjmp_buf` used to recover via `siglongjmp()` if the stack walker crashes. If none is available, the sampler must drop the sample and report `SAMPLES_DROPPED_THREAD_LOCAL`.
+- **Samplers**: Every sampler must set up the `ProfiledThread` TLS before sampling, and skip the sample if it isn't available. Signal-based or allocation-adjacent hooks sampler (e.g. malloc/free interceptors) must use `ProfiledThread::acquireCurrent()`; non-signal-based samplers use `ProfiledThread::initCurrentThreadSignalSafe()`.
+- **Samplers**: Every sampler must set up the `ProfiledThread` TLS before sampling, and skip the sample if it isn't available. Signal-based samplers and allocation-adjacent hooks (e.g. malloc/free interceptors) must use `ProfiledThread::acquireCurrent()`; non-signal-based samplers use `ProfiledThread::initCurrentThreadSignalSafe()`.
+- **JNI/JVMTI callbacks**: Use `ProfiledThread::initCurrentThreadSignalSafe()` to set up `ProfiledThread` for the thread.
 
 ### Atomic Memory Ordering (Critical for arm64)
 arm64 has a weakly-ordered memory model (unlike x86 TSO). Incorrect ordering causes real lockups on arm64 that never reproduce on x86.

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

@@ -52,14 +52,20 @@
 constexpr size_t KNUTH_MULTIPLICATIVE_CONSTANT = 0x9e3779b97f4a7c15ULL;
 
 #ifdef DEBUG
-#define DEBUG_ONLY(s) s
-#define TEST_LOG(fmt, ...) do { \
-  fprintf(stdout, "[TEST::INFO] " fmt "\n", ##__VA_ARGS__); \
-  fflush(stdout); \
-} while (0)
+    #define DEBUG_ONLY(s) s
+    #define TEST_LOG(fmt, ...) do { \
+      fprintf(stdout, "[TEST::INFO] " fmt "\n", ##__VA_ARGS__); \
+      fflush(stdout); \
+    } while (0)
+#else
+    #define DEBUG_ONLY(s)
+    #define TEST_LOG(fmt, ...) // No-op in non-debug mode
+#endif
+
+#ifdef __FAULT_INJECTION__
+    #define FAULT_INJECTION_ONLY(s) s
 #else
-#define DEBUG_ONLY(s)
-#define TEST_LOG(fmt, ...) // No-op in non-debug mode
+    #define FAULT_INJECTION_ONLY(s)
 #endif
 
 // Lightweight stderr warning that does not depend on the Log subsystem.

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

@@ -19,7 +19,7 @@
 #include "guards.h"
 #include "otel_context.h"
 #include "profiler.h"
-#include "threadLocalData.h"
+#include "threadLocalData.inline.h"
 #include <cstring>
 
 /**

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 Datadog, Inc
+ * Copyright 2023, 2026 Datadog, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -134,6 +134,15 @@
   X(SAFECOPY_FAILED, "safecopy_failed")                                       \
   X(SAFEFETCH_FAILED, "safefetch_failed")                                     \
   X(STACKWALK_LONGJMP_RECOVERED, "stackwalk_longjmp_recovered")               \
+  /* Strict subset of SAMPLES_DROPPED_THREAD_LOCAL, not an independent count: \
+   * ThreadLocalDataPool::claim() increments this on capacity exhaustion, and \
+   * every acquireCurrent() caller that gets nullptr back -- for this or any  \
+   * other reason -- separately increments SAMPLES_DROPPED_THREAD_LOCAL too.  \
+   * So every pool-exhaustion drop bumps both counters together; the two     \
+   * should be subtracted (thread_local_pool_exhausted from                  \
+   * samples_dropped_thread_local) to isolate non-pool priming drops, never  \
+   * summed. */                                                              \
+  X(SAMPLES_DROPPED_TLS_POOL_EXHAUSTED, "thread_local_pool_exhausted")        \
   /* writeElement() guards against a corrupted/dangling JfrMetadata tree.     \
    * Root cause is still unconfirmed, so these counters are the durable       \
    * signal for spotting a recurrence. */                                     \

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

@@ -27,6 +27,7 @@
 #include "log.h"
 #include "profiler.h"
 #include "signalCookie.h"
+#include "threadLocalData.inline.h"
 #include "threadState.inline.h"
 #include <assert.h>
 #include <errno.h>
@@ -205,40 +206,40 @@ Error CTimerJvmti::start(Arguments &args) {
 }
 
 void CTimerJvmti::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
-  SIGNAL_HANDLER_GUARD();
+  int saved_errno = errno;
   if (!OS::shouldProcessSignal(siginfo, SI_TIMER, SignalCookie::cpu())) {
     Counters::increment(CTIMER_SIGNAL_FOREIGN);
     OS::forwardForeignSignal(signo, siginfo, ucontext);
     return;
   }
   Counters::increment(CTIMER_SIGNAL_OWN);
 
+  SIGNAL_HANDLER_GUARD_OR_DROP_WITH_ERRNO(saved_errno);
   InflightGuard inflight;
+  ProfiledThread *current = SIGNAL_HANDLER_CURRENT_THREAD();
+  assert(!current->isDeepCrashHandler());
 
-  CriticalSection cs;
+  CriticalSection cs(current);
   if (!cs.entered()) {
+    errno = saved_errno;
     return;
   }
-  int saved_errno = errno;
   if (!__atomic_load_n(&_enabled, __ATOMIC_ACQUIRE)) {
     errno = saved_errno;
     return;
   }
   int tid = 0;
-  ProfiledThread *current = ProfiledThread::current();
-  assert(current == nullptr || !current->isDeepCrashHandler());
-  if (current != nullptr && JVMThread::current() == nullptr
+
+  if (JVMThread::current() == nullptr
       && current->inInitWindow()) {
     current->tickInitWindow();
     errno = saved_errno;
     return;
   }
-  if (current != NULL) {
-    current->noteCPUSample(Profiler::instance()->recordingEpoch());
-    tid = current->tid();
-  } else {
-    tid = OS::threadId();
-  }
+
+  current->noteCPUSample(Profiler::instance()->recordingEpoch());
+  tid = current->tid();
+
   Shims::instance().setSighandlerTid(tid);
 
   ExecutionEvent event;
@@ -254,7 +255,8 @@ void CTimerJvmti::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
 }
 
 void CTimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
-  SIGNAL_HANDLER_GUARD();
+  int saved_errno = errno;
+
   // Reject signals that did not originate from our timer_create timers.
   // This guards against Go's process-wide setitimer(ITIMER_PROF) and other
   // foreign SIGPROF sources that would otherwise drive our handler onto
@@ -266,39 +268,33 @@ void CTimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   }
   Counters::increment(CTIMER_SIGNAL_OWN);
 
+  SIGNAL_HANDLER_GUARD_OR_DROP_WITH_ERRNO(saved_errno);
   InflightGuard inflight;
+  ProfiledThread* current = SIGNAL_HANDLER_CURRENT_THREAD();
+  assert(current != nullptr);
 
   // Atomically try to enter critical section - prevents all reentrancy races
-  CriticalSection cs;
+  CriticalSection cs(current);
   if (!cs.entered()) {
     return;  // Another critical section is active, defer profiling
   }
-  // Save the current errno value
-  int saved_errno = errno;
   // we want to ensure memory order because of the possibility the instance gets
   // cleared
   if (!__atomic_load_n(&_enabled, __ATOMIC_ACQUIRE)) {
     return;
   }
-  int tid = 0;
-  ProfiledThread *current = ProfiledThread::current();
-  assert(current == nullptr || !current->isDeepCrashHandler());
+  assert(!current->isDeepCrashHandler());
   // Guard against the race window between Profiler::registerThread() and
   // thread_native_entry setting JVM TLS (PROF-13072): skip at most one signal
   // per thread. Pure native threads (where JVMThread::current() is always null)
   // are allowed through once the one-shot window expires.
-  if (current != nullptr && JVMThread::current() == nullptr
-      && current->inInitWindow()) {
+  if (JVMThread::current() == nullptr && current->inInitWindow()) {
     current->tickInitWindow();
     errno = saved_errno;
     return;
   }
-  if (current != NULL) {
-    current->noteCPUSample(Profiler::instance()->recordingEpoch());
-    tid = current->tid();
-  } else {
-    tid = OS::threadId();
-  }
+  current->noteCPUSample(Profiler::instance()->recordingEpoch());
+  int tid = current->tid();
   Shims::instance().setSighandlerTid(tid);
 
   ExecutionEvent event;

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

@@ -22,7 +22,7 @@
 
 #include "counters.h"          // Counters::increment (FAULTS_INJECTED)
 #include "os.h"                // OS::page_size
-#include "threadLocalData.h"   // ProfiledThread::current / nextFiRandom
+#include "threadLocalData.inline.h"   // ProfiledThread::current / nextFiRandom
 #include <atomic>
 #include <sys/mman.h>
 

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

@@ -34,8 +34,9 @@
 //
 //   return INJECT_FAULT_BOOL_LIKELY(dlopen(name, flags) != nullptr);
 //
-// The three tiers name their firing frequency: RARE 0.01%, UNLIKELY 0.1%,
-// LIKELY 1%.  See faultInjection.cpp for the poison-address and PRNG details.
+// The four tiers name their firing frequency: RARE 0.01%, UNLIKELY 0.1%,
+// LIKELY 1%, HIGH 10%.  See faultInjection.cpp for the poison-address and PRNG
+// details.
 
 #ifndef _FAULT_INJECTION_H
 #define _FAULT_INJECTION_H
@@ -56,6 +57,7 @@ namespace faultinj {
 constexpr u64 PROB_RARE     = 1844674407370955ULL;    // 1e-4  (0.01%)
 constexpr u64 PROB_UNLIKELY = 18446744073709552ULL;   // 1e-3  (0.1%)
 constexpr u64 PROB_LIKELY   = 184467440737095520ULL;  // 1e-2  (1%)
+constexpr u64 PROB_HIGH     = 1844674407370955162ULL; // 1e-1  (10%)
 
 // Called once at profiler startup (off the signal path) to mmap the PROT_NONE
 // guard region used by poisonAddress().  Safe to call before any injection.
@@ -106,31 +108,29 @@ inline T injectValue(T orig, T faulty, u64 threshold, const char* fn) {
     ::faultinj::injectAddress((ptr), ::faultinj::PROB_UNLIKELY, __func__)
 #define INJECT_FAULT_ADDRESS_LIKELY(ptr) \
     ::faultinj::injectAddress((ptr), ::faultinj::PROB_LIKELY, __func__)
+#define INJECT_FAULT_ADDRESS_HIGH(ptr) \
+    ::faultinj::injectAddress((ptr), ::faultinj::PROB_HIGH, __func__)
 
 #define INJECT_FAULT_BOOL_RARE(v) \
     ::faultinj::injectValue((v), false, ::faultinj::PROB_RARE, __func__)
 #define INJECT_FAULT_BOOL_UNLIKELY(v) \
     ::faultinj::injectValue((v), false, ::faultinj::PROB_UNLIKELY, __func__)
 #define INJECT_FAULT_BOOL_LIKELY(v) \
     ::faultinj::injectValue((v), false, ::faultinj::PROB_LIKELY, __func__)
+#define INJECT_FAULT_BOOL_HIGH(v) \
+    ::faultinj::injectValue((v), false, ::faultinj::PROB_HIGH, __func__)
 
 #else  // __FAULT_INJECTION__ not defined — strict identity, zero cost.
 
 #define INJECT_FAULT_ADDRESS_RARE(ptr)     (ptr)
 #define INJECT_FAULT_ADDRESS_UNLIKELY(ptr) (ptr)
 #define INJECT_FAULT_ADDRESS_LIKELY(ptr)   (ptr)
-
-#define INJECT_FAULT_INT_RARE(v)     (v)
-#define INJECT_FAULT_INT_UNLIKELY(v) (v)
-#define INJECT_FAULT_INT_LIKELY(v)   (v)
-
-#define INJECT_FAULT_LONG_RARE(v)     (v)
-#define INJECT_FAULT_LONG_UNLIKELY(v) (v)
-#define INJECT_FAULT_LONG_LIKELY(v)   (v)
+#define INJECT_FAULT_ADDRESS_HIGH(ptr)     (ptr)
 
 #define INJECT_FAULT_BOOL_RARE(v)     (v)
 #define INJECT_FAULT_BOOL_UNLIKELY(v) (v)
 #define INJECT_FAULT_BOOL_LIKELY(v)   (v)
+#define INJECT_FAULT_BOOL_HIGH(v)     (v)
 
 #define NO_INJECTION_ASSERT(a) (assert(a))
 

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

@@ -19,6 +19,7 @@
 #include "arch.h"
 #include "arguments.h"
 #include "buffers.h"
+#include "common.h"
 #include "countingAllocator.h"
 #include "counters.h"
 #include "dictionary.h"

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

@@ -17,7 +17,19 @@
 #include "guards.h"
 #include "common.h"
 #include "os.h"
-#include "threadLocalData.h"
+#include "threadLocalData.inline.h"
+
+#include <cassert>
+
+void blockProfilingForExit() {
+    sigset_t prof_signals, old_signals;
+    sigemptyset(&prof_signals);
+    sigaddset(&prof_signals, SIGPROF);     // Used by ITimer and CTimer
+    sigaddset(&prof_signals, SIGVTALRM);   // Used by WallClock
+
+    int rc= pthread_sigmask(SIG_BLOCK, &prof_signals, &old_signals);
+    assert(rc == 0);
+}
 
 // Signal-context tracking — backed by ProfiledThread::_signal_depth; see
 // the comment block in guards.h for the rationale (initial-exec TLS was
@@ -30,15 +42,17 @@ int getInSignalDepth() {
 
 bool isInTrackedSignalContext() {
     ProfiledThread *pt = ProfiledThread::current();
-    // null ProfiledThread = no thread context; the SignalHandlerScope
-    // never ran, so we have no positive evidence of a signal frame.
+    // null ProfiledThread = no thread context;
+    // the SignalHandlerScope never ran, so we have no positive evidence
+    // of a signal frame.
     // See header comment for the rationale of returning false here.
     return pt != nullptr && pt->signalDepth() != 0;
 }
 
-SignalHandlerScope::SignalHandlerScope() : _active(true) {
-    ProfiledThread *pt = ProfiledThread::current();
+SignalHandlerScope::SignalHandlerScope(bool shouldRunPriming) : _current(nullptr), _active(true) {
+    ProfiledThread *pt = shouldRunPriming ? ProfiledThread::acquireCurrent() : ProfiledThread::current();
     if (pt != nullptr) {
+        _current = pt;
         pt->enterSignalScope();
     } else {
         // No thread context: nothing to update; mark inactive so destructor
@@ -49,17 +63,15 @@ SignalHandlerScope::SignalHandlerScope() : _active(true) {
 
 SignalHandlerScope::~SignalHandlerScope() {
     if (!_active) return;
-    ProfiledThread *pt = ProfiledThread::current();
-    if (pt != nullptr) {
-        pt->exitSignalScope();
+    if (_current != nullptr) {
+        _current->exitSignalScope();
     }
 }
 
 void SignalHandlerScope::release() {
     if (!_active) return;
-    ProfiledThread *pt = ProfiledThread::current();
-    if (pt != nullptr) {
-        pt->exitSignalScope();
+    if (_current != nullptr) {
+        _current->exitSignalScope();
     }
     _active = false;
 }
@@ -71,46 +83,24 @@ void signalHandlerUnwindAfterLongjmp() {
     }
 }
 
-// Static bitmap storage for fallback cases
-uint64_t CriticalSection::_fallback_bitmap[CriticalSection::FALLBACK_BITMAP_WORDS] = {};
 
-CriticalSection::CriticalSection() : _entered(false), _using_fallback(false), _word_index(0), _bit_mask(0), _thread_ptr(nullptr) {
-    _thread_ptr = ProfiledThread::current();
+CriticalSection::CriticalSection(ProfiledThread* pt) : _entered(false), _thread_ptr(pt) {
+    // acquireCurrent() falls back to ThreadLocalDataPool::acquire() (a
+    // pre-allocated, async-signal-safe pool) when the calling thread has
+    // not been primed yet. _thread_ptr can still legitimately be nullptr
+    // here if priming is unsupported (e.g. macOS) and the pool is
+    // exhausted; treat that as "did not enter" rather than dereferencing.
+    if (_thread_ptr == nullptr) {
+        _thread_ptr = ProfiledThread::acquireCurrent();
+    }
+
     if (_thread_ptr != nullptr) {
-        // Primary path: Use ProfiledThread storage (fast and memory-efficient)
         _entered = _thread_ptr->tryEnterCriticalSection();
-    } else {
-        // Fallback path: Use hash-based bitmap for stress tests and edge cases
-        _using_fallback = true;
-        int tid = OS::threadId();
-
-        // Hash TID to distribute across bitmap words, reducing clustering
-        // We are OK with false collision for the fallback - it should be used only for testing when we don't have full profiler initialized
-        _word_index = hash_tid(tid) % FALLBACK_BITMAP_WORDS;
-        uint32_t bit_index = tid % 64;
-        _bit_mask = 1ULL << bit_index;
-
-        // Use ACQUIRE ordering to ensure visibility of protected data after acquiring critical section
-        uint64_t old_word = __atomic_fetch_or(&_fallback_bitmap[_word_index], _bit_mask, __ATOMIC_ACQUIRE);
-        _entered = !(old_word & _bit_mask);  // Success if bit was previously 0
     }
 }
 
 CriticalSection::~CriticalSection() {
-    if (_entered) {
-        if (_using_fallback) {
-            // Clear the bit atomically for fallback bitmap
-            // Use RELEASE ordering to ensure protected data writes are visible before releasing
-            __atomic_fetch_and(&_fallback_bitmap[_word_index], ~_bit_mask, __ATOMIC_RELEASE);
-        } else {
-            // Release ProfiledThread flag using the pointer captured at construction
-            if (_thread_ptr != nullptr) {
-                _thread_ptr->exitCriticalSection();
-            }
-        }
+    if (_entered && _thread_ptr != nullptr) {
+        _thread_ptr->exitCriticalSection();
     }
 }
-
-uint32_t CriticalSection::hash_tid(int tid) {
-    return static_cast<uint32_t>(tid * KNUTH_MULTIPLICATIVE_CONSTANT);
-}