vmalloc: initial implementation

Signed-off-by: Andy-Python-Programmer <[email protected]>

Author: Andy-Python-Programmer

src/aero_kernel/src/mem/alloc.rs

@@ -19,21 +19,13 @@
 
 use core::alloc;
 use core::alloc::{GlobalAlloc, Layout};
-use core::ptr::NonNull;
-
-use linked_list_allocator::Heap;
 
 use crate::utils::sync::Mutex;
-use crate::AERO_SYSTEM_ALLOCATOR;
 
 use super::paging::FRAME_ALLOCATOR;
-use super::AddressSpace;
+use super::vmalloc;
 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;
-
 #[repr(C)]
 struct SlabHeader {
     ptr: *mut Slab,
@@ -125,7 +117,6 @@ impl Slab {
 
 struct ProtectedAllocator {
     slabs: [Slab; 10],
-    linked_list_heap: Heap,
 }
 
 struct Allocator {
@@ -148,8 +139,6 @@ impl Allocator {
                     Slab::new(512),
                     Slab::new(1024),
                 ],
-
-                linked_list_heap: Heap::empty(),
             }),
         }
     }
@@ -165,71 +154,21 @@ impl Allocator {
         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,
-                            )
-                        }
-                        .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()
+            let size = align_up(layout.size() as _, layout.align() as _) / Size4KiB::SIZE;
+
+            vmalloc::get_vmalloc()
+                .alloc(size as usize)
+                .map(|addr| addr.as_mut_ptr::<u8>())
+                .unwrap_or(core::ptr::null_mut())
         }
     }
 
     fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        let mut inner = self.inner.lock_irq();
-        let address = ptr as usize;
-
-        if address >= HEAP_START && address < HEAP_END {
-            unsafe {
-                inner
-                    .linked_list_heap
-                    .deallocate(NonNull::new_unchecked(ptr), layout);
-            }
+        let _inner = self.inner.lock_irq();
+        let address = VirtAddr::new(ptr as u64);
 
+        if address >= vmalloc::VMALLOC_START && address < vmalloc::VMALLOC_END {
+            vmalloc::get_vmalloc().dealloc(address, layout.size() / Size4KiB::SIZE as usize);
             return;
         }
 
@@ -395,30 +334,7 @@ fn alloc_error_handler(layout: alloc::Layout) -> ! {
 
 /// Initialize the heap at the [HEAP_START].
 pub fn init_heap() {
-    unsafe {
-        let mut address_space = AddressSpace::this();
-        let mut offset_table = address_space.offset_page_table();
-
-        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,
-            )
-            .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);
-    }
+    vmalloc::init();
 
     #[cfg(feature = "kmemleak")]
     kmemleak::MEM_LEAK_CATCHER.init();

src/aero_kernel/src/mem/mod.rs

@@ -20,6 +20,7 @@
 pub mod alloc;
 pub mod paging;
 pub mod pti;
+mod vmalloc;
 
 use core::alloc::Layout;
 

src/aero_kernel/src/mem/paging/page.rs

@@ -118,6 +118,12 @@ impl<S: PageSize> Page<S> {
     pub fn range_inclusive(start: Self, end: Self) -> PageRangeInclusive<S> {
         PageRangeInclusive { start, end }
     }
+
+    /// Returns a range of pages, exclusive `end`.
+    #[inline]
+    pub const fn range(start: Self, end: Self) -> PageRange<S> {
+        PageRange { start, end }
+    }
 }
 
 impl<S: NotGiantPageSize> Page<S> {

src/aero_kernel/src/mem/vmalloc.rs

@@ -0,0 +1,170 @@
+/*
+ * Copyright (C) 2021-2022 The Aero Project Developers.
+ *
+ * This file is part of The Aero Project.
+ *
+ * Aero is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * Aero is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with Aero. If not, see <https://www.gnu.org/licenses/>.
+ */
+
+//! Due to internal-fragmentation in the buddy frame allocator, we cannot allocate large
+//! amount of contiguous physical memory. We instead use [`vmalloc`] to allocate virtually
+//! contiguous memory.
+//!
+//! An area is reserved for [`vmalloc`] in the kernel address space, starting at [`VMALLOC_VIRT_START`] and
+//! ending at [`VMALLOC_VIRT_END`].
+
+use alloc::collections::LinkedList;
+use spin::Once;
+
+use crate::utils::sync::{Mutex, MutexGuard};
+
+use super::{paging::*, AddressSpace};
+
+pub(super) const VMALLOC_MAX_SIZE: usize = 128 * 1024 * 1024; // 128 GiB
+pub(super) const VMALLOC_START: VirtAddr = VirtAddr::new(0xfffff80000000000);
+pub(super) const VMALLOC_END: VirtAddr =
+    VirtAddr::new(0xfffff80000000000 + VMALLOC_MAX_SIZE as u64);
+
+static VMALLOC: Once<Mutex<Vmalloc>> = Once::new();
+
+struct VmallocArea {
+    addr: VirtAddr,
+    size: usize,
+}
+
+impl VmallocArea {
+    fn new(addr: VirtAddr, size: usize) -> Self {
+        Self { addr, size }
+    }
+}
+
+pub(super) struct Vmalloc {
+    free_list: LinkedList<VmallocArea>,
+}
+
+impl Vmalloc {
+    fn new() -> Self {
+        let mut this = Self {
+            free_list: LinkedList::new(),
+        };
+
+        this.free_list
+            .push_front(VmallocArea::new(VMALLOC_START, VMALLOC_MAX_SIZE));
+        this
+    }
+
+    pub(super) fn alloc(&mut self, mut npages: usize) -> Option<VirtAddr> {
+        npages += 1; // allocate a guard page
+        log::debug!("vmalloc: (npages={npages})");
+
+        let size_bytes = npages * Size4KiB::SIZE as usize;
+
+        let area = self
+            .free_list
+            .iter_mut()
+            .find(|area| area.size >= size_bytes)?;
+
+        let address = area.addr.clone();
+
+        if area.size > size_bytes {
+            area.addr = area.addr + size_bytes;
+            area.size -= size_bytes;
+        } else {
+            // the size of the area is exactly the size we need, so remove it from
+            // the free list.
+            log::warn!("todo: implement this")
+        }
+
+        log::debug!("{:?}", address);
+
+        let mut address_space = AddressSpace::this();
+        let mut offset_table = address_space.offset_page_table();
+
+        let page_range = {
+            let start_page: Page = Page::containing_address(address);
+            let end_page = Page::containing_address(address + size_bytes);
+
+            Page::range(start_page, end_page)
+        };
+
+        // map the pages at the allocated address.
+        for page in page_range {
+            let frame: PhysFrame<Size4KiB> = unsafe {
+                FRAME_ALLOCATOR
+                    .allocate_frame()
+                    .expect("vmalloc: physical memory exhausted")
+            };
+
+            unsafe {
+                offset_table.map_to(
+                    page,
+                    frame,
+                    PageTableFlags::PRESENT | PageTableFlags::WRITABLE,
+                )
+            }
+            .unwrap()
+            .flush();
+        }
+
+        Some(address)
+    }
+
+    pub(super) fn dealloc(&mut self, addr: VirtAddr, mut npages: usize) {
+        npages += 1; // deallocate the vmalloc guard page
+
+        log::debug!("vmdealloc: (npages={npages})");
+
+        let size = npages * Size4KiB::SIZE as usize;
+
+        // check if this block can be merged into another block.
+        let merge = self
+            .free_list
+            .iter_mut()
+            .find(|area| addr + size == area.addr);
+
+        if let Some(merge) = merge {
+            merge.addr = addr;
+            merge.size += size;
+        } else {
+            // the block cannot be merged, so add it to the free list.
+            self.free_list.push_back(VmallocArea::new(addr, size));
+        }
+
+        let mut address_space = AddressSpace::this();
+        let mut offset_table = address_space.offset_page_table();
+
+        let page_range = {
+            let start_page: Page = Page::containing_address(addr);
+            let end_page = Page::containing_address(addr + size);
+
+            Page::range(start_page, end_page)
+        };
+
+        for page in page_range {
+            // unmap the page at the address which in turn will deallocate
+            // the frame (refcnt == 0).
+            offset_table.unmap(page).unwrap().1.flush();
+        }
+    }
+}
+
+pub fn init() {
+    VMALLOC.call_once(|| Mutex::new(Vmalloc::new()));
+}
+
+/// ## Panics
+/// * If the `vmalloc` allocator is not initialized.
+pub(super) fn get_vmalloc() -> MutexGuard<'static, Vmalloc> {
+    VMALLOC.get().expect("get_vmalloc: not initialized").lock()
+}