fix(libkos): heap leaked padding/remainders + never coalesced (UI churn OOM'd tally); address-ordered free list, merge on free, split align padding; Heap<Pages> host-tested

crates/libkos/src/alloc_impl.rs

@@ -1,3 +1,10 @@
+//! The heap: a free list of blocks kept in address order and merged with
+//! its neighbours whenever a block is freed, so memory freed comes back in
+//! one piece and the heap grows only as far as what is live at once needs.
+//! Every block is a multiple of `UNIT` bytes and starts on a `UNIT`
+//! boundary; a stricter alignment is met by splitting the front of a block
+//! off as a free block of its own, so nothing is ever lost.
+
 use core::alloc::{GlobalAlloc, Layout};
 use core::cell::UnsafeCell;
 use core::ops::{Deref, DerefMut};
@@ -5,32 +12,163 @@ use core::ptr;
 use core::sync::atomic::{AtomicBool, Ordering};
 
 const PAGE_SIZE: usize = 4096;
+/// The granule: big enough to hold a free block's header.
+const UNIT: usize = 16;
 
 struct FreeBlock {
     size: usize,
     next: *mut FreeBlock,
 }
 
-pub struct KosAllocator {
+/// Where the heap's memory comes from: `n` fresh pages, page aligned, or
-    /// A program's threads share the heap: one at a time in the free list.
+/// null.
-    locked: AtomicBool,
+pub trait Pages {
-    inner: UnsafeCell<AllocatorInner>,
+    fn get(&mut self, n: usize) -> *mut u8;
 }
 
-/// How much heap a process gets is its page budget's business (the kernel
+/// Pages from the kernel; a process's page budget bounds how many.
-/// enforces it), not the allocator's.
+pub struct KernelPages;
-struct AllocatorInner {
+
-    free_list: *mut FreeBlock,
+impl Pages for KernelPages {
+    fn get(&mut self, n: usize) -> *mut u8 {
+        let Some(count) = u16::try_from(n).ok().and_then(types::bounded::PageCount::new) else {
+            return ptr::null_mut();
+        };
+        crate::syscall::request_pages(count).map_or(ptr::null_mut(), |va| va.raw() as *mut u8)
+    }
+}
+
+fn round(n: usize) -> usize {
+    n.max(UNIT).div_ceil(UNIT) * UNIT
+}
+
+pub struct Heap<P: Pages> {
+    /// Free blocks, lowest address first; no two adjacent.
+    free: *mut FreeBlock,
+    pages: P,
+}
+
+impl<P: Pages> Heap<P> {
+    pub const fn new(pages: P) -> Self {
+        Self { free: ptr::null_mut(), pages }
+    }
+
+    /// `size` bytes aligned to `align`, or null.
+    pub fn alloc(&mut self, size: usize, align: usize) -> *mut u8 {
+        let size = round(size);
+        let align = align.max(UNIT);
+        let found = self.take(size, align);
+        if !found.is_null() {
+            return found;
+        }
+        // Room for the worst alignment padding too.
+        let n = (size + align - UNIT).div_ceil(PAGE_SIZE);
+        let fresh = self.pages.get(n);
+        if fresh.is_null() {
+            return ptr::null_mut();
+        }
+        self.give(fresh, n * PAGE_SIZE);
+        self.take(size, align)
+    }
+
+    /// Returns a block `alloc` handed out for `size`.
+    pub fn dealloc(&mut self, at: *mut u8, size: usize) {
+        self.give(at, round(size));
+    }
+
+    /// First fit: the lowest block that holds `size` at `align`.
+    fn take(&mut self, size: usize, align: usize) -> *mut u8 {
+        let mut prev: *mut FreeBlock = ptr::null_mut();
+        let mut cur = self.free;
+        while !cur.is_null() {
+            // SAFETY: every block on the list is free memory this heap owns,
+            // at least `UNIT` bytes, holding a header.
+            let (start, block_size, next) = unsafe { (cur as usize, (*cur).size, (*cur).next) };
+            let aligned = start.next_multiple_of(align);
+            let front = aligned - start;
+            if front + size <= block_size {
+                let back = block_size - front - size;
+                // Unlink, then hand the unused front and back pieces back.
+                // SAFETY: `prev` is on the list, before `cur`.
+                unsafe {
+                    match prev.is_null() {
+                        true => self.free = next,
+                        false => (*prev).next = next,
+                    }
+                }
+                if front > 0 {
+                    self.give(start as *mut u8, front);
+                }
+                if back > 0 {
+                    self.give((aligned + size) as *mut u8, back);
+                }
+                return aligned as *mut u8;
+            }
+            prev = cur;
+            cur = next;
+        }
+        ptr::null_mut()
+    }
+
+    /// Puts `size` bytes at `at` on the list, merged with any neighbour.
+    fn give(&mut self, at: *mut u8, size: usize) {
+        let start = at as usize;
+        let mut prev: *mut FreeBlock = ptr::null_mut();
+        let mut next = self.free;
+        while !next.is_null() && (next as usize) < start {
+            prev = next;
+            // SAFETY: on the list.
+            next = unsafe { (*next).next };
+        }
+        let block = at as *mut FreeBlock;
+        // SAFETY: `at..at+size` is free memory this heap owns (just freed
+        // or just received), `UNIT` aligned and at least `UNIT` long; `prev`
+        // and `next` are on the list, around it.
+        unsafe {
+            (*block).size = size;
+            (*block).next = next;
+            if !next.is_null() && start + size == next as usize {
+                (*block).size += (*next).size;
+                (*block).next = (*next).next;
+            }
+            match prev.is_null() {
+                true => self.free = block,
+                false if prev as usize + (*prev).size == start => {
+                    (*prev).size += (*block).size;
+                    (*prev).next = (*block).next;
+                }
+                false => (*prev).next = block,
+            }
+        }
+    }
+
+    /// Free bytes, and how many pieces they are in.
+    pub fn free_bytes(&self) -> (usize, usize) {
+        let (mut bytes, mut pieces) = (0, 0);
+        let mut cur = self.free;
+        while !cur.is_null() {
+            // SAFETY: on the list.
+            unsafe {
+                bytes += (*cur).size;
+                cur = (*cur).next;
+            }
+            pieces += 1;
+        }
+        (bytes, pieces)
+    }
+}
+
+pub struct KosAllocator {
+    /// A program's threads share the heap: one at a time in it.
+    locked: AtomicBool,
+    heap: UnsafeCell<Heap<KernelPages>>,
 }
 
 unsafe impl Sync for KosAllocator {}
 
 impl KosAllocator {
     pub const fn new() -> Self {
-        Self {
+        Self { locked: AtomicBool::new(false), heap: UnsafeCell::new(Heap::new(KernelPages)) }
-            locked: AtomicBool::new(false),
-            inner: UnsafeCell::new(AllocatorInner { free_list: ptr::null_mut() }),
-        }
     }
 
     fn lock(&self) -> Locked<'_> {
@@ -47,20 +185,21 @@ impl Default for KosAllocator {
     }
 }
 
-/// The free list, held until dropped.
+/// The heap, held until dropped.
 struct Locked<'a>(&'a KosAllocator);
 
 impl Deref for Locked<'_> {
-    type Target = AllocatorInner;
+    type Target = Heap<KernelPages>;
-    fn deref(&self) -> &AllocatorInner {
+    fn deref(&self) -> &Heap<KernelPages> {
-        // Only one `Locked` exists at a time.
+        // SAFETY: only one `Locked` exists at a time.
-        unsafe { &*self.0.inner.get() }
+        unsafe { &*self.0.heap.get() }
     }
 }
 
 impl DerefMut for Locked<'_> {
-    fn deref_mut(&mut self) -> &mut AllocatorInner {
+    fn deref_mut(&mut self) -> &mut Heap<KernelPages> {
-        unsafe { &mut *self.0.inner.get() }
+        // SAFETY: as above.
+        unsafe { &mut *self.0.heap.get() }
     }
 }
 
@@ -70,99 +209,92 @@ impl Drop for Locked<'_> {
     }
 }
 
-impl AllocatorInner {
-    /// `n` pages in one piece, so an allocation larger than a page fits.
-    fn request_pages(&mut self, n: usize) -> *mut u8 {
-        let Some(count) = u16::try_from(n).ok().and_then(types::bounded::PageCount::new) else {
-            return ptr::null_mut();
-        };
-        let Some(va) = crate::syscall::request_pages(count) else {
-            return ptr::null_mut();
-        };
-        va.raw() as *mut u8
-    }
-
-    fn alloc_from_free_list(&mut self, size: usize, align: usize) -> *mut u8 {
-        let mut prev: *mut FreeBlock = ptr::null_mut();
-        let mut current = self.free_list;
-
-        while !current.is_null() {
-            let block = unsafe { &mut *current };
-            let addr = current as usize;
-            let aligned = (addr + align - 1) & !(align - 1);
-            let padding = aligned - addr;
-
-            if block.size >= size + padding {
-                let remaining = block.size - size - padding;
-
-                if remaining >= core::mem::size_of::<FreeBlock>() {
-                    let new_block = (aligned + size) as *mut FreeBlock;
-                    unsafe {
-                        (*new_block).size = remaining;
-                        (*new_block).next = block.next;
-                    }
-                    if prev.is_null() {
-                        self.free_list = new_block;
-                    } else {
-                        unsafe {
-                            (*prev).next = new_block;
-                        }
-                    }
-                } else {
-                    if prev.is_null() {
-                        self.free_list = block.next;
-                    } else {
-                        unsafe {
-                            (*prev).next = block.next;
-                        }
-                    }
-                }
-
-                return aligned as *mut u8;
-            }
-
-            prev = current;
-            current = block.next;
-        }
-
-        ptr::null_mut()
-    }
-
-    fn add_to_free_list(&mut self, ptr: *mut u8, size: usize) {
-        let block = ptr as *mut FreeBlock;
-        unsafe {
-            (*block).size = size;
-            (*block).next = self.free_list;
-        }
-        self.free_list = block;
-    }
-}
-
 unsafe impl GlobalAlloc for KosAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let mut inner = self.lock();
+        self.lock().alloc(layout.size(), layout.align())
-        let size = layout.size().max(core::mem::size_of::<FreeBlock>());
-        let align = layout.align();
-
-        let ptr = inner.alloc_from_free_list(size, align);
-        if !ptr.is_null() {
-            return ptr;
-        }
-
-        // Alignment can cost up to `align` bytes of the block.
-        let pages_needed = (size + align).div_ceil(PAGE_SIZE);
-        let block = inner.request_pages(pages_needed);
-        if block.is_null() {
-            return ptr::null_mut();
-        }
-        inner.add_to_free_list(block, pages_needed * PAGE_SIZE);
-
-        inner.alloc_from_free_list(size, align)
     }
 
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+    unsafe fn dealloc(&self, at: *mut u8, layout: Layout) {
-        let mut inner = self.lock();
+        self.lock().dealloc(at, layout.size());
-        let size = layout.size().max(core::mem::size_of::<FreeBlock>());
+    }
-        inner.add_to_free_list(ptr, size);
+}
+
+#[cfg(test)]
+mod tests {
+    extern crate std;
+
+    use super::*;
+    use std::vec::Vec;
+
+    /// Pages from one aligned buffer, counted.
+    struct Arena {
+        base: *mut u8,
+        pages: usize,
+        used: usize,
+    }
+
+    impl Pages for Arena {
+        fn get(&mut self, n: usize) -> *mut u8 {
+            if self.used + n > self.pages {
+                return ptr::null_mut();
+            }
+            // SAFETY: inside the buffer.
+            let at = unsafe { self.base.add(self.used * PAGE_SIZE) };
+            self.used += n;
+            at
+        }
+    }
+
+    fn arena(pages: usize) -> Heap<Arena> {
+        let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
+        // SAFETY: a fresh allocation, leaked for the test's life.
+        let base = unsafe { std::alloc::alloc(layout) };
+        Heap::new(Arena { base, pages, used: 0 })
+    }
+
+    #[test]
+    fn freed_memory_merges_back_into_one_piece() {
+        let mut h = arena(4);
+        let blocks: Vec<*mut u8> = (0..20).map(|i| h.alloc(40 + i * 8, 8)).collect();
+        assert!(blocks.iter().all(|b| !b.is_null()));
+        // Freed out of order, they still end up as one block.
+        for i in (0..20).step_by(2).chain((1..20).step_by(2)) {
+            h.dealloc(blocks[i], 40 + i * 8);
+        }
+        assert_eq!(h.free_bytes(), (PAGE_SIZE, 1));
+    }
+
+    #[test]
+    fn alignment_is_met_and_its_padding_comes_back() {
+        let mut h = arena(4);
+        let small = h.alloc(16, 16);
+        let big = h.alloc(100, 256);
+        assert_eq!(big as usize % 256, 0);
+        h.dealloc(small, 16);
+        h.dealloc(big, 100);
+        assert_eq!(h.free_bytes(), (PAGE_SIZE, 1), "no padding was lost");
+    }
+
+    /// Many rounds of making and freeing short-lived things of varied sizes
+    /// (as a UI does every frame) must not make the heap grow.
+    #[test]
+    fn churn_does_not_grow_the_heap() {
+        let mut h = arena(64);
+        let mut live: Vec<(*mut u8, usize)> = Vec::new();
+        let mut seed = 12345u32;
+        for round in 0..20_000 {
+            seed = seed.wrapping_mul(1_103_515_245).wrapping_add(12345);
+            let size = 8 + (seed >> 16) as usize % 6000;
+            let at = h.alloc(size, if seed & 1 == 0 { 8 } else { 64 });
+            assert!(!at.is_null(), "out of memory at round {round}");
+            // SAFETY: just allocated, `size` bytes.
+            unsafe { ptr::write_bytes(at, 0xAB, size) };
+            live.push((at, size));
+            if live.len() > 12 {
+                let (at, size) = live.remove((seed as usize >> 8) % live.len());
+                h.dealloc(at, size);
+            }
+        }
+        assert!(h.pages.used <= 24, "used {} pages for ~12 blocks of < 6 KiB", h.pages.used);
     }
 }