feat: implement dynamic alignment for ArenaAllocator

Author: shruti2522

.github/workflows/rust.yml

@@ -67,7 +67,7 @@ jobs:
       - name: Set up Miri
         run: cargo +nightly miri setup
       - name: Run tests under Miri
-        run: cargo +nightly miri test -p oscars
+        run: cargo +nightly miri test -p oscars --all-features
 
   docs:
     name: Documentation

Cargo.lock

@@ -8,6 +8,7 @@ allocator-api2 = { version = "0.4.0", optional = true }
 hashbrown = "0.16.1"
 oscars_derive = { path = "../oscars_derive", version = "0.1.0" }
 rustc-hash = "2.1.1"
+thin-vec = { version = "0.2", optional = true }
 
 [dev-dependencies]
 criterion = { version = "0.5", features = ["html_reports"] }
@@ -32,3 +33,4 @@ default = ["mark_sweep"]
 std = []
 mark_sweep = []
 gc_allocator = ["dep:allocator-api2", "mark_sweep"]
+thin-vec = ["dep:thin-vec", "mark_sweep"]

oscars/Cargo.toml

@@ -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)
     }
 }
 
@@ -341,9 +342,12 @@ impl<'arena> Arena<'arena> {
         value_ref: &T,
     ) -> Result<ArenaAllocationData, ArenaAllocError> {
         let size = core::mem::size_of::<ArenaHeapItem<T>>();
-        let alignment = core::mem::align_of_val(value_ref);
+        let alignment = core::mem::align_of::<ArenaHeapItem<T>>();
 
-        assert!(alignment <= self.layout.align());
+        // The arena's buffer must be at least as aligned as the value we are storing.
+        if alignment > self.layout.align() {
+            return Err(ArenaAllocError::AlignmentNotPossible);
+        }
 
         // Safety: This is safe as `current_offset` must be less then the length
         // of the buffer.
@@ -396,6 +400,22 @@ 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"
+        );
+        // Zero the buffer so stale object graphs are not observable after recycling.
+        // SAFETY: buffer is valid for the full layout size and was allocated with
+        // the same layout in try_init.
+        unsafe { core::ptr::write_bytes(self.buffer.as_ptr(), 0, self.layout.size()) };
+        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/alloc.rs

@@ -1,5 +1,7 @@
 //! An Arena allocator that manages multiple backing arenas
 
+use core::mem;
+
 use rust_alloc::alloc::LayoutError;
 use rust_alloc::collections::LinkedList;
 
@@ -48,19 +50,33 @@ const DEFAULT_ARENA_SIZE: usize = 4096;
 /// Default upper limit of 2MB (2 ^ 21)
 const DEFAULT_HEAP_THRESHOLD: usize = 2_097_152;
 
+/// Minimum guaranteed alignment for every arena buffer.
+const DEFAULT_MIN_ALIGNMENT: usize = 8;
+
+/// 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,
+    min_alignment: 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> {
     fn default() -> Self {
         Self {
             heap_threshold: DEFAULT_HEAP_THRESHOLD,
             arena_size: DEFAULT_ARENA_SIZE,
+            min_alignment: DEFAULT_MIN_ALIGNMENT,
             arenas: LinkedList::default(),
+            recycled_arenas: core::array::from_fn(|_| None),
+            recycled_count: 0,
         }
     }
 }
@@ -74,17 +90,24 @@ impl<'alloc> ArenaAllocator<'alloc> {
         self.heap_threshold = heap_threshold;
         self
     }
+    /// Override the baseline alignment for every new arena buffer.
+    pub fn with_min_alignment(mut self, min_alignment: usize) -> Self {
+        self.min_alignment = min_alignment;
+        self
+    }
 
     pub fn arenas_len(&self) -> usize {
         self.arenas.len()
     }
 
     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) {
@@ -94,22 +117,26 @@ impl<'alloc> ArenaAllocator<'alloc> {
 
 impl<'alloc> ArenaAllocator<'alloc> {
     pub fn try_alloc<T>(&mut self, value: T) -> Result<ArenaPointer<'alloc, T>, ArenaAllocError> {
+        // Determine the minimum alignment this type requires.
+        let required_alignment = mem::align_of::<alloc::ArenaHeapItem<T>>();
+
         let active = match self.get_active_arena() {
             Some(arena) => arena,
             None => {
-                // TODO: don't hard code alignment
-                //
-                // TODO: also, we need a min-alignment
-                self.initialize_new_arena()?;
+                self.initialize_new_arena(required_alignment)?;
                 self.get_active_arena().expect("must exist, we just set it")
             }
         };
 
         match active.get_allocation_data(&value) {
             // SAFETY: TODO
             Ok(data) => unsafe { Ok(active.alloc_unchecked::<T>(value, data)) },
-            Err(ArenaAllocError::OutOfMemory) => {
-                self.initialize_new_arena()?;
+            // The active arena is either full or was created with an alignment
+            // that is too small for this type. Either way, close it and spin up
+            // a fresh arena that satisfies the alignment requirement.
+            Err(ArenaAllocError::OutOfMemory | ArenaAllocError::AlignmentNotPossible) => {
+                active.close();
+                self.initialize_new_arena(required_alignment)?;
                 let new_active = self.get_active_arena().expect("must exist");
                 new_active.try_alloc(value)
             }
@@ -127,8 +154,28 @@ impl<'alloc> ArenaAllocator<'alloc> {
             .transpose()
     }
 
-    pub fn initialize_new_arena(&mut self) -> Result<(), ArenaAllocError> {
-        let new_arena = Arena::try_init(self.arena_size, 16)?;
+    /// Initialize a fresh arena, attempting to reuse a recycled one first.
+    pub fn initialize_new_arena(
+        &mut self,
+        required_alignment: usize,
+    ) -> Result<(), ArenaAllocError> {
+        let alignment = self.min_alignment.max(required_alignment);
+
+        // 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.
+                // Only reuse if its original alignment satisfies the current requirement,
+                // otherwise drop it and fall through to a fresh OS allocation.
+                if recycled.layout.align() >= alignment {
+                    self.arenas.push_front(recycled);
+                    return Ok(());
+                }
+            }
+        }
+
+        let new_arena = Arena::try_init(self.arena_size, alignment)?;
         self.arenas.push_front(new_arena);
         Ok(())
     }
@@ -138,8 +185,14 @@ 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)
+        for arena in self.arenas.extract_if(|a| a.run_drop_check()) {
+            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: arena drops here, returning memory to the OS
         }
     }
 

oscars/src/alloc/arena2/mod.rs

@@ -89,6 +89,71 @@ 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);
+    // recycled_count == 1 proves the arena was parked in the recycle slot, not freed to the OS.
+    assert_eq!(allocator.recycled_count, 1);
+
+    // 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);
+    // recycled_count == 0 proves the recycled slot was consumed rather than a new OS allocation.
+    assert_eq!(allocator.recycled_count, 0);
+}
+
+#[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
@@ -119,3 +184,62 @@ fn as_ptr_clears_not_flips_tag_bit() {
         ptr_a.as_mut().mark_dropped();
     }
 }
+
+// === test for Dynamic Alignment === //
+
+#[test]
+fn test_over_aligned_type() {
+    #[repr(C, align(512))]
+    struct HighlyAligned {
+        _data: [u8; 128],
+    }
+
+    let mut allocator = ArenaAllocator::default().with_arena_size(4096);
+    let ptr = allocator
+        .try_alloc(HighlyAligned { _data: [0; 128] })
+        .unwrap();
+
+    let addr = ptr.as_ptr().as_ptr() as usize;
+    assert_eq!(addr % 512, 0);
+    assert_eq!(allocator.arenas_len(), 1);
+}
+
+#[test]
+fn test_alignment_upgrade_after_small_alloc() {
+    #[repr(C, align(512))]
+    struct BigAlign([u8; 16]);
+
+    let mut allocator = ArenaAllocator::default().with_arena_size(4096);
+
+    // force the first arena to use 8-byte alignment
+    let _small = allocator.try_alloc(0u8).unwrap();
+    assert_eq!(allocator.arenas_len(), 1);
+
+    let ptr = allocator.try_alloc(BigAlign([0; 16])).unwrap();
+
+    let addr = ptr.as_ptr().as_ptr() as usize;
+    assert_eq!(addr % 512, 0);
+    assert_eq!(allocator.arenas_len(), 2);
+}
+
+#[test]
+fn test_alignment_upgrade_on_full_arena() {
+    #[repr(C, align(512))]
+    struct BigAlign([u8; 16]);
+
+    let mut allocator = ArenaAllocator::default().with_arena_size(4096);
+
+    // fill the first arena
+    let mut count = 0usize;
+    while allocator.arenas_len() < 2 {
+        let _ = allocator.try_alloc(0u64).unwrap();
+        count += 1;
+        assert!(count < 1024);
+    }
+
+    let ptr = allocator.try_alloc(BigAlign([0; 16])).unwrap();
+
+    let addr = ptr.as_ptr().as_ptr() as usize;
+    assert_eq!(addr % 512, 0);
+    assert_eq!(allocator.arenas_len(), 3);
+}