obkrnl/malloc/
mod.rs

1use self::vm::VmHeap;
2use crate::context::current_thread;
3use crate::lock::Mutex;
4use crate::vm::Vm;
5use alloc::boxed::Box;
6use core::alloc::{GlobalAlloc, Layout};
7use core::cell::{RefCell, UnsafeCell};
8use core::hint::unreachable_unchecked;
9use core::ptr::{NonNull, null_mut};
10use talc::{ClaimOnOom, Span, Talc};
11
12mod vm;
13
14/// Implementation of [`GlobalAlloc`] for objects belong to kernel space.
15///
16/// This allocator has 2 stages. The first stage will allocate a memory from a static buffer (AKA
17/// arena). This stage will be primary used for bootstrapping the kernel. The second stage will be
18/// activated once the required subsystems has been initialized.
19///
20/// The first stage is **not** thread safe so stage 2 must be activated before start a new CPU.
21pub struct KernelHeap {
22    stage: UnsafeCell<Stage>,
23    primitive_ptr: *const u8,
24    primitive_end: *const u8,
25}
26
27impl KernelHeap {
28    /// # Safety
29    /// The specified memory must be valid for reads and writes and it must be exclusively available
30    /// to [`KernelHeap`].
31    pub const unsafe fn new<const L: usize>(primitive: *mut [u8; L]) -> Self {
32        let primitive_ptr = primitive.cast();
33
34        // SAFETY: The safety requirement of our function satify the safety requirement of
35        // ClaimOnOom::new().
36        let primitive = unsafe { Talc::new(ClaimOnOom::new(Span::from_array(primitive))) };
37
38        Self {
39            stage: UnsafeCell::new(Stage::One(RefCell::new(primitive))),
40            primitive_ptr,
41            // SAFETY: L is a length of primitive_ptr so the resulting pointer is valid.
42            primitive_end: unsafe { primitive_ptr.add(L) },
43        }
44    }
45
46    /// # Safety
47    /// This must be called by main CPU and can be called only once.
48    pub unsafe fn activate_stage2(&self, vm: &'static Vm) {
49        // Setup VM  heap using primitive heap.
50        let vm = Box::new(VmHeap::new(vm));
51
52        // What we are doing here is highly unsafe. Do not edit the code after this unless you know
53        // what you are doing!
54        let stage = self.stage.get();
55        let primitive = match unsafe { stage.read() } {
56            Stage::One(v) => Mutex::new(v.into_inner()),
57            // SAFETY: The safety requirement of our function make this unreachable.
58            Stage::Two(_, _) => unsafe { unreachable_unchecked() },
59        };
60
61        // Switch to stage 2 WITHOUT dropping the value contained in Stage::One. What we did here is
62        // moving the value from Stage::One to Stage::Two.
63        unsafe { stage.write(Stage::Two(vm, primitive)) };
64    }
65}
66
67unsafe impl GlobalAlloc for KernelHeap {
68    #[inline(never)]
69    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
70        // If stage 2 has not activated yet then this function is not allowed to access the CPU
71        // context due to it can be called before the context has been activated.
72
73        // SAFETY: GlobalAlloc::alloc required layout to be non-zero.
74        match unsafe { &*self.stage.get() } {
75            Stage::One(primitive) => unsafe {
76                primitive
77                    .borrow_mut()
78                    .malloc(layout)
79                    .map_or(null_mut(), |v| v.as_ptr())
80            },
81            Stage::Two(vm, primitive) => match current_thread().active_heap_guard() {
82                0 => unsafe { vm.alloc(layout) },
83                _ => unsafe {
84                    primitive
85                        .lock()
86                        .malloc(layout)
87                        .map_or(null_mut(), |v| v.as_ptr())
88                },
89            },
90        }
91    }
92
93    #[inline(never)]
94    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
95        // If stage 2 has not activated yet then this function is not allowed to access the CPU
96        // context due to it can be called before the context has been activated.
97
98        // SAFETY: GlobalAlloc::dealloc required ptr to be the same one that returned
99        // from our GlobalAlloc::alloc and layout to be the same one that passed to it.
100        match unsafe { &*self.stage.get() } {
101            Stage::One(primitive) => unsafe {
102                primitive
103                    .borrow_mut()
104                    .free(NonNull::new_unchecked(ptr), layout)
105            },
106            Stage::Two(vm, primitive) => {
107                if ptr.cast_const() >= self.primitive_ptr && ptr.cast_const() < self.primitive_end {
108                    unsafe { primitive.lock().free(NonNull::new_unchecked(ptr), layout) }
109                } else {
110                    // SAFETY: ptr is not owned by primitive heap so with the requirements of
111                    // GlobalAlloc::dealloc the ptr will be owned by VM heap for sure.
112                    unsafe { vm.dealloc(ptr, layout) };
113                }
114            }
115        }
116    }
117}
118
119// We impose restriction on the user to activate stage 2 before going multi-threaded.
120unsafe impl Send for KernelHeap {}
121unsafe impl Sync for KernelHeap {}
122
123/// Stage of [KernelHeap].
124enum Stage {
125    One(RefCell<Talc<ClaimOnOom>>),
126    Two(Box<VmHeap>, Mutex<Talc<ClaimOnOom>>),
127}
obkrnl/malloc/mod.rs

obkrnl/malloc/
mod.rs
