Add tests for draining QSBR page lists during GC

Add _testinternalcapi.get_qsbr_page_list_length(), which counts the pages on
all threads' QSBR page lists, and use it to check that gc.collect() frees
empty pages parked by the calling thread and by an idle thread. Also check
that the GC skips thread states that have not initialized their mimalloc
state yet.

Author: Matt Page

Lib/test/test_free_threading/test_gc.py

@@ -10,6 +10,11 @@
 from test import support
 from test.support import threading_helper
 
+try:
+    import _testinternalcapi
+except ImportError:
+    _testinternalcapi = None
+
 
 class MyObj:
     pass
@@ -210,5 +215,85 @@ def reader():
             pass
 
 
+def allocate_and_free():
+    # Free enough small objects to leave many mimalloc pages empty. Empty
+    # pages in the object heaps wait on QSBR before they are freed.
+    junk = [[] for _ in range(100_000)]
+    del junk
+
+
+@threading_helper.requires_working_threading()
+@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
+class TestQSBRPageList(TestCase):
+    # gh-158972: gc.collect() should free empty pages waiting on QSBR, even
+    # if the thread that parked them is idle.
+
+    def parked_pages(self):
+        return _testinternalcapi.get_qsbr_page_list_length()
+
+    def test_gc_collect_frees_parked_pages(self):
+        gc.collect()
+        allocate_and_free()
+        parked = self.parked_pages()
+        if parked < 10:
+            self.skipTest(f"only {parked} pages were parked")
+        gc.collect()
+        self.assertLess(self.parked_pages(), parked // 4)
+
+    def test_gc_collect_frees_pages_parked_by_idle_thread(self):
+        done = threading.Lock()
+        done.acquire()
+        resume = threading.Lock()
+        resume.acquire()
+
+        def worker():
+            allocate_and_free()
+            done.release()
+            # Block without allocating so the worker doesn't free its parked
+            # pages itself.
+            resume.acquire()
+
+        t = Thread(target=worker)
+        t.start()
+        try:
+            done.acquire()
+            parked = self.parked_pages()
+            if parked < 10:
+                self.skipTest(f"only {parked} pages were parked")
+            gc.collect()
+            self.assertLess(self.parked_pages(), parked // 4)
+        finally:
+            resume.release()
+            t.join()
+
+    def test_gc_collect_while_starting_threads(self):
+        # A new thread is added to the interpreter before it initializes its
+        # mimalloc state, so the GC can see thread states without a usable
+        # page list.
+        stop = threading.Event()
+
+        def collector():
+            while not stop.is_set():
+                gc.collect()
+
+        deadline = time.monotonic() + 2.0
+
+        def starter():
+            while time.monotonic() < deadline:
+                threads = [Thread(target=lambda: None) for _ in range(16)]
+                with threading_helper.start_threads(threads):
+                    pass
+
+        c = Thread(target=collector)
+        c.start()
+        try:
+            starters = [Thread(target=starter) for _ in range(4)]
+            with threading_helper.start_threads(starters):
+                pass
+        finally:
+            stop.set()
+            c.join()
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_testinternalcapi.c

@@ -2885,6 +2885,28 @@ get_long_lived_total(PyObject *self, PyObject *Py_UNUSED(ignored))
     return PyLong_FromInt64(PyInterpreterState_Get()->gc.long_lived_total);
 }
 
+// Return the total number of mimalloc pages on all threads' QSBR page lists.
+static PyObject *
+get_qsbr_page_list_length(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    Py_ssize_t count = 0;
+    _PyEval_StopTheWorld(interp);
+    _Py_FOR_EACH_TSTATE_BEGIN(interp, p) {
+        struct _mimalloc_thread_state *m = &((_PyThreadStateImpl *)p)->mimalloc;
+        if (!_Py_atomic_load_int(&m->initialized)) {
+            continue;
+        }
+        struct llist_node *node;
+        llist_for_each(node, &m->page_list) {
+            count++;
+        }
+    }
+    _Py_FOR_EACH_TSTATE_END(interp);
+    _PyEval_StartTheWorld(interp);
+    return PyLong_FromSsize_t(count);
+}
+
 #endif
 
 static PyObject *
@@ -3442,6 +3464,7 @@ static PyMethodDef module_functions[] = {
     {"get_tlbc", get_tlbc, METH_O, NULL},
     {"get_tlbc_id", get_tlbc_id, METH_O, NULL},
     {"get_long_lived_total", get_long_lived_total, METH_NOARGS},
+    {"get_qsbr_page_list_length", get_qsbr_page_list_length, METH_NOARGS},
 #endif
 #ifdef _Py_TIER2
     {"uop_symbols_test", _Py_uop_symbols_test, METH_NOARGS},