alloc: small slab allocator

* Slab allocator for small objects (under 1/8 of page)

Signed-off-by: Andy-Python-Programmer <[email protected]>

Author: Andy-Python-Programmer

src/Cargo.lock

@@ -1,2 +1,5 @@
 [workspace]
 members = ["aero_kernel", "aero_syscall", "aero_proc"]
+
+[profile.release]
+debug = true

src/Cargo.toml

@@ -33,7 +33,9 @@ log = "0.4.14"
 xmas-elf = "0.8.0"
 hashbrown = "0.12.3"
 rustc-demangle = "0.1.20"
-intrusive-collections = "0.9.2"
+# intrusive-collections:
+#   `nightly`: Get access to const variants of the functions.
+intrusive-collections = { version = "0.9.2", features = ["nightly"] }
 serde_json = { version = "1.0", default-features = false, features = ["alloc"] }
 lai = { git = "https://github.com/aero-os/lai-rs" }
 uapi = { path = "../uapi" }

src/aero_kernel/Cargo.toml

@@ -45,7 +45,10 @@
     maybe_uninit_write_slice,
     slice_ptr_get,
     maybe_uninit_as_bytes,
-    pointer_is_aligned
+    pointer_is_aligned,
+    const_trait_impl,
+    int_roundings,
+    const_ptr_is_null
 )]
 #![deny(trivial_numeric_casts, unused_allocation)]
 #![test_runner(crate::tests::test_runner)]

src/aero_kernel/src/main.rs

@@ -20,33 +20,56 @@
 use core::alloc;
 use core::alloc::{GlobalAlloc, Layout};
 
-use super::paging::FRAME_ALLOCATOR;
+use super::slab::{SlabHeader, SmallSlab};
 use super::vmalloc;
 use crate::mem::paging::*;
+use crate::utils::sync::IrqGuard;
 
-struct Allocator {}
+struct Allocator {
+    zones: [SmallSlab; 9],
+}
 
 impl Allocator {
     const fn new() -> Self {
-        Self {}
+        Self {
+            zones: [
+                SmallSlab::new(8),
+                SmallSlab::new(16),
+                SmallSlab::new(32),
+                SmallSlab::new(64),
+                SmallSlab::new(128),
+                SmallSlab::new(256),
+                SmallSlab::new(512),
+                SmallSlab::new(1024),
+                SmallSlab::new(2048),
+            ],
+        }
     }
 
     fn alloc(&self, layout: Layout) -> *mut u8 {
+        let _guard = IrqGuard::new();
+
         let size = align_up(layout.size() as _, layout.align() as _);
 
-        if size <= Size4KiB::SIZE {
-            let frame: PhysFrame<Size4KiB> = FRAME_ALLOCATOR.allocate_frame().unwrap();
-            frame.start_address().as_hhdm_virt().as_mut_ptr()
-        } else if size <= Size2MiB::SIZE {
-            let frame: PhysFrame<Size2MiB> = FRAME_ALLOCATOR.allocate_frame().unwrap();
-            frame.start_address().as_hhdm_virt().as_mut_ptr()
+        for slab in self.zones.iter() {
+            if size as usize <= slab.size() {
+                return slab.alloc();
+            }
+        }
+
+        if size <= Size2MiB::SIZE {
+            FRAME_ALLOCATOR
+                .alloc(size as usize)
+                .unwrap()
+                .as_hhdm_virt()
+                .as_mut_ptr()
         } else {
             let size = align_up(size, Size4KiB::SIZE) / Size4KiB::SIZE;
 
             vmalloc::get_vmalloc()
                 .alloc(size as usize)
                 .map(|addr| addr.as_mut_ptr::<u8>())
-                .unwrap_or(core::ptr::null_mut())
+                .unwrap()
         }
     }
 
@@ -58,7 +81,13 @@ impl Allocator {
             return;
         }
 
-        // TODO: free the slab.
+        let _guard = IrqGuard::new();
+        if layout.size() <= 1024 {
+            let slab_header = (ptr as usize & !(0xfff)) as *mut SlabHeader;
+
+            let slab_header = unsafe { &mut *slab_header };
+            slab_header.ptr.dealloc(ptr);
+        }
     }
 }
 

src/aero_kernel/src/mem/alloc.rs

@@ -20,6 +20,7 @@
 pub mod alloc;
 pub mod paging;
 pub mod pti;
+mod slab;
 mod vmalloc;
 
 use core::alloc::Layout;

src/aero_kernel/src/mem/mod.rs

@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2021-2023 The Aero Project Developers.
+ *
+ * This file is part of The Aero Project.
+ *
+ * Aero is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Aero is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Aero. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+use core::ptr::NonNull;
+
+use intrusive_collections::UnsafeRef;
+
+use crate::mem::paging::*;
+use crate::utils::sync::Mutex;
+
+#[repr(C)]
+pub struct SlabHeader {
+    pub ptr: UnsafeRef<SmallSlab>,
+}
+
+const_assert_eq!(core::mem::size_of::<SlabHeader>(), 8);
+
+/// For small slabs, the [`BufCtl`]s are stored inline.
+struct BufCtl(Option<NonNull<BufCtl>>);
+
+impl BufCtl {
+    const NULL: Self = Self(None);
+
+    /// Constructs a [`BufCtl`] from a raw pointer.
+    const fn from_ptr(ptr: *mut BufCtl) -> Self {
+        assert!(!ptr.is_null());
+
+        // SAFETY: We have verified above that the pointer is non-null.
+        Self(Some(unsafe { NonNull::new_unchecked(ptr) }))
+    }
+}
+
+const_assert_eq!(core::mem::size_of::<BufCtl>(), 8);
+
+/// Used for allocations smaller than `1/8` of a page.
+pub struct SmallSlab {
+    /// Size of the slab.
+    size: usize,
+    first_free: Mutex<BufCtl>,
+}
+
+impl SmallSlab {
+    pub const fn new(size: usize) -> Self {
+        assert!(size.is_power_of_two());
+
+        Self {
+            size,
+            first_free: Mutex::new(BufCtl::NULL),
+        }
+    }
+
+    pub fn alloc(&self) -> *mut u8 {
+        let mut first_free = self.first_free.lock_irq();
+
+        if let Some(entry) = first_free.0 {
+            *first_free = BufCtl(unsafe { entry.as_ref() }.0);
+            entry.as_ptr().cast()
+        } else {
+            drop(first_free);
+
+            self.expand();
+            self.alloc()
+        }
+    }
+
+    pub fn dealloc(&self, ptr: *mut u8) {
+        assert!(!ptr.is_null());
+
+        let mut first_free = self.first_free.lock_irq();
+
+        let mut new_head = BufCtl::from_ptr(ptr.cast());
+        new_head.0 = first_free.0;
+        *first_free = new_head;
+    }
+
+    fn expand(&self) {
+        let frame: PhysFrame<Size4KiB> = FRAME_ALLOCATOR.allocate_frame().expect("slab: OOM");
+
+        let ptr = frame.start_address().as_hhdm_virt().as_mut_ptr::<u8>();
+        let header_size =
+            align_up(core::mem::size_of::<SlabHeader>() as u64, self.size as u64) as usize;
+
+        let avaliable_size = Size4KiB::SIZE as usize - header_size;
+        let slab_ptr = unsafe { &mut *ptr.cast::<SlabHeader>() };
+
+        // SAFETY: We are constructing an [`UnsafeRef`] from ourselves which is a valid reference.
+        slab_ptr.ptr = unsafe { UnsafeRef::from_raw(self as *const _) };
+
+        let first_free = unsafe { ptr.add(header_size).cast() };
+        *self.first_free.lock_irq() = BufCtl::from_ptr(first_free);
+
+        // Initialize the free-list:
+        //
+        // For objects smaller than 1/8 of a page, A slab is built by allocating a 4KiB page,
+        // placing the slab header at the end, and dividing the rest into equal-size buffers:
+        //
+        // ------------------------------------------------------
+        // | buffer | buffer | buffer | buffer | slab header
+        // ------------------------------------------------------
+        //                          4KiB
+        let max = (avaliable_size / self.size) - 1;
+        let fact = self.size / 8;
+
+        for i in 0..max {
+            unsafe {
+                let entry = first_free.add(i * fact);
+                let next = first_free.add((i + 1) * fact);
+
+                (&mut *entry).0 = Some(NonNull::new_unchecked(next));
+            }
+        }
+
+        unsafe {
+            let entry = &mut *first_free.add(max * fact);
+            *entry = BufCtl::NULL;
+        }
+    }
+
+    pub fn size(&self) -> usize {
+        self.size
+    }
+}
+
+unsafe impl Send for SmallSlab {}
+unsafe impl Sync for SmallSlab {}

src/aero_kernel/src/mem/slab.rs

@@ -37,7 +37,7 @@ unsafe impl Allocator for DmaAllocator {
         // XXX: The DMA buffer must be aligned to a page boundary.
         let size_bytes = layout.size();
 
-        let phys = FRAME_ALLOCATOR.alloc(size_bytes).ok_or(AllocError)?;
+        let phys = FRAME_ALLOCATOR.alloc_zeroed(size_bytes).ok_or(AllocError)?;
         let virt = phys.as_hhdm_virt();
 
         // SAFETY: The frame is aligned and non-null.