第六节 集成字节内存分配器

现在 buddy 内存分配器也已经准备好,下面就是启用它成为正式的字节内存分配器,完全替换掉早期内存分配器的功能,内核从本节开始将获得完整的内存管理功能。

分两步实现 GlobalAllocator 对正式的字节内存分配器的支持。

第一步,封装上节实现的基于 buddy 算法的字节内存分配器:

// axalloc/src/buddy.rs
use core::ptr::NonNull;
use crate::{Layout, AllocResult, AllocError};
use buddy_allocator::Heap;

pub struct BuddyByteAllocator {
    inner: Heap<32>,
}
impl BuddyByteAllocator {
    pub const fn new() -> Self {
        Self {
            inner: Heap::<32>::new(),
        }
    }
}
impl BuddyByteAllocator {
    pub fn init(&mut self, start: usize, size: usize) {
        unsafe { self.inner.init(start, size) };
    }
}

// axalloc/Cargo.toml
[dependencies]
buddy_allocator = { path = "../buddy_allocator" }

第 7 行:简单封装上节实现的内存分配器,最大 Order 设置为 32。

第 17~19 行:初始化字节分配器,指定一段内存范围用于当前的字节分配功能。

实现最基本的基于字节分配/释放的方法:

// axalloc/src/buddy.rs
impl BuddyByteAllocator {
    pub fn alloc_bytes(&mut self, layout: Layout) -> AllocResult<NonNull<u8>> {
        self.inner.alloc(layout).map_err(|_| AllocError::NoMemory)
    }

    pub fn dealloc_bytes(&mut self, pos: NonNull<u8>, layout: Layout) {
        self.inner.dealloc(pos, layout)
    }
}

简单调用内部分配器的对应功能。

第二步,扩展 GlobalAllocator 启用字节分配:

// axalloc/src/lib.rs
mod buddy;
use buddy::BuddyByteAllocator;
const MIN_HEAP_SIZE: usize = 0x8000; // 32 K

struct GlobalAllocator {
    early_alloc: SpinRaw<EarlyAllocator>,
    page_alloc: SpinRaw<BitmapPageAllocator>,
    byte_alloc: SpinRaw<BuddyByteAllocator>,
    finalized: BootOnceCell<bool>,
}

impl GlobalAllocator {
    pub const fn new() -> Self {
        Self {
            early_alloc: SpinRaw::new(EarlyAllocator::uninit_new()),
            page_alloc: SpinRaw::new(BitmapPageAllocator::new()),
            byte_alloc: SpinRaw::new(BuddyByteAllocator::new()),
            finalized: BootOnceCell::new(),
        }
    }
}

第 9 和 18 行:引入和初始化字节分配器。

impl GlobalAllocator {
    fn alloc_bytes(&self, layout: Layout) -> *mut u8 {
        let ret = if self.finalized.is_init() {
            self.byte_alloc.lock().alloc_bytes(layout)
        } else {
            self.early_alloc.lock().alloc_bytes(layout)
        };

        if let Ok(ptr) = ret {
            ptr.as_ptr()
        } else {
            alloc::alloc::handle_alloc_error(layout)
        }
    }
    fn dealloc_bytes(&self, ptr: *mut u8, layout: Layout) {
        if self.finalized.is_init() {
            self.byte_alloc.lock().dealloc_bytes(
                NonNull::new(ptr).expect("dealloc null ptr"),
                layout
            )
        } else {
            self.early_alloc.lock().dealloc_bytes(
                NonNull::new(ptr).expect("dealloc null ptr"),
                layout
            )
        }
    }
}

第 2~14 行:GlobalAllocator 面向字节分配的接口,根据 finalized 决定调用早期还是正式的内存分配器。

第 15~17 行:相应的,GlobalAllocator 面向字节释放的接口,同样基于 finalized 区分处于那个阶段。

下面在 GlobalAllocator 中启用字节分配:

实际上 GlobalAllocator 的两个子分配器 - 字节分配与页分配,它们之间并非并列关系,而是存在一个层次间关联。

字节分配器与页分配器关系
  1. 系统的全部内存由页分配器负责管理,它提供一组页作为字节分配器的初始内存。
  2. 字节分配器的现有内存不足以满足请求时,它会向页分配器要求追加更多的内存页再进行分配,直至页分配器无法满足时才会真正失败。

先来实现第 1 条功能,扩展 final_init(...),之前只是处理了页分配,下面来启用字节分配。

// axalloc/src/lib.rs
impl GlobalAllocator {
    pub fn final_init(&self, start: usize, size: usize) {
        self.page_alloc.lock().init(start, size);
        let layout = Layout::from_size_align(MIN_HEAP_SIZE, PAGE_SIZE).unwrap();
        let heap_ptr = self.alloc_pages(layout) as usize;
        self.byte_alloc.lock().init(heap_ptr, MIN_HEAP_SIZE);
        self.finalized.init(true);
    }
}

第 5~7 行:从页分配器申请了一组页,用于初始化字节分配器。

验证一下当前实现的字节分配功能是否正常,在组件 axorigin 中申请一个长度为 4K 的长字节:

// axorigin/src/main.rs
pub fn main(_hartid: usize, _dtb: usize) {
    ... ...
    try_alloc_pages();
    try_alloc_long_string();
    ... ...
}

fn try_alloc_long_string() {
    use core::alloc::Layout;
    extern crate alloc;

    const LENGTH: usize = 0x1000;
    let layout = Layout::from_size_align(LENGTH, 1).unwrap();
    let p = unsafe { alloc::alloc::alloc(layout) };
    println!("Allocate long string: [{:?}].", p);
    unsafe { alloc::alloc::dealloc(p, layout) };
    println!("Release long string ok!");
}

执行 make run,显示结果:

ArceOS is starting ...

Now: 0.111465 Hello, ArceOS![from String] Allocate pages: [0xffffffc08026c000]. Release pages ok! Allocate long string: [0xffffffc0801f9000]. Release long string ok! Elapsed: 0.008842

测试正常!

继续实现第 2 条 - 当字节分配内存不足时,从页分配器中申请追加更多页面扩展可分配的范围,除非页分配器的内存也耗尽才会导致失败。主要是对 GlobalAllocator 的字节分配函数进行简单的改造:

// axalloc/src/lib.rs
use log::info;

impl GlobalAllocator {
    fn alloc_bytes(&self, layout: Layout) -> *mut u8 {
        if !self.finalized.is_init() {
            return self.early_alloc.lock().alloc_bytes(layout).unwrap().as_ptr();
        }

        loop {
            let mut balloc = self.byte_alloc.lock();
            if let Ok(ptr) = balloc.alloc_bytes(layout) {
                return ptr.as_ptr();
            } else {
                let old_size = balloc.total_bytes();
                let expand_size = old_size
                    .max(layout.size())
                    .next_power_of_two()
                    .max(PAGE_SIZE);
                let layout = Layout::from_size_align(expand_size, PAGE_SIZE).unwrap();
                let heap_ptr = self.alloc_pages(layout) as usize;
                info!(
                    "expand heap memory: [{:#x}, {:#x})",
                    heap_ptr,
                    heap_ptr + expand_size
                );
                let _ = balloc.add_memory(heap_ptr, expand_size);
            }
        }
    }
}

第 12 行:先尝试从现有的内存范围内申请内存,如果内存不足,将会触发第 15~27 行的流程。

第 15~27 行:从页分配器中申请追加更多的页面用作字节分配。从第 15 到 19 行可以看出,每次请求量都是翻倍,但至少要申请一页。

上面实现涉及两个 buddy 分配器没有实现的功能,补上:

// axalloc/src/buddy.rs
impl BuddyByteAllocator {
    pub fn add_memory(&mut self, start: usize, size: usize) -> AllocResult {
        unsafe { self.inner.add_to_heap(start, start + size) };
        Ok(())
    }
    pub fn total_bytes(&self) -> usize {
        self.inner.stats_total_bytes()
    }
}

// buddy_allocator/src/lib.rs
impl<const ORDER: usize> Heap<ORDER> {
    pub fn stats_total_bytes(&self) -> usize {
        self.total
    }
}

// axalloc/Cargo.toml
[dependencies]
log = "0.4"

下面来测试一下,在 axorigin 中尝试申请大于 1M 的内存,完整验证新的字节内存分配功能。

// axorigin/src/main.rs
const LENGTH: usize = 0x200000;

只需要修改一行,把长度从 0x1000 扩大为 0x200000,再次执行 make run,观察结果:

ArceOS is starting ...

Now: 0.111744 Hello, ArceOS![from String] Allocate pages: [0xffffffc08026d000]. Release pages ok! Allocate long string: [0xffffffc080600000]. Release long string ok! Elapsed: 0.011159

申请大于 1M 内存(即超过 early 分配器的最大分配能力),仍然执行成功!