Adopt ThreadLocal to store ProfiledThread

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

@@ -26,6 +26,7 @@
 #include "engine.h"
 #include "hotspot/vmStructs.inline.h"
 #include "incbin.h"
+#include "jvmSupport.h"
 #include "jvmThread.h"
 #include "os.h"
 #include "otel_process_ctx.h"
@@ -68,6 +69,11 @@ class JniString {
 
 extern "C" DLLEXPORT jboolean JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_init0(JNIEnv *env, jclass unused) {
+  Error error = Profiler::instance()->init();
+  if (error) {
+    return JNI_FALSE;
+  }
+
   // JavaVM* has already been stored when the native library was loaded so we can pass nullptr here
   return VM::initProfilerBridge(nullptr, true);
 }
@@ -129,6 +135,19 @@ Java_com_datadoghq_profiler_JavaProfiler_getSamples(JNIEnv *env,
   return (jlong)Profiler::instance()->total_samples();
 }
 
+// Init or get current profiled thread.
+// Calling thread's thread local may not be initialized due to race. 
+// Especially, during the early startup phase.
+// This problem can be eliminated with native agent.
+static ProfiledThread* initOrGetCurrentThread() {
+  ProfiledThread* current = ProfiledThread::current();
+  if (current == nullptr) {
+    current = ProfiledThread::initCurrentThreadSignalSafe();
+  }
+  return current;
+}
+
+
 // some duplication between add and remove, though we want to avoid having an extra branch in the hot path
 
 // JavaCritical is faster JNI, but more restrictive - parameters and return value have to be
@@ -137,7 +156,7 @@ Java_com_datadoghq_profiler_JavaProfiler_getSamples(JNIEnv *env,
 // still compatible in the event of signature changes in the future.
 extern "C" DLLEXPORT void JNICALL
 JavaCritical_com_datadoghq_profiler_JavaProfiler_filterThreadAdd0() {
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = initOrGetCurrentThread();
   assert(current != nullptr);
   int tid = current->tid();
   if (unlikely(tid < 0)) {
@@ -167,7 +186,7 @@ JavaCritical_com_datadoghq_profiler_JavaProfiler_filterThreadAdd0() {
 
 extern "C" DLLEXPORT void JNICALL
 JavaCritical_com_datadoghq_profiler_JavaProfiler_filterThreadRemove0() {
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = initOrGetCurrentThread();
   assert(current != nullptr);
   int tid = current->tid();
   if (unlikely(tid < 0)) {
@@ -319,7 +338,7 @@ Java_com_datadoghq_profiler_JavaProfiler_recordQueueEnd0(
 
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_parkEnter0(JNIEnv *env, jclass unused) {
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = initOrGetCurrentThread();
   if (current == nullptr) {
     return;
   }
@@ -337,10 +356,11 @@ Java_com_datadoghq_profiler_JavaProfiler_parkEnter0(JNIEnv *env, jclass unused)
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_parkExit0(
     JNIEnv *env, jclass unused, jlong blocker, jlong unblockingSpanId) {
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = initOrGetCurrentThread();
   if (current == nullptr) {
     return;
   }
+
   u64 park_block_token = 0;
   if (!current->parkExit(park_block_token) || park_block_token == 0) {
     return;
@@ -370,7 +390,7 @@ Java_com_datadoghq_profiler_JavaProfiler_blockEnter0(
   if (!decodeJavaBlockState(state, decoded)) {
     return 0;
   }
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = initOrGetCurrentThread();
   if (current == nullptr) {
     return 0;
   }
@@ -392,7 +412,7 @@ Java_com_datadoghq_profiler_JavaProfiler_blockExit0(
   if (block_token == 0) {
     return;
   }
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = initOrGetCurrentThread();
   if (current == nullptr) {
     return;
   }
@@ -685,7 +705,7 @@ Java_com_datadoghq_profiler_OTelContext_readProcessCtx0(JNIEnv *env, jclass unus
 
 extern "C" DLLEXPORT jobject JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_initializeContextTLS0(JNIEnv* env, jclass unused, jlongArray metadata) {
-  ProfiledThread* thrd = ProfiledThread::current();
+  ProfiledThread* thrd = initOrGetCurrentThread();
   assert(thrd != nullptr);
 
   if (!thrd->isContextInitialized()) {

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

@@ -32,12 +32,12 @@ bool JVMSupport::initialize() {
         return false;
     }
 
-    // Add ProfiledThread key checking here in next PR
-    return true;
+    // Check ProfiledThread key, it is critical for storing per-thread metadata
+    return ProfiledThread::isThreadKeyValid();
 }
 
 bool JVMSupport::isInitialized() {
-    return JVMThread::isInitialized();
+    return JVMThread::isInitialized() && ProfiledThread::isThreadKeyValid();
 }
 
 JVMSupport::JMethodIDLoadStats JVMSupport::getLoadState() {

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

@@ -1235,6 +1235,17 @@ Error Profiler::checkState() {
   return Error::OK;
 }
 
+Error Profiler::init() {
+  MutexLocker ml(_state_lock);
+  Error error = checkState();
+  if (error) {
+    return error;
+  }
+
+  ProfiledThread::initCurrentThread();
+  return Error::OK;
+}
+
 Error Profiler::start(Arguments &args, bool reset) {
   MutexLocker ml(_state_lock);
   Error error = checkState();
@@ -1356,7 +1367,7 @@ Error Profiler::start(Arguments &args, bool reset) {
   // Minor optim: Register the current thread (start thread won't be called)
   if (_thread_filter.enabled()) {
     _thread_filter.clearActive();
-    ProfiledThread *current = ProfiledThread::current();
+    ProfiledThread *current = ProfiledThread::initCurrentThread();
     assert(current != nullptr);
     int slot_id = current->filterSlotId();
     if (slot_id < 0) {

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

@@ -285,6 +285,7 @@ class alignas(alignof(SpinLock)) Profiler {
     return __atomic_load_n(&_epoch, __ATOMIC_RELAXED);
   }
 
+  Error init();
   Error run(Arguments &args);
   Error runInternal(Arguments &args, std::ostream &out);
   Error restart(Arguments &args);

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

@@ -11,56 +11,29 @@
 #include <cstring>
 #include <time.h>
 
-pthread_key_t ProfiledThread::_tls_key;
-bool ProfiledThread::_tls_key_initialized = false;
 
-void ProfiledThread::initTLSKey() {
-  static pthread_once_t tls_initialized = PTHREAD_ONCE_INIT;
-  pthread_once(&tls_initialized, doInitTLSKey);
-}
-
-void ProfiledThread::doInitTLSKey() {
-  pthread_key_create(&_tls_key, freeKey);
-  // Must be set AFTER pthread_key_create so signal handlers see a valid key.
-  // Store-release pairs with the acquire loads in currentSignalSafe() and release()
-  // to prevent hardware load-load reordering on weakly-ordered architectures (aarch64):
-  // a plain volatile write is not sufficient there.
-  __atomic_store_n(&_tls_key_initialized, true, __ATOMIC_RELEASE);
-}
+ThreadLocal<ProfiledThread*>  ProfiledThread::_current_thread;
 
-inline void ProfiledThread::freeKey(void *key) {
-  ProfiledThread *tls_ref = (ProfiledThread *)(key);
-  if (tls_ref != NULL) {
-    SignalBlocker blocker;
-    delete tls_ref;
+ProfiledThread* ProfiledThread::initCurrentThread() {
+  ProfiledThread* tls = current();
+  if (tls == nullptr) {
+    int tid = OS::threadId();
+    tls = ProfiledThread::forTid(tid);
+    _current_thread.set(tls);
   }
+  return tls;
 }
 
-void ProfiledThread::initCurrentThread() {
-  // JVMTI callback path - does NOT use buffer
-  // Allocate dedicated ProfiledThread for Java threads (not from buffer)
-  // This MUST happen here to prevent lazy allocation in signal handler
-  initTLSKey();
-
-  if (pthread_getspecific(_tls_key) != NULL) {
-    return; // Already initialized
-  }
-
-  int tid = OS::threadId();
-  ProfiledThread *tls = ProfiledThread::forTid(tid);
-  pthread_setspecific(_tls_key, (const void *)tls);
+ProfiledThread* ProfiledThread::initCurrentThreadSignalSafe() {
+  SignalBlocker blocker;
+  return initCurrentThread();
 }
 
 void ProfiledThread::release() {
-  if (!__atomic_load_n(&_tls_key_initialized, __ATOMIC_ACQUIRE)) {
-    return;
-  }
-  pthread_key_t key = _tls_key;
-  ProfiledThread *tls = (ProfiledThread *)pthread_getspecific(key);
-  if (tls != NULL) {
-    SignalBlocker blocker;
-    pthread_setspecific(key, NULL);
-    delete tls;
+  ProfiledThread* pt = _current_thread.get();
+  if (pt != nullptr) {
+    _current_thread.clear();
+    delete pt;
   }
 }
 
@@ -72,27 +45,6 @@ int ProfiledThread::currentTid() {
   return OS::threadId();
 }
 
-ProfiledThread *ProfiledThread::current() {
-  initTLSKey();
-
-  ProfiledThread *tls = (ProfiledThread *)pthread_getspecific(_tls_key);
-  if (tls == NULL) {
-    // Lazy allocation - safe since current() is never called from signal handlers
-    int tid = OS::threadId();
-    tls = ProfiledThread::forTid(tid);
-    pthread_setspecific(_tls_key, (const void *)tls);
-  }
-  return tls;
-}
-
-ProfiledThread *ProfiledThread::currentSignalSafe() {
-  // Signal-safe: never allocate, just return existing TLS or null.
-  // Use _tls_key_initialized instead of key != 0 because pthread_key_create
-  // can legitimately return key 0 (common on musl where keys start at 0).
-  return __atomic_load_n(&_tls_key_initialized, __ATOMIC_ACQUIRE) ? (ProfiledThread *)pthread_getspecific(_tls_key) : nullptr;
-}
-
-
 Context ProfiledThread::snapshotContext(size_t numAttrs) {
   Context ctx = {};
   u64 span_id = 0, root_span_id = 0;

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

@@ -9,6 +9,7 @@
 #include "context.h"
 #include "otel_context.h"
 #include "os.h"
+#include "threadLocal.h"
 #include "threadState.h"
 #include "unwindStats.h"
 #include <atomic>
@@ -53,12 +54,7 @@ class ProfiledThread : public ThreadLocalData {
   // Even with a 5-level cap we can still encounter highly recursive signal handlers.
   static constexpr u32 CRASH_HANDLER_NESTING_LIMIT = 5;
 private:
-  static pthread_key_t _tls_key;
-  static bool _tls_key_initialized;
-
-  static void initTLSKey();
-  static void doInitTLSKey();
-  static inline void freeKey(void *key);
+  static ThreadLocal<ProfiledThread*>  _current_thread;
 
   // longjmp buffer. Used by hotspot only at this moment.
   // Published in walkVM() and consumed in checkFault() from an asynchronous
@@ -108,32 +104,45 @@ class ProfiledThread : public ThreadLocalData {
   virtual ~ProfiledThread() { }
 public:
   static ProfiledThread *forTid(int tid) { return new ProfiledThread(tid); }
+  static bool isThreadKeyValid() {
+    return _current_thread.isKeyValid();
+  }
 
-  static void initCurrentThread();
-  static void release();
 #ifdef UNIT_TEST
   // Simulates the moment inside release() after pthread_setspecific(NULL) but
   // before delete — the race window the clearCurrentThreadTLS fix covers.
   // Returns the detached pointer so the caller can delete it after assertions.
   static ProfiledThread* clearCurrentThreadTLS() {
-    if (__atomic_load_n(&_tls_key_initialized, __ATOMIC_ACQUIRE)) {
-      ProfiledThread *pt = (ProfiledThread *)pthread_getspecific(_tls_key);
-      pthread_setspecific(_tls_key, nullptr);
-      return pt;
-    }
-    return nullptr;
+    assert(isThreadKeyValid() && "Should not reach here - profiling should have been disabled");
+    ProfiledThread* pt = _current_thread.get();
+    _current_thread.clear();
+    return pt;
   }
   // Deletes a ProfiledThread returned by clearCurrentThreadTLS().
   // Needed because the destructor is private.
   static void deleteForTest(ProfiledThread *pt) { delete pt; }
 #endif
+  // initCurrentThread() and release() are not async-signal-safe: 
+  // must call outside of a signal handler with signal blocked
+  static ProfiledThread* initCurrentThread();
+  static void release();
 
-  static ProfiledThread *current();
-  static ProfiledThread *currentSignalSafe(); // Signal-safe version that never allocates
-  static int currentTid();
+  // This version blocks signals, so that initialization cannot
+  // be interrupted by signals
+  static ProfiledThread* initCurrentThreadSignalSafe();
 
-  inline int tid() { return _tid; }
+  // This call is async-signal-safe
+  static inline ProfiledThread *current() {
+    assert(isThreadKeyValid() && "Should not reach here - profiling should have been disabled");
+    return _current_thread.get();
+  }
+
+  static inline ProfiledThread *currentSignalSafe() { // Signal-safe version that never allocates
+    return current();
+  }
 
+  static int currentTid();
+  inline int tid() { return _tid; }
   inline u64 noteCPUSample(u32 recording_epoch) {
     _recording_epoch = recording_epoch;
     return ++_cpu_epoch;

ddprof-lib/src/test/cpp/ddprof_ut.cpp

@@ -382,8 +382,7 @@ static DdprofGlobalSetup ddprof_global_setup;
 
         if (pid == 0) {
             // ---- child process (fork isolates TLS from other tests) ----
-            ProfiledThread::initCurrentThread();
-            ProfiledThread* pt = ProfiledThread::currentSignalSafe();
+            ProfiledThread* pt = ProfiledThread::initCurrentThread();
             if (pt == nullptr) _exit(2);
 
             // Baseline: entering critical section works.