gh-116522: Refactor `_PyThreadState_DeleteExcept` (#117131)

Split `_PyThreadState_DeleteExcept` into two functions:

- `_PyThreadState_RemoveExcept` removes all thread states other than one
  passed as an argument. It returns the removed thread states as a
  linked list.

- `_PyThreadState_DeleteList` deletes those dead thread states. It may
  call destructors, so we want to "start the world" before calling
  `_PyThreadState_DeleteList` to avoid potential deadlocks.

5 files changed, 41 insertions, 23 deletions

diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
index 9aa4392..35e266a 100644
--- a/Include/internal/pycore_pystate.h
+++ b/Include/internal/pycore_pystate.h
@@ -218,7 +218,8 @@ extern PyThreadState * _PyThreadState_New(
     PyInterpreterState *interp,
     int whence);
 extern void _PyThreadState_Bind(PyThreadState *tstate);
-extern void _PyThreadState_DeleteExcept(PyThreadState *tstate);
+extern PyThreadState * _PyThreadState_RemoveExcept(PyThreadState *tstate);
+extern void _PyThreadState_DeleteList(PyThreadState *list);
 extern void _PyThreadState_ClearMimallocHeaps(PyThreadState *tstate);
 
 // Export for '_testinternalcapi' shared extension
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
index 8867916..a4b635e 100644
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -664,6 +664,14 @@ PyOS_AfterFork_Child(void)
         goto fatal_error;
     }
 
+    // Remove the dead thread states. We "start the world" once we are the only
+    // thread state left to undo the stop the world call in `PyOS_BeforeFork`.
+    // That needs to happen before `_PyThreadState_DeleteList`, because that
+    // may call destructors.
+    PyThreadState *list = _PyThreadState_RemoveExcept(tstate);
+    _PyEval_StartTheWorldAll(&_PyRuntime);
+    _PyThreadState_DeleteList(list);
+
     status = _PyImport_ReInitLock(tstate->interp);
     if (_PyStatus_EXCEPTION(status)) {
         goto fatal_error;
diff --git a/Python/ceval_gil.c b/Python/ceval_gil.c
index 78c13d6..d88ac65 100644
--- a/Python/ceval_gil.c
+++ b/Python/ceval_gil.c
@@ -579,9 +579,8 @@ PyEval_ReleaseThread(PyThreadState *tstate)
 }
 
 #ifdef HAVE_FORK
-/* This function is called from PyOS_AfterFork_Child to destroy all threads
-   which are not running in the child process, and clear internal locks
-   which might be held by those threads. */
+/* This function is called from PyOS_AfterFork_Child to re-initialize the
+   GIL and pending calls lock. */
 PyStatus
 _PyEval_ReInitThreads(PyThreadState *tstate)
 {
@@ -598,8 +597,6 @@ _PyEval_ReInitThreads(PyThreadState *tstate)
     struct _pending_calls *pending = &tstate->interp->ceval.pending;
     _PyMutex_at_fork_reinit(&pending->mutex);
 
-    /* Destroy all threads except the current one */
-    _PyThreadState_DeleteExcept(tstate);
     return _PyStatus_OK();
 }
 #endif
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 683534d..1d315b8 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -1934,8 +1934,11 @@ Py_FinalizeEx(void)
        will be called in the current Python thread. Since
        _PyRuntimeState_SetFinalizing() has been called, no other Python thread
        can take the GIL at this point: if they try, they will exit
-       immediately. */
-    _PyThreadState_DeleteExcept(tstate);
+       immediately. We start the world once we are the only thread state left,
+       before we call destructors. */
+    PyThreadState *list = _PyThreadState_RemoveExcept(tstate);
+    _PyEval_StartTheWorldAll(runtime);
+    _PyThreadState_DeleteList(list);
 
     /* At this point no Python code should be running at all.
        The only thread state left should be the main thread of the main
diff --git a/Python/pystate.c b/Python/pystate.c
index 3ef4051..47d327a 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -1763,15 +1763,17 @@ PyThreadState_DeleteCurrent(void)
 }
 
 
-/*
- * Delete all thread states except the one passed as argument.
- * Note that, if there is a current thread state, it *must* be the one
- * passed as argument.  Also, this won't touch any other interpreters
- * than the current one, since we don't know which thread state should
- * be kept in those other interpreters.
- */
-void
-_PyThreadState_DeleteExcept(PyThreadState *tstate)
+// Unlinks and removes all thread states from `tstate->interp`, with the
+// exception of the one passed as an argument. However, it does not delete
+// these thread states. Instead, it returns the removed thread states as a
+// linked list.
+//
+// Note that if there is a current thread state, it *must* be the one
+// passed as argument.  Also, this won't touch any interpreters other
+// than the current one, since we don't know which thread state should
+// be kept in those other interpreters.
+PyThreadState *
+_PyThreadState_RemoveExcept(PyThreadState *tstate)
 {
     assert(tstate != NULL);
     PyInterpreterState *interp = tstate->interp;
@@ -1783,8 +1785,7 @@ _PyThreadState_DeleteExcept(PyThreadState *tstate)
 
     HEAD_LOCK(runtime);
     /* Remove all thread states, except tstate, from the linked list of
-       thread states.  This will allow calling PyThreadState_Clear()
-       without holding the lock. */
+       thread states. */
     PyThreadState *list = interp->threads.head;
     if (list == tstate) {
         list = tstate->next;
@@ -1799,11 +1800,19 @@ _PyThreadState_DeleteExcept(PyThreadState *tstate)
     interp->threads.head = tstate;
     HEAD_UNLOCK(runtime);
 
-    _PyEval_StartTheWorldAll(runtime);
+    return list;
+}
+
+// Deletes the thread states in the linked list `list`.
+//
+// This is intended to be used in conjunction with _PyThreadState_RemoveExcept.
+void
+_PyThreadState_DeleteList(PyThreadState *list)
+{
+    // The world can't be stopped because we PyThreadState_Clear() can
+    // call destructors.
+    assert(!_PyRuntime.stoptheworld.world_stopped);
 
-    /* Clear and deallocate all stale thread states.  Even if this
-       executes Python code, we should be safe since it executes
-       in the current thread, not one of the stale threads. */
     PyThreadState *p, *next;
     for (p = list; p; p = next) {
         next = p->next;