RustでWASMを使ったプラグインを呼び出す

目的

Rustではコンパイル時にコードが固まってしまうため、動的にプログラムを動かすことができません。いままではluaのエンジンをRust内で動かしてluaスクリプトを実行していました。

しかしRustに慣れているとluaがしんどく感じるので、なんとかRustでできないか考えました。

RustではWASMのランタイムがあってWASMのプログラムを実行できるのでプラグインをRustで書いて実行してみました。

コード

構成

- plugin
  - src
    - binding.rs # cargo component buildで自動生成
    - lib.rs
  - wit
    - world.wit
  -  Cargo.toml
- wasm-engine
  - src
    - main.rs
  - Cargo.toml

プラグイン

プラグインではStringの引数としてJSONを受けて、そのJSONにあるtokenを取得して𝕏APIのusers/meを呼び出してその結果をステータスコードとともにJSONの文字列として返します。
エラーはなるべく呼び出し元に返すようにしています。

作成

cargo install cargo-component
cargo component new plugin --lib

Cargo.toml

[package]
name = "plugin"
version = "0.1.0"
edition = "2024"

[dependencies]
wit-bindgen-rt = { version = "0.41.0", features = ["bitflags"] }
waki = { version = "0.5", features = ["json"] }
serde_json = "1"

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

[package.metadata.component]
package = "component:plugin"

[package.metadata.component.dependencies]

wit/world.iwt

package component:plugin;

interface plugin {
    execute: func(json: string) -> string;
}

world plugin-provider {
    export plugin;
}

src/lib.rs

#[allow(warnings)]
mod bindings;

use std::time::Duration;

use waki::{Client, RequestBuilder, Response};

use crate::bindings::exports::component::plugin::plugin::Guest;
use serde_json::json;

struct Component;

fn get_json(src: &str) -> Result<serde_json::Value, String> {
    match serde_json::from_str(src) {
        Ok(json) => Ok(json),
        Err(e) => Err(json!({
            "error": e.to_string()
        })
        .to_string()),
    }
}

fn get_token(json: &serde_json::Value) -> Result<String, String> {
    match json["token"].as_str() {
        Some(token) => Ok(token.to_string()),
        None => Err(json!({
            "error": "token not found"
        })
        .to_string()),
    }
}

fn make_header(token: &str) -> Vec<(&'static str, String)> {
    let mut headers = Vec::new();
    headers.push(("Authorization", format!("Bearer {}", token)));
    headers
}

fn make_response(builder: RequestBuilder) -> Result<Response, String> {
    let res = builder.send();
    match res {
        Ok(res) => Ok(res),
        Err(e) => Err(json!({
            "error": e.to_string()
        })
        .to_string()),
    }
}

fn make_result(response: Response) -> String {
    let status_code = response.status_code();
    let body = match response.json::<serde_json::Value>() {
        Ok(body) => body,
        Err(e) => {
            return json!({
                "status_code": status_code,
                "error": e.to_string()
            })
            .to_string();
        }
    };
    json!({
        "status_code": status_code,
        "body": body
    })
    .to_string()
}

impl Guest for Component {
    fn execute(json: String) -> String {
        let json = match get_json(&json) {
            Ok(json) => json,
            Err(e) => return e,
        };

        let token = match get_token(&json) {
            Ok(token) => token,
            Err(e) => return e,
        };
        let url = "https://api.twitter.com/2/users/me";
        let builder = Client::new()
            .get(url)
            .headers(make_header(&token))
            .connect_timeout(Duration::from_secs(5));
        let res = match make_response(builder) {
            Ok(res) => res,
            Err(e) => return e,
        };
        make_result(res)
    }
}

bindings::export!(Component with_types_in bindings);

コンパイル

cargo add target wasm32-wasip2
cargo component build --target wasm32-wasip2 --release

呼び出し側

プラグインを呼び出す側のエンジンを作成します。
引数でwasmファイルを与えます。
MyStateはadd_to_linker_asyncするたびにコンパイラが必要なものを教えてくれるので、それに合わせて作成しています。

Cargo.toml

[package]
name = "wasm-engine"
version = "0.1.0"
edition = "2024"

[dependencies]
anyhow = "1.0.97"
clap = { version = "4.5.35", features = ["derive"] }
reqwest = { version = "0.12", features = ["json", "multipart", "rustls-tls"], default-features = false }
serde_json = "1"
tokio = { version = "1", features = ["macros", "rt-multi-thread"]}
wasmtime = { version = "31.0.0", features = [] }
wasmtime-wasi = "31.0.0"
wasmtime-wasi-http = "31.0.0"

src/main.rs

use std::time::Instant;

use clap::Parser;
use serde_json::json;
use wasmtime::component::{Component, Linker, ResourceTable, TypedFunc};
use wasmtime::{Config, Engine, Store};
use wasmtime_wasi::{IoView, WasiCtx, WasiView};
use wasmtime_wasi_http::{WasiHttpCtx, WasiHttpView};

struct MyState {
    table: ResourceTable,
    ctx: WasiCtx,
    http_ctx: WasiHttpCtx,
}

impl IoView for MyState {
    fn table(&mut self) -> &mut ResourceTable {
        &mut self.table
    }
}

impl WasiView for MyState {
    fn ctx(&mut self) -> &mut WasiCtx {
        &mut self.ctx
    }
}

impl WasiHttpView for MyState {
    fn ctx(&mut self) -> &mut WasiHttpCtx {
        &mut self.http_ctx
    }
}

#[derive(Parser, Debug)]
struct Args {
    wasm_file: String,
}

#[tokio::main]
async fn main() {
    let args = Args::parse();
    if let Err(e) = start(args).await {
        println!("{}", e);
    }
}

async fn start(args: Args) -> anyhow::Result<()> {
    let start = Instant::now();

    // 非同期をサポートするためにConfigの準備
    let mut config = Config::new();
    config.async_support(true);
    let engine = Engine::new(&config)?;

    let compnent = Component::from_file(&engine, &args.wasm_file)?;
    let mut linker: Linker<MyState> = Linker::new(&engine);

    // リンクで重複を許すための設定
    linker.allow_shadowing(true);
    wasmtime_wasi::add_to_linker_async(&mut linker)?;
    wasmtime_wasi_http::add_to_linker_async(&mut linker)?;

    // 環境変数からトークンを取得
    let params = json!({"token": std::env::var("ACCESS_TOKEN").unwrap()}).to_string();

    // 事前準備完了
    println!("Execution time: {:?}", start.elapsed());

    // スレッドで実行
    let mut handlers = vec![];
    for i in 0..3 {
        let handler = tokio::spawn({
            let mut store: Store<MyState> = Store::new(
                &engine,
                MyState {
                    table: ResourceTable::new(),
                    ctx: WasiCtx::builder().build(),
                    http_ctx: WasiHttpCtx::new(),
                },
            );
            let params = params.clone();
            let linker = linker.clone();
            let compnent = compnent.clone();
            async move {
                let instance = linker.instantiate_async(&mut store, &compnent).await?;
                let plugin_index = instance
                    .get_export(&mut store, None, "component:plugin/plugin")
                    .unwrap();
                let execute_index = instance
                    .get_export(&mut store, Some(&plugin_index), "execute")
                    .unwrap();
                let execute: TypedFunc<(String,), (String,)> =
                    instance.get_typed_func(&mut store, execute_index).unwrap();
                let (result,) = execute.call_async(&mut store, (params,)).await?;
                execute.post_return_async(store).await?;
                let result = serde_json::from_str::<serde_json::Value>(&result).unwrap();
                println!("Result: {} {:?}", i, result);
                Ok::<(), anyhow::Error>(())
            }
        });
        handlers.push(handler);
    }
    for handler in handlers {
        if let Err(e) = handler.await {
            println!("Error: {:?}", e);
        }
    }
    println!("Execution time: {:?}", start.elapsed());

    Ok(())
}

実行

ACCESS_TOKEN=xxxxx cargo run ../plugin/target/wasm32-wasip2/release/plugin.wasm

結果

Execution time: 1.421843859s
Result: 0 Object {"body": Object {"data": Object {"id": String("19522946"), "name": String("青柳康平"), "username": String("aoyagikouhei")}}, "status_code": Number(200)}
Result: 1 Object {"body": Object {"data": Object {"id": String("19522946"), "name": String("青柳康平"), "username": String("aoyagikouhei")}}, "status_code": Number(200)}
Result: 2 Object {"body": Object {"data": Object {"id": String("19522946"), "name": String("青柳康平"), "username": String("aoyagikouhei")}}, "status_code": Number(200)}
Execution time: 1.699192153s

まとめ

エンジン側からアクセストークンを渡して、プラグイン側で𝕏APIを呼んで、エンジンで実行結果を受け取ることができました。

スレッドで実行するためにストアを毎回生成してLinkerとComponentをcloneしてみましたが、速度はあまり遅くないです。