Snapshot support on top of the upstream sync

include/mimalloc.h

@@ -367,6 +367,11 @@ typedef struct mi_heap_area_s {
 typedef bool (mi_cdecl mi_block_visit_fun)(const mi_heap_t* heap, const mi_heap_area_t* area, void* block, size_t block_size, void* arg);
 
 mi_decl_export bool   mi_heap_visit_blocks(mi_heap_t* heap, bool visit_blocks, mi_block_visit_fun* visitor, void* arg);
+mi_decl_export void mi_os_hint_floor(void* floor) mi_attr_noexcept;
+mi_decl_export bool mi_startup_snapshot_hints_enabled(void) mi_attr_noexcept;   // deterministic placement is on for this process (snapshot-capable executable, or MIMALLOC_DETERMINISTIC_HINT=1)
+mi_decl_export void mi_arenas_seal_existing(void) mi_attr_noexcept; // snapshot restore: no new allocations inside pre-existing (snapshot) arenas
+mi_decl_export void mi_arenas_freeze_pages(void) mi_attr_noexcept;   // snapshot build: pages written into the snapshot are never freed into again
+mi_decl_export void mi_arenas_visit_free_ranges(mi_heap_t* heap, void (*visit)(void* start, size_t size, void* arg), void* arg) mi_attr_noexcept;  // visit maximal runs of arena slices that belong to no page
 mi_decl_export bool   mi_heap_visit_abandoned_blocks(mi_heap_t* heap, bool visit_blocks, mi_block_visit_fun* visitor, void* arg);
 
 
@@ -443,6 +448,8 @@ mi_decl_export bool mi_subproc_visit_heaps(mi_subproc_id_t subproc, mi_heap_visi
 
 struct mi_theap_s;
 typedef struct mi_theap_s mi_theap_t;
+mi_decl_export void mi_theap_freeze(mi_theap_t* theap) mi_attr_noexcept;  // snapshot restore: never collect/purge/idle-sweep this theap again
+mi_decl_export void mi_theap_adopt_current_thread(mi_theap_t* theap) mi_attr_noexcept;  // snapshot restore: this theap's thread state now belongs to the calling thread
 
 mi_decl_export mi_theap_t* mi_heap_theap(mi_heap_t* heap);
 mi_decl_export mi_theap_t* mi_theap_set_default(mi_theap_t* theap);

include/mimalloc/internal.h

@@ -917,7 +917,7 @@ static inline bool mi_page_all_free(const mi_page_t* page) {
 // Page hole purging  (see the "Page hole purging" section in `page.c`)
 // ------------------------------------------------------
 
-void          _mi_page_purge_holes(mi_page_t* page);
+void          _mi_page_purge_holes(mi_page_t* page, mi_tld_t* tld);   // `tld`: the one being swept (owns the sweep state)
 void          _mi_page_purged_reset(mi_page_t* page);
 bool          _mi_page_unpurge_run(mi_page_t* page);
 void          _mi_page_unpurge_all(mi_page_t* page);
@@ -926,12 +926,11 @@ void          _mi_page_unpurge_unformed_upto(mi_page_t* page, uintptr_t end);
 size_t        _mi_page_unformed_purged_bytes(const mi_page_t* page);            // the bytes of this page's unformed tail that are discarded right now
 bool          _mi_page_purge_os_page_blocks(size_t os_page_size, size_t block_size, uintptr_t page_start,
                                             size_t capacity, size_t k, size_t* first, size_t* last);
-bool          _mi_page_purge_holes_in_progress(void);
 void          _mi_page_holes_count_page_freed(void);
 void          _mi_page_holes_count_ineligible(const mi_page_t* page);
 void          _mi_page_holes_reset_ineligible(void);
-void          _mi_page_purge_holes_begin(void);
-void          _mi_page_purge_holes_end(void);
+void          _mi_page_purge_holes_begin(mi_tld_t* tld);
+void          _mi_page_purge_holes_end(mi_tld_t* tld);
 void          _mi_page_purge_holes_sweep_begin(mi_tld_t* tld);   // once per idle sweep, before its passes
 
 // ------------------------------------------------------

include/mimalloc/types.h

@@ -363,6 +363,9 @@ typedef size_t mi_page_flags_t;
 #define MI_THREADID_ABANDONED           MI_ZU(0)
 #define MI_THREADID_ABANDONED_MAPPED    (MI_ZU(1) << MI_PAGE_FLAG_BITS)
 #define MI_THREADID_DETACHED            (MI_ZU(2) << MI_PAGE_FLAG_BITS)
+// Snapshot: pages captured in a snapshot are owned by nobody and never change again; frees of their blocks are dropped
+// (see mi_free_block_mt) and no thread id can ever equal this value.
+#define MI_THREADID_FROZEN              (~MI_ZU(0) & ~MI_PAGE_FLAG_MASK)
 
 // Thread free list.
 // Points to a list of blocks that are freed by other threads.
@@ -580,6 +583,7 @@ struct mi_theap_s {
 
   long                  page_full_retain;                    // how many full pages can be retained per queue (before abandoning them)
   bool                  allow_page_reclaim;                  // `true` if this theap can reclaim abandoned pages
+  bool                  frozen;                              // pages are an immutable snapshot: never allocate from, collect, or purge them
   bool                  allow_page_abandon;                  // `true` if this theap can abandon pages to reduce memory footprint
   bool                  is_detached;                         // `true` if `tld->thread_id == MI_THREADID_DETACHED`
   bool                  prof_force_slow;                     // if profiling is enabled: keep `pages_free_direct` poisoned so every malloc routes through `_mi_malloc_generic`
@@ -720,6 +724,14 @@ struct mi_tld_s {
   _Atomic(uint32_t)     park_swept;           // this park's sweep is done: don't claim it again until the thread re-parks
   mi_tld_t*             subproc_next;         // list of tlds in the subproc, so the scavenger can find parked threads
   size_t                holes_sweep_seq;      // idle sweeps run over THIS tld's heaps (paces `purge_holes_full_every`)
+  // State of the hole sweep currently running over this tld's pages (owner or scavenger; one sweeper at a time).
+  // Kept here, reached through the page or the tld in hand, rather than in `__thread` variables: on macOS a first
+  // access to a thread-local from inside the allocator makes dyld allocate the thread's TLV block with malloc, which
+  // re-enters the allocator before the variable exists (unbounded recursion on any new thread's first allocation).
+  bool                  holes_sweeping;       // `_mi_page_purge_holes_begin/end`: a sweep is rewriting this tld's pages
+  bool                  holes_sweep_full;     // this sweep ignores `page->swept_state`
+  size_t                holes_sweep_skipped;  // per-sweep counters, folded into the process-wide ones at the end
+  size_t                holes_sweep_visited;
   mi_msecs_t            holes_sweep_last;     // when this tld's heaps were last swept (paces `purge_holes_min_interval`)
 };
 

src/arena.c

@@ -1371,7 +1371,7 @@ static bool mi_arena_page_purge_holes_at(size_t slice_index, size_t slice_count,
     _mi_page_holes_count_page_freed();
     return true;
   }
-  _mi_page_purge_holes(page);
+  _mi_page_purge_holes(page, parg->tld);
   mi_bitmap_set(bitmap, slice_index);     // back in the map *before* unowning: unown may free the page
   mi_abandoned_page_unown(page, NULL);
   return true;
@@ -1383,7 +1383,7 @@ static bool mi_arena_page_purge_holes_at(size_t slice_index, size_t slice_count,
 void _mi_arenas_purge_abandoned_holes(mi_heap_t* heap, mi_tld_t* tld) {
   if (heap == NULL) return;
   if (!mi_option_is_enabled(mi_option_purge_holes)) return;
-  _mi_page_purge_holes_begin();
+  _mi_page_purge_holes_begin(tld);
   mi_forall_arenas(heap, ((mi_arena_t*)NULL), 0, arena) {
     mi_arena_pages_t* const arena_pages = mi_heap_arena_pages(heap, arena);
     if (arena_pages != NULL) {
@@ -1398,7 +1398,7 @@ void _mi_arenas_purge_abandoned_holes(mi_heap_t* heap, mi_tld_t* tld) {
     }
   }
   mi_forall_arenas_end();
-  _mi_page_purge_holes_end();
+  _mi_page_purge_holes_end(tld);
 }
 
 // The read-only counterpart of the sweep above: account for the holes in the abandoned pages
@@ -1469,6 +1469,72 @@ void _mi_arenas_holes_committed(mi_heap_t* heap, mi_holes_report_t* rep) {
 }
 
 
+// snapshot restore: every arena that exists right now holds snapshot memory. Make them exclusive (to nobody) so no
+// theap on any thread places new blocks in their free space; new memory comes from arenas created after this call.
+void mi_arenas_seal_existing(void) mi_attr_noexcept {
+  mi_subproc_t* subproc = _mi_subproc_main();
+  const size_t n = mi_arenas_get_count(subproc);
+  for (size_t i = 0; i < n; i++) {
+    mi_arena_t* arena = mi_arena_from_index(subproc, i);
+    if (arena != NULL) { arena->is_exclusive = true; }
+  }
+}
+
+static bool mi_arena_abandoned_bit_clear(size_t slice_index, size_t slice_count, mi_arena_t* arena, void* arg) {
+  MI_UNUSED(slice_count); MI_UNUSED(arena);
+  mi_bitmap_clear((mi_bitmap_t*)arg, slice_index);
+  return true;
+}
+
+static bool mi_arena_page_freeze(size_t slice_index, size_t slice_count, mi_arena_t* arena, void* arg) {
+  MI_UNUSED(slice_count); MI_UNUSED(arg);
+  mi_page_t* page = mi_arena_page_at_slice(arena, slice_index);
+  mi_atomic_store_release(&page->xthread_id, (mi_threadid_t)MI_THREADID_FROZEN);
+  return true;
+}
+
+// Snapshot build: every page that exists right now is about to be written into the snapshot. Owned by MI_THREADID_FROZEN,
+// frees of their blocks in the restored process (or in what is left of this one) take the cross-thread path and are dropped
+// there, so snapshot pages are never dirtied by the allocator. Done here, in the builder, so the restore writes nothing.
+void mi_arenas_freeze_pages(void) mi_attr_noexcept {
+  mi_subproc_t* subproc = _mi_subproc_main();
+  const size_t n = mi_arenas_get_count(subproc);
+  mi_lock(&subproc->heaps_lock) {
+    for (mi_heap_t* heap = subproc->heaps; heap != NULL; heap = heap->next) {   // page bitmaps are kept per heap
+      for (size_t i = 0; i < n; i++) {
+        mi_arena_t* arena = mi_arena_from_index(subproc, i);
+        if (arena == NULL) continue;
+        mi_arena_pages_t* arena_pages = mi_heap_arena_pages(heap, arena);
+        if (arena_pages == NULL) continue;
+        (void)_mi_bitmap_forall_set(arena_pages->pages, &mi_arena_page_freeze, arena, NULL);
+        // A frozen page must not be findable as abandoned either: reclaiming it (or sweeping its holes) would allocate into
+        // and rewrite snapshot memory. Everything abandoned at this point is snapshot memory, so the heap's counts go to zero.
+        for (size_t bin = 0; bin < MI_ARENA_BIN_COUNT; bin++) {
+          mi_bitmap_t* const abandoned = arena_pages->pages_abandoned[bin];
+          (void)_mi_bitmap_forall_set(abandoned, &mi_arena_abandoned_bit_clear, arena, abandoned);
+        }
+      }
+      for (size_t bin = 0; bin < MI_BIN_COUNT; bin++) { mi_atomic_store_relaxed(&heap->abandoned_count[bin], 0); }
+    }
+  }
+}
+
+// Visit every maximal run of free slices (belonging to no page) across all arenas of `heap`'s subproc.
+void mi_arenas_visit_free_ranges(mi_heap_t* heap, void (*visit)(void* start, size_t size, void* arg), void* arg) mi_attr_noexcept {
+  if (heap == NULL || visit == NULL) return;
+  mi_forall_arenas(heap, ((mi_arena_t*)NULL), 0, arena) {
+    const size_t slice_count = arena->slice_count;
+    size_t i = 0;
+    while (i < slice_count) {
+      if (!mi_bbitmap_is_setN(arena->slices_free, i, 1)) { i++; continue; }
+      size_t j = i; while (j < slice_count && mi_bbitmap_is_setN(arena->slices_free, j, 1)) j++;
+      visit(mi_arena_slice_start(arena, i), mi_size_of_slices(j - i), arg);
+      i = j;
+    }
+  }
+  mi_forall_arenas_end();
+}
+
 /* -----------------------------------------------------------
   Arena free
 ----------------------------------------------------------- */

src/free.c

@@ -63,6 +63,9 @@ static void mi_decl_noinline mi_free_try_collect_mt(mi_page_t* page, mi_block_t*
 // Free a block multi-threaded
 static inline void mi_free_block_mt(mi_page_t* page, mi_block_t* block, bool was_guarded, bool allow_collect) mi_attr_noexcept
 {
+  // A page frozen into a snapshot (thread id MI_THREADID_FROZEN, so every free of its blocks lands here) is never written
+  // again: dropping the block keeps the page clean and file-backed. Costs one compare on a path that is already the slow one.
+  if mi_unlikely(mi_page_thread_id(page) == MI_THREADID_FROZEN) return;
   // todo: we cannot safely check for double free in _mt -- should check when collecting the thread_free list
   if (!was_guarded) { mi_check_padding(page, block); }   // checking padding is safe for mt
   // adjust stats (after padding check )

src/init.c

@@ -111,7 +111,9 @@ static mi_decl_cache_align mi_tld_t mi_tld_detached = {
   NULL,                   // park_theap0
   MI_ATOMIC_VAR_INIT(0),  // park_swept
   NULL,                   // subproc_next
-  0, 0                    // holes_sweep_seq / _last
+  0,                      // holes_sweep_seq
+  false, false, 0, 0,     // holes_sweeping / holes_sweep_full / _skipped / _visited
+  0                       // holes_sweep_last
 };
 
 mi_decl_hidden mi_decl_cache_align const mi_theap_t _mi_theap_empty = {
@@ -130,6 +132,7 @@ mi_decl_hidden mi_decl_cache_align const mi_theap_t _mi_theap_empty = {
   NULL, NULL,             // hnext, hprev
   0,                      // full page retain
   false,                  // allow reclaim
+  false,                  // frozen
   true,                   // allow abandon
   true,                   // is_detached
   false, 0,               // prof_force_slow, prof_countdown (fork)

src/os.c

@@ -4,10 +4,12 @@ This is free software; you can redistribute it and/or modify it under the
 terms of the MIT license. A copy of the license can be found in the file
 "LICENSE" at the root of this distribution.
 -----------------------------------------------------------------------------*/
+#include <errno.h>   // ENOENT
 #include "mimalloc.h"
 #include "mimalloc/internal.h"
 #include "mimalloc/atomic.h"
 #include "mimalloc/prim.h"
+#include <stdlib.h>
 #include "mimalloc/prim-tls.h"  // _mi_theap_default for random
 
 /* -----------------------------------------------------------
@@ -128,12 +130,71 @@ bool _mi_os_commit(mi_subproc_t* subproc, void* addr, size_t size, bool* is_zero
 #define MI_HINT_AREA  ((uintptr_t)4 << 40)  // upto (2+4) 6TiB  (since before win8 there is "only" 8TiB available to processes)
 #define MI_HINT_MAX   ((uintptr_t)30 << 40) // wrap after 30TiB (area after 32TiB is used for huge OS pages)
 
+// snapshot support (see mi_arenas_freeze_pages / mi_os_hint_floor). The embedder can name a function, callable from the
+// very first allocation, that says whether this executable can build or map a snapshot (`-DMI_STARTUP_SNAPSHOT_HOST_FN=name`,
+// `int name(void)`); such executables get deterministic address hints without any environment. It can also name the
+// environment variable that marks a snapshot *build* (`-DMI_STARTUP_SNAPSHOT_BUILD_ENV="VAR"`): a build keeps its heap at the
+// default base (that heap becomes the snapshot); any other snapshot-capable process keeps its own early heap above snapshot space.
+#if defined(MI_STARTUP_SNAPSHOT_HOST_FN)
+#ifdef __cplusplus
+extern "C" int MI_STARTUP_SNAPSHOT_HOST_FN(void);
+#else
+extern int MI_STARTUP_SNAPSHOT_HOST_FN(void);
+#endif
+#define mi_startup_snapshot_capable() (MI_STARTUP_SNAPSHOT_HOST_FN() != 0)
+#else
+#define mi_startup_snapshot_capable() (0)
+#endif
+#if !defined(MI_STARTUP_SNAPSHOT_BUILD_ENV)
+#define MI_STARTUP_SNAPSHOT_BUILD_ENV "MIMALLOC_SNAPSHOT_BUILD"
+#endif
+#define MI_STARTUP_SNAPSHOT_RESTORER_FLOOR  ((uintptr_t)0x21000000000ULL)   // 64GiB above the default hint base
+static mi_decl_cache_align _Atomic(uintptr_t) aligned_base; // = 0  (hint bump pointer; file scope so mi_os_hint_floor can move it)
+static _Atomic(uintptr_t) hint_floor; // = 0: lowest address hinted allocations may use (also where the pointer restarts when it wraps or is invalid)
+
+// Snapshot restore: every hinted OS allocation from now on lands at or above `floor` (the snapshot occupies the area below). Also repairs a
+// pointer that is unusable (0, past MI_HINT_MAX — e.g. inherited from another process's data segment).
+void mi_os_hint_floor(void* floor) mi_attr_noexcept {
+  uintptr_t f = _mi_align_up((uintptr_t)floor, MI_HINT_ALIGN);
+  uintptr_t curf = mi_atomic_load_acquire(&hint_floor);
+  while (curf < f && !mi_atomic_cas_weak_acq_rel(&hint_floor, &curf, f)) { }
+  uintptr_t cur = mi_atomic_load_acquire(&aligned_base);
+  while ((cur < f || cur > MI_HINT_MAX) && !mi_atomic_cas_weak_acq_rel(&aligned_base, &cur, f)) { }
+}
+
+// Deterministic hinting: on when the executable is snapshot-capable (MI_STARTUP_SNAPSHOT_HOST_FN) or requested through the
+// environment. Decided once; the embedder's own reservations ask the same question (`mi_startup_snapshot_hints_enabled`).
+static bool mi_startup_snapshot_hints_decide(void) {
+  static int decided = -1;
+  if (decided < 0) {
+    // _mi_getenv (0 = found), not getenv: this runs inside the very first allocation, and the C runtime's getenv allocates on some platforms.
+    char buf[64];
+    const int from_env = (_mi_getenv("MIMALLOC_DETERMINISTIC_HINT", buf, sizeof(buf)) == 0 && buf[0] == '1') ? 1 : 0;
+    // MIMALLOC_HINT_FLOOR=0x...: start hinted OS allocations at/above this address (a snapshot will be mapped below it)
+    if (_mi_getenv("MIMALLOC_HINT_FLOOR", buf, sizeof(buf)) == 0 && buf[0] != 0) {
+      unsigned long long v = strtoull(buf, NULL, 0);
+      if (v) mi_os_hint_floor((void*)(uintptr_t)v);
+    }
+    // Snapshot-capable and not building: this process may be about to map a snapshot over the default hint area, so nothing it
+    // allocates before then (libc scratch, dyld) may live there.
+    // Presence is what matters here; the value (an output path) may well not fit `buf`, which _mi_getenv reports as EAGAIN, not ENOENT.
+    else if (mi_startup_snapshot_capable() && _mi_getenv(MI_STARTUP_SNAPSHOT_BUILD_ENV, buf, sizeof(buf)) == ENOENT) mi_os_hint_floor((void*)MI_STARTUP_SNAPSHOT_RESTORER_FLOOR);
+    decided = (from_env || mi_startup_snapshot_capable()) ? 1 : 0;
+  }
+  return decided == 1;
+}
+
+bool mi_startup_snapshot_hints_enabled(void) mi_attr_noexcept {
+  return mi_startup_snapshot_hints_decide();
+}
+
 void* _mi_os_get_aligned_hint(size_t try_alignment, size_t size)
 {
-  static mi_decl_cache_align _Atomic(uintptr_t) aligned_base; // = 0
 
   // todo: perhaps only do alignment hints if THP is enabled?
-  if (try_alignment <= mi_os_mem_config.alloc_granularity || try_alignment > MI_HINT_ALIGN) return NULL;
+  const int deterministic_all = mi_startup_snapshot_hints_enabled();
+  if (!deterministic_all && (try_alignment <= mi_os_mem_config.alloc_granularity || try_alignment > MI_HINT_ALIGN)) return NULL;
+  if (deterministic_all && try_alignment > MI_HINT_ALIGN) return NULL;
   if (mi_os_mem_config.virtual_address_bits < 46) return NULL;  // < 64TiB virtual address space
   size = _mi_align_up(size, MI_HINT_ALIGN);
   #if (MI_SECURE>=1)
@@ -142,13 +203,19 @@ void* _mi_os_get_aligned_hint(size_t try_alignment, size_t size)
   size += MI_HINT_ALIGN;              // put in virtual gaps between hinted blocks; this splits VLA's but increases guarded areas.
 
   uintptr_t hint = mi_atomic_add_acq_rel(&aligned_base, size);
-  if (hint == 0 || hint > MI_HINT_MAX) {   // wrap or initialize
+  if (hint == 0 || hint > MI_HINT_MAX || hint < mi_atomic_load_relaxed(&hint_floor)) {   // wrap or initialize (never below the floor)
     uintptr_t init = MI_HINT_BASE;
+    { uintptr_t fl = mi_atomic_load_relaxed(&hint_floor); if (fl > init) init = fl; }
+    // Experiment: MIMALLOC_DETERMINISTIC_HINT=1 disables hint randomization (snapshot determinism tests).
+    const int deterministic = deterministic_all;
+    MI_UNUSED(deterministic);   // only consulted in the randomizing configurations below
     #if (MI_SECURE>=1 || defined(NDEBUG))  // security: randomize start of aligned allocations unless in debug mode
+    if (!deterministic) {
     mi_theap_t* const theap = _mi_theap_default();     // don't use `mi_theap_get_default()` as that can cause allocation recursively (issue #1267)
     if (!mi_theap_is_initialized(theap)) return NULL;  // no hint as we lack randomness at this point
     const uintptr_t r = _mi_theap_random_next(theap);
     init = init + ((MI_HINT_ALIGN * ((r>>17) & 0xFFFFF)) % MI_HINT_AREA);  // (randomly 20 bits)*4MiB == 0 to 4TiB
+    }
     #endif
     uintptr_t expected = hint + size;
     mi_atomic_cas_strong_acq_rel(&aligned_base, &expected, init);
@@ -163,6 +230,8 @@ void* _mi_os_get_aligned_hint(size_t try_alignment, size_t size) {
    MI_UNUSED(try_alignment); MI_UNUSED(size);
    return NULL;
 }
+bool mi_startup_snapshot_hints_enabled(void) mi_attr_noexcept { return false; }
+void mi_os_hint_floor(void* floor) mi_attr_noexcept { MI_UNUSED(floor); }
 #endif