GH-126910: Add GNU backtrace support for unwinding JIT frames

Author: diegorusso

Include/internal/pycore_jit_publish.h

@@ -0,0 +1,27 @@
+#ifndef Py_INTERNAL_JIT_PUBLISH_H
+#define Py_INTERNAL_JIT_PUBLISH_H
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#include <stddef.h>
+
+typedef struct _PyJitCodeRegistration _PyJitCodeRegistration;
+
+#ifdef _Py_JIT
+
+/* Return a teardown handle when any backend stores registration state.
+ * A NULL result is valid when publication succeeded only through backends
+ * with no unregister step, such as perf map output.
+ */
+_PyJitCodeRegistration *_PyJit_RegisterCode(const void *code_addr,
+                                            size_t code_size,
+                                            const char *entry,
+                                            const char *filename);
+
+void _PyJit_UnregisterCode(_PyJitCodeRegistration *registration);
+
+#endif  // _Py_JIT
+
+#endif  // Py_INTERNAL_JIT_PUBLISH_H

Include/internal/pycore_jit_unwind.h

@@ -10,9 +10,12 @@
 
 #if defined(_Py_JIT) && defined(__linux__) && defined(__ELF__)
 #  define PY_HAVE_JIT_GDB_UNWIND
+#  define PY_HAVE_JIT_GNU_BACKTRACE_UNWIND
 #endif
 
-#if defined(PY_HAVE_PERF_TRAMPOLINE) || defined(PY_HAVE_JIT_GDB_UNWIND)
+#if defined(PY_HAVE_PERF_TRAMPOLINE) \
+    || defined(PY_HAVE_JIT_GDB_UNWIND) \
+    || defined(PY_HAVE_JIT_GNU_BACKTRACE_UNWIND)
 
 #if defined(PY_HAVE_JIT_GDB_UNWIND)
 extern PyMutex _Py_jit_debug_mutex;
@@ -63,6 +66,13 @@ void *_PyJitUnwind_GdbRegisterCode(const void *code_addr,
 
 void _PyJitUnwind_GdbUnregisterCode(void *handle);
 
-#endif  // defined(PY_HAVE_PERF_TRAMPOLINE) || defined(PY_HAVE_JIT_GDB_UNWIND)
+#if defined(PY_HAVE_JIT_GNU_BACKTRACE_UNWIND)
+void *_PyJitUnwind_GnuBacktraceRegisterCode(const void *code_addr,
+                                            size_t code_size);
+
+void _PyJitUnwind_GnuBacktraceUnregisterCode(void *handle);
+#endif
+
+#endif  // JIT unwind support
 
 #endif  // Py_INTERNAL_JIT_UNWIND_H

Include/internal/pycore_optimizer.h

@@ -9,6 +9,7 @@ extern "C" {
 #endif
 
 #include "pycore_typedefs.h"      // _PyInterpreterFrame
+#include "pycore_jit_publish.h"
 #include "pycore_uop.h"           // _PyUOpInstruction
 #include "pycore_uop_ids.h"
 #include "pycore_stackref.h"      // _PyStackRef
@@ -198,7 +199,7 @@ typedef struct _PyExecutorObject {
     uint32_t code_size;
     size_t jit_size;
     void *jit_code;
-    void *jit_gdb_handle;
+    _PyJitCodeRegistration *jit_registration;
     _PyExitData exits[1];
 } _PyExecutorObject;
 

Lib/test/test_frame_pointer_unwind.py

@@ -70,7 +70,7 @@ def _frame_pointers_expected(machine):
     return None
 
 
-def _build_stack_and_unwind():
+def _build_stack_and_unwind(unwinder):
     import operator
 
     def build_stack(n, unwinder, warming_up_caller=False):
@@ -89,7 +89,7 @@ def build_stack(n, unwinder, warming_up_caller=False):
             result = operator.call(build_stack, n - 1, unwinder, warming_up)
         return result
 
-    stack = build_stack(10, _testinternalcapi.manual_frame_pointer_unwind)
+    stack = build_stack(10, unwinder)
     return stack
 
 
@@ -112,8 +112,9 @@ def _classify_stack(stack, jit_enabled):
     return annotated, python_frames, jit_frames, other_frames
 
 
-def _annotate_unwind():
-    stack = _build_stack_and_unwind()
+def _annotate_unwind(unwinder_name="manual_frame_pointer_unwind"):
+    unwinder = getattr(_testinternalcapi, unwinder_name)
+    stack = _build_stack_and_unwind(unwinder)
     jit_enabled = hasattr(sys, "_jit") and sys._jit.is_enabled()
     jit_backend = _testinternalcapi.get_jit_backend()
     ranges = _testinternalcapi.get_jit_code_ranges() if jit_enabled else []
@@ -132,13 +133,14 @@ def _annotate_unwind():
         "jit_frames": jit_frames,
         "other_frames": other_frames,
         "jit_backend": jit_backend,
+        "unwinder": unwinder_name,
     })
 
 
-def _manual_unwind_length(**env):
+def _unwind_result(unwinder_name, **env):
     code = (
         "from test.test_frame_pointer_unwind import _annotate_unwind; "
-        "print(_annotate_unwind());"
+        f"print(_annotate_unwind({unwinder_name!r}));"
     )
     run_env = os.environ.copy()
     run_env.update(env)
@@ -197,7 +199,7 @@ def test_manual_unwind_respects_frame_pointers(self):
 
         for env, using_jit in envs:
             with self.subTest(env=env):
-                result = _manual_unwind_length(**env)
+                result = _unwind_result("manual_frame_pointer_unwind", **env)
                 jit_frames = result["jit_frames"]
                 python_frames = result.get("python_frames", 0)
                 jit_backend = result.get("jit_backend")
@@ -240,5 +242,51 @@ def test_manual_unwind_respects_frame_pointers(self):
                     )
 
 
+@support.requires_gil_enabled("test requires the GIL enabled")
+@unittest.skipIf(support.is_wasi, "test not supported on WASI")
+@unittest.skipUnless(sys.platform == "linux", "GNU backtrace unwinding test requires Linux")
+class GnuBacktraceUnwindTests(unittest.TestCase):
+
+    def setUp(self):
+        super().setUp()
+        try:
+            _testinternalcapi.gnu_backtrace_unwind()
+        except RuntimeError as exc:
+            if "not supported" in str(exc):
+                self.skipTest("gnu backtrace unwinding not supported on this platform")
+            raise
+
+    def test_gnu_backtrace_unwinds_through_jit_frames(self):
+        jit_available = hasattr(sys, "_jit") and sys._jit.is_available()
+        envs = [({"PYTHON_JIT": "0"}, False)]
+        if jit_available:
+            envs.append(({"PYTHON_JIT": "1"}, True))
+
+        for env, using_jit in envs:
+            with self.subTest(env=env):
+                result = _unwind_result("gnu_backtrace_unwind", **env)
+                python_frames = result.get("python_frames", 0)
+                jit_frames = result.get("jit_frames", 0)
+                jit_backend = result.get("jit_backend")
+
+                self.assertGreater(
+                    python_frames,
+                    0,
+                    f"expected to find Python frames in GNU backtrace with env {env}",
+                )
+                if using_jit and jit_backend == "jit":
+                    self.assertGreater(
+                        jit_frames,
+                        0,
+                        f"expected GNU backtrace to include JIT frames with env {env}",
+                    )
+                else:
+                    self.assertEqual(
+                        jit_frames,
+                        0,
+                        f"unexpected JIT frames counted in GNU backtrace with env {env}",
+                    )
+
+
 if __name__ == "__main__":
     unittest.main()

Makefile.pre.in

@@ -470,6 +470,7 @@ PYTHON_OBJS=	\
 		Python/instruction_sequence.o \
 		Python/intrinsics.o \
 		Python/jit.o \
+		Python/jit_publish.o \
 		$(JIT_OBJS) \
 		Python/legacy_tracing.o \
 		Python/lock.o \

Modules/_testinternalcapi.c

@@ -47,6 +47,9 @@
 #if defined(HAVE_DLADDR) && !defined(__wasi__)
 #  include <dlfcn.h>
 #endif
+#if defined(HAVE_EXECINFO_H)
+#  include <execinfo.h>
+#endif
 #ifdef MS_WINDOWS
 #  include <windows.h>
 #  include <intrin.h>
@@ -58,6 +61,7 @@
 
 
 static const uintptr_t min_frame_pointer_addr = 0x1000;
+#define MAX_UNWIND_FRAMES 200
 
 
 static PyObject *
@@ -328,7 +332,6 @@ get_jit_backend(PyObject *self, PyObject *Py_UNUSED(args))
 static PyObject *
 manual_unwind_from_fp(uintptr_t *frame_pointer)
 {
-    Py_ssize_t max_depth = 200;
     int stack_grows_down = _Py_STACK_GROWS_DOWN;
 
     if (frame_pointer == NULL) {
@@ -340,14 +343,20 @@ manual_unwind_from_fp(uintptr_t *frame_pointer)
         return NULL;
     }
 
-    for (Py_ssize_t depth = 0;
-         depth < max_depth && frame_pointer != NULL;
-         depth++)
-    {
+    Py_ssize_t depth = 0;
+    while (frame_pointer != NULL) {
         uintptr_t fp_addr = (uintptr_t)frame_pointer;
         if ((fp_addr % sizeof(uintptr_t)) != 0) {
             break;
         }
+        if (depth >= MAX_UNWIND_FRAMES) {
+            Py_DECREF(result);
+            PyErr_Format(
+                PyExc_RuntimeError,
+                "manual frame pointer unwind returned more than %d frames",
+                MAX_UNWIND_FRAMES);
+            return NULL;
+        }
         uintptr_t return_addr = frame_pointer[1];
 
         PyObject *addr_obj = PyLong_FromUnsignedLongLong(return_addr);
@@ -361,6 +370,7 @@ manual_unwind_from_fp(uintptr_t *frame_pointer)
             return NULL;
         }
         Py_DECREF(addr_obj);
+        depth++;
 
         uintptr_t *next_fp = (uintptr_t *)frame_pointer[0];
         // Stop if the frame pointer is extremely low.
@@ -383,6 +393,49 @@ manual_unwind_from_fp(uintptr_t *frame_pointer)
 
     return result;
 }
+
+#if defined(HAVE_EXECINFO_H) && defined(HAVE_BACKTRACE)
+static PyObject *
+gnu_backtrace_unwind(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    void *addresses[MAX_UNWIND_FRAMES + 1];
+    int frame_count = backtrace(addresses, (int)Py_ARRAY_LENGTH(addresses));
+    if (frame_count < 0) {
+        PyErr_SetString(PyExc_RuntimeError, "backtrace() failed");
+        return NULL;
+    }
+    if (frame_count > MAX_UNWIND_FRAMES) {
+        PyErr_Format(
+            PyExc_RuntimeError,
+            "backtrace() returned more than %d frames",
+            MAX_UNWIND_FRAMES);
+        return NULL;
+    }
+
+    PyObject *result = PyList_New(frame_count);
+    if (result == NULL) {
+        return NULL;
+    }
+    for (int i = 0; i < frame_count; i++) {
+        PyObject *addr_obj = PyLong_FromUnsignedLongLong((uintptr_t)addresses[i]);
+        if (addr_obj == NULL) {
+            Py_DECREF(result);
+            return NULL;
+        }
+        PyList_SET_ITEM(result, i, addr_obj);
+    }
+    return result;
+}
+#else
+static PyObject *
+gnu_backtrace_unwind(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyErr_SetString(PyExc_RuntimeError,
+                    "gnu_backtrace_unwind is not supported on this platform");
+    return NULL;
+}
+#endif
+
 #if defined(__GNUC__) || defined(__clang__)
 static PyObject *
 manual_frame_pointer_unwind(PyObject *self, PyObject *args)
@@ -2915,6 +2968,7 @@ static PyMethodDef module_functions[] = {
     {"classify_stack_addresses", classify_stack_addresses, METH_VARARGS},
     {"get_jit_code_ranges", get_jit_code_ranges, METH_NOARGS},
     {"get_jit_backend", get_jit_backend, METH_NOARGS},
+    {"gnu_backtrace_unwind", gnu_backtrace_unwind, METH_NOARGS},
     {"manual_frame_pointer_unwind", manual_frame_pointer_unwind, METH_NOARGS},
     {"test_bswap", test_bswap, METH_NOARGS},
     {"test_popcount", test_popcount, METH_NOARGS},

PCbuild/_freeze_module.vcxproj

@@ -237,6 +237,7 @@
     <ClCompile Include="..\Python\intrinsics.c" />
     <ClCompile Include="..\Python\instrumentation.c" />
     <ClCompile Include="..\Python\jit.c" />
+    <ClCompile Include="..\Python\jit_publish.c" />
     <ClCompile Include="..\Python\legacy_tracing.c" />
     <ClCompile Include="..\Python\lock.c" />
     <ClCompile Include="..\Python\marshal.c" />

PCbuild/_freeze_module.vcxproj.filters

@@ -265,6 +265,9 @@
     <ClCompile Include="..\Python\jit.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\jit_publish.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="..\Objects\lazyimportobject.c">
       <Filter>Source Files</Filter>
     </ClCompile>

PCbuild/pythoncore.vcxproj

@@ -648,6 +648,7 @@
     <ClCompile Include="..\Python\instruction_sequence.c" />
     <ClCompile Include="..\Python\instrumentation.c" />
     <ClCompile Include="..\Python\jit.c" />
+    <ClCompile Include="..\Python\jit_publish.c" />
     <ClCompile Include="..\Python\legacy_tracing.c" />
     <ClCompile Include="..\Python\lock.c" />
     <ClCompile Include="..\Python\marshal.c" />

PCbuild/pythoncore.vcxproj.filters

@@ -1481,6 +1481,9 @@
     <ClCompile Include="..\Python\jit.c">
       <Filter>Python</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\jit_publish.c">
+      <Filter>Python</Filter>
+    </ClCompile>
     <ClCompile Include="..\Python\legacy_tracing.c">
       <Filter>Source Files</Filter>
     </ClCompile>