diff --git a/ddprof-lib/src/main/cpp/javaApi.cpp b/ddprof-lib/src/main/cpp/javaApi.cpp index c83a7bfb6..7cc91f75c 100644 --- a/ddprof-lib/src/main/cpp/javaApi.cpp +++ b/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()) { diff --git a/ddprof-lib/src/main/cpp/jvmSupport.cpp b/ddprof-lib/src/main/cpp/jvmSupport.cpp index cff98c33b..bf463e8b3 100644 --- a/ddprof-lib/src/main/cpp/jvmSupport.cpp +++ b/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() { diff --git a/ddprof-lib/src/main/cpp/profiler.cpp b/ddprof-lib/src/main/cpp/profiler.cpp index dde305272..5c9da3776 100644 --- a/ddprof-lib/src/main/cpp/profiler.cpp +++ b/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) { diff --git a/ddprof-lib/src/main/cpp/profiler.h b/ddprof-lib/src/main/cpp/profiler.h index 1e1b2825e..7dbaa763f 100644 --- a/ddprof-lib/src/main/cpp/profiler.h +++ b/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); diff --git a/ddprof-lib/src/main/cpp/threadLocalData.cpp b/ddprof-lib/src/main/cpp/threadLocalData.cpp index 1843514bd..b6288148c 100644 --- a/ddprof-lib/src/main/cpp/threadLocalData.cpp +++ b/ddprof-lib/src/main/cpp/threadLocalData.cpp @@ -11,56 +11,29 @@ #include #include -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::_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; diff --git a/ddprof-lib/src/main/cpp/threadLocalData.h b/ddprof-lib/src/main/cpp/threadLocalData.h index e8b9a30c4..513a82882 100644 --- a/ddprof-lib/src/main/cpp/threadLocalData.h +++ b/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 @@ -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 _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; diff --git a/ddprof-lib/src/test/cpp/ddprof_ut.cpp b/ddprof-lib/src/test/cpp/ddprof_ut.cpp index 733bbb443..c3926d4a4 100644 --- a/ddprof-lib/src/test/cpp/ddprof_ut.cpp +++ b/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.