Fix PyThreadState_EnsureFromView() and add a test.

Author: ZeroIntensity

Include/pystate.h

@@ -134,6 +134,7 @@ PyAPI_FUNC(void) PyInterpreterView_Close(PyInterpreterView *view);
 PyAPI_FUNC(PyInterpreterView *) PyInterpreterView_FromMain(void);
 
 PyAPI_FUNC(PyThreadState *) PyThreadState_Ensure(PyInterpreterGuard *guard);
+PyAPI_FUNC(PyThreadState *) PyThreadState_EnsureFromView(PyInterpreterView *view);
 PyAPI_FUNC(void) PyThreadState_Release(PyThreadState *tstate);
 
 #ifndef Py_LIMITED_API

Modules/_testcapimodule.c

@@ -2767,11 +2767,14 @@ test_interp_view_after_shutdown(PyObject *self, PyObject *unused)
     PyThreadState *save_tstate = PyThreadState_Swap(NULL);
     PyThreadState *interp_tstate = Py_NewInterpreter();
     if (interp_tstate == NULL) {
+        PyThreadState_Swap(save_tstate);
         return PyErr_NoMemory();
     }
 
     PyInterpreterView *view = PyInterpreterView_FromCurrent();
     if (view == NULL) {
+        Py_EndInterpreter(interp_tstate);
+        PyThreadState_Swap(save_tstate);
         return PyErr_NoMemory();
     }
 
@@ -2789,6 +2792,64 @@ test_interp_view_after_shutdown(PyObject *self, PyObject *unused)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_thread_state_ensure_view(PyObject *self, PyObject *unused)
+{
+    // For simplicity's sake, we assume that functions won't fail due to being
+    // out of memory.
+    PyThreadState *save_tstate = PyThreadState_Swap(NULL);
+    PyThreadState *interp_tstate = Py_NewInterpreter();
+    assert(interp_tstate != NULL);
+    assert(PyInterpreterState_Get() == PyThreadState_GetInterpreter(interp_tstate));
+
+    PyInterpreterView *main_view = PyInterpreterView_FromMain();
+    assert(main_view != NULL);
+
+    PyInterpreterView *view = PyInterpreterView_FromCurrent();
+    assert(view != NULL);
+
+    Py_BEGIN_ALLOW_THREADS;
+    PyThreadState *tstate = PyThreadState_EnsureFromView(view);
+    assert(tstate != NULL);
+    assert(PyThreadState_Get() == interp_tstate);
+
+    // Test a nested call
+    PyThreadState *tstate2 = PyThreadState_EnsureFromView(view);
+    assert(PyThreadState_Get() == interp_tstate);
+
+    // We're in a new interpreter now. PyThreadState_EnsureFromView() should
+    // now create a new thread state.
+    PyThreadState *main_tstate = PyThreadState_EnsureFromView(main_view);
+    assert(main_tstate == interp_tstate); // The old thread state
+    assert(PyInterpreterState_Get() == PyInterpreterState_Main());
+
+    // Going back to the old interpreter should create a new thread state again.
+    PyThreadState *tstate3 = PyThreadState_EnsureFromView(view);
+    assert(PyInterpreterState_Get() == PyThreadState_GetInterpreter(interp_tstate));
+    assert(PyThreadState_Get() != interp_tstate);
+    PyThreadState_Release(tstate3);
+    PyThreadState_Release(main_tstate);
+
+    // We're back in the original interpreter. PyThreadState_EnsureFromView() should
+    // no longer create a new thread state.
+    assert(PyThreadState_Get() == interp_tstate);
+    PyThreadState *tstate4 = PyThreadState_EnsureFromView(view);
+    assert(PyThreadState_Get() == interp_tstate);
+    PyThreadState_Release(tstate4);
+    PyThreadState_Release(tstate2);
+    PyThreadState_Release(tstate);
+    assert(PyThreadState_GetUnchecked() == NULL);
+    Py_END_ALLOW_THREADS;
+
+    assert(PyThreadState_Get() == interp_tstate);
+    PyInterpreterView_Close(view);
+    PyInterpreterView_Close(main_view);
+    Py_EndInterpreter(interp_tstate);
+    PyThreadState_Swap(save_tstate);
+
+    Py_RETURN_NONE;
+}
+
 
 static PyObject*
 test_soft_deprecated_macros(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(args))
@@ -2927,6 +2988,7 @@ static PyMethodDef TestMethods[] = {
     {"test_thread_state_ensure_nested", test_thread_state_ensure_nested, METH_NOARGS},
     {"test_thread_state_ensure_crossinterp", test_thread_state_ensure_crossinterp, METH_NOARGS},
     {"test_interp_view_after_shutdown", test_interp_view_after_shutdown, METH_NOARGS},
+    {"test_thread_state_ensure_view", test_thread_state_ensure_view, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Python/pylifecycle.c

@@ -2340,7 +2340,7 @@ make_pre_finalization_calls(PyThreadState *tstate, int subinterpreters)
                               || interp_has_atexit_callbacks(interp)
                               || interp_has_pending_calls(interp)
                               || has_subinterpreters
-                              || interp->finalization_guards.countdown > 0);
+                              || _Py_atomic_load_ssize_acquire(&interp->finalization_guards.countdown) > 0);
         if (!should_continue) {
             break;
         }

Python/pystate.c

@@ -3344,29 +3344,23 @@ _PyInterpreterGuard_GetInterpreter(PyInterpreterGuard *guard)
     return guard->interp;
 }
 
-static PyInterpreterGuard *
-try_acquire_interp_guard(PyInterpreterState *interp)
+static int
+try_acquire_interp_guard(PyInterpreterState *interp, PyInterpreterGuard *guard)
 {
     assert(interp != NULL);
     _PyRWMutex_RLock(&interp->finalization_guards.lock);
 
     if (_PyInterpreterState_GetFinalizing(interp) != NULL) {
         _PyRWMutex_RUnlock(&interp->finalization_guards.lock);
         assert(_Py_atomic_load_ssize_relaxed(&interp->finalization_guards.countdown) == 0);
-        return NULL;
-    }
-
-    PyInterpreterGuard *guard = PyMem_RawMalloc(sizeof(PyInterpreterGuard));
-    if (guard == NULL) {
-        _PyRWMutex_RUnlock(&interp->finalization_guards.lock);
-        return NULL;
+        return -1;
     }
 
     _Py_atomic_add_ssize(&interp->finalization_guards.countdown, 1);
     _PyRWMutex_RUnlock(&interp->finalization_guards.lock);
 
     guard->interp = interp;
-    return guard;
+    return 0;
 }
 
 PyInterpreterGuard *
@@ -3375,8 +3369,14 @@ PyInterpreterGuard_FromCurrent(void)
     PyInterpreterState *interp = _PyInterpreterState_GET();
     assert(interp != NULL);
 
-    PyInterpreterGuard *guard = try_acquire_interp_guard(interp);
+    PyInterpreterGuard *guard = PyMem_RawMalloc(sizeof(PyInterpreterGuard));
     if (guard == NULL) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+
+    if (try_acquire_interp_guard(interp, guard) < 0) {
+        PyMem_RawFree(guard);
         PyErr_SetString(PyExc_PythonFinalizationError,
                         "cannot acquire finalization guard anymore");
         return NULL;
@@ -3396,11 +3396,6 @@ PyInterpreterGuard_Close(PyInterpreterGuard *guard)
     _PyRWMutex_RUnlock(&interp->finalization_guards.lock);
 
     assert(old > 0);
-    if (old <= 0) {
-        Py_FatalError("interpreter has negative guard count, likely due"
-                      " to an extra PyInterpreterGuard_Close() call");
-    }
-
     PyMem_RawFree(guard);
 }
 
@@ -3440,18 +3435,32 @@ PyInterpreterGuard_FromView(PyInterpreterView *view)
     int64_t interp_id = view->id;
     assert(interp_id >= 0);
 
+    // This allocation has to happen before we acquire the runtime lock, because
+    // PyMem_RawMalloc() might call some weird callback (such as tracemalloc)
+    // that tries to re-entrantly acquire the lock.
+    PyInterpreterGuard *guard = PyMem_RawMalloc(sizeof(PyInterpreterGuard));
+    if (guard == NULL) {
+        return NULL;
+    }
+
     // Interpreters cannot be deleted while we hold the runtime lock.
     _PyRuntimeState *runtime = &_PyRuntime;
     HEAD_LOCK(runtime);
     PyInterpreterState *interp = interp_look_up_id(runtime, interp_id);
     if (interp == NULL) {
         HEAD_UNLOCK(runtime);
-        return 0;
+        PyMem_RawFree(guard);
+        return NULL;
     }
 
-    PyInterpreterGuard *guard = try_acquire_interp_guard(interp);
+    int result = try_acquire_interp_guard(interp, guard);
     HEAD_UNLOCK(runtime);
 
+    if (result < 0) {
+        PyMem_RawFree(guard);
+        return NULL;
+    }
+
     assert(guard == NULL || guard->interp != NULL);
     return guard;
 }
@@ -3528,20 +3537,24 @@ PyThreadState_EnsureFromView(PyInterpreterView *view)
         return NULL;
     }
 
-    PyThreadState *tstate = PyThreadState_Ensure(guard);
-    if (tstate == NULL) {
+    PyThreadState *result_tstate = PyThreadState_Ensure(guard);
+    if (result_tstate == NULL) {
         PyInterpreterGuard_Close(guard);
         return NULL;
     }
 
+    PyThreadState *tstate = current_fast_get();
+    assert(tstate != NULL);
+
     if (tstate->ensure.owned_guard != NULL) {
         assert(tstate->ensure.owned_guard->interp == guard->interp);
         PyInterpreterGuard_Close(guard);
     } else {
+        assert(tstate->ensure.owned_guard == NULL);
         tstate->ensure.owned_guard = guard;
     }
 
-    return tstate;
+    return result_tstate;
 }
 
 void