Rustで実装するmalloc

この記事は、NTT Communications Advent Calendar 2021 21日目の記事です。

はじめに

こんにちは、イノベーションセンターの鈴ヶ嶺(@suzu_3_14159265)です。普段は、クラウド・ハイブリッドクラウド・エッジデバイスなどを利用したAI/MLシステムに関する業務に従事しています。本日は、Rustで動的メモリ確保(dynamic memory allocation)のmallocを実装してPythonやvimを動かしてみようという内容をお届けします。

また、去年もRustネタのアドベントカレンダーを書いているのでぜひ見ていただけると嬉しいです！

• NTTコミュニケーションズ Advent Calendar 2020 Rustで実装するNetflow Collector

実装するmallocのアルゴリズム

今回実装するmallocのアルゴリズムは小さなメモリサイズにはSegregated Free Listを用いて、大きなメモリサイズにはmmapを用いる方針で実装します。まず、最初に動的メモリ確保における重要な課題であるメモリ断片化の概要について説明し、それぞれのアルゴリズムの手法について説明します。

課題：メモリ断片化

上記の図ように、時間が経過するにつれてメモリの解放や格納などが行われた結果として未使用領域が飛び飛びに配置され、合計としての空き容量は存在するが連続して要求されたメモリ領域を提供できない課題をメモリ断片化と言います。このような虫食い状態のメモリ配置では、有効な資源(メモリ)を活用できません。このメモリ断片化の課題を次のSegregated Free Listによって解決します。

Segregated Free List

Segregated Free Listは上記の図のように、複数のリストをそれぞれのメモリサイズごとに管理する方式です。これによりメモリ断片化が起きないbest fitなメモリ割り当てを実現します。また、未使用メモリを探索する計算量もそれぞれの管理するリストからunlinkするだけなのでO(1)で済みます。割り当てるメモリ領域がない場合は、sbrkシステムコールによるヒープ拡張を用いてメモリを確保してListに追加するように実装する必要があります。しかし、このSegregated Free Listのみでは全ての要求されうるメモリサイズのリストを用意する必要があり現実的ではないため、今回は512Byte以上のメモリについては後述するmmapを用いて別管理する方法を取ります。

mmap

今回、512byte以上の大きなメモリサイズについてはmmapで別管理してメモリを確保、解放する方法を採用しました。本来のmmapはファイルをメモリにマップするためのシステムコールですが、ファイルディスクリプタ fd　に-1を指定し、無名マッピング MAP_ANONYMOUS　を利用することでメモリ確保APIとしても使用可能です。この命令を用いて直接kernelからメモリを取得して割り当てます。

非スレッドセーフ

また、今回マルチスレッドには未対応な実装となっているためいくつかのプログラムは普通にSegmentation faultで落ちるので注意してください。glibcなどで実装されている詳しいmallocのアルゴリズムについては知りたい方はDoug Leaのmallocという通称 dlmalloc をおすすめします。1つのファイルに実装がまとめられており比較的読みやすいものになっています。

ちなみに以下は、dlmallocのコンパイル方法です。

wget http://gee.cs.oswego.edu/pub/misc/malloc-2.8.4.c
gcc -fPIC -shared -o dlmalloc.so malloc-2.8.4.c 

Rustで共有ライブラリを作成する方法

wrap jemalloc

まず、 cargo new --lib malloc-rs で新しくライブラリプロジェクトを作成します。そして以下のように Cargo.toml の crate-type に cdylib を追加しましょう。これで cargo build　と実行することで*.so が生成され、Rustを用いて他言語から使用可能なライブラリが作成可能になります。

Cargo.toml

[package]
name = "malloc-rs"
version = "0.1.0"
edition = "2021"

[lib]
crate-type = ["cdylib"]

[dependencies]
libc = "0.2.112"
jemalloc-sys = "0.3.2"

試しに中身をjemallocにwrapして置き換えるようにコードを書くと次のようになります。実装する関数はmalloc, free, calloc, reallocとなります。jemallocはメモリ断片化の回避と並行処理に対してスケーラブルに動作するmalloc実装です。ちなみに余談ですがRust1.31.1までmallocの実装にjemallocが標準として使用されていました。しかし、システム言語であるのにシステムのメモリアロケータを使用しないのは不自然な点やバイナリサイズの増加が問題として浮かび上がってきたため現在ではシステムのアロケータが使用されるようになっています。

以下のコードを見ていただいたら分かるように、extern "C" と記述することでC APIとして公開できます。そのままだとmangleされるので #[no_mangele] と付け加えましょう。

実際の動作には、LinuxでLD_PRELOAD という環境変数に共有ライブラリを指定して、動的ライブラリを差し替えることができる仕組みを使います。今回、mallocを呼ぶと標準出力で call malloc と出力されるような実装にしたので実際に動作させて確認してみます。 またここでは、writeシステムコールを使用して出力しています。Rustで一般的に使われる println! やlibcの printf などは内部でmallocを使用しているためSegmentation faultで落ちるので注意してください。

src/lib.rs

use libc::size_t;

extern crate jemalloc_sys;
extern crate libc;

#[no_mangle]
pub unsafe extern "C" fn malloc(size: size_t) -> *mut libc::c_void {
    let message = "call malloc\n";
    let buf = message.as_ptr() as *const libc::c_void;
    let buf_len = message.len();
    libc::write(1, buf, buf_len);
    jemalloc_sys::malloc(size)
}

#[no_mangle]
pub unsafe extern "C" fn realloc(p: *mut libc::c_void, size: size_t) -> *mut libc::c_void {
    jemalloc_sys::realloc(p, size)
}

#[no_mangle]
pub unsafe extern "C" fn calloc(number: size_t, size: size_t) -> *mut libc::c_void {
    jemalloc_sys::calloc(number, size)
}

#[no_mangle]
pub unsafe extern "C" fn free(p: *mut libc::c_void) {
    jemalloc_sys::free(p)
}

実行結果

cargo new --lib malloc-rs
# vim malloc-rs/src/lib.rs
cargo build
export LD_PRELOAD=`pwd`/target/debug/libmalloc_rs.so
ls # 適当なコマンドを実行してみる

mallocの実装

Headerはやや冗長ですが、3つのメンバのメモリサイズのsize , mmapされたメモリかis_mmap, 次の要素を指し示す next を設定しました。今回のケースにおいて、mmapされたメモリかどうかはsizeで判断可能(512Byteより大きい)です。また本来であれば、割り当ては8Byteにアライメントされるためsizeの下位3bitに空き容量があります。それも今回は、実装の見通しやすさを優先して次の図のようにしました。

全体的な実装は以下のようになります。

• get_headerは、割り当てたメモリのポインタからHeaderサイズを引いてHeaderを取得します。
• init_mallocで、Segregated Free Listを初期化します。
• add_listは、Segregated Free Listに割り当て可能なメモリがない場合にヒープからsbrkを用いてメモリを確保してリストを追加します
• find_chunkは、Segregated Free Listからsize(Byte)のメモリを探索します

extern crate libc;

use libc::{c_void, size_t};

/// malloc Header
struct Header {
    /// size of buffer
    size: size_t,

    /// flag of mmap
    is_mmap: size_t,

    // next header
    next: *mut Header,
}

/// 8byte aligment
const ALIGN: usize = 8;

/// max byte in segregated free list
const MAX_BYTE: usize = 512;

/// init size of one free list
const INIT_LIST_SIZE: usize = 512;

/// add size of free list
const ADD_LIST_SIZE: usize = 512;

/// number of free list
const NUM_LIST: usize = MAX_BYTE / ALIGN + 1;

/// init size of heap(sbrk)
const INIT_HEAP_SIZE: usize = NUM_LIST * (INIT_LIST_SIZE + std::mem::size_of::<Header>());

/// flag of call init_malloc
static mut IS_INIT_MALLOC: bool = false;

/// segregated free list
static mut FREE_LISTS: [*mut Header; (NUM_LIST)] = [std::ptr::null_mut(); (NUM_LIST)];

fn get_align(size: usize) -> usize {
    (size + ALIGN - 1) / ALIGN * ALIGN
}

unsafe fn get_header(p: *mut c_void) -> *mut Header {
    let header = p.sub(std::mem::size_of::<Header>()) as *mut Header;
    header
}

/// init malloc function
/// Setup the initial value of the segregated free list using sbrk from heap.
unsafe fn init_malloc() -> Result<(), *mut c_void> {
    IS_INIT_MALLOC = true;

    let current_p = libc::sbrk(0);
    let ret = libc::sbrk((INIT_HEAP_SIZE as isize).try_into().unwrap());

    if ret != current_p {
        // fail sbrk
        return Err(ret);
    }

    // init segregated free list
    let mut p = ret;
    for i in 1..NUM_LIST {
        FREE_LISTS[i] = p as *mut Header;

        let num_header = INIT_LIST_SIZE / (i * ALIGN);
        for j in 0..num_header {
            let mut header = p as *mut Header;
            let size = i * ALIGN;
            (*header).size = size;
            (*header).is_mmap = 0;
            (*header).next = std::ptr::null_mut();

            let next_p = p.add(size + std::mem::size_of::<Header>());
            if j != (num_header - 1) {
                (*header).next = next_p as *mut Header;
            } else {
                // last element
                (*header).next = std::ptr::null_mut();
            }

            p = next_p;
        }
    }
    Ok(())
}

/// add segregated free list
/// When there is no more memory in the segregated free list, use sbrk to add memory from the heap.
unsafe fn add_list(size: usize) -> Result<*mut Header, *mut c_void> {
    let current_p = libc::sbrk(0);
    let num_header = ADD_LIST_SIZE / size;
    let ret = libc::sbrk(
        (num_header * (size + std::mem::size_of::<Header>()))
            .try_into()
            .unwrap(),
    );

    if ret != current_p {
        // fail sbrk
        return Err(ret);
    }

    let mut p = ret;
    for j in 0..num_header {
        let mut header = p as *mut Header;
        (*header).size = size;
        (*header).is_mmap = 0;
        (*header).next = std::ptr::null_mut();

        let next_p = p.add(size + std::mem::size_of::<Header>());
        if j != (num_header - 1) {
            (*header).next = next_p as *mut Header;
        } else {
            // last element
            (*header).next = std::ptr::null_mut();
        }

        p = next_p;
    }

    Ok(ret as *mut Header)
}

/// find header function
/// Get a header of a given size from the segregated free list.
unsafe fn find_chunk(size: usize) -> Result<*mut Header, *mut c_void> {
    // index of segregated free list
    let index = size / 8;

    if FREE_LISTS[index] == std::ptr::null_mut() {
        let new_list_ret = add_list(size);

        match new_list_ret {
            Ok(new_list) => {
                FREE_LISTS[index] = new_list;
            }
            Err(err) => {
                return Err(err);
            }
        }
    }

    let header = FREE_LISTS[index];

    // unlink chunk
    FREE_LISTS[index] = (*header).next;

    Ok(header)
}

/// malloc function
#[no_mangle]
pub unsafe extern "C" fn malloc(size: size_t) -> *mut c_void {
    if size == 0 {
        return std::ptr::null_mut();
    }

    if !IS_INIT_MALLOC {
        if init_malloc().is_err() {
            return std::ptr::null_mut();
        }
    }

    let size_align = get_align(size);

    if size_align <= MAX_BYTE {
        // get memory from segregated free list
        let header_ret = find_chunk(size_align);
        if header_ret.is_err() {
            return std::ptr::null_mut();
        }
        let header = header_ret.unwrap();
        return (header as *mut c_void).add(std::mem::size_of::<Header>());
    }

    let mmap_size = std::mem::size_of::<Header>() + size;

    let p = libc::mmap(
        ::std::ptr::null_mut(),
        mmap_size,
        libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
        libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
        -1,
        0,
    );

    if p == libc::MAP_FAILED {
        return std::ptr::null_mut();
    }

    let mut header = p as *mut Header;
    (*header).size = mmap_size;
    (*header).is_mmap = 1;

    p.add(std::mem::size_of::<Header>())
}

/// realloc function
#[no_mangle]
pub unsafe extern "C" fn realloc(p: *mut c_void, size: size_t) -> *mut c_void {
    let size_align = get_align(size);
    if p == std::ptr::null_mut() {
        return malloc(size_align);
    }

    let new_p = malloc(size_align);
    let header = get_header(p);

    let memcpy_size = if (*header).size < size_align {
        (*header).size
    } else {
        size_align
    };
    libc::memcpy(new_p, p, memcpy_size);

    free(p);
    return new_p;
}

/// calloc function
#[no_mangle]
pub unsafe extern "C" fn calloc(number: size_t, size: size_t) -> *mut c_void {
    let new_p = malloc(size * number);
    libc::memset(new_p, 0, size * number);
    return new_p;
}

/// free function
#[no_mangle]
pub unsafe extern "C" fn free(p: *mut c_void) {
    if p == std::ptr::null_mut() {
        return;
    }

    let header = get_header(p);
    let size = (*header).size;
    if (*header).is_mmap == 1 {
        // free mmap
        let munmap_ret = libc::munmap(p.sub(std::mem::size_of::<Header>()), size);
        debug_assert!(munmap_ret == 0);
    } else {
        // reuse in segregated free list
        let index = size / ALIGN;
        let first_header = FREE_LISTS[index];
        FREE_LISTS[index] = header;
        (*header).next = first_header;
    }
}

実行結果

実行には、シングルスレッドで動作するプログラムとしてPythonとvimを実行してみました。

Python(numpy)

ここでは以下のように、numpyでランダムに大きなサイズの配列を作成し続ける処理で問題ないかを確認しました。

cargo build
export LD_PRELOAD=`pwd`/target/debug/libmalloc_rs.so
python3 <<EOF
import numpy as np
import random
while True:
    print(np.sum(np.random.random(random.randint(1000, 10000))))
EOF

上記の結果から、動作には問題なく正常にmalloc, freeができていることが分かります。

invader.vim

次にvimで動作するインベーダーゲームのmattn/invader-vimを動作させてみました。

cargo build
export LD_PRELOAD=`pwd`/target/debug/libmalloc_rs.so
wget https://raw.githubusercontent.com/mattn/invader-vim/master/plugin/invader.vim
vim
#:source ./invader.vim
#:Invader

問題なく動作していることが分かります。（というかvimでインベーダーが動かせるのすごいなぁ…）

まとめ

コードを見て分かるとおりunsafeを多用していたりpure-Rustな実装ではないため、あまりRustの恩恵を感じることができないものになってしまいました。一方で、Resultなどの高機能なパターンマッチをmallocを実装する上で使用可能な点や、libcなどのライブラリとのスムーズな連携は素晴らしい点だと感じることができました。

Rustは現在では、rallocのようにpure-Rust実装のメモリアロケータが実装されており、またLinux Kernelにカーネル内の第二言語としてサポートを提案するようなRFCの議論もあるため個人的に今後の活躍を期待しています。

今回実装したコードは以下に置いておきました。

• suzusuzu/malloc-rs

それでは、明日の記事もお楽しみに！

参考

• Glibc malloc internal
• A Memory Allocator by Doug Lea
• dlmalloc version 2.8.4.c
• Dynamic Memory Allocation in the Heap
• jemalloc
• mattn/invader-vim