LoginSignup
27
4
ElixirAdvent Calendar 2020Day 9
@kobae964

Rustler で行儀の良い NIF を書く

Last updated at Posted at 2020-12-08

概要

Rustler を使って、Elixir から呼び出せる NIF を Rust で書き、性能測定を行いました。

何をやった?

NIF (natively implemented functions) は、Erlang ランタイムから呼び出すことのできる、 C などのローレベルな言語で書かれた関数のことです。
NIF (に限らず一般に FFI) は、主に C で高速な計算を行いたい場合、または低レベルな機能を直接呼び出したい場合などに使われます。
Discord などが Elixir から Rust で書かれた関数を呼び出すことで、パフォーマンスクリティカルな機能を実装することに成功しています。 https://blog.discord.com/using-rust-to-scale-elixir-for-11-million-concurrent-users-c6f19fc029d3
今回は Rustler と呼ばれるライブラリを使って、Elixir から呼び出せる NIF を Rust で書き、性能測定を行いました。

そもそも yielding な NIF になぜしなければいけないのか?

Erlang は cooperative なタスクスケジューリング方式をとっています。これによりコンテクストスイッチのコストが削減できます。

cooperative なタスクスケジューリングは、実行されるタスクが定期的に実行をやめ、他のタスクに実行の順番を譲ることを前提に、preemptive なスケジューリングより単純で高速なスケジューリングができることが特徴です。以下に両方を比較した表を掲載します。

項目 cooperative preemptive
コンテクストスイッチの速度 高速 (保存すべきデータが少ないため) 低速 (レジスタの情報などをすべて保存する必要があるため)
タスクの実装の手間 大きい (中断できるポイントをキリの良いところに挟む必要がある) 小さい (中断ポイントを挟む必要はない)
割り込みを気にするべき度合い 低い (中断ポイントの間は割り込みが起こらないことが保証される) 高い (いつ中断されるか分からない)
それぞれのタスクを信頼すべきか? yes (全てのタスクが正しく実行の順番を譲る必要がある) no (「悪い」タスクがあっても強制的に実行の順番を奪える)
関連: Rust の futures, Go の goroutine

問題設定

NIF を使うためには、NIF を使って解くべき問題をうまく設定する必要があります。
今回は最小シュタイナー木を NIF を使って計算させることにしました。これは以下のような問題です。

シュタイナー木とは、エッジの集合$E$とノードの集合$V$から成る無向グラフ $G=(V,E)$ において、$V$の部分集合$T$が与えられたとき、$T$に含まれるノードすべてを含む木のことである。グラフのエッジに重みがつけられている時、重みの和が最小であるシュタイナー木を求めよ。

これは Dreyfus-Wagner のアルゴリズム を使うことで解けます。ここでは詳しい解説は行いませんが、計算量は $n = |V|, m = |E|, k = |T|$ として $O(3^k n + 2^k m \log n)$ 時間です。指数時間アルゴリズムですから、$k$ の増加とともにかかる時間が急増します。

実装

https://rhye.org/post/native-scheduling-erlang/ を参考にしました。
ここでは要点だけ説明します。

(1) 状態を定義する

計算途中の状態を表す構造体を定義し、計算を中断するときは全ての途中計算をその構造体の中に入れることにします。

state.rs
#[derive(Debug)]
pub(crate) struct State {
    // 入力
    pub n: usize,
    pub edges: Vec<(usize, usize)>,
    pub terms: Vec<usize>,
    // 距離などを覚えておく、動的計画法のためのメモリ領域
    pub dp: Vec<Vec<usize>>,
    /// pre[i][s] = (0, 0): (i, s) is the initial state
    /// pre[i][s] = (1, x): x <= s and dp[i][s] = dp[i][x] + dp[i][s - x]
    /// pre[i][s] = (2, y): dp[i][s] = dp[y][s] + cost(y, i)
    /// pre[i][s] = (3, 0): dp[i][s] is not reached
    pub pre: Vec<Vec<(i32, usize)>>,
    // 今何をやっているか?
    pub phase: usize,
    // ループインデックス
    pub loop_index: usize,
}

(2) 状態を変更しながら、次の中断ポイントまで進める関数を書く

core.rs に compute という名前の関数を定義し、「次の中断ポイントまで計算を進める」という操作を行わせます。
関数全体は長すぎるので到底掲載できません。そのため簡略化した関数を掲載します。コメントなどを読んで察してください。

core.rs
pub(crate) enum Ret {
    // 計算が正常に終了した。
    Ok(usize, Vec<(usize, usize)>),
    // 計算が以上終了した。
    Error(Error),
    /// Aborting the computation. state is mutated. The caller should save state for later invocations.
    Yielding,
}

pub(crate) fn compute(state: &mut State) -> Ret {
    let n = state.n;
    let k = state.terms.len();
    let terms = &state.terms;
    :
    :
    :
    if state.phase == 2 {
        // 一回ループを回す
        if state.loop_index < 1 << k {
            :
            : 実際の計算
            :
            state.loop_index += 1;
            return Ret::Yielding;
        }
    }
}

(3) compute をラップする関数を作り、そこでスケジューリングの諸々を行う

compute を呼び出して、返り値が Yielding であったならば (中断ポイントに到達したならば)、経過時間を調べ、Erlang に経過時間を rustler::schedule::consume_timeslice で報告し、返り値によって場合分けをします。

  • 返り値が true だった場合 (1 ms 経っており、Erlang 側が実行順を譲ってほしいと言った場合) 次の呼び出しを enif_schedule_nif でスケジュールします。Rustler は NifReturned::Reschedule というものを提供しているため、それで代用できます。
  • 返り値が false だった場合 (1 ms 経っておらず、Erlang 側がまだ実行順を譲る必要がないと言った場合) 自分自身を再帰呼び出しして、処理を続行します。
lib.rs
unsafe fn steiner_tree_yielding(env: Env<'_>, ptr: *mut State) -> Term<'_> {
    let start = Instant::now();
    let result;
    {
        let state = &mut *ptr;
        result = compute(state);
    }

    match result {
        Ret::Ok(ans, picked_edges) => {
            // destroy the state
            destroy_state(ptr);
            (atoms::ok(), (ans, picked_edges)).encode(env)
        }
        Ret::Error(e) => {
            // destroy the state
            destroy_state(ptr);
            (atoms::error(), e).encode(env)
        }
        Ret::Yielding => {
            let elapsed = start.elapsed();
            // We are given 1ms timeslice. This value is hardcoded.
            let elapsed_micro = elapsed.as_micros().min(1000000000) as i32;
            let percentage = (elapsed_micro / 10).max(1).min(100);
            let should_yield = rustler::schedule::consume_timeslice(env, percentage);
            if should_yield {
                let encoded = encode_state_ptr_as_NIF_TERM(ptr);
                let result = NifReturned::Reschedule {
                    fun_name: CString::new("steiner_tree_interrupted").unwrap(),
                    flags: SchedulerFlags::Normal,
                    fun: steiner_tree_interrupted,
                    args: vec![encoded],
                };
                // NOTE: result.apply(env).encode(env) won't work here:
                // NIF_TERM is just an alias of usize, so it would return the Erlang representation of an integer 0.
                Term::new(env, result.apply(env))
            } else {
                steiner_tree_yielding(env, ptr)
            }
        }
    }
}

以上で基幹となる部分の説明は終わりです。ここで定義した関数を Elixir 側に export する方法は、[バージョンアップに追従] Elixirから簡単にRustを呼び出せるRustler #1 準備編 [rustler 0.21.1]などを見ていただければと思います。

性能

https://rhye.org/post/native-scheduling-erlang/ と同じく、シュタイナー木を延々と求め続ける 10 個のタスクを spawn_link し、それとは別に 1 秒ごとにログを出力しようとするタスクを spawn_link しました。シュタイナー木を求めるタスクは、一度求め終わったら [1] や [8] などの、自身の番号を出力します。また、1 秒ごとにログを出力するタスクは、実際には前の呼び出しから何ミリ秒経っているかと、それと 1000 ミリ秒とのずれを出力します。

テストに使ったグラフおよび集合 $T$ は以下のようなものです。

  • $V = \{0, \ldots, n-1\}, E = \{(0, 1), (1, 2), \ldots, (n-2,n-1)\}, T = V$

このグラフでは $k = n$ が成立します。つまり $n$ の大小で実行時間を調整できるということです。

以下は GitHub Actions の自動テストで実行したときのログです。ログは OTP 23.1 / Elixir 1.11.2 のものですが、他の環境でも似たような結果でした。

yield する方 (n = 11)

[3]
[yielding]: Time since last schedule:1011.994114 ms, jitter: -11.994113999999968 ms
[2][3]
[yielding]: Time since last schedule:1009.711821 ms, jitter: -9.711820999999986 ms
[2][3][4]
[yielding]: Time since last schedule:1010.182276 ms, jitter: -10.182276000000002 ms
[2][3][6][7][8][9][5][1][10]
[yielding]: Time since last schedule:1009.218805 ms, jitter: -9.218804999999975 ms
[4][2][3]
[yielding]: Time since last schedule:1008.509958 ms, jitter: -8.509957999999983 ms
[6][4]
[yielding]: Time since last schedule:1008.631686 ms, jitter: -8.631685999999945 ms
[8][9][5][1][10][7][2][3]
[yielding]: Time since last schedule:1008.409084 ms, jitter: -8.409084000000007 ms
[6][4]
[yielding]: Time since last schedule:1007.722641 ms, jitter: -7.722640999999953 ms
[2][3]
[yielding]: Time since last schedule:1009.856148 ms, jitter: -9.856147999999962 ms
[8][10][5][9][1][7][6][4]
[yielding]: Time since last schedule:1007.9521 ms, jitter: -7.952099999999973 ms
[2][3]
[yielding]: Time since last schedule:1007.853612 ms, jitter: -7.853611999999998 ms
[6][4][10][5][8][9][1][2][7]
[yielding]: Time since last schedule:1007.58261 ms, jitter: -7.582610000000045 ms
[3]
[yielding]: Time since last schedule:1006.607348 ms, jitter: -6.607348000000002 ms
[6][4]
[yielding]: Time since last schedule:1009.876556 ms, jitter: -9.87655600000005 ms
[5][9][10][8][1][2][7][3]

yield しない方 (n = 8)

[1][2][3][4][5][6][7][8][9][10][1][2][3][4][6][5][7][8][9][10][1][2][4][3][5][6][7][8][9][10][2][1][4][3][5][6][7][8][10][9][1][2][4][3][6][5][8][7][9][10][2][1][4][3][6][5][8][7][9][10][2][1][4][3][5][6][7][8][10][9][1][2][3][4][5][6][7][8][10][9][2][1][4][3][6][5][7][8][9][10][1][2][4][3][5][6][8][7][10][9][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6]
[nonyielding]: Time since last schedule:1657.828671 ms, jitter: -657.828671 ms
[7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][7][8][9][10][1][2][3][4][5][6][8][7][9][1][10][3][2][5][4][7][6][1][8][3][9][5][10][7][2][1][4][3][6][5][8][7][9][1][10][3][2][5][4][7][6][1][8][3][9][5][10][7][2][1][4][3][6][5][8][7][9][1][10][3][2][5][4][7][6][1][8][3][9][5][10][7][2][1][4][3][6][5][8][7][9][1][10][3][2][5][4][7]
[nonyielding]: Time since last schedule:1210.907527 ms, jitter: -210.90752700000007 ms
[6][1][8][3][9][5][10][7][2][1][4][3][6][5][8][7][9][1][10][3][2][5][4][7][6][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][3][9][5][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7]
[nonyielding]: Time since last schedule:1082.74562 ms, jitter: -82.74561999999992 ms
[8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7]
[nonyielding]: Time since last schedule:1287.225716 ms, jitter: -287.2257159999999 ms
[10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7]
[nonyielding]: Time since last schedule:1145.399656 ms, jitter: -145.39965600000005 ms
[10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][4][7]
[nonyielding]: Time since last schedule:1087.274331 ms, jitter: -87.27433100000007 ms
[6][1][8][3][9][5][10][7][2][1][4][3][6][5][8][7][9][1][10][3][2][5][4][7][6][1][8][3][9][5][10][7][2][1][4][3][6][5][8][7][9][1][10][3][2][5][4][7][6][1][8][3][9][5][10][7][2][1][4][3][6][5][8][7][9][1][10][3][2][5][4][7][6][1][8][3][9][5][10][7][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7]
[nonyielding]: Time since last schedule:1081.941069 ms, jitter: -81.94106899999997 ms
[8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7]
[nonyielding]: Time since last schedule:1482.162191 ms, jitter: -482.1621909999999 ms
[10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][3][2][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7]
[nonyielding]: Time since last schedule:1090.608636 ms, jitter: -90.60863599999993 ms
[4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7]
[nonyielding]: Time since last schedule:1153.194424 ms, jitter: -153.19442400000003 ms
[4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7]
[nonyielding]: Time since last schedule:1437.213033 ms, jitter: -437.213033 ms
[10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7][8][1][9][3][10][5][2][7][4][1][6][3][8][5][9][7][10][1][2][3][4][5][6][7]

yield しない方が n が小さい分より多くのタスクを実行できていますが、ログ出力における 1000 ミリ秒とのずれが目立ちます。yield する方はほとんど正確に 1000 ミリ秒ごとにログ出力できています。
(両方の条件で n を揃えていないのは、yield しない方で n を 11 にしてしまうと、タスクが重くなりすぎて 1 秒ごとに出力するタスクが全く実行されなくなってしまうためです。)

テストは https://github.com/koba-e964/rustler_steiner/blob/main/lib/perftest.ex にあります。

苦労した点

問題設定

入力が小さく、必要な時間が長く、また入力の中身にほとんど依存せずに必要な時間が決まるような問題、およびそれを解くアルゴリズムを選ぶ必要がありました。シュタイナー木問題はこれらの要求を全て満たしていました。

Rustler の欠点

enif_consume_timeslice と enif_schedule_nif が扱いにくい

Rustler の公式ドキュメントには両者のドキュメントがありません。enif_schedule_nif に至っては rustler 側の safe binding そのものがなく、rustler-sys が提供している extern "C" unsafe な関数を直接呼び出す必要がありました。Erlang が提供している enif_schedule_nif は C の関数ポインタを受け取るのですが、rustler-sys が提供しているのはそれを直接呼び出す以下のような関数です:

pub fn enif_schedule_nif(
    env: *mut ErlNifEnv,
    fun_name: *const c_uchar,
    flags: c_int,
    fp: unsafe extern "C" fn(env: *mut ErlNifEnv, argc: c_int, argv: *const ERL_NIF_TERM) -> ERL_NIF_TERM,
    argc: c_int,
    argv: *const ERL_NIF_TERM
) -> ERL_NIF_TERM;

これの safe binding を提供してほしい気持ちがありますが、いまだに議論中 (rusterlium/rustler#107) であるため、難易度が高いものと推察することができます。

Erlang の NIF そのものの欠点

enif_consume_timeslice が扱いにくい

Rust binding には一切関係なく、以下のような疑問があります。

  • なぜ利用者側で経過時間 (与えられたタイムスライスに対する経過時間の割合) を計算して呼び出さなければならないのか? それは Erlang 側の仕事では?
  • そもそも与えられたタイムスライスというのは何秒なのか? 1 ms 程度と近似できるという記述はあったが、正確な値はわからない。1
  • 時間の精度が粗すぎる。1 ms の 1% 単位でしか消費した時間を報告できないので、例えば呼び出し開始時から 1 us しか経っていなくとも、1 ms の 1% である 10us 経過したように報告する必要がある。2

The scheduling timeslice is not an exact entity, but can usually be approximated to about 1 millisecond.

今後の課題

resource を使って状態を管理するというのはやりたいです。
現状の実装では単にメモリを enif_alloc を使って確保し、 enif_free を使って解放しています。確保した領域へのポインタを Erlang 側に渡さないといけませんが、現在これはポインタを usize に変換して、それを数値としてみなして Erlang の項へと変換し、Erlang 側に渡しています。状態を復元するときはこれの逆を行います。
resource を使う場合は、resource 特有の型が割り当てられ、他のデータ型とは区別されるため、より安全になります。

コード

https://github.com/koba-e964/rustler_steiner にあります。

  1. 以下のような記述が https://erlang.org/doc/man/erl_nif.html#enif_consume_timeslice にあります。 ↩

  2. 利用者側で 10us 以上経たない限り報告しないようにして、問題を緩和することはできます。 ↩

27
4
1

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
27
4