gh-124622: Add PyThreadState_Ensure() function

Doc/c-api/init.rst

@@ -1279,6 +1279,7 @@ with sub-interpreters:
       Hangs the current thread, rather than terminating it, if called while the
       interpreter is finalizing.
 
+
 .. c:function:: void PyGILState_Release(PyGILState_STATE)
 
    Release any resources previously acquired.  After this call, Python's state will
@@ -1499,6 +1500,37 @@ All of the following functions must be called after :c:func:`Py_Initialize`.
    .. versionadded:: 3.8
 
 
+.. c:function:: int PyThreadState_Ensure(PyInterpreterState *interp, const char **errmsg)
+
+   Similar to :c:func:`PyGILState_Ensure`, except that it returns with a status
+   code even in the case of failure, and takes an interpreter state.
+   Specifically, it returns a status code (``>= 0``) when the operation
+   succeeded, or sets *\*errmsg* (if *errmsg* is not NULL) and returns ``-1``
+   on failure.
+
+   On success, the thread state must be released by
+   :c:func:`PyThreadState_Release`.
+
+   In the case of failure, it is *unsafe* to use the Python API following the
+   call. Releasing the obtained *state* via :c:func:`PyGILState_Release` must
+   only be done in the case of success.
+
+   .. versionadded:: next
+
+
+.. c:function:: void PyThreadState_Release(int state)
+
+   Release any resources previously acquired.  After this call, Python's state
+   will be the same as it was prior to the corresponding
+   :c:func:`PyThreadState_Ensure` call (but generally this state will be
+   unknown to the caller).
+
+   Every call to :c:func:`PyThreadState_Ensure` must be matched by a call to
+   :c:func:`PyThreadState_Release` on the same thread.
+
+   .. versionadded:: next
+
+
 .. c:function:: PyObject* PyUnstable_InterpreterState_GetMainModule(PyInterpreterState *interp)
 
    Return a :term:`strong reference` to the ``__main__`` :ref:`module object <moduleobjects>`

Doc/whatsnew/3.14.rst

@@ -1372,6 +1372,11 @@ New features
   and get an attribute of the module.
   (Contributed by Victor Stinner in :gh:`128911`.)
 
+* Add :c:func:`PyThreadState_Ensure` and :c:func:`PyThreadState_Release`
+  functions: similar to :c:func:`PyGILState_Ensure` and
+  :c:func:`PyGILState_Release`, but :c:func:`PyThreadState_Ensure` returns
+  ``-1`` on failure. Patch by Victor Stinner.
+  (Contributed by Victor Stinner in :gh:`124622`.)
 
 Limited C API changes
 ---------------------

Include/cpython/pystate.h

@@ -140,6 +140,8 @@ struct _ts {
     PyObject *dict;  /* Stores per-thread state */
 
     int gilstate_counter;
+    // PyThreadState_Ensure() call depth
+    int ensure_depth;
 
     PyObject *async_exc; /* Asynchronous exception to raise */
     unsigned long thread_id; /* Thread id where this tstate was created */

Include/internal/pycore_ceval.h

@@ -327,6 +327,12 @@ 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 _PyEval_AcquireLockOrFail(
+    PyThreadState *tstate,
+    const char **errmsg);
+extern int _PyEval_RestoreThreadOrFail(
+    PyThreadState *tstate,
+    const char **errmsg);
 
 #ifdef __cplusplus
 }

Include/internal/pycore_pystate.h

@@ -313,6 +313,10 @@ _Py_AssertHoldsTstateFunc(const char *func)
 #define _Py_AssertHoldsTstate()
 #endif
 
+extern int _PyThreadState_AttachOrFail(
+    PyThreadState *tstate,
+    const char **errmsg);
+
 #ifdef __cplusplus
 }
 #endif

Include/pystate.h

@@ -111,6 +111,14 @@ PyAPI_FUNC(PyGILState_STATE) PyGILState_Ensure(void);
 */
 PyAPI_FUNC(void) PyGILState_Release(PyGILState_STATE);
 
+#if !defined(Py_LIMITED_API) || Py_LIMITED_API+0 >= 0x030e0000
+/* New in 3.14 */
+PyAPI_FUNC(int) PyThreadState_Ensure(
+    PyInterpreterState *interp,
+    const char **errmsg);
+PyAPI_FUNC(void) PyThreadState_Release(int state);
+#endif
+
 /* Helper/diagnostic function - get the current thread state for
    this thread.  May return NULL if no GILState API has been used
    on the current thread.  Note that the main thread always has such a

Misc/NEWS.d/next/C_API/2025-02-05-15-01-23.gh-issue-124622.jnROuN.rst

@@ -0,0 +1,4 @@
+Add :c:func:`PyThreadState_Ensure` and :c:func:`PyThreadState_Release`
+functions: similar to :c:func:`PyGILState_Ensure` and
+:c:func:`PyGILState_Release`, but :c:func:`PyThreadState_Ensure` returns ``-1``
+on failure. Patch by Victor Stinner.

Misc/stable_abi.toml

@@ -2545,3 +2545,7 @@
     added = '3.14'
 [function.Py_PACK_VERSION]
     added = '3.14'
+[function.PyThreadState_Ensure]
+    added = '3.14'
+[function.PyThreadState_Release]
+    added = '3.14'

Modules/_testcapimodule.c

@@ -1263,19 +1263,24 @@ typedef struct {
     PyThread_type_lock start_event;
     PyThread_type_lock exit_event;
     PyObject *callback;
+    PyInterpreterState *interp;
 } test_c_thread_t;
 
 static void
 temporary_c_thread(void *data)
 {
     test_c_thread_t *test_c_thread = data;
-    PyGILState_STATE state;
     PyObject *res;
 
     PyThread_release_lock(test_c_thread->start_event);
 
     /* Allocate a Python thread state for this thread */
-    state = PyGILState_Ensure();
+    const char *errmsg;
+    int state = PyThreadState_Ensure(test_c_thread->interp, &errmsg);
+    if (state < 0) {
+        fprintf(stderr, "ERROR: PyThreadState_Ensure() failed: %s", errmsg);
+        abort();
+    }
 
     res = PyObject_CallNoArgs(test_c_thread->callback);
     Py_CLEAR(test_c_thread->callback);
@@ -1288,7 +1293,7 @@ temporary_c_thread(void *data)
     }
 
     /* Destroy the Python thread state for this thread */
-    PyGILState_Release(state);
+    PyThreadState_Release(state);
 
     PyThread_release_lock(test_c_thread->exit_event);
 }
@@ -1310,6 +1315,7 @@ call_in_temporary_c_thread(PyObject *self, PyObject *args)
     test_c_thread.start_event = PyThread_allocate_lock();
     test_c_thread.exit_event = PyThread_allocate_lock();
     test_c_thread.callback = NULL;
+    test_c_thread.interp = PyInterpreterState_Get();
     if (!test_c_thread.start_event || !test_c_thread.exit_event) {
         PyErr_SetString(PyExc_RuntimeError, "could not allocate lock");
         goto exit;
@@ -1370,6 +1376,106 @@ join_temporary_c_thread(PyObject *self, PyObject *Py_UNUSED(ignored))
     Py_RETURN_NONE;
 }
 
+static PyObject*
+_test_tstate_ensure(int subinterpreter, int clear_current_before)
+{
+    PyThreadState *oldts = PyThreadState_Get();
+
+    PyInterpreterState *interp;
+    PyThreadState *subinterp_tstate = NULL;
+    PyThreadState *expected_tstate = NULL;
+    if (subinterpreter) {
+        // Create a sub-interpreter
+        subinterp_tstate = Py_NewInterpreter();
+        assert(PyThreadState_Get() == subinterp_tstate);
+        interp = PyThreadState_GetInterpreter(subinterp_tstate);
+        expected_tstate = subinterp_tstate;
+    }
+    else {
+        interp = PyThreadState_GetInterpreter(oldts);
+        expected_tstate = oldts;
+    }
+
+    if (clear_current_before) {
+        PyThreadState_Swap(NULL);
+        assert(PyThreadState_GetUnchecked() == NULL);
+    }
+
+    // First call
+    const char *errmsg;
+    int state1 = PyThreadState_Ensure(interp, &errmsg);
+    if (state1 < 0) {
+        fprintf(stderr, "ERROR: PyThreadState_Ensure() failed: %s", errmsg);
+        abort();
+    }
+    PyThreadState *ensure1 = PyThreadState_GetUnchecked();
+    if (clear_current_before) {
+        assert(ensure1 != expected_tstate);
+    }
+    else {
+        assert(ensure1 == expected_tstate);
+    }
+
+    {
+        // Second call
+        int state2 = PyThreadState_Ensure(interp, &errmsg);
+        if (state2 < 0) {
+            fprintf(stderr, "ERROR: PyThreadState_Ensure() failed: %s", errmsg);
+            abort();
+        }
+        PyThreadState *ensure2 = PyThreadState_GetUnchecked();
+        assert(ensure2 == ensure1);
+
+        PyThreadState_Release(state2);
+    }
+    PyThreadState_Release(state1);
+
+    if (!clear_current_before) {
+        assert(PyThreadState_GetUnchecked() == expected_tstate);
+    }
+    else {
+        assert(PyThreadState_GetUnchecked() == NULL);
+    }
+
+    if (subinterpreter) {
+        PyThreadState_Swap(subinterp_tstate);
+        Py_EndInterpreter(subinterp_tstate);
+        assert(PyThreadState_GetUnchecked() == NULL);
+    }
+    PyThreadState_Swap(oldts);
+
+    Py_RETURN_NONE;
+}
+
+
+static PyObject *
+test_tstate_ensure(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    return _test_tstate_ensure(0, 0);
+}
+
+
+static PyObject *
+test_tstate_ensure_clear(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    return _test_tstate_ensure(0, 1);
+}
+
+
+static PyObject *
+test_tstate_ensure_subinterp(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    return _test_tstate_ensure(1, 0);
+}
+
+
+static PyObject *
+test_tstate_ensure_subinterp_clear(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    return _test_tstate_ensure(1, 1);
+}
+
+
 /* marshal */
 
 static PyObject*
@@ -2537,6 +2643,10 @@ static PyMethodDef TestMethods[] = {
     {"call_in_temporary_c_thread", call_in_temporary_c_thread, METH_VARARGS,
      PyDoc_STR("set_error_class(error_class) -> None")},
     {"join_temporary_c_thread", join_temporary_c_thread, METH_NOARGS},
+    {"test_tstate_ensure", test_tstate_ensure, METH_NOARGS},
+    {"test_tstate_ensure_clear", test_tstate_ensure_clear, METH_NOARGS},
+    {"test_tstate_ensure_subinterp", test_tstate_ensure_subinterp, METH_NOARGS},
+    {"test_tstate_ensure_subinterp_clear", test_tstate_ensure_subinterp_clear, METH_NOARGS},
     {"pymarshal_write_long_to_file",
         pymarshal_write_long_to_file, METH_VARARGS},
     {"pymarshal_write_object_to_file",