From 7673aa2517fa44080c51df8833aa7e79dad12ea8 Mon Sep 17 00:00:00 2001
From: daanx <daan@effp.org>
Date: Mon, 25 Nov 2024 18:41:57 -0800
Subject: [PATCH] ensure forced abandoned pages can be accessed after free

---
 include/mimalloc/types.h |  3 ++-
 src/page.c               |  6 ++----
 src/segment.c            | 38 ++++++++++++++++++++++----------------
 3 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/include/mimalloc/types.h b/include/mimalloc/types.h
index f7bca137..44074450 100644
--- a/include/mimalloc/types.h
+++ b/include/mimalloc/types.h
@@ -416,7 +416,8 @@ typedef struct mi_segment_s {
   // segment fields
   struct mi_segment_s* next;             // must be the first (non-constant) segment field  -- see `segment.c:segment_init`
   struct mi_segment_s* prev;
-  bool                 was_reclaimed;    // true if it was reclaimed (used to limit on-free reclamation)
+  bool                 was_reclaimed;    // true if it was reclaimed (used to limit reclaim-on-free reclamation)
+  bool                 dont_free;        // can be temporarily true to ensure the segment is not freed
 
   size_t               abandoned;        // abandoned pages (i.e. the original owning thread stopped) (`abandoned <= used`)
   size_t               abandoned_visits; // count how often this segment is visited for reclaiming (to force reclaim if it is too long)
diff --git a/src/page.c b/src/page.c
index 43ac7c4e..c681d6d0 100644
--- a/src/page.c
+++ b/src/page.c
@@ -411,10 +411,8 @@ void _mi_page_force_abandon(mi_page_t* page) {
 
   // ensure this page is no longer in the heap delayed free list
   _mi_heap_delayed_free_all(heap);
-  // TODO: can we still access the page as it may have been
-  // freed and the memory decommitted?
-  // A way around this is to explicitly unlink this page from
-  // the heap delayed free list.
+  // We can still access the page meta-info even if it is freed as we ensure 
+  // in `mi_segment_force_abandon` that the segment is not freed (yet)
   if (page->capacity == 0) return; // it may have been freed now
 
   // and now unlink it from the page queue and abandon (or free)
diff --git a/src/segment.c b/src/segment.c
index 16764da8..74abcdbc 100644
--- a/src/segment.c
+++ b/src/segment.c
@@ -652,6 +652,10 @@ static mi_segment_t* mi_segment_alloc(size_t required, mi_page_kind_t page_kind,
 static void mi_segment_free(mi_segment_t* segment, bool force, mi_segments_tld_t* tld) {
   MI_UNUSED(force);
   mi_assert(segment != NULL);
+  
+  // in `mi_segment_force_abandon` we set this to true to ensure the segment's memory stays valid
+  if (segment->dont_free) return;
+
   // don't purge as we are freeing now
   mi_segment_remove_all_purges(segment, false /* don't force as we are about to free */, tld);
   mi_segment_remove_from_free_queue(segment, tld);
@@ -1056,32 +1060,34 @@ static mi_segment_t* mi_segment_try_reclaim(mi_heap_t* heap, size_t block_size,
 static void mi_segment_force_abandon(mi_segment_t* segment, mi_segments_tld_t* tld)
 {
   mi_assert_internal(segment->abandoned < segment->used);
+  mi_assert_internal(!segment->dont_free);
+  
+  // ensure the segment does not get free'd underneath us (so we can check if a page has been freed in `mi_page_force_abandon`)
+  segment->dont_free = true;
 
   // for all pages
   for (size_t i = 0; i < segment->capacity; i++) {
     mi_page_t* page = &segment->pages[i];
     if (page->segment_in_use) {
-      // ensure used count is up to date and collect potential concurrent frees
-      _mi_page_free_collect(page, false);
-      {
-        // abandon the page if it is still in-use (this will free it if possible as well)
-        mi_assert_internal(segment->used > 0);
-        if (segment->used == segment->abandoned+1) {
-          // the last page.. abandon and return as the segment will be abandoned after this
-          // and we should no longer access it.
-          _mi_page_force_abandon(page);
-          return;
-        }
-        else {
-          // abandon and continue
-          _mi_page_force_abandon(page);
-        }
+      // abandon the page if it is still in-use (this will free the page if possible as well (but not our segment))
+      mi_assert_internal(segment->used > 0);
+      if (segment->used == segment->abandoned+1) {
+        // the last page.. abandon and return as the segment will be abandoned after this
+        // and we should no longer access it.
+        segment->dont_free = false;
+        _mi_page_force_abandon(page);
+        return;
+      }
+      else {
+        // abandon and continue
+        _mi_page_force_abandon(page);
       }
     }
   }
+  segment->dont_free = false;
   mi_assert(segment->used == segment->abandoned);
   mi_assert(segment->used == 0);
-  if (segment->used == 0) {
+  if (segment->used == 0) {  // paranoia
     // all free now
     mi_segment_free(segment, false, tld);
   }