Refactor stackwalker

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

@@ -16,7 +16,6 @@
 #include "incbin.h"
 #include "jfrMetadata.h"
 #include "jniHelper.h"
-#include "jvm.h"
 #include "os.h"
 #include "profiler.h"
 #include "rustDemangler.h"

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

@@ -0,0 +1,18 @@
+#ifndef _FRAMES_H
+#define _FRAMES_H
+
+#include <jni.h>
+#include "vmEntry.h"
+
+inline int makeFrame(ASGCT_CallFrame *frames, jint type, jmethodID id) {
+  frames[0].bci = type;
+  frames[0].method_id = id;
+  return 1;
+}
+
+inline int makeFrame(ASGCT_CallFrame *frames, jint type,
+                            const char *id) {
+  return makeFrame(frames, type, (jmethodID)id);
+}
+
+#endif // _FRAMES_H

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright The async-profiler authors
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef _HOTSPOT_HOTSPOTSUPPORT_H
+#define _HOTSPOT_HOTSPOTSUPPORT_H
+
+#include "stackWalker.h"
+
+class ProfiledThread;
+
+class HotspotSupport {
+private:
+    static int 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);
+    static int walkVM(void* ucontext, ASGCT_CallFrame* frames, int max_depth, 
+                      StackWalkFeatures features, EventType event_type,
+                      int lock_index, bool* truncated = nullptr);
+
+    static int getJavaTraceAsync(void *ucontext, ASGCT_CallFrame *frames,
+                                 int max_depth, StackContext *java_ctx,
+                                 bool *truncated);
+
+public:
+    static void checkFault(ProfiledThread* thrd = nullptr);
+    static int walkJavaStack(StackWalkRequest& request);
+};
+
+#endif // _HOTSPOT_HOTSPOTSUPPORT_H

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

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "jitCodeCache.h"
+
+#include "hotspot/vmStructs.h"
+
+SpinLock JitCodeCache::_stubs_lock;
+CodeCache JitCodeCache::_runtime_stubs("[stubs]");
+std::atomic<const void *> JitCodeCache::_call_stub_begin = { nullptr };
+std::atomic<const void *> JitCodeCache::_call_stub_end = { nullptr };
+
+// CompiledMethodLoad is also needed to enable DebugNonSafepoints info by
+// default
+void JNICALL JitCodeCache::CompiledMethodLoad(jvmtiEnv *jvmti, jmethodID method,
+                                jint code_size, const void *code_addr,
+                                jint map_length,
+                                const jvmtiAddrLocationMap *map,
+                                const void *compile_info) {
+  CodeHeap::updateBounds(code_addr, (const char *)code_addr + code_size);
+}
+
+void JNICALL JitCodeCache::DynamicCodeGenerated(jvmtiEnv *jvmti, const char *name,
+                                  const void *address, jint length) {
+  _stubs_lock.lock();
+  _runtime_stubs.add(address, length, name, true);
+  _stubs_lock.unlock();
+
+  if (name[0] == 'I' && strcmp(name, "Interpreter") == 0) {
+    CodeHeap::setInterpreterStart(address);
+  } else if (strcmp(name, "call_stub") == 0) {
+    _call_stub_begin.store(address, std::memory_order_relaxed);
+    // This fence ensures that _call_stub_begin is visible before _call_stub_end, so that isCallStub() works correctly
+    std::atomic_thread_fence(std::memory_order_release);
+    _call_stub_end.store((const char *)address + length, std::memory_order_relaxed);
+  }
+
+  CodeHeap::updateBounds(address, (const char *)address + length);
+}
+
+CodeBlob* JitCodeCache::findRuntimeStub(const void *address) {
+  CodeBlob *stub = nullptr;
+  _stubs_lock.lockShared();
+  if (_runtime_stubs.contains(address)) {
+    stub = _runtime_stubs.findBlobByAddress(address);
+  }
+  _stubs_lock.unlockShared();
+  return stub;
+}

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef _HOTSPOT_JITCODECACHE_H
+#define _HOTSPOT_JITCODECACHE_H
+
+#include <atomic>
+
+#include "codeCache.h"
+#include "spinLock.h"
+
+// The class tracks JIT-compiled code and ranges
+class JitCodeCache {
+private:
+    static SpinLock _stubs_lock;
+    static CodeCache _runtime_stubs;
+    static std::atomic<const void *> _call_stub_begin;
+    static std::atomic<const void *> _call_stub_end;
+
+public:
+    static void JNICALL CompiledMethodLoad(jvmtiEnv *jvmti, jmethodID method,
+                                         jint code_size, const void *code_addr,
+                                         jint map_length,
+                                         const jvmtiAddrLocationMap *map,
+                                         const void *compile_info);
+    static void JNICALL DynamicCodeGenerated(jvmtiEnv *jvmti, const char *name,
+                                           const void *address, jint length);
+
+    static inline bool isCallStub(const void *address) {
+        const void* stub_end = _call_stub_end.load(std::memory_order_acquire);
+        return stub_end != nullptr &&
+            address >= _call_stub_begin.load(std::memory_order_relaxed) &&
+            address < stub_end;    
+    }
+
+    static CodeBlob* findRuntimeStub(const void *address);
+};
+
+#endif // _HOTSPOT_JITCODECACHE_H

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

@@ -41,6 +41,7 @@ inline bool crashProtectionActive();
 
 template <typename T>
 inline T* cast_to(const void* ptr) {
+    assert(VM::isHotspot()); // This should only be used in HotSpot-specific code
     assert(T::type_size() > 0); // Ensure type size has been initialized
     assert(crashProtectionActive() || ptr == nullptr || SafeAccess::isReadableRange(ptr, T::type_size()));
     return reinterpret_cast<T*>(const_cast<void*>(ptr));

ddprof-lib/src/main/cpp/j9/j9Ext.cpp → ddprof-lib/src/main/cpp/j9/j9Support.cpp

@@ -15,23 +15,23 @@
  * limitations under the License.
  */
 
-#include "j9Ext.h"
+#include "j9/j9Support.h"
 #include "os.h"
 #include <string.h>
 
-jvmtiEnv *J9Ext::_jvmti;
+jvmtiEnv *J9Support::_jvmti;
 
-void *(*J9Ext::_j9thread_self)() = NULL;
+void *(*J9Support::_j9thread_self)() = NULL;
 
-jvmtiExtensionFunction J9Ext::_GetOSThreadID = NULL;
-jvmtiExtensionFunction J9Ext::_GetJ9vmThread = NULL;
-jvmtiExtensionFunction J9Ext::_GetStackTraceExtended = NULL;
-jvmtiExtensionFunction J9Ext::_GetAllStackTracesExtended = NULL;
+jvmtiExtensionFunction J9Support::_GetOSThreadID = NULL;
+jvmtiExtensionFunction J9Support::_GetJ9vmThread = NULL;
+jvmtiExtensionFunction J9Support::_GetStackTraceExtended = NULL;
+jvmtiExtensionFunction J9Support::_GetAllStackTracesExtended = NULL;
 
-int J9Ext::InstrumentableObjectAlloc_id = -1;
+int J9Support::InstrumentableObjectAlloc_id = -1;
 
 // Look for OpenJ9-specific JVM TI extension
-bool J9Ext::initialize(jvmtiEnv *jvmti, const void *j9thread_self) {
+bool J9Support::initialize(jvmtiEnv *jvmti, const void *j9thread_self) {
   _jvmti = jvmti;
   _j9thread_self = (void *(*)())j9thread_self;
 

ddprof-lib/src/main/cpp/j9/j9Ext.h → ddprof-lib/src/main/cpp/j9/j9Support.h

@@ -15,8 +15,8 @@
  * limitations under the License.
  */
 
-#ifndef _J9_J9EXT_H
-#define _J9_J9EXT_H
+#ifndef _J9_J9SUPPORT_H
+#define _J9_J9SUPPORT_H
 
 #include <jvmti.h>
 
@@ -70,7 +70,7 @@ static inline int sanitizeJ9FrameType(jint j9_type) {
   return FRAME_JIT_COMPILED;
 }
 
-class J9Ext {
+class J9Support {
   friend class JVMThread;
   friend class J9WallClock;
 private:
@@ -181,4 +181,4 @@ class J9Ext {
     }
 };
 
-#endif // _J9_J9EXT_H
+#endif // _J9_J9SUPPORT_H

ddprof-lib/src/main/cpp/j9/j9WallClock.cpp

@@ -16,7 +16,7 @@
  */
 
 #include "j9WallClock.h"
-#include "j9Ext.h"
+#include "j9/j9Support.h"
 #include "profiler.h"
 #include "threadState.h"
 #include <stdlib.h>
@@ -74,7 +74,7 @@ void J9WallClock::timerLoop() {
 
     jvmtiStackInfoExtended *stack_infos;
     jint thread_count;
-    if (J9Ext::GetAllStackTracesExtended(
+    if (J9Support::GetAllStackTracesExtended(
             _max_stack_depth, (void **)&stack_infos, &thread_count) == 0) {
       for (int i = 0; i < thread_count; i++) {
         jvmtiStackInfoExtended *si = &stack_infos[i];
@@ -95,7 +95,7 @@ void J9WallClock::timerLoop() {
           frames[j].bci = FrameType::encode(sanitizeJ9FrameType(fi->type), fi->location);
         }
 
-        int tid = J9Ext::GetOSThreadID(si->thread);
+        int tid = J9Support::GetOSThreadID(si->thread);
         if (tid == -1) {
           // clearly an invalid TID; skip the thread
           continue;

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "jvmSupport.h"
+
+#include "asyncSampleMutex.h"
+#include "frames.h"
+#include "os.h"
+#include "profiler.h"
+#include "thread.h"
+
+#include "hotspot/hotspotSupport.h"
+
+#include <jni.h>
+
+int JVMSupport::walkJavaStack(StackWalkRequest& request) {
+    if (VM::isHotspot()) {
+        return HotspotSupport::walkJavaStack(request);
+    } else if (VM::isOpenJ9() || VM::isZing()) {
+        assert(request.event_type == BCI_CPU || request.event_type == BCI_WALL);
+        return asyncGetCallTrace(request.frames, request.max_depth, request.ucontext);
+    }
+    assert(false && "Unsupported JVM");
+    return 0;
+}
+
+int JVMSupport::asyncGetCallTrace(ASGCT_CallFrame *frames, int max_depth, void* ucontext) {
+    JNIEnv *jni = VM::jni();
+    if (jni == nullptr) {
+        return 0;
+    }
+
+    AsyncSampleMutex mutex(ProfiledThread::currentSignalSafe());
+    if (!mutex.acquired()) {
+        return 0;
+    }
+
+    JitWriteProtection jit(false);
+    // AsyncGetCallTrace writes to ASGCT_CallFrame array
+    ASGCT_CallTrace trace = {jni, 0, frames};
+    VM::_asyncGetCallTrace(&trace, max_depth, ucontext);
+    if (trace.num_frames > 0) {
+        return trace.num_frames;
+    }
+
+    const char* err_string = Profiler::asgctError(trace.num_frames);
+    if (err_string == NULL) {
+        // No Java stack, because thread is not in Java context
+        return 0;
+    }
+
+    Profiler::instance()->incFailure(-trace.num_frames);
+    return makeFrame(frames, BCI_ERROR, err_string);
+}

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef _JVMSUPPORT_H
+#define _JVMSUPPORT_H
+
+#include "stackWalker.h"
+#include "vmEntry.h"
+
+// Stack recovery techniques used to workaround AsyncGetCallTrace flaws.
+// Can be disabled with 'safemode' option.
+enum StackRecovery {
+  UNKNOWN_JAVA = (1 << 0),
+  POP_STUB = (1 << 1),
+  POP_METHOD = (1 << 2),
+  LAST_JAVA_PC = (1 << 4),
+  GC_TRACES = (1 << 5),
+  PROBE_SP = 0x100,
+};
+
+
+class JVMSupport {
+    static int asyncGetCallTrace(ASGCT_CallFrame *frames, int max_depth, void* ucontext);
+public:
+    static int walkJavaStack(StackWalkRequest& request);
+};
+
+#endif // _JVMSUPPORT_H