GH-126910: add test for manual frame unwinding

Author: diegorusso

Include/internal/pycore_jit.h

@@ -26,6 +26,7 @@ typedef _Py_CODEUNIT *(*jit_func)(
 int _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size_t length);
 void _PyJIT_Free(_PyExecutorObject *executor);
 void _PyJIT_Fini(void);
+PyAPI_FUNC(int) _PyJIT_AddressInJitCode(PyInterpreterState *interp, uintptr_t addr);
 
 #endif  // _Py_JIT
 

Lib/test/test_frame_pointer_unwind.py

@@ -0,0 +1,198 @@
+import json
+import os
+import platform
+import subprocess
+import sys
+import sysconfig
+import unittest
+
+from test import support
+from test.support import import_helper
+
+
+_testinternalcapi = import_helper.import_module("_testinternalcapi")
+
+
+if not support.has_subprocess_support:
+    raise unittest.SkipTest("test requires subprocess support")
+
+
+def _frame_pointers_expected(machine):
+    cflags = " ".join(
+        value for value in (
+            sysconfig.get_config_var("PY_CORE_CFLAGS"),
+            sysconfig.get_config_var("CFLAGS"),
+        )
+        if value
+    )
+    if "no-omit-frame-pointer" in cflags:
+        return True
+    if machine in {"aarch64", "arm64"}:
+        return "-fomit-frame-pointer" not in cflags
+    if machine == "x86_64":
+        return False
+    # MSVC ignores /Oy and /Oy- on x64/ARM64.
+    if sys.platform == "win32" and machine == "arm64":
+        # Windows ARM64 guidelines recommend frame pointers (x29) for stack walking.
+        return True
+    if sys.platform == "win32" and machine == "x86_64":
+        # Windows x64 uses unwind metadata; frame pointers are not required.
+        return None
+    return None
+
+
+def _build_stack_and_unwind():
+    import operator
+
+    def build_stack(n, unwinder, warming_up_caller=False):
+        if warming_up_caller:
+            return
+        if n == 0:
+            return unwinder()
+        warming_up = True
+        while warming_up:
+            # Can't branch on JIT state inside JITted code, so compute here.
+            warming_up = (
+                hasattr(sys, "_jit")
+                and sys._jit.is_enabled()
+                and not sys._jit.is_active()
+            )
+            result = operator.call(build_stack, n - 1, unwinder, warming_up)
+        return result
+
+    stack = build_stack(10, _testinternalcapi.manual_frame_pointer_unwind)
+    return stack
+
+
+def _classify_stack(stack, jit_enabled):
+    labels = _testinternalcapi.classify_stack_addresses(stack, jit_enabled)
+
+    annotated = []
+    jit_frames = 0
+    python_frames = 0
+    other_frames = 0
+    for idx, (frame, tag) in enumerate(zip(stack, labels)):
+        addr = int(frame)
+        if tag == "jit":
+            jit_frames += 1
+        elif tag == "python":
+            python_frames += 1
+        else:
+            other_frames += 1
+        annotated.append((idx, addr, tag))
+    return annotated, python_frames, jit_frames, other_frames
+
+
+def _annotate_unwind():
+    stack = _build_stack_and_unwind()
+    jit_enabled = hasattr(sys, "_jit") and sys._jit.is_enabled()
+    ranges = _testinternalcapi.get_jit_code_ranges() if jit_enabled else []
+    if jit_enabled and ranges:
+        print("JIT ranges:")
+        for start, end in ranges:
+            print(f"  {int(start):#x}-{int(end):#x}")
+    annotated, python_frames, jit_frames, other_frames = _classify_stack(
+        stack, jit_enabled
+    )
+    for idx, addr, tag in annotated:
+        print(f"#{idx:02d} {addr:#x} -> {tag}")
+    return json.dumps({
+        "length": len(stack),
+        "python_frames": python_frames,
+        "jit_frames": jit_frames,
+        "other_frames": other_frames,
+    })
+
+
+def _manual_unwind_length(**env):
+    code = (
+        "from test.test_frame_pointer_unwind import _annotate_unwind; "
+        "print(_annotate_unwind());"
+    )
+    run_env = os.environ.copy()
+    run_env.update(env)
+    proc = subprocess.run(
+        [sys.executable, "-c", code],
+        env=run_env,
+        capture_output=True,
+        text=True,
+    )
+    # Surface the output for debugging/visibility when running this test
+    if proc.stdout:
+        print(proc.stdout, end="")
+    if proc.returncode:
+        raise RuntimeError(
+            f"unwind helper failed (rc={proc.returncode}): {proc.stderr or proc.stdout}"
+        )
+    stdout_lines = proc.stdout.strip().splitlines()
+    if not stdout_lines:
+        raise RuntimeError("unwind helper produced no output")
+    try:
+        return json.loads(stdout_lines[-1])
+    except ValueError as exc:
+        raise RuntimeError(
+            f"unexpected output from unwind helper: {proc.stdout!r}"
+        ) from exc
+
+
+class FramePointerUnwindTests(unittest.TestCase):
+
+    def setUp(self):
+        super().setUp()
+        machine = platform.machine().lower()
+        expected = _frame_pointers_expected(machine)
+        if expected is None:
+            self.skipTest(f"unsupported architecture for frame pointer check: {machine}")
+        try:
+            _testinternalcapi.manual_frame_pointer_unwind()
+        except RuntimeError as exc:
+            if "not supported" in str(exc):
+                self.skipTest("manual frame pointer unwinding not supported on this platform")
+            raise
+        self.machine = machine
+        self.frame_pointers_expected = expected
+
+    def test_manual_unwind_respects_frame_pointers(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 = _manual_unwind_length(**env)
+                jit_frames = result["jit_frames"]
+                python_frames = result.get("python_frames", 0)
+                if self.frame_pointers_expected:
+                    self.assertGreater(
+                        python_frames,
+                        0,
+                        f"expected to find Python frames on {self.machine} with env {env}",
+                    )
+                    if using_jit:
+                        self.assertGreater(
+                            jit_frames,
+                            0,
+                            f"expected to find JIT frames on {self.machine} with env {env}",
+                        )
+                    else:
+                        self.assertEqual(
+                            jit_frames,
+                            0,
+                            f"unexpected JIT frames counted on {self.machine} with env {env}",
+                        )
+                else:
+                    self.assertEqual(
+                        python_frames,
+                        1,
+                        f"unexpected Python frames counted on {self.machine} with env {env}",
+                    )
+                    self.assertEqual(
+                        jit_frames,
+                        0,
+                        f"unexpected JIT frames counted on {self.machine} with env {env}",
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()