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
This commit is contained in:
@@ -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<AllocatorInner>,
|
||||
/// 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<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 {
|
||||
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<KernelPages>;
|
||||
fn deref(&self) -> &Heap<KernelPages> {
|
||||
// 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<KernelPages> {
|
||||
// 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();
|
||||
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;
|
||||
self.lock().alloc(layout.size(), layout.align())
|
||||
}
|
||||
|
||||
// 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) {
|
||||
let mut inner = self.lock();
|
||||
let size = layout.size().max(core::mem::size_of::<FreeBlock>());
|
||||
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<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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user