wasi-threadsの概要とComponentModelでの利用の現状とその先 #

2024-01-22

JSで書かれたツールをRustで書き直し、Wasmで提供する。ということを考える中でwasi-threadsをベースにrayonを使用したComponentを作りたい。という思いがでてきたのだが、それが可能なのか、もし可能であればどうやるのかが分からなかったため調査した。

上記をwasmtimeとNode.jsで使用できることを目標とする。

目次

TL;DR #

• wasi-threadsをベースにrayonの使用は実現できそう
• Node.jsにおけるwasi対応はまだまだexperimentalで、自分でshimを用意する必要がありそう
• wasi-threadsを使用したwasmをComponent化するのは現状難しそう
• 将来的にはshared-everything-threadsによりthread周りがbuilt-inになりComponent Modelへフィードバックされそう

wasi-threadsについて #

proposalを見る #

wasi-threadsは以下に記載がある。現在はphase1のようだ。

Wasmにthreadsは存在するが、shared memory,atomic,wait/notifyについて定義されいるだけでspawnについては定義されていない。wasi-threadsはこのspawnを定義しようというものだ。

witを見ればわかるが、interfaceとしては以下のみのシンプルなものとなっている。

thread-spawn: func(
    /// A value being passed to a start function (`wasi_thread_start()`).
    start-arg: start-arg,
) -> thread-spawn-result

ゴールとしてはpthreadsのサポートだが、full POSIX compatibityを目指すものでもないと書いてある。 次は簡単なサンプルを動かしてみる。

minimumなサンプルを動かしてみる #

Rustで以下のようなコードを用意し、wasm32-wasi-preview1-threadsを target に build する。

use std::thread;

fn main() {
    let mut x = 10;
    thread::scope(|s| {
        s.spawn(|| {
            x += 20;
        });
    });
    println!("{x}");
}

以下で動く。(WASMTIME_NEW_CLIはcliのinterfaceが変わるようでひとまず 0 としている)

❯ WASMTIME_NEW_CLI=0 wasmtime run wasi-threads.wasm --wasm-features=threads --wasi-modules=experimental-wasi-threads

30

wasmtimeでは簡単に動くがNode.jsだと自分でthread-spawnを用意する必要がありそうでこれは、後述する。

wasmtimeでサンプルがどのように動いているのか #

上記サンプルがどのように動いているかはspecとwasmtimeのwasi-threads実装を見るのがよい。

細かい点を省略・簡略化して要点だけ抜き出すと以下だ。

pub fn spawn(&self, host: T, thread_start_arg: i32) -> Result<i32> {
    let instance_pre = self.instance_pre.clone();
    // ...omitted
    let wasi_thread_id = random_thread_id();
    let builder = thread::Builder::new() // ...omitted
    builder.spawn(move || {
        let result = catch_unwind(AssertUnwindSafe(|| {
            let mut store = Store::new(&instance_pre.module().engine(), host);
            let instance = instance_pre.instantiate(&mut store).unwrap();
            let thread_entry_point = instance
                .get_typed_func::<(i32, i32), ()>(&mut store, "wasi_thread_start")
                .unwrap();
            match thread_entry_point.call(&mut store, (wasi_thread_id, thread_start_arg)) {
              // ...omitted
            }
        }));
    })
}

大枠は以下の手順となっている。

• 1. host側でthreadをspawnする
• 2. Storeからwasm instanceを取り出す
• 3. 2 で取り出したinstanceからexportされているwasi_thread_start関数を取り出す
• 4. 3 で得られたentrypointにスレッドの ID とthread_start_argを渡す

thread_start_argはwitを見返せばわかるがthread-spawnの引数として渡される、thread内で実行すべきwasm内の関数へのポインタと引数だ。

順番は前後するがwasi-threadsのREADMEにもhostが何をすべきかは記載されている。

Upon a call to wasi_thread_spawn, the WASI host must:

1. instantiate the module again — this child instance will be used for the new thread
2. in the child instance, import all of the same WebAssembly objects, including the above mentioned shared memories, as the parent
3. optionally, spawn a new host-level thread (other spawning mechanisms are possible)
4. calculate a positive, non-duplicate thread ID, tid, and return it to the caller; any error in the previous steps is indicated by returning a negative error code.
5. in the new thread, call the child instance's exported entry function with the thread ID and the start argument: wasi_thread_start(tid, start_arg)

大体齟齬はないように見える。wat側も軽く眺めたがwasi_thread_startが呼ばれるとatomicでthreadの起動を待ち合わせ、start_argから呼ぶべき関数を特定し、call_indirectで呼び出しているようだった。

wasmにおけるthread、あまりイメージが湧いていなかったが、今回で解像度が上がった。

Node.jsでサンプルを動かす PoC #

wasmtimeでの動かし方やどう動いているかは理解した。次はNode.jsで動かすことを考えてみる。

Node.jsではworker threadsを使用することになると思うのだが、その際wasm instanceを共有できないため、このspecは実現が難しいのでは？と思っていたのだが、Twitter(X)でいろいろアドバイスをいただき、単にmemoryだけ共有しつつworker側でinstantiateし直せばいいのでは。という結論にひとまず至った。

確かにhostがすべきことの最初のステップがinstantiate the module againと書いてある。これが同じ挙動を指しているのかは分からないが。

とりあえず作ってみたPoCが以下だ。

全文は上記なのだが、少し分解し記載する。

main側 #

以下がwasmからみてmain thread側。 wasmをinstantiateするところまで難しいところはないがmemoryがsharedである点に注意。

const { Worker } = require("node:worker_threads");
const { readFile } = require("node:fs/promises");
const { WASI } = require("wasi");
const { argv, env } = require("node:process");
const { join } = require("node:path");

const wasi = new WASI({ version: "preview1", args: argv, env });
const file = readFile(join(__dirname, "wasi-threads.wasm"));

(async () => {
  const wasm = await WebAssembly.compile(await file);
  const opts = { initial: 17, maximum: 17, shared: true };
  const memory = new WebAssembly.Memory(opts);
  const instance = await WebAssembly.instantiate(wasm, {
    ...wasi.getImportObject(),
    wasi: {
      "thread-spawn": (start_arg) => {
        const worker = new Worker(WORKER_FILE, { workerData: { memory } });
        worker.postMessage(start_arg);
      },
    },
    env: { memory },
  });
  wasi.start(instance);
})();

instantiate時にはwasi-threadsを使用するために不足しているものとしてwasi.thread-spawnを自分で差し込む必要がある。 これによりRust(wasm)側でspawnが呼ばれとこの関数が呼び出される。 まずはnew Workerでworkerを作成し、memoryを共有しておく。

その後postMessageでstart_argを共有する。このstart_argは前述したように呼び出すべき関数の情報と引数を持ったデータ構造へのポインタだ。すなわちリニアメモリのindexとなる。

wasi: {
  "thread-spawn": (start_arg) => {
    const worker = new Worker(WORKER_FILE, { workerData: { memory } });
    worker.postMessage(start_arg);
  },
},

worker側 #

対してworker側は以下のようになる。 ほぼ同じ処理なのだが、wasi_thread_startを読んでいるのがポイントだ。

main側からmessageを受信するとinstantiateし、そこから取り出したwasi_thread_startを読んでいる。 これはwasi-threadsのspecにより定められた関数でwasmからexportされることになっている。

引数はthread_idとstart_argだ。thread_idはhostのthread_idを送る必要はないようだ。wasmtimeでもランダムなidを採番していた。

const { Worker, workerData, parentPort } = require("node:worker_threads");
const { readFile } = require("node:fs/promises");
const { WASI } = require("wasi");
const { argv, env } = require("node:process");
const { join } = require("node:path");

const wasi = new WASI({ version: "preview1", args: argv, env });

const file = readFile(join(__dirname, "wasi-threads.wasm"));

parentPort.on("message", async (start_arg) => {
  const wasm = await WebAssembly.compile(await file);
  const { memory } = workerData;
  const instance = await WebAssembly.instantiate(wasm, {
    ...wasi.getImportObject(),
    wasi: {
      "thread-spawn": (arg) => {
        const worker = new Worker(WORKER_FILE, { workerData: { memory } });
        worker.postMessage(arg);
      },
    },
    env: { memory },
  });
  // thread id and start_arg
  instance.exports.wasi_thread_start(1, start_arg);
  process.exit(0);
});

本番で動かすにはまだま考慮事項がありそうだが、ひとまずサンプルは動いた。 rayonも試しに使用してみたがこちらも動いてそうだ。

また、Xで教えてもらったこちらも参考になりそうなので後でコードを読んでみたい。ざっと眺めた感じ、似たようなことをやっているようには見える。

wasi-threadsとComponent Modelについて #

ここまででwasi-threadsの概要とwasmtime,Node.jsでの動作について確認できたのでComponent Modelについての調査を記載しいていく。 結論からいうと、現状今回作成したwasmをComponent化するのは難しそうだ。

具体的には以下のコマンドでエラーが発生する。

wasm-tools component new wasi-threads.wasm -o component.wasm --adapt ./wasi_snapshot_preview1.wasm

error: failed to encode a component from module

Caused by:
    0: module is only allowed to import functions
(base)

これはwasmが(import "env" "memory" (memory (;0;) 17 17 shared))を要求しており、かつ、Component Modelとしてはfunctionしかimportできないため。つまりはwasm-toolsのバグや未実装などではなく、仕様レベルで検討が必要な話となる。

この点に関しては質問したのだが、以下の回答をもらっており、つまりはshared-everything-threadsというproposalで議論されていく。とのことだ。

The current plan is to extend threading support in core wasm through the shared-everything-threads proposal. The meat of this proposal is adding shared attributes to everything (not just memory), allowing whole instances to be shared (avoiding the O(n^2) function table usage with previous approaches) as well as a new solution for thread-local storage. There's also possibly a thread.spawn instruction for actually creating a new thread, but it's contentious (since it's a lot more work for browsers) so it's possible instead that we'll need to add thread.spawn as a canonical built-in to the Component Model (which would be importable by core wasm with the shared func type added by the shared-everything-threads proposal).

proposalは上記。まだ生煮えではありそうで、理解は今後の課題とするが、軽く眺めた感じ以下のようなことに言及しているように見える。

• wasm threadsではatomicやnotifyなどを定義したがspawnなどは定義しなかった
• これによりwasi-threadsなどが提案されたが、これはwasm側の仕様として検討すべき
• 現在やテーブルが共有できないことによるオーバーヘッドや参照がスレッド間で共有できないことによりGC言語間でのthread運用に問題がある
• そのため、たとえばpthreadsに特化しない、あらゆる言語を考慮したthreadを定義する。
• thread.spawn,thread.hw_concurrencyなどをbuilt-inで用意し、Component Modelに修正を提案する
• sharedというannotationを用意し、共有状態を静的に示す

つまりはComponent Modelでwasi-threadsは利用できないし、そもそもthreads周りはまだまだ改良が入りそうで、その結果はComponent Modelにフィードバックされていきそう。と解釈した。

まとめ #

いまいちぼんやりしたイメージだったwasi-threadsの概要がこれで把握できた。また、hostやwasmがどのような動きをしているかの動作イメージの解像度も上がったように思う。

と、同時に現状のspecでは問題や実現できないことがありそうで、まだまだ改良が入りComponent Modelにもフィードバックが取り込まれていきそうという情報もキャッチアップできた。

次はshared-everything-threadsの動向を見ながら理解を進めていきたい。

以上。