Store full_page_bytes in mi_heap_t.

Should should avoid memory contention.  Avoid casting *intptr_t
to *Py_ssize_t.  Include large and huge pages in count (promote
eagerly to MI_BIN_FULL).  Add comment noting about abandoned
pages potentially being lost (their byte count never being
subtracted).

Author: nascheme

Include/internal/mimalloc/mimalloc/types.h

@@ -517,9 +517,11 @@ typedef struct mi_abandoned_pool_s {
   // still be read.
   mi_decl_cache_align _Atomic(size_t)           abandoned_readers; // = 0
 
-  // Total bytes (block_size * capacity) of pages currently in MI_BIN_FULL
-  // state whose pool association is this pool.
+#if MI_FULL_PAGE_BYTES
+  // Bytes (block_size * capacity) of full pages currently abandoned to this
+  // pool.
   mi_decl_cache_align _Atomic(intptr_t)         full_page_bytes; // = 0
+#endif
 } mi_abandoned_pool_t;
 
 
@@ -592,6 +594,11 @@ struct mi_heap_s {
   uint8_t               tag;                                 // custom identifier for this heap
   uint8_t               debug_offset;                        // number of bytes to preserve when filling freed or uninitialized memory
   bool                  page_use_qsbr;                       // should freeing pages be delayed using QSBR
+#if MI_FULL_PAGE_BYTES
+  // Bytes (block_size * capacity) of pages currently in MI_BIN_FULL state
+  // owned by this heap.
+  _Atomic(intptr_t)     full_page_bytes;
+#endif
 };
 
 

Include/internal/pycore_gc.h

@@ -337,6 +337,8 @@ extern int _PyGC_VisitStackRef(union _PyStackRef *ref, visitproc visit, void *ar
 #ifdef Py_GIL_DISABLED
 extern void _PyGC_VisitObjectsWorldStopped(PyInterpreterState *interp,
                                            gcvisitobjects_t callback, void *arg);
+// Estimate of bytes allocated by mimalloc.
+PyAPI_FUNC(Py_ssize_t) _PyGC_GetHeapBytes(PyInterpreterState *interp);
 #endif
 
 #ifdef __cplusplus

Include/internal/pycore_mimalloc.h

@@ -36,6 +36,11 @@ typedef enum {
 #  define MI_TSAN 1
 #endif
 
+#ifdef Py_GIL_DISABLED
+// Track full-page byte totals on each mi_heap_t and mi_abandoned_pool_t.
+#  define MI_FULL_PAGE_BYTES 1
+#endif
+
 #ifdef __cplusplus
 extern "C++" {
 #endif

Lib/test/test_gc.py

@@ -1271,8 +1271,58 @@ def test():
         assert_python_ok("-c", code_inside_function)
 
 
-    @unittest.skipUnless(Py_GIL_DISABLED, "requires free-threaded GC")
-    @unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
+
+@unittest.skipUnless(Py_GIL_DISABLED, "requires free-threaded GC")
+@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
+class FreeThreadingTests(unittest.TestCase):
+    # Tests that are specific to the free-threading GC.
+
+    def test_gc_heap_bytes_large_allocs(self):
+        # The free-threaded GC threshold uses _PyGC_GetHeapBytes(), which
+        # sums mimalloc's full_page_bytes counters.  Large/huge pages
+        # (>MI_MEDIUM_OBJ_SIZE_MAX, MI_BIN_HUGE) get eagerly promoted to
+        # MI_BIN_FULL by `_mi_malloc_generic` -- without that, mimalloc
+        # would never count these pages, and a cycle holding a large
+        # buffer would not register as memory pressure.
+        gc.collect()
+        baseline = _testinternalcapi.get_gc_heap_bytes()
+        size = 1 << 20  # 1 MiB
+        k = 5
+        data = [bytearray(size) for _ in range(k)]
+        after_alloc = _testinternalcapi.get_gc_heap_bytes()
+        # All k pages should be counted.  Page size rounds up the request,
+        # so the increase should be at least k * size.
+        self.assertGreaterEqual(after_alloc - baseline, k * size)
+        del data
+        gc.collect()
+        after_free = _testinternalcapi.get_gc_heap_bytes()
+        # Freeing the lone block in each huge page un-fulls it.  Allow some
+        # slop for unrelated allocations triggered by gc.collect().
+        self.assertLess(abs(after_free - baseline), size)
+
+    def test_gc_heap_bytes_many_small_allocs(self):
+        # Filling small pages should also bump the counter.  Small/medium
+        # transitions are lazy (only when a page actually becomes full), so
+        # use enough allocations to fill many pages.
+        gc.collect()
+        baseline = _testinternalcapi.get_gc_heap_bytes()
+        n = 100_000
+        objs = [bytes(4) for i in range(n)]
+        after_alloc = _testinternalcapi.get_gc_heap_bytes()
+        print('small after alloc', baseline, after_alloc)
+        self.assertGreater(after_alloc - baseline, 1 << 20)
+        del objs
+        gc.collect()
+        after_free = _testinternalcapi.get_gc_heap_bytes()
+        print('small after free', baseline, after_free)
+        # Should drop substantially once the pages empty out.
+        self.assertLess(after_free - baseline, (after_alloc - baseline) // 2)
+
+    def test_gc_heap_bytes_nonneg(self):
+        # Counter is intptr_t and only increases or decreases via paired
+        # hooks; it must never go negative.
+        self.assertGreaterEqual(_testinternalcapi.get_gc_heap_bytes(), 0)
+
     def test_tuple_untrack_counts(self):
         # This ensures that the free-threaded GC is counting untracked tuples
         # in the "long_lived_total" count.  This is required to avoid

Modules/_testinternalcapi.c

@@ -2636,6 +2636,12 @@ get_long_lived_total(PyObject *self, PyObject *Py_UNUSED(ignored))
     return PyLong_FromInt64(PyInterpreterState_Get()->gc.long_lived_total);
 }
 
+static PyObject *
+get_gc_heap_bytes(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    return PyLong_FromSsize_t(_PyGC_GetHeapBytes(PyInterpreterState_Get()));
+}
+
 #endif
 
 static PyObject *
@@ -3001,6 +3007,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_gc_heap_bytes", get_gc_heap_bytes, METH_NOARGS},
 #endif
 #ifdef _Py_TIER2
     {"uop_symbols_test", _Py_uop_symbols_test, METH_NOARGS},

Objects/mimalloc/heap.c

@@ -270,6 +270,11 @@ static void mi_heap_reset_pages(mi_heap_t* heap) {
   _mi_memcpy_aligned(&heap->pages, &_mi_heap_empty.pages, sizeof(heap->pages));
   heap->thread_delayed_free = NULL;
   heap->page_count = 0;
+#if MI_FULL_PAGE_BYTES
+  // All pages have been removed (destroyed, or transferred via
+  // mi_heap_absorb which already moved the bytes to the destination heap).
+  mi_atomic_store_relaxed(&heap->full_page_bytes, (intptr_t)0);
+#endif
 }
 
 // called from `mi_heap_destroy` and `mi_heap_delete` to free the internal heap resources.
@@ -427,6 +432,14 @@ static void mi_heap_absorb(mi_heap_t* heap, mi_heap_t* from) {
   }
   mi_assert_internal(from->page_count == 0);
 
+#if MI_FULL_PAGE_BYTES
+  // The page-state hooks didn't fire for these transfers, so move the
+  // full_page_bytes accounting in bulk.  mi_heap_reset_pages(from) below
+  // will zero `from->full_page_bytes`.
+  intptr_t bytes = mi_atomic_load_relaxed(&from->full_page_bytes);
+  mi_atomic_addi(&heap->full_page_bytes, bytes);
+#endif
+
   // and do outstanding delayed frees in the `from` heap
   // note: be careful here as the `heap` field in all those pages no longer point to `from`,
   // turns out to be ok as `_mi_heap_delayed_free` only visits the list and calls a

Objects/mimalloc/init.c

@@ -104,7 +104,10 @@ mi_decl_cache_align const mi_heap_t _mi_heap_empty = {
   false,
   0,
   0,
-  0
+  0,
+#if MI_FULL_PAGE_BYTES
+  MI_ATOMIC_VAR_INIT(0),  // full_page_bytes
+#endif
 };
 
 #define tld_empty_stats  ((mi_stats_t*)((uint8_t*)&tld_empty + offsetof(mi_tld_t,stats)))

Objects/mimalloc/page-queue.c

@@ -151,7 +151,7 @@ static mi_page_queue_t* mi_heap_page_queue_of(mi_heap_t* heap, const mi_page_t*
   uint8_t bin = (mi_page_is_in_full(page) ? MI_BIN_FULL : mi_bin(page->xblock_size));
   mi_assert_internal(bin <= MI_BIN_FULL);
   mi_page_queue_t* pq = &heap->pages[bin];
-  mi_assert_internal(mi_page_is_in_full(page) || page->xblock_size == pq->block_size);
+  mi_assert_internal(bin >= MI_BIN_HUGE || page->xblock_size == pq->block_size);
   return pq;
 }
 
@@ -264,7 +264,9 @@ static void mi_page_queue_enqueue_from(mi_page_queue_t* to, mi_page_queue_t* fro
                      (page->xblock_size == to->block_size && mi_page_queue_is_full(from)) ||
                      (page->xblock_size == from->block_size && mi_page_queue_is_full(to)) ||
                      (page->xblock_size > MI_LARGE_OBJ_SIZE_MAX && mi_page_queue_is_huge(to)) ||
-                     (page->xblock_size > MI_LARGE_OBJ_SIZE_MAX && mi_page_queue_is_full(to)));
+                     (page->xblock_size > MI_LARGE_OBJ_SIZE_MAX && mi_page_queue_is_full(to)) ||
+                     (mi_page_queue_is_huge(from) && mi_page_queue_is_full(to)) ||
+                     (mi_page_queue_is_full(from) && mi_page_queue_is_huge(to)));
 
   mi_heap_t* heap = mi_page_heap(page);
   if (page->prev != NULL) page->prev->next = page->next;

Objects/mimalloc/page.c

@@ -255,6 +255,78 @@ void _mi_page_free_collect(mi_page_t* page, bool force) {
   mi_assert_internal(!force || page->local_free == NULL);
 }
 
+/* -----------------------------------------------------------
+  Full-page byte accounting (MI_FULL_PAGE_BYTES)
+
+  Maintain `mi_heap_t.full_page_bytes` (bytes of MI_BIN_FULL pages owned by
+  the heap) and `mi_abandoned_pool_t.full_page_bytes` (bytes of MI_BIN_FULL
+  pages currently abandoned to that pool).  Page weight is
+  `mi_page_block_size(page) * page->capacity`.  Capacity is stable while a
+  page is in the full queue (`mi_page_extend_free` only runs on non-full
+  queues), so inc and dec see the same value.
+
+  State machine:
+    to-full        : heap += size
+    from-full      : heap -= size
+    abandon a full : heap -= size; pool += size
+    reclaim a full : pool -= size; heap += size
+    free a full    : heap -= size
+
+  The in_full bit is unconditionally cleared by `mi_page_queue_remove`, so
+  `_mi_page_abandon` re-sets it after queue_remove to preserve the "this
+  page's bytes were transferred to the pool" marker through abandonment.
+  `_mi_page_reclaim` then routes such pages straight to MI_BIN_FULL, so
+  `mi_page_queue_push` keeps the bit set; subsequent unfull/free fires the
+  matching dec.
+
+  Large/huge pages (block_size > MI_MEDIUM_OBJ_SIZE_MAX) are 1-block pages
+  in MI_BIN_HUGE; mimalloc never walks that queue on a subsequent alloc, so
+  it would never call `mi_page_to_full` on them.  `_mi_malloc_generic`
+  therefore eagerly calls `mi_page_to_full` on a freshly-filled huge page
+  (see the MI_FULL_PAGE_BYTES block at the bottom of that function).
+  Inc/dec then proceed identically to small/medium pages.
+
+  Known minor leak: if a page abandoned-while-full later becomes empty and
+  then freed, the +size we added on abandon is never subtracted.
+----------------------------------------------------------- */
+
+#if MI_FULL_PAGE_BYTES
+static inline intptr_t mi_page_full_size(mi_page_t* page) {
+  return (intptr_t)(mi_page_block_size(page) * (size_t)page->capacity);
+}
+
+static void mi_page_full_inc(mi_page_t* page) {
+  mi_atomic_addi(&mi_page_heap(page)->full_page_bytes, mi_page_full_size(page));
+}
+
+static void mi_page_full_dec(mi_page_t* page) {
+  mi_atomic_addi(&mi_page_heap(page)->full_page_bytes, -mi_page_full_size(page));
+}
+
+// Called from `_mi_page_abandon` *before* the page's heap pointer is cleared.
+// Transfers the page's bytes from its heap to the pool that will own the
+// abandoned page.  No-op if the page is not currently in MI_BIN_FULL.
+static void mi_page_full_abandon(mi_page_t* page) {
+  if (!mi_page_is_in_full(page)) return;
+  intptr_t bytes = mi_page_full_size(page);
+  mi_heap_t* heap = mi_page_heap(page);
+  mi_atomic_addi(&heap->full_page_bytes, -bytes);
+  mi_atomic_addi(&heap->tld->segments.abandoned->full_page_bytes, bytes);
+}
+
+// Called from `_mi_page_reclaim` when a page abandoned-while-full is
+// returning to a heap.  in_full=true here means "this page's bytes are
+// currently in the pool counter from abandon".  Transfer them: pool -= size,
+// new-heap += size.  The caller routes the page directly into MI_BIN_FULL,
+// so the in_full bit (and matching dec hook on free/unfull) survives.
+static void mi_page_full_reclaim(mi_page_t* page) {
+  if (!mi_page_is_in_full(page)) return;
+  intptr_t bytes = mi_page_full_size(page);
+  mi_heap_t* heap = mi_page_heap(page);
+  mi_atomic_addi(&heap->tld->segments.abandoned->full_page_bytes, -bytes);
+  mi_atomic_addi(&heap->full_page_bytes, bytes);
+}
+#endif // MI_FULL_PAGE_BYTES
 
 
 /* -----------------------------------------------------------
@@ -271,8 +343,24 @@ void _mi_page_reclaim(mi_heap_t* heap, mi_page_t* page) {
   mi_assert_internal(_mi_page_segment(page)->kind != MI_SEGMENT_HUGE);
   #endif
 
-  // TODO: push on full queue immediately if it is full?
-  mi_page_queue_t* pq = mi_page_queue(heap, mi_page_block_size(page));
+  mi_page_queue_t* pq;
+#if MI_FULL_PAGE_BYTES
+  // If the page was abandoned full (in_full preserved as marker), route
+  // it directly to MI_BIN_FULL.  Pushing to the size-bucket queue would
+  // rely on a later alloc walking that queue to promote it via
+  // mi_page_to_full -- which happens for small/medium bins but never for
+  // MI_BIN_HUGE, so a reclaimed full huge page would otherwise leave the
+  // pool counter without re-crediting any heap.  mi_page_full_reclaim
+  // does the pool-to-heap transfer.
+  if (mi_page_is_in_full(page)) {
+    pq = &heap->pages[MI_BIN_FULL];
+  } else {
+    pq = mi_page_queue(heap, mi_page_block_size(page));
+  }
+  mi_page_full_reclaim(page);
+#else
+  pq = mi_page_queue(heap, mi_page_block_size(page));
+#endif
   mi_page_queue_push(heap, pq, page);
   _PyMem_mi_page_reclaimed(page);
   mi_assert_expensive(_mi_page_is_valid(page));
@@ -360,8 +448,8 @@ void _mi_page_unfull(mi_page_t* page) {
   mi_assert_internal(mi_page_is_in_full(page));
   if (!mi_page_is_in_full(page)) return;
 
-#ifdef Py_GIL_DISABLED
-  _PyMem_mi_page_full_dec(page);
+#if MI_FULL_PAGE_BYTES
+  mi_page_full_dec(page);
 #endif
 
   mi_heap_t* heap = mi_page_heap(page);
@@ -378,8 +466,8 @@ static void mi_page_to_full(mi_page_t* page, mi_page_queue_t* pq) {
   mi_assert_internal(!mi_page_is_in_full(page));
 
   if (mi_page_is_in_full(page)) return;
-#ifdef Py_GIL_DISABLED
-  _PyMem_mi_page_full_inc(page);
+#if MI_FULL_PAGE_BYTES
+  mi_page_full_inc(page);
 #endif
   mi_page_queue_enqueue_from(&mi_page_heap(page)->pages[MI_BIN_FULL], pq, page);
   _mi_page_free_collect(page,false);  // try to collect right away in case another thread freed just before MI_USE_DELAYED_FREE was set
@@ -398,6 +486,13 @@ void _mi_page_abandon(mi_page_t* page, mi_page_queue_t* pq) {
 
   mi_heap_t* pheap = mi_page_heap(page);
 
+#if MI_FULL_PAGE_BYTES
+  // Capture in_full while the heap pointer is still valid; transfer the
+  // bytes from heap counter to pool counter.  Must run before
+  // mi_page_queue_remove, which clears the in_full bit unconditionally.
+  bool was_in_full = mi_page_is_in_full(page);
+  mi_page_full_abandon(page);
+#endif
 #ifdef Py_GIL_DISABLED
   if (page->qsbr_node.next != NULL) {
     // remove from QSBR queue, but keep the goal
@@ -413,6 +508,15 @@ void _mi_page_abandon(mi_page_t* page, mi_page_queue_t* pq) {
   mi_assert_internal(mi_page_thread_free_flag(page)==MI_NEVER_DELAYED_FREE);
   mi_page_set_heap(page, NULL);
 
+#if MI_FULL_PAGE_BYTES
+  // Preserve the in_full marker through abandonment so `_mi_page_reclaim`'s
+  // `mi_page_full_reclaim` call can transfer the bytes back to the
+  // reclaiming heap.  Nothing reads in_full on a heap-less page.
+  if (was_in_full) {
+    mi_page_set_in_full(page, true);
+  }
+#endif
+
 #if (MI_DEBUG>1) && !MI_TRACK_ENABLED
   // check there are no references left..
   for (mi_block_t* block = (mi_block_t*)pheap->thread_delayed_free; block != NULL; block = mi_block_nextx(pheap, block, pheap->keys)) {
@@ -442,12 +546,16 @@ void _mi_page_free(mi_page_t* page, mi_page_queue_t* pq, bool force) {
 #ifdef Py_GIL_DISABLED
   mi_assert_internal(page->qsbr_goal == 0);
   mi_assert_internal(page->qsbr_node.next == NULL);
-  // Defensive: a full page whose last block is freed locally goes through
+#endif
+#if MI_FULL_PAGE_BYTES
+  // A full page whose last block is freed locally goes through
   // _mi_page_retire -> _PyMem_mi_page_maybe_free -> _mi_page_free without
-  // ever calling _mi_page_unfull, so the per-thread full-page counter must
-  // be decremented here to maintain the invariant.
+  // ever calling _mi_page_unfull, so the heap's full_page_bytes counter
+  // must be decremented here to maintain the invariant.  `heap` is non-NULL
+  // for any page reaching _mi_page_free (abandoned pages take the
+  // segment-level cleanup path instead).
   if (mi_page_is_in_full(page)) {
-    _PyMem_mi_page_full_dec(page);
+    mi_page_full_dec(page);
   }
 #endif
 
@@ -977,14 +1085,28 @@ void* _mi_malloc_generic(mi_heap_t* heap, size_t size, bool zero, size_t huge_al
   mi_assert_internal(mi_page_block_size(page) >= size);
 
   // and try again, this time succeeding! (i.e. this should never recurse through _mi_page_malloc)
+  void* p;
   if mi_unlikely(zero && page->xblock_size == 0) {
     // note: we cannot call _mi_page_malloc with zeroing for huge blocks; we zero it afterwards in that case.
-    void* p = _mi_page_malloc(heap, page, size, false);
+    p = _mi_page_malloc(heap, page, size, false);
     mi_assert_internal(p != NULL);
     _mi_memzero_aligned(p, mi_page_usable_block_size(page));
-    return p;
   }
   else {
-    return _mi_page_malloc(heap, page, size, zero);
+    p = _mi_page_malloc(heap, page, size, zero);
+  }
+
+#if MI_FULL_PAGE_BYTES
+  // Eagerly promote a freshly-filled huge page (1 block per page, in
+  // MI_BIN_HUGE) to MI_BIN_FULL so its bytes get counted.  See the
+  // "Full-page byte accounting" comment block above.
+  if (p != NULL && !mi_page_immediate_available(page)) {
+    mi_page_queue_t* page_pq = mi_page_queue_of(page);
+    if (mi_page_queue_is_huge(page_pq) && !mi_page_is_in_full(page)) {
+      mi_page_to_full(page, page_pq);
+    }
   }
+#endif
+
+  return p;
 }