diff --git a/libkernel/src/memory/allocators/frame.rs b/libkernel/src/memory/allocators/frame.rs index 736b77f..de4a9e1 100644 --- a/libkernel/src/memory/allocators/frame.rs +++ b/libkernel/src/memory/allocators/frame.rs @@ -1,6 +1,8 @@ use crate::memory::page::PageFrame; use intrusive_collections::{LinkedListLink, UnsafeRef, intrusive_adapter}; +use super::slab::slab::Slab; + #[derive(Clone, Copy, Debug)] pub struct AllocatedInfo { /// Current ref count of the allocated block. @@ -26,6 +28,8 @@ pub enum FrameState { AllocatedHead(AllocatedInfo), /// The frame is a tail page of an allocated block. AllocatedTail(TailInfo), + /// The frame is being used by the slab allocator. + Slab(Slab), /// The frame is part of the kernel's own image. Kernel, } diff --git a/libkernel/src/memory/allocators/phys.rs b/libkernel/src/memory/allocators/phys.rs index 865baa3..18ccd67 100644 --- a/libkernel/src/memory/allocators/phys.rs +++ b/libkernel/src/memory/allocators/phys.rs @@ -2,7 +2,13 @@ use crate::{ CpuOps, error::{KernelError, Result}, memory::{ - PAGE_SHIFT, address::AddressTranslator, allocators::frame::FrameState, page::PageFrame, + PAGE_SHIFT, + address::AddressTranslator, + allocators::{ + frame::FrameState, + slab::{SLAB_FRAME_ALLOC_ORDER, SLAB_SIZE_BYTES}, + }, + page::PageFrame, region::PhysMemoryRegion, }, sync::spinlock::SpinLockIrq, @@ -16,6 +22,7 @@ use log::info; use super::{ frame::{AllocatedInfo, Frame, FrameAdapter, FrameList, TailInfo}, + slab::slab::Slab, smalloc::Smalloc, }; @@ -23,13 +30,29 @@ use super::{ // 2^MAX_ORDER pages. const MAX_ORDER: usize = 10; -struct FrameAllocatorInner { +pub(super) struct FrameAllocatorInner { frame_list: FrameList, free_pages: usize, free_lists: [LinkedList; MAX_ORDER + 1], } impl FrameAllocatorInner { + pub(super) fn free_slab(&mut self, frame: UnsafeRef) { + assert!(matches!(frame.state, FrameState::Slab(_))); + + // SAFETY: The caller should guarntee exclusive ownership of this frame + // as it's being passed back to the FA. + let frame = unsafe { &mut *UnsafeRef::into_raw(frame) }; + + // Restore frame state data for slabs. + frame.state = FrameState::AllocatedHead(AllocatedInfo { + ref_count: 1, + order: SLAB_FRAME_ALLOC_ORDER as _, + }); + + self.free_frames(PhysMemoryRegion::new(frame.pfn.pa(), SLAB_SIZE_BYTES)); + } + /// Frees a previously allocated block of frames. /// The PFN can point to any page within the allocated block. fn free_frames(&mut self, region: PhysMemoryRegion) { @@ -129,7 +152,7 @@ impl FrameAllocatorInner { } pub struct FrameAllocator { - inner: SpinLockIrq, + pub(super) inner: SpinLockIrq, } pub struct PageAllocation<'a, CPU: CpuOps> { @@ -147,6 +170,25 @@ impl PageAllocation<'_, CPU> { pub fn region(&self) -> &PhysMemoryRegion { &self.region } + + /// Leak the allocation as a slab allocation, for it to be picked back up + /// again once the slab has been free'd. + /// + /// Returns an `UnsafeRef` to `Frame` that was converted to a slab for use + /// in the slab lists. + pub(super) fn into_slab(self, slab_info: Slab) -> *const Frame { + let mut inner = self.inner.lock_save_irq(); + + let frame = inner.get_frame_mut(self.region.start_address().to_pfn()); + + debug_assert!(matches!(frame.state, FrameState::AllocatedHead(_))); + + frame.state = FrameState::Slab(slab_info); + + self.leak(); + + frame as _ + } } impl Clone for PageAllocation<'_, CPU> { @@ -294,7 +336,7 @@ impl FrameAllocator { /// # Safety /// It's unsafe because it deals with raw pointers and takes ownership of /// the metadata memory. It should only be called once. - pub unsafe fn init>(mut smalloc: Smalloc) -> Self { + pub unsafe fn init>(mut smalloc: Smalloc) -> (Self, FrameList) { let highest_addr = smalloc .iter_memory() .map(|r| r.end_address()) @@ -379,9 +421,12 @@ impl FrameAllocator { allocator.free_pages ); - FrameAllocator { - inner: SpinLockIrq::new(allocator), - } + ( + FrameAllocator { + inner: SpinLockIrq::new(allocator), + }, + frame_list, + ) } } @@ -394,7 +439,9 @@ pub mod tests { use super::*; use crate::{ memory::{ - address::{IdentityTranslator, PA}, allocators::smalloc::RegionList, region::PhysMemoryRegion + address::{IdentityTranslator, PA}, + allocators::smalloc::RegionList, + region::PhysMemoryRegion, }, test::MockCpuOps, }; @@ -407,10 +454,14 @@ pub mod tests { pub struct TestFixture { pub allocator: FrameAllocator, + pub frame_list: FrameList, base_ptr: *mut u8, layout: Layout, } + unsafe impl Send for TestFixture {} + unsafe impl Sync for TestFixture {} + impl TestFixture { /// Creates a new test fixture. /// @@ -459,10 +510,11 @@ pub mod tests { .unwrap(); } - let allocator = unsafe { FrameAllocator::init(smalloc) }; + let (allocator, frame_list) = unsafe { FrameAllocator::init(smalloc) }; Self { allocator, + frame_list, base_ptr, layout, } @@ -550,7 +602,13 @@ pub mod tests { // The middle pages should be marked as Kernel let reserved_pfn = PageFrame::from_pfn( - fixture.allocator.inner.lock_save_irq().frame_list.base_page().value() + fixture + .allocator + .inner + .lock_save_irq() + .frame_list + .base_page() + .value() + (pages_in_max_block * 2 + 4), ); diff --git a/libkernel/src/memory/allocators/slab/allocator.rs b/libkernel/src/memory/allocators/slab/allocator.rs new file mode 100644 index 0000000..c8b504b --- /dev/null +++ b/libkernel/src/memory/allocators/slab/allocator.rs @@ -0,0 +1,499 @@ +/// A slab allocator for Moss. +use super::{ + SLAB_FRAME_ALLOC_ORDER, SLAB_MAX_OBJ_SHIFT, alloc_order, + slab::{Slab, SlabState}, +}; +use crate::{ + CpuOps, + memory::{ + address::{AddressTranslator, VA}, + allocators::{ + frame::{Frame, FrameAdapter, FrameList, FrameState}, + phys::PageAllocGetter, + }, + }, + sync::spinlock::SpinLockIrq, +}; +use core::marker::PhantomData; +use intrusive_collections::{LinkedList, UnsafeRef}; + +const MAX_FREE_SLABS: usize = 32; + +/// Slab manager for a specific size class. +/// +/// Manages a collection of slabs for a particular object size (size class). Two +/// main linked lists are managed: a 'free' list and a 'partial' list. +/// +/// - The 'partial' list is a collection of slabs which are partially full. We +/// allocate from this list first for new allocations. +/// - The 'free' list is a collection of slabs which have no objects allocated +/// from them yet. We cache them in the hope that they will be used later on, +/// without the need to lock the FrameAllocator (FA) to get more physical +/// memory. When a particular size is reached (`MAX_FREE_SLABS`), we batch free +/// half of the slabs back to the FA. +/// +/// # Full Slabs +/// +/// There is no 'full' list. Full slabs are unlinked from both the 'partial' and +/// 'free' lists. They are allowed to float "in the ether" (referenced only by +/// the global [FrameList]). When freeing an object from a 'full' slab, the +/// allocator detects the state transition and re-links the frame into the +/// 'partial'/'free' list. +/// +/// # Safety and Ownership +/// +/// The FA and the `SlabManager` share a list of frame metadata via [FrameList]. To +/// share this list safely, we implement an implicit ownership model: +/// +/// 1. When the FA allocates a frame, it initializes the metadata. +/// 2. We convert this frame into a slab allocation via +/// [PageAllocation::as_slab]. +/// 3. Once that function returns, this `SlabManager` is considered the +/// exclusive owner of that frame's metadata. +/// +/// It is therefore safe for the `SlabManager` to obtain a mutable reference to +/// the metadata of any frame in its possession, as the `SpinLock` protecting +/// this struct guarantees exclusive access to the specific size class that +/// "owns" the frame. Ownership is eventually returned to the FA via +/// [FrameAllocatorInner::free_slab]. +pub struct SlabManager, T: AddressTranslator<()>> { + free: LinkedList, + partial: LinkedList, + free_list_sz: usize, + obj_shift: usize, + frame_list: FrameList, + phantom1: PhantomData, + phantom2: PhantomData, + phantom3: PhantomData, +} + +impl, T: AddressTranslator<()>> SlabManager { + fn new(obj_shift: usize, frame_list: FrameList) -> Self { + Self { + free: LinkedList::new(FrameAdapter::new()), + partial: LinkedList::new(FrameAdapter::new()), + free_list_sz: 0, + obj_shift, + frame_list, + phantom1: PhantomData, + phantom2: PhantomData, + phantom3: PhantomData, + } + } + + /// Try to allocate a new object using free and partial slabs. Does *not* + /// allocate any physical memory for the allocation. + /// + /// # Returns + /// - `None` if there are no free or patial slabs available. + /// - `Some(ptr)` if the allocation was successful. + pub fn try_alloc(&mut self) -> Option<*mut u8> { + // Lets start with partial list. + if let Some(frame) = self.partial.pop_front().map(|x| { + // SAFETY: We hold a mutable reference for self, therefore we can be + // sure that we have exclusive access to all Frames owned by this + // SlabAllocInner. + unsafe { &mut *UnsafeRef::into_raw(x) } + }) && let FrameState::Slab(ref mut slab) = frame.state + && let Some(ptr) = slab.alloc_object() + { + if slab.state() == SlabState::Partial { + // If the slab is still partial, re-insert back into the partial + // list for further allocations. + self.partial + .push_front(unsafe { UnsafeRef::from_raw(frame as *mut _) }); + } + + return Some(ptr); + } + + if let Some(frame) = self.free.pop_front().map(|x| { + // SAFETY: As above. + unsafe { &mut *UnsafeRef::into_raw(x) } + }) { + let (ptr, state) = match frame.state { + FrameState::Slab(ref mut slab) => { + let ptr = slab.alloc_object().unwrap(); + let state = slab.state(); + + (ptr, state) + } + _ => unreachable!("Frame state should be slab"), + }; + + if state == SlabState::Partial { + // SAFETY: The frame is now owned by the list and no other refs + // to it will exist. + self.partial + .push_front(unsafe { UnsafeRef::from_raw(frame as *const _) }); + } + + return Some(ptr); + } + + None + } + + /// Allocate an object for the given size class. Uses up partial and free + /// slabs first; if none are avilable allocate a new slab from the frame + /// allocator. + pub fn alloc(&mut self) -> *mut u8 { + // Fast path, first. + if let Some(ptr) = self.try_alloc() { + return ptr; + } + + // Slow path, allocate a new frame. + let new_alloc = A::global_page_alloc() + .alloc_frames(SLAB_FRAME_ALLOC_ORDER as _) + .expect("OOM - cannot allocate physical frame"); + + let mut slab = Slab::new::(&new_alloc, self.obj_shift); + + let obj = slab.alloc_object().expect("Slab should be empty"); + let state = slab.state(); + let frame = new_alloc.into_slab(slab); + + // We now have ownership of the frame. + if state == SlabState::Partial { + self.partial + // SAFETY: Since we called `as_slab` above, we now have + // exclusive ownership of the page frame. + .push_front(unsafe { UnsafeRef::from_raw(frame) }); + } + + obj + } + + /// Free the given allocation. + pub fn free(&mut self, ptr: *mut u8) { + // Find the frame. + let va = VA::from_ptr_mut(ptr.cast()); + + let frame = self.frame_list.get_frame(va.to_pa::().to_pfn()); + + let (frame, state) = { + // Get the slab allocation data for this object. + // + // SAFETY: Since we hold the lock for slabs of this size class (&mut + // self), we are guaranteed exclusive ownership of all slab frames + // of this object size. + fn do_free_obj( + frame: *mut Frame, + ptr: *mut u8, + frame_list: &FrameList, + obj_shift: usize, + ) -> (*mut Frame, SlabState) { + match (unsafe { &mut (*frame) }).state { + FrameState::AllocatedTail(ref tail_info) => { + let head_frame = frame_list.get_frame(tail_info.head); + do_free_obj(head_frame, ptr, frame_list, obj_shift) + } + FrameState::Slab(ref mut slab) => { + if slab.obj_shift() != obj_shift { + panic!("Slab allocator: Layout mismatch on free"); + } + slab.put_object(ptr); + (frame, slab.state()) + } + _ => unreachable!("Slab allocation"), + } + } + + do_free_obj(frame, ptr, &self.frame_list, self.obj_shift) + }; + + // SAFETY: As above + let frame = unsafe { &mut *frame }; + + match state { + SlabState::Free => { + if self.free_list_sz == MAX_FREE_SLABS { + let mut num_freed = 0; + let mut fa = A::global_page_alloc().inner.lock_save_irq(); + + // batch free some free slabs. + for _ in 0..(MAX_FREE_SLABS >> 1) { + let frame = self.free.pop_front().expect("Should have free slabs"); + + fa.free_slab(frame); + + num_freed += 1; + } + + self.free_list_sz -= num_freed; + } + + if frame.link.is_linked() { + // The frame *must* be linked in the partial list if linked. + unsafe { self.partial.cursor_mut_from_ptr(frame as _) }.remove(); + } + + self.free + .push_front(unsafe { UnsafeRef::from_raw(frame as _) }); + + self.free_list_sz += 1; + } + SlabState::Partial => { + if !frame.link.is_linked() { + // We must have free'd an object on a previously full slab. + // Insert into the partial list. + self.partial + .push_front(unsafe { UnsafeRef::from_raw(frame as _) }); + } + } + SlabState::Full => unreachable!("we've just free'd an object"), + } + } +} + +pub struct SlabAllocator, T: AddressTranslator<()>> { + managers: [SpinLockIrq, CPU>; SLAB_MAX_OBJ_SHIFT as usize + 1], +} + +unsafe impl, T: AddressTranslator<()>> Send + for SlabAllocator +{ +} +unsafe impl, T: AddressTranslator<()>> Sync + for SlabAllocator +{ +} + +impl, T: AddressTranslator<()>> SlabAllocator { + pub fn new(frame_list: FrameList) -> Self { + Self { + managers: core::array::from_fn(|n| { + SpinLockIrq::new(SlabManager::new(n, frame_list.clone())) + }), + } + } + + pub fn allocator_for_layout( + &self, + layout: core::alloc::Layout, + ) -> Option<&SpinLockIrq, CPU>> { + Some(&self.managers[alloc_order(layout)?]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + memory::{ + address::{IdentityTranslator, PA}, + allocators::{ + phys::{FrameAllocator, tests::TestFixture}, + slab::SLAB_SIZE_BYTES, + }, + }, + sync::once_lock::OnceLock, + test::MockCpuOps, + }; + use core::alloc::Layout; + + type TstSlabAlloc = SlabAllocator; + + static FIXTURE: OnceLock = OnceLock::new(); + + struct SlabTestAllocGetter {} + + impl PageAllocGetter for SlabTestAllocGetter { + fn global_page_alloc() -> &'static FrameAllocator { + &FIXTURE.get().expect("Test not initalised").allocator + } + } + + /// Initializes the global allocator for the test suite. + fn init_allocator() -> &'static TestFixture { + FIXTURE.get_or_init(|| { + // Allocate 32MB for the test heap to ensure we don't run out during large file tests + TestFixture::new(&[(0, 32 * 1024 * 1024)], &[]) + }) + } + + fn create_allocator_fixture() -> TstSlabAlloc { + let fixture = init_allocator(); + + let frame_list = fixture.frame_list.clone(); + + SlabAllocator::new(frame_list) + } + + #[test] + fn alloc_free_basic() { + let allocator = create_allocator_fixture(); + + // 64-byte allocation + let layout = Layout::from_size_align(64, 64).unwrap(); + + unsafe { + let alloc = allocator.allocator_for_layout(layout).unwrap(); + let ptr = alloc.lock_save_irq().alloc(); + assert!(!ptr.is_null()); + assert_eq!(ptr as usize % 64, 0, "Alignment not respected"); + + // Write to it to ensure valid pointer. + *ptr = 0xAA; + + alloc.lock_save_irq().free(ptr); + } + } + + #[test] + fn slab_creation_and_partial_list() { + let allocator = create_allocator_fixture(); + + // 1024 byte allocation. + let layout = Layout::from_size_align(1024, 1024).unwrap(); + let alloc = allocator.allocator_for_layout(layout).unwrap(); + + // Initial State: No slabs + { + let inner = alloc.lock_save_irq(); + assert!(inner.partial.is_empty()); + assert!(inner.free.is_empty()); + } + + // Alloc one object + let ptr = alloc.lock_save_irq().alloc(); + + { + let inner = alloc.lock_save_irq(); + // Should now have 1 slab in partial + assert_eq!(inner.partial.iter().count(), 1); + assert!(inner.free.is_empty()); + + assert!(matches!( + unsafe { + &(*inner + .frame_list + .get_frame(PA::from_value(ptr as usize).to_pfn())) + } + .state, + FrameState::Slab(_) + )); + } + + // Free the object + alloc.lock_save_irq().free(ptr); + + { + let inner = alloc.lock_save_irq(); + // Should move to Free list + assert!(inner.partial.is_empty()); + assert_eq!(inner.free.iter().count(), 1); + } + } + + #[test] + fn slab_exhaustion_and_floating_slabs() { + let allocator = create_allocator_fixture(); + + // Slab Capacity = 4 objects at 4k. + let layout = Layout::from_size_align(4096, 4096).unwrap(); + let alloc = allocator.allocator_for_layout(layout).unwrap(); + + let mut ptrs = Vec::new(); + + // Alloc 4 objects (Fill the first slab) + { + let mut alloc = alloc.lock_save_irq(); + for _ in 0..4 { + ptrs.push(alloc.alloc()); + } + } + + { + let mut inner = alloc.lock_save_irq(); + // The slab is now full. It should have been removed from partial. + assert!(inner.partial.is_empty()); + assert!(inner.free.is_empty()); + + // try_alloc should return `None`. + assert!(inner.try_alloc().is_none()); + } + + // Alloc 1 more object (Triggers new slab) + let ptr_new = alloc.lock_save_irq().alloc(); + ptrs.push(ptr_new); + + { + // Now we have a second slab, which is Partial (1/4 used) + assert_eq!(alloc.lock_save_irq().partial.iter().count(), 1); + } + + // Free an object from the first (Full/Floating) slab. + let ptr_full_slab = ptrs[0]; + alloc.lock_save_irq().free(ptr_full_slab); + + { + // Both slabs should now be in partial. + assert_eq!(alloc.lock_save_irq().partial.iter().count(), 2); + } + } + + #[test] + fn batch_freeing_threshold() { + let allocator = create_allocator_fixture(); + + let layout = Layout::from_size_align(64, 64).unwrap(); + let mut alloc = allocator + .allocator_for_layout(layout) + .unwrap() + .lock_save_irq(); + + let mut all_ptrs = Vec::new(); + + // Create 33 separate slabs. + let objs_per_slab = SLAB_SIZE_BYTES / 64; + + // Allocate 33 * 256 objects + for _ in 0..(MAX_FREE_SLABS + 1) { + for _ in 0..objs_per_slab { + all_ptrs.push(alloc.alloc()); + } + } + + // At this point, we have 33 full slabs. None are in partial/free lists. + assert!(alloc.partial.is_empty()); + assert!(alloc.free.is_empty()); + + // Free everything. + // When we free the last object of a slab, it goes to Free list. + // When Free list hits 32, the *next* one triggers a batch free (pops 16). + for ptr in all_ptrs { + alloc.free(ptr); + } + + assert_eq!(alloc.free_list_sz, 17); + assert_eq!(alloc.free.iter().count(), 17); + } + + #[test] + #[should_panic(expected = "Layout mismatch")] + fn layout_mismatch_panic() { + let allocator = create_allocator_fixture(); + + // Alloc with size 64 + let layout_alloc = Layout::from_size_align(64, 64).unwrap(); + // Free with size 32 (Different lock, different inner allocator) + let layout_free = Layout::from_size_align(32, 32).unwrap(); + + let mut alloc_alloc = allocator + .allocator_for_layout(layout_alloc) + .unwrap() + .lock_save_irq(); + let mut alloc_free = allocator + .allocator_for_layout(layout_free) + .unwrap() + .lock_save_irq(); + + let ptr = alloc_alloc.alloc(); + // This should panic because the slab metadata inside the page + // says "Size 64", but we are calling free on the "Size 32" inner allocator. + // The code has a check: `if slab.obj_shift() != obj_shift { panic! }` + alloc_free.free(ptr); + } +} diff --git a/libkernel/src/memory/allocators/slab/mod.rs b/libkernel/src/memory/allocators/slab/mod.rs index 773e542..3a7e9c3 100644 --- a/libkernel/src/memory/allocators/slab/mod.rs +++ b/libkernel/src/memory/allocators/slab/mod.rs @@ -5,5 +5,22 @@ pub(super) const SLAB_FRAME_ALLOC_ORDER: usize = 2; pub(super) const SLAB_SIZE_BYTES: usize = PAGE_SIZE << SLAB_FRAME_ALLOC_ORDER; const SLAB_MAX_OBJ_SHIFT: u32 = SLAB_SIZE_BYTES.ilog2() - 1; +pub mod allocator; #[allow(clippy::module_inception)] pub(super) mod slab; + +/// Returns the index into the slab/cache list for a given layout. +fn alloc_order(layout: core::alloc::Layout) -> Option { + // We must take alignemnt into account too. + let size = core::cmp::max(layout.size(), layout.align()); + + let alloc_order = size.next_power_of_two().ilog2() as usize; + + if alloc_order > SLAB_MAX_OBJ_SHIFT as usize { + return None; + } + + // Since slabs use a `u16` as the 'next_free' pointer, our minimum order + // must be 1. + Some(if alloc_order == 0 { 1 } else { alloc_order }) +} diff --git a/src/arch/arm64/boot/mod.rs b/src/arch/arm64/boot/mod.rs index 6723605..f482bf9 100644 --- a/src/arch/arm64/boot/mod.rs +++ b/src/arch/arm64/boot/mod.rs @@ -105,7 +105,7 @@ fn arch_init_stage2(frame: *mut ExceptionState) -> *mut ExceptionState { .take() .expect("Smalloc should not have been taken yet"); - let page_alloc = unsafe { FrameAllocator::init(smalloc) }; + let (page_alloc, _) = unsafe { FrameAllocator::init(smalloc) }; if PAGE_ALLOC.set(page_alloc).is_err() { panic!("Cannot setup physical memory allocator");