feat: Adopt ThreadLocal for storing ProfiledThread

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

@@ -9,11 +9,10 @@
 #include "log.h"
 #include "os.h"
 #include "common.h"
-#include "thread.h"
 #include "vmEntry.h" // For BCI_ERROR constant
 #include "arch.h" // For LP64_ONLY macro and COMMA macro
 #include "guards.h" // For table swap critical sections
-#include "thread.h"
+#include "threadLocalData.h"
 #include <string.h>
 #include <atomic>
 

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

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

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

@@ -121,7 +121,8 @@
   /* Delegated stacks dropped at slot-lock. Rec-lock drops from all recording  \
    * paths (delegated and direct) go into SAMPLES_DROPPED_REC_LOCK. */         \
   X(JVMTI_STACKS_DROPPED_LOCK, "jvmti_stacks_dropped_lock")                   \
-  X(SAMPLES_DROPPED_REC_LOCK, "samples_dropped_rec_lock")
+  X(SAMPLES_DROPPED_REC_LOCK, "samples_dropped_rec_lock")                     \
+  X(SAMPLES_DROPPED_THREAD_LOCAL, "samples_dropped_thread_local")
 #define X_ENUM(a, b) a,
 typedef enum CounterId : int {
   DD_COUNTER_TABLE(X_ENUM) DD_NUM_COUNTERS

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

@@ -17,7 +17,7 @@
 #include "guards.h"
 #include "common.h"
 #include "os.h"
-#include "thread.h"
+#include "threadLocalData.h"
 
 // Signal-context tracking — backed by ProfiledThread::_signal_depth; see
 // the comment block in guards.h for the rationale (initial-exec TLS was
@@ -103,9 +103,13 @@ CriticalSection::~CriticalSection() {
             // Use RELEASE ordering to ensure protected data writes are visible before releasing
             __atomic_fetch_and(&_fallback_bitmap[_word_index], ~_bit_mask, __ATOMIC_RELEASE);
         } else {
+            // Allows thread exiting inside CriticalSection
+            ProfiledThread* pt = ProfiledThread::currentSignalSafe();
+            assert(pt == nullptr || pt == _thread_ptr);
+
             // Release ProfiledThread flag using the pointer captured at construction
-            if (_thread_ptr != nullptr) {
-                _thread_ptr->exitCriticalSection();
+            if (pt != nullptr) {
+                pt->exitCriticalSection();
             }
         }
     }

ddprof-lib/src/main/cpp/hotspot/hotspotSupport.cpp

@@ -6,20 +6,20 @@
 
 #include <climits>
 #include <cstdlib>
-#include <setjmp.h>
 #include "asyncSampleMutex.h"
+#include "frames.h"
+#include "guards.h"
 #include "hotspot/hotspotSupport.h"
 #include "hotspot/jitCodeCache.h"
 #include "hotspot/vmStructs.inline.h"
 #include "jvmSupport.inline.h"
-#include "guards.h"
+#include "profiler.h"
 #include "stackWalker.inline.h"
-#include "frames.h"
+#include "threadLocal.h"
 
 using StackWalkValidation::inDeadZone;
 using StackWalkValidation::aligned;
 using StackWalkValidation::MAX_FRAME_SIZE;
-using StackWalkValidation::sameStack;
 
 // Initialize once, they survive on profiler restart
 static jobject JAVA_PLATFORM_CLASSLOADER = nullptr;
@@ -227,15 +227,46 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
 __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontext, ASGCT_CallFrame* frames, int max_depth,
                         StackWalkFeatures features, EventType event_type,
                         const void* pc, uintptr_t sp, uintptr_t fp, int lock_index, bool* truncated) {
+
     // VMStructs is only available for hotspot JVM 
     assert(VM::isHotspot());
+
+    ProfiledThread* prof_thread = ProfiledThread::currentSignalSafe();
+    if (prof_thread == nullptr) {
+        Counters::increment(SAMPLES_DROPPED_THREAD_LOCAL);
+        return 0;
+    }
+
     HotspotStackFrame frame(ucontext);
     uintptr_t bottom = (uintptr_t)&frame + MAX_WALK_SIZE;
 
     Profiler* profiler = Profiler::instance();
     int bcp_offset = InterpreterFrame::bcp_offset();
 
+
     jmp_buf crash_protection_ctx;
+    // Chaining jmp_buf
+    // A non-signal-based-sampler can be interrupted by signal based sampler,
+    // then we end up multiple HotspotSupport::walkVM() calls on stack,
+    // each one sets up jmp_buf, they need to be chained to jump back to
+    // correct location.
+    jmp_buf* prev_jmp_buf = prof_thread->getJmpCtx();
+    // Should be preserved across setjmp/longjmp
+    volatile int depth = 0;
+    int actual_max_depth = truncated ? max_depth + 1 : max_depth;
+
+    if (setjmp(crash_protection_ctx) != 0) {
+        // checkFault() does a longjmp from inside segvHandler, bypassing
+        // segvHandler's SignalHandlerScope destructor.  Compensate.
+        SIGNAL_HANDLER_UNWIND_AFTER_LONGJMP();
+        prof_thread->setJmpCtx(prev_jmp_buf);
+        if (depth < max_depth) {
+            fillFrame(frames[depth++], BCI_ERROR, "break_not_walkable");
+        }
+        return depth;
+    }
+
+    prof_thread->setJmpCtx(&crash_protection_ctx);
     VMThread* vm_thread = VMThread::current();
     if (vm_thread != NULL && !vm_thread->isThreadAccessible()) {
         Counters::increment(WALKVM_THREAD_INACCESSIBLE);
@@ -246,39 +277,16 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
     } else {
         Counters::increment(WALKVM_VMTHREAD_OK);
     }
-    void* saved_exception = vm_thread != NULL ? vm_thread->exception() : NULL;
 
-    // Should be preserved across setjmp/longjmp
-    volatile int depth = 0;
-    int actual_max_depth = truncated ? max_depth + 1 : max_depth;
     bool fp_chain_fallback = false;
     int fp_chain_depth = 0;
 
-    ProfiledThread* profiled_thread = ProfiledThread::currentSignalSafe();
-
     VMJavaFrameAnchor* anchor = NULL;
     if (vm_thread != NULL) {
         anchor = vm_thread->anchor();
         if (anchor == NULL) {
             Counters::increment(WALKVM_ANCHOR_NULL);
         }
-        vm_thread->exception() = &crash_protection_ctx;
-        if (profiled_thread != nullptr) {
-            profiled_thread->setCrashProtectionActive(true);
-        }
-        if (setjmp(crash_protection_ctx) != 0) {
-            // checkFault() does a longjmp from inside segvHandler, bypassing
-            // segvHandler's SignalHandlerScope destructor.  Compensate.
-            SIGNAL_HANDLER_UNWIND_AFTER_LONGJMP();
-            if (profiled_thread != nullptr) {
-                profiled_thread->setCrashProtectionActive(false);
-            }
-            vm_thread->exception() = saved_exception;
-            if (depth < max_depth) {
-                fillFrame(frames[depth++], BCI_ERROR, "break_not_walkable");
-            }
-            return depth;
-        }
     }
 
     const void* prev_native_pc = NULL;
@@ -616,7 +624,8 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
                 if (features.vtable_target && nm->isVTableStub() && depth == 0) {
                     uintptr_t receiver = frame.jarg0();
                     if (receiver != 0) {
-                        VMSymbol* symbol = VMKlass::fromOop(receiver)->name();
+                        VMKlass* klass = VMKlass::fromOop(receiver);
+                        VMSymbol* symbol = klass != nullptr ? klass->name() : nullptr;
                         // Store the raw VMSymbol* in the frame's method_id
                         // slot. BCI_VTABLE_RECEIVER (vmEntry.h) repurposes
                         // method_id for this pointer — same precedent as
@@ -922,12 +931,7 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
     }
 
     done:
-    if (profiled_thread != nullptr) {
-        profiled_thread->setCrashProtectionActive(false);
-    }
-    if (vm_thread != NULL) {
-        vm_thread->exception() = saved_exception;
-    }
+    prof_thread->setJmpCtx(prev_jmp_buf);
 
     // Drop unknown leaf frame - it provides no useful information and breaks
     // aggregation by lumping unrelated samples under a single "unknown" entry
@@ -953,33 +957,25 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
 }
 
 void HotspotSupport::checkFault(ProfiledThread* thrd) {
-    if (!JVMThread::isInitialized()) {
-        // JVM has not been loaded or has not been initialized yet
+    // Should not get to here (?)
+    if (thrd == nullptr) {
         return;
     }
 
-    VMThread* vm_thread = VMThread::current();
-    if (vm_thread == NULL || !vm_thread->isThreadAccessible()) {
+    if (!JVMThread::isInitialized()) {
+        // JVM has not been loaded or has not been initialized yet
         return;
     }
 
-    // Prefer the semantic crash protection flag (reliable regardless of stack frame sizes).
-    // Fall back to sameStack heuristic when ProfiledThread TLS is unavailable (e.g. during
-    // early init or in crash recovery tests). sameStack uses a fixed 8KB threshold which
-    // can fail with ASAN-inflated frames, but the crashProtectionActive path handles that.
-    bool protected_walk = (thrd != nullptr && thrd->isCrashProtectionActive())
-                       || sameStack(vm_thread->exception(), &vm_thread);
-    if (!protected_walk) {
+    // Check if longjmp is setup for this thread
+    if (!thrd->isProtected()) {
         return;
     }
 
-    if (thrd != nullptr) {
-        thrd->resetCrashHandler();
-    }
-    longjmp(*(jmp_buf*)vm_thread->exception(), 1);
+    thrd->resetCrashHandler();
+    longjmp(*thrd->getJmpCtx(), 1);
 }
 
-
 int HotspotSupport::getJavaTraceAsync(void *ucontext, ASGCT_CallFrame *frames,
                                 int max_depth, StackContext *java_ctx,
                                 bool *truncated) {
@@ -1189,7 +1185,6 @@ int HotspotSupport::getJavaTraceAsync(void *ucontext, ASGCT_CallFrame *frames,
   return trace.frames - frames + 1;
 }
 
-
 int HotspotSupport::walkJavaStack(StackWalkRequest& request) {
   CStack cstack = Profiler::instance()->cstackMode();
   StackWalkFeatures features = Profiler::instance()->stackWalkFeatures();

ddprof-lib/src/main/cpp/hotspot/hotspotSupport.h

@@ -9,7 +9,6 @@
 
 #include "hotspot/hotspotStackFrame.h"
 #include "hotspot/jitCodeCache.h"
-#include "profiler.h"
 #include "stackFrame.h"
 #include "stackWalker.h"
 

ddprof-lib/src/main/cpp/hotspot/vmStructs.cpp

@@ -15,6 +15,7 @@
 #include "jvmThread.h"
 #include "safeAccess.h"
 #include "spinLock.h"
+#include "threadLocalData.h"
 #include "threadState.h"
 
 CodeCache* VMStructs::_libjvm = nullptr;

ddprof-lib/src/main/cpp/hotspot/vmStructs.h

@@ -16,7 +16,6 @@
 #include "counters.h"
 #include "jvmThread.h"
 #include "safeAccess.h"
-#include "thread.h"
 #include "threadState.h"
 #include "vmEntry.h"
 
@@ -45,12 +44,8 @@ class VMNMethod;
 inline bool crashProtectionActive();
 
 template <typename T>
-inline T* cast_to(const void* ptr) {
-    assert(VM::isHotspot()); // This should only be used in HotSpot-specific code
-    assert(T::type_size() > 0); // Ensure type size has been initialized
-    assert(crashProtectionActive() || ptr == nullptr || SafeAccess::isReadableRange(ptr, T::type_size()));
-    return reinterpret_cast<T*>(const_cast<void*>(ptr));
-}
+inline T* cast_to(const void* ptr);
+
 
 template <typename T>
 T* cast_or_null(const void* ptr) {
@@ -677,7 +672,13 @@ DECLARE(VMKlass)
             if (_compact_object_headers) {
                 uintptr_t mark = *(uintptr_t*)oop;
                 if (mark & MONITOR_BIT) {
-                    mark = *(uintptr_t*)(mark ^ MONITOR_BIT);
+                    // TOCTOU: MonitorDeflationThread may free the ObjectMonitor between
+                    // reading the mark word and dereferencing the monitor pointer. Use
+                    // safeFetch64 so a concurrent deflation/free does not crash here.
+                    mark = (uintptr_t)SafeAccess::safeFetch64((int64_t*)(mark ^ MONITOR_BIT), 0);
+                    if (mark == 0) {
+                        return nullptr;
+                    }
                 }
                 narrow_klass = mark >> _markWord_klass_shift;
             } else {
@@ -843,17 +844,6 @@ DECLARE(VMThread)
         return *(void**) at(_thread_exception_offset);
     }
 
-    // Returns true if setjmp crash protection is currently active for this thread.
-    // Reads the exception field via direct pointer arithmetic, deliberately bypassing
-    // at() and its crashProtectionActive() assertion to avoid infinite recursion.
-    // Safe because 'this' is the current live thread (we are in its signal handler).
-    static bool isExceptionActive() {
-        if (_thread_exception_offset < 0) return false;
-        void* vt = JVMThread::current();
-        if (vt == nullptr) return false;
-        return *(const void* const*)((const char*)vt + _thread_exception_offset) != nullptr;
-    }
-
     NOADDRSANITIZE VMJavaFrameAnchor* anchor() {
         if (!isJavaThread(this)) return NULL;
         assert(_thread_anchor_offset >= 0);
@@ -1221,18 +1211,4 @@ class InterpreterFrame : VMStructs {
     }
 };
 
-// Defined here (after VMThread) so the VMThread::isExceptionActive() fallback
-// is accessible. The forward declaration at the top of this file allows cast_to()
-// to reference it before VMThread is declared.
-inline bool crashProtectionActive() {
-    ProfiledThread* pt = ProfiledThread::currentSignalSafe();
-    if (pt != nullptr && pt->isCrashProtectionActive()) return true;
-    // Fallback for threads without ProfiledThread TLS (e.g. JVM internal threads):
-    // if walkVM has set up setjmp protection via vm_thread->exception(), the assert
-    // is equally redundant — any bad read will be caught by the SIGSEGV handler.
-    // Uses VMThread::isExceptionActive() which reads the field directly without
-    // going through at() to avoid recursive assertion.
-    return JVMThread::key() != pthread_key_t(-1) && VMThread::isExceptionActive();
-}
-
 #endif // _HOTSPOT_VMSTRUCTS_H

ddprof-lib/src/main/cpp/hotspot/vmStructs.inline.h

@@ -7,8 +7,26 @@
 #ifndef _HOTSPOT_VMSTRUCTS_INLINE_H
 #define _HOTSPOT_VMSTRUCTS_INLINE_H
 
+#include "hotspot/hotspotSupport.h"
 #include "hotspot/vmStructs.h"
 #include "jvmThread.h"
+#include "threadLocalData.h"
+
+inline bool crashProtectionActive() {
+    ProfiledThread* pt = ProfiledThread::currentSignalSafe();
+    if (pt != nullptr) {
+        return pt->isProtected();
+    }
+    return false;
+}
+
+template <typename T>
+inline T* cast_to(const void* ptr) {
+    assert(VM::isHotspot()); // This should only be used in HotSpot-specific code
+    assert(T::type_size() > 0); // Ensure type size has been initialized
+    assert(crashProtectionActive() || ptr == nullptr || SafeAccess::isReadableRange(ptr, T::type_size()));
+    return reinterpret_cast<T*>(const_cast<void*>(ptr));
+}
 
 VMThread* VMThread::current() {
     assert(VM::isHotspot());

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

@@ -21,7 +21,7 @@
 #include "os.h"
 #include "profiler.h"
 #include "stackWalker.h"
-#include "thread.h"
+#include "threadLocalData.h"
 #include "threadState.inline.h"
 #include "guards.h"
 #include <sys/time.h>

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

@@ -30,7 +30,7 @@
 #include "os.h"
 #include "otel_process_ctx.h"
 #include "profiler.h"
-#include "thread.h"
+#include "threadLocalData.h"
 #include "tsc.h"
 #include "vmEntry.h"
 #include <errno.h>

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

@@ -9,7 +9,7 @@
 #include "frames.h"
 #include "os.h"
 #include "profiler.h"
-#include "thread.h"
+#include "threadLocalData.h"
 #include "vmEntry.h"
 
 #include "hotspot/hotspotSupport.h"
@@ -29,6 +29,11 @@ void JVMSupport::setLoadState(JMethodIDLoadStats state) {
     __atomic_store(&jmethodID_load_state, &state, __ATOMIC_RELEASE);
 }
 
+// If any of the two keys is invalid, profiler should not start
+bool JVMSupport::checkFatalError() {
+    return !JVMThread::hasValidKey() || !ProfiledThread::hasValidKey();     
+} 
+
 void JVMSupport::initExecution(Arguments& args, jvmtiEnv* jvmti, JNIEnv* jni) {
     JMethodIDLoadStats current_state = getLoadState();
     // Already setup by previous execution