buddy_slab_allocator/slab/
mod.rs

1//! Slab allocator — bitmap-based with lock-free cross-CPU freeing.
2//!
3//! The [`SlabAllocator`] is a standalone component that manages object allocation
4//! within pre-supplied slab pages.  It does **not** allocate pages itself; instead
5//! it returns [`SlabAllocResult::NeedsSlab`] to request pages from the caller.
6//!
7//! Cross-CPU frees go through the lock-free [`SlabPageHeader::remote_free`] path.
8
9pub mod cache;
10pub mod page;
11pub mod size_class;
12
13use core::{alloc::Layout, ptr::NonNull};
14
15use ax_kspin::SpinRaw as SpinMutex;
16use cache::{CacheDeallocResult, SlabCache};
17pub use page::SlabPageHeader;
18pub use size_class::SizeClass;
19
20use crate::error::{AllocError, AllocResult};
21
22/// Result of a slab allocation attempt.
23pub enum SlabAllocResult {
24    /// Object successfully allocated.
25    Allocated(NonNull<u8>),
26    /// The slab cache for this size class has no free objects.
27    /// The caller should allocate `pages` pages from the buddy allocator,
28    /// call [`SlabAllocator::add_slab`], and retry.
29    NeedsSlab { size_class: SizeClass, pages: usize },
30}
31
32/// Result of a slab deallocation.
33pub enum SlabDeallocResult {
34    /// Object freed, nothing else to do.
35    Done,
36    /// The slab page at `base` became empty and should be returned to the buddy.
37    FreeSlab { base: usize, pages: usize },
38}
39
40/// Result of a pool-mediated slab deallocation.
41pub enum SlabPoolDeallocResult {
42    /// Object freed on the local CPU path.
43    Done,
44    /// Object was queued onto the owner's remote-free list.
45    RemoteQueued,
46    /// The slab page at `base` became empty and should be returned to the buddy.
47    FreeSlab { base: usize, pages: usize },
48}
49
50/// Object-safe slab interface used by [`crate::GlobalAllocator`] EII hooks.
51pub trait SlabTrait: Sync {
52    /// Logical CPU id this slab belongs to.
53    fn cpu_id(&self) -> usize;
54
55    /// Page size used by this slab.
56    fn page_size(&self) -> usize;
57
58    /// Allocate one object.
59    fn alloc(&self, layout: Layout) -> AllocResult<SlabAllocResult>;
60
61    /// Register a freshly allocated slab page.
62    fn add_slab(&self, size_class: SizeClass, base: usize, bytes: usize);
63
64    /// Free an object on the owner CPU path.
65    fn dealloc_local(&self, ptr: NonNull<u8>, layout: Layout) -> SlabDeallocResult;
66
67    /// Free an object on the remote CPU path.
68    fn dealloc_remote(&self, ptr: NonNull<u8>) {
69        let owner_cpu = u16::try_from(self.cpu_id()).expect("CPU id exceeds slab owner range");
70        unsafe { SlabPageHeader::remote_free_object(ptr, owner_cpu, self.page_size()) };
71    }
72}
73
74/// Object-safe slab-pool interface used by [`crate::GlobalAllocator`] EII hooks.
75pub trait SlabPoolTrait: Sync {
76    /// Return the slab belonging to the current CPU.
77    fn current_slab(&self) -> &dyn SlabTrait;
78
79    /// Return the owner slab for the given CPU.
80    fn owner_slab(&self, cpu_idx: usize) -> &dyn SlabTrait;
81
82    /// Logical CPU id of the current CPU.
83    fn current_cpu_id(&self) -> usize {
84        self.current_slab().cpu_id()
85    }
86
87    /// Allocate one object from the current CPU's slab.
88    fn alloc(&self, layout: Layout) -> AllocResult<SlabAllocResult> {
89        self.current_slab().alloc(layout)
90    }
91
92    /// Register a freshly allocated slab page in the current CPU's slab.
93    fn add_slab(&self, size_class: SizeClass, base: usize, bytes: usize) {
94        self.current_slab().add_slab(size_class, base, bytes)
95    }
96
97    /// Free an object, routing to local or remote slab ownership as needed.
98    fn dealloc(&self, ptr: NonNull<u8>, layout: Layout, owner_cpu: usize) -> SlabPoolDeallocResult {
99        if owner_cpu == self.current_cpu_id() {
100            match self.current_slab().dealloc_local(ptr, layout) {
101                SlabDeallocResult::Done => SlabPoolDeallocResult::Done,
102                SlabDeallocResult::FreeSlab { base, pages } => {
103                    SlabPoolDeallocResult::FreeSlab { base, pages }
104                }
105            }
106        } else {
107            self.owner_slab(owner_cpu).dealloc_remote(ptr);
108            SlabPoolDeallocResult::RemoteQueued
109        }
110    }
111}
112
113/// Convenience helpers for callback-style slab access.
114pub trait SlabPoolExt: SlabPoolTrait {
115    /// Access the current CPU's slab via a callback.
116    fn with_current_slab<R>(&self, f: impl FnOnce(&dyn SlabTrait) -> R) -> R {
117        f(self.current_slab())
118    }
119
120    /// Access the given owner's slab via a callback.
121    fn with_owner_slab<R>(&self, cpu_idx: usize, f: impl FnOnce(&dyn SlabTrait) -> R) -> R {
122        f(self.owner_slab(cpu_idx))
123    }
124}
125
126impl<T: ?Sized + SlabPoolTrait> SlabPoolExt for T {}
127
128/// Standalone slab allocator (one per CPU or standalone use).
129pub struct SlabAllocator<const PAGE_SIZE: usize = 0x1000> {
130    caches: [SlabCache; SizeClass::COUNT],
131}
132
133/// Default per-CPU slab wrapper used by EII integrators.
134pub struct PerCpuSlab<const PAGE_SIZE: usize = 0x1000> {
135    cpu_id: u16,
136    inner: SpinMutex<SlabAllocator<PAGE_SIZE>>,
137}
138
139/// Default static slab-pool wrapper used by EII integrators.
140pub struct StaticSlabPool<const PAGE_SIZE: usize = 0x1000, const N: usize = 1> {
141    slabs: [PerCpuSlab<PAGE_SIZE>; N],
142    current_cpu_id: fn() -> usize,
143}
144
145impl<const PAGE_SIZE: usize> PerCpuSlab<PAGE_SIZE> {
146    /// Create an empty per-CPU slab wrapper for `cpu_id`.
147    pub const fn new(cpu_id: u16) -> Self {
148        Self {
149            cpu_id,
150            inner: SpinMutex::new(SlabAllocator::new()),
151        }
152    }
153
154    /// Reset the inner slab allocator to an empty state.
155    pub fn reset(&self) {
156        *self.inner.lock() = SlabAllocator::new();
157    }
158
159    /// Return this slab's logical CPU id.
160    pub const fn cpu_id(&self) -> usize {
161        self.cpu_id as usize
162    }
163
164    /// Allocate one object.
165    pub fn alloc(&self, layout: Layout) -> AllocResult<SlabAllocResult> {
166        self.inner.lock().alloc(layout)
167    }
168
169    /// Register a freshly allocated slab page.
170    pub fn add_slab(&self, size_class: SizeClass, base: usize, bytes: usize) {
171        self.inner
172            .lock()
173            .add_slab(size_class, base, bytes, self.cpu_id);
174    }
175
176    /// Free an object on the owner CPU path.
177    pub fn dealloc_local(&self, ptr: NonNull<u8>, layout: Layout) -> SlabDeallocResult {
178        self.inner.lock().dealloc(ptr, layout)
179    }
180
181    /// Queue an object onto this slab's remote-free list.
182    pub fn dealloc_remote(&self, ptr: NonNull<u8>) {
183        unsafe { SlabPageHeader::remote_free_object(ptr, self.cpu_id, PAGE_SIZE) };
184    }
185}
186
187impl<const PAGE_SIZE: usize, const N: usize> StaticSlabPool<PAGE_SIZE, N> {
188    /// Create a static slab pool from pre-built per-CPU slabs and a CPU-id hook.
189    pub const fn new(slabs: [PerCpuSlab<PAGE_SIZE>; N], current_cpu_id: fn() -> usize) -> Self {
190        Self {
191            slabs,
192            current_cpu_id,
193        }
194    }
195}
196
197impl<const PAGE_SIZE: usize> SlabAllocator<PAGE_SIZE> {
198    /// Create a new (empty) slab allocator.  No pages are owned yet.
199    pub const fn new() -> Self {
200        Self {
201            caches: [
202                SlabCache::new(SizeClass::Bytes8),
203                SlabCache::new(SizeClass::Bytes16),
204                SlabCache::new(SizeClass::Bytes32),
205                SlabCache::new(SizeClass::Bytes64),
206                SlabCache::new(SizeClass::Bytes128),
207                SlabCache::new(SizeClass::Bytes256),
208                SlabCache::new(SizeClass::Bytes512),
209                SlabCache::new(SizeClass::Bytes1024),
210                SlabCache::new(SizeClass::Bytes2048),
211            ],
212        }
213    }
214}
215
216impl<const PAGE_SIZE: usize> Default for SlabAllocator<PAGE_SIZE> {
217    fn default() -> Self {
218        Self::new()
219    }
220}
221
222impl<const PAGE_SIZE: usize> SlabAllocator<PAGE_SIZE> {
223    /// Try to allocate an object matching `layout`.
224    ///
225    /// If the matching cache is exhausted, [`SlabAllocResult::NeedsSlab`] is returned
226    /// so the caller can supply pages and retry.
227    pub fn alloc(&mut self, layout: Layout) -> AllocResult<SlabAllocResult> {
228        let sc = SizeClass::from_layout(layout).ok_or(AllocError::InvalidParam)?;
229        let cache = &mut self.caches[sc.index()];
230
231        match cache.alloc_object::<PAGE_SIZE>() {
232            Some(addr) => {
233                // SAFETY: `addr` is non-null, aligned, and within a live slab page.
234                let ptr = unsafe { NonNull::new_unchecked(addr as *mut u8) };
235                Ok(SlabAllocResult::Allocated(ptr))
236            }
237            None => Ok(SlabAllocResult::NeedsSlab {
238                size_class: sc,
239                pages: sc.slab_pages(PAGE_SIZE),
240            }),
241        }
242    }
243
244    /// Free an object previously allocated with [`alloc`](Self::alloc).
245    ///
246    /// This is the **local** (owner-CPU) path.  Cross-CPU frees should go through
247    /// [`SlabPageHeader::remote_free`] directly (see [`GlobalAllocator`]).
248    pub fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) -> SlabDeallocResult {
249        let sc = SizeClass::from_layout(layout).expect("layout exceeds slab size");
250        let cache = &mut self.caches[sc.index()];
251
252        match cache.dealloc_object::<PAGE_SIZE>(ptr.as_ptr() as usize) {
253            CacheDeallocResult::Done => SlabDeallocResult::Done,
254            CacheDeallocResult::FreeSlab { base, pages } => {
255                SlabDeallocResult::FreeSlab { base, pages }
256            }
257        }
258    }
259
260    /// Supply a freshly allocated slab page to the given size class.
261    ///
262    /// `base` is the virtual address of the page(s), `bytes` = pages × PAGE_SIZE.
263    pub fn add_slab(&mut self, size_class: SizeClass, base: usize, bytes: usize, owner_cpu: u16) {
264        self.caches[size_class.index()].add_slab(base, bytes, owner_cpu);
265    }
266}
267
268impl<const PAGE_SIZE: usize> SlabTrait for PerCpuSlab<PAGE_SIZE> {
269    fn cpu_id(&self) -> usize {
270        PerCpuSlab::cpu_id(self)
271    }
272
273    fn page_size(&self) -> usize {
274        PAGE_SIZE
275    }
276
277    fn alloc(&self, layout: Layout) -> AllocResult<SlabAllocResult> {
278        PerCpuSlab::alloc(self, layout)
279    }
280
281    fn add_slab(&self, size_class: SizeClass, base: usize, bytes: usize) {
282        PerCpuSlab::add_slab(self, size_class, base, bytes)
283    }
284
285    fn dealloc_local(&self, ptr: NonNull<u8>, layout: Layout) -> SlabDeallocResult {
286        PerCpuSlab::dealloc_local(self, ptr, layout)
287    }
288}
289
290impl<const PAGE_SIZE: usize, const N: usize> SlabPoolTrait for StaticSlabPool<PAGE_SIZE, N> {
291    fn current_slab(&self) -> &dyn SlabTrait {
292        &self.slabs[(self.current_cpu_id)()]
293    }
294
295    fn owner_slab(&self, cpu_idx: usize) -> &dyn SlabTrait {
296        &self.slabs[cpu_idx]
297    }
298}
buddy_slab_allocator/slab/mod.rs

buddy_slab_allocator/slab/
mod.rs
