From 9f2161805899644dc8cf9acf468b4d07a2e63a54 Mon Sep 17 00:00:00 2001 From: Matthew Leach Date: Sun, 8 Feb 2026 21:46:27 +0000 Subject: [PATCH] libkernel: memory: slab: heap: new Add a new module which implements the high-level heap logic (implements `GlobalAlloc`) for the slab allocator. Also implement a multi-threaded stress test of the heap ensuring concurrent allocation is unique (doesn't corrupt other allocs) and all memory is free'd back to the `FrameAllocator` when all caches are purged. --- .../src/memory/allocators/slab/allocator.rs | 9 +- libkernel/src/memory/allocators/slab/cache.rs | 19 +- libkernel/src/memory/allocators/slab/heap.rs | 397 ++++++++++++++++++ libkernel/src/memory/allocators/slab/mod.rs | 1 + src/arch/arm64/boot/mod.rs | 6 +- src/arch/arm64/memory/heap.rs | 171 ++------ 6 files changed, 465 insertions(+), 138 deletions(-) create mode 100644 libkernel/src/memory/allocators/slab/heap.rs diff --git a/libkernel/src/memory/allocators/slab/allocator.rs b/libkernel/src/memory/allocators/slab/allocator.rs index c8b504b..9fe3413 100644 --- a/libkernel/src/memory/allocators/slab/allocator.rs +++ b/libkernel/src/memory/allocators/slab/allocator.rs @@ -57,9 +57,9 @@ const MAX_FREE_SLABS: usize = 32; /// "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, + pub(super) free: LinkedList, + pub(super) partial: LinkedList, + pub(super) free_list_sz: usize, obj_shift: usize, frame_list: FrameList, phantom1: PhantomData, @@ -248,7 +248,8 @@ impl, T: AddressTranslator<()>> SlabManager } pub struct SlabAllocator, T: AddressTranslator<()>> { - managers: [SpinLockIrq, CPU>; SLAB_MAX_OBJ_SHIFT as usize + 1], + pub(super) managers: + [SpinLockIrq, CPU>; SLAB_MAX_OBJ_SHIFT as usize + 1], } unsafe impl, T: AddressTranslator<()>> Send diff --git a/libkernel/src/memory/allocators/slab/cache.rs b/libkernel/src/memory/allocators/slab/cache.rs index 8b64fd3..655545f 100644 --- a/libkernel/src/memory/allocators/slab/cache.rs +++ b/libkernel/src/memory/allocators/slab/cache.rs @@ -1,4 +1,7 @@ -use super::{alloc_order, allocator::SlabManager}; +use super::{ + alloc_order, + allocator::{SlabAllocator, SlabManager}, +}; use crate::{ CpuOps, memory::{ @@ -135,4 +138,18 @@ impl SlabCache { pub fn get_cache(&mut self, layout: core::alloc::Layout) -> Option<&mut PtrCache> { Some(&mut self.caches[alloc_order(layout)?]) } + + /// Flush all cache lines back into the slab allocator. + pub fn purge_into, T: AddressTranslator<()>>( + &mut self, + slab_alloc: &SlabAllocator, + ) { + for (line, slab) in self.caches.iter_mut().zip(slab_alloc.managers.iter()) { + let mut slab = slab.lock_save_irq(); + for i in 0..line.next_free { + slab.free(line.ptrs[i]); + } + line.next_free = 0; + } + } } diff --git a/libkernel/src/memory/allocators/slab/heap.rs b/libkernel/src/memory/allocators/slab/heap.rs new file mode 100644 index 0000000..b345456 --- /dev/null +++ b/libkernel/src/memory/allocators/slab/heap.rs @@ -0,0 +1,397 @@ +use super::{allocator::SlabAllocator, cache::SlabCache}; +use crate::{ + CpuOps, + memory::{ + PAGE_SIZE, + address::{AddressTranslator, VA}, + allocators::phys::PageAllocGetter, + page::ClaimedPage, + region::PhysMemoryRegion, + }, +}; +use core::{alloc::GlobalAlloc, marker::PhantomData, ops::DerefMut}; + +pub trait SlabGetter, T: AddressTranslator<()>> { + fn global_slab_alloc() -> &'static SlabAllocator; +} + +pub trait SlabCacheStorage { + fn store(ptr: *mut SlabCache); + fn get() -> impl DerefMut; +} + +pub struct KHeap +where + CPU: CpuOps, + S: SlabCacheStorage, + PG: PageAllocGetter, + T: AddressTranslator<()>, + SG: SlabGetter, +{ + phantom1: PhantomData, + phantom2: PhantomData, + phantom3: PhantomData, + phantom4: PhantomData, + phantom5: PhantomData, +} + +impl Default for KHeap +where + CPU: CpuOps, + S: SlabCacheStorage, + PG: PageAllocGetter, + T: AddressTranslator<()>, + SG: SlabGetter, +{ + fn default() -> Self { + Self::new() + } +} + +impl KHeap +where + CPU: CpuOps, + S: SlabCacheStorage, + PG: PageAllocGetter, + T: AddressTranslator<()>, + SG: SlabGetter, +{ + pub const fn new() -> Self { + Self { + phantom1: PhantomData, + phantom2: PhantomData, + phantom3: PhantomData, + phantom4: PhantomData, + phantom5: PhantomData, + } + } + + /// Calculates the Frame Allocator order required for a large allocation. + fn calculate_huge_order(layout: core::alloc::Layout) -> usize { + // Ensure we cover the size, rounding UP to the nearest page. + let size = core::cmp::max(layout.size(), layout.align()); + let pages_needed = size.div_ceil(PAGE_SIZE); + pages_needed.next_power_of_two().ilog2() as usize + } + + pub fn init_for_this_cpu() { + let page: ClaimedPage = + ClaimedPage::alloc_zeroed().expect("Cannot allocate heap page"); + + // SAFETY: We just successfully allocated the above page and the + // lifetime of the returned pointer will be for the entire lifetime of + // the kernel ('sttaic). + let slab_cache = unsafe { SlabCache::from_page(page) }; + + // Store the slab_cache pointer in the storage. + S::store(slab_cache); + } +} + +unsafe impl GlobalAlloc for KHeap +where + CPU: CpuOps, + S: SlabCacheStorage, + PG: PageAllocGetter, + T: AddressTranslator<()>, + SG: SlabGetter, +{ + unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 { + let mut cache = S::get(); + + let Some(cache_line) = cache.get_cache(layout) else { + // Allocation is too big for SLAB. Defer to using the frame + // allocator directly. + return PG::global_page_alloc() + .alloc_frames(Self::calculate_huge_order(layout) as _) + .unwrap() + .leak() + .start_address() + .to_va::() + .cast::() + .as_ptr_mut(); + }; + + if let Some(ptr) = cache_line.alloc() { + // Fast path, cache-hit. + return ptr; + } + + // Fall back to the slab allocator. + let mut slab = SG::global_slab_alloc() + .allocator_for_layout(layout) + .unwrap() + .lock_save_irq(); + + let ptr = slab.alloc(); + + // Fill up our cache with objects from the (maybe freshly allocated) + // slab. + cache_line.fill_from(&mut slab); + + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) { + let mut cache = S::get(); + + let Some(cache_line) = cache.get_cache(layout) else { + // If the allocation didn't fit in the slab, we must have used the + // FA directly. + let allocated_region = PhysMemoryRegion::new( + VA::from_ptr_mut(ptr as _).to_pa::(), + PAGE_SIZE << Self::calculate_huge_order(layout), + ); + + unsafe { + PG::global_page_alloc().alloc_from_region(allocated_region); + } + + return; + }; + + if cache_line.free(ptr).is_ok() { + return; + } + + // The cache is full. Return some memory back to the slab allocator. + let mut slab = SG::global_slab_alloc() + .allocator_for_layout(layout) + .unwrap() + .lock_save_irq(); + + slab.free(ptr); + + cache_line.drain_into(&mut slab); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + memory::{ + address::{IdentityTranslator, PA}, + allocators::{ + phys::{FrameAllocator, PageAllocGetter, tests::TestFixture}, + slab::{allocator::SlabAllocator, cache::SlabCache}, + }, + }, + test::MockCpuOps, + }; + use rand::{Rng, rng}; + use std::{ + cell::RefCell, + ops::{Deref, DerefMut}, + sync::{Arc, Barrier, OnceLock}, + thread, + }; + + static FIXTURE: OnceLock = OnceLock::new(); + static SLAB_ALLOCATOR: OnceLock< + SlabAllocator, + > = OnceLock::new(); + + fn get_fixture() -> &'static TestFixture { + FIXTURE.get_or_init(|| TestFixture::new(&[(0, 512 * 1024 * 1024)], &[])) + } + + struct TestAllocGetter; + impl PageAllocGetter for TestAllocGetter { + fn global_page_alloc() -> &'static FrameAllocator { + &get_fixture().allocator + } + } + + struct TestSlabGetter; + impl SlabGetter for TestSlabGetter { + fn global_slab_alloc() + -> &'static SlabAllocator { + SLAB_ALLOCATOR.get_or_init(|| { + let fixture = get_fixture(); + SlabAllocator::new(fixture.frame_list.clone()) + }) + } + } + + thread_local! { + static TLS_CACHE: RefCell> = RefCell::new(None); + } + + struct ThreadLocalCacheStorage; + + struct ThreadCacheGuard { + ptr: *mut SlabCache, + } + + impl Deref for ThreadCacheGuard { + type Target = SlabCache; + fn deref(&self) -> &Self::Target { + unsafe { &*self.ptr } + } + } + + impl DerefMut for ThreadCacheGuard { + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut *self.ptr } + } + } + + impl SlabCacheStorage for ThreadLocalCacheStorage { + fn store(ptr: *mut SlabCache) { + TLS_CACHE.with(|c| { + *c.borrow_mut() = Some(ptr); + }); + } + + fn get() -> impl Deref + DerefMut { + let ptr = TLS_CACHE.with(|c| { + c.borrow() + .expect("Thread cache not initialized for this thread") + }); + ThreadCacheGuard { ptr } + } + } + + type TestHeap = KHeap< + MockCpuOps, + ThreadLocalCacheStorage, + TestAllocGetter, + IdentityTranslator, + TestSlabGetter, + >; + + #[test] + fn heap_stress_test() { + let _ = get_fixture(); + let _ = TestSlabGetter::global_slab_alloc(); + + let num_threads = 8; + let ops_per_thread = 100_000; + let barrier = Arc::new(Barrier::new(num_threads)); + + // Track allocated memory usage to verify leak detection later + let initial_free_pages = get_fixture().allocator.free_pages(); + println!("Initial Free Pages: {}", initial_free_pages); + + let mut handles = vec![]; + + for t_idx in 0..num_threads { + let barrier = barrier.clone(); + + handles.push(thread::spawn(move || { + TestHeap::init_for_this_cpu(); + + barrier.wait(); + + let heap = TestHeap::new(); + let mut rng = rng(); + + // Track allocations: (Ptr, Layout, PatternByte) + let mut allocations: Vec<(*mut u8, core::alloc::Layout, u8)> = Vec::new(); + + for _ in 0..ops_per_thread { + // Randomly decide to Alloc (70%) or Free (30%) + // Bias towards Alloc to build up memory pressure + if rng.random_bool(0.6) || allocations.is_empty() { + // Allocation path + + // Random size: biased to small (slab), occasional huge + let size = if rng.random_bool(0.95) { + rng.random_range(1..=2048) // Small (Slabs) + } else { + rng.random_range(4096..=16384) // Large (Pages) + }; + + // Random alignment (power of 2) + let align = 1 << rng.random_range(0..=6); + let layout = core::alloc::Layout::from_size_align(size, align).unwrap(); + + unsafe { + let ptr = heap.alloc(layout); + assert!(!ptr.is_null(), "Allocation failed"); + assert_eq!(ptr as usize % align, 0, "Alignment violation"); + + // Write Pattern + let pattern: u8 = rng.random(); + std::ptr::write_bytes(ptr, pattern, size); + + allocations.push((ptr, layout, pattern)); + } + } else { + // Free Path. + + // Remove a random allocation from our list + let idx = rng.random_range(0..allocations.len()); + let (ptr, layout, pattern) = allocations.swap_remove(idx); + + unsafe { + // Verify Pattern + let slice = std::slice::from_raw_parts(ptr, layout.size()); + for (i, &byte) in slice.iter().enumerate() { + assert_eq!( + byte, pattern, + "Memory Corruption detected in thread {} at byte {}", + t_idx, i + ); + } + + heap.dealloc(ptr, layout); + } + } + } + + // Free everything. + for (ptr, layout, pattern) in allocations { + unsafe { + let slice = std::slice::from_raw_parts(ptr, layout.size()); + for &byte in slice.iter() { + assert_eq!(byte, pattern, "Corruption detected during cleanup"); + } + heap.dealloc(ptr, layout); + } + } + + // Purge the per-cpu caches. + let slab = SLAB_ALLOCATOR.get().unwrap(); + ThreadLocalCacheStorage::get().purge_into(&slab); + + let addr = ThreadLocalCacheStorage::get().deref() as *const SlabCache; + + // Return the slab cache page. + unsafe { + FIXTURE + .get() + .unwrap() + .allocator + .alloc_from_region(PhysMemoryRegion::new( + PA::from_value(addr as usize), + PAGE_SIZE, + )); + } + })) + } + + // Wait for all threads + for h in handles { + h.join().unwrap(); + } + + // Purge the all slab free lsts (the partial list should be empty). + for slab_man in SLAB_ALLOCATOR.get().unwrap().managers.iter() { + let mut frame_alloc = FIXTURE.get().unwrap().allocator.inner.lock_save_irq(); + + let mut slab = slab_man.lock_save_irq(); + + assert!(slab.partial.is_empty()); + + while let Some(slab) = slab.free.pop_front() { + frame_alloc.free_slab(slab); + } + } + + let final_free = get_fixture().allocator.free_pages(); + + assert_eq!(initial_free_pages, final_free); + } +} diff --git a/libkernel/src/memory/allocators/slab/mod.rs b/libkernel/src/memory/allocators/slab/mod.rs index 5cd1159..ea76b71 100644 --- a/libkernel/src/memory/allocators/slab/mod.rs +++ b/libkernel/src/memory/allocators/slab/mod.rs @@ -7,6 +7,7 @@ const SLAB_MAX_OBJ_SHIFT: u32 = SLAB_SIZE_BYTES.ilog2() - 1; pub mod allocator; pub mod cache; +pub mod heap; #[allow(clippy::module_inception)] pub(super) mod slab; diff --git a/src/arch/arm64/boot/mod.rs b/src/arch/arm64/boot/mod.rs index 17873de..15810c8 100644 --- a/src/arch/arm64/boot/mod.rs +++ b/src/arch/arm64/boot/mod.rs @@ -2,7 +2,7 @@ use super::{ exceptions::{ExceptionState, secondary_exceptions_init}, memory::{ fixmap::FIXMAPS, - heap::{KHeap, SLAB_ALLOC}, + heap::{KernelHeap, SLAB_ALLOC}, mmu::setup_kern_addr_space, }, proc::vdso::vdso_init, @@ -118,7 +118,7 @@ fn arch_init_stage2(frame: *mut ExceptionState) -> *mut ExceptionState { panic!("Cannot setup slab allocator"); } - KHeap::init_for_this_cpu(); + KernelHeap::init_for_this_cpu(); // Don't trap wfi/wfe in el0. SCTLR_EL1.modify(SCTLR_EL1::NTWE::DontTrap + SCTLR_EL1::NTWI::DontTrap); @@ -156,7 +156,7 @@ fn arch_init_secondary(ctx_frame: *mut ExceptionState) -> *mut ExceptionState { SCTLR_EL1.modify(SCTLR_EL1::NTWE::DontTrap + SCTLR_EL1::NTWI::DontTrap); // Setup heap per-cpu data. - KHeap::init_for_this_cpu(); + KernelHeap::init_for_this_cpu(); // Enable interrupts and exceptions. secondary_exceptions_init(); diff --git a/src/arch/arm64/memory/heap.rs b/src/arch/arm64/memory/heap.rs index cbe397f..d0b179a 100644 --- a/src/arch/arm64/memory/heap.rs +++ b/src/arch/arm64/memory/heap.rs @@ -1,38 +1,40 @@ +use crate::{ + arch::ArchImpl, + memory::{PageOffsetTranslator, page::PgAllocGetter}, + sync::OnceLock, +}; use core::{ - alloc::GlobalAlloc, arch::asm, ops::{Deref, DerefMut}, ptr, }; - -use crate::{ - arch::ArchImpl, - memory::{ - PAGE_ALLOC, PageOffsetTranslator, - page::{ClaimedPage, PgAllocGetter}, - }, - sync::OnceLock, -}; use libkernel::{ CpuOps, - memory::{ - PAGE_SIZE, - address::VA, - allocators::slab::{allocator::SlabAllocator, cache::SlabCache}, - region::PhysMemoryRegion, + memory::allocators::slab::{ + allocator::SlabAllocator, + cache::SlabCache, + heap::{KHeap, SlabCacheStorage, SlabGetter}, }, }; -pub static SLAB_ALLOC: OnceLock> = - OnceLock::new(); +type SlabAlloc = SlabAllocator; -struct PerCpuCache { - ptr: *mut SlabCache, +pub static SLAB_ALLOC: OnceLock = OnceLock::new(); + +pub struct StaticSlabGetter {} + +impl SlabGetter for StaticSlabGetter { + fn global_slab_alloc() -> &'static SlabAlloc { + SLAB_ALLOC.get().unwrap() + } +} + +pub struct PerCpuCache { flags: usize, } impl PerCpuCache { - fn get() -> Self { + fn get_ptr() -> *mut SlabCache { let mut cache: *mut SlabCache = ptr::null_mut(); unsafe { asm!("mrs {}, TPIDR_EL1", out(reg) cache, options(nostack, nomem)) }; @@ -41,9 +43,22 @@ impl PerCpuCache { panic!("Attempted to use alloc/free before CPU initalisation!"); } + cache + } +} + +impl SlabCacheStorage for PerCpuCache { + fn store(ptr: *mut SlabCache) { + #[allow(clippy::pointers_in_nomem_asm_block)] + unsafe { + asm!("msr TPIDR_EL1, {}", in(reg) ptr, options(nostack, nomem)); + }; + } + + fn get() -> impl DerefMut { let flags = ArchImpl::disable_interrupts(); - Self { ptr: cache, flags } + Self { flags } } } @@ -54,7 +69,7 @@ impl Deref for PerCpuCache { // SAFETY: The pointer uses a CPU-banked register for access. We've // disabled interrupts so we know we cannot be preempted, therefore // mutable access to the cache is safe. - unsafe { &(*self.ptr) } + unsafe { &(*Self::get_ptr()) } } } @@ -63,7 +78,7 @@ impl DerefMut for PerCpuCache { // SAFETY: The pointer uses a CPU-banked register for access. We've // disabled interrupts so we know we cannot be preempted, therefore // mutable access to the cache is safe. - unsafe { &mut (*self.ptr) } + unsafe { &mut (*Self::get_ptr()) } } } @@ -73,112 +88,8 @@ impl Drop for PerCpuCache { } } -pub struct KHeap {} - -impl KHeap { - /// Calculates the Frame Allocator order required for a large allocation. - fn calculate_huge_order(layout: core::alloc::Layout) -> usize { - // Ensure we cover the size, rounding UP to the nearest page. - let size = core::cmp::max(layout.size(), layout.align()); - let pages_needed = size.div_ceil(PAGE_SIZE); - pages_needed.next_power_of_two().ilog2() as usize - } - - pub fn init_for_this_cpu() { - let page = ClaimedPage::alloc_zeroed().expect("Cannot allocate heap page"); - - // SAFETY: We just successfully allocated the above page and the - // lifetime of the returned pointer will be for the entire lifetime of - // the kernel ('sttaic). - let slab_cache = unsafe { SlabCache::from_page(page) }; - - // Store the slab_cache pointer in the CPU-banked register `TPIDR_EL1`. - #[allow(clippy::pointers_in_nomem_asm_block)] - unsafe { - asm!("msr TPIDR_EL1, {}", in(reg) slab_cache, options(nostack, nomem)); - } - } -} - -unsafe impl GlobalAlloc for KHeap { - unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 { - let mut cache = PerCpuCache::get(); - - let Some(cache_line) = cache.get_cache(layout) else { - // Allocation is too big for SLAB. Defer to using the frame - // allocator directly. - return PAGE_ALLOC - .get() - .unwrap() - .alloc_frames(Self::calculate_huge_order(layout) as _) - .unwrap() - .leak() - .start_address() - .to_va::() - .cast::() - .as_ptr_mut(); - }; - - if let Some(ptr) = cache_line.alloc() { - // Fast path, cache-hit. - return ptr; - } - - // Fall back to the slab allocator. - let mut slab = SLAB_ALLOC - .get() - .expect("Slab alocator not initalised") - .allocator_for_layout(layout) - .unwrap() - .lock_save_irq(); - - let ptr = slab.alloc(); - - // Fill up our cache with objects from the (maybe freshly allocated) - // slab. - cache_line.fill_from(&mut slab); - - ptr - } - - unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) { - let mut cache = PerCpuCache::get(); - - let Some(cache_line) = cache.get_cache(layout) else { - // If the allocation didn't fit in the slab, we must have used the - // FA directly. - let allocated_region = PhysMemoryRegion::new( - VA::from_ptr_mut(ptr as _).to_pa::(), - PAGE_SIZE << Self::calculate_huge_order(layout), - ); - - unsafe { - PAGE_ALLOC - .get() - .unwrap() - .alloc_from_region(allocated_region); - } - - return; - }; - - if cache_line.free(ptr).is_ok() { - return; - } - - // The cache is full. Return some memory back to the slab allocator. - let mut slab = SLAB_ALLOC - .get() - .expect("Slab alocator not initalised") - .allocator_for_layout(layout) - .unwrap() - .lock_save_irq(); - - slab.free(ptr); - - cache_line.drain_into(&mut slab); - } -} +pub type KernelHeap = + KHeap; #[global_allocator] -static K_HEAP: KHeap = KHeap {}; +static K_HEAP: KernelHeap = KernelHeap::new();