diff --git a/crates/foundation/src/alloc/slab.rs b/crates/foundation/src/alloc/slab.rs index 9fb65f2..b718591 100644 --- a/crates/foundation/src/alloc/slab.rs +++ b/crates/foundation/src/alloc/slab.rs @@ -1,8 +1,8 @@ //! A slab allocator -use core::{alloc::Layout, cell::Cell, hint::unlikely, mem::offset_of, num::NonZero, ptr::NonNull}; - -const UNLINKED: NonNull<()> = unsafe { NonNull::new_unchecked(!0 as *mut ()) }; +use core::{ + alloc::Layout, cell::Cell, fmt::Debug, hint::unlikely, num::NonZero, pin::Pin, ptr::NonNull, +}; use liballoc::alloc::Allocator; @@ -15,7 +15,20 @@ pub struct Slab { element_size: usize, /// Pointer to the first chunk in the slab. head: Option>>, + /// Pointer to the first chunk in the slab that is full. + full_head: Option>>, alloc: A, + _pd: core::marker::PhantomPinned, +} + +impl Debug for Slab { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("Slab") + .field("element_size", &self.element_size) + .field("head", &self.head) + .field("full_head", &self.full_head) + .finish() + } } struct SlabChunk { @@ -39,7 +52,18 @@ enum SlotResult { None, } +enum FreeSlotResult { + Empty, + WasFull, + NotEmpty, +} + impl SlabChunk { + unsafe fn slab_pinned_mut<'a>(&self) -> Pin<&'a mut Slab> { + // SAFETY: `slab` is created from a `Pin<&mut Slab>` + unsafe { Pin::new_unchecked(self.slab.as_ptr().as_mut_unchecked()) } + } + fn pop_free_slot(&self) -> SlotResult { let Some(slot) = self.free.get() else { return SlotResult::None; @@ -55,14 +79,18 @@ impl SlabChunk { } /// returns `true` if the slab is now empty and can be freed - fn push_free_slot(&self, slot: NonNull) -> bool { + fn push_free_slot(&self, slot: NonNull) -> FreeSlotResult { let slot = slot.cast::(); let next = self.free.get(); unsafe { slot.as_ptr().write(ChunkSlot(next)) }; self.free.set(Some(slot)); self.count.update(|count| count - 1); - self.count.get() == 0 + match (next, self.count.get()) { + (None, _) => FreeSlotResult::WasFull, + (_, 0) => FreeSlotResult::Empty, + _ => FreeSlotResult::NotEmpty, + } } } @@ -79,57 +107,72 @@ impl Slab { "element_size must be a power of two" ); - Self { + let slab = Self { element_size, head: None, - alloc, - } + full_head: None, + alloc: alloc.clone(), + _pd: core::marker::PhantomPinned, + }; + + #[cfg(test)] + std::eprintln!("Slab::new({element_size}) -> {slab:#?}",); + + slab } - fn first_slot_offset(&self) -> usize { - mem::align_up(core::mem::size_of::>(), self.element_size) + fn first_slot_offset(element_size: usize) -> usize { + mem::align_up(core::mem::size_of::>(), element_size) } - fn count_and_layout(&self) -> (usize, Layout) { + fn count_and_layout(element_size: usize) -> (usize, Layout) { let (count, size, align) = { - let one_page_count = (PAGE_SIZE - self.first_slot_offset()) / self.element_size; + let one_page_count = (PAGE_SIZE - Self::first_slot_offset(element_size)) / element_size; if one_page_count < 3 { let count = 3; - let size = - (self.first_slot_offset() + count * self.element_size).next_power_of_two(); + let size = (Self::first_slot_offset(element_size) + count * element_size) + .next_power_of_two(); assert!(size.is_multiple_of(PAGE_SIZE)); assert!(size >= PAGE_SIZE); - assert!(size.is_multiple_of(self.element_size)); + //assert!(size.is_multiple_of(self.element_size)); (count, size, size) } else { let count = one_page_count; let size = PAGE_SIZE; - (count, size, self.element_size) + let align = element_size.max(PAGE_SIZE); + (count, size, align) } }; + #[cfg(test)] + std::eprintln!("element_size: {element_size} count: {count}, size: {size}, align: {align}"); + (count, unsafe { Layout::from_size_align_unchecked(size, align) }) } - fn alloc_chunk(&mut self) -> NonNull> { + fn alloc_chunk(self: Pin<&mut Self>) -> NonNull> { // we want to limit chunks to 1 page unless the element size is so // large that we can fit fewer than 3 elements in a page. - let (count, layout) = self.count_and_layout(); + let (count, layout) = Self::count_and_layout(self.element_size); let Some(bytes) = self.alloc.allocate(layout).ok() else { panic!() }; let chunk = bytes.as_non_null_ptr().cast::>(); + + #[cfg(test)] + std::eprintln!("Slab::alloc_chunk(chunk: {chunk:#?}, layout: {layout:?})"); + unsafe { let first_slot = chunk .as_ptr() - .byte_add(self.first_slot_offset()) + .byte_add(Self::first_slot_offset(self.element_size)) .cast::(); for i in 0..(count - 1) { @@ -143,7 +186,8 @@ impl Slab { chunk.write(SlabChunk { next: self.head, - slab: NonNull::from(self), + // SAFETY: we only access `slab` via `SlabChunk::slab_pinned_mut` + slab: NonNull::from(Pin::into_inner_unchecked(self)), free: Cell::new(Some(NonNull::new_unchecked(first_slot))), count: Cell::new(0), }); @@ -153,26 +197,37 @@ impl Slab { } #[cold] - fn alloc_chunk_cold(&mut self) -> NonNull> { + fn alloc_chunk_cold(self: Pin<&mut Self>) -> NonNull> { self.alloc_chunk() } - fn alloc_slot(&mut self) -> NonNull<[u8]> { - let mut chunk = match self.head { + fn alloc_slot(mut self: Pin<&mut Self>) -> NonNull<[u8]> { + #[cfg(test)] + std::eprintln!("Slab::alloc_slot({self:?})"); + + let chunk = match self.head { Some(chunk) => chunk, None => { - let chunk = self.alloc_chunk_cold(); - self.head = Some(chunk); + let chunk = self.as_mut().alloc_chunk_cold(); + unsafe { + self.as_mut().get_unchecked_mut().head = Some(chunk); + } chunk } }; - let chunk = unsafe { chunk.as_mut() }; + let chunk = unsafe { chunk.as_ptr().as_mut_unchecked() }; let ptr = match chunk.pop_free_slot() { SlotResult::Some(non_null) => non_null, SlotResult::Last(non_null) => { - self.head = chunk.next.replace(UNLINKED.cast()); + unsafe { + let mut_ref = self.as_mut().get_unchecked_mut(); + mut_ref.head = chunk.next.take(); + chunk.next = mut_ref.full_head; + mut_ref.full_head = Some(chunk.into()); + } + non_null } SlotResult::None => { @@ -183,49 +238,97 @@ impl Slab { ptr.cast_slice(self.element_size) } - fn free_slot(&mut self, slot: NonNull) { - let (_, layout) = self.count_and_layout(); - let mut chunk = slot + fn free_slot(element_size: usize, slot: NonNull) { + let (_, layout) = Self::count_and_layout(element_size); + + let chunk_ptr = slot .map_addr(|addr| unsafe { NonZero::new_unchecked(mem::align_down(addr.get(), layout.align())) }) .cast::>(); - let chunk = unsafe { chunk.as_mut() }; - let linked = chunk.next != Some(UNLINKED.cast()); + #[cfg(test)] + std::eprintln!("Slab::free_slot(chunk: {chunk_ptr:#?})"); - if chunk.push_free_slot(slot) { - if linked { - let mut head = self.head.expect("chunk is linked, so head exists"); - while let Some(next) = unsafe { head.as_ref().next } { - if next == chunk.into() { - unsafe { head.as_mut().next = chunk.next }; - break; + let chunk = unsafe { chunk_ptr.as_ptr().as_mut_unchecked() }; + + let mut slab = unsafe { chunk.slab_pinned_mut() }; + + match chunk.push_free_slot(slot) { + FreeSlotResult::Empty => { + // chunk is empty: + // first unlink the chunk from the slab.. + if slab.head == Some(chunk.into()) { + unsafe { slab.as_mut().get_unchecked_mut().head = chunk.next }; + } else { + let mut head = slab.head.expect("chunk is linked, so head exists"); + while let Some(next) = unsafe { (&raw const (*head.as_ptr()).next).read() } { + if next == chunk.into() { + unsafe { (&raw mut (*head.as_ptr()).next).write(chunk.next) }; + break; + } + head = next; } - head = next; + } + + // ..then free it. + #[cfg(test)] + std::eprintln!("free_slot::drop({chunk_ptr:?}, layout: {layout:?})"); + + unsafe { slab.alloc.deallocate(chunk_ptr.cast(), layout) }; + } + FreeSlotResult::WasFull => { + // chunk was full: + + // unlink it from the full list + if slab.full_head == Some(chunk.into()) { + unsafe { slab.as_mut().get_unchecked_mut().full_head = chunk.next }; + } else { + let mut head = slab.full_head.expect("chunk is linked, so head exists"); + while let Some(next) = unsafe { (&raw const (*head.as_ptr()).next).read() } { + if next == chunk.into() { + unsafe { (&raw mut (*head.as_ptr()).next).write(chunk.next) }; + break; + } + head = next; + } + } + + // link it to the head of the slab + chunk.next = slab.head; + unsafe { + slab.as_mut().get_unchecked_mut().head = Some(chunk.into()); } } - - unsafe { - self.alloc - .deallocate(NonNull::from_mut(chunk).cast(), layout) - }; - } else if !linked { - chunk.next = self.head; - self.head = Some(chunk.into()); + FreeSlotResult::NotEmpty => { + // chunk is not empty, and was not full, so nothing to do + } } } } impl Drop for Slab { fn drop(&mut self) { - let (_, layout) = self.count_and_layout(); + let (_, layout) = Self::count_and_layout(self.element_size); let mut chunk = self.head.take(); while let Some(mut chunk_ptr) = chunk { let chunk_ref = unsafe { chunk_ptr.as_mut() }; chunk = chunk_ref.next.take(); + #[cfg(test)] + std::eprintln!("Slab::drop({chunk_ptr:?}, layout: {layout:?})"); + + unsafe { + self.alloc.deallocate(chunk_ptr.cast(), layout); + } + } + + let mut chunk = self.full_head.take(); + while let Some(mut chunk_ptr) = chunk { + let chunk_ref = unsafe { chunk_ptr.as_mut() }; + chunk = chunk_ref.next.take(); + unsafe { self.alloc.deallocate(chunk_ptr.cast(), layout); } @@ -242,18 +345,19 @@ pub struct SlabAllocator { impl SlabAllocator { pub fn new(alloc: A) -> Self { - let slabs = [ - Slab::new(16, alloc.clone()), - Slab::new(32, alloc.clone()), - Slab::new(64, alloc.clone()), - Slab::new(128, alloc.clone()), - Slab::new(256, alloc.clone()), - Slab::new(512, alloc.clone()), - Slab::new(1024, alloc.clone()), - Slab::new(2048, alloc.clone()), - ]; - - Self { slabs, alloc } + Self { + slabs: [ + Slab::new(16, alloc.clone()), + Slab::new(32, alloc.clone()), + Slab::new(64, alloc.clone()), + Slab::new(128, alloc.clone()), + Slab::new(256, alloc.clone()), + Slab::new(512, alloc.clone()), + Slab::new(1024, alloc.clone()), + Slab::new(2048, alloc.clone()), + ], + alloc, + } } fn slab_index_for_size(size: usize) -> Option { @@ -268,6 +372,14 @@ impl SlabAllocator { // subtracting 4 (since 2^4 = 16) let index = (size.next_power_of_two().trailing_zeros() - 16usize.trailing_zeros()) as usize; + #[cfg(test)] + std::eprintln!( + "slab_index_for_size: size: {size}, next_power_of_two: {}, trailing_zeros: {}, index: {}", + size.next_power_of_two(), + size.next_power_of_two().trailing_zeros(), + index, + ); + // we have 8 slabs, so the index must be less than 8 if index < SLAB_ALLOCATOR_BUCKETS { Some(index) @@ -276,7 +388,10 @@ impl SlabAllocator { } } - pub fn alloc(&mut self, layout: Layout) -> Option> { + pub fn alloc(self: Pin<&mut Self>, layout: Layout) -> Option> { + #[cfg(test)] + std::eprintln!("SlabAllocator::alloc({:?})", layout); + if unlikely(layout.size() == 0) { return Some(NonNull::dangling().cast_slice(0)); } @@ -284,9 +399,12 @@ impl SlabAllocator { let size = layout.size().max(layout.align()); match Self::slab_index_for_size(size) { - Some(slab_index) => { - Some(unsafe { self.slabs.get_unchecked_mut(slab_index).alloc_slot() }) - } + Some(slab_index) => unsafe { + Some( + self.map_unchecked_mut(|alloc| alloc.slabs.get_unchecked_mut(slab_index)) + .alloc_slot(), + ) + }, None => { // allocate directly from the backing allocator self.alloc.allocate(layout).ok() @@ -294,7 +412,10 @@ impl SlabAllocator { } } - pub fn dealloc(&mut self, ptr: NonNull, layout: Layout) { + pub fn dealloc(self: Pin<&mut Self>, ptr: NonNull, layout: Layout) { + #[cfg(test)] + std::eprintln!("SlabAllocator::dealloc({:?})", layout); + if unlikely(layout.size() == 0) { return; } @@ -302,15 +423,8 @@ impl SlabAllocator { let size = layout.size().max(layout.align()); match Self::slab_index_for_size(size) { - Some(_) => unsafe { - let chunk = ptr - .map_addr(|addr| { - NonZero::new_unchecked(mem::align_down(addr.get(), layout.align())) - }) - .cast::>(); - - let mut slab = (&raw const (*chunk.as_ptr()).slab).read(); - slab.as_mut().free_slot(ptr); + Some(slab_index) => unsafe { + Slab::::free_slot(self.slabs.get_unchecked(slab_index).element_size, ptr); }, None => { // deallocate directly to the backing allocator @@ -320,30 +434,16 @@ impl SlabAllocator { } } -impl Drop for SlabAllocator { - fn drop(&mut self) {} -} - #[cfg(test)] mod tests { use super::*; + use core::{mem::forget, pin}; use std::{prelude::rust_2024::*, sync::Mutex}; - struct SlabAllocatorWrapper(Mutex>); + struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex>>); - impl SlabAllocatorWrapper { - fn new(alloc: A) -> Self { - Self(Mutex::new(SlabAllocator::new(alloc))) - } - } - impl SlabAllocatorWrapper { - fn new_system() -> SlabAllocatorWrapper { - SlabAllocatorWrapper::new(std::alloc::System) - } - } - - unsafe impl Allocator for SlabAllocatorWrapper { + unsafe impl<'a, A: Allocator + Clone> Allocator for SlabAllocatorWrapper<'a, A> { fn allocate( &self, layout: core::alloc::Layout, @@ -351,29 +451,33 @@ mod tests { self.0 .lock() .unwrap() + .as_mut() .alloc(layout) .ok_or(core::alloc::AllocError) } unsafe fn deallocate(&self, ptr: NonNull, layout: core::alloc::Layout) { - self.0.lock().unwrap().dealloc(ptr, layout) + self.0.lock().unwrap().as_mut().dealloc(ptr, layout) } } #[test] fn test_slab_allocator() { - let alloc = SlabAllocatorWrapper::new_system(); + let mut slab = pin::pin!(SlabAllocator::new(std::alloc::Global)); + let alloc = SlabAllocatorWrapper(Mutex::new(slab.as_mut())); - let layout = Layout::from_size_align(16, 8).unwrap(); + let layout = Layout::from_size_align(4, 4).unwrap(); let ptr = alloc.allocate(layout).unwrap(); assert_eq!(ptr.len(), 16); assert_eq!(ptr.addr().get() % 8, 0); - let mut boxed = Box::new_in(42u32, &alloc); - assert_eq!(*boxed, 42); - assert_ne!(Box::as_non_null(&mut boxed), ptr.cast()); - - unsafe { alloc.deallocate(ptr.cast(), layout) }; - _ = boxed; + { + let mut boxed = Box::new_in(42u32, &alloc); + assert_eq!(*boxed, 42); + assert_ne!(Box::as_non_null(&mut boxed), ptr.cast()); + unsafe { alloc.deallocate(ptr.cast(), layout) }; + drop(boxed); + std::eprintln!("test_slab_allocator: dropped boxed"); + } } }