Add thread state references to PyThreadState_Ensure() and

PyThreadState_Release()

Per the new updates to PEP-788.

Author: ZeroIntensity

Include/cpython/pystate.h

@@ -213,9 +213,6 @@ struct _ts {
         /* Number of nested PyThreadState_Ensure() calls on this thread state */
         Py_ssize_t counter;
 
-        /* Thread state that was active before PyThreadState_Ensure() was called. */
-        PyThreadState *prior_tstate;
-
         /* Should this thread state be deleted upon calling
            PyThreadState_Release() (with the counter at 1)?
 
@@ -311,6 +308,10 @@ PyAPI_FUNC(void) PyInterpreterWeakRef_Close(PyInterpreterWeakRef wref);
 } while (0)
 
 
-PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterRef interp_ref);
+/* Thread references */
+
+typedef uintptr_t PyThreadRef;
+
+PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterRef interp_ref, PyThreadRef *thread_ref);
 
-PyAPI_FUNC(void) PyThreadState_Release(void);
+PyAPI_FUNC(void) PyThreadState_Release(PyThreadRef thread_ref);

Modules/_testcapimodule.c

@@ -2608,17 +2608,18 @@ test_thread_state_ensure_nested(PyObject *self, PyObject *unused)
     PyInterpreterRef ref = get_strong_ref();
     PyThreadState *save_tstate = PyThreadState_Swap(NULL);
     assert(PyGILState_GetThisThreadState() == save_tstate);
+    PyThreadRef refs[10];
 
     for (int i = 0; i < 10; ++i) {
         // Test reactivation of the detached tstate.
-        if (PyThreadState_Ensure(ref) < 0) {
+        if (PyThreadState_Ensure(ref, &refs[i]) < 0) {
             PyInterpreterRef_Close(ref);
             return PyErr_NoMemory();
         }
 
         // No new thread state should've been created.
         assert(PyThreadState_Get() == save_tstate);
-        PyThreadState_Release();
+        PyThreadState_Release(refs[i]);
     }
 
     assert(PyThreadState_GetUnchecked() == NULL);
@@ -2627,7 +2628,7 @@ test_thread_state_ensure_nested(PyObject *self, PyObject *unused)
     // If the (detached) gilstate matches the interpreter, then it shouldn't
     // create a new thread state.
     for (int i = 0; i < 10; ++i) {
-        if (PyThreadState_Ensure(ref) < 0) {
+        if (PyThreadState_Ensure(ref, &refs[i]) < 0) {
             // This will technically leak other thread states, but it doesn't
             // matter because this is a test.
             PyInterpreterRef_Close(ref);
@@ -2639,7 +2640,7 @@ test_thread_state_ensure_nested(PyObject *self, PyObject *unused)
 
     for (int i = 0; i < 10; ++i) {
         assert(PyThreadState_Get() == save_tstate);
-        PyThreadState_Release();
+        PyThreadState_Release(refs[i]);
     }
 
     assert(PyThreadState_GetUnchecked() == NULL);
@@ -2672,7 +2673,9 @@ test_thread_state_ensure_crossinterp(PyObject *self, PyObject *unused)
        interp = interpreters.create()
        interp.exec(some_func)
        */
-    if (PyThreadState_Ensure(ref) < 0) {
+    PyThreadRef thread_ref;
+    PyThreadRef other_thread_ref;
+    if (PyThreadState_Ensure(ref, &thread_ref) < 0) {
         PyInterpreterRef_Close(ref);
         return PyErr_NoMemory();
     }
@@ -2683,16 +2686,16 @@ test_thread_state_ensure_crossinterp(PyObject *self, PyObject *unused)
     assert(PyGILState_GetThisThreadState() == ensured_tstate);
 
     // Now though, we should reactivate the thread state
-    if (PyThreadState_Ensure(ref) < 0) {
+    if (PyThreadState_Ensure(ref, &other_thread_ref) < 0) {
         PyInterpreterRef_Close(ref);
         return PyErr_NoMemory();
     }
 
     assert(PyThreadState_Get() == ensured_tstate);
-    PyThreadState_Release();
+    PyThreadState_Release(other_thread_ref);
 
     // Ensure that we're restoring the prior thread state
-    PyThreadState_Release();
+    PyThreadState_Release(thread_ref);
     assert(PyThreadState_Get() == interp_tstate);
     assert(PyGILState_GetThisThreadState() == interp_tstate);
 

Programs/_testembed.c

@@ -2332,19 +2332,20 @@ static void
 do_tstate_ensure(void *arg)
 {
     ThreadData *data = (ThreadData *)arg;
-    int res = PyThreadState_Ensure(data->ref);
+    PyThreadRef refs[4];
+    int res = PyThreadState_Ensure(data->ref, &refs[0]);
     assert(res == 0);
-    PyThreadState_Ensure(data->ref);
-    PyThreadState_Ensure(data->ref);
+    PyThreadState_Ensure(data->ref, &refs[1]);
+    PyThreadState_Ensure(data->ref, &refs[2]);
     PyGILState_STATE gstate = PyGILState_Ensure();
-    PyThreadState_Ensure(data->ref);
+    PyThreadState_Ensure(data->ref, &refs[3]);
     res = PyRun_SimpleString(THREAD_CODE);
-    PyThreadState_Release();
+    PyThreadState_Release(refs[3]);
     PyGILState_Release(gstate);
-    PyThreadState_Release();
-    PyThreadState_Release();
+    PyThreadState_Release(refs[2]);
+    PyThreadState_Release(refs[1]);
     assert(res == 0);
-    PyThreadState_Release();
+    PyThreadState_Release(refs[0]);
     PyInterpreterRef_Close(data->ref);
     data->done = 1;
 }

Python/pystate.c

@@ -3319,10 +3319,11 @@ PyInterpreterRef_Main(PyInterpreterRef *strong_ptr)
 }
 
 int
-PyThreadState_Ensure(PyInterpreterRef interp_ref)
+PyThreadState_Ensure(PyInterpreterRef interp_ref, PyThreadRef *thread_ref)
 {
     PyInterpreterState *interp = ref_as_interp(interp_ref);
     PyThreadState *attached_tstate = current_fast_get();
+    *thread_ref = 0;
     if (attached_tstate != NULL && attached_tstate->interp == interp) {
         /* Yay! We already have an attached thread state that matches. */
         ++attached_tstate->ensure.counter;
@@ -3347,7 +3348,7 @@ PyThreadState_Ensure(PyInterpreterRef interp_ref)
     fresh_tstate->ensure.delete_on_release = 1;
 
     if (attached_tstate != NULL) {
-        fresh_tstate->ensure.prior_tstate = PyThreadState_Swap(fresh_tstate);
+        *thread_ref = (PyThreadRef)PyThreadState_Swap(fresh_tstate);
     } else {
         _PyThreadState_Attach(fresh_tstate);
     }
@@ -3356,15 +3357,17 @@ PyThreadState_Ensure(PyInterpreterRef interp_ref)
 }
 
 void
-PyThreadState_Release(void)
+PyThreadState_Release(PyThreadRef thread_ref)
 {
     PyThreadState *tstate = current_fast_get();
     _Py_EnsureTstateNotNULL(tstate);
     Py_ssize_t remaining = --tstate->ensure.counter;
     if (remaining < 0) {
         Py_FatalError("PyThreadState_Release() called more times than PyThreadState_Ensure()");
     }
-    PyThreadState *to_restore = tstate->ensure.prior_tstate;
+    // The thread reference might be NULL
+    assert(thread_ref >= 0);
+    PyThreadState *to_restore = (PyThreadState *)thread_ref;
     if (remaining == 0) {
         if (tstate->ensure.delete_on_release) {
             PyThreadState_Clear(tstate);