Add a test for PyThreadState_Ensure() across interpreters.

Author: ZeroIntensity

Include/cpython/pystate.h

@@ -288,8 +288,9 @@ PyAPI_FUNC(PyInterpreterState *) PyInterpreterState_Lookup(int64_t interp_id);
 /* Release a reference to an interpreter incremented by PyInterpreterState_Hold() */
 PyAPI_FUNC(void) PyInterpreterState_Release(PyInterpreterState *interp);
 
-// Export for '_testcapi' shared extension
+// Exports for '_testcapi' shared extension
 PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);
+PyAPI_FUNC(void) _PyInterpreterState_Incref(PyInterpreterState *interp);
 
 PyAPI_FUNC(int) PyThreadState_SetDaemon(int daemon);
 

Modules/_testcapimodule.c

@@ -2591,6 +2591,36 @@ test_interp_lookup(PyObject *self, PyObject *unused)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_interp_ensure(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    PyThreadState *tstate = Py_NewInterpreter();
+    PyInterpreterState *subinterp = PyThreadState_GetInterpreter(tstate);
+
+    for (int i = 0; i < 10; ++i) {
+        _PyInterpreterState_Incref(interp);
+        int res = PyThreadState_Ensure(interp);
+        assert(res == 0);
+        assert(PyInterpreterState_Get() == interp);
+    }
+
+    for (int i = 0; i < 10; ++i) {
+        _PyInterpreterState_Incref(subinterp);
+        int res = PyThreadState_Ensure(subinterp);
+        assert(res == 0);
+        assert(PyInterpreterState_Get() == subinterp);
+    }
+
+    for (int i = 0; i < 20; ++i) {
+        PyThreadState_Release();
+    }
+
+    PyThreadState_Swap(save_tstate);
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -2687,6 +2717,7 @@ static PyMethodDef TestMethods[] = {
     {"toggle_reftrace_printer", toggle_reftrace_printer, METH_O},
     {"test_interp_refcount", test_interp_refcount, METH_NOARGS},
     {"test_interp_lookup", test_interp_lookup, METH_NOARGS},
+    {"test_interp_ensure", test_interp_ensure, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Python/pystate.c

@@ -3230,12 +3230,19 @@ _PyInterpreterState_Refcount(PyInterpreterState *interp)
     return refcount;
 }
 
+void
+_PyInterpreterState_Incref(PyInterpreterState *interp)
+{
+    assert(interp != NULL);
+    assert(_Py_atomic_load_ssize_relaxed(&interp->threads.finalizing.countdown) >= 0);
+    _Py_atomic_add_ssize(&interp->threads.finalizing.countdown, 1);
+}
+
 PyInterpreterState *
 PyInterpreterState_Hold(void)
 {
     PyInterpreterState *interp = PyInterpreterState_Get();
-    assert(_Py_atomic_load_ssize_relaxed(&interp->threads.finalizing.countdown) >= 0);
-    _Py_atomic_add_ssize(&interp->threads.finalizing.countdown, 1);
+    _PyInterpreterState_Incref(interp);
     return interp;
 }