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

@@ -1785,7 +1785,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/jitCodeCache.cpp

@@ -6,6 +6,7 @@
 #include "jitCodeCache.h"
 
 #include "hotspot/vmStructs.h"
+#include "threadLocalData.h"
 
 SpinLock JitCodeCache::_stubs_lock;
 CodeCache JitCodeCache::_runtime_stubs("[stubs]");
@@ -19,11 +20,13 @@ void JNICALL JitCodeCache::CompiledMethodLoad(jvmtiEnv *jvmti, jmethodID method,
                                 jint map_length,
                                 const jvmtiAddrLocationMap *map,
                                 const void *compile_info) {
+  ProfiledThread::initCurrentThreadSignalSafe();
   CodeHeap::updateBounds(code_addr, (const char *)code_addr + code_size);
 }
 
 void JNICALL JitCodeCache::DynamicCodeGenerated(jvmtiEnv *jvmti, const char *name,
                                   const void *address, jint length) {
+  ProfiledThread::initCurrentThreadSignalSafe();
   _stubs_lock.lock();
   _runtime_stubs.add(address, length, name, true);
   _stubs_lock.unlock();

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

@@ -597,6 +597,8 @@ void JNICALL VMStructs::NativeMethodBind(jvmtiEnv *jvmti, JNIEnv *jni, jthread t
     static int delayedCounter = 0;
     static void **delayed = (void **)malloc(512 * sizeof(void *) * 2);
 
+    ProfiledThread::initCurrentThreadSignalSafe();
+
     if (_memory_usage_func == NULL) {
         if (jvmti != NULL && jni != NULL) {
             checkNativeBinding(jvmti, jni, method, address);
@@ -684,7 +686,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);
 }
@@ -137,8 +144,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 = ProfiledThread::initCurrentThreadSignalSafe();
+  if(current == nullptr) {
+    return;
+  }
+
   int tid = current->tid();
   if (unlikely(tid < 0)) {
     return;
@@ -167,8 +178,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 = ProfiledThread::initCurrentThreadSignalSafe();
+  if(current == nullptr) {
+    return;
+  }
+
   int tid = current->tid();
   if (unlikely(tid < 0)) {
     return;
@@ -204,6 +219,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
+  ProfiledThread::initCurrentThreadSignalSafe();
+
   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 +280,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
+  ProfiledThread::initCurrentThreadSignalSafe();
+
   int tid = ProfiledThread::currentTid();
   if (tid < 0) {
     return;
@@ -285,6 +307,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
+  ProfiledThread::initCurrentThreadSignalSafe();
+
   int tid = ProfiledThread::currentTid();
   if (tid < 0) {
     return;
@@ -319,7 +345,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 = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
     return;
   }
@@ -337,10 +363,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 = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
     return;
   }
+
   u64 park_block_token = 0;
   if (!current->parkExit(park_block_token) || park_block_token == 0) {
     return;
@@ -370,7 +397,7 @@ Java_com_datadoghq_profiler_JavaProfiler_blockEnter0(
   if (!decodeJavaBlockState(state, decoded)) {
     return 0;
   }
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
     return 0;
   }
@@ -392,7 +419,7 @@ Java_com_datadoghq_profiler_JavaProfiler_blockExit0(
   if (block_token == 0) {
     return;
   }
-  ProfiledThread *current = ProfiledThread::current();
+  ProfiledThread *current = ProfiledThread::initCurrentThreadSignalSafe();
   if (current == nullptr) {
     return;
   }
@@ -685,7 +712,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 = ProfiledThread::initCurrentThreadSignalSafe();
   assert(thrd != nullptr);
 
   if (!thrd->isContextInitialized()) {
@@ -850,7 +878,7 @@ extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_setTraceContext0(JNIEnv* env, jclass unused,
     jlong localRootSpanId, jlong spanId, jlong traceIdHigh, jlong traceIdLow,
     jint slot0, jint enc0, jbyteArray utf0, jint slot1, jint enc1, jbyteArray utf1) {
-  ProfiledThread* thrd = ProfiledThread::current();
+  ProfiledThread* thrd = ProfiledThread::initCurrentThreadSignalSafe();
   if (thrd == nullptr) {
     return;
   }
@@ -906,7 +934,7 @@ Java_com_datadoghq_profiler_JavaProfiler_setTraceContext0(JNIEnv* env, jclass un
 // detached (valid=0), mirroring the DBB clear path (setContext(0,0,0,0) + clearContextValue*).
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_clearTraceContext0(JNIEnv* env, jclass unused) {
-  ProfiledThread* thrd = ProfiledThread::current();
+  ProfiledThread* thrd = ProfiledThread::initCurrentThreadSignalSafe();
   if (thrd == nullptr) {
     return;
   }
@@ -931,7 +959,7 @@ Java_com_datadoghq_profiler_JavaProfiler_clearTraceContext0(JNIEnv* env, jclass
 extern "C" DLLEXPORT jboolean JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_setContextValue0(JNIEnv* env, jclass unused,
     jint slot, jint encoding, jbyteArray utf8) {
-  ProfiledThread* thrd = ProfiledThread::current();
+  ProfiledThread* thrd = ProfiledThread::initCurrentThreadSignalSafe();
   if (thrd == nullptr || slot < 0 || slot >= (jint)DD_TAGS_CAPACITY) {
     return JNI_FALSE;
   }

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

@@ -20,6 +20,12 @@
 volatile JVMSupport::JMethodIDLoadStats JVMSupport::jmethodID_load_state = JVMSupport::No_loaded;
 Mutex JVMSupport::_initialization_lock;
 
+
+// This method must be called after JVM has been properly initialized, e.g. after JVMTI::VMinit()
+// callback.
+// Currently, there are two paths lead to this call
+// - JVMTI::VMInit() callback (vmEntry.cpp)
+// - JavaProfiler.getInstance() via JNI down call - JVM must have been initialized
 bool JVMSupport::initialize() {
     MutexLocker locker(_initialization_lock);
 
@@ -32,12 +38,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 +101,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;
     }