Fix PyUnstable_ThreadState_ResetStack()

Add c_stack_init_start and c_stack_init_size to _PyThreadStateImpl to
cache PyUnstable_ThreadState_ResetStack() values. Set these members
at the first call, and then reuse these values.

Author: vstinner

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_start;
+    size_t c_stack_init_size;
+
     PyObject *asyncio_running_loop; // Strong reference
     PyObject *asyncio_running_task; // Strong reference
 

Modules/_testinternalcapi.c

@@ -2436,8 +2436,12 @@ 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_start = ts->c_stack_init_start;
+    size_t init_size = ts->c_stack_init_size;
+
     // Test the minimum stack size
     size_t size = _PyOS_STACK_MARGIN_BYTES * 3;
     void *start = (void*)(_Py_get_machine_stack_pointer() - size);
@@ -2455,7 +2459,11 @@ test_threadstate_set_stack(PyObject *self, PyObject *Py_UNUSED(args))
     assert(PyErr_ExceptionMatches(PyExc_ValueError));
     PyErr_Clear();
 
+    // Test PyUnstable_ThreadState_ResetStack()
     PyUnstable_ThreadState_ResetStack(tstate);
+    assert(ts->c_stack_init_start == init_start);
+    assert(ts->c_stack_init_size == init_size);
+
     Py_RETURN_NONE;
 }
 

Python/ceval.c

@@ -489,8 +489,9 @@ PyUnstable_ThreadState_SetStack(PyThreadState *tstate,
 }
 
 
-void
-PyUnstable_ThreadState_ResetStack(PyThreadState *tstate)
+// Get the stack start address and stack size (in bytes)
+static void
+get_stack(uintptr_t *start, size_t *size)
 {
 #ifdef WIN32
     ULONG_PTR low, high;
@@ -499,15 +500,14 @@ PyUnstable_ThreadState_ResetStack(PyThreadState *tstate)
     ULONG guarantee = 0;
     SetThreadStackGuarantee(&guarantee);
 
-    uintptr_t start = (uintptr_t)low + guarantee;
-    size_t size = (uintptr_t)high - start;
-    tstate_set_stack(tstate, (void*)start, size);
+    *start = (uintptr_t)low + guarantee;
+    *size = (uintptr_t)high - start;
 
 #elif defined(__APPLE__)
     pthread_t this_thread = pthread_self();
     void *top = pthread_get_stackaddr_np(this_thread); // top of the stack
-    size_t size = pthread_get_stacksize_np(this_thread);
-    tstate_set_stack(tstate, (char*)top - size, size);
+    *size = pthread_get_stacksize_np(this_thread);
+    *start = (uintptr_t)top - *size;
 
 #else
    // XXX musl supports HAVE_PTHRED_GETATTR_NP, but the resulting stack size
@@ -525,23 +525,40 @@ PyUnstable_ThreadState_ResetStack(PyThreadState *tstate)
         err |= pthread_attr_destroy(&attr);
     }
     if (err == 0) {
-        uintptr_t base = ((uintptr_t)stack_addr) + guard_size;
-        uintptr_t start = base;
-        size_t pystack_size = (base + stack_size) - start;
-        tstate_set_stack(tstate, (void*)start, pystack_size);
+        uintptr_t base = (uintptr_t)stack_addr + guard_size;
+        *start = base;
+        *size = (base + stack_size) - *start;
     }
     else
 #  endif
     {
         uintptr_t here_addr = _Py_get_machine_stack_pointer();
         uintptr_t top = _Py_SIZE_ROUND_UP(here_addr, 4096);
-        uintptr_t start = top - Py_C_STACK_SIZE;
-        size_t pystack_size = top - start;
-        tstate_set_stack(tstate, (void*)start, pystack_size);
+        *size = Py_C_STACK_SIZE;
+        *start = top - *size;
     }
 #endif
 }
 
+void
+PyUnstable_ThreadState_ResetStack(PyThreadState *tstate)
+{
+    _PyThreadStateImpl *ts = (_PyThreadStateImpl *)tstate;
+    if (ts->c_stack_init_start != 0) {
+        tstate_set_stack(tstate,
+                         (void*)ts->c_stack_init_start,
+                         ts->c_stack_init_size);
+        return;
+    }
+
+    uintptr_t start;
+    size_t size;
+    get_stack(&start, &size);
+    tstate_set_stack(tstate, (void*)start, size);
+    ts->c_stack_init_start = start;
+    ts->c_stack_init_size = size;
+}
+
 
 /* The function _Py_EnterRecursiveCallTstate() only calls _Py_CheckRecursiveCall()
    if the recursion_depth reaches recursion_limit. */

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_start = 0;
+    _tstate->c_stack_init_size = 0;
+
     _tstate->asyncio_running_loop = NULL;
     _tstate->asyncio_running_task = NULL;