Add a test for interpreter reference counts.

Author: ZeroIntensity

Include/cpython/pystate.h

@@ -270,7 +270,10 @@ PyAPI_FUNC(void) _PyInterpreterState_SetEvalFrameFunc(
  * incremented reference count. PyInterpreterState_Delete() won't delete the
  * full interpreter structure until the reference is released by
  * PyThreadState_Ensure() or PyInterpreterState_Release(). */
-PyAPI_FUNC(PyInterpreterState *) PyInterpreterState_Hold();
+PyAPI_FUNC(PyInterpreterState *) PyInterpreterState_Hold(void);
 
 /* Release a reference to an interpreter incremented by PyInterpreterState_Hold() */
 PyAPI_FUNC(void) PyInterpreterState_Release(PyInterpreterState *interp);
+
+// Export for '_testcapi' shared extension
+PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);

Modules/_testcapimodule.c

@@ -2546,6 +2546,30 @@ toggle_reftrace_printer(PyObject *ob, PyObject *arg)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_interp_refcount(PyObject *self, PyObject *unused)
+{
+    PyThreadState *save = PyThreadState_Get();
+    PyThreadState *tstate = Py_NewInterpreter();
+    assert(tstate == PyThreadState_Get());
+    PyInterpreterState *interp = PyThreadState_GetInterpreter(tstate);
+    assert(interp != NULL);
+
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterState *held = PyInterpreterState_Hold();
+    assert(_PyInterpreterState_Refcount(interp) == 2);
+    PyInterpreterState_Release(held);
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+
+    held = PyInterpreterState_Hold();
+    Py_EndInterpreter(tstate);
+    PyThreadState_Swap(save);
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterState_Release(held);
+
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -2640,6 +2664,7 @@ static PyMethodDef TestMethods[] = {
     {"test_atexit", test_atexit, METH_NOARGS},
     {"code_offset_to_line", _PyCFunction_CAST(code_offset_to_line), METH_FASTCALL},
     {"toggle_reftrace_printer", toggle_reftrace_printer, METH_O},
+    {"test_interp_refcount", test_interp_refcount, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Python/pystate.c

@@ -3196,6 +3196,15 @@ _Py_GetMainConfig(void)
     return _PyInterpreterState_GetConfig(interp);
 }
 
+Py_ssize_t
+_PyInterpreterState_Refcount(PyInterpreterState *interp)
+{
+    assert(interp != NULL);
+    Py_ssize_t refcount = _Py_atomic_load_ssize_relaxed(&interp->refcount);
+    assert(refcount > 0);
+    return refcount;
+}
+
 PyInterpreterState *
 PyInterpreterState_Hold(void)
 {
@@ -3208,9 +3217,6 @@ PyInterpreterState_Hold(void)
 void
 PyInterpreterState_Release(PyInterpreterState *interp)
 {
-    PyThreadState *tstate = PyThreadState_Get();
-    if (tstate->interp != interp) {
-        Py_FatalError("thread state has the wrong interpreter");
-    }
+    assert(interp != NULL);
     decref_interpreter(interp);
 }