gh-114940: Add a Per-Interpreter Lock For the List of Thread States

Include/internal/pycore_interp.h

@@ -127,6 +127,7 @@ struct _is {
 
     uintptr_t last_restart_version;
     struct pythreads {
+        _PyRecursiveMutex mutex;
         uint64_t next_unique_id;
         /* The linked list of threads, newest first. */
         PyThreadState *head;

Include/internal/pycore_lock.h

@@ -151,10 +151,12 @@ _PyOnceFlag_CallOnce(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg)
     return _PyOnceFlag_CallOnceSlow(flag, fn, arg);
 }
 
+#define _PyThread_ident_t unsigned long long
+
 // A recursive mutex. The mutex should zero-initialized.
 typedef struct {
     PyMutex mutex;
-    unsigned long long thread;  // i.e., PyThread_get_thread_ident_ex()
+    _PyThread_ident_t thread;  // i.e., PyThread_get_thread_ident_ex()
     size_t level;
 } _PyRecursiveMutex;
 
@@ -164,6 +166,37 @@ extern PyLockStatus _PyRecursiveMutex_LockTimed(_PyRecursiveMutex *m, PyTime_t t
 PyAPI_FUNC(void) _PyRecursiveMutex_Unlock(_PyRecursiveMutex *m);
 extern int _PyRecursiveMutex_TryUnlock(_PyRecursiveMutex *m);
 
+PyAPI_FUNC(_PyThread_ident_t) PyThread_get_thread_ident_ex(void);
+
+static inline void
+_PyRecursiveMutex_LockFlags(_PyRecursiveMutex *m, _PyLockFlags flags)
+{
+    uint8_t expected = _Py_UNLOCKED;
+    if (_Py_atomic_compare_exchange_uint8(&m->mutex._bits, &expected, _Py_LOCKED)) {
+        m->thread = PyThread_get_thread_ident_ex();
+        assert(m->level == 0);
+    }
+    else {
+        _PyRecursiveMutex_LockTimed(m, -1, flags);
+    }
+}
+
+#ifdef HAVE_FORK
+PyAPI_FUNC(_PyThread_ident_t) PyThread_get_thread_ident_ex(void);
+
+static inline void
+_PyRecursiveMutex_at_fork_reinit(_PyRecursiveMutex *m,
+                                 _PyThread_ident_t parent)
+{
+    if (m->thread == parent) {
+        m->thread = PyThread_get_thread_ident_ex();
+    }
+    else {
+        memset(m, 0, sizeof(*m));
+    }
+}
+#endif
+
 // A readers-writer (RW) lock. The lock supports multiple concurrent readers or
 // a single writer. The lock is write-preferring: if a writer is waiting while
 // the lock is read-locked then, new readers will be blocked. This avoids

Include/internal/pycore_pystate.h

@@ -269,14 +269,18 @@ extern int _PyOS_InterruptOccurred(PyThreadState *tstate);
 #define HEAD_UNLOCK(runtime) \
     PyMutex_Unlock(&(runtime)->interpreters.mutex)
 
+#define THREADS_HEAD_LOCK(interp) \
+    _PyRecursiveMutex_LockFlags(&(interp)->threads.mutex, _Py_LOCK_DONT_DETACH)
+#define THREADS_HEAD_UNLOCK(interp) \
+    _PyRecursiveMutex_Unlock(&(interp)->threads.mutex)
+
 #define _Py_FOR_EACH_TSTATE_UNLOCKED(interp, t) \
     for (PyThreadState *t = interp->threads.head; t; t = t->next)
 #define _Py_FOR_EACH_TSTATE_BEGIN(interp, t) \
-    HEAD_LOCK(interp->runtime); \
+    THREADS_HEAD_LOCK(interp); \
     _Py_FOR_EACH_TSTATE_UNLOCKED(interp, t)
 #define _Py_FOR_EACH_TSTATE_END(interp) \
-    HEAD_UNLOCK(interp->runtime)
-
+    THREADS_HEAD_UNLOCK(interp)
 
 // Get the configuration of the current interpreter.
 // The caller must hold the GIL.

Include/internal/pycore_runtime.h

@@ -164,6 +164,11 @@ typedef struct pyruntimestate {
         _Py_AuditHookEntry *head;
     } audit_hooks;
 
+#ifdef HAVE_FORK
+    // Guarded by multiple global locks.
+    PyThread_ident_t fork_tid;
+#endif
+
     struct _py_object_runtime_state object_state;
     struct _Py_float_runtime_state float_state;
     struct _Py_unicode_runtime_state unicode_state;

Modules/posixmodule.c

@@ -624,11 +624,13 @@ PyOS_BeforeFork(void)
     _PyImport_AcquireLock(interp);
     _PyEval_StopTheWorldAll(&_PyRuntime);
     HEAD_LOCK(&_PyRuntime);
+    _PyRuntime.fork_tid = PyThread_get_thread_ident_ex();
 }
 
 void
 PyOS_AfterFork_Parent(void)
 {
+    _PyRuntime.fork_tid = 0;
     HEAD_UNLOCK(&_PyRuntime);
     _PyEval_StartTheWorldAll(&_PyRuntime);
 
@@ -648,6 +650,7 @@ PyOS_AfterFork_Child(void)
     if (_PyStatus_EXCEPTION(status)) {
         goto fatal_error;
     }
+    _PyRuntime.fork_tid = 0;
 
     PyThreadState *tstate = _PyThreadState_GET();
     _Py_EnsureTstateNotNULL(tstate);

Objects/object.c

@@ -119,11 +119,11 @@ get_reftotal(PyInterpreterState *interp)
        since we can't determine which interpreter updated it. */
     Py_ssize_t total = REFTOTAL(interp);
 #ifdef Py_GIL_DISABLED
-    _Py_FOR_EACH_TSTATE_UNLOCKED(interp, p) {
-        /* This may race with other threads modifications to their reftotal */
+    _Py_FOR_EACH_TSTATE_BEGIN(interp, p) {
         _PyThreadStateImpl *tstate_impl = (_PyThreadStateImpl *)p;
         total += _Py_atomic_load_ssize_relaxed(&tstate_impl->reftotal);
     }
+    _Py_FOR_EACH_TSTATE_END(interp);
 #endif
     return total;
 }
@@ -317,10 +317,7 @@ _Py_GetLegacyRefTotal(void)
 Py_ssize_t
 _PyInterpreterState_GetRefTotal(PyInterpreterState *interp)
 {
-    HEAD_LOCK(&_PyRuntime);
-    Py_ssize_t total = get_reftotal(interp);
-    HEAD_UNLOCK(&_PyRuntime);
-    return total;
+    return get_reftotal(interp);
 }
 
 #endif /* Py_REF_DEBUG */

Objects/obmalloc.c

@@ -1439,13 +1439,14 @@ get_mimalloc_allocated_blocks(PyInterpreterState *interp)
 {
     size_t allocated_blocks = 0;
 #ifdef Py_GIL_DISABLED
-    _Py_FOR_EACH_TSTATE_UNLOCKED(interp, t) {
+    _Py_FOR_EACH_TSTATE_BEGIN(interp, t) {
         _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)t;
         for (int i = 0; i < _Py_MIMALLOC_HEAP_COUNT; i++) {
             mi_heap_t *heap = &tstate->mimalloc.heaps[i];
             mi_heap_visit_blocks(heap, false, &count_blocks, &allocated_blocks);
         }
     }
+    _Py_FOR_EACH_TSTATE_END(interp);
 
     mi_abandoned_pool_t *pool = &interp->mimalloc.abandoned_pool;
     for (uint8_t tag = 0; tag < _Py_MIMALLOC_HEAP_COUNT; tag++) {

Python/gc_free_threading.c

@@ -304,7 +304,8 @@ gc_visit_heaps_lock_held(PyInterpreterState *interp, mi_block_visit_fun *visitor
     Py_ssize_t offset_pre = offset_base + 2 * sizeof(PyObject*);
 
     // visit each thread's heaps for GC objects
-    _Py_FOR_EACH_TSTATE_UNLOCKED(interp, p) {
+    int failed = 0;
+    _Py_FOR_EACH_TSTATE_BEGIN(interp, p) {
         struct _mimalloc_thread_state *m = &((_PyThreadStateImpl *)p)->mimalloc;
         if (!_Py_atomic_load_int(&m->initialized)) {
             // The thread may not have called tstate_mimalloc_bind() yet.
@@ -314,14 +315,20 @@ gc_visit_heaps_lock_held(PyInterpreterState *interp, mi_block_visit_fun *visitor
         arg->offset = offset_base;
         if (!mi_heap_visit_blocks(&m->heaps[_Py_MIMALLOC_HEAP_GC], true,
                                   visitor, arg)) {
-            return -1;
+            failed = 1;
+            break;
         }
         arg->offset = offset_pre;
         if (!mi_heap_visit_blocks(&m->heaps[_Py_MIMALLOC_HEAP_GC_PRE], true,
                                   visitor, arg)) {
-            return -1;
+            failed = 1;
+            break;
         }
     }
+    _Py_FOR_EACH_TSTATE_END(interp);
+    if (failed) {
+        return -1;
+    }
 
     // visit blocks in the per-interpreter abandoned pool (from dead threads)
     mi_abandoned_pool_t *pool = &interp->mimalloc.abandoned_pool;