From 44e47dfba55a825ce7c340b1460cb493165353c9 Mon Sep 17 00:00:00 2001
From: Sam Gross <colesbury@gmail.com>
Date: Mon, 15 Jan 2024 10:37:36 -0500
Subject: [PATCH] gh-112532: Fix memory block count for free-threaded build
 (gh-113995)

This fixes `_PyInterpreterState_GetAllocatedBlocks()` and
`_Py_GetGlobalAllocatedBlocks()` in the free-threaded builds. The
gh-113263 change that introduced multiple mimalloc heaps per-thread
broke the logic for counting the number of allocated blocks. For subtle
reasons, this led to reported reference count leaks in the refleaks
buildbots.
---
 Objects/obmalloc.c | 45 +++++++++++++++++++++++++++------------------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c
index 883adcb1c19..88ad305ad77 100644
--- a/Objects/obmalloc.c
+++ b/Objects/obmalloc.c
@@ -626,12 +626,16 @@ _PyMem_PymallocEnabled(void)
 static int
 _PyMem_MimallocEnabled(void)
 {
+#ifdef Py_GIL_DISABLED
+    return 1;
+#else
     if (_PyMem_DebugEnabled()) {
         return (_PyMem_Debug.obj.alloc.malloc == _PyObject_MiMalloc);
     }
     else {
         return (_PyObject.malloc == _PyObject_MiMalloc);
     }
+#endif
 }
 #endif  // WITH_MIMALLOC
 
@@ -1041,20 +1045,35 @@ static bool count_blocks(
     *(size_t *)allocated_blocks += area->used;
     return 1;
 }
+
+static Py_ssize_t
+get_mimalloc_allocated_blocks(PyInterpreterState *interp)
+{
+    size_t allocated_blocks = 0;
+#ifdef Py_GIL_DISABLED
+    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++) {
+            mi_heap_t *heap = &tstate->mimalloc.heaps[i];
+            mi_heap_visit_blocks(heap, false, &count_blocks, &allocated_blocks);
+        }
+    }
+    // TODO(sgross): count blocks in abandoned segments.
+#else
+    // TODO(sgross): this only counts the current thread's blocks.
+    mi_heap_t *heap = mi_heap_get_default();
+    mi_heap_visit_blocks(heap, false, &count_blocks, &allocated_blocks);
+#endif
+    return allocated_blocks;
+}
 #endif
 
 Py_ssize_t
 _PyInterpreterState_GetAllocatedBlocks(PyInterpreterState *interp)
 {
 #ifdef WITH_MIMALLOC
-    // TODO(sgross): this only counts the current thread's blocks.
     if (_PyMem_MimallocEnabled()) {
-        size_t allocated_blocks = 0;
-
-        mi_heap_t *heap = mi_heap_get_default();
-        mi_heap_visit_blocks(heap, false, &count_blocks, &allocated_blocks);
-
-        return allocated_blocks;
+        return get_mimalloc_allocated_blocks(interp);
     }
 #endif
 
@@ -1105,7 +1124,7 @@ _PyInterpreterState_FinalizeAllocatedBlocks(PyInterpreterState *interp)
 
 static Py_ssize_t get_num_global_allocated_blocks(_PyRuntimeState *);
 
-/* We preserve the number of blockss leaked during runtime finalization,
+/* We preserve the number of blocks leaked during runtime finalization,
    so they can be reported if the runtime is initialized again. */
 // XXX We don't lose any information by dropping this,
 // so we should consider doing so.
@@ -1121,16 +1140,6 @@ _Py_FinalizeAllocatedBlocks(_PyRuntimeState *runtime)
 static Py_ssize_t
 get_num_global_allocated_blocks(_PyRuntimeState *runtime)
 {
-#ifdef WITH_MIMALLOC
-    if (_PyMem_MimallocEnabled()) {
-        size_t allocated_blocks = 0;
-
-        mi_heap_t *heap = mi_heap_get_default();
-        mi_heap_visit_blocks(heap, false, &count_blocks, &allocated_blocks);
-
-        return allocated_blocks;
-    }
-#endif
     Py_ssize_t total = 0;
     if (_PyRuntimeState_GetFinalizing(runtime) != NULL) {
         PyInterpreterState *interp = _PyInterpreterState_Main();