Probe all pages when extending stack limits

python · markshannon · Feb 19, 2025 · Feb 10, 2025 · Feb 11, 2025 · Feb 11, 2025
commit 5e5db03e7f72497cccfbe15bf16c16b2d73f15cc
diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
@@ -222,7 +222,7 @@ static inline void _Py_LeaveRecursiveCallTstate(PyThreadState *tstate) {
     (void)tstate;
 }
 
-PyAPI_FUNC(void) _Py_InitializeRecursionCheck(PyThreadState *tstate);
+PyAPI_FUNC(void) _Py_UpdateRecursionLimits(PyThreadState *tstate);
 
 static inline int _Py_ReachedRecursionLimit(PyThreadState *tstate, int margin_count)  {
     char here;
@@ -231,7 +231,7 @@ static inline int _Py_ReachedRecursionLimit(PyThreadState *tstate, int margin_co
         return 0;
     }
     if (tstate->c_stack_hard_limit == 0) {
-        _Py_InitializeRecursionCheck(tstate);
+        _Py_UpdateRecursionLimits(tstate);
     }
     return here_addr <= tstate->c_stack_soft_limit + margin_count * PYOS_STACK_MARGIN_BYTES;
 }
@@ -327,7 +327,7 @@ void _Py_unset_eval_breaker_bit_all(PyInterpreterState *interp, uintptr_t bit);
 
 PyAPI_FUNC(PyObject *) _PyFloat_FromDouble_ConsumeInputs(_PyStackRef left, _PyStackRef right, double value);
 
-extern int _PyOS_CheckStack(int words);
+extern void _Py_StackProbe(uintptr_t from, int bytes, int *probed);
 
 #ifdef __cplusplus
 }

@@ -21,13 +21,18 @@ PyAPI_FUNC(void) PyErr_DisplayException(PyObject *);
 /* Stuff with no proper home (yet) */
 PyAPI_DATA(int) (*PyOS_InputHook)(void);
 
-/* Stack size, in "pointers" (so we get extra safety margins
-   on 64-bit platforms).  On a 32-bit platform, this translates
-   to an 16k margin. */
-#define PYOS_STACK_MARGIN 4096
+/* Stack size, in "pointers". This must be large enough, so 
+ * no two calls to check recursion depth are more than this far 
+ * apart. In practice, that means it must be larger than the C
+ * stack consumption of PyEval_EvalDefault */
+#if defined(Py_DEBUG) && defined(WIN32)
+#  define PYOS_STACK_MARGIN 3072
+#else
+#  define PYOS_STACK_MARGIN 2048
+#endif
 #define PYOS_STACK_MARGIN_BYTES (PYOS_STACK_MARGIN * sizeof(void *))
 
-#if defined(WIN32) && !defined(MS_WIN64) && !defined(_M_ARM) && defined(_MSC_VER) && _MSC_VER >= 1300
+#if defined(WIN32)
 /* Enable stack checking under Microsoft C */
 // When changing the platforms, ensure PyOS_CheckStack() docs are still correct
 #define USE_STACKCHECK

@@ -322,10 +322,8 @@
 
 #if defined(__s390x__)
 #  define Py_C_STACK_SIZE 320000
-#elif defined(_WIN32) && defined(_M_ARM64)
-#  define Py_C_STACK_SIZE 400000
 #elif defined(_WIN32)
-#  define Py_C_STACK_SIZE 1200000
+   // Don't define Py_C_STACK_SIZE, use probing instead
 #elif defined(__ANDROID__)
 #  define Py_C_STACK_SIZE 1200000
 #elif defined(__sparc__)
@@ -342,32 +340,30 @@
 #endif
 
 void
-_Py_InitializeRecursionCheck(PyThreadState *tstate)
+_Py_UpdateRecursionLimits(PyThreadState *tstate)
 {
     char here;
     uintptr_t here_addr = (uintptr_t)&here;
-    tstate->c_stack_top = here_addr;
+    int to_probe = PYOS_STACK_MARGIN_BYTES * 2;
 #ifdef USE_STACKCHECK
-    if (_PyOS_CheckStack(PYOS_STACK_MARGIN * 2) == 0) {
-        tstate->c_stack_soft_limit = here_addr - PYOS_STACK_MARGIN_BYTES;
-        return;
-    }
-    int margin = PYOS_STACK_MARGIN;
-    assert(tstate->c_stack_soft_limit != UINTPTR_MAX);
-    if (_PyOS_CheckStack(margin)) {
-        margin = PYOS_STACK_MARGIN/2;
-    }
-    else {
-        if (_PyOS_CheckStack(PYOS_STACK_MARGIN*3/2) == 0) {
-            margin = PYOS_STACK_MARGIN*3/2;
-        }
+    if (tstate->c_stack_top == 0) {
+        assert(tstate->c_stack_soft_limit == UINTPTR_MAX);
+        tstate->c_stack_top = _Py_SIZE_ROUND_UP(here_addr, 4096);
+        tstate->c_stack_soft_limit = tstate->c_stack_top;
+        to_probe = PYOS_STACK_MARGIN_BYTES * 4;
+    }
+    int depth;
+    uintptr_t implicit_hard_limit = tstate->c_stack_soft_limit - PYOS_STACK_MARGIN_BYTES;
+    _Py_StackProbe(implicit_hard_limit, to_probe, &depth);
+    tstate->c_stack_soft_limit = implicit_hard_limit - depth + PYOS_STACK_MARGIN_BYTES * 2; 
+    if (depth != to_probe) {
+        tstate->c_stack_hard_limit = tstate->c_stack_soft_limit - PYOS_STACK_MARGIN_BYTES;
     }
-    tstate->c_stack_hard_limit = here_addr - margin * sizeof(void *);
-    tstate->c_stack_soft_limit = tstate->c_stack_hard_limit + PYOS_STACK_MARGIN_BYTES;
 #else
-    assert(tstate->c_stack_soft_limit == UINTPTR_MAX);
-    tstate->c_stack_soft_limit = here_addr - Py_C_STACK_SIZE;
-    tstate->c_stack_hard_limit = here_addr - (Py_C_STACK_SIZE + PYOS_STACK_MARGIN_BYTES);
+    assert(tstate->c_stack_top == 0);
+    tstate->c_stack_top = _Py_SIZE_ROUND_UP(here_addr, 4096);
+    tstate->c_stack_soft_limit = tstate->c_stack_top - Py_C_STACK_SIZE;
+    tstate->c_stack_hard_limit = tstate->c_stack_top - (Py_C_STACK_SIZE + PYOS_STACK_MARGIN_BYTES);
 #endif
 }
 
@@ -380,15 +376,15 @@
     uintptr_t here_addr = (uintptr_t)&here;
     assert(tstate->c_stack_soft_limit != 0);
     if (tstate->c_stack_hard_limit == 0) {
-        _Py_InitializeRecursionCheck(tstate);
+        _Py_UpdateRecursionLimits(tstate);
     }
     if (here_addr >= tstate->c_stack_soft_limit) {
         return 0;
     }
     assert(tstate->c_stack_hard_limit != 0);
     if (here_addr < tstate->c_stack_hard_limit) {
         /* Overflowing while handling an overflow. Give up. */
-        int kbytes_used = (tstate->c_stack_top - here_addr)/1024;
+        int kbytes_used = (int)(tstate->c_stack_top - here_addr)/1024;
         char buffer[80];
         snprintf(buffer, 80, "Unrecoverable stack overflow (used %d kB)%s", kbytes_used, where);
         Py_FatalError(buffer);
@@ -397,7 +393,7 @@
         return 0;
     }
     else {
-        int kbytes_used = (tstate->c_stack_top - here_addr)/1024;
+        int kbytes_used = (int)(tstate->c_stack_top - here_addr)/1024;
         tstate->recursion_headroom++;
         _PyErr_Format(tstate, PyExc_RecursionError,
                     "Stack overflow (used %d kB)%s",

@@ -1493,6 +1493,7 @@ init_threadstate(_PyThreadStateImpl *_tstate,
     tstate->py_recursion_limit = interp->ceval.recursion_limit;
     tstate->py_recursion_remaining = interp->ceval.recursion_limit;
     tstate->c_stack_soft_limit = UINTPTR_MAX;
+    tstate->c_stack_top = 0;
     tstate->c_stack_hard_limit = 0;
 
     tstate->exc_info = &tstate->exc_state;

@@ -1552,14 +1552,24 @@ _Py_SourceAsString(PyObject *cmd, const char *funcname, const char *what, PyComp
 #include <malloc.h>
 #include <excpt.h>
 
-int _PyOS_CheckStack(int words)
+void _Py_StackProbe(uintptr_t from, int bytes, int *probed)
 {
+    assert((bytes & 4095) == 0);
+    int depth = 4096;
+    char here;
+    uintptr_t here_addr = (uintptr_t)&here;
     __try
     {
-        /* alloca throws a stack overflow exception if there's
-           not enough space left on the stack */
-        alloca(words * sizeof(void *));
-        return 0;
+        while (depth <= bytes) {
+            uintptr_t probe_point = from - depth + 64;
+            if (probe_point < here_addr) {
+                /* alloca throws a stack overflow exception if there's
+                 * not enough space left on the stack */
+                alloca(here_addr - probe_point);
+            }
+            *probed = depth;
+            depth += 4096;
+        }
     }
     __except (GetExceptionCode() == STATUS_STACK_OVERFLOW ? EXCEPTION_EXECUTE_HANDLER : EXCEPTION_CONTINUE_SEARCH)
     {
@@ -1569,15 +1579,25 @@ int _PyOS_CheckStack(int words)
             Py_FatalError("Could not reset the stack!");
         }
     }
-    return 1;
 }
 
 /*
  * Return non-zero when we run out of memory on the stack; zero otherwise.
  */
 int PyOS_CheckStack(void)
 {
-    return _PyOS_CheckStack(PYOS_STACK_MARGIN);
+    char here;
+    uintptr_t here_addr = (uintptr_t)&here;
+    uintptr_t from = _Py_SIZE_ROUND_UP(here_addr, 4096);
+    int depth;
+    _Py_StackProbe(from, PYOS_STACK_MARGIN_BYTES*2, &depth);
+    assert(depth <= PYOS_STACK_MARGIN_BYTES*2);
+    if (depth == PYOS_STACK_MARGIN_BYTES*2) {
+        return 0;
+    }
+    else {
+        return -1;
+    }
 }
 
 #endif /* WIN32 && _MSC_VER */