故事是這樣子的,雖然在上上篇的 Context Switch 我們曾經預告過因為 static 的問題, 可能要停更一段時間,但後來沒停更,static 問題很快就用 spin 提供的 no_std Mutex 解決了。
不過,no_std Mutex 雖然解決 static 不用手爆 constructor 的問題,同時卻產生了另一個問題:stack overflow。

因為呼叫 static variable constructor 會用到 os process 的 stack,static variable size 一大,os process 的 stack 就會被吃乾抹淨。
所以問題來了,我們必須要儘可能的縮小 static variable 的尺寸,像上篇把 process 的 context 全塞在裡面是不可能的 (為了讓 code 動起來我把 process 數量降到剩 2 個),因此下一步就很清楚了,必須挑戰作業系統裡一大難題:記憶體管理

這個問題說難不難,我最早看到用 rust 寫 OS 的 blog_os, 作者自己就有公佈一款 linked list allocator, 理論上我可以直接在 Cargo 上引用這個 repository,然後插到我的 project 即可。
當然,除了直接用人家的 project ,我們還是要看一下 Rust 的記憶體管理是怎麼實作的,這篇我們就實作一個只能 alloc 不能 dealloc 的記憶體管理 (其實這樣說起來,我在 Mutex 都自己實作才對,不該直接用 spin 提供的)。

Global Allocator

rust 在設計上,記憶體管理是透過 Global Allocator 來處理,一般有 std 下會使用 std::alloc 的 System allocator,像我們這樣沒有 Global Allocator 的程式,是無法使用 heap 的,現下的 OS 在呼叫 Box 的時候:

extern crate alloc;
use alloc::boxed::Box;

let b = Box::new(64);

會出現如下錯誤:

error: no global memory allocator found but one is required;  
link to std or add `#[global_allocator]` to a static item that implements the GlobalAlloc trait.

錯誤訊息說得很明白了,要提供一個實作 GlobalAlloc trait 的 struct,為它加上 #[global_allocator] 這個 attribute。

extern crate alloc;
use alloc::alloc::{GlobalAlloc, Layout};

pub struct HeapAllocator;

GlobalAlloc trait 裡面定義了四個函式,必要的函式為前兩個 alloc 與 dealloc,rust 會自動幫我們推出 alloc_zeroed 跟 realloc:

pub unsafe trait GlobalAlloc {
  unsafe fn alloc(&self, layout: Layout) -> *mut u8;
  unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);

  unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { ... }
  unsafe fn realloc(
        &self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { ... }
}

為我們的空殼實作 GlobalAlloc,trait 本身就是 unsafe 還真鮮; Layout 是 rust 定義用來表示一塊記憶體的資料結構,儲存了這塊記憶體的 size, align 等資訊,有趣的是,rust dealloc 的時候也會帶 layout 進來, 這意味著跟 C 的 free 不一樣, GlobalAlloc 在背後還是有某種管理機制在 realloc 的時候可以知道記憶體的尺寸資訊:

use core::ptr::null_mut;

unsafe impl {
  unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
    null_mut()
  }   

  unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
    panic!("dealloc should be never called");
  }   
}

除了這個還要加上 memory alloc error 的實作,或者在 main.rs (lib.rs) 裡面加上 #![feature(default_alloc_error_handler)] 讓編譯器自行弄成 panic!

#[alloc_error_handler]
fn alloc_error_handler(layout: Layout) -> ! {
  panic!("allocation error {:?}", layout);
}

最後是在 main.rs 將我們的空殼 HeapAllocator 設定為 #[global_allocator]

#[global_allocator]
static ALLOCATOR: HeapAllocator = HeapAllocator;

如此一來上面使用 Box 的程式就能通過編譯了;在執行時則會 panic,因為 Box 會對 alloc 傳回的 null pointer 存取。

實作 Linked List

一般最簡便的 heap 管理就是用 linked_list,這次也好好跟 linked_list_allocator 的作者學會怎麼打造 rust 的 linked_list。
第一步先打造我們的 Hole,用 Option 來打造 list 應該是基本中的基本,Option 內是另一個 Hole 的 static reference。

pub struct Hole {
  size: usize,
  next: Option<&'static mut Hole>,
}

HoleList 的 first 是 dummy Hole,size 為 0;在生成的時候會如下,在 hole_addr 上建一個 Hole pointer,對 pointer 寫入 Hole struct 即可。
ptr 寫入這個動作是 unsafe 的,不過我們整個 new 函式都是 unsafe 所以沒問題。

pub struct HoleList {
  first: Hole,
}

impl HoleList {
  pub unsafe fn new(hole_addr: usize, hole_size: usize) -> Self {
    let ptr = hole_addr as *mut Hole;
    ptr.write( Hole {
      size: hole_size,
      next: None,
    });
    Self {
      first: Hole {
        size: 0,
        next: Some(&mut *ptr),
      }
    }
  }
}

因為 Hole 本身有 reference 不適合複製,另外實作一個 HoleInfo 來代表一塊記憶體:

struct HoleInfo {
  addr: usize,
  size: usize,
}

impl Hole {
  pub fn info(&self) -> HoleInfo {
    HoleInfo {
      addr: self as *const _ as usize,
      size: self.size,
    }
  }
}

HoleList 之外我們再包一層 Heap 用來記錄整個記憶體使用的狀況:

struct Heap {
  bottom: usize,
  size: usize,
  used: usize,
  holes: HoleList,
}

impl Heap {
  pub const fn empty() -> Self {
    Self {
      bottom: 0,
      size: 0,
      used: 0,
      holes: HoleList::empty(),
    }
  }
}

原本空無一物的 HeapAllocator 也要改造一番,為了避免 malloc 的競態條件,heap 外要加上 Mutex。

pub struct HeapAllocator {
  heap: Mutex<Heap>,
}

impl HeapAllocator {
  pub const fn empty() -> Self {
    Self {
      heap: Mutex::new(Heap::empty()),
    }
  }
}

整體的架構依序為:HeapAllocator -> Heap -> HoleList -> [Hole]

實作 alloc

global allocate 函式,實作很單純,因為我們已經知道不會有記憶體被 free,要做的就只有拿出下一個元素(單一區塊管理所有記憶體), 檢查 size 夠不夠,如果夠的話就在 Hole 的尾端把空出記憶體;不夠的話就是回傳 Err
實際的 linked_list 則是在 Hole 沒有足夠記憶體的時候,會去拿出下一個 Hole 試著 alloc 看看,直到 list 尾端就回傳 ErrAllocation 用來表示一塊分配的記憶體,未來處理 align 的時候會在裡面加上更多東西。

pub fn allocate(previous: &mut Hole, layout: Layout) -> Result<Allocation, ()> {
  let mut current = previous.next.as_mut().unwrap().info();
  let required_size = layout.size();
  if current.size < required_size {
    Err(())
  } else {
    let allocated_info = HoleInfo {
      addr: current.addr + current.size - required_size,
      size: required_size,
    };
    current.size -= required_size;
    Ok(Allocation {
      info: allocated_info,
    })
  }
}

如此一來,HoleList 就是把自己的 List 拿出來,呼叫 global allocate 函式得到 Allocation,轉型為 NonNull pointer , 這裡可以用 NonNull 因為 Null 的狀況已經由 Result Err 來表示了:

// HoleList
pub fn allocate(&mut self, layout: Layout) -> Result<NonNull<u8>, ()> {
  allocate(&mut self.first, layout).map(|allocation| {
    NonNull::new(allocation.info.addr as *mut u8).unwrap()
  })
}

在實作 GlobalAlloc 的 alloc 函式裡面呼叫 allocate,如果是 Err 的狀況會由 map_or 轉型為 Null pointer。

unsafe impl GlobalAlloc for HeapAllocator {
  unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    self.heap
      .lock()
      .allocate(layout)
      .map_or(null_mut(), |allocation| allocation.as_ptr())
  }
}

初始化 Heap

我們 heap 的尺寸和 xv6 一樣設定為 128MB(扣掉 kernel code 的空間,因此實際空間會略小於 128MB),heap 的起始位置則由 linker script 決定,在 linker script .data, .bss 等區間之後,加上:

PROVIDE(_END = .);

初始化記憶體就能透過 extern "C" 取得 _END 的值,再把這個 heap 設定寫入 ALLOCATOR 的 heap 中。

pub fn init_heap() {
  extern "C" {
    static _END: usize;
  }

  let heap_start: usize = unsafe {
    &_END as *const usize as usize
  };
  let heap_end = memorylayout::PHYSTOP as usize;
  let heap_size = heap_end - heap_start;
  unsafe {
    ALLOCATOR
      .heap
      .lock()
      .init(heap_start, heap_size)
  }
}

結語

我們實作了一款只會 allocate 不會 free 的 allocator,但至少,現在我們可以在 code 裡面使用各種使用 heap 的資料結構,例如上述的 Box。
簡而言之我們的 HeapAllocator 大概可以用下面這段概述:

我之前是幫人弄記憶體管理的,而我 malloc 的原則是:「○林涼 malloc 爆」
沒錯,就是○林涼 malloc 爆,老子才不管甚麼要不要 free 的,每次 malloc 就是一大塊。1K malloc 1M,1M malloc 1G。
後來作業系統跳 panic 跟我說記憶體不足,你有頭緒嗎?
我TMD怎麼會知道。

第二,如果你跟 xv6 的實作比較的話,可以明顯看到 rust 優於 C 的一點,rust 的 locking 是強制性的,只要我們把東西包到 Mutex 裡面, 你一定要用 lock 再去存取內部的資料;相對的 xv6 實作則是:

struct {
  struct spinlock lock;
  struct run *freelist;
} kmem;

存取 freelist 之前記得先跟 lock 去 acquire lock。
啊沒有 lock 會怎麼樣?系統爆掉啊。
誰的錯?當然是程式設計師的。

Rust 在這方面 比較 貼心一點,雖然 Mutex 的實作相對麻煩一些,但只要做完,實作的限制就會如通往建功嶼的海中道路般浮現; 你的實作會被限制到只剩安全的方式,不這麼做 compiler 就會跳出來把你攔住,叫你回去修練 21 小時之後再來, 或是非得讓你狂下 unsafe 或 unwrap 弄出有危險的程式它才放行。

套句 jserv 教授的話:

C 語言假定了程式設計師都是聰明人,很清楚知道自己在做什麼

而愚者用了 C 就必須要為自己的愚蠢付出代價。