ostd/mm/heap/
mod.rs

1// SPDX-License-Identifier: MPL-2.0
2
3//! Manages the kernel heap using slab or buddy allocation strategies.
4
5use core::{
6    alloc::{AllocError, GlobalAlloc, Layout},
7    ptr::NonNull,
8};
9
10use crate::mm::Vaddr;
11
12mod slab;
13mod slot;
14mod slot_list;
15
16pub use self::{
17    slab::{Slab, SlabMeta},
18    slot::{HeapSlot, SlotInfo},
19    slot_list::SlabSlotList,
20};
21
22/// The trait for the global heap allocator.
23///
24/// By providing the slab ([`Slab`]) and heap slot ([`HeapSlot`])
25/// mechanisms, OSTD allows users to implement their own kernel heap in a safe
26/// manner, as an alternative to the unsafe [`core::alloc::GlobalAlloc`].
27///
28/// To provide the global heap allocator, use [`crate::global_heap_allocator`]
29/// to mark a static variable that implements this trait. Use
30/// [`crate::global_heap_allocator_slot_map`] to specify the sizes of
31/// slots for different layouts. This latter restriction may be lifted in the
32/// future.
33pub trait GlobalHeapAllocator: Sync {
34    /// Allocates a [`HeapSlot`] according to the layout.
35    ///
36    /// OSTD calls this method to allocate memory from the global heap.
37    ///
38    /// The returned [`HeapSlot`] must be valid for the layout, i.e., the size
39    /// must be at least the size of the layout and the alignment must be at
40    /// least the alignment of the layout. Furthermore, the size of the
41    /// returned [`HeapSlot`] must match the size returned by the function
42    /// marked with [`crate::global_heap_allocator_slot_map`].
43    fn alloc(&self, layout: Layout) -> Result<HeapSlot, AllocError>;
44
45    /// Deallocates a [`HeapSlot`].
46    ///
47    /// OSTD calls this method to deallocate memory back to the global heap.
48    ///
49    /// Each deallocation must correspond to exactly one previous allocation. The provided
50    /// [`HeapSlot`] must match the one returned from the original allocation.
51    fn dealloc(&self, slot: HeapSlot) -> Result<(), AllocError>;
52}
53
54unsafe extern "Rust" {
55    /// The reference to the global heap allocator generated by the
56    /// [`crate::global_heap_allocator`] attribute.
57    static __GLOBAL_HEAP_ALLOCATOR_REF: &'static dyn GlobalHeapAllocator;
58
59    /// Gets the size and type of heap slots to serve allocations of the layout.
60    /// See [`crate::global_heap_allocator_slot_map`].
61    fn __GLOBAL_HEAP_SLOT_INFO_FROM_LAYOUT(layout: Layout) -> Option<SlotInfo>;
62}
63
64/// Gets the reference to the user-defined global heap allocator.
65fn get_global_heap_allocator() -> &'static dyn GlobalHeapAllocator {
66    // SAFETY: This up-call is redirected safely to Rust code by OSDK.
67    unsafe { __GLOBAL_HEAP_ALLOCATOR_REF }
68}
69
70/// Gets the size and type of heap slots to serve allocations of the layout.
71///
72/// This function is defined by the OSTD user and should be idempotent, as we
73/// require it to be implemented as a `const fn`.
74///
75/// See [`crate::global_heap_allocator_slot_map`].
76fn slot_size_from_layout(layout: Layout) -> Option<SlotInfo> {
77    // SAFETY: This up-call is redirected safely to Rust code by OSDK.
78    unsafe { __GLOBAL_HEAP_SLOT_INFO_FROM_LAYOUT(layout) }
79}
80
81macro_rules! abort_with_message {
82    ($($arg:tt)*) => {
83        log::error!($($arg)*);
84        crate::panic::abort();
85    };
86}
87
88#[alloc_error_handler]
89fn handle_alloc_error(layout: core::alloc::Layout) -> ! {
90    abort_with_message!("Heap allocation error, layout = {:#x?}", layout);
91}
92
93#[global_allocator]
94static HEAP_ALLOCATOR: AllocDispatch = AllocDispatch;
95
96struct AllocDispatch;
97
98// TODO: Somehow restrict unwinding in the user-provided global allocator.
99// Panicking should be fine, but we shouldn't unwind on panics.
100unsafe impl GlobalAlloc for AllocDispatch {
101    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
102        let Some(required_slot) = slot_size_from_layout(layout) else {
103            abort_with_message!("Heap allocation size not found for layout = {:#x?}", layout);
104        };
105
106        let res = get_global_heap_allocator().alloc(layout);
107        let Ok(slot) = res else {
108            return core::ptr::null_mut();
109        };
110
111        if required_slot.size() != slot.size()
112            || slot.size() < layout.size()
113            || !(slot.as_ptr() as Vaddr).is_multiple_of(layout.align())
114        {
115            abort_with_message!(
116                "Heap allocation mismatch: slot ptr = {:p}, size = {:x}; layout = {:#x?}; required_slot = {:#x?}",
117                slot.as_ptr(),
118                slot.size(),
119                layout,
120                required_slot,
121            );
122        }
123
124        slot.as_ptr()
125    }
126
127    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
128        // Now we restore the `HeapSlot` from the pointer and the layout.
129        let Some(required_slot) = slot_size_from_layout(layout) else {
130            abort_with_message!(
131                "Heap deallocation size not found for layout = {:#x?}",
132                layout
133            );
134        };
135
136        // SAFETY: The validity of the pointer is guaranteed by the caller. The
137        // size must match the size of the slot when it was allocated, since we
138        // require `slot_size_from_layout` to be idempotent.
139        let slot = unsafe { HeapSlot::new(NonNull::new_unchecked(ptr), required_slot) };
140        let res = get_global_heap_allocator().dealloc(slot);
141
142        if res.is_err() {
143            abort_with_message!(
144                "Heap deallocation error, ptr = {:p}, layout = {:#x?}, required_slot = {:#x?}",
145                ptr,
146                layout,
147                required_slot,
148            );
149        }
150    }
151}
ostd/mm/heap/mod.rs

ostd/mm/heap/
mod.rs
