From 31633f4473966b3bcd470440bab7f348711be48f Mon Sep 17 00:00:00 2001
From: Sam Gross <colesbury@gmail.com>
Date: Fri, 9 Feb 2024 09:23:12 -0500
Subject: [PATCH] gh-115184: Fix refleak tracking issues in free-threaded build
 (#115188)

Fixes a few issues related to refleak tracking in the free-threaded build:

- Count blocks in abandoned segments
- Call `_mi_page_free_collect` earlier during heap traversal in order to get an accurate count of blocks in use.
- Add missing refcount tracking in `_Py_DecRefSharedDebug` and `_Py_ExplicitMergeRefcount`.
- Pause threads in  `get_num_global_allocated_blocks` to ensure that traversing the mimalloc heaps is safe.
---
 Objects/mimalloc/heap.c |  2 +-
 Objects/object.c        | 11 +++++++----
 Objects/obmalloc.c      |  9 ++++++++-
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/Objects/mimalloc/heap.c b/Objects/mimalloc/heap.c
index 164b28f0fab..154dad0b128 100644
--- a/Objects/mimalloc/heap.c
+++ b/Objects/mimalloc/heap.c
@@ -538,7 +538,6 @@ bool _mi_heap_area_visit_blocks(const mi_heap_area_t* area, mi_page_t *page, mi_
   mi_assert(page != NULL);
   if (page == NULL) return true;
 
-  _mi_page_free_collect(page,true);
   mi_assert_internal(page->local_free == NULL);
   if (page->used == 0) return true;
 
@@ -635,6 +634,7 @@ bool _mi_heap_area_visit_blocks(const mi_heap_area_t* area, mi_page_t *page, mi_
 typedef bool (mi_heap_area_visit_fun)(const mi_heap_t* heap, const mi_heap_area_ex_t* area, void* arg);
 
 void _mi_heap_area_init(mi_heap_area_t* area, mi_page_t* page) {
+  _mi_page_free_collect(page,true);
   const size_t bsize = mi_page_block_size(page);
   const size_t ubsize = mi_page_usable_block_size(page);
   area->reserved = page->reserved * bsize;
diff --git a/Objects/object.c b/Objects/object.c
index bbf7f98ae3d..37a4b7a417e 100644
--- a/Objects/object.c
+++ b/Objects/object.c
@@ -346,6 +346,9 @@ _Py_DecRefSharedDebug(PyObject *o, const char *filename, int lineno)
     if (should_queue) {
         // TODO: the inter-thread queue is not yet implemented. For now,
         // we just merge the refcount here.
+#ifdef Py_REF_DEBUG
+        _Py_IncRefTotal(_PyInterpreterState_GET());
+#endif
         Py_ssize_t refcount = _Py_ExplicitMergeRefcount(o, -1);
         if (refcount == 0) {
             _Py_Dealloc(o);
@@ -399,10 +402,6 @@ _Py_ExplicitMergeRefcount(PyObject *op, Py_ssize_t extra)
     Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&op->ob_ref_shared);
     do {
         refcnt = Py_ARITHMETIC_RIGHT_SHIFT(Py_ssize_t, shared, _Py_REF_SHARED_SHIFT);
-        if (_Py_REF_IS_MERGED(shared)) {
-            return refcnt;
-        }
-
         refcnt += (Py_ssize_t)op->ob_ref_local;
         refcnt += extra;
 
@@ -410,6 +409,10 @@ _Py_ExplicitMergeRefcount(PyObject *op, Py_ssize_t extra)
     } while (!_Py_atomic_compare_exchange_ssize(&op->ob_ref_shared,
                                                 &shared, new_shared));
 
+#ifdef Py_REF_DEBUG
+    _Py_AddRefTotal(_PyInterpreterState_GET(), extra);
+#endif
+
     _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, 0);
     _Py_atomic_store_uintptr_relaxed(&op->ob_tid, 0);
     return refcnt;
diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c
index bea4ea85332..6a12c3dca38 100644
--- a/Objects/obmalloc.c
+++ b/Objects/obmalloc.c
@@ -1073,7 +1073,12 @@ get_mimalloc_allocated_blocks(PyInterpreterState *interp)
             mi_heap_visit_blocks(heap, false, &count_blocks, &allocated_blocks);
         }
     }
-    // TODO(sgross): count blocks in abandoned segments.
+
+    mi_abandoned_pool_t *pool = &interp->mimalloc.abandoned_pool;
+    for (uint8_t tag = 0; tag < _Py_MIMALLOC_HEAP_COUNT; tag++) {
+        _mi_abandoned_pool_visit_blocks(pool, tag, false, &count_blocks,
+                                        &allocated_blocks);
+    }
 #else
     // TODO(sgross): this only counts the current thread's blocks.
     mi_heap_t *heap = mi_heap_get_default();
@@ -1189,6 +1194,7 @@ get_num_global_allocated_blocks(_PyRuntimeState *runtime)
         }
     }
     else {
+        _PyEval_StopTheWorldAll(&_PyRuntime);
         HEAD_LOCK(runtime);
         PyInterpreterState *interp = PyInterpreterState_Head();
         assert(interp != NULL);
@@ -1208,6 +1214,7 @@ get_num_global_allocated_blocks(_PyRuntimeState *runtime)
             }
         }
         HEAD_UNLOCK(runtime);
+        _PyEval_StartTheWorldAll(&_PyRuntime);
 #ifdef Py_DEBUG
         assert(got_main);
 #endif