Zenn
🥺

メモリと仲良しになろう！[アロケータ実装編]

2025/03/06に公開
AI
Go
Rust
プログラミング
C
tech

はじめに

本記事は「Writing an OS in Rust」のアロケータ実装編を個人的に咀嚼するために作成しました。
不明瞭な部分がありましたら原文をご参照いただけると幸いです。

誰のための記事か？(for whom)

  • メモリ上でどのようにデータが載るのかが分からない人
  • どのようにアロケーションが行われるか分からない人
  • アロケータの実装に興味がある人

なぜ学ぶのか？ (why)

  • 自作 OS 開発をする上での基礎的な知識として必要であるため
  • メモリアロケーションは高価な処理であるということを理解するため

前回のおさらい (アロケータとは？)

前回書いた記事からそのまま引用します。

アロケータとは「ヒープメモリにある引き出しの管理人さん」のことです。

あるデータをメモリに格納するとき、好き勝手な引き出しを開けて収納するとどうなるでしょうか。誰かが既に引き出しに収納しているかもしれませんね。では引き出しに収納しているデータを勝手に捨てたらどうでしょうか。データの持ち主が引き出しを開けに帰ってきて「自分のお道具セットがない！」なんてあったらあまりに可哀想です。従って新しいデータを格納したい場合は誰も使っていない引き出しを探す必要があります。しかしどこが空いているかなんて分かりませんよね。

そこで、どの引き出しが利用中でどの引き出しが空いているのかを管理する人が必要です。まさにその人こそが「アロケータ」です。

1. アロケータの紹介

アロケータにはいくつか種類があります。実は万能なアロケータは存在しないのです。用途によってどのアロケータが最適なのかが変わってくるためです。その中でも色々なアロケータの中で取り入れられている重要な考え方が詰まったアロケータが下記の3つです。

  • Bump Allocator
  • Linked List Allocator
  • Fixed Size Block Allocator

2. アロケータの実装

「アロケータの実装」というワードを聞くだけだとなんだか漠然としていて難しそうに見えます。しかし、アロケータの責任はたったこれだけです。

  1. 空きメモリの中から、データが入るために十分なスペースがあるアドレスを探すこと
  2. 解放命令が下されるまで、格納したデータを保護すること

とてもシンプルですね。いかに処理の無駄を削ぎ落としていかに速度を向上させるかを考えなければ、全く難しいものではありません。単純なアルゴリズムであればたった数行で完成してしまいます。

手始めに「アホなアロケータ」を作る

アロケータの責任範囲に対する理解を深めるには、まず「アホなアロケータ」を作ることが良いです。
ではどのようにアホにするか。パッと思いつくアイデアとしては、ずっと単一のアドレスを返すことです。(2番目のアロケータの責任を無視しますが一旦良しとします)

試しに実装してみましょう。

アロケータが持つ状態

  • heap_start: ヒープメモリの先頭位置
noob_allocator.rs
struct NoobAllocator {
    heap_start: usize,
}

アロケータの初期化

  • ヒープメモリの先頭アドレスを受け取る
  • heap_start にヒープメモリの先頭アドレスを代入する
noob_allocator.rs
unsafe fn init(&mut self, heap_start: usize, _heap_size: usize) {
    self.heap_start = heap_start;
}

アロケートアルゴリズム

  • ヒープメモリの先頭アドレスを return する
noob_allocator.rs
unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
    let allocator = self.lock();
    allocator.heap_start as *mut u8
}

デアロケートアルゴリズム

アホなアロケータは空き領域を探さず単一のアドレスを返すので、デアロケーションする必要はありません。従って何も実装する必要はありません。

noob_allocator.rs
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
    // なにもしない
}

動かしてみる

下記のように41という数値をヒープメモリに載せてみましょう。(Rust では Box を使うことで強制的にヒープメモリに格納してくれます)

tests/heap_allocation.rs
#[test_case]
fn simple_allocation() {
    let heap_value_1 = Box::new(41);
    assert_eq!(*heap_value_1, 41);
}
出力
heap_allocation::simple_allocation...   [ok]

アホなアロケータは単一のアドレスを返しているだけですが、テストが PASS しました！
次の例ではどうでしょうか。2回ヒープアロケーションが走るテストとなっています。

tests/heap_allocation.rs
#[test_case]
fn simple_allocation() {
    let heap_value_1 = Box::new(41);
    let heap_value_2 = Box::new(13);
    assert_eq!(*heap_value_1, 41);
    assert_eq!(*heap_value_2, 13);
}
出力
heap_allocation::simple_allocation...   [failed]

Error: panicked at tests/heap_allocation.rs:44:5:
assertion `left == right` failed
  left: 13
 right: 41

なんと失敗しました！なぜならばアロケータが利用中のアドレスを空きアドレスとして返却しているからです。つまり競合相手が現れるとデータを上書きされてしまいます。

これで少しはアロケータに必要なモノが何であるかの体感値を手に入れていただけたかと思います。それでは Writing an OS in Rust で紹介されている実装を元に、複数回のアロケーションが走っても正常に動作するように修正してみましょう。

実装の前に前回のおさらい (アラインメント)

アロケータ実装で最も注意するポイントは、格納先のアドレス調整です。前回の記事を思い出してみてください。

メモリを扱う上でアラインメントは絶対的正義です。

アラインメントのことは思い出していただけたでしょうか。どんなことがあってもアラインして格納しなければなりません。格納するデータが「何 byte 境界アラインされなければいけないのか」を考慮しなければならないことを頭に入れつつ読み進めてください。

アラインメントはアロケーション実装する上で頻繁に登場するので関数として定義しておくのがベストです。アルゴリズムとしては下記の通りです。

  • address と align を受け取る
  • address を align で割った余りを計算する (A)
  • (A) を remainder とする
  • remainder が 0 かどうかチェックする
    • 0 に一致するのであれば
      • address を return する
    • 0 に一致しないのであれば
      • address から remainder を引いて align を足す (B)
      • (B) を return する
/// `addr` を `align` の倍数になるようにする。
fn align_up(address: usize, align: usize) -> usize {
    let remainder = address % align;
    if remainder == 0 {
        address
    } else {
        // align - remainder > 0 なので、address + (align - remainder) は必ず address より大きい
        address + (align - remainder)
    }
}

Bump Allocator

Bump Allocator は最も単純で分かりやすいアルゴリズムを採用しています。こいつはアロケート対象のデータを貰って、低アドレスから高アドレスに向かって安直に積み上げていきます。アロケーションする度に、空きメモリの先頭アドレス (next) を更新して、次回アロケーション時のために覚えておきます。

Pros / Cons

  • Pros
    • 空きメモリを探す必要が全くないので非常に高速
  • Cons
    • アロケートされている区画の個数が0にならないとメモリの掃除ができない

アロケータが持つ状態

  • heap_start: ヒープメモリの先頭位置
  • heap_end: ヒープメモリの末尾位置
  • next: 空きメモリの先頭アドレス
  • allocations: アロケートされている区画の個数
bump_allocator.rs
struct BumpAllocator {
    heap_start: usize,
    heap_end: usize,
    next: usize,
    allocations: usize,
}

アロケータの初期化

  • ヒープメモリの先頭アドレスとサイズを受け取る
  • heap_start にヒープメモリの先頭アドレスを代入する
  • heap_end に heap_start とサイズを足し合わせたものを代入する
  • next に heap_start を代入する
bump_allocator.rs
unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
    self.heap_start = heap_start;
    self.heap_end = heap_start + heap_size;
    self.next = self.heap_start
}

アロケーションアルゴリズム

  • next を基準にして、受け取ったデータが求めるアライン境界を満たしたアドレスに調整する (A)
  • (A) を start_address とする
  • start_address とレイアウトサイズをもとにデータの末尾アドレスを計算する (B)
  • (B) を end_address とする
  • end_address が heap_end を超えているかチェックする
    • 超えているのであれば
      • null を return する
    • 超えていないのであれば
      • next を end_address に更新する
      • allocations をインクリメントする
      • start_address を return する
bump_allocator.rs
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    let mut bump = self.lock();

    let alloc_start = align_up(bump.next, layout.align());
    let alloc_end = match alloc_start.checked_add(layout.size()) {
        Some(end) => end,
        None => return ptr::null_mut(),
    };

    if alloc_end > bump.heap_end {
        ptr::null_mut()
    } else {
        bump.next = alloc_end;
        bump.allocations += 1;
        alloc_start as *mut u8
    }
}

デアロケーションアルゴリズム

  • allocations をデクリメントする
  • allocations が0かどうかチェックする
    • true ならば
      • next に heap_start を代入することでリセットする
bump_allocator.rs
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
    let mut bump = self.lock();
    
    bump.allocations -= 1;
    if bump.allocations == 0 {
        bump.next = bump.heap_start;
    }
}

Linked List Allocator

Linked List Allocator は、空きメモリ区画の1つ1つを Node として見なし、全体に散らばった空きメモリ全体を Node の列としてみなすアルゴリズムを採用しています。これによって Bump Allocator では開放済みメモリを掃除できずにメモリを逼迫してしまう問題がありましたが、Linked List Allocator はどのメモリがどの程度空いているのかを Node として保持しているので、メモリを無駄にしないという強みがあります。

Pros / Cons

  • Pros
    • メモリの無駄が最小限に抑えられる
  • Cons
    • 適切なサイズの空きメモリを探すのに時間がかかる
    • 空きメモリの断片化が発生する
      • ただし Node を整列し結合する処理を挟まなかった場合に限る

Node の定義

アロケータが持つ状態を見る前に、Linked List Allocator の根幹を支える Node の定義から始めましょう。Node は空きメモリ区画を指すので定義は下記のようになります。

  • size: 空きメモリ区画のサイズ
  • next: 次の空きメモリ区画の先頭アドレス (Optional)
    • 末尾 Node の場合、次の空きメモリ区画の先頭アドレスが存在しないため Optional
linked_list_allocator.rs
struct Node {
    size: usize,
    next: Option<&'static mut Node>,
}

impl Node {
    const fn new(size: usize) -> Self {
        Node { size, next: None }
    }

    fn start_address(&self) -> usize {
        self as *const Self as usize
    }

    fn end_address(&self) -> usize {
        self.start_address() + self.size
    }
}

アロケータが持つ状態

ちょっと驚くかもしれませんが、ヒープメモリの先頭アドレスとサイズを状態として保持する必要がありません。それもそのはず、空きメモリ区画を表現するのが Node なのでアロケータはわざわざどこからどこまでが空いているかなんて覚えておく必要がないのです。したがって定義は下記の通りになります。

  • root_node: Node の先頭
linked_list_allocator.rs
struct LinkedListAllocator {
    root_node: Node,
}

空きメモリ区画を追加する関数の定義

アロケータの初期化方法を実装する前に、便利な関数を定義していきましょう。空きメモリ全体は単体の Node ではなく複数の Node で表現するので、空きメモリ区画を追加する関数が必要です。

少し面白いのですが、Linked List Allocator は、空きメモリを表す Node をその空きメモリの先頭に書き込むということをします。「？？？」となりますよね。なぜ空きメモリにわざわざ Node を書き込むのでしょうか。それは可変長データを使えないからです。よくよく考えてみれば可変長データを使うためにはアロケータが必要です。アロケータがないからアロケータを作っているのですから、可変長データだなんて便利なものは使えないのです。

これを踏まえて実装をみていきましょう。

  • address (空きメモリ区画の先頭アドレス) と size (空きメモリ区画のサイズ) を受け取る
    • ただし address は Node のアラインメント要件を満たしているとする
    • ただし size は Node のサイズ以上であるという要件を満たしているとする
  • Node を size で初期化する (A)
  • (A) を new_node とする
  • root_node の next をコピーする (B)
  • (B) を new_node の next に代入する
  • root_node の next を空にする
  • address に new_node を書き込む
  • root_node の next に address を代入する
linked_list_allocator.rs
/// `address` は `Node` がのアラインメント要件を満たしているとする
/// `size` は `Node` のサイズ以上であるという要件を満たしているとする
unsafe fn add_free_region(&mut self, address: usize, size: usize) {
    let mut new_node = Node::new(size);
    new_node.next = self.root_node.next.take();
    let node_ptr = address as *mut Node;

    unsafe {
        node_ptr.write(new_node);
        self.root_node.next = Some(&mut *node_ptr)
    }
}

Node にデータを書き込めるかの判定をする関数の定義

格納したいデータがやってきた時、Node に入るかどうかを判定するロジックはちょっとだけ複雑なので関数として切り出しておきましょう。Node をデータに置き換えた時に、余ってしまったメモリを無駄にしないように、余ったメモリに Node が入るかどうかも検査項目に加えます。

  • node と size (データサイズ) と align(データのアラインメント) を受けとる
  • node の先頭アドレスを align でアラインする (A)
  • (A) を alloc_start とする
  • alloc_start に size を足す (B)
  • (B) を alloc_end とする
  • alloc_end が node の末尾アドレスより大きいかチェックする
    • もしも true ならばエラーを返す
  • node の末尾アドレスから alloc_end を引く (C)
  • (C) を excess_size とする
  • excess_size が0より大きく、かつ Node の size より小さいかチェックする
    • もしも true ならばエラーを返す
  • alloc_start を返す
linked_list_allocator.rs
fn alloc_from_region(node: &Node, size: usize, align: usize) -> Result<usize, ()> {
    let alloc_start = align_up(node.start_address(), align);
    let alloc_end = alloc_start.checked_add(size).ok_or(())?;

    if alloc_end > node.end_address() {
        return Err(());
    }

    let excess_size = node.end_address() - alloc_end;
    if excess_size > 0 && excess_size < size_of::<Node>() {
        return Err(());
    }

    Ok(alloc_start)
}

データを格納できる Node を探す関数を定義する

先ほど定義した Node にデータを書き込めるかの判定をする関数を利用して、データの格納が可能な Node を探し出せるようにしましょう。

  • size (データサイズ) と align(データのアラインメント) を受け取る
  • root_node を current とする
  • current の next が存在しなくなるまで下記処理をループする
    • next を region とする
    • region に size と align を渡して格納可能かチェックする
      • 格納可能ならば
        • region の next をコピーする (A)
        • (A) を next とする
        • current の next をコピーしたものと alloc_start のタプルを ret とする
        • current の next に next を代入する
      • 格納不可能ならば
        • current に current の next を代入する
linked_list_allocator.rs
unsafe fn find_region(
    &mut self,
    size: usize,
    align: usize,
) -> Option<(&'static mut Node, usize)> {
    let mut current = &mut self.root_node;
    while let Some(ref mut region) = current.next {
        if let Ok(alloc_start) = Self::alloc_from_region(region, size, align) {
            let next = region.next.take();
            let ret = Some((current.next.take().unwrap(), alloc_start));
            current.next = next;
            return ret;
        } else {
            current = current.next.as_mut().unwrap()
        }
    }

    None
}

格納データを Node のアラインメントに合わせる関数を定義する

デアロケーションする際に Node に置き換える必要があります。そのためアロケーションとデアロケーションする際にデータを Node のアラインメント要件を満たすようにレイアウトを修正する必要があります。

linked_list_allocator.rs
fn size_align(layout: Layout) -> (usize, usize) {
    let layout = layout
        .align_to(align_of::<Node>())
        .expect("adjust layout alignment failed")
        .pad_to_align();
    let size = layout.size().max(size_of::<Node>());
    (size, layout.align())
}

アロケータの初期化

諸々必要な関数の定義が完了しましたので、これらを使ってアロケータの実装に差し掛かっていきましょう。まず、アロケータの初期化をする時はヒープの先頭とそのサイズをもらって、それをまるっと Node にしてしまえば良いですね。

  • heap_start (ヒープメモリの先頭アドレス) と heap_size (ヒープメモリのサイズ) を受け取る
  • heap_start と heap_size を使って空きメモリ区画を追加する
linked_list_allocator.rs
unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
    unsafe { self.add_free_region(heap_start, heap_size) };
}

アロケーションアルゴリズム

  • データの layout を受け取る
  • layout と Node のレイアウト要件を満たすような size と align を求める
  • 利用可能な領域から、size と align を満たす領域とその開始アドレス (alloc_start) を検索する
    • もしも該当する領域が見つからなければ、null を返す
  • alloc_start に size を足して、データの末尾アドレスを計算する (B)
  • (B) を alloc_end とする
  • 対象領域の末尾アドレスから alloc_end を引いて、余剰領域を求める (C)
  • (C) を excess_size とする
  • excess_size が 0 より大きいかチェックする
    • もしも true ならば、excess_size を空き領域として登録する
  • alloc_start を返す
linked_list_allocator.rs
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    let (size, align) = LinkedListAllocator::size_align(layout);
    let mut allocator = self.lock();

    if let Some((region, alloc_start)) = unsafe { allocator.find_region(size, align) } {
        let alloc_end = alloc_start.checked_add(size).expect("overflow");
        let excess_size = region.end_address() - alloc_end;
        if excess_size > 0 {
            unsafe {
                allocator.add_free_region(alloc_end, excess_size);
            }
        }
        alloc_start as *mut u8
    } else {
        null_mut()
    }
}

デアロケーションアルゴリズム

  • ptr (開放するアドレス) と layout (開放するデータレイアウト) を受け取る
  • layout と Node のレイアウト要件を満たすような size を求める
  • ptr から size 分を開放領域として追加する
linked_list_allocator.rs
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    let (size, _) = LinkedListAllocator::size_align(layout);

    unsafe {
        self.lock().add_free_region(ptr as usize, size);
    }
}

Fixed Size Block Allocator

Fixed Size Block Allocator は、固定サイズの空きメモリと、フォールバックアロケータを持っているアロケータです。フォールバックアロケータとは、データが入る固定サイズブロックが見つからなかった場合に用いられる代替のアロケータです。Linked List Allocator では細かなメモリの断片化が問題となっていましたが、Fixed Size Block Allocator では固定サイズの空きメモリを設けることでこの問題を解決します。

Pros / Cons

  • Pros
    • よく使われるサイズのブロックが増えるのでメモリの無駄が少なくなる
    • データがフィットする固定サイズブロックが見つかった場合は高速にアロケーションできる
  • Cons
    • 小さな固定サイズブロックが増えすぎると、大きなデータをアロケーションできなくなる

ブロックを表す Node を定義する

ブロックも Linked List Allocator と同様に Node で表現しましょう。

fixed_size_block_allocator.rs
struct Node {
    next: Option<&'static mut Node>,
}

固定サイズを定義する

固定サイズを定義していきましょう。小さいものは 8byte から、大きいものは 2048byte にしてみます。このサイズはブロックのアラインメントでも利用します。

fixed_size_block_allocator.rs
const BLOCK_SIZE: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048];

また、アロケートしたいデータが入るブロックサイズを導出できるようにしたいので、そのための関数を定義しましょう。

  • layout (データのレイアウト) を受け取る
  • BLOCK_SIZE の中から layout の size がスッポリ入る index を探す
    • 見つかった場合
      • index を return する
    • 見つからなかった場合
      • 無 を return する
fixed_size_block_allocator.rs
fn list_index(layout: &Layout) -> Option<usize> {
    BLOCK_SIZE.iter().position(|&s| s >= layout.size())
}

アロケータが持つ状態

  • list_heads: 固定サイズブロックの先頭 Node の列
  • fallback_allocator: 固定サイズブロックが見つからなかった場合に利用するアロケータ
fixed_size_block_allocator.rs
struct FixedSizeBlockAllocator {
    list_heads: [Option<&'static mut Node>; BLOCK_SIZE.len()],
    fallback_allocator: linked_list_allocator::Heap,
}

アロケータの初期化

この初期化で面白い点は固定サイズブロックの初期化はしない点です。どうしてでしょうか。そのサイズブロックが頻繁に使われるかが分からないためです。そのため最初は固定サイズブロックは使わず、フォールバックアロケータを利用してアロケートします。そしてデアロケートが行われて初めて固定サイズブロックを作成します。

  • heap_start (ヒープの開始アドレス) と heap_size (ヒープサイズ) を受け取る
  • heap_start と heap_size で fallback_allocator を初期化する
fixed_size_block_allocator.rs
unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
    self.fallback_allocator.init(heap_start, heap_size);
}

アロケーションアルゴリズム

  • layout (データのレイアウト) を受け取る
  • データが格納できる BLOCK_SIZE の index を探す
    • index が存在する場合
      • list_heads の index 番目が存在するか確認する
        • 存在する場合
          • list_heads の index 版目を node とする
          • list_heads に node の next を代入する
          • node の next を空にする
          • node のアドレスを return する
        • 存在しない場合
          • BLOCK_SIZE の index 番目を block_size とする
          • BLOCK_SIZE の index 番目を block_align とする
          • block_size と block_align からレイアウトを作成する
          • 作成したレイアウトを渡してフォールバックアロケータにアロケートしてもらい、得られたアドレスを return する
    • index が存在しない場合
      • layout をフォールバックアロケータに渡してアロケートしてもらい、得られたアドレスを return する
fixed_size_block_allocator.rs
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    let mut allocator = self.lock();
    match list_index(&layout) {
        Some(index) => match allocator.list_heads[index].take() {
            Some(node) => {
                allocator.list_heads[index] = node.next.take();
                node as *mut Node as *mut u8
            }
            None => {
                let block_size = BLOCK_SIZE[index];
                let block_align = block_size;
                let layout =
                    Layout::from_size_align(block_size, block_align).expect("invalid layout");
                allocator.fallback_alloc(layout)
            }
        },
        None => allocator.fallback_alloc(layout),
    }
}

デアロケーションアルゴリズム

  • ptr (開放するアドレス) と layout (開放するデータレイアウト) を受け取る
  • データが格納できる BLOCK_SIZE の index を探す
    • index が存在する場合
      • 新しい Node を作成する (A)
      • (A) を new_node とする
      • list_head の index 番目をコピーして new_node の next に代入する
      • ptr に new_node を書き込む
      • list_head の index 番目に new_node のアドレスを代入する
    • index が存在しない場合
      • ptr と layout をフォールバックアロケータに渡してデアロケートしてもらう
fixed_size_block_allocator.rs
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    let mut allocator = self.lock();
    match list_index(&layout) {
        Some(index) => {
            let new_node = Node {
                next: allocator.list_heads[index].take(),
            };
            let new_node_ptr = ptr as *mut Node;
            new_node_ptr.write(new_node);
            unsafe { allocator.list_heads[index] = Some(&mut *new_node_ptr) };
        }
        None => {
            let ptr = NonNull::new(ptr).expect("dealloc null pointer");
            allocator.fallback_allocator.deallocate(ptr, layout);
        }
    }
}

3. まとめ

Writing an OS in Rust の実装を通してアロケータに対する理解度が深まったと思います。実装してみれば分かると思いますが、うまくキャッシュを効かせたり、処理の負荷を減らす工夫を施さないとパフォーマンスに影響がありそうなのが肌感覚で分かったと思います。
この記事を通してメモリアロケーションは高価な処理であるということが分かれば、プログラムを書く時にアロケーションをいかに減らすかが頭の中にチラつくはずです。

さらにアロケータを実装することでメモリに関する知識や、実行時にメモリがどのように利用されているのが鮮明に分かるようになったはずです。よいメモリライフを手に入れてください！

よろしければ GitHub のフォローもお願いします!!🥺
https://github.com/tomoikey

Discussion