Check C stack depth before stack allocating JIT optimizer struct

Author: markshannon

Lib/test/test_call.py

@@ -1075,6 +1075,7 @@ def c_py_recurse(m):
                 c_py_recurse(100_000)
 
 
+
 class TestFunctionWithManyArgs(unittest.TestCase):
     def test_function_with_many_args(self):
         for N in (10, 500, 1000):

Lib/test/test_capi/test_opt.py

@@ -10,11 +10,12 @@
 
 from test.support import (script_helper, requires_specialization,
                           import_helper, Py_GIL_DISABLED, requires_jit_enabled,
-                          reset_code)
+                          reset_code, infinite_recursion)
 
 _testinternalcapi = import_helper.import_module("_testinternalcapi")
 
 from _testinternalcapi import TIER2_THRESHOLD
+from _testcapi import pyobject_vectorcall
 
 
 @contextlib.contextmanager
@@ -2502,6 +2503,37 @@ def testfunc(n):
         self.assertIn("_BINARY_OP", uops)
 
 
+@requires_jit_enabled
+class TestStackUse(unittest.TestCase):
+
+    def test_jit_at_depth(self):
+
+        def run_at_depth(n, func):
+            if n:
+                pyobject_vectorcall(run_at_depth, (n-1, func), None)
+            else:
+                func()
+
+        def find_stack_limit():
+            nonlocal depth
+            depth += 1
+            pyobject_vectorcall(find_stack_limit, (), None)
+
+        with infinite_recursion():
+            depth = 0
+            try:
+                find_stack_limit()
+            except RecursionError:
+                pass
+
+            def loop_func():
+                r = 0
+                for i in range(TIER2_THRESHOLD):
+                    r += i
+                return r
+            run_at_depth(depth-1, loop_func)
+
+
 def global_identity(x):
     return x
 

Python/optimizer_analysis.c

@@ -695,6 +695,11 @@ _Py_uop_analyze_and_optimize(
 )
 {
     OPT_STAT_INC(optimizer_attempts);
+    /* Make sure we have enough C stack space for the optimizer */
+    int margin = 1 + sizeof(JitOptContext)/_PyOS_STACK_MARGIN_BYTES;
+    if (_Py_ReachedRecursionLimitWithMargin(_PyThreadState_GET(), margin)) {
+        return 0;
+    }
 
     int err = remove_globals(frame, buffer, length, dependencies);
     if (err <= 0) {