RustでV8を再利用してWasmランタイムを作ってみた

TL;DR

ChromeやNode.jsで使われているJavaScriptエンジンであるV8は、WebAssembly(Wasm)の実行環境も含んでいます。頑張れば、このWebAssemblyの実行環境を利用して、非Web向けのWasmランタイムが作れるんじゃないかと思ったので、実際に作ってみました。作ったものは以下で公開しています。

V8について

V8は、オープンソースで開発されている高速なJavaScriptエンジンです。
Wasmが発表された頃の2015~2016年には、V8の最適化コンパイラを再利用してWasmランタイムを実装する取り組みがなされ、現在では、LiftoffとTurboFanという２つのコンパイラによって、Wasmの実行を高速化しています。[1] それぞれの役割は以下の通りです。

• Liftoff: Wasm命令に対してワンパスで簡単な最適化と機械語へのコンパイルを行う軽量なコンパイラ
• TurboFan: ホットな関数に注目して、重めの最適化をかけるコンパイラ

これらの2つのコンパイラがあることで、トレードオフの関係にあるウォームアップ時間と性能のバランスの良い点を実現することができます。

V8 API

V8はライブラリとしてC++ APIを提供しており、ChromeやNode.jsはこれを利用して、各ソフトウェアでJSを実行しています。しかし、V8はビルドシステムとして、Bazelを利用しているので、CMakeやRustのプロジェクトに組み込もうとするとかなり大変です。
そこで登場するのが、rusty_v8です。これV8 APIのRustバインディングを提供するクレート、神々の努力によってcargo build一発でビルドして、Rustのプログラムから利用することができます。

非Web向けWasmランタイムの構成

Wasmランタイムには2種類あり、ブラウザやNode.jsからWasmを利用するJS(Web)埋め込みと、OSの上でJSを使わずに動かす非Web埋め込みがあります。Wasm自体は、簡素な仕様になっており、ファイルやネットワークの機能を使うには、埋め込み側の関数をインポートして使います。JS埋め込みならJSの関数が、非Web埋め込みならシステムコールのラッパの仕様であるWASIで定められるネイティブ関数を利用することになります。

V8はもちろんJS埋め込みを想定して開発されているので、当然そのままだとOSの機能を使うことは出来ません。しかし、V8にはFunctionTemplateというJS側からネイティブ関数を呼び出すための機能があるので、これを利用すればOSとWasmの橋渡しをすることが出来ます。

図でまとめるとこんな感じになります。

コードを書いていく！

これから作るランタイムの流れをJSで書くと以下のようになります。

// 1. ファイルからバイト列を読み込んでWebAssembly.Moduleを作る
var module = WebAssembly.Module(WASM_FILE);
// 2. ネイティブ関数にWASIを実装して、JSの関数としてモジュールのインポートに渡す
var importObject = {
    wasi_snapshot_preview1: {
        fd_write: NATIVE_FUNCTION
    }
};
// 3. モジュールをインスタンス化して実行環境を作成する
var instance = WebAssembly.Instance(module, importObject);
// 4. Wasmモジュールのエントリーポイントを呼び出す
instance.exports._start();

1. Wasmファイルを読み込んでモジュールをインスタンス化する

まずは、Wasmファイルの内容を読み込んでWebAssembly.Moduleクラスのインスタンスを作成します。このクラス自体は、モジュールの読み込みや検証を行うだけで、実行環境を作ったり、実行したりはしません。

これはとても簡単に書くことが出来て、Rustの標準ライブラリでファイルを読み込んで、その内容をV8 APIにわたすだけで済みます。コードはこんな感じ。

let mut isolate = v8::Isolate::new(Default::default());
let scope = &mut v8::HandleScope::new(&mut isolate);
let context = v8::Context::new(scope, Default::default());
let scope = &mut v8::ContextScope::new(scope, context);
        
let wasm_module = std::fs::read("hello.wasm").expect("Failed to read file");
let module = v8::WasmModuleObject::compile(scope, &wasm_module).unwrap();

v8::IsolateはV8のインスタンスで、通常プログラムに一つだけ存在します。
V8では、JSの値やオブジェクトはHandleと呼ばれる構造体経由で扱うことになり、v8::HandleScopeは、そのハンドルをまとめて扱えるようにしたものです。
v8::Contextは、Isolate内に複数作ることができるJSの実行環境です。

そして、v8::WasmModuleObject::compileはV8のAPIで、この関数を呼び出すと、WebAssembly.Moduleインスタンスのハンドルが返ってきます。

2. インポートオブジェクトの作成

先ほど述べたように、WebAssembly自体にはファイルなどの入出力の機能がないので、そのような機能を持つ関数を外部からインポートすることになります。今回実行するWasmファイルは以下のHello worldを表示するプログラムになりますが、このモジュールではwasi_snapshot_preview1.fd_writeという関数をインポートして、画面に文字を書き込みます。

(module
  (type (;0;) (func (param i32 i32 i32 i32) (result i32)))
  (import "wasi_snapshot_preview1" "fd_write" (func $fd_write (type 0)))
  (export "memory" (memory 0))
  (memory $0 1)
  (data (i32.const 16) "Hello, World\n")

  (func $printHello
    (call $fd_write
            (i32.const 1)
            (i32.const 0)
            (i32.const 1)
            (i32.const 128)
    )
    drop
  )
  (func (export "_start")
    (i32.store (i32.const 0) (i32.const 16)) 
    (i32.store (i32.const 4) (i32.const 13))
    (call $printHello)
  )
)

したがって、まず、wasi_snapshot_preview1.fd_writeをネイティブ関数で実装してから、V8のFunctionTemplate機能でJSの関数にキャストします。そして、通常のJSのオブジェクトの子要素としてセットしてあげればよいです。

// 2. ネイティブ関数にWASIを実装して、JSの関数としてモジュールのインポートに渡す
var importObject = {
    wasi_snapshot_preview1: {
        fd_write: NATIVE_FUNCTION
    }
};

それでは、WASI関数の一つであるwasi_snapshot_preview1.fd_writeを実装していきます。WASIを自前で実装してもいいのですが、今回はできるだけ楽をするために、wasi-commonという外部のクレートに頼ります。(線形メモリを利用するのでwiggleも必要になります)

以下のように、グローバル変数でWASIのコンテキストを作り、wasi-commonの標準入出力をシステムのものを使うように設定します。wasi-commonはこのコンテキストを使って、fd_writeなどのWASIの実装を呼び出します。

use std::{
    cell::UnsafeCell,
    sync::{Mutex, OnceLock},
};
use wasi_common::sync::WasiCtxBuilder;
use wasi_common::WasiCtx;

static WASI_CTX: OnceLock<Mutex<WasiCtx>> = OnceLock::new();

fn get_wasi_ctx_mut() -> &'static Mutex<WasiCtx> {
    WASI_CTX.get_or_init(|| {
        let mut builder = WasiCtxBuilder::new();
        let mut builder = builder.inherit_stdin().inherit_stdout().inherit_stderr();
        Mutex::new(builder.build())
    })
}

続いて、V8に渡すfd_write関数を書いていきます。この関数はwasi-commonで提供されているfd_write関数を呼び出すラッパになっています。v8に渡すため、変数scopeとargsから、それぞれHandleScopeや関数呼び出し時に渡される際の引数(JSの値)を取得することが出来ます。

最初に、グローバル名前空間からgInstanceという名前のオブジェクトを取得していますが、これは後でWebAssembly.Instanceを作成した後にセットするグローバル変数です。その後に、gInstance.exports.memoryのようにして、WebAssemblyの線形メモリにアクセスします。これをwasi-commonのfd_writeに渡すことで、wasi-commonから線形メモリにアクセスして、メモリの内容を読み込むことができるようになります。

use std::{
    cell::UnsafeCell,
    sync::{Mutex, OnceLock},
};
use tokio::runtime::Runtime as TokioRuntime;
use wasi_common::snapshots::preview_1::wasi_snapshot_preview1 as preview1;

fn fd_write(
    scope: &mut v8::HandleScope,
    args: v8::FunctionCallbackArguments,
    mut rv: v8::ReturnValue,
) {
    // get global object
    let context = scope.get_current_context();
    let global = context.global(scope);
    // access to global memory
    let str_instance = v8::String::new(scope, "gInstance").unwrap();
    let instance = global.get(scope, str_instance.into()).unwrap();
    let instance = instance.to_object(scope).unwrap();

    // access to instance.exports.memory.buffer
    let str_exports = v8::String::new(scope, "exports").unwrap();
    let exports = instance.get(scope, str_exports.into()).unwrap();
    let exports = exports.to_object(scope).unwrap();
    let str_memory = v8::String::new(scope, "memory").unwrap();
    let memory = exports.get(scope, str_memory.into()).unwrap();
    let memory = memory.to_object(scope).unwrap();

    // cast as ArrayBuffer
    let str_buffer = v8::String::new(scope, "buffer").unwrap();
    let array_buffer = memory.get(scope, str_buffer.into()).unwrap();
    let array_buffer = array_buffer.cast::<v8::ArrayBuffer>();
    let backing_store = array_buffer.get_backing_store();
    let memory: &mut [u8] = unsafe {
        std::slice::from_raw_parts_mut(
            backing_store.data().unwrap().as_ptr() as *mut u8,
            backing_store.byte_length(),
        )
    };
    let memory = unsafe { &*(memory as *mut [u8] as *mut [UnsafeCell<u8>]) };
    let mut memory = wiggle::GuestMemory::Shared(memory);

    let arg0 = args.get(0);
    let arg0 = arg0.integer_value(scope).unwrap_or_default() as i32;
    let arg1 = args.get(1);
    let arg1 = arg1.integer_value(scope).unwrap_or_default() as i32;
    let arg2 = args.get(2);
    let arg2 = arg2.integer_value(scope).unwrap_or_default() as i32;
    let arg3 = args.get(3);
    let arg3 = arg3.integer_value(scope).unwrap_or_default() as i32;

    let mut wasi_ctx = get_wasi_ctx_mut().lock().unwrap();
    let result = TokioRuntime::new()
        .unwrap()
        .block_on(preview1::fd_write(
            &mut *wasi_ctx,
            &mut memory,
            arg0,
            arg1,
            arg2,
            arg3,
        ))
        .unwrap();

    rv.set(v8::Integer::new(scope, result).into());
}

これでWASIをネイティブ関数に実装することが出来たので、このネイティブ関数をFunctionTemplateを使用してv8::Functionにキャストし、インポートオブジェクトとして設定します。

let import_wasi_p1 = v8::Object::new(scope);
let func_template = v8::FunctionTemplate::new(scope, fd_write);
let func = func_template.get_function(scope).unwrap();
let str_fd_write = v8::String::new(scope, "fd_write").unwrap();
import_wasi_p1.set(scope, str_fd_write.into(), func.into());

let str_wasi_p1 = v8::String::new(scope, "wasi_snapshot_preview1").unwrap();
let import_object = v8::Object::new(scope);
import_object.set(scope, str_wasi_p1.into(), import_wasi_p1.into());

これで、WebAssembly.Instanceを作成する準備ができました。あとは、このコンストラクタに1で作成したWebAssembly.Moduleとインポートオブジェクトを渡せばよいです。

3. WebAssembly.Instanceの作成

グローバル名前空間からWebAssemblyを取得して、Instanceフィールドにアクセスします。
これをコンストラクタとして、new_instanceとして呼び出せば、WebAssembly.Instanceのインスタンスが作成できます。

let str_wasm = v8::String::new(scope, "WebAssembly").unwrap();
let global_wasm = global
    .get(scope, str_wasm.into())
    .unwrap()
    .to_object(scope)
    .unwrap();
let str2 = v8::String::new(scope, "Instance").unwrap();
let instance_ctor = global_wasm.get(scope, str2.into()).unwrap();
let instance_ctor = instance_ctor.cast::<v8::Function>();
let instance = instance_ctor
    .new_instance(scope, &[module.into(), import_object.into()])
    .unwrap();

ついでに、2でこのグローバル変数からインスタンスを参照してメモリにアクセスする必要があったので、gInstanceという名前のグローバル変数としてセットしてあげます。

let global = context.global(scope);
let str_ginstance = v8::String::new(scope, "gInstance").unwrap();
global.set(scope, str_ginstance.into(), instance.into());

4. モジュールを実行

あとはWebAssemblyモジュールのエントリーポイントとなる関数を呼び出すだけです。通常、WebAssemmlyにコンパイルする際に、コンパイラによって_startという名前の関数が作成されます。この関数がエントリーポイントとなるので、以下のようにgInstance.exports._start()として呼び出します。

let str_exports = v8::String::new(scope, "exports").unwrap();
let exports = instance.get(scope, str_exports.into()).unwrap();
let exports = exports.to_object(scope).unwrap();
let str_start = v8::String::new(scope, "_start").unwrap();
let start = exports.get(scope, str_start.into()).unwrap();
let start = start.cast::<v8::Function>();

// call instance.exports._start()
let ret = start.call(scope, exports.into(), &[]).unwrap();

ここまでのコードで実行してみます。

うまく動いていることが確認できました！

最後に

V8からWASIを利用して、非web埋め込みのwasmランタイムを実装することが出来ました。しかし、wasi関数の呼び出しごとにJSを経由してしまい、余計なオーバーヘッドが発生してしまいます。V8をフォークしたり、APIを追加することでこの辺りは高速化できるかもしれませんが現状ではとても大変そうです。

この記事に興味を持ったら、ぜひV8にコントリビュートして、より高速なwasmランタイムを作りましょう！

参考

[1] https://v8.dev/docs/wasm-compilation-pipeline