GH-126910 jit_unwind: add GDB JIT interface and test

Author: diegorusso

Include/internal/pycore_ceval.h

@@ -108,6 +108,10 @@ extern PyStatus _PyPerfTrampoline_AfterFork_Child(void);
 #ifdef PY_HAVE_PERF_TRAMPOLINE
 extern _PyPerf_Callbacks _Py_perfmap_callbacks;
 extern _PyPerf_Callbacks _Py_perfmap_jit_callbacks;
+extern void _PyPerfJit_WriteNamedCode(const void *code_addr,
+                                      unsigned int code_size,
+                                      const char *entry,
+                                      const char *filename);
 #endif
 
 static inline PyObject*

Include/internal/pycore_jit_unwind.h

@@ -11,13 +11,18 @@ size_t _PyJitUnwind_EhFrameSize(int absolute_addr);
 /*
  * Build DWARF .eh_frame data for JIT code; returns size written or 0 on error.
  * absolute_addr selects the FDE address encoding:
- * - 0: PC-relative offsets.
- * - nonzero: absolute addresses.
+ * - 0: PC-relative offsets (perf jitdump synthesized DSO).
+ * - nonzero: absolute addresses (GDB JIT in-memory ELF).
  */
 size_t _PyJitUnwind_BuildEhFrame(uint8_t *buffer, size_t buffer_size,
                                  const void *code_addr, size_t code_size,
                                  int absolute_addr);
 
+void _PyJitUnwind_GdbRegisterCode(const void *code_addr,
+                                  unsigned int code_size,
+                                  const char *entry,
+                                  const char *filename);
+
 #endif  // PY_HAVE_PERF_TRAMPOLINE
 
 #endif  // Py_CORE_JIT_UNWIND_H

Lib/test/test_gdb/gdb_jit_sample.py

@@ -0,0 +1,20 @@
+# Sample script for use by test_gdb.test_jit
+
+import operator
+import sys
+
+
+def jit_bt_hot(depth, warming_up_caller=False):
+    if warming_up_caller:
+        return
+    if depth == 0:
+        id(42)
+        return
+
+    warming_up = True
+    while warming_up:
+        warming_up = sys._jit.is_enabled() & (not sys._jit.is_active())
+        operator.call(jit_bt_hot, depth - 1, warming_up)
+
+
+jit_bt_hot(10)

Lib/test/test_gdb/test_jit.py

@@ -0,0 +1,35 @@
+import os
+import re
+import sys
+import unittest
+
+from .util import setup_module, DebuggerTests
+
+
+JIT_SAMPLE_SCRIPT = os.path.join(os.path.dirname(__file__), "gdb_jit_sample.py")
+
+
+def setUpModule():
+    setup_module()
+
+
+@unittest.skipUnless(
+    hasattr(sys, "_jit") and sys._jit.is_available(),
+    "requires a JIT-enabled build",
+)
+class JitBacktraceTests(DebuggerTests):
+    def test_bt_unwinds_through_jit_frames(self):
+        gdb_output = self.get_stack_trace(
+            script=JIT_SAMPLE_SCRIPT,
+            cmds_after_breakpoint=["bt"],
+            PYTHON_JIT="1",
+        )
+        self.assertIn("py::jit_executor:<jit>", gdb_output)
+        self.assertIn("py::jit_shim:<jit>", gdb_output)
+        self.assertRegex(
+            gdb_output,
+            re.compile(
+                r"py::jit_executor:<jit>.*(_PyEval_EvalFrameDefault|_PyEval_Vector)",
+                re.DOTALL,
+            ),
+        )

Lib/test/test_gdb/util.py

@@ -160,7 +160,8 @@ def get_stack_trace(self, source=None, script=None,
                         breakpoint=BREAKPOINT_FN,
                         cmds_after_breakpoint=None,
                         import_site=False,
-                        ignore_stderr=False):
+                        ignore_stderr=False,
+                        **env_vars):
         '''
         Run 'python -c SOURCE' under gdb with a breakpoint.
 
@@ -239,7 +240,7 @@ def get_stack_trace(self, source=None, script=None,
             args += [script]
 
         # Use "args" to invoke gdb, capturing stdout, stderr:
-        out, err = run_gdb(*args, PYTHONHASHSEED=PYTHONHASHSEED)
+        out, err = run_gdb(*args, PYTHONHASHSEED=PYTHONHASHSEED, **env_vars)
 
         if not ignore_stderr:
             for line in err.splitlines():

Python/jit.c

@@ -15,6 +15,7 @@
 #include "pycore_interpframe.h"
 #include "pycore_interpolation.h"
 #include "pycore_intrinsics.h"
+#include "pycore_jit_unwind.h"
 #include "pycore_lazyimportobject.h"
 #include "pycore_list.h"
 #include "pycore_long.h"
@@ -60,6 +61,28 @@ jit_error(const char *message)
     PyErr_Format(PyExc_RuntimeWarning, "JIT %s (%d)", message, hint);
 }
 
+static void
+jit_record_code(const void *code_addr, size_t code_size,
+                const char *entry, const char *filename)
+{
+#ifdef PY_HAVE_PERF_TRAMPOLINE
+    _PyPerf_Callbacks callbacks;
+    _PyPerfTrampoline_GetCallbacks(&callbacks);
+    if (callbacks.write_state == _Py_perfmap_jit_callbacks.write_state) {
+        _PyPerfJit_WriteNamedCode(
+            code_addr, (unsigned int)code_size, entry, filename);
+        return;
+    }
+    _PyJitUnwind_GdbRegisterCode(
+        code_addr, (unsigned int)code_size, entry, filename);
+#else
+    (void)code_addr;
+    (void)code_size;
+    (void)entry;
+    (void)filename;
+#endif
+}
+
 static size_t _Py_jit_shim_size = 0;
 
 static int
@@ -731,6 +754,10 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     }
     executor->jit_code = memory;
     executor->jit_size = total_size;
+    jit_record_code(memory,
+                    code_size + state.trampolines.size,
+                    "jit_executor",
+                    "<jit>");
     return 0;
 }
 
@@ -781,6 +808,10 @@ compile_shim(void)
         return NULL;
     }
     _Py_jit_shim_size = total_size;
+    jit_record_code(memory,
+                    code_size + state.trampolines.size,
+                    "jit_shim",
+                    "<jit>");
     return (_PyJitEntryFuncPtr)memory;
 }