gh-114940: Use fine-grained mutex protection for PyInterpreterState.threads

Include/internal/pycore_interp.h

@@ -140,6 +140,8 @@ struct _is {
            or the size specified by the THREAD_STACK_SIZE macro. */
         /* Used in Python/thread.c. */
         size_t stacksize;
+        /* The lock that protects this struct. */
+        PyMutex mutex;
     } threads;
 
     /* Reference to the _PyRuntime global variable. This field exists

Include/internal/pycore_pystate.h

@@ -269,6 +269,11 @@ extern int _PyOS_InterruptOccurred(PyThreadState *tstate);
 #define HEAD_UNLOCK(runtime) \
     PyMutex_Unlock(&(runtime)->interpreters.mutex)
 
+#define INTERP_HEAD_LOCK(interpstate) \
+    PyMutex_LockFlags(&(interpstate)->threads.mutex, _Py_LOCK_DONT_DETACH)
+#define INTERP_HEAD_UNLOCK(interpstate) \
+    PyMutex_Unlock(&(interpstate)->threads.mutex)
+
 // Get the configuration of the current interpreter.
 // The caller must hold the GIL.
 // Export for test_peg_generator.

Objects/codeobject.c

@@ -2871,20 +2871,24 @@ get_indices_in_use(PyInterpreterState *interp, struct flag_set *in_use)
     assert(interp->stoptheworld.world_stopped);
     assert(in_use->flags == NULL);
     int32_t max_index = 0;
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *p = interp->threads.head; p != NULL; p = p->next) {
         int32_t idx = ((_PyThreadStateImpl *) p)->tlbc_index;
         if (idx > max_index) {
             max_index = idx;
         }
     }
+    INTERP_HEAD_UNLOCK(interp);
     in_use->size = (size_t) max_index + 1;
     in_use->flags = PyMem_Calloc(in_use->size, sizeof(*in_use->flags));
     if (in_use->flags == NULL) {
         return -1;
     }
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *p = interp->threads.head; p != NULL; p = p->next) {
         in_use->flags[((_PyThreadStateImpl *) p)->tlbc_index] = 1;
     }
+    INTERP_HEAD_UNLOCK(interp);
     return 0;
 }
 

Objects/object.c

@@ -119,11 +119,13 @@ get_reftotal(PyInterpreterState *interp)
        since we can't determine which interpreter updated it. */
     Py_ssize_t total = REFTOTAL(interp);
 #ifdef Py_GIL_DISABLED
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *p = interp->threads.head; p != NULL; p = p->next) {
         /* This may race with other threads modifications to their reftotal */
         _PyThreadStateImpl *tstate_impl = (_PyThreadStateImpl *)p;
         total += _Py_atomic_load_ssize_relaxed(&tstate_impl->reftotal);
     }
+    INTERP_HEAD_UNLOCK(interp);
 #endif
     return total;
 }

Objects/obmalloc.c

@@ -1405,6 +1405,7 @@ get_mimalloc_allocated_blocks(PyInterpreterState *interp)
 {
     size_t allocated_blocks = 0;
 #ifdef Py_GIL_DISABLED
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *t = interp->threads.head; t != NULL; t = t->next) {
         _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)t;
         for (int i = 0; i < _Py_MIMALLOC_HEAP_COUNT; i++) {
@@ -1418,6 +1419,7 @@ get_mimalloc_allocated_blocks(PyInterpreterState *interp)
         _mi_abandoned_pool_visit_blocks(pool, tag, false, &count_blocks,
                                         &allocated_blocks);
     }
+    INTERP_HEAD_UNLOCK(interp);
 #else
     // TODO(sgross): this only counts the current thread's blocks.
     mi_heap_t *heap = mi_heap_get_default();

Python/ceval.c

@@ -2402,11 +2402,11 @@ void
 PyEval_SetProfileAllThreads(Py_tracefunc func, PyObject *arg)
 {
     PyThreadState *this_tstate = _PyThreadState_GET();
-    PyInterpreterState* interp = this_tstate->interp;
+    PyInterpreterState *interp = this_tstate->interp;
 
     _PyRuntimeState *runtime = &_PyRuntime;
     HEAD_LOCK(runtime);
-    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
+    PyThreadState *ts = PyInterpreterState_ThreadHead(interp);
     HEAD_UNLOCK(runtime);
 
     while (ts) {
@@ -2433,11 +2433,11 @@ void
 PyEval_SetTraceAllThreads(Py_tracefunc func, PyObject *arg)
 {
     PyThreadState *this_tstate = _PyThreadState_GET();
-    PyInterpreterState* interp = this_tstate->interp;
+    PyInterpreterState *interp = this_tstate->interp;
 
     _PyRuntimeState *runtime = &_PyRuntime;
     HEAD_LOCK(runtime);
-    PyThreadState* ts = PyInterpreterState_ThreadHead(interp);
+    PyThreadState *ts = PyInterpreterState_ThreadHead(interp);
     HEAD_UNLOCK(runtime);
 
     while (ts) {

Python/gc_free_threading.c

@@ -350,9 +350,9 @@ gc_visit_heaps(PyInterpreterState *interp, mi_block_visit_fun *visitor,
     assert(interp->stoptheworld.world_stopped);
 
     int err;
-    HEAD_LOCK(&_PyRuntime);
+    INTERP_HEAD_LOCK(interp);
     err = gc_visit_heaps_lock_held(interp, visitor, arg);
-    HEAD_UNLOCK(&_PyRuntime);
+    INTERP_HEAD_UNLOCK(interp);
     return err;
 }
 
@@ -374,7 +374,7 @@ gc_visit_stackref(_PyStackRef stackref)
 static void
 gc_visit_thread_stacks(PyInterpreterState *interp)
 {
-    HEAD_LOCK(&_PyRuntime);
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *p = interp->threads.head; p != NULL; p = p->next) {
         for (_PyInterpreterFrame *f = p->current_frame; f != NULL; f = f->previous) {
             PyObject *executable = PyStackRef_AsPyObjectBorrow(f->f_executable);
@@ -390,7 +390,7 @@ gc_visit_thread_stacks(PyInterpreterState *interp)
             }
         }
     }
-    HEAD_UNLOCK(&_PyRuntime);
+    INTERP_HEAD_UNLOCK(interp);
 }
 
 static void
@@ -429,14 +429,14 @@ process_delayed_frees(PyInterpreterState *interp)
 
     // Merge the queues from other threads into our own queue so that we can
     // process all of the pending delayed free requests at once.
-    HEAD_LOCK(&_PyRuntime);
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *p = interp->threads.head; p != NULL; p = p->next) {
         _PyThreadStateImpl *other = (_PyThreadStateImpl *)p;
         if (other != current_tstate) {
             llist_concat(&current_tstate->mem_free_queue, &other->mem_free_queue);
         }
     }
-    HEAD_UNLOCK(&_PyRuntime);
+    INTERP_HEAD_UNLOCK(interp);
 
     _PyMem_ProcessDelayed((PyThreadState *)current_tstate);
 }
@@ -1226,7 +1226,7 @@ gc_collect_internal(PyInterpreterState *interp, struct collection_state *state,
         state->gcstate->old[i-1].count = 0;
     }
 
-    HEAD_LOCK(&_PyRuntime);
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *p = interp->threads.head; p != NULL; p = p->next) {
         _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)p;
 
@@ -1236,7 +1236,7 @@ gc_collect_internal(PyInterpreterState *interp, struct collection_state *state,
         // merge refcounts for all queued objects
         merge_queued_objects(tstate, state);
     }
-    HEAD_UNLOCK(&_PyRuntime);
+    INTERP_HEAD_UNLOCK(interp);
 
     process_delayed_frees(interp);
 
@@ -1991,13 +1991,13 @@ PyUnstable_GC_VisitObjects(gcvisitobjects_t callback, void *arg)
 void
 _PyGC_ClearAllFreeLists(PyInterpreterState *interp)
 {
-    HEAD_LOCK(&_PyRuntime);
+    INTERP_HEAD_LOCK(interp);
     _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)interp->threads.head;
     while (tstate != NULL) {
         _PyObject_ClearFreeLists(&tstate->freelists, 0);
         tstate = (_PyThreadStateImpl *)tstate->base.next;
     }
-    HEAD_UNLOCK(&_PyRuntime);
+    INTERP_HEAD_UNLOCK(interp);
 }
 
 #endif  // Py_GIL_DISABLED

Python/instrumentation.c

@@ -1006,13 +1006,12 @@ set_global_version(PyThreadState *tstate, uint32_t version)
 
 #ifdef Py_GIL_DISABLED
     // Set the version on all threads in free-threaded builds.
-    _PyRuntimeState *runtime = &_PyRuntime;
-    HEAD_LOCK(runtime);
+    INTERP_HEAD_LOCK(interp);
     for (tstate = interp->threads.head; tstate;
          tstate = PyThreadState_Next(tstate)) {
         set_version_raw(&tstate->eval_breaker, version);
     };
-    HEAD_UNLOCK(runtime);
+    INTERP_HEAD_UNLOCK(interp);
 #else
     // Normal builds take the current version from instrumentation_version when
     // attaching a thread, so we only have to set the current thread's version.

Python/pystate.c

@@ -773,7 +773,6 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
 {
     assert(interp != NULL);
     assert(tstate != NULL);
-    _PyRuntimeState *runtime = interp->runtime;
 
     /* XXX Conditions we need to enforce:
 
@@ -790,18 +789,16 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
     }
 
     // Clear the current/main thread state last.
-    HEAD_LOCK(runtime);
+    INTERP_HEAD_LOCK(interp);
     PyThreadState *p = interp->threads.head;
-    HEAD_UNLOCK(runtime);
     while (p != NULL) {
         // See https://github.com/python/cpython/issues/102126
         // Must be called without HEAD_LOCK held as it can deadlock
         // if any finalizer tries to acquire that lock.
         PyThreadState_Clear(p);
-        HEAD_LOCK(runtime);
         p = p->next;
-        HEAD_UNLOCK(runtime);
     }
+    INTERP_HEAD_UNLOCK(interp);
     if (tstate->interp == interp) {
         /* We fix tstate->_status below when we for sure aren't using it
            (e.g. no longer need the GIL). */
@@ -1851,13 +1848,13 @@ _PyThreadState_RemoveExcept(PyThreadState *tstate)
 {
     assert(tstate != NULL);
     PyInterpreterState *interp = tstate->interp;
-    _PyRuntimeState *runtime = interp->runtime;
 
 #ifdef Py_GIL_DISABLED
+    _PyRuntimeState *runtime = interp->runtime;
     assert(runtime->stoptheworld.world_stopped);
 #endif
 
-    HEAD_LOCK(runtime);
+    INTERP_HEAD_LOCK(interp);
     /* Remove all thread states, except tstate, from the linked list of
        thread states. */
     PyThreadState *list = interp->threads.head;
@@ -1872,7 +1869,7 @@ _PyThreadState_RemoveExcept(PyThreadState *tstate)
     }
     tstate->prev = tstate->next = NULL;
     interp->threads.head = tstate;
-    HEAD_UNLOCK(runtime);
+    INTERP_HEAD_UNLOCK(interp);
 
     return list;
 }
@@ -2339,7 +2336,6 @@ _PyEval_StartTheWorld(PyInterpreterState *interp)
 int
 PyThreadState_SetAsyncExc(unsigned long id, PyObject *exc)
 {
-    _PyRuntimeState *runtime = &_PyRuntime;
     PyInterpreterState *interp = _PyInterpreterState_GET();
 
     /* Although the GIL is held, a few C API functions can be called
@@ -2348,7 +2344,7 @@ PyThreadState_SetAsyncExc(unsigned long id, PyObject *exc)
      * list of thread states we're traversing, so to prevent that we lock
      * head_mutex for the duration.
      */
-    HEAD_LOCK(runtime);
+    INTERP_HEAD_LOCK(interp);
     for (PyThreadState *tstate = interp->threads.head; tstate != NULL; tstate = tstate->next) {
         if (tstate->thread_id != id) {
             continue;
@@ -2363,13 +2359,13 @@ PyThreadState_SetAsyncExc(unsigned long id, PyObject *exc)
          */
         Py_XINCREF(exc);
         PyObject *old_exc = _Py_atomic_exchange_ptr(&tstate->async_exc, exc);
-        HEAD_UNLOCK(runtime);
+        INTERP_HEAD_UNLOCK(interp);
 
         Py_XDECREF(old_exc);
         _Py_set_eval_breaker_bit(tstate, _PY_ASYNC_EXCEPTION_BIT);
         return 1;
     }
-    HEAD_UNLOCK(runtime);
+    INTERP_HEAD_UNLOCK(interp);
     return 0;
 }