Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion oscars/src/alloc/arena2/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -396,6 +397,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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thought: we should probably zero out the memory here to as a proper reset to be "good citizens" so to speak.

I'm not convinced that blocks this PR though.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added write_bytes(0) over the full layout size to zero the buffer. This prevents stale object graphs from being observable through any future partial walk of a recycled arena

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> {
Expand Down
33 changes: 30 additions & 3 deletions oscars/src/alloc/arena2/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand All @@ -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,
}
}
}
Expand All @@ -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) {
Expand Down Expand Up @@ -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(())
Expand All @@ -138,8 +159,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
}
}

Expand Down
65 changes: 65 additions & 0 deletions oscars/src/alloc/arena2/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this show us that the arena is recycled? How do we know it's recycled the same memory?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

heap_size() equality proves that the OS footprint didn't grow, but you're right it doesnt prove that the recycled slot was actually consumed. Sorry missed that 😅

I have added two recycled_count assertions to the test to verify that,
recycled_count == 1 after drop_dead_arenas() - confirms the arena was saved for reuse instead of being freed.
recycled_count == 0 after the second alloc loop - confirms initialize_new_arena reused that saved arena instead of asking the OS for new memory

Together these two checks confirm that the recycling process worked correctly.

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
Expand Down
23 changes: 23 additions & 0 deletions oscars/src/alloc/mempool3/alloc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
50 changes: 43 additions & 7 deletions oscars/src/alloc/mempool3/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 ()>,
}

Expand All @@ -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,
}
}
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -267,16 +297,22 @@ 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.
for pool in self.slot_pools.extract_if(.., |p| p.run_drop_check()) {
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());
Expand Down
61 changes: 61 additions & 0 deletions oscars/src/alloc/mempool3/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
1 change: 1 addition & 0 deletions oscars/src/collectors/mark_sweep/cell.rs
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,7 @@ impl<T: ?Sized> GcRefCell<T> {
}

// returns a raw pointer to the inner value or `None` if currently mutably borrowed
#[allow(dead_code)]
pub(crate) fn get_raw(&self) -> Option<*mut T> {
match self.borrow.get().borrowed() {
BorrowState::Writing => None,
Expand Down
9 changes: 3 additions & 6 deletions oscars/src/collectors/mark_sweep/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -163,10 +163,6 @@ impl MarkSweepGarbageCollector {
// memory so we can still inspect the trace color on ephemerons;
// use sweep_color since alive objects were marked with it.
self.sweep_trace_color(sweep_color);

// finally tell the allocator to reclaim raw OS memory
// from arenas that are completely empty now
self.allocator.borrow_mut().drop_empty_pools();
}

// Force drops all elements in the internal tracking queues and clears
Expand Down Expand Up @@ -235,8 +231,9 @@ impl MarkSweepGarbageCollector {
let new_color = sweep_color.flip();
self.trace_color.set(new_color);

// NOTE: It would actually be interesting to reuse the pools that are empty rather
// than drop the page and reallocate when a new page is needed ... TBD
// Reclaim OS memory from pool pages that became fully empty during the sweep above.
// Empty pool pages are parked in a recycle list rather than immediately freed to the OS,
// allowing the next try_alloc to pull from that list and avoid OS allocation thrashing.
self.allocator.borrow_mut().drop_empty_pools();

// Drain pending queues while `is_collecting` is still true so that any
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ impl<K: Trace, V: Trace> Ephemeron<K, V> {
}
}

#[allow(dead_code)] // TODO: figure out what to do with this
pub fn key(&self) -> &K {
self.key.value()
}
Expand Down
Loading