gh-157838: Merge biased refcounts on behalf of detached threads

Include/internal/pycore_pystate.h

@@ -150,6 +150,23 @@ extern void _PyThreadState_Detach(PyThreadState *tstate);
 // to the "detached" state.
 extern void _PyThreadState_Suspend(PyThreadState *tstate);
 
+#ifdef Py_GIL_DISABLED
+// Try to atomically transition a *different* thread's state from "detached"
+// to "suspended". On success, the target thread cannot attach until
+// _PyThreadState_ResumeDetached() is called, and the caller may safely
+// perform operations that are normally only permitted for the owning thread
+// (such as merging the biased reference counts of objects it owns).
+//
+// The caller must not run arbitrary Python code, allocate GC objects, or
+// stop the world while holding the thread in the suspended state.
+// Returns 1 on success, 0 if the thread was not in the "detached" state.
+extern int _PyThreadState_TrySuspendDetached(PyThreadState *tstate);
+
+// Undo a successful _PyThreadState_TrySuspendDetached(): switch the thread
+// back to "detached" and wake it if it is waiting to attach.
+extern void _PyThreadState_ResumeDetached(PyThreadState *tstate);
+#endif
+
 // Mark the thread state as "shutting down". This is used during interpreter
 // and runtime finalization. The thread may no longer attach to the
 // interpreter and will instead block via _PyThreadState_HangThread().

Python/brc.c

@@ -48,6 +48,47 @@ find_thread_state(struct _brc_bucket *bucket, uintptr_t thread_id)
     return NULL;
 }
 
+// Merge the refcounts of all objects in `to_merge` without deallocating.
+// Objects whose merged refcount is zero are moved to `to_dealloc`, which
+// borrows the chunks of `to_merge` so that this never allocates.
+static void
+merge_queued_refcounts(_PyObjectStack *to_merge, _PyObjectStack *to_dealloc)
+{
+    assert(to_dealloc->head == NULL);
+    _PyObjectStackChunk *buf = to_merge->head;
+    to_merge->head = NULL;
+    while (buf != NULL) {
+        Py_ssize_t n = 0;
+        for (Py_ssize_t i = 0; i < buf->n; i++) {
+            PyObject *ob = buf->objs[i];
+            // Subtract one when merging because the queue had a reference.
+            if (_Py_ExplicitMergeRefcount(ob, -1) == 0) {
+                buf->objs[n++] = ob;
+            }
+        }
+        _PyObjectStackChunk *prev = buf->prev;
+        if (n == 0) {
+            buf->n = 0;
+            _PyObjectStackChunk_Free(buf);
+        }
+        else {
+            buf->n = n;
+            buf->prev = to_dealloc->head;
+            to_dealloc->head = buf;
+        }
+        buf = prev;
+    }
+}
+
+static void
+dealloc_merged_objects(_PyObjectStack *to_dealloc)
+{
+    PyObject *ob;
+    while ((ob = _PyObjectStack_Pop(to_dealloc)) != NULL) {
+        _Py_Dealloc(ob);
+    }
+}
+
 // Enqueue an object to be merged by the owning thread. This steals a
 // reference to the object.
 void
@@ -93,6 +134,27 @@ _Py_brc_queue_object(PyObject *ob)
         return;
     }
 
+    if (_PyThreadState_TrySuspendDetached(&tstate->base)) {
+        // The owning thread is detached (e.g. blocked on a lock or in a
+        // system call) and may not run Python code again for a long time,
+        // so merge its queue on its behalf instead of waiting for it. While
+        // it is held in the "suspended" state it cannot attach and therefore
+        // cannot touch ob_ref_local or ob_tid.
+        //
+        // Only the merges happen while the thread is suspended: they are
+        // plain field updates. Deallocations run arbitrary code (and may
+        // stop the world, which would deadlock on the suspended thread), so
+        // they are deferred until after the thread is resumed and the
+        // bucket mutex is released.
+        _PyObjectStack to_dealloc = {0};
+        merge_queued_refcounts(&tstate->brc.objects_to_merge, &to_dealloc);
+        _PyThreadState_ResumeDetached(&tstate->base);
+        PyMutex_Unlock(&bucket->mutex);
+
+        dealloc_merged_objects(&to_dealloc);
+        return;
+    }
+
     // Notify owning thread
     _Py_set_eval_breaker_bit(&tstate->base, _PY_EVAL_EXPLICIT_MERGE_BIT);
 

Python/pystate.c

@@ -2380,6 +2380,30 @@ _PyThreadState_SetShuttingDown(PyThreadState *tstate)
 #endif
 }
 
+#ifdef Py_GIL_DISABLED
+int
+_PyThreadState_TrySuspendDetached(PyThreadState *tstate)
+{
+    assert(tstate != _PyThreadState_GET());
+    int expected = _Py_THREAD_DETACHED;
+    if (_Py_atomic_load_int_relaxed(&tstate->state) != expected) {
+        return 0;
+    }
+    return _Py_atomic_compare_exchange_int(&tstate->state, &expected,
+                                           _Py_THREAD_SUSPENDED);
+}
+
+void
+_PyThreadState_ResumeDetached(PyThreadState *tstate)
+{
+    assert(tstate != _PyThreadState_GET());
+    assert(_Py_atomic_load_int_relaxed(&tstate->state) == _Py_THREAD_SUSPENDED);
+    _Py_atomic_store_int(&tstate->state, _Py_THREAD_DETACHED);
+    // Wake the thread if it is parked in tstate_wait_attach().
+    _PyParkingLot_UnparkAll(&tstate->state);
+}
+#endif
+
 // Decrease stop-the-world counter of remaining number of threads that need to
 // pause. If we are the final thread to pause, notify the requesting thread.
 static void