Adopt ThreadLocal to store ProfiledThread

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

@@ -41,7 +41,7 @@ void ContextApi::initializeContextTLS(ProfiledThread* thrd) {
 }
 
 bool ContextApi::get(u64& span_id, u64& root_span_id) {
-    ProfiledThread* thrd = ProfiledThread::currentSignalSafe();
+    ProfiledThread* thrd = ProfiledThread::current();
     if (thrd == nullptr || !thrd->isContextInitialized()) {
         return false;
     }
@@ -59,7 +59,7 @@ bool ContextApi::get(u64& span_id, u64& root_span_id) {
 }
 
 Context ContextApi::snapshot() {
-    ProfiledThread* thrd = ProfiledThread::currentSignalSafe();
+    ProfiledThread* thrd = ProfiledThread::current();
     if (thrd == nullptr) {
         return {};
     }

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

@@ -225,7 +225,7 @@ void CTimerJvmti::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
     return;
   }
   int tid = 0;
-  ProfiledThread *current = ProfiledThread::currentSignalSafe();
+  ProfiledThread *current = ProfiledThread::current();
   assert(current == nullptr || !current->isDeepCrashHandler());
   if (current != nullptr && JVMThread::current() == nullptr
       && current->inInitWindow()) {
@@ -281,7 +281,7 @@ void CTimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
     return;
   }
   int tid = 0;
-  ProfiledThread *current = ProfiledThread::currentSignalSafe();
+  ProfiledThread *current = ProfiledThread::current();
   assert(current == nullptr || !current->isDeepCrashHandler());
   // Guard against the race window between Profiler::registerThread() and
   // thread_native_entry setting JVM TLS (PROF-13072): skip at most one signal

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

@@ -1771,7 +1771,7 @@ void Recording::writeCurrentContext(Buffer *buf) {
   buf->putVar64(rootSpanId);
 
   size_t numAttrs = Profiler::instance()->numContextAttributes();
-  ProfiledThread* thrd = hasContext ? ProfiledThread::currentSignalSafe() : nullptr;
+  ProfiledThread* thrd = hasContext ? ProfiledThread::current() : nullptr;
   for (size_t i = 0; i < numAttrs; i++) {
     buf->putVar32(thrd != nullptr ? thrd->getOtelTagEncoding(i) : 0);
   }

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

@@ -24,20 +24,20 @@
 // rejected because of the static TLS surplus on Graal).
 
 int getInSignalDepth() {
-    ProfiledThread *pt = ProfiledThread::currentSignalSafe();
+    ProfiledThread *pt = ProfiledThread::current();
     return pt != nullptr ? static_cast<int>(pt->signalDepth()) : 0;
 }
 
 bool isInTrackedSignalContext() {
-    ProfiledThread *pt = ProfiledThread::currentSignalSafe();
+    ProfiledThread *pt = ProfiledThread::current();
     // 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::currentSignalSafe();
+    ProfiledThread *pt = ProfiledThread::current();
     if (pt != nullptr) {
         pt->enterSignalScope();
     } else {
@@ -49,23 +49,23 @@ SignalHandlerScope::SignalHandlerScope() : _active(true) {
 
 SignalHandlerScope::~SignalHandlerScope() {
     if (!_active) return;
-    ProfiledThread *pt = ProfiledThread::currentSignalSafe();
+    ProfiledThread *pt = ProfiledThread::current();
     if (pt != nullptr) {
         pt->exitSignalScope();
     }
 }
 
 void SignalHandlerScope::release() {
     if (!_active) return;
-    ProfiledThread *pt = ProfiledThread::currentSignalSafe();
+    ProfiledThread *pt = ProfiledThread::current();
     if (pt != nullptr) {
         pt->exitSignalScope();
     }
     _active = false;
 }
 
 void signalHandlerUnwindAfterLongjmp() {
-    ProfiledThread *pt = ProfiledThread::currentSignalSafe();
+    ProfiledThread *pt = ProfiledThread::current();
     if (pt != nullptr) {
         pt->exitSignalScope();
     }
@@ -75,7 +75,7 @@ void signalHandlerUnwindAfterLongjmp() {
 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::currentSignalSafe();
+    _thread_ptr = ProfiledThread::current();
     if (_thread_ptr != nullptr) {
         // Primary path: Use ProfiledThread storage (fast and memory-efficient)
         _entered = _thread_ptr->tryEnterCriticalSection();

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

@@ -232,7 +232,7 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
     // VMStructs is only available for hotspot JVM 
     assert(VM::isHotspot());
 
-    ProfiledThread* prof_thread = ProfiledThread::currentSignalSafe();
+    ProfiledThread* prof_thread = ProfiledThread::current();
     if (prof_thread == nullptr) {
         Counters::increment(SAMPLES_DROPPED_THREAD_LOCAL);
         return 0;
@@ -1198,7 +1198,7 @@ int HotspotSupport::walkJavaStack(StackWalkRequest& request) {
     if (cstack >= CSTACK_VM) {
       java_frames = walkVM(ucontext, frames, max_depth, features, eventTypeFromBCI(request.event_type), lock_index, truncated);
     } else {
-        AsyncSampleMutex mutex(ProfiledThread::currentSignalSafe());
+        AsyncSampleMutex mutex(ProfiledThread::current());
         if (mutex.acquired()) {
             java_frames = getJavaTraceAsync(ucontext, frames, max_depth, java_ctx, truncated);
             if (java_frames > 0 && java_ctx->pc != NULL && VMStructs::hasMethodStructs()) {
@@ -1223,7 +1223,7 @@ int HotspotSupport::walkJavaStack(StackWalkRequest& request) {
         java_frames = walkVM(ucontext, frames, max_depth, features, eventTypeFromBCI(request.event_type), lock_index, truncated);
     } else {
         // Async events
-        AsyncSampleMutex mutex(ProfiledThread::currentSignalSafe());
+        AsyncSampleMutex mutex(ProfiledThread::current());
         if (mutex.acquired()) {
             java_frames = getJavaTraceAsync(ucontext, frames, max_depth, java_ctx, truncated);
             if (java_frames > 0 && java_ctx->pc != NULL && VMStructs::hasMethodStructs()) {

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

@@ -684,7 +684,7 @@ bool VMThread::isJavaThread(VMThread* vm_thread) {
 
     // JVMTI ThreadStart callback may have set the flag, which is reliable.
     // Or we may already compute and cache it, so use it instead.
-    ProfiledThread *prof_thread = ProfiledThread::currentSignalSafe();
+    ProfiledThread *prof_thread = ProfiledThread::current();
     if (prof_thread != nullptr) {
         ProfiledThread::ThreadType type = prof_thread->threadType();
         if (type != ProfiledThread::ThreadType::TYPE_UNKNOWN) {

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

@@ -13,7 +13,7 @@
 #include "threadLocalData.h"
 
 inline bool crashProtectionActive() {
-    ProfiledThread* pt = ProfiledThread::currentSignalSafe();
+    ProfiledThread* pt = ProfiledThread::current();
     if (pt != nullptr) {
         return pt->isProtected();
     }

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

@@ -49,7 +49,7 @@ void ITimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
     return;  // Another critical section is active, defer profiling
   }
   int tid = 0;
-  ProfiledThread *current = ProfiledThread::currentSignalSafe();
+  ProfiledThread *current = ProfiledThread::current();
   if (current != NULL) {
     current->noteCPUSample(Profiler::instance()->recordingEpoch());
     tid = current->tid();
@@ -116,7 +116,7 @@ void ITimerJvmti::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
     errno = saved_errno;
     return;
   }
-  ProfiledThread *current = ProfiledThread::currentSignalSafe();
+  ProfiledThread *current = ProfiledThread::current();
   if (current != nullptr && JVMThread::current() == nullptr
       && current->inInitWindow()) {
     current->tickInitWindow();

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,12 @@ class JniString {
 
 extern "C" DLLEXPORT jboolean JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_init0(JNIEnv *env, jclass unused) {
+  Error error = Profiler::instance()->init();
+  if (error) {
+    throwNew(env, "java/lang/IllegalStateException", error.message());
+    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 +136,22 @@ 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 call could be expensive, if TLS has not yet set.
+// Note: the racy can be avoided with native agent, remove this method
+//       once converted to native agent.
+// TODO: Priming is not fully restored at this moment, this method only
+//       handles a few special cases.
+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,8 +160,12 @@ 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();
-  assert(current != nullptr);
+  // Initialize thread TLS if it has not yet done
+  ProfiledThread *current = initOrGetCurrentThread();
+  if(current == nullptr) {
+    return;
+  }
+
   int tid = current->tid();
   if (unlikely(tid < 0)) {
     return;
@@ -167,8 +194,12 @@ JavaCritical_com_datadoghq_profiler_JavaProfiler_filterThreadAdd0() {
 
 extern "C" DLLEXPORT void JNICALL
 JavaCritical_com_datadoghq_profiler_JavaProfiler_filterThreadRemove0() {
-  ProfiledThread *current = ProfiledThread::current();
-  assert(current != nullptr);
+  // Initialize thread TLS if it has not yet done
+  ProfiledThread *current = initOrGetCurrentThread();
+  if(current == nullptr) {
+    return;
+  }
+
   int tid = current->tid();
   if (unlikely(tid < 0)) {
     return;
@@ -204,6 +235,10 @@ Java_com_datadoghq_profiler_JavaProfiler_recordTrace0(
     JNIEnv *env, jclass unused, jlong rootSpanId, jstring endpoint,
     jstring operation, jint sizeLimit) {
   JniString endpoint_str(env, endpoint);
+
+  // Initialize thread TLS if it has not yet done
+  initOrGetCurrentThread();
+
   u32 endpointLabel = Profiler::instance()->stringLabelMap()->bounded_lookup(
       endpoint_str.c_str(), endpoint_str.length(), sizeLimit);
   // StringDictionary reserves 0 as "no entry"; valid IDs start at 1.
@@ -261,6 +296,9 @@ Java_com_datadoghq_profiler_JavaProfiler_describeDebugCounters0(
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_recordSettingEvent0(
     JNIEnv *env, jclass unused, jstring name, jstring value, jstring unit) {
+  // Initialize thread TLS if it has not yet done
+  initOrGetCurrentThread();
+
   int tid = ProfiledThread::currentTid();
   if (tid < 0) {
     return;
@@ -285,6 +323,10 @@ extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_recordQueueEnd0(
     JNIEnv *env, jclass unused, jlong startTime, jlong endTime, jstring task,
     jstring scheduler, jthread origin, jstring queueType, jint queueLength) {
+
+  // Initialize thread TLS if it has not yet done
+  initOrGetCurrentThread();
+
   int tid = ProfiledThread::currentTid();
   if (tid < 0) {
     return;
@@ -341,6 +383,7 @@ Java_com_datadoghq_profiler_JavaProfiler_parkExit0(
   if (current == nullptr) {
     return;
   }
+
   u64 park_block_token = 0;
   if (!current->parkExit(park_block_token) || park_block_token == 0) {
     return;
@@ -685,7 +728,8 @@ 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();
+  // Initialize thread TLS if it has not yet done
+  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() {
@@ -95,7 +95,7 @@ int JVMSupport::asyncGetCallTrace(ASGCT_CallFrame *frames, int max_depth, void*
         return 0;
     }
 
-    AsyncSampleMutex mutex(ProfiledThread::currentSignalSafe());
+    AsyncSampleMutex mutex(ProfiledThread::current());
     if (!mutex.acquired()) {
         return 0;
     }

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

@@ -64,13 +64,14 @@ class RoutineInfo {
 // Unregister the current thread from the profiler and release its TLS under a
 // single SignalBlocker to close the race window between unregisterThread()
 // returning and release() acquiring its internal guard (PROF-14603).  Without
-// this, a SIGVTALRM delivered in that window could call currentSignalSafe()
+// this, a SIGVTALRM delivered in that window could call current()
 // and dereference a now-freed ProfiledThread.  Kept noinline so the
 // SignalBlocker's sigset_t does not appear in the caller's stack frame on
 // musl/aarch64 where the deopt blob may corrupt the wrapper's stack guard.
 __attribute__((noinline))
-static void unregister_and_release(int tid) {
+static void unregister_and_release() {
     SignalBlocker blocker;
+    int tid = ProfiledThread::currentTid();
     Profiler::unregisterThread(tid);
     ProfiledThread::release();
 }
@@ -83,7 +84,7 @@ static void unregister_and_release(int tid) {
 // appear in the caller's frame on platforms with stack-protector canaries.
 __attribute__((noinline))
 static void cleanup_unregister(void*) {
-    unregister_and_release(ProfiledThread::currentTid());
+    unregister_and_release();
 }
 
 // Thread-cleanup wrapper that avoids the static-libgcc / forced-unwind crash.
@@ -254,8 +255,8 @@ static void delete_routine_info(RoutineInfo* thr) {
 __attribute__((noinline))
 static void init_tls_and_register() {
     SignalBlocker blocker;
-    ProfiledThread::initCurrentThread();
-    if (ProfiledThread *pt = ProfiledThread::currentSignalSafe()) {
+    ProfiledThread* pt = ProfiledThread::initCurrentThread();
+    if (pt != nullptr) {
         pt->startInitWindow();
     }
     Profiler::registerThread(ProfiledThread::currentTid());
@@ -371,8 +372,10 @@ static void* start_routine_wrapper(void* args) {
         routine = thr->routine();
         params = thr->args();
         delete thr;
-        ProfiledThread::initCurrentThread();
-        ProfiledThread::currentSignalSafe()->startInitWindow();
+        ProfiledThread* pt = ProfiledThread::initCurrentThread();
+        if (pt != nullptr) {
+          pt->startInitWindow();
+        }
         Profiler::registerThread(ProfiledThread::currentTid());
     }
     // Use POSIX cleanup instead of C++ RAII to handle pthread_exit(): see run_with_cleanup.

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

@@ -741,7 +741,7 @@ void PerfEvents::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   if (!cs.entered()) {
     return;  // Another critical section is active, defer profiling
   }
-  ProfiledThread *current = ProfiledThread::currentSignalSafe();
+  ProfiledThread *current = ProfiledThread::current();
   if (current != NULL) {
     current->noteCPUSample(Profiler::instance()->recordingEpoch());
   }