第五节 实现字节内存分配器
上节我们完成了从早期的页分配器到正式页分配器的切换，下面来处理字节分配器，我们将基于经典的伙伴算法（buddy）实现这个内存分配器。关于伙伴算法，简单来说，一个内存块可以分成对等大小、地址连续的两个内存块，它们称为伙伴。当进行内存分配时，如果没有正好符合要求的空闲内存块，则需要对更大的空闲内存块逐级平分，直到划分出符合要求的最小内存块；内存释放时，尝试与它的伙伴进行合并，直至不能合并为止。这个算法兼顾了分配速度和碎片化问题。它的实现原理如下图所示：
从图中可见，buddy 内存分配器的实现结构简单，就是通过数组 + 链表来维护空闲的内存块。
数组索引称 Order，每个 Order 维护一组具有相同大小的空闲内存块，它们通过链表结构连接在一起。Order 与本级所维护的空闲块大小的对应关系：空闲块大小 = 2^order^ 字节。在分配时，每一级 Order 上的内存块都可以平分为一对伙伴内存块，挂到更低一级的 Order 链表上；反之在释放内存块时，查看它的伙伴内存块是否也是空闲，如果是则合并成上一级的大块，挂到上一级 Order 的链表中。
这里提到的链表是第三章第六节引入的侵入式链表 linked_list，其特点是：链表的节点直接嵌入到空闲内存块内部。
先来看 buddy 内存分配器的主体结构 Heap：
// buddy_system_allocator/src/lib.rs
#![no_std]
use core::cmp::{min, max};
use core::alloc::Layout;
use core::mem::size_of;
use core::ptr::NonNull;
pub struct Heap<const ORDER: usize> {
free_list: [linked_list::LinkedList; ORDER],
used: usize,
allocated: usize,
total: usize,
}
impl<const ORDER: usize> Heap<ORDER> {
pub const fn new() -> Self {
Heap {
free_list: [linked_list::LinkedList::new(); ORDER],
used: 0,
allocated: 0,
total: 0,
}
}
pub const fn empty() -> Self {
Self::new()
}
}
// buddy_allocator/Cargo.toml
[dependencies]
linked_list = { path = "../linked_list" }第 10 行：free_list 的类型是长度为 Order 的数组，数组元素的类型是侵入式链表 LinkedList。对应上图的设计结构。
// buddy_system_allocator/src/lib.rs
impl<const ORDER: usize> Heap<ORDER> {
pub unsafe fn add_to_heap(&mut self, mut start: usize, mut end: usize) {
start = (start + size_of::<usize>() - 1) & (!size_of::<usize>() + 1);
end &= !size_of::<usize>() + 1;
assert!(start <= end);
let mut total = 0;
let mut current_start = start;
while current_start + size_of::<usize>() <= end {
let lowbit = current_start & (!current_start + 1);
let size = min(lowbit, prev_power_of_two(end - current_start));
total += size;
self.free_list[size.trailing_zeros() as usize].push(current_start as *mut usize);
current_start += size;
}
self.total += total;
}
pub unsafe fn init(&mut self, start: usize, size: usize) {
self.add_to_heap(start, start + size);
}
}
pub(crate) fn prev_power_of_two(num: usize) -> usize {
1 << (usize::BITS as usize - num.leading_zeros() as usize - 1)
}第 3~18 行，初始化内存分配器的核心方法 add_to_heap，根据新增的内存段的大小找到合适的 Order，插入对应链表。
// buddy_system_allocator/src/lib.rs
impl<const ORDER: usize> Heap<ORDER> {
pub fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, ()> {
let size = max(
layout.size().next_power_of_two(),
max(layout.align(), size_of::<usize>()),
);
let class = size.trailing_zeros() as usize;
for i in class..self.free_list.len() {
if !self.free_list[i].is_empty() {
for j in (class + 1..i + 1).rev() {
if let Some(block) = self.free_list[j].pop() {
unsafe {
self.free_list[j - 1]
.push((block as usize + (1 << (j - 1))) as *mut usize);
self.free_list[j - 1].push(block);
}
} else {
return Err(());
}
}
let result = NonNull::new(
self.free_list[class]
.pop()
.expect("current block should have free space now")
as *mut u8,
);
if let Some(result) = result {
self.used += layout.size();
self.allocated += size;
return Ok(result);
} else {
return Err(());
}
}
}
Err(())
}
pub fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
let size = max(
layout.size().next_power_of_two(),
max(layout.align(), size_of::<usize>()),
);
let class = size.trailing_zeros() as usize;
unsafe {
self.free_list[class].push(ptr.as_ptr() as *mut usize);
let mut current_ptr = ptr.as_ptr() as usize;
let mut current_class = class;
while current_class < self.free_list.len() {
let buddy = current_ptr ^ (1 << current_class);
let mut flag = false;
for block in self.free_list[current_class].iter_mut() {
if block.value() as usize == buddy {
block.pop();
flag = true;
break;
}
}
if flag {
self.free_list[current_class].pop();
current_ptr = min(current_ptr, buddy);
current_class += 1;
self.free_list[current_class].push(current_ptr as *mut usize);
} else {
break;
}
}
}
self.used -= layout.size();
self.allocated -= size;
}
}第 9~37 行：当分配内存块时，逐级向上找到能满足分配的最小块，取出后不断平分下去，直至大小刚好大于申请目标的大小。此外，每次平分剩余的后半部分都挂到对应大小的链表上。
第 51~69 行：释放内存块时正好相反，先尝试与它的伙伴块进行合并，逐级向上尝试，直到不能合并为止，最后挂到对应 Order 链表上。
下面引入测试用例验证一下 buddy 内存分配器的功能：
// buddy_allocator/tests/test_heap.rs
use core::alloc::Layout;
use core::mem::size_of;
use buddy_allocator::Heap;
#[test]
fn test_empty_heap() {
let mut heap = Heap::<32>::new();
assert!(heap.alloc(Layout::from_size_align(1, 1).unwrap()).is_err());
}
#[test]
fn test_heap_add() {
let mut heap = Heap::<32>::new();
assert!(heap.alloc(Layout::from_size_align(1, 1).unwrap()).is_err());
let space: [usize; 100] = [0; 100];
unsafe {
heap.add_to_heap(space.as_ptr() as usize, space.as_ptr().add(100) as usize);
}
let addr = heap.alloc(Layout::from_size_align(1, 1).unwrap());
assert!(addr.is_ok());
}
#[test]
fn test_heap_oom() {
let mut heap = Heap::<32>::new();
let space: [usize; 100] = [0; 100];
unsafe {
heap.add_to_heap(space.as_ptr() as usize, space.as_ptr().add(100) as usize);
}
assert!(heap
.alloc(Layout::from_size_align(100 * size_of::<usize>(), 1).unwrap())
.is_err());
assert!(heap.alloc(Layout::from_size_align(1, 1).unwrap()).is_ok());
}
#[test]
fn test_heap_alloc_and_free() {
let mut heap = Heap::<32>::new();
assert!(heap.alloc(Layout::from_size_align(1, 1).unwrap()).is_err());
let space: [usize; 100] = [0; 100];
unsafe {
heap.add_to_heap(space.as_ptr() as usize, space.as_ptr().add(100) as usize);
}
for _ in 0..100 {
let addr = heap.alloc(Layout::from_size_align(1, 1).unwrap()).unwrap();
heap.dealloc(addr, Layout::from_size_align(1, 1).unwrap());
}
}执行 make test，测试通过！