pystate: refcount threads to handle race between interpreter shutdown and thread exit

Author: colesbury

Include/cpython/pystate.h

@@ -132,6 +132,8 @@ struct _ts {
 
     mi_heap_t *heaps[Py_NUM_HEAPS];
 
+    Py_ssize_t refcount;
+
     /* Has been initialized to a safe state.
 
        In order to be effective, this must be set to 0 during or right

Include/internal/pycore_pystate.h

@@ -169,6 +169,8 @@ PyAPI_FUNC(PyThreadState *) _PyThreadState_UnlinkExcept(
     int already_dead);
 PyAPI_FUNC(void) _PyThreadState_DeleteGarbage(PyThreadState *garbage);
 
+extern void _PyThreadState_Exit(PyThreadState *tstate);
+
 static inline void
 _PyThreadState_Signal(PyThreadState *tstate, uintptr_t bit)
 {

Python/ceval_gil.c

@@ -263,7 +263,7 @@ take_gil(PyThreadState *tstate)
            This code path can be reached by a daemon thread after Py_Finalize()
            completes. In this case, tstate is a dangling pointer: points to
            PyThreadState freed memory. */
-        PyThread_exit_thread();
+        _PyThreadState_Exit(tstate);
     }
 
     assert(is_tstate_valid(tstate));
@@ -310,7 +310,7 @@ take_gil(PyThreadState *tstate)
                     _PyThreadState_Unsignal(gil->holder, EVAL_DROP_GIL);
                 }
                 MUTEX_UNLOCK(gil->mutex);
-                PyThread_exit_thread();
+                _PyThreadState_Exit(tstate);
             }
             assert(is_tstate_valid(tstate));
 
@@ -350,7 +350,7 @@ take_gil(PyThreadState *tstate)
            wait_for_thread_shutdown() from Py_Finalize(). */
         MUTEX_UNLOCK(gil->mutex);
         drop_gil(ceval, ceval2, tstate);
-        PyThread_exit_thread();
+        _PyThreadState_Exit(tstate);
     }
     assert(is_tstate_valid(tstate));
 

Python/pystate.c

@@ -992,6 +992,26 @@ free_threadstate(PyThreadState *tstate)
     }
 }
 
+static void
+_PyThreadState_DecRef(PyThreadState *tstate)
+{
+    if (tstate != &tstate->interp->_initial_thread.tstate) {
+        if (_Py_atomic_add_ssize(&tstate->refcount, -1) == 1) {
+            free_threadstate(tstate);
+        }
+    }
+}
+
+void
+_PyThreadState_Exit(PyThreadState *tstate)
+{
+    if (_PyThreadState_GetStatus(tstate) == _Py_THREAD_ATTACHED) {
+        _Py_atomic_store_int(&tstate->status, _Py_THREAD_DETACHED);
+    }
+    _PyThreadState_DecRef(tstate);
+    PyThread_exit_thread();
+}
+
 /* Get the thread state to a minimal consistent state.
    Further init happens in pylifecycle.c before it can be used.
    All fields not initialized here are expected to be zeroed out,
@@ -1061,6 +1081,7 @@ init_threadstate(PyThreadState *tstate,
     if (_PyRuntime.stop_the_world_requested) {
         tstate->status = _Py_THREAD_GC;
     }
+    tstate->refcount = 2;
     tstate->_initialized = 1;
 }
 
@@ -1562,7 +1583,12 @@ void
 _PyThreadState_DeleteExcept(_PyRuntimeState *runtime, PyThreadState *tstate)
 {
     PyThreadState *garbage = _PyThreadState_UnlinkExcept(runtime, tstate, 0);
-    _PyThreadState_DeleteGarbage(garbage);
+    PyThreadState *next;
+    for (PyThreadState *p = garbage; p; p = next) {
+        next = p->next;
+        PyThreadState_Clear(p);
+        _PyThreadState_DecRef(tstate);
+    }
 }
 
 PyThreadState *