recycle empty arenas in arena2 and mempool3 allocator

Author: shruti2522

oscars/src/alloc/arena2/alloc.rs

@@ -115,7 +115,7 @@ impl<T> TaggedPtr<T> {
     }
 
     fn as_ptr(&self) -> *mut T {
-        self.0.map_addr(|addr| addr ^ MASK)
+        self.0.map_addr(|addr| addr & !MASK)
     }
 }
 
@@ -203,7 +203,8 @@ impl<'arena, T> ArenaPointer<'arena, T> {
     ///
     /// safe because the gc collector owns the arena and keeps it alive
     pub(crate) unsafe fn extend_lifetime(self) -> ArenaPointer<'static, T> {
-        ArenaPointer(self.0.extend_lifetime(), PhantomData)
+        // SAFETY: upheld by caller
+        ArenaPointer(unsafe { self.0.extend_lifetime() }, PhantomData)
     }
 }
 
@@ -396,6 +397,18 @@ impl<'arena> Arena<'arena> {
         }
         result
     }
+
+    /// Reset arena to its initial empty state, reusing the existing OS buffer.
+    /// Must only be called when `run_drop_check()` is true (all items dropped).
+    pub fn reset(&self) {
+        debug_assert!(
+            self.run_drop_check(),
+            "reset() called on an arena with live items"
+        );
+        self.flags.set(ArenaState::default());
+        self.last_allocation.set(core::ptr::null_mut());
+        self.current_offset.set(0);
+    }
 }
 
 impl<'arena> Drop for Arena<'arena> {

oscars/src/alloc/arena2/mod.rs

@@ -48,11 +48,18 @@ const DEFAULT_ARENA_SIZE: usize = 4096;
 /// Default upper limit of 2MB (2 ^ 21)
 const DEFAULT_HEAP_THRESHOLD: usize = 2_097_152;
 
+/// Maximum number of idle arenas held (4 idle pages x 4KB = 16KB of OS memory pressure buffered)
+const MAX_RECYCLED_ARENAS: usize = 4;
+
 #[derive(Debug)]
 pub struct ArenaAllocator<'alloc> {
     heap_threshold: usize,
     arena_size: usize,
     arenas: LinkedList<Arena<'alloc>>,
+    // empty arenas kept alive to avoid OS reallocation on the next cycle
+    recycled_arenas: [Option<Arena<'alloc>>; MAX_RECYCLED_ARENAS],
+    // number of idle arenas currently held
+    recycled_count: usize,
 }
 
 impl<'alloc> Default for ArenaAllocator<'alloc> {
@@ -61,6 +68,8 @@ impl<'alloc> Default for ArenaAllocator<'alloc> {
             heap_threshold: DEFAULT_HEAP_THRESHOLD,
             arena_size: DEFAULT_ARENA_SIZE,
             arenas: LinkedList::default(),
+            recycled_arenas: core::array::from_fn(|_| None),
+            recycled_count: 0,
         }
     }
 }
@@ -80,11 +89,13 @@ impl<'alloc> ArenaAllocator<'alloc> {
     }
 
     pub fn heap_size(&self) -> usize {
+        // recycled arenas hold no live objects, exclude them from GC pressure
         self.arenas_len() * self.arena_size
     }
 
     pub fn is_below_threshold(&self) -> bool {
-        self.heap_size() <= self.heap_threshold - self.arena_size
+        // saturating_sub avoids underflow when heap_threshold < arena_size
+        self.heap_size() <= self.heap_threshold.saturating_sub(self.arena_size)
     }
 
     pub fn increase_threshold(&mut self) {
@@ -128,6 +139,16 @@ impl<'alloc> ArenaAllocator<'alloc> {
     }
 
     pub fn initialize_new_arena(&mut self) -> Result<(), ArenaAllocError> {
+        // Check the recycle list first to avoid an OS allocation.
+        if self.recycled_count > 0 {
+            self.recycled_count -= 1;
+            if let Some(recycled) = self.recycled_arenas[self.recycled_count].take() {
+                // arena.reset() was already called when it was parked
+                self.arenas.push_front(recycled);
+                return Ok(());
+            }
+        }
+
         let new_arena = Arena::try_init(self.arena_size, 16)?;
         self.arenas.push_front(new_arena);
         Ok(())
@@ -138,8 +159,19 @@ impl<'alloc> ArenaAllocator<'alloc> {
     }
 
     pub fn drop_dead_arenas(&mut self) {
-        for dead_arenas in self.arenas.extract_if(|a| a.run_drop_check()) {
-            drop(dead_arenas)
+        let dead_arenas: rust_alloc::vec::Vec<Arena> =
+            self.arenas.extract_if(|a| a.run_drop_check()).collect();
+
+        for arena in dead_arenas {
+            if self.recycled_count < MAX_RECYCLED_ARENAS {
+                //reset in place and park in the reserve.
+                arena.reset();
+                self.recycled_arenas[self.recycled_count] = Some(arena);
+                self.recycled_count += 1;
+            } else {
+                // Reserve is full so free this arena to the OS.
+                drop(arena);
+            }
         }
     }
 

oscars/src/alloc/arena2/tests.rs

@@ -88,3 +88,95 @@ fn arc_drop() {
 
     assert_eq!(allocator.arenas_len(), 0);
 }
+
+#[test]
+fn recycled_arena_avoids_realloc() {
+    let mut allocator = ArenaAllocator::default().with_arena_size(512);
+
+    let mut ptrs = Vec::new();
+    for i in 0..16 {
+        ptrs.push(allocator.try_alloc(i).unwrap().as_ptr());
+    }
+    assert_eq!(allocator.arenas_len(), 1);
+    // heap_size counts only live arenas, so capture it while one is active.
+    let heap_while_live = allocator.heap_size();
+    assert_eq!(heap_while_live, 512);
+
+    for mut ptr in ptrs {
+        unsafe { ptr.as_mut().mark_dropped() };
+    }
+    allocator.drop_dead_arenas();
+
+    // After recycling, the arena is parked, no live arenas, so heap_size is 0.
+    assert_eq!(allocator.arenas_len(), 0);
+    assert_eq!(allocator.heap_size(), 0);
+
+    // Allocate again, must reuse the recycled arena without growing OS footprint.
+    // heap_size returns to the same value as when a live arena was present.
+    for i in 16..32 {
+        let _ = allocator.try_alloc(i).unwrap();
+    }
+    assert_eq!(allocator.arenas_len(), 1);
+    assert_eq!(allocator.heap_size(), heap_while_live);
+}
+
+#[test]
+fn max_recycled_cap_respected() {
+    let mut allocator = ArenaAllocator::default().with_arena_size(128);
+
+    let mut ptrs_per_arena: Vec<Vec<NonNull<ArenaHeapItem<u64>>>> = Vec::new();
+
+    for _ in 0..5 {
+        let mut ptrs = Vec::new();
+        let target_len = allocator.arenas_len() + 1;
+        while allocator.arenas_len() < target_len {
+            ptrs.push(allocator.try_alloc(0u64).unwrap().as_ptr());
+        }
+        ptrs_per_arena.push(ptrs);
+    }
+    assert_eq!(allocator.arenas_len(), 5);
+
+    for ptrs in ptrs_per_arena {
+        for mut ptr in ptrs {
+            unsafe { ptr.as_mut().mark_dropped() };
+        }
+    }
+
+    allocator.drop_dead_arenas();
+
+    assert_eq!(allocator.arenas_len(), 0);
+    assert_eq!(allocator.heap_size(), 0);
+    // The recycled list holds exactly max_recycled pages.
+    assert_eq!(allocator.recycled_count, 4);
+}
+
+// === test for TaggedPtr::as_ptr === //
+
+// `TaggedPtr::as_ptr` must use `addr & !MASK` to unconditionally clear the high
+// bit rather than XORing it out. The XOR approach worked for tagged items
+// but incorrectly flipped the bit on untagged items, corrupting the pointer.
+#[test]
+fn as_ptr_clears_not_flips_tag_bit() {
+    let mut allocator = ArenaAllocator::default();
+
+    let mut ptr_a = allocator.try_alloc(1u64).unwrap().as_ptr();
+    let mut ptr_b = allocator.try_alloc(2u64).unwrap().as_ptr();
+    let _ptr_c = allocator.try_alloc(3u64).unwrap().as_ptr();
+    assert_eq!(allocator.arenas_len(), 1);
+
+    // Mark B and C as dropped, leave A live.
+    unsafe {
+        ptr_b.as_mut().mark_dropped();
+    }
+
+    let states = allocator.arena_drop_states();
+    assert_eq!(
+        states[0].as_slice(),
+        &[false, true, false],
+        "item_drop_states must correctly report live/dropped status for all nodes"
+    );
+
+    unsafe {
+        ptr_a.as_mut().mark_dropped();
+    }
+}

oscars/src/alloc/mempool3/alloc.rs

@@ -277,6 +277,29 @@ impl SlotPool {
     pub fn run_drop_check(&self) -> bool {
         self.live.get() == 0
     }
+
+    /// Reset this pool to the empty state it had after `try_init`, reusing the
+    /// existing OS buffer. Must only be called when `run_drop_check()` is true.
+    ///
+    /// After `reset()` the pool is ready for `alloc_slot` without any further
+    /// OS interaction, avoiding a round trip through the global allocator.
+    pub fn reset(&self) {
+        debug_assert_eq!(
+            self.live.get(),
+            0,
+            "reset() called on a non-empty SlotPool (live = {})",
+            self.live.get()
+        );
+        // Clear the bitmap so all slots become free again.
+        //
+        // SAFETY: buffer is valid for at least `bitmap_bytes` and was
+        // originally zero initialised in try_init with the same length.
+        unsafe {
+            core::ptr::write_bytes(self.buffer.as_ptr(), 0, self.bitmap_bytes);
+        }
+        self.bump.set(0);
+        self.free_list.set(None);
+    }
 }
 
 impl Drop for SlotPool {

oscars/src/alloc/mempool3/mod.rs

@@ -56,6 +56,10 @@ pub struct PoolAllocator<'alloc> {
     pub(crate) free_cache: Cell<usize>,
     // per size class cached index of the last pool used by alloc_slot
     pub(crate) alloc_cache: [Cell<usize>; 12],
+    // empty slot pools kept alive to avoid OS reallocation on the next cycle
+    pub(crate) recycled_pools: Vec<SlotPool>,
+    // maximum number of idle pages held across all size classes
+    pub(crate) max_recycled: usize,
     _marker: core::marker::PhantomData<&'alloc ()>,
 }
 
@@ -82,6 +86,9 @@ impl<'alloc> Default for PoolAllocator<'alloc> {
                 Cell::new(usize::MAX),
                 Cell::new(usize::MAX),
             ],
+            recycled_pools: Vec::new(),
+            // one idle page per size class keeps memory pressure manageable
+            max_recycled: SIZE_CLASSES.len(),
             _marker: core::marker::PhantomData,
         }
     }
@@ -155,6 +162,29 @@ impl<'alloc> PoolAllocator<'alloc> {
         }
 
         // need a new pool for this size class
+        // try the recycle list first
+        // to avoid a round trip through the OS allocator
+        if let Some(pos) = self
+            .recycled_pools
+            .iter()
+            .rposition(|p| p.slot_size == slot_size)
+        {
+            let pool = self.recycled_pools.swap_remove(pos);
+            // pool.reset() was already called in drop_empty_pools when it was parked
+            let slot_ptr = pool.alloc_slot().ok_or(PoolAllocError::OutOfMemory)?;
+            let insert_idx = self.slot_pools.len();
+            self.slot_pools.push(pool);
+            self.alloc_cache[sc_idx].set(insert_idx);
+
+            // SAFETY: slot_ptr was successfully allocated for this size class
+            return unsafe {
+                let dst = slot_ptr.as_ptr() as *mut PoolItem<T>;
+                dst.write(PoolItem(value));
+                Ok(PoolPointer::from_raw(NonNull::new_unchecked(dst)))
+            };
+        }
+
+        // Recycle list had no match, allocate a fresh page from the OS.
         let total = self.page_size.max(slot_size * 4);
         let new_pool = SlotPool::try_init(slot_size, total, 16)?;
         self.current_heap_size += new_pool.layout.size();
@@ -267,16 +297,27 @@ impl<'alloc> PoolAllocator<'alloc> {
         false
     }
 
-    /// drop empty slot pools and bump pages
+    /// Reclaim slot pool pages that became empty after a GC sweep.
+    ///
+    /// Empty pages are parked in a recycle list (up to `max_recycled`)
+    /// to avoid global allocator round trips on the next allocation.
     pub fn drop_empty_pools(&mut self) {
-        self.slot_pools.retain(|p| {
-            if p.run_drop_check() {
-                self.current_heap_size = self.current_heap_size.saturating_sub(p.layout.size());
-                false
+        // Drain fully empty slot pools into the recycle list.
+        let empties: Vec<SlotPool> = self
+            .slot_pools
+            .extract_if(.., |p| p.run_drop_check())
+            .collect();
+
+        for pool in empties {
+            if self.recycled_pools.len() < self.max_recycled {
+                pool.reset();
+                self.recycled_pools.push(pool);
             } else {
-                true
+                self.current_heap_size = self.current_heap_size.saturating_sub(pool.layout.size());
             }
-        });
+        }
+
+        // Bump pages have no size class affinity so we always free them.
         self.bump_pages.retain(|p| {
             if p.run_drop_check() {
                 self.current_heap_size = self.current_heap_size.saturating_sub(p.layout.size());

oscars/src/alloc/mempool3/tests.rs

@@ -173,3 +173,64 @@ fn slot_count_tight_capacity() {
     assert_eq!(bitmap_bytes, 8);
     assert_eq!(slot_count, 7);
 }
+
+/// Verify that recycled empty slot pools are reused on the next `try_alloc`
+/// without allocating new OS memory, the heap_size should be unchanged.
+#[test]
+fn recycled_pool_avoids_realloc() {
+    let mut allocator = PoolAllocator::default().with_page_size(4096);
+
+    let ptrs: Vec<_> = (0u64..16)
+        .map(|i| allocator.try_alloc(i).unwrap().as_ptr())
+        .collect();
+    assert_eq!(allocator.slot_pools.len(), 1);
+    let heap_after_first_alloc = allocator.current_heap_size;
+
+    for ptr in ptrs {
+        allocator.free_slot(ptr.cast::<u8>());
+    }
+    allocator.drop_empty_pools();
+
+    assert_eq!(allocator.slot_pools.len(), 0);
+    assert_eq!(allocator.recycled_pools.len(), 1);
+    assert_eq!(allocator.current_heap_size, heap_after_first_alloc);
+
+    let heap_before_second_alloc = allocator.current_heap_size;
+    for i in 16u64..32 {
+        let _ = allocator.try_alloc(i).unwrap();
+    }
+
+    assert_eq!(allocator.slot_pools.len(), 1);
+    assert_eq!(allocator.recycled_pools.len(), 0);
+    assert_eq!(allocator.current_heap_size, heap_before_second_alloc);
+}
+
+/// Verify that when more pools become empty than `max_recycled` allows,
+/// the overflow is freed to the OS.
+#[test]
+fn max_recycled_cap_respected() {
+    let mut allocator = PoolAllocator::default().with_page_size(32);
+    allocator.max_recycled = 0;
+
+    let p1 = allocator.try_alloc(1u64).unwrap().as_ptr();
+    let px = allocator.try_alloc(2u64).unwrap().as_ptr();
+    let py = allocator.try_alloc(3u64).unwrap().as_ptr();
+    assert_eq!(allocator.slot_pools.len(), 1);
+
+    let p2 = allocator.try_alloc(4u64).unwrap().as_ptr();
+    assert_eq!(allocator.slot_pools.len(), 2);
+
+    let heap_before = allocator.current_heap_size;
+
+    allocator.free_slot(p1.cast::<u8>());
+    allocator.free_slot(px.cast::<u8>());
+    allocator.free_slot(py.cast::<u8>());
+    allocator.free_slot(p2.cast::<u8>());
+
+    allocator.max_recycled = 1;
+    allocator.drop_empty_pools();
+
+    assert_eq!(allocator.slot_pools.len(), 0);
+    assert_eq!(allocator.recycled_pools.len(), 1);
+    assert!(allocator.current_heap_size < heap_before);
+}