From 419b715a1facf3a210fadd74c39ee00de53be7f0 Mon Sep 17 00:00:00 2001 From: Gabriel Kaszewski Date: Wed, 28 Oct 2026 13:01:52 +0100 Subject: [PATCH] fix(libkos): heap leaked padding/remainders + never coalesced (UI churn OOM'd tally); address-ordered free list, merge on free, split align padding; Heap host-tested --- crates/libkos/src/alloc_impl.rs | 350 ++++++++++++++++++++++---------- 1 file changed, 241 insertions(+), 109 deletions(-) diff --git a/crates/libkos/src/alloc_impl.rs b/crates/libkos/src/alloc_impl.rs index 6346c17..dd1c519 100644 --- a/crates/libkos/src/alloc_impl.rs +++ b/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 { - /// A program's threads share the heap: one at a time in the free list. - locked: AtomicBool, - inner: UnsafeCell, +/// Where the heap's memory comes from: `n` fresh pages, page aligned, or +/// null. +pub trait Pages { + fn get(&mut self, n: usize) -> *mut u8; } -/// How much heap a process gets is its page budget's business (the kernel -/// enforces it), not the allocator's. -struct AllocatorInner { - free_list: *mut FreeBlock, +/// Pages from the kernel; a process's page budget bounds how many. +pub struct KernelPages; + +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 { + /// Free blocks, lowest address first; no two adjacent. + free: *mut FreeBlock, + pages: P, +} + +impl Heap

{ + 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>, } unsafe impl Sync for KosAllocator {} impl KosAllocator { pub const fn new() -> Self { - Self { - locked: AtomicBool::new(false), - inner: UnsafeCell::new(AllocatorInner { free_list: ptr::null_mut() }), - } + Self { locked: AtomicBool::new(false), heap: UnsafeCell::new(Heap::new(KernelPages)) } } 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; - fn deref(&self) -> &AllocatorInner { - // Only one `Locked` exists at a time. - unsafe { &*self.0.inner.get() } + type Target = Heap; + fn deref(&self) -> &Heap { + // SAFETY: only one `Locked` exists at a time. + unsafe { &*self.0.heap.get() } } } impl DerefMut for Locked<'_> { - fn deref_mut(&mut self) -> &mut AllocatorInner { - unsafe { &mut *self.0.inner.get() } + fn deref_mut(&mut self) -> &mut Heap { + // 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::() { - 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(); - let size = layout.size().max(core::mem::size_of::()); - 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) + self.lock().alloc(layout.size(), layout.align()) } - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - let mut inner = self.lock(); - let size = layout.size().max(core::mem::size_of::()); - inner.add_to_free_list(ptr, size); + unsafe fn dealloc(&self, at: *mut u8, layout: Layout) { + self.lock().dealloc(at, layout.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 { + 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); } }