janis/curiOS
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

foundation: slab allocator

Browse source
This commit is contained in:
janis 2026-08-05 21:33:50 +02:00
parent 67551afc56
commit 86d80132be
Signed by: janis
SSH key fingerprint: SHA256:bB1qbbqmDXZNT0KKD5c2Dfjg53JGhj7B3CFcLIzSqq8
1 changed files with 205 additions and 101 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

272
crates/foundation/src/alloc/slab.rs
View file

@ -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<A: Allocator + Clone> {
element_size: usize,
/// Pointer to the first chunk in the slab.
head: Option<NonNull<SlabChunk<A>>>,
/// Pointer to the first chunk in the slab that is full.
full_head: Option<NonNull<SlabChunk<A>>>,
alloc: A,
_pd: core::marker::PhantomPinned,
}
impl<A: Allocator + Clone> Debug for Slab<A> {
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<A: Allocator + Clone> {
@ -39,7 +52,18 @@ enum SlotResult {
None,
}
enum FreeSlotResult {
Empty,
WasFull,
NotEmpty,
}
impl<A: Allocator + Clone> SlabChunk<A> {
unsafe fn slab_pinned_mut<'a>(&self) -> Pin<&'a mut Slab<A>> {
// SAFETY: `slab` is created from a `Pin<&mut Slab<A>>`
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<A: Allocator + Clone> SlabChunk<A> {
}
/// returns `true` if the slab is now empty and can be freed
fn push_free_slot(&self, slot: NonNull<u8>) -> bool {
fn push_free_slot(&self, slot: NonNull<u8>) -> FreeSlotResult {
let slot = slot.cast::<ChunkSlot>();
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<A: Allocator + Clone> Slab<A> {
"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::<SlabChunk<A>>(), self.element_size)
fn first_slot_offset(element_size: usize) -> usize {
mem::align_up(core::mem::size_of::<SlabChunk<A>>(), 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<SlabChunk<A>> {
fn alloc_chunk(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
// 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::<SlabChunk<A>>();
#[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::<ChunkSlot>();
for i in 0..(count - 1) {
@ -143,7 +186,8 @@ impl<A: Allocator + Clone> Slab<A> {
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<A: Allocator + Clone> Slab<A> {
}
#[cold]
fn alloc_chunk_cold(&mut self) -> NonNull<SlabChunk<A>> {
fn alloc_chunk_cold(self: Pin<&mut Self>) -> NonNull<SlabChunk<A>> {
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<A: Allocator + Clone> Slab<A> {
ptr.cast_slice(self.element_size)
}
fn free_slot(&mut self, slot: NonNull<u8>) {
let (_, layout) = self.count_and_layout();
let mut chunk = slot
fn free_slot(element_size: usize, slot: NonNull<u8>) {
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::<SlabChunk<A>>();
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 } {
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 { head.as_mut().next = chunk.next };
unsafe { (&raw mut (*head.as_ptr()).next).write(chunk.next) };
break;
}
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 {
self.alloc
.deallocate(NonNull::from_mut(chunk).cast(), layout)
};
} else if !linked {
chunk.next = self.head;
self.head = Some(chunk.into());
slab.as_mut().get_unchecked_mut().head = Some(chunk.into());
}
}
FreeSlotResult::NotEmpty => {
// chunk is not empty, and was not full, so nothing to do
}
}
}
}
impl<A: Allocator + Clone> Drop for Slab<A> {
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,7 +345,8 @@ pub struct SlabAllocator<A: Allocator + Clone> {
impl<A: Allocator + Clone> SlabAllocator<A> {
pub fn new(alloc: A) -> Self {
let slabs = [
Self {
slabs: [
Slab::new(16, alloc.clone()),
Slab::new(32, alloc.clone()),
Slab::new(64, alloc.clone()),
@ -251,9 +355,9 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
Slab::new(512, alloc.clone()),
Slab::new(1024, alloc.clone()),
Slab::new(2048, alloc.clone()),
];
Self { slabs, alloc }
],
alloc,
}
}
fn slab_index_for_size(size: usize) -> Option<usize> {
@ -268,6 +372,14 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
// 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<A: Allocator + Clone> SlabAllocator<A> {
}
}
pub fn alloc(&mut self, layout: Layout) -> Option<NonNull<[u8]>> {
pub fn alloc(self: Pin<&mut Self>, layout: Layout) -> Option<NonNull<[u8]>> {
#[cfg(test)]
std::eprintln!("SlabAllocator::alloc({:?})", layout);
if unlikely(layout.size() == 0) {
return Some(NonNull::dangling().cast_slice(0));
}
@ -284,9 +399,12 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
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<A: Allocator + Clone> SlabAllocator<A> {
}
}
pub fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
pub fn dealloc(self: Pin<&mut Self>, ptr: NonNull<u8>, layout: Layout) {
#[cfg(test)]
std::eprintln!("SlabAllocator::dealloc({:?})", layout);
if unlikely(layout.size() == 0) {
return;
}
@ -302,15 +423,8 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
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::<SlabChunk<A>>();
let mut slab = (&raw const (*chunk.as_ptr()).slab).read();
slab.as_mut().free_slot(ptr);
Some(slab_index) => unsafe {
Slab::<A>::free_slot(self.slabs.get_unchecked(slab_index).element_size, ptr);
},
None => {
// deallocate directly to the backing allocator
@ -320,30 +434,16 @@ impl<A: Allocator + Clone> SlabAllocator<A> {
}
}
impl<A: Allocator + Clone> Drop for SlabAllocator<A> {
fn drop(&mut self) {}
}
#[cfg(test)]
mod tests {
use super::*;
use core::{mem::forget, pin};
use std::{prelude::rust_2024::*, sync::Mutex};
struct SlabAllocatorWrapper<A: Allocator + Clone>(Mutex<SlabAllocator<A>>);
struct SlabAllocatorWrapper<'a, A: Allocator + Clone>(Mutex<Pin<&'a mut SlabAllocator<A>>>);
impl<A: Allocator + Clone> SlabAllocatorWrapper<A> {
fn new(alloc: A) -> Self {
Self(Mutex::new(SlabAllocator::new(alloc)))
}
}
impl SlabAllocatorWrapper<std::alloc::System> {
fn new_system() -> SlabAllocatorWrapper<std::alloc::System> {
SlabAllocatorWrapper::new(std::alloc::System)
}
}
unsafe impl<A: Allocator + Clone> Allocator for SlabAllocatorWrapper<A> {
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<u8>, 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;
drop(boxed);
std::eprintln!("test_slab_allocator: dropped boxed");
}
}
}