refactor: refactor slab allocator

Cargo.lock

@@ -158,6 +158,7 @@ dependencies = [
  "itertools",
  "pointers",
  "posix_types",
+ "slab_allocator",
 ]
 
 [[package]]
@@ -327,6 +328,15 @@ version = "1.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
 
+[[package]]
+name = "slab_allocator"
+version = "0.1.0"
+dependencies = [
+ "eonix_mm",
+ "eonix_sync",
+ "intrusive_list",
+]
+
 [[package]]
 name = "syn"
 version = "2.0.89"

Cargo.toml

@@ -12,6 +12,7 @@ atomic_unique_refcell = { path = "./crates/atomic_unique_refcell", features = [
     "no_std",
 ] }
 buddy_allocator = { path = "./crates/buddy_allocator" }
+slab_allocator = { path = "./crates/slab_allocator" }
 eonix_mm = { path = "./crates/eonix_mm" }
 eonix_preempt = { path = "./crates/eonix_preempt" }
 eonix_runtime = { path = "./crates/eonix_runtime" }

crates/intrusive_list/src/lib.rs

@@ -41,6 +41,14 @@ impl List {
             node
         })
     }
+
+    pub fn is_empty(&self) -> bool {
+        self.count == 0
+    }
+
+    pub fn head(&mut self) -> Option<&mut Link> {
+        self.head.next_mut()
+    }
 }
 
 impl Link {

crates/slab_allocator/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "slab_allocator"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+eonix_mm = { path = "../eonix_mm" }
+eonix_sync = { path = "../eonix_sync" }
+intrusive_list = { path = "../intrusive_list" }
+

crates/slab_allocator/src/lib.rs

@@ -0,0 +1,71 @@
+#![no_std]
+
+mod slab_cache;
+
+use core::cmp::max;
+
+use eonix_mm::paging::{PageAlloc, RawPage};
+use eonix_sync::Spin;
+use intrusive_list::Link;
+use slab_cache::SlabCache;
+
+pub trait SlabRawPage: RawPage {
+    /// Get the container raw page struct of the list link.
+    ///
+    /// # Safety
+    /// The caller MUST ensure that the link points to a `RawPage`.
+    unsafe fn from_link(link: &mut Link) -> Self;
+
+    /// Get the list link of the raw page.
+    ///
+    /// # Safety
+    /// The caller MUST ensure that at any time, only one mutable reference
+    /// to the link exists.
+    unsafe fn get_link(&self) -> &mut Link;
+
+    fn slab_init(&self, object_size: u32);
+
+    fn is_emtpy(&self) -> bool;
+
+    fn is_full(&self) -> bool;
+
+    // which slab page the ptr belong
+    fn in_which(ptr: *mut u8) -> Self;
+
+    fn alloc_slot(&self) -> *mut u8;
+
+    fn dealloc_slot(&self, ptr: *mut u8);
+}
+
+pub struct SlabAllocator<T, A, const SLAB_CACHE_COUNT: usize> {
+    slabs: [Spin<SlabCache<T, A>>; SLAB_CACHE_COUNT],
+}
+
+unsafe impl<T, A, const SLAB_CACHE_COUNT: usize> Send for SlabAllocator<T, A, SLAB_CACHE_COUNT> {}
+unsafe impl<T, A, const SLAB_CACHE_COUNT: usize> Sync for SlabAllocator<T, A, SLAB_CACHE_COUNT> {}
+
+impl<Raw, Allocator, const SLAB_CACHE_COUNT: usize> SlabAllocator<Raw, Allocator, SLAB_CACHE_COUNT>
+where
+    Raw: SlabRawPage,
+    Allocator: PageAlloc<RawPage = Raw> + Clone,
+{
+    pub fn new_in(alloc: Allocator) -> Self {
+        Self {
+            slabs: core::array::from_fn(|i| {
+                Spin::new(SlabCache::new_in(1 << (i + 3), alloc.clone()))
+            }),
+        }
+    }
+
+    pub fn alloc(&self, mut size: usize) -> *mut u8 {
+        size = max(8, size);
+        let idx = size.next_power_of_two().trailing_zeros() - 3;
+        self.slabs[idx as usize].lock().alloc()
+    }
+
+    pub fn dealloc(&self, ptr: *mut u8, mut size: usize) {
+        size = max(8, size);
+        let idx = size.next_power_of_two().trailing_zeros() - 3;
+        self.slabs[idx as usize].lock().dealloc(ptr);
+    }
+}

crates/slab_allocator/src/slab_cache.rs

@@ -0,0 +1,89 @@
+use super::SlabRawPage;
+use core::marker::PhantomData;
+use eonix_mm::paging::{PageAlloc, PAGE_SIZE};
+use intrusive_list::List;
+
+pub(crate) struct SlabCache<T, A> {
+    empty_list: List,
+    partial_list: List,
+    full_list: List,
+    object_size: u32,
+    alloc: A,
+    _phantom: PhantomData<T>,
+}
+
+impl<Raw, Allocator> SlabCache<Raw, Allocator>
+where
+    Raw: SlabRawPage,
+    Allocator: PageAlloc<RawPage = Raw>,
+{
+    pub(crate) const fn new_in(object_size: u32, alloc: Allocator) -> Self {
+        // avoid uncessary branch in alloc and dealloc
+        assert!(object_size <= PAGE_SIZE as u32 / 2);
+
+        Self {
+            empty_list: List::new(),
+            partial_list: List::new(),
+            full_list: List::new(),
+            alloc,
+            object_size: object_size,
+            _phantom: PhantomData,
+        }
+    }
+
+    pub(crate) fn alloc(&mut self) -> *mut u8 {
+        if !self.partial_list.is_empty() {
+            let page_ptr = unsafe {
+                Raw::from_link(
+                    self.partial_list
+                        .head()
+                        .expect("partial pages should not be empty"),
+                )
+            };
+
+            let ptr = page_ptr.alloc_slot();
+            if page_ptr.is_full() {
+                self.partial_list.remove(unsafe { page_ptr.get_link() });
+                self.full_list.insert(unsafe { page_ptr.get_link() });
+            }
+            return ptr;
+        }
+
+        if !self.empty_list.is_empty() {
+            let page_ptr = unsafe {
+                Raw::from_link(
+                    self.empty_list
+                        .head()
+                        .expect("empty pages should not be empty"),
+                )
+            };
+
+            let ptr = page_ptr.alloc_slot();
+            self.empty_list.remove(unsafe { page_ptr.get_link() });
+            self.partial_list.insert(unsafe { page_ptr.get_link() });
+            return ptr;
+        }
+
+        let new_page = self.alloc.alloc().expect("slab_cache get page fail!");
+        new_page.slab_init(self.object_size);
+        let ptr = new_page.alloc_slot();
+        self.partial_list.insert(unsafe { new_page.get_link() });
+        ptr
+    }
+
+    pub(crate) fn dealloc(&mut self, ptr: *mut u8) {
+        let page_ptr = Raw::in_which(ptr);
+
+        if page_ptr.is_full() {
+            self.full_list.remove(unsafe { page_ptr.get_link() });
+            self.partial_list.insert(unsafe { page_ptr.get_link() });
+        }
+
+        page_ptr.dealloc_slot(ptr);
+
+        if page_ptr.is_emtpy() {
+            self.partial_list.remove(unsafe { page_ptr.get_link() });
+            self.empty_list.insert(unsafe { page_ptr.get_link() });
+        }
+    }
+}

src/kernel/mem.rs

@@ -2,6 +2,7 @@ pub mod paging;
 
 mod access;
 mod address;
+mod allocator;
 mod mm_area;
 mod mm_list;
 mod page_alloc;

src/kernel/mem/access.rs

@@ -1,7 +1,9 @@
 use core::{num::NonZero, ptr::NonNull};
-use eonix_mm::address::{Addr as _, PAddr};
+use eonix_mm::address::{Addr as _, PAddr, VAddr};
 use eonix_mm::paging::{PageAccess, PageBlock, PFN};
 
+use super::page_alloc::RawPagePtr;
+
 const PHYS_OFFSET: usize = 0xffff_ff00_0000_0000;
 
 /// A block of memory starting at a non-zero address and having a specific length.
@@ -156,3 +158,22 @@ impl PageAccess for KernelPageAccess {
         }
     }
 }
+
+pub trait RawPageAccess {
+    /// Translate the address belonged RawPage ptr
+    /// Use it with care.
+    ///
+    /// # Panic
+    /// If the address is not properly aligned.
+    ///
+    /// # Safety
+    /// the address must be kernel accessible pointer
+    unsafe fn as_raw_page(&self) -> RawPagePtr;
+}
+
+impl RawPageAccess for VAddr {
+    unsafe fn as_raw_page(&self) -> RawPagePtr {
+        let pfn: PFN = PAddr::from(self.addr() - PHYS_OFFSET).into();
+        RawPagePtr::from(pfn)
+    }
+}

src/kernel/mem/allocator.rs

@@ -0,0 +1,53 @@
+use core::alloc::{GlobalAlloc, Layout};
+use eonix_mm::paging::{PageAlloc, PAGE_SIZE_BITS};
+use eonix_sync::LazyLock;
+use slab_allocator::{SlabAllocator, SlabRawPage};
+
+use super::page_alloc::RawPagePtr;
+use super::GlobalPageAlloc;
+
+static SLAB_ALLOCATOR: LazyLock<SlabAllocator<RawPagePtr, GlobalPageAlloc, 9>> =
+    LazyLock::new(|| SlabAllocator::new_in(GlobalPageAlloc));
+
+struct Allocator;
+
+unsafe impl GlobalAlloc for Allocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let size = layout.size().next_power_of_two();
+
+        let result = if size <= 2048 {
+            SLAB_ALLOCATOR.alloc(size)
+        } else {
+            let page_num = size >> PAGE_SIZE_BITS;
+            let order = page_num.next_power_of_two().trailing_zeros();
+            let raw_page = GlobalPageAlloc
+                .alloc_order(order)
+                .expect("allocate page failed!");
+            raw_page.real_ptr().as_ptr()
+        };
+
+        if result.is_null() {
+            core::ptr::null_mut()
+        } else {
+            result as *mut u8
+        }
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        let size = layout.size().next_power_of_two();
+
+        if size <= 2048 {
+            SLAB_ALLOCATOR.dealloc(ptr, size)
+        } else {
+            let page_ptr: RawPagePtr = SlabRawPage::in_which(ptr);
+            page_ptr
+                .as_mut()
+                .refcount
+                .fetch_sub(1, core::sync::atomic::Ordering::Relaxed);
+            GlobalPageAlloc.dealloc(page_ptr);
+        };
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: Allocator = Allocator;

src/kernel/mem/page_alloc.rs

@@ -9,7 +9,9 @@ use eonix_mm::{
 };
 use eonix_sync::Spin;
 use intrusive_list::List;
-use raw_page::{PageFlags, RawPage, RawPagePtr};
+use raw_page::{PageFlags, RawPage};
+
+pub use raw_page::RawPagePtr;
 
 const COSTLY_ORDER: u32 = 3;
 const BATCH_SIZE: u32 = 64;

src/kernel/mem/page_alloc/raw_page.rs

@@ -3,14 +3,99 @@ use core::{
     ptr::NonNull,
     sync::atomic::{AtomicU32, AtomicUsize, Ordering},
 };
-use eonix_mm::paging::{RawPage as RawPageTrait, PFN};
+use eonix_mm::{
+    address::{PAddr, VAddr},
+    paging::{RawPage as RawPageTrait, PAGE_SIZE, PFN},
+};
 use intrusive_list::{container_of, Link};
+use slab_allocator::SlabRawPage;
+
+use crate::kernel::mem::access::RawPageAccess;
+use crate::kernel::mem::PhysAccess;
 
 const PAGE_ARRAY: NonNull<RawPage> =
     unsafe { NonNull::new_unchecked(0xffffff8040000000 as *mut _) };
 
 pub struct PageFlags(AtomicU32);
 
+pub struct SlabPageInner {
+    pub object_size: u32,
+    pub allocated_count: u32,
+    pub free_next: Option<NonNull<usize>>,
+}
+
+pub struct BuddyPageInner {}
+
+pub enum PageType {
+    Buddy(BuddyPageInner),
+    Slab(SlabPageInner),
+}
+
+impl PageType {
+    // slab
+    pub fn new_slab(&mut self, object_size: u32) {
+        *self = PageType::Slab(SlabPageInner {
+            object_size,
+            allocated_count: 0,
+            free_next: None,
+        })
+    }
+
+    pub fn object_size(&self) -> u32 {
+        assert!(matches!(self, PageType::Slab(_)));
+
+        match self {
+            PageType::Slab(inner) => inner.object_size,
+            _ => unsafe { core::hint::unreachable_unchecked() },
+        }
+    }
+
+    pub fn allocated_count(&self) -> u32 {
+        assert!(matches!(self, PageType::Slab(_)));
+
+        match self {
+            PageType::Slab(inner) => inner.allocated_count,
+            _ => unsafe { core::hint::unreachable_unchecked() },
+        }
+    }
+
+    pub fn allocated_count_add(&mut self, val: u32) {
+        assert!(matches!(self, PageType::Slab(_)));
+
+        match self {
+            PageType::Slab(inner) => inner.allocated_count += val,
+            _ => unsafe { core::hint::unreachable_unchecked() },
+        }
+    }
+
+    pub fn allocated_count_sub(&mut self, val: u32) {
+        assert!(matches!(self, PageType::Slab(_)));
+
+        match self {
+            PageType::Slab(inner) => inner.allocated_count -= val,
+            _ => unsafe { core::hint::unreachable_unchecked() },
+        }
+    }
+
+    pub fn free_next(&self) -> Option<NonNull<usize>> {
+        assert!(matches!(self, PageType::Slab(_)));
+
+        match self {
+            PageType::Slab(inner) => inner.free_next,
+            _ => unsafe { core::hint::unreachable_unchecked() },
+        }
+    }
+
+    pub fn set_free_next(&mut self, free_next: Option<NonNull<usize>>) {
+        assert!(matches!(self, PageType::Slab(_)));
+
+        match self {
+            PageType::Slab(inner) => inner.free_next = free_next,
+            _ => unsafe { core::hint::unreachable_unchecked() },
+        }
+    }
+}
+
 pub struct RawPage {
     /// This can be used for LRU page swap in the future.
     ///
@@ -21,6 +106,8 @@ pub struct RawPage {
     pub order: u32,
     pub flags: PageFlags,
     pub refcount: AtomicUsize,
+
+    pub type_: PageType,
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -80,6 +167,12 @@ impl RawPagePtr {
     pub const fn refcount(&self) -> &AtomicUsize {
         &self.as_ref().refcount
     }
+
+    // return the ptr point to the actually raw page
+    pub fn real_ptr<T>(&self) -> NonNull<T> {
+        let pfn = unsafe { PFN::from(RawPagePtr(NonNull::new_unchecked(self.as_ptr()))) };
+        unsafe { PAddr::from(pfn).as_ptr::<T>() }
+    }
 }
 
 impl From<RawPagePtr> for PFN {
@@ -148,3 +241,80 @@ impl BuddyRawPage for RawPagePtr {
         self.flags().clear(PageFlags::FREE);
     }
 }
+
+impl SlabRawPage for RawPagePtr {
+    unsafe fn from_link(link: &mut Link) -> Self {
+        let raw_page_ptr = container_of!(link, RawPage, link);
+        Self(raw_page_ptr)
+    }
+
+    unsafe fn get_link(&self) -> &mut Link {
+        &mut self.as_mut().link
+    }
+
+    fn is_emtpy(&self) -> bool {
+        self.as_ref().type_.allocated_count() == 0
+    }
+
+    fn is_full(&self) -> bool {
+        self.as_ref().type_.free_next().is_none()
+    }
+
+    fn alloc_slot(&self) -> *mut u8 {
+        let ptr = self.as_ref().type_.free_next();
+
+        match ptr {
+            Some(ptr) => {
+                let next_free = unsafe { ptr.read() as *mut usize };
+                self.as_mut().type_.set_free_next(NonNull::new(next_free));
+                self.as_mut().type_.allocated_count_add(1);
+                return ptr.as_ptr() as *mut u8;
+            }
+            None => unreachable!(),
+        }
+    }
+
+    fn in_which(ptr: *mut u8) -> RawPagePtr {
+        let vaddr = VAddr::from(ptr as usize & !(PAGE_SIZE - 1));
+
+        unsafe { vaddr.as_raw_page() }
+    }
+
+    fn dealloc_slot(&self, ptr: *mut u8) {
+        let ptr = ptr as *mut usize;
+
+        if let Some(last_free) = self.as_ref().type_.free_next() {
+            unsafe { *ptr = last_free.as_ptr() as usize }
+        } else {
+            unsafe { *ptr = 0 }
+        }
+
+        self.as_mut().type_.allocated_count_sub(1);
+        self.as_mut().type_.set_free_next(NonNull::new(ptr));
+    }
+
+    fn slab_init(&self, object_size: u32) {
+        assert!(object_size >= core::mem::size_of::<usize>() as u32);
+
+        self.as_mut().type_.new_slab(object_size);
+
+        let mut slot_count = PAGE_SIZE / object_size as usize;
+        let mut ptr = self.real_ptr::<usize>().as_ptr();
+        self.as_mut().type_.set_free_next(NonNull::new(ptr));
+
+        // SAFETY: carefully ptr operate
+        unsafe {
+            loop {
+                if slot_count == 1 {
+                    *ptr = 0;
+                    break;
+                }
+
+                let next_ptr = ptr.byte_add(object_size as usize);
+                *ptr = next_ptr as usize;
+                ptr = next_ptr;
+                slot_count -= 1;
+            }
+        }
+    }
+}

src/lib.rs

@@ -25,7 +25,6 @@ mod rcu;
 mod sync;
 
 use alloc::{ffi::CString, sync::Arc};
-use core::alloc::{GlobalAlloc, Layout};
 use elf::ParsedElf32;
 use eonix_mm::{address::PAddr, paging::PFN};
 use eonix_runtime::{run::FutureRun, scheduler::Scheduler, task::Task};
@@ -61,37 +60,9 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
     arch::freeze()
 }
 
-extern "C" {
-    fn _do_allocate(size: usize) -> *mut core::ffi::c_void;
-    fn _do_deallocate(ptr: *mut core::ffi::c_void, size: core::ffi::c_size_t) -> i32;
-    fn init_pci();
-}
-
-struct Allocator;
-unsafe impl GlobalAlloc for Allocator {
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let result = _do_allocate(layout.size());
-
-        if result.is_null() {
-            core::ptr::null_mut()
-        } else {
-            result as *mut u8
-        }
-    }
-
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        match _do_deallocate(ptr as *mut core::ffi::c_void, layout.size()) {
-            0 => (),
-            _ => panic!("Failed to deallocate memory"),
-        }
-    }
-}
-
-#[global_allocator]
-static ALLOCATOR: Allocator = Allocator;
-
 extern "C" {
     fn init_allocator();
+    fn init_pci();
 }
 
 #[no_mangle]