Use the upstream cursor API to visit abandoned segments

Author: corona10

Include/internal/mimalloc/mimalloc/internal.h

@@ -249,7 +249,6 @@ void       _mi_abandoned_reclaim_all(mi_heap_t* heap, mi_segments_tld_t* tld);
 void       _mi_abandoned_collect(mi_heap_t* heap, bool force, mi_segments_tld_t* tld);
 bool       _mi_segment_attempt_reclaim(mi_heap_t* heap, mi_segment_t* segment);
 bool       _mi_segment_visit_blocks(mi_segment_t* segment, int heap_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg);
-bool       _mi_abandoned_visit_blocks_unlocked(mi_subproc_t* subproc, int heap_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg);  // CPython
 
 // "page.c"
 void*       _mi_malloc_generic(mi_heap_t* heap, size_t size, bool zero, size_t huge_alignment, mi_page_t** ppage)  mi_attr_noexcept mi_attr_malloc;

Include/internal/pycore_mimalloc.h

@@ -66,6 +66,27 @@ struct _mimalloc_thread_state {
     int initialized;
     struct llist_node page_list;
 };
+
+// Visit the blocks in the abandoned segments (from exited threads) of a
+// mimalloc sub-process. Same as mi_abandoned_visit_blocks() but without
+// requiring mi_option_visit_abandoned: the caller must guarantee that no
+// other thread abandons or reclaims segments concurrently (STW).
+static inline bool
+_PyMem_mi_visit_abandoned_blocks(mi_subproc_t *subproc, int heap_tag,
+                                 bool visit_blocks,
+                                 mi_block_visit_fun *visitor, void *arg)
+{
+    mi_arena_field_cursor_t current;
+    _mi_arena_field_cursor_init(NULL, subproc, true /* visit all */, &current);
+    mi_segment_t *segment;
+    bool ok = true;
+    while (ok && (segment = _mi_arena_segment_clear_abandoned_next(&current)) != NULL) {
+        ok = _mi_segment_visit_blocks(segment, heap_tag, visit_blocks, visitor, arg);
+        _mi_arena_segment_mark_abandoned(segment);
+    }
+    _mi_arena_field_cursor_done(&current);
+    return ok;
+}
 #endif
 
 #endif // Py_INTERNAL_MIMALLOC_H

Objects/mimalloc/arena-abandon.c

@@ -345,41 +345,3 @@ bool mi_abandoned_visit_blocks(mi_subproc_id_t subproc_id, int heap_tag, bool vi
   return ok;
 }
 
-
-/* -----------------------------------------------------------
-  CPython: visit all abandoned segments of a sub-process without
-  taking any locks or claiming the segments.
-
-  Unlike `mi_abandoned_visit_blocks`, this does not require
-  `mi_option_visit_abandoned` to be enabled. The caller must
-  guarantee that no other thread abandons or reclaims segments
-  concurrently (in CPython: a stop-the-world pause), and the
-  visitor must not allocate or free mimalloc managed memory.
------------------------------------------------------------ */
-
-bool _mi_abandoned_visit_blocks_unlocked(mi_subproc_t* subproc, int heap_tag, bool visit_blocks, mi_block_visit_fun* visitor, void* arg) {
-  // segments in arenas are tracked in the `blocks_abandoned` bitmaps
-  const size_t max_arena = mi_arena_get_count();
-  for (size_t arena_idx = 0; arena_idx < max_arena; arena_idx++) {
-    mi_arena_t* arena = mi_arena_from_index(arena_idx);
-    if (arena == NULL) continue;
-    for (size_t field_idx = 0; field_idx < arena->field_count; field_idx++) {
-      size_t field = mi_atomic_load_relaxed(&arena->blocks_abandoned[field_idx]);
-      while (field != 0) {
-        const size_t bit_idx = mi_ctz(field);
-        field &= ~((size_t)1 << bit_idx);
-        const mi_bitmap_index_t bitmap_idx = mi_bitmap_index_create(field_idx, bit_idx);
-        mi_segment_t* const segment = (mi_segment_t*)mi_arena_block_start(arena, bitmap_idx);
-        mi_assert_internal(mi_atomic_load_relaxed(&segment->thread_id) == 0);
-        if (segment->subproc != subproc) continue;
-        if (!_mi_segment_visit_blocks(segment, heap_tag, visit_blocks, visitor, arg)) return false;
-      }
-    }
-  }
-  // segments outside arenas (OS allocated) are kept in a per sub-process list
-  for (mi_segment_t* segment = subproc->abandoned_os_list; segment != NULL; segment = segment->abandoned_os_next) {
-    mi_assert_internal(segment->subproc == subproc);
-    if (!_mi_segment_visit_blocks(segment, heap_tag, visit_blocks, visitor, arg)) return false;
-  }
-  return true;
-}

Objects/obmalloc.c

@@ -1825,8 +1825,8 @@ get_mimalloc_allocated_blocks(PyInterpreterState *interp)
 
     mi_subproc_t *subproc = interp->mimalloc.subproc;
     for (int tag = 0; tag < _Py_MIMALLOC_HEAP_COUNT; tag++) {
-        _mi_abandoned_visit_blocks_unlocked(subproc, tag, false, &count_blocks,
-                                            &allocated_blocks);
+        _PyMem_mi_visit_abandoned_blocks(subproc, tag, false, &count_blocks,
+                                         &allocated_blocks);
     }
 #else
     // TODO(sgross): this only counts the current thread's blocks.

Python/gc_free_threading.c

@@ -403,13 +403,13 @@ gc_visit_heaps_lock_held(PyInterpreterState *interp, mi_block_visit_fun *visitor
     // per-interpreter mimalloc sub-process.
     mi_subproc_t *subproc = interp->mimalloc.subproc;
     arg->offset = offset_base;
-    if (!_mi_abandoned_visit_blocks_unlocked(subproc, _Py_MIMALLOC_HEAP_GC,
-                                             true, visitor, arg)) {
+    if (!_PyMem_mi_visit_abandoned_blocks(subproc, _Py_MIMALLOC_HEAP_GC,
+                                          true, visitor, arg)) {
         return -1;
     }
     arg->offset = offset_pre;
-    if (!_mi_abandoned_visit_blocks_unlocked(subproc, _Py_MIMALLOC_HEAP_GC_PRE,
-                                             true, visitor, arg)) {
+    if (!_PyMem_mi_visit_abandoned_blocks(subproc, _Py_MIMALLOC_HEAP_GC_PRE,
+                                          true, visitor, arg)) {
         return -1;
     }
     return 0;