Add tests for the ghost unwinding

Author: pablogsal

setup.py

@@ -337,14 +337,25 @@ def build_js_files(self):
     name="memray._test_utils",
     sources=[
         "src/memray/_memray_test_utils.pyx",
+        "src/memray/_memray/ghost_stack_test_utils.cpp",
+        *GHOST_STACK_SOURCES,
     ],
     language="c++",
     extra_compile_args=["-std=c++17", "-Wall", *EXTRA_COMPILE_ARGS],
     extra_link_args=["-std=c++17", *EXTRA_LINK_ARGS],
+    extra_objects=[*GHOST_STACK_OBJECTS],
     define_macros=DEFINE_MACROS,
     undef_macros=UNDEF_MACROS,
 )
 
+MEMRAY_TEST_EXTENSION.include_dirs = [
+    "src",
+    str(GHOST_STACK_LOCATION / "include"),
+]
+
+if IS_LINUX:
+    MEMRAY_TEST_EXTENSION.libraries = ["unwind"]
+
 MEMRAY_INJECT_EXTENSION = Extension(
     name="memray._inject",
     sources=[

src/memray/_memray/ghost_stack/src/ghost_stack.cpp

@@ -198,8 +198,9 @@ class GhostStackImpl {
             trampolines_installed_ = false;
 
             // Increment epoch to signal state change
-            uint64_t new_epoch = epoch_.fetch_add(1, std::memory_order_release) + 1;
-            LOG_DEBUG("  New epoch=%lu (entries preserved for stale trampolines)\n", (unsigned long)new_epoch);
+            epoch_.fetch_add(1, std::memory_order_release);
+            LOG_DEBUG("  New epoch=%lu (entries preserved for stale trampolines)\n",
+                      (unsigned long)epoch_.load(std::memory_order_acquire));
         }
         LOG_DEBUG("=== reset EXIT ===\n");
     }
@@ -240,8 +241,8 @@ class GhostStackImpl {
                   (unsigned long)epoch_.load(std::memory_order_acquire));
 
         // Increment epoch FIRST to signal any in-flight operations
-        uint64_t new_epoch = epoch_.fetch_add(1, std::memory_order_release) + 1;
-        LOG_DEBUG("  New epoch=%lu\n", (unsigned long)new_epoch);
+        epoch_.fetch_add(1, std::memory_order_release);
+        LOG_DEBUG("  New epoch=%lu\n", (unsigned long)epoch_.load(std::memory_order_acquire));
 
         entries_.clear();
         tail_.store(0, std::memory_order_release);
@@ -347,9 +348,8 @@ class GhostStackImpl {
             // Only update tail_ if we find a match - don't corrupt it during search
             for (size_t i = tail; i > 0; --i) {
                 if (entries_[i - 1].stack_pointer == sp) {
-                    size_t skipped = tail - (i - 1);
                     LOG_DEBUG("longjmp detected: found matching SP at index %zu (skipped %zu frames)\n",
-                              i - 1, skipped);
+                              i - 1, tail - (i - 1));
 
                     // Update tail_ to skip all the frames that were bypassed by longjmp
                     tail_.store(i - 1, std::memory_order_release);
@@ -723,4 +723,4 @@ uintptr_t ghost_exception_handler(void* exception) {
     return ret;
 }
 
-} // extern "C"
+} // extern 

src/memray/_memray/ghost_stack_test_utils.cpp

@@ -0,0 +1,85 @@
+#include "ghost_stack_test_utils.h"
+
+#include <Python.h>
+#include <cstdint>
+
+#ifdef MEMRAY_HAS_GHOST_STACK
+#    include "ghost_stack.h"
+#    define UNW_LOCAL_ONLY
+#    include <libunwind.h>
+#endif
+
+extern "C" {
+
+PyObject*
+ghost_stack_test_backtrace(void)
+{
+#ifdef MEMRAY_HAS_GHOST_STACK
+    void* frames[256];
+    size_t n = ghost_stack_backtrace(frames, 256);
+    PyObject* result = PyList_New(static_cast<Py_ssize_t>(n));
+    if (!result) return nullptr;
+    for (size_t i = 0; i < n; i++) {
+        PyObject* addr = PyLong_FromUnsignedLongLong(reinterpret_cast<uintptr_t>(frames[i]));
+        if (!addr) {
+            Py_DECREF(result);
+            return nullptr;
+        }
+        PyList_SET_ITEM(result, static_cast<Py_ssize_t>(i), addr);
+    }
+    return result;
+#else
+    Py_RETURN_NONE;
+#endif
+}
+
+PyObject*
+libunwind_test_backtrace(void)
+{
+#ifdef MEMRAY_HAS_GHOST_STACK
+    void* frames[256];
+    int n = unw_backtrace(frames, 256);
+    if (n < 0) n = 0;
+    PyObject* result = PyList_New(static_cast<Py_ssize_t>(n));
+    if (!result) return nullptr;
+    for (int i = 0; i < n; i++) {
+        PyObject* addr = PyLong_FromUnsignedLongLong(reinterpret_cast<uintptr_t>(frames[i]));
+        if (!addr) {
+            Py_DECREF(result);
+            return nullptr;
+        }
+        PyList_SET_ITEM(result, static_cast<Py_ssize_t>(i), addr);
+    }
+    return result;
+#else
+    Py_RETURN_NONE;
+#endif
+}
+
+void
+ghost_stack_test_reset(void)
+{
+#ifdef MEMRAY_HAS_GHOST_STACK
+    ghost_stack_reset();
+#endif
+}
+
+void
+ghost_stack_test_init(void)
+{
+#ifdef MEMRAY_HAS_GHOST_STACK
+    ghost_stack_init(nullptr);
+#endif
+}
+
+int
+ghost_stack_test_has_support(void)
+{
+#ifdef MEMRAY_HAS_GHOST_STACK
+    return 1;
+#else
+    return 0;
+#endif
+}
+
+}  // extern "C"

src/memray/_memray/ghost_stack_test_utils.h

@@ -0,0 +1,29 @@
+#pragma once
+
+#include <Python.h>
+#include <stddef.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Returns a Python list of frame addresses from ghost_stack_backtrace
+// Returns Py_None if MEMRAY_HAS_GHOST_STACK is not defined
+PyObject* ghost_stack_test_backtrace(void);
+
+// Returns a Python list of frame addresses from unw_backtrace (libunwind)
+// Returns Py_None if MEMRAY_HAS_GHOST_STACK is not defined
+PyObject* libunwind_test_backtrace(void);
+
+// Reset ghost_stack shadow stack
+void ghost_stack_test_reset(void);
+
+// Initialize ghost_stack
+void ghost_stack_test_init(void);
+
+// Check if ghost_stack support is available
+int ghost_stack_test_has_support(void);
+
+#ifdef __cplusplus
+}
+#endif

src/memray/_memray_test_utils.pyx

@@ -285,3 +285,50 @@ cdef class PrimeCaches:
         return self
     def __exit__(self, *args):
         sys.setprofile(self.old_profile)
+
+
+# Ghost stack test utilities
+cdef extern from "_memray/ghost_stack_test_utils.h":
+    object ghost_stack_test_backtrace()
+    object libunwind_test_backtrace()
+    void ghost_stack_test_reset()
+    void ghost_stack_test_init()
+    int ghost_stack_test_has_support()
+
+
+def has_ghost_stack_support():
+    """Check if ghost_stack support is available."""
+    return ghost_stack_test_has_support() != 0
+
+
+cdef class GhostStackTestContext:
+    """Context manager for ghost_stack testing.
+
+    Usage:
+        with GhostStackTestContext() as ctx:
+            frames = ctx.backtrace()
+            libunwind_frames = ctx.libunwind_backtrace()
+    """
+
+    def __enter__(self):
+        # init is defensive in case ghost_stack wasn't initialized globally;
+        # reset clears any stale shadow stack state from previous operations
+        ghost_stack_test_init()
+        ghost_stack_test_reset()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        ghost_stack_test_reset()
+        return False
+
+    def backtrace(self):
+        """Capture ghost_stack frames and return as list of addresses."""
+        return ghost_stack_test_backtrace()
+
+    def libunwind_backtrace(self):
+        """Capture libunwind frames for comparison."""
+        return libunwind_test_backtrace()
+
+    def reset(self):
+        """Reset ghost_stack shadow stack."""
+        ghost_stack_test_reset()