v3

Author: zhengyu123

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

@@ -40,6 +40,13 @@ static jmethodID getMethodId(VMMethod* method) {
     return NULL;
 }
 
+static ThreadLocal<jmp_buf*> jmp_ctx;
+
+void HotspotSupport::initThread() {
+    // Ensure the slot is allocated
+    jmp_ctx.set(nullptr);
+}
+
 /**
  * Converts a BCI_* frame type value to the corresponding EventType enum value.
  *
@@ -146,8 +153,6 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
     }
 }
 
-static ThreadLocal<jmp_buf*> jmp_ctx;
-
 __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) {

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

@@ -34,6 +34,10 @@ class HotspotSupport {
         return HotspotStackFrame::unwindAtomicStub(frame, pc);
     }
 
+    // Per-thread initialization.
+    // *Must* be called before signal is enabled for the thread
+    static void initThread();
+
     static inline bool isJitCode(const void* p) {
         return JitCodeCache::isJitCode(p);
     }

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

@@ -1212,7 +1212,7 @@ inline bool crashProtectionActive() {
     // 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();
+    return JVMThread::isInitialized() && VMThread::isExceptionActive();
 }
 
 #endif // _HOTSPOT_VMSTRUCTS_H

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

@@ -4,29 +4,33 @@
  */
 
 #include "jvmThread.h"
+#include "hotspot/hotspotSupport.h"
 #include "hotspot/vmStructs.inline.h"
 #include "j9/j9Support.h"
 #include "zing/zingSupport.h"
 #include "vmEntry.h"
 
-pthread_key_t JVMThread::_thread_key = pthread_key_t(-1);
 jfieldID JVMThread::_tid = nullptr;
+ThreadLocal<JVMThread*> JVMThread::_jvm_thread;
 
 bool JVMThread::initialize() {
   void* current_thread = currentThreadSlow();
-  if (current_thread == nullptr) {
-    return false;
-  }
-
-  for (int i = 0; i < 1024; i++) {
-    if (pthread_getspecific((pthread_key_t)i) == current_thread) {
-        _thread_key = pthread_key_t(i);
-        break;
-    }
-  }
-  // _tid is initialized in currentThreadSlow()
+  _jvm_thread.initialize(current_thread);
+  // _tid is side-effect of currentThreadSlow()
   assert(_tid != nullptr);
-  return _thread_key != pthread_key_t(-1);
+  return _jvm_thread.isKeyValid();
+}
+
+/**
+ * Per-thread initialization
+ * 
+ * This method *must* be called before signal is enabled
+ * for the thread.
+ */
+void JVMThread::initThread() {
+    if (VM::isHotspot()) {
+      HotspotSupport::initThread();
+    }
 }
 
 int JVMThread::nativeThreadId(JNIEnv* jni, jthread thread) {

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

@@ -6,48 +6,48 @@
 #ifndef _JVMTHREAD_H
 #define _JVMTHREAD_H
 
-#include <cassert>
 #include <jni.h>
 #include <jvmti.h>
-#include <pthread.h>
+
+#include "threadLocal.h"
 
 /**
  * JVMThread represents a native JVM thread that is JVM implementation agnostic
  */
 class JVMThread {
 private:
-    static pthread_key_t _thread_key;
     static jfieldID _tid;
+    static ThreadLocal<JVMThread*> _jvm_thread;
 
 public:
-    static bool isInitialized() {
-        return _thread_key != pthread_key_t(-1);
-    }
-
     /*
      * The initialization happens in early startup, in single-threaded mode,
      * no synchronization is needed
      */
     static bool initialize();
+    static bool isInitialized() {
+        return _tid != nullptr && _jvm_thread.isKeyValid();
+    }
+
+    // Per-thread initialization
+    static void initThread();
+
     static inline void* current() {
-        assert(isInitialized());
-        return pthread_getspecific(_thread_key);
+        return _jvm_thread.get();
     }
 
     static inline pthread_key_t key() {
-        return _thread_key;
+        return _jvm_thread.key();
     }
 
     static int nativeThreadId(JNIEnv* jni, jthread thread);
-
     static inline jlong javaThreadId(JNIEnv* env, jthread thread) {
        return env->GetLongField(thread, _tid);
     }
 
     static inline bool hasJavaThreadId() {
         return _tid != nullptr;
     }
-
 private:
     static void* currentThreadSlow();
 };

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

@@ -8,6 +8,7 @@
 #ifdef __linux__
 #include "counters.h"
 #include "guards.h"
+#include "jvmThread.h"
 #include "nativeSocketSampler.h"
 #include "profiler.h"
 
@@ -258,6 +259,7 @@ static void init_tls_and_register() {
         pt->startInitWindow();
     }
     Profiler::registerThread(ProfiledThread::currentTid());
+    JVMThread::initThread();
 }
 
 // Wrapper around the real start routine.
@@ -373,6 +375,7 @@ static void* start_routine_wrapper(void* args) {
         ProfiledThread::initCurrentThread();
         ProfiledThread::currentSignalSafe()->startInitWindow();
         Profiler::registerThread(ProfiledThread::currentTid());
+        JVMThread::initThread();
     }
     // Use POSIX cleanup instead of C++ RAII to handle pthread_exit(): see run_with_cleanup.
     // cleanup_unregister has already run on run_with_cleanup's normal return path.

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

@@ -16,6 +16,15 @@
  * Due to some restrictions of the language implementations, especially, on musl/aarch64,
  * they cannot be safely used in profiler.
  * 
+ * pthread_(get/set)specific() are not async-signal-safe, according to 
+ * https://man7.org/linux/man-pages/man7/signal-safety.7.html
+ * 
+ * In POSIX implementation, pthread_setspecific() call can trigger memory allocation
+ * if the slot is not available.
+ * Because we depend on the APIs to maintain per-thread data, we need to workaround the
+ * problem - call pthread_setspecific() at least once before signal is enabled for the
+ * thread (ideally, the value is set before signal is enabled).
+ * 
  * How to use?
  * A ThreadLocal should be declared as a static variable, e.g.
  * 
@@ -40,52 +49,60 @@
  * 
  */
 
+#include <unistd.h>
+
 // The function to create value if it does not exist
 typedef void* (*CREATE_FUNC)(void);
 // Cleanup the value when deleting the key
 typedef void (*CLEAN_FUNC)(void*);
+
+static constexpr pthread_key_t INVLID_KEY = pthread_key_t(-1);
+
 template <typename T, CREATE_FUNC C = nullptr, CLEAN_FUNC F = nullptr>
 class ThreadLocal {
 protected:
     pthread_key_t _key;
-    bool _key_valid;
 
 public:
     ThreadLocal(const ThreadLocal&) = delete;
     ThreadLocal& operator=(const ThreadLocal&) = delete;
 
-    ThreadLocal() {
+    ThreadLocal() : _key(INVLID_KEY) {
         static_assert(sizeof(T) == sizeof(void*),
                       "ThreadLocal<T> requires sizeof(T)==sizeof(void*); use a pointer type or add a specialization");
-        _key_valid = pthread_key_create(&_key, F) == 0;
+        pthread_key_create(&_key, F);
         // What to do if we can not create a key?
         // We probably want to shutdown profiler gracefully, instead of
         // aborting user application - We will need this mechanism globally,
         // defer to a separate task.
-        assert(_key_valid);
+        assert(isKeyValid());
     }
 
     ~ThreadLocal() {
-        if(_key_valid) {
+        if(isKeyValid()) {
             pthread_key_delete(_key);
         } else {
             assert(false && "Invalid pthread key");
         }
     }
 
+    bool isKeyValid() const {
+        return _key != INVLID_KEY;
+    }
+
     /**
      * set(nullptr) will result in the value being recreated when get() is called
      * when CREATE_FUNC is not nullptr.
      * Note: caller is responsible to free old value, which mirrors thread_local
      */
     void set(T value) {
-        assert(_key_valid && "Invalid pthread key");
+        assert(isKeyValid() && "Invalid pthread key");
         int err = pthread_setspecific(_key, reinterpret_cast<const void*>(value));
         assert(err == 0);
     }
 
     T get() {
-        assert(_key_valid && "Invalid pthread key");
+        assert(isKeyValid() && "Invalid pthread key");
         void* p = pthread_getspecific(_key);
         if (p == nullptr && C != nullptr) {
             p = C();
@@ -96,7 +113,7 @@ class ThreadLocal {
 
     // Clear the value
     void clear() {
-        assert(_key_valid && "Invalid pthread key");
+        assert(isKeyValid() && "Invalid pthread key");
         void* p = pthread_getspecific(_key);
         if (p == nullptr) return;
         int err = pthread_setspecific(_key, nullptr);
@@ -112,41 +129,44 @@ template <>
 class ThreadLocal<double> {
 protected:
     pthread_key_t _key;
-    bool _key_valid;
 
 public:
     ThreadLocal(const ThreadLocal&) = delete;
     ThreadLocal& operator=(const ThreadLocal&) = delete;
 
-    ThreadLocal() {
+    ThreadLocal() : _key(INVLID_KEY) {
         // Only support 64-bit platforms, double and void* are the same size
         static_assert(sizeof(void*) == 8);
         static_assert(sizeof(double) == 8);
-        _key_valid = pthread_key_create(&_key, nullptr) == 0;
+        pthread_key_create(&_key, nullptr);
         // What to do if we can not create a key?
-        assert(_key_valid && "Invalid pthread key");
+        assert(isKeyValid() && "Invalid pthread key");
     }
 
     ~ThreadLocal() {
-        if(_key_valid) {
+        if(isKeyValid()) {
             pthread_key_delete(_key);
         } else {
-            assert(_key_valid && "Invalid pthread key");
+            assert(isKeyValid() && "Invalid pthread key");
         }
     }
 
+    bool isKeyValid() const {
+        return _key != INVLID_KEY;
+    }
+
     // double <--> u64 cast, preserve bit format
     // Can use std::bit_cast after upgrade C++ version to 20
     void set(double value) {
-        assert(_key_valid && "Invalid pthread key");
+        assert(isKeyValid() && "Invalid pthread key");
         u64 val;
         memcpy(&val, &value, sizeof(value));
         int err = pthread_setspecific(_key, reinterpret_cast<const void*>(val));
         assert(err == 0);
     }
 
-    double get() {
-        assert(_key_valid && "Invalid pthread key");
+    double get() const {
+        assert(isKeyValid() && "Invalid pthread key");
         void* p = pthread_getspecific(_key);
         if (p == nullptr) {
             return 0.0;
@@ -159,10 +179,69 @@ class ThreadLocal<double> {
     }
 
     void clear() {
-        assert(_key_valid && "Invalid pthread key");
+        assert(isKeyValid() && "Invalid pthread key");
         int err = pthread_setspecific(_key, nullptr);
         assert(err == 0);
     }
 };
 
+class JVMThread;
+
+/**
+ * This thread local mirrors JVM's Thread::current(). The value is set by JVM
+ * and it is read-only variable.
+ */
+template <>
+class ThreadLocal<JVMThread*> {
+protected:
+    pthread_key_t _key;
+
+public:
+    ThreadLocal(const ThreadLocal&) = delete;
+    ThreadLocal& operator=(const ThreadLocal&) = delete;
+
+    ThreadLocal() : _key(INVLID_KEY) {
+    }
+
+    void initialize(void* current_thread) {
+        // Called from known JavaThread, it should never be nullptr.
+        if (current_thread == nullptr) {
+            assert(false && "Should not reach here");
+        }
+
+        long max_keys = sysconf(_SC_THREAD_KEYS_MAX);
+
+        for (long i = 0; i < max_keys; i++) {
+            if (pthread_getspecific((pthread_key_t)i) == current_thread) {
+                _key = pthread_key_t(i);
+                break;
+            }
+        }
+
+        assert(isKeyValid() && "Invalid thread key");
+    }
+
+    bool isKeyValid() const {
+        return _key != INVLID_KEY;
+    }
+
+    pthread_key_t key() const {
+        return _key;
+    }
+
+    void set(JVMThread* value) {
+        assert(false && "Should not reach here, value is set by JVM");
+    }
+
+    void* get() const {
+        assert(isKeyValid() && "Invalid pthread key");
+        return pthread_getspecific(_key);
+    }
+
+    void clear() {
+        assert(false && "Should not reach here");
+    }
+};
+
+
 #endif // _THREADLOCAL_H