Restore priming

AGENTS.md

@@ -288,6 +288,11 @@ The profiler uses a sophisticated double-buffered storage system for call traces
 - **Atomic Operations**: Instance ID management and counter updates use atomics
 - **Memory Allocation**: Minimize malloc() in hot paths, use pre-allocated containers
 
+### Sampler Safety
+- **Stack walker**: `HotspotSupport::walkVM()`, `StackWalker::walkDwarf()`, and `StackWalker::walkFP()` must be protected by `sigsetjmp()`/`siglongjmp()`.
+- **Samplers**: Every sampler must set up the `ProfiledThread` thread-local before sampling, and skip the sample if it isn't available. Signal-based samplers 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.
 - **Cross-thread reads**: Always use `__ATOMIC_ACQUIRE` for loads that must see stores from another thread. Never use `__ATOMIC_RELAXED` for cross-thread visibility unless you can prove no ordering dependency exists.

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

@@ -134,6 +134,7 @@
   X(SAFECOPY_FAILED, "safecopy_failed")                                       \
   X(SAFEFETCH_FAILED, "safefetch_failed")                                     \
   X(STACKWALK_LONGJMP_RECOVERED, "stackwalk_longjmp_recovered")               \
+  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>
@@ -225,19 +226,19 @@ void CTimerJvmti::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
     return;
   }
   int tid = 0;
-  ProfiledThread *current = ProfiledThread::current();
-  assert(current == nullptr || !current->isDeepCrashHandler());
-  if (current != nullptr && JVMThread::current() == nullptr
+  ProfiledThread *current = SIGNAL_HANDLER_CURRENT_THREAD();
+  assert(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();
   }
   Shims::instance().setSighandlerTid(tid);
 
@@ -267,6 +268,8 @@ void CTimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   Counters::increment(CTIMER_SIGNAL_OWN);
 
   InflightGuard inflight;
+  ProfiledThread* current = SIGNAL_HANDLER_CURRENT_THREAD();
+  assert(current != nullptr);
 
   // Atomically try to enter critical section - prevents all reentrancy races
   CriticalSection cs;
@@ -280,25 +283,18 @@ void CTimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   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,39 @@ 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_ADDRESS_HIGH(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_INT_HIGH(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_LONG_HIGH(v)     (v)
 
 #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.cpp

@@ -28,6 +28,7 @@
 #include "unwindStats.h"
 #include "symbols.h"
 #include "threadFilter.h"
+#include "threadLocalData.inline.h"
 #include "threadState.h"
 #include "tsc.h"
 #include "hotspot/vmStructs.h"

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

@@ -17,7 +17,9 @@
 #include "guards.h"
 #include "common.h"
 #include "os.h"
-#include "threadLocalData.h"
+#include "threadLocalData.inline.h"
+
+#include <cassert>
 
 // Signal-context tracking — backed by ProfiledThread::_signal_depth; see
 // the comment block in guards.h for the rationale (initial-exec TLS was
@@ -30,15 +32,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() : _active(true), _current(nullptr) {
+    ProfiledThread *pt = ProfiledThread::acquireCurrent();
     if (pt != nullptr) {
+        _current = pt;
         pt->enterSignalScope();
     } else {
         // No thread context: nothing to update; mark inactive so destructor
@@ -49,9 +53,8 @@ SignalHandlerScope::SignalHandlerScope() : _active(true) {
 
 SignalHandlerScope::~SignalHandlerScope() {
     if (!_active) return;
-    ProfiledThread *pt = ProfiledThread::current();
-    if (pt != nullptr) {
-        pt->exitSignalScope();
+    if (_current != nullptr) {
+        _current->exitSignalScope();
     }
 }
 
@@ -71,46 +74,16 @@ 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) {
+CriticalSection::CriticalSection() : _entered(false), _thread_ptr(nullptr) {
     _thread_ptr = ProfiledThread::current();
-    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
-    }
+    assert(_thread_ptr != nullptr);
+    _entered = _thread_ptr->tryEnterCriticalSection();
 }
 
 CriticalSection::~CriticalSection() {
+    assert(_thread_ptr != nullptr);
     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();
-            }
-        }
+        _thread_ptr->exitCriticalSection();
     }
 }
-
-uint32_t CriticalSection::hash_tid(int tid) {
-    return static_cast<uint32_t>(tid * KNUTH_MULTIPLICATIVE_CONSTANT);
-}

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

@@ -22,6 +22,8 @@
 #include <signal.h>
 #include <pthread.h>
 
+#include "counters.h"
+
 class ProfiledThread;
 
 // ---------------------------------------------------------------------------
@@ -42,8 +44,8 @@ class ProfiledThread;
 // pthread_getspecific (POSIX guarantees it does not allocate; returns
 // nullptr when unset).
 //
-// When ProfiledThread is null on a thread we don't yet have a thread
-// context — uninstrumented JVM-internal threads (VM Thread, JIT, GC) fall
+// When ProfiledThread is null or via thread priming on a thread
+// — uninstrumented JVM-internal threads (VM Thread, JIT, GC) fall
 // into this bucket too, and they can receive signals.  The
 // SignalHandlerScope guard is a no-op on those threads (nothing to
 // update), so isInTrackedSignalContext() returns false: production code
@@ -80,14 +82,32 @@ class SignalHandlerScope {
     void release();
     SignalHandlerScope(const SignalHandlerScope&)            = delete;
     SignalHandlerScope& operator=(const SignalHandlerScope&) = delete;
+
+    bool isActive() const { return _active; }
+    ProfiledThread* current() const { return _current; }
 private:
+    ProfiledThread* _current;
     bool _active;
 };
 
 // Declare a scope guard local that increments the depth on entry and
 // decrements on scope exit.  Use as the very first statement in every
-// installed signal handler.
-#define SIGNAL_HANDLER_GUARD() SignalHandlerScope _signal_handler_scope
+// installed sampler signal handler
+#define SIGNAL_HANDLER_GUARD()                              \
+      SignalHandlerScope _signal_handler_scope;             \
+      if (!_signal_handler_scope.isActive()) {              \
+        Counters::increment(SAMPLES_DROPPED_THREAD_LOCAL);  \
+        return;                                             \
+      }
+
+// Declare a scope guard local that increments the depth on entry and
+// decrements on scope exit.  Use as the very first statement in every
+// installed non-sampler signal handler
+#define SIGNAL_HANDLER_GUARD_NO_SAMPLE()                    \
+      SignalHandlerScope _signal_handler_scope;
+
+// Cheaper way to retrieve current ProfiledThread inside the scope
+#define SIGNAL_HANDLER_CURRENT_THREAD() _signal_handler_scope.current()
 
 // Manually release the most recent SIGNAL_HANDLER_GUARD() before chaining to
 // another handler that may siglongjmp through us (e.g. J9's SIGSEGV null-pointer
@@ -136,17 +156,7 @@ void signalHandlerUnwindAfterLongjmp();
  */
 class CriticalSection {
 private:
-    static constexpr size_t FALLBACK_BITMAP_WORDS = 1024;  // 8KB for 64K bits
-    // Atomic bitmap for thread-safe critical section tracking without TLS
-    // Must be atomic because multiple signal handlers can run concurrently across
-    // different threads and attempt to set/clear bits simultaneously. Compare-and-swap
-    // operations ensure race-free bit manipulation even during signal interruption.
-    static uint64_t _fallback_bitmap[FALLBACK_BITMAP_WORDS];
-
     bool _entered;          // Track if this instance successfully entered
-    bool _using_fallback;   // Track which storage mechanism we're using
-    uint32_t _word_index;   // For fallback bitmap cleanup
-    uint64_t _bit_mask;     // For fallback bitmap cleanup
     ProfiledThread* _thread_ptr; // ProfiledThread captured at construction
 
 public:
@@ -161,10 +171,6 @@ class CriticalSection {
 
     // Check if this instance successfully entered the critical section
     bool entered() const { return _entered; }
-
-private:
-    // Hash function to distribute thread IDs across bitmap words
-    static uint32_t hash_tid(int tid);
 };
 
 /**