Centralize Java thread detection and type-safe thread state handling

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

@@ -188,7 +188,7 @@ __attribute__((no_sanitize("address"))) int HotspotSupport::walkVM(void* ucontex
     // Show extended frame types and stub frames for execution-type events
     bool details = event_type <= MALLOC_SAMPLE || features.mixed;
 
-    if (details && vm_thread != NULL && vm_thread->cachedIsJavaThread()) {
+    if (details && vm_thread != NULL && VMThread::isJavaThread(vm_thread)) {
         anchor = vm_thread->anchor();
     }
 
@@ -777,30 +777,12 @@ int HotspotSupport::getJavaTraceAsync(void *ucontext, ASGCT_CallFrame *frames,
     return 0;
   }
 
-  int state = vm_thread->state();
-  /**  from OpenJDK
-   https://github.com/openjdk/jdk/blob/7455bb23c1d18224e48e91aae4f11fe114d04fab/src/hotspot/share/utilities/globalDefinitions.hpp#L1030
-  enum JavaThreadState {
-      _thread_uninitialized     =  0, // should never happen (missing initialization) 
-      _thread_new               =  2, // just starting up, i.e., in process of being initialized 
-      _thread_new_trans         =  3, // corresponding transition state (not used, included for completeness) 
-      _thread_in_native         =  4, // running in native code 
-      _thread_in_native_trans   =  5, // corresponding transition state 
-      _thread_in_vm             =  6, // running in VM
-      _thread_in_vm_trans       =  7, // corresponding transition state
-      _thread_in_Java           =  8, // running in Java or in stub code
-      _thread_in_Java_trans     =  9, // corresponding transition state (not used, included for completeness)
-      _thread_blocked           = 10, // blocked in vm 
-      _thread_blocked_trans     = 11, // corresponding transition state
-      _thread_max_state         = 12  // maximum thread state+1 - used forstatistics allocation
-  };
-  **/
-
-  bool in_java = (state == 8 || state == 9);
+  JVMJavaThreadState state = vm_thread->state();
+  bool in_java = (state == _thread_in_Java || state == _thread_in_Java_trans);
   if (in_java && java_ctx->sp != 0) {
     // skip ahead to the Java frames before calling AGCT
     frame.restore((uintptr_t)java_ctx->pc, java_ctx->sp, java_ctx->fp);
-  } else if (state != 0) {
+  } else if (state != _thread_uninitialized) {
     VMJavaFrameAnchor* a = vm_thread->anchor();
     if (a == nullptr || a->lastJavaSP() == 0) {
       // we haven't found the top Java frame ourselves, and the lastJavaSP wasn't
@@ -810,7 +792,7 @@ int HotspotSupport::getJavaTraceAsync(void *ucontext, ASGCT_CallFrame *frames,
       return 0;
     }
   }
-  bool blocked_in_vm = (state == 10 || state == 11);
+  bool blocked_in_vm = (state == _thread_blocked || state == _thread_blocked_trans);
   // avoid unwinding during deoptimization
   if (blocked_in_vm && vm_thread->osThreadState() == OSThreadState::RUNNABLE) {
     Counters::increment(AGCT_BLOCKED_IN_VM);

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

@@ -629,19 +629,55 @@ void* VMThread::initialize(jthread thread) {
     return vm_thread;
 }
 
-static ExecutionMode convertJvmExecutionState(int state) {
+bool VMThread::isJavaThread(VMThread* vm_thread) {
+    // Not a JVM thread - native thread, e.g. thread launched by JNI code
+    if (vm_thread == nullptr) {
+        return false;
+    }
+
+    // Must be called from current thread
+    assert(vm_thread == VMThread::current());
+
+    // 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();
+    if (prof_thread != nullptr) {
+        ProfiledThread::ThreadType type = prof_thread->threadType();
+        if (type != ProfiledThread::ThreadType::TYPE_UNKNOWN) {
+            return type == ProfiledThread::ThreadType::TYPE_JAVA_THREAD;
+        }
+    }
+
+    // jvmti ThreadStart does not callback to JVM internal threads, e.g. Compiler threads, which are also JavaThreads,
+    // let's check the vtable pointer to make sure it is a Java thread. 
+    // A Java thread should have the same vtable as the one we got from a known Java thread during initialization
+    bool is_java_thread = vm_thread->hasJavaThreadVtable();
+    // Cache the thread type for future quick check
+    if (prof_thread != nullptr) {
+        prof_thread->setJavaThread(is_java_thread);
+    }
+    if (!is_java_thread) {
+        Counters::increment(WALKVM_CACHED_NOT_JAVA);
+    }
+    return is_java_thread;
+}
+
+static ExecutionMode convertJvmExecutionState(JVMJavaThreadState state) {
   switch (state) {
-  case 4:
-  case 5:
+  case _thread_in_native:
+  case _thread_in_native_trans:
     return ExecutionMode::NATIVE;
-  case 6:
-  case 7:
+  case _thread_uninitialized:
+  case _thread_new:
+  case _thread_new_trans:
+  case _thread_in_vm:
+  case _thread_in_vm_trans:
     return ExecutionMode::JVM;
-  case 8:
-  case 9:
+  case _thread_in_Java:
+  case _thread_in_Java_trans:
     return ExecutionMode::JAVA;
-  case 10:
-  case 11:
+  case _thread_blocked:
+  case _thread_blocked_trans:
     return ExecutionMode::SAFEPOINT;
   default:
     return ExecutionMode::UNKNOWN;
@@ -650,29 +686,15 @@ static ExecutionMode convertJvmExecutionState(int state) {
 
 ExecutionMode VMThread::getExecutionMode() {
   assert(VM::isHotspot());
-
   VMThread* vm_thread = VMThread::current();
-  // Not a JVM thread - native thread, e.g. thread launched by JNI code
   if (vm_thread == nullptr) {
     return ExecutionMode::NATIVE;
   }
 
-  ProfiledThread *prof_thread = ProfiledThread::currentSignalSafe();
-  bool is_java_thread = prof_thread != nullptr && prof_thread->isJavaThread();
-
-  // A Java thread that JVM tells us via jvmti `ThreadStart()` callback.
-  if (is_java_thread) {
-    int raw_thread_state = vm_thread->state();
-
-    // Java threads: [4, 12) = [_thread_in_native, _thread_max_state)
-    // JVM internal threads: 0 or outside this range
-    is_java_thread = raw_thread_state >= 4 && raw_thread_state < 12;
-
-    return is_java_thread ? convertJvmExecutionState(raw_thread_state)
-                        : ExecutionMode::JVM;
+  if (isJavaThread(vm_thread)) {
+    JVMJavaThreadState thread_state = vm_thread->state();
+    return convertJvmExecutionState(thread_state);
   } else {
-    // It is a JVM internal thread, may or may not be a Java thread, 
-    // e.g. Compiler thread or GC thread, etc
     return ExecutionMode::JVM;
   }
 }
@@ -682,13 +704,8 @@ OSThreadState VMThread::getOSThreadState() {
   if (vm_thread == nullptr) {
     return OSThreadState::UNKNOWN;
   }
-  int raw_thread_state = vm_thread->state();
-  bool is_java_thread = raw_thread_state >= 4 && raw_thread_state < 12;
-  OSThreadState state = OSThreadState::UNKNOWN;
-  if (is_java_thread) {
-    state = vm_thread->osThreadState();
-  }
-  return state;
+  // All hotspot JVM threads have osThread fields
+  return vm_thread->osThreadState();
 }
 
 int VMThread::osThreadId() {
@@ -704,7 +721,7 @@ JNIEnv* VMThread::jni() {
     if (_env_offset < 0) {
         return VM::jni();  // fallback for non-HotSpot JVM
     }
-    return isJavaThread() ? (JNIEnv*) at(_env_offset) : NULL;
+    return isJavaThread(this) ? (JNIEnv*) at(_env_offset) : NULL;
 }
 
 jmethodID VMMethod::id() {
@@ -983,23 +1000,21 @@ OSThreadState VMThread::osThreadState() {
     return OSThreadState::UNKNOWN;
 }
 
-int VMThread::state() {
-    if (!cachedIsJavaThread()) return 0;
+JVMJavaThreadState VMThread::state() {
+    assert(isJavaThread(this) && "Must be a Java thread");
+    int state = 0;
     int offset = VMStructs::thread_state_offset();
     if (offset >= 0) {
-        int* state = (int*)at(offset);
-        if (state == nullptr) {
-            return 0;
-        } else {
-            int value = SafeAccess::safeFetch32(state, 0);
+        int* state_addr = (int*)at(offset);
+        if (state_addr != nullptr) {
+            state = SafeAccess::safeFetch32(state_addr, 0);
             // Checking for bad data
-            if (value > _thread_max_state) {
-                value = 0;
+            if (state >= _thread_max_state || state < 0) {
+                state = 0;
             }
-            return value;
         }
     }
-    return 0;
+    return static_cast<JVMJavaThreadState>(state);
 }
 
 bool HeapUsage::is_jmx_attempted = false;

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

@@ -699,56 +699,19 @@ DECLARE(VMThread)
 
     static inline VMThread* current();
     static inline VMThread* fromJavaThread(JNIEnv* env, jthread thread);
+    static bool isJavaThread(VMThread* thread);
     static ExecutionMode getExecutionMode();
     static OSThreadState getOSThreadState();
 
     int osThreadId();
     JNIEnv* jni();
 
-    const void** vtable() {
-        assert(SafeAccess::isReadable(this));
-        return *(const void***)this;
-    }
-
-    // This thread is considered a JavaThread if at least 2 of the selected 3 vtable entries
-    // match those of a known JavaThread (which is either application thread or AttachListener).
-    // Indexes were carefully chosen to work on OpenJDK 8 to 25, both product an debug builds.
-    bool isJavaThread() {
-        const void** vtbl = vtable();
-        return (vtbl[1] == _java_thread_vtbl[1]) +
-               (vtbl[3] == _java_thread_vtbl[3]) +
-               (vtbl[5] == _java_thread_vtbl[5]) >= 2;
-    }
-
-    // Cached version of isJavaThread(). On first call per thread, computes the
-    // vtable check and caches the result in ProfiledThread for O(1) subsequent
-    // lookups. This is needed because JVMTI ThreadStart only fires for application
-    // threads, not for JVM-internal JavaThread subclasses (CompilerThread, etc.).
-    bool cachedIsJavaThread() {
-        ProfiledThread* pt = ProfiledThread::currentSignalSafe();
-        if (pt != NULL) {
-            if (!pt->isJavaThreadKnown()) {
-                pt->cacheJavaThread(isJavaThread());
-            }
-            bool result = pt->isJavaThread();
-            if (!result) Counters::increment(WALKVM_CACHED_NOT_JAVA);
-            return result;
-        }
-        bool result = isJavaThread();
-        if (!result) Counters::increment(WALKVM_CACHED_NOT_JAVA);
-        return result;
-    }
-
     OSThreadState osThreadState();
 
-    int state();
-
-    bool inJava() {
-        return state() == 8;
-    }
+    JVMJavaThreadState state();
 
     bool inDeopt() {
-        if (!cachedIsJavaThread()) return false;
+        if (!isJavaThread(this)) return false;
         assert(_thread_vframe_offset >= 0);
         return SafeAccess::loadPtr((void**) at(_thread_vframe_offset), nullptr) != NULL;
     }
@@ -789,14 +752,17 @@ DECLARE(VMThread)
     }
 
     NOADDRSANITIZE VMJavaFrameAnchor* anchor() {
-        if (!cachedIsJavaThread()) return NULL;
+        if (!isJavaThread(this)) return NULL;
         assert(_thread_anchor_offset >= 0);
         return VMJavaFrameAnchor::cast(at(_thread_anchor_offset));
     }
 
     inline VMMethod* compiledMethod();
 private:
     static inline int nativeThreadId(JNIEnv* jni, jthread thread);
+    inline void** vtable();
+    inline bool hasJavaThreadVtable();
+
 DECLARE_END
 
 DECLARE(VMConstMethod)

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

@@ -31,14 +31,29 @@ int VMThread::nativeThreadId(JNIEnv* jni, jthread thread) {
     return -1;
 }
 
+void** VMThread::vtable() {
+    assert(SafeAccess::isReadable(this));
+    return *(void***)this;
+}
+
+// This thread is considered a JavaThread if at least 2 of the selected 3 vtable entries
+// match those of a known JavaThread (which is either application thread or AttachListener).
+// Indexes were carefully chosen to work on OpenJDK 8 to 25, both product an debug builds.
+bool VMThread::hasJavaThreadVtable() {
+    void** vtbl = vtable();
+    return (SafeAccess::load(&vtbl[1]) == _java_thread_vtbl[1]) +
+           (SafeAccess::load(&vtbl[3]) == _java_thread_vtbl[3]) +
+           (SafeAccess::load(&vtbl[5]) == _java_thread_vtbl[5]) >= 2;
+}
+
 VMNMethod* VMMethod::code() {
     assert(_method_code_offset >= 0);
     const void* code_ptr = *(const void**) at(_method_code_offset);
     return VMNMethod::cast(code_ptr);
 }
 
 VMMethod* VMThread::compiledMethod() {
-    if (!cachedIsJavaThread()) return NULL;
+    if (!isJavaThread(this)) return NULL;
     assert(_comp_method_offset >= 0);
     assert(_comp_env_offset >= 0);
     assert(_comp_task_offset >= 0);

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

@@ -68,7 +68,7 @@ static CTimer ctimer;
 void Profiler::onThreadStart(jvmtiEnv *jvmti, JNIEnv *jni, jthread thread) {
   ProfiledThread::initCurrentThread();
   ProfiledThread *current = ProfiledThread::current();
-  current->setJavaThread();
+  current->setJavaThread(true);
   int tid = current->tid();
   if (_thread_filter.enabled()) {
     int slot_id = _thread_filter.registerThread();

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

@@ -19,7 +19,15 @@
 #include <sys/types.h>
 #include <vector>
 
-class ProfiledThread : public ThreadLocalData {     
+class ProfiledThread : public ThreadLocalData {    
+public:
+  enum ThreadType : u32 {
+    TYPE_UNKNOWN = 0,
+    TYPE_JAVA_THREAD = 0x1,
+    TYPE_NOT_JAVA_THREAD = 0x2,
+    TYPE_MASK = TYPE_JAVA_THREAD | TYPE_NOT_JAVA_THREAD
+  };
+
 private:
   // We are allowing several levels of nesting because we can be
   // eg. in a crash handler when wallclock signal kicks in,
@@ -34,10 +42,6 @@ class ProfiledThread : public ThreadLocalData {
   static volatile int _running_buffer_pos;
   static ProfiledThread** _buffer;
 
-  // Misc flags
-  static constexpr u32 FLAG_JAVA_THREAD = 0x01;
-  static constexpr u32 FLAG_JAVA_THREAD_KNOWN = 0x02;
-
   // Free slot recycling - lock-free stack of available buffer slots
   // Note: Using plain int with GCC atomic builtins instead of std::atomic
   // because std::atomic is not guaranteed async-signal-safe (may use mutexes)
@@ -201,32 +205,17 @@ class ProfiledThread : public ThreadLocalData {
     return &_otel_ctx_record;
   }
 
-  // JavaThread status cache — avoids repeated vtable checks in VMThread::isJavaThread().
-  // JVMTI ThreadStart only fires for application threads, not for JVM-internal
-  // JavaThread subclasses (CompilerThread, etc.), so we cache the vtable result
-  // here for O(1) subsequent lookups via VMThread::cachedIsJavaThread().
-
-  // Called from JVMTI ThreadStart callback for threads known to be Java threads.
-  inline void setJavaThread() {
-    _misc_flags |= (FLAG_JAVA_THREAD | FLAG_JAVA_THREAD_KNOWN);
-  }
-
-  // Returns the cached JavaThread status. Only valid after setJavaThread() or
-  // cacheJavaThread() has been called (check isJavaThreadKnown() first).
-  inline bool isJavaThread() const {
-    return (_misc_flags & FLAG_JAVA_THREAD) != 0;
-  }
-
-  // Returns true if the JavaThread status has been determined and cached.
-  inline bool isJavaThreadKnown() const {
-    return (_misc_flags & FLAG_JAVA_THREAD_KNOWN) != 0;
+  // Record java thread state
+  inline void setJavaThread(bool is_java) {
+    if (is_java) {
+      _misc_flags = ((_misc_flags & ~TYPE_MASK) | TYPE_JAVA_THREAD);
+    } else {
+      _misc_flags = ((_misc_flags & ~TYPE_MASK) | TYPE_NOT_JAVA_THREAD);
+    }
   }
 
-  // Caches the result of VMThread::isJavaThread() vtable check.
-  // Sets FLAG_JAVA_THREAD_KNOWN unconditionally; sets FLAG_JAVA_THREAD if isJava is true.
-  // Written from the signal handler on the owning thread — no cross-thread visibility concern.
-  inline void cacheJavaThread(bool isJava) {
-    _misc_flags |= FLAG_JAVA_THREAD_KNOWN | (isJava ? FLAG_JAVA_THREAD : 0);
+  inline enum ThreadType threadType() const {
+    return static_cast<ThreadType>(_misc_flags & TYPE_MASK);
   }
 
   inline bool isCrashProtectionActive() const { return _crash_protection_active; }