diff --git a/src/aero_kernel/src/arch/x86_64/mod.rs b/src/aero_kernel/src/arch/x86_64/mod.rs index 029c1180063..fb81e56246c 100644 --- a/src/aero_kernel/src/arch/x86_64/mod.rs +++ b/src/aero_kernel/src/arch/x86_64/mod.rs @@ -170,10 +170,10 @@ extern "C" fn x86_64_aero_main(boot_info: &'static StivaleStruct) -> ! { gdt::init_boot(); log::info!("loaded bootstrap GDT"); - let mut offset_table = paging::init(mmap_tag).unwrap(); + paging::init(mmap_tag).unwrap(); log::info!("loaded paging"); - alloc::init_heap(&mut offset_table).expect("failed to initialize the kernel heap"); + alloc::init_heap(); log::info!("loaded heap"); paging::init_vm_frames(); diff --git a/src/aero_kernel/src/mem/alloc.rs b/src/aero_kernel/src/mem/alloc.rs index 3f3d3487246..24bd8f696e0 100644 --- a/src/aero_kernel/src/mem/alloc.rs +++ b/src/aero_kernel/src/mem/alloc.rs @@ -21,93 +21,238 @@ use core::alloc; use core::alloc::{GlobalAlloc, Layout}; use core::ptr::NonNull; -use crate::utils::sync::Mutex; -use linked_list_allocator::{align_up, Heap}; +use linked_list_allocator::Heap; -use crate::mem::paging::*; +use crate::utils::sync::Mutex; use crate::AERO_SYSTEM_ALLOCATOR; use super::paging::FRAME_ALLOCATOR; use super::AddressSpace; +use crate::mem::paging::*; const HEAP_MAX_SIZE: usize = 128 * 1024 * 1024; // 128 GiB const HEAP_START: usize = 0xfffff80000000000; const HEAP_END: usize = HEAP_START + HEAP_MAX_SIZE; -pub struct LockedHeap(Mutex); +#[repr(C)] +struct SlabHeader { + ptr: *mut Slab, +} -impl LockedHeap { - /// Creates a new uninitialized instance of the kernel - /// global allocator. - #[inline] - pub const fn new_uninit() -> Self { - Self(Mutex::new(Heap::empty())) +/// The slab is the primary unit of currency in the slab allocator. +struct Slab { + size: usize, + first_free: usize, +} + +impl Slab { + const fn new(size: usize) -> Self { + Self { + size, + first_free: 0, + } } - /// Allocate memory as described by the given `layout`. - /// - /// Returns a pointer to newly-allocated memory, or null to indicate - /// allocation failure. - unsafe fn allocate(&self, layout: alloc::Layout) -> Result, ()> { - // SAFETY: We we need to be careful to not cause a deadlock as the interrupt - // handlers utilize the heap and might interrupt an in-progress allocation. So, we - // lock the interrupts during the allocation. - let mut heap = self.0.lock_irq(); + fn init(&mut self) { + unsafe { + let frame: PhysFrame = FRAME_ALLOCATOR + .allocate_frame() + .expect("slab_init: failed to allocate frame"); + + self.first_free = frame.start_address().as_u64() as usize; + self.first_free += crate::PHYSICAL_MEMORY_OFFSET.as_u64() as usize; + } + + let hdr_size = core::mem::size_of::() as u64; + let aligned_hdr_size = align_up(hdr_size, self.size as u64) as usize; - heap.allocate_first_fit(layout).or_else(|_| { - let heap_top = heap.top(); - let size = align_up(layout.size(), 0x1000); + let avl_size = Size4KiB::SIZE as usize - aligned_hdr_size; - // Check if our heap has not increased beyond the maximum allowed size. - if heap_top + size > HEAP_END { - panic!("The heap size has increased more then {:#x}", HEAP_END) + let slab_ptr = unsafe { &mut *(self.first_free as *mut SlabHeader) }; + slab_ptr.ptr = self as *mut Slab; + + self.first_free += aligned_hdr_size; + + let arr_ptr = self.first_free as *mut usize; + let array = unsafe { core::slice::from_raw_parts_mut(arr_ptr, avl_size) }; + + // A slab is built by allocating a 4KiB page, placing the slab data at + // the end, and dividing the rest into equal-size buffers: + // + // ------------------------------------------------------ + // | buffer | buffer | buffer | buffer | slab header + // ------------------------------------------------------ + // one page + let max = avl_size / self.size - 1; + let fact = self.size / 8; + + for i in 0..max { + unsafe { + array[i * fact] = array.as_ptr().add((i + 1) * fact) as usize; } + } + + array[max * fact] = 0; + } + + fn alloc(&mut self) -> *mut u8 { + if self.first_free == 0 { + self.init(); + } + + let old_free = self.first_free as *mut usize; + + unsafe { + self.first_free = *old_free; + } + + old_free as *mut u8 + } + + fn dealloc(&mut self, ptr: *mut u8) { + if ptr == core::ptr::null_mut() { + panic!("dealloc: attempted to free a nullptr") + } + + let new_head = ptr as *mut usize; + + unsafe { + *new_head = self.first_free; + } + + self.first_free = new_head as usize; + } +} + +struct ProtectedAllocator { + slabs: [Slab; 10], + linked_list_heap: Heap, +} + +struct Allocator { + inner: Mutex, +} - // Else we just have to extend the heap. - let mut address_space = AddressSpace::this(); - let mut offset_table = address_space.offset_page_table(); - - let page_range = { - let heap_start = VirtAddr::new(heap_top as _); - let heap_end = heap_start + size - 1u64; - - let heap_start_page: Page = Page::containing_address(heap_start); - let heap_end_page = Page::containing_address(heap_end); - - Page::range_inclusive(heap_start_page, heap_end_page) - }; - - for page in page_range { - let frame = unsafe { - FRAME_ALLOCATOR - .allocate_frame() - .expect("Failed to allocate frame to extend heap") - }; - - unsafe { - offset_table.map_to( - page, - frame, - PageTableFlags::PRESENT | PageTableFlags::WRITABLE, - &mut FRAME_ALLOCATOR, - ) - } - .expect("Failed to map frame to extend the heap") - .flush(); +impl Allocator { + const fn new() -> Self { + Self { + inner: Mutex::new(ProtectedAllocator { + slabs: [ + Slab::new(8), + Slab::new(16), + Slab::new(24), + Slab::new(32), + Slab::new(48), + Slab::new(64), + Slab::new(128), + Slab::new(256), + Slab::new(512), + Slab::new(1024), + ], + + linked_list_heap: Heap::empty(), + }), + } + } + + fn alloc(&self, layout: Layout) -> *mut u8 { + let mut inner = self.inner.lock_irq(); + + let slab = inner + .slabs + .iter_mut() + .find(|slab| slab.size >= layout.size()); + + if let Some(slab) = slab { + slab.alloc() + } else { + inner + .linked_list_heap + .allocate_first_fit(layout) + .or_else(|_| { + let heap_top = inner.linked_list_heap.top(); + let size = align_up(layout.size() as u64, 0x1000); + + // Check if our heap has not increased beyond the maximum allowed size. + if heap_top + size as usize > HEAP_END { + panic!("the heap size has increased more then {:#x}", HEAP_END) + } + + // Else we just have to extend the heap. + let mut address_space = AddressSpace::this(); + let mut offset_table = address_space.offset_page_table(); + + let page_range = { + let heap_start = VirtAddr::new(heap_top as _); + let heap_end = heap_start + size - 1u64; + + let heap_start_page: Page = Page::containing_address(heap_start); + let heap_end_page = Page::containing_address(heap_end); + + Page::range_inclusive(heap_start_page, heap_end_page) + }; + + for page in page_range { + let frame = unsafe { + FRAME_ALLOCATOR + .allocate_frame() + .expect("Failed to allocate frame to extend heap") + }; + + unsafe { + offset_table.map_to( + page, + frame, + PageTableFlags::PRESENT | PageTableFlags::WRITABLE, + &mut FRAME_ALLOCATOR, + ) + } + .expect("Failed to map frame to extend the heap") + .flush(); + } + + unsafe { + inner.linked_list_heap.extend(size as usize); // Now extend the heap. + inner.linked_list_heap.allocate_first_fit(layout) // And try again. + } + }) + .expect("alloc: memory exhausted") + .as_ptr() + } + } + + fn dealloc(&self, ptr: *mut u8, layout: Layout) { + let mut inner = self.inner.lock_irq(); + + let size = layout.size(); + + if size >= Size4KiB::SIZE as usize { + unsafe { + inner + .linked_list_heap + .deallocate(NonNull::new_unchecked(ptr), layout); } - heap.extend(size); // Now extend the heap. - heap.allocate_first_fit(layout) // And try again. - }) + return; + } + + let slab_header = (ptr as usize & !(0xfff)) as *mut SlabHeader; + + let slab_header = unsafe { &mut *slab_header }; + let slab = unsafe { &mut *slab_header.ptr }; + + slab.dealloc(ptr); } +} + +pub struct LockedHeap(Allocator); - /// Initializes an empty heap. - /// - /// ## Safety - /// This function should only be called once and the provided `start` address - /// should be a valid address. - unsafe fn init(&self, start: usize, size: usize) { - self.0.lock().init(start, size); +impl LockedHeap { + /// Creates a new uninitialized instance of the kernel + /// global allocator. + #[inline] + pub const fn new_uninit() -> Self { + Self(Allocator::new()) } } @@ -220,7 +365,10 @@ unsafe impl GlobalAlloc for LockedHeap { // necessary and sufficient. debug_assert!(layout.size() < usize::MAX - (layout.align() - 1)); - let ptr = self.allocate(layout).unwrap().as_ptr(); + // SAFETY: We we need to be careful to not cause a deadlock as the interrupt + // handlers utilize the heap and might interrupt an in-progress allocation. So, we + // lock the interrupts during the allocation. + let ptr = self.0.alloc(layout); #[cfg(feature = "kmemleak")] kmemleak::MEM_LEAK_CATCHER.track_caller(ptr, layout); @@ -235,9 +383,7 @@ unsafe impl GlobalAlloc for LockedHeap { #[cfg(feature = "kmemleak")] kmemleak::MEM_LEAK_CATCHER.unref(ptr); - self.0 - .lock_irq() - .deallocate(NonNull::new_unchecked(ptr), layout) + self.0.dealloc(ptr, layout) } } @@ -251,29 +397,33 @@ fn alloc_error_handler(layout: alloc::Layout) -> ! { } /// Initialize the heap at the [HEAP_START]. -pub fn init_heap(offset_table: &mut OffsetPageTable) -> Result<(), MapToError> { - let frame: PhysFrame = unsafe { - FRAME_ALLOCATOR - .allocate_frame() - .ok_or(MapToError::FrameAllocationFailed)? - }; - +pub fn init_heap() { unsafe { - offset_table.map_to( - Page::containing_address(VirtAddr::new(HEAP_START as _)), - frame, - PageTableFlags::PRESENT | PageTableFlags::WRITABLE, - &mut FRAME_ALLOCATOR, - ) - }? - .flush(); + let mut address_space = AddressSpace::this(); + let mut offset_table = address_space.offset_page_table(); - unsafe { - AERO_SYSTEM_ALLOCATOR.init(HEAP_START, 4096); + let frame: PhysFrame = FRAME_ALLOCATOR + .allocate_frame() + .expect("init_heap: failed to allocate frame for the linked list allocator"); + + offset_table + .map_to( + Page::containing_address(VirtAddr::new(HEAP_START as _)), + frame, + PageTableFlags::PRESENT | PageTableFlags::WRITABLE, + &mut FRAME_ALLOCATOR, + ) + .expect("init_heap: failed to initialize the heap") + .flush(); + + AERO_SYSTEM_ALLOCATOR + .0 + .inner + .lock_irq() + .linked_list_heap + .init(HEAP_START, Size4KiB::SIZE as usize); } #[cfg(feature = "kmemleak")] kmemleak::MEM_LEAK_CATCHER.init(); - - Ok(()) }