gh-139653: Add PyUnstable_ThreadState_SetStack()

Add PyUnstable_ThreadState_SetStack() and
PyUnstable_ThreadState_ResetStack() functions to set the stack base
address and stack size of a Python thread state.

Author: vstinner

Doc/c-api/init.rst

@@ -1353,6 +1353,32 @@ All of the following functions must be called after :c:func:`Py_Initialize`.
    .. versionadded:: 3.11
 
 
+.. c:function:: int PyUnstable_ThreadState_SetStack(PyThreadState *tstate, void *stack_start_addr, size_t stack_size)
+
+   Set the stack start address and stack size of a Python thread state.
+
+   *stack_size* must be greater than ``0``.
+
+   On success, return ``0``.
+   On failure, set an exception and return ``-1``.
+
+   .. seealso::
+      The :c:func:`PyUnstable_ThreadState_ResetStack` function.
+
+   .. versionadded:: next
+
+
+.. c:function:: void PyUnstable_ThreadState_ResetStack(PyThreadState *tstate)
+
+   Reset the stack start address and stack size of a Python thread state to
+   the operating system defaults.
+
+   .. seealso::
+      The :c:func:`PyUnstable_ThreadState_SetStack` function.
+
+   .. versionadded:: next
+
+
 .. c:function:: PyInterpreterState* PyInterpreterState_Get(void)
 
    Get the current interpreter.

Doc/whatsnew/3.14.rst

@@ -2994,6 +2994,11 @@ New features in the C API
   as arguments to C API functions.
   (Contributed by Sam Gross in :gh:`133164`.)
 
+* Add :c:func:`PyUnstable_ThreadState_SetStack` and
+  :c:func:`PyUnstable_ThreadState_ResetStack` functions to set the stack base
+  address and stack size of a Python thread state.
+  (Contributed by Victor Stinner in :gh:`139653`.)
+
 
 Limited C API changes
 ---------------------

Include/cpython/pystate.h

@@ -252,6 +252,16 @@ PyAPI_FUNC(int) PyGILState_Check(void);
 */
 PyAPI_FUNC(PyObject*) _PyThread_CurrentFrames(void);
 
+// Set the stack start address and stack size of a Python thread state
+PyAPI_FUNC(int) PyUnstable_ThreadState_SetStack(
+    PyThreadState *tstate,
+    void *stack_start_addr,  // Stack start address
+    size_t stack_size);      // Stack size (in bytes)
+
+// Reset the stack start address and stack size of a Python thread state
+PyAPI_FUNC(void) PyUnstable_ThreadState_ResetStack(
+    PyThreadState *tstate);
+
 /* Routines for advanced debuggers, requested by David Beazley.
    Don't use unless you know what you are doing! */
 PyAPI_FUNC(PyInterpreterState *) PyInterpreterState_Main(void);

Include/internal/pycore_tstate.h

@@ -37,6 +37,10 @@ typedef struct _PyThreadStateImpl {
     uintptr_t c_stack_soft_limit;
     uintptr_t c_stack_hard_limit;
 
+    // PyUnstable_ThreadState_ResetStack() values
+    uintptr_t c_stack_init_base;
+    uintptr_t c_stack_init_top;
+
     PyObject *asyncio_running_loop; // Strong reference
     PyObject *asyncio_running_task; // Strong reference
 

Misc/NEWS.d/next/C_API/2025-10-06-22-17-47.gh-issue-139653.6-1MOd.rst

@@ -0,0 +1,3 @@
+Add :c:func:`PyUnstable_ThreadState_SetStack` and
+:c:func:`PyUnstable_ThreadState_ResetStack` functions to set the stack base
+address and stack size of a Python thread state. Patch by Victor Stinner.

Modules/_testinternalcapi.c

@@ -2418,6 +2418,55 @@ set_vectorcall_nop(PyObject *self, PyObject *func)
     Py_RETURN_NONE;
 }
 
+static void
+check_threadstate_set_stack(PyThreadState *tstate, void *start, size_t size)
+{
+    assert(PyUnstable_ThreadState_SetStack(tstate, start, size) == 0);
+    assert(!PyErr_Occurred());
+
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    assert(ts->c_stack_hard_limit == (uintptr_t)start + _PyOS_STACK_MARGIN_BYTES);
+    assert(ts->c_stack_top == (uintptr_t)start + size);
+    assert(ts->c_stack_soft_limit >= ts->c_stack_hard_limit);
+    assert(ts->c_stack_soft_limit < ts->c_stack_top);
+}
+
+
+static PyObject *
+test_threadstate_set_stack(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    assert(!PyErr_Occurred());
+
+    uintptr_t init_base = ts->c_stack_init_base;
+    size_t init_top = ts->c_stack_init_top;
+
+    // Test the minimum stack size
+    size_t size = _PyOS_STACK_MARGIN_BYTES * 3;
+    void *start = (void*)(_Py_get_machine_stack_pointer() - size);
+    check_threadstate_set_stack(tstate, start, size);
+
+    // Test a larger size
+    size = 7654321;
+    start = (void*)(_Py_get_machine_stack_pointer() - size);
+    check_threadstate_set_stack(tstate, start, size);
+
+    // Test invalid size (too small)
+    size = 5;
+    start = (void*)(_Py_get_machine_stack_pointer() - size);
+    assert(PyUnstable_ThreadState_SetStack(tstate, start, size) == -1);
+    assert(PyErr_ExceptionMatches(PyExc_ValueError));
+    PyErr_Clear();
+
+    // Test PyUnstable_ThreadState_ResetStack()
+    PyUnstable_ThreadState_ResetStack(tstate);
+    assert(ts->c_stack_init_base == init_base);
+    assert(ts->c_stack_init_top == init_top);
+
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef module_functions[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_recursion_depth", get_recursion_depth, METH_NOARGS},
@@ -2527,6 +2576,7 @@ static PyMethodDef module_functions[] = {
 #endif
     {"simple_pending_call", simple_pending_call, METH_O},
     {"set_vectorcall_nop", set_vectorcall_nop, METH_O},
+    {"test_threadstate_set_stack", test_threadstate_set_stack, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Python/ceval.c

@@ -439,7 +439,7 @@ int pthread_attr_destroy(pthread_attr_t *a)
 #endif
 
 static void
-hardware_stack_limits(uintptr_t *top, uintptr_t *base)
+hardware_stack_limits(uintptr_t *base, uintptr_t *top)
 {
 #ifdef WIN32
     ULONG_PTR low, high;
@@ -482,23 +482,86 @@ hardware_stack_limits(uintptr_t *top, uintptr_t *base)
 #endif
 }
 
-void
-_Py_InitializeRecursionLimits(PyThreadState *tstate)
+static void
+tstate_set_stack(PyThreadState *tstate,
+                 uintptr_t base, uintptr_t top)
 {
-    uintptr_t top;
-    uintptr_t base;
-    hardware_stack_limits(&top, &base);
+    assert(base < top);
+    assert((top - base) >= (_PyOS_STACK_MARGIN_BYTES * 3));
+
 #ifdef _Py_THREAD_SANITIZER
     // Thread sanitizer crashes if we use more than half the stack.
     uintptr_t stacksize = top - base;
-    base += stacksize/2;
+    base += stacksize / 2;
 #endif
     _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
     _tstate->c_stack_top = top;
     _tstate->c_stack_hard_limit = base + _PyOS_STACK_MARGIN_BYTES;
     _tstate->c_stack_soft_limit = base + _PyOS_STACK_MARGIN_BYTES * 2;
+
+#ifndef NDEBUG
+    // Sanity checks
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    assert(ts->c_stack_hard_limit <= ts->c_stack_soft_limit);
+    assert(ts->c_stack_soft_limit < ts->c_stack_top);
+#endif
+}
+
+
+void
+_Py_InitializeRecursionLimits(PyThreadState *tstate)
+{
+    uintptr_t base, top;
+    hardware_stack_limits(&base, &top);
+    assert(top != 0);
+
+    tstate_set_stack(tstate, base, top);
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    ts->c_stack_init_base = base;
+    ts->c_stack_init_top = top;
+
+    // Test the stack pointer
+#if !defined(NDEBUG) && !defined(__wasi__)
+    uintptr_t here_addr = _Py_get_machine_stack_pointer();
+    assert(ts->c_stack_soft_limit < here_addr);
+    assert(here_addr < ts->c_stack_top);
+#endif
+}
+
+
+int
+PyUnstable_ThreadState_SetStack(PyThreadState *tstate,
+                                void *stack_start_addr, size_t stack_size)
+{
+    if (stack_size < (_PyOS_STACK_MARGIN_BYTES * 3)) {
+        PyErr_Format(PyExc_ValueError,
+                     "stack_size must be at least %zu bytes",
+                     _PyOS_STACK_MARGIN_BYTES * 3);
+        return -1;
+    }
+
+    uintptr_t base = (uintptr_t)stack_start_addr;
+    uintptr_t top = base + stack_size;
+    tstate_set_stack(tstate, base, top);
+    return 0;
 }
 
+
+void
+PyUnstable_ThreadState_ResetStack(PyThreadState *tstate)
+{
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    if (ts->c_stack_init_top != 0) {
+        tstate_set_stack(tstate,
+                         ts->c_stack_init_base,
+                         ts->c_stack_init_top);
+        return;
+    }
+
+    _Py_InitializeRecursionLimits(tstate);
+}
+
+
 /* The function _Py_EnterRecursiveCallTstate() only calls _Py_CheckRecursiveCall()
    if the recursion_depth reaches recursion_limit. */
 int

Python/pystate.c

@@ -1490,6 +1490,9 @@ init_threadstate(_PyThreadStateImpl *_tstate,
     _tstate->c_stack_top = 0;
     _tstate->c_stack_hard_limit = 0;
 
+    _tstate->c_stack_init_base = 0;
+    _tstate->c_stack_init_top = 0;
+
     _tstate->asyncio_running_loop = NULL;
     _tstate->asyncio_running_task = NULL;