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_sync::{RawSpinLockGuard, SpinLock};
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: SpinLock<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: SpinLock::new(SlabAllocator::new()),
151        }
152    }
153
154    #[inline]
155    fn inner(&self) -> RawSpinLockGuard<'_, SlabAllocator<PAGE_SIZE>> {
156        // SAFETY: per-CPU slab access preserves the legacy raw-lock contract:
157        // the pool selects one owner CPU and callers exclude local re-entry.
158        unsafe { self.inner.lock_raw() }
159    }
160
161    /// Reset the inner slab allocator to an empty state.
162    pub fn reset(&self) {
163        *self.inner() = SlabAllocator::new();
164    }
165
166    /// Return this slab's logical CPU id.
167    pub const fn cpu_id(&self) -> usize {
168        self.cpu_id as usize
169    }
170
171    /// Allocate one object.
172    pub fn alloc(&self, layout: Layout) -> AllocResult<SlabAllocResult> {
173        self.inner().alloc(layout)
174    }
175
176    /// Register a freshly allocated slab page.
177    pub fn add_slab(&self, size_class: SizeClass, base: usize, bytes: usize) {
178        self.inner().add_slab(size_class, base, bytes, self.cpu_id);
179    }
180
181    /// Free an object on the owner CPU path.
182    pub fn dealloc_local(&self, ptr: NonNull<u8>, layout: Layout) -> SlabDeallocResult {
183        self.inner().dealloc(ptr, layout)
184    }
185
186    /// Queue an object onto this slab's remote-free list.
187    pub fn dealloc_remote(&self, ptr: NonNull<u8>) {
188        unsafe { SlabPageHeader::remote_free_object(ptr, self.cpu_id, PAGE_SIZE) };
189    }
190}
191
192impl<const PAGE_SIZE: usize, const N: usize> StaticSlabPool<PAGE_SIZE, N> {
193    /// Create a static slab pool from pre-built per-CPU slabs and a CPU-id hook.
194    pub const fn new(slabs: [PerCpuSlab<PAGE_SIZE>; N], current_cpu_id: fn() -> usize) -> Self {
195        Self {
196            slabs,
197            current_cpu_id,
198        }
199    }
200}
201
202impl<const PAGE_SIZE: usize> SlabAllocator<PAGE_SIZE> {
203    /// Create a new (empty) slab allocator.  No pages are owned yet.
204    pub const fn new() -> Self {
205        Self {
206            caches: [
207                SlabCache::new(SizeClass::Bytes8),
208                SlabCache::new(SizeClass::Bytes16),
209                SlabCache::new(SizeClass::Bytes32),
210                SlabCache::new(SizeClass::Bytes64),
211                SlabCache::new(SizeClass::Bytes128),
212                SlabCache::new(SizeClass::Bytes256),
213                SlabCache::new(SizeClass::Bytes512),
214                SlabCache::new(SizeClass::Bytes1024),
215                SlabCache::new(SizeClass::Bytes2048),
216            ],
217        }
218    }
219}
220
221impl<const PAGE_SIZE: usize> Default for SlabAllocator<PAGE_SIZE> {
222    fn default() -> Self {
223        Self::new()
224    }
225}
226
227impl<const PAGE_SIZE: usize> SlabAllocator<PAGE_SIZE> {
228    /// Try to allocate an object matching `layout`.
229    ///
230    /// If the matching cache is exhausted, [`SlabAllocResult::NeedsSlab`] is returned
231    /// so the caller can supply pages and retry.
232    pub fn alloc(&mut self, layout: Layout) -> AllocResult<SlabAllocResult> {
233        let sc = SizeClass::from_layout(layout).ok_or(AllocError::InvalidParam)?;
234        let cache = &mut self.caches[sc.index()];
235
236        match cache.alloc_object::<PAGE_SIZE>() {
237            Some(addr) => {
238                // SAFETY: `addr` is non-null, aligned, and within a live slab page.
239                let ptr = unsafe { NonNull::new_unchecked(addr as *mut u8) };
240                Ok(SlabAllocResult::Allocated(ptr))
241            }
242            None => Ok(SlabAllocResult::NeedsSlab {
243                size_class: sc,
244                pages: sc.slab_pages(PAGE_SIZE),
245            }),
246        }
247    }
248
249    /// Free an object previously allocated with [`alloc`](Self::alloc).
250    ///
251    /// This is the **local** (owner-CPU) path.  Cross-CPU frees should go through
252    /// [`SlabPageHeader::remote_free`] directly (see [`GlobalAllocator`]).
253    pub fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) -> SlabDeallocResult {
254        let sc = SizeClass::from_layout(layout).expect("layout exceeds slab size");
255        let cache = &mut self.caches[sc.index()];
256
257        match cache.dealloc_object::<PAGE_SIZE>(ptr.as_ptr() as usize) {
258            CacheDeallocResult::Done => SlabDeallocResult::Done,
259            CacheDeallocResult::FreeSlab { base, pages } => {
260                SlabDeallocResult::FreeSlab { base, pages }
261            }
262        }
263    }
264
265    /// Supply a freshly allocated slab page to the given size class.
266    ///
267    /// `base` is the virtual address of the page(s), `bytes` = pages × PAGE_SIZE.
268    pub fn add_slab(&mut self, size_class: SizeClass, base: usize, bytes: usize, owner_cpu: u16) {
269        self.caches[size_class.index()].add_slab(base, bytes, owner_cpu);
270    }
271}
272
273impl<const PAGE_SIZE: usize> SlabTrait for PerCpuSlab<PAGE_SIZE> {
274    fn cpu_id(&self) -> usize {
275        PerCpuSlab::cpu_id(self)
276    }
277
278    fn page_size(&self) -> usize {
279        PAGE_SIZE
280    }
281
282    fn alloc(&self, layout: Layout) -> AllocResult<SlabAllocResult> {
283        PerCpuSlab::alloc(self, layout)
284    }
285
286    fn add_slab(&self, size_class: SizeClass, base: usize, bytes: usize) {
287        PerCpuSlab::add_slab(self, size_class, base, bytes)
288    }
289
290    fn dealloc_local(&self, ptr: NonNull<u8>, layout: Layout) -> SlabDeallocResult {
291        PerCpuSlab::dealloc_local(self, ptr, layout)
292    }
293}
294
295impl<const PAGE_SIZE: usize, const N: usize> SlabPoolTrait for StaticSlabPool<PAGE_SIZE, N> {
296    fn current_slab(&self) -> &dyn SlabTrait {
297        &self.slabs[(self.current_cpu_id)()]
298    }
299
300    fn owner_slab(&self, cpu_idx: usize) -> &dyn SlabTrait {
301        &self.slabs[cpu_idx]
302    }
303}
buddy_slab_allocator/slab/mod.rs

buddy_slab_allocator/slab/
mod.rs
