この記事は、WebAssembly Advent Calendar 2024 の 11 日目の記事です。
wasi-threadsの実例として’紹介させていただければと思います。
以前Visual Regression Test をサポートする reg-actions をリリースしたという記事を書きました。
今回はreg-actionsの画像比較・レポート作成部分をRustで書き直しwebpのサポートとwasi-threads対応を行ったためその過程をまとめ、共有します。
コードはまだmainに合流できていないところもありますが、以下にあります。
reg-actions概要
reg-actionsは以下のような設定を行うことで、プルリクエスト毎に指定ディレクトリの画像の差分を検出し、結果をコメントにて通知するgithub actionsです。
name: CI
on: [push, pull_request]
jobs:
build:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: reg-viz/reg-actions@v3
with:
github-token: "${{ secrets.GITHUB_TOKEN }}"
image-directory-path: "./images"
今回のスコープ
reg-actionsはreg-cliをwrapしたものです。
reg-cliはflowtypeを使用しており、その他のスタックや仕組みも古いものとなっており、その刷新を最優先とし、その過程でのwebp対応とprocessをforkしていた箇所のwasi-threadsへの置き換えをもう一つの目標としました。
また、dependenciesが古くても動くが一番の懸念はセキュリティです。wasmを使用することで攻撃しにくくする意図もありました。こういった観点からどのような変化があったかはまた別の記事にまとめたい思いがあります。
対応内容
pixelmatchの移植
まずは画像の比較部分の移植を行いました。具体的にはreg-cliが依存しているmapbox/pixelmatchというライブラリの移植です。
このライブラリはrgbaの byte 列を受け取って pixel 毎の diff をとるライブラリで移植自体は難しくありませんでした。
以下のように利用でき、pixel 単位での diff を算出できます。
fn main() {
let img1 = image::open("./examples/4a.png").unwrap();
let img2 = image::open("./examples/4b.png").unwrap();
let result = pixelmatch(
&img1.to_rgba(),
&img2.to_rgba(),
img1.dimensions(),
Some(PixelmatchOption {
threshold: 0.1,
include_anti_alias: true,
..PixelmatchOption::default()
}),
)
.unwrap();
}
image-diff-jsの移植
次にpixelmatchを wrap した、画像のdecodingやサイズ調整を行っているimage-diff-jsの移植を行いました。
基本的にはimage-rsが多くの種類の画像に対応しているので基本image-rsに任せれば OK なのですが、この作業の開始時にはimage-rsがwebpに対応していなかったため、自身でlibwebpを併用しました。(現時点では制限はありますが、webpをサポートしています。)
内容は古くなっていることに注意が必要ですが、当時の試行錯誤は以下に記載しています。
これにより以下のように画像のdecodeと比較が行えるようになりました。
あとはこれをreg-cliから使えるようにすれば良さそうです。
fn main() {
let img1 = std::fs::read(actual_dir.join(path))?;
let img2 = std::fs::read(expected_dir.join(path))?;
let res = image_diff_rs::diff(
img1,
img2,
&DiffOption::default()
).unwrap();
}
reg-cliの移植
上記で画像のdecodeや比較はできるようになったのでreg-cliの移植に取り掛かることにしました。
reg-cliは 1000 枚以上の画像を比較するユースケースを想定してため、process poolを用意し、それらが順次画像を読んで比較していく構成となっています。
開発時にはまだworker threadがなかったこともあり、以下のようなイメージでforkし比較処理を行っています。
import { fork } from 'child_process';
import path from 'path';
import type { DiffCreatorParams, DiffResult } from './diff';
const processes = range(concurrency).map(() => fork(path.resolve(__dirname, './diff.js')));
return bluebird
.map(
images,
image => {
return p.run({
...dirs,
image,
matchingThreshold,
thresholdRate,
thresholdPixel,
enableAntialias,
});
},
{ concurrency },
);
これをwasi-threadsとrayonに置き換えられないかと考えたのがこの移植を開始したきっかけです。結果的に以下のように置き換えられました。
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(options.concurrency.unwrap_or_else(|| 4))
.build()
.unwrap();
let result = pool
.install(|| {
targets.par_iter().map(|path| {
let img1 = std::fs::read(actual_dir.join(path))?;
let img2 = std::fs::read(expected_dir.join(path))?;
let res = image_diff_rs::diff(
img1,
img2,
&DiffOption::default(),
)?;
Ok((path.clone(), res))
})
})
.collect::<Result<Vec<(PathBuf, DiffOutput)>, CompareError>>()?;
wasi-threads
Rustでは以下のようにtargetをwasm32-wasip1-threadsを指定してbuildすることでwasi-threadsが使用できます。
cargo build --release --target=wasm32-wasip1-threads
その際、threadをspawnしている箇所がwasi-sdkを介してwasi.thread-spawnが呼ばれるようになります。これをNode.js側で用意してやる方針で進めました。
簡略化するとmain側では以下のようにthread-spawnの呼び出しでWorkerを作成し、引数やMemoryなどを送信します。Memoryはshared: trueとするのを忘れないようにしてください。
let nextId = 1;
const opts = { initial: 17, maximum: 20, shared: true };
const memory = new WebAssembly.Memory(opts);
let instance = await WebAssembly.instantiate(wasm, {
...new WASI().getImportObject(),
wasi: {
"thread-spawn": (startArg) => {
const worker = new Worker("./worker.js");
const tid = nextTid++;
worker.postMessage({ startArg, tid, memory });
return tid;
},
},
env: { memory },
});
Worker側はメッセージをうけてwasmを受け取ったMemoryとともにinstantiateし、wasi_thread_startにthreadIdとthread-spawnで受け取った値を渡し、処理を開始します。
const handler = async ({ startArg, tid, memory }) => {
const wasm = await WebAssembly.compile(await file);
let instance = await WebAssembly.instantiate(wasm, { ... });
instance.exports.wasi_thread_start(tid, startArg);
};
parentPort.addListener("message", handler);
これが基本的な流れになりそうです。ただ、いくつか注意点もあり、それらはwasm32-wasip1-threadsでrayonを使ったコードをNode.jsで動かすにも記載いたしました。もし興味があれば覗いて見てください。
rayon
ここまで準備ができるとNode.jsでもrayonが使用できます。SiteIsolationが必要となるもののbrowserでも動きそうに見えます。これは個人的には悲願です。
前述したように以下のようにpar_iter()を用いてforkしていた箇所を置き換えました。
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(options.concurrency.unwrap_or_else(|| 4))
.build()
.unwrap();
let result = pool
.install(|| {
targets.par_iter().map(|path| {
let img1 = std::fs::read(actual_dir.join(path))?;
let img2 = std::fs::read(expected_dir.join(path))?;
let res = image_diff_rs::diff(
img1,
img2,
&DiffOption::default(),
)?;
Ok((path.clone(), res))
})
})
.collect::<Result<Vec<(PathBuf, DiffOutput)>, CompareError>>()?;
ここまででコアな部分は完了になります。
あとはglobで対象となるファイルを読んだりレポートを作成したりを粛々とrustに移植しました。特に特筆する点もなさそうなので詳細は割愛します。
作業中ですがコードは以下です。
build
wasm32-wasip1-threadsをtargetに、と前述しましたが、今回はlibwebpを使用していることもあり、それだけではだめでwasi-sdkを使用するようにしました。
これでbuildができるようになります。
export WASI_VERSION_FULL=24.0
export WASI_SDK_PATH=`pwd`/wasi-sdk-${WASI_VERSION_FULL}
CFLAGS="--sysroot ${WASI_SDK_PATH}/share/wasi-sysroot" cargo build --release --target=wasm32-wasip1-threads
これによりwasmが生成されるのでこれをJSから読んで結果を受け取れれば良さそうです。
wasm - JS
前述したようにJSからwasmを実行し、結果を受け取る必要があるのですが、それは以下のように細工しました。
具体的にはmainは空にしておきつつ、wasm_mainをJSから呼ばせ、その返り値のpointerから結果のJSONを取得する案です。
#[cfg(all(target_os = "wasi", target_env = "p1"))]
pub fn main() {
// NOP
}
#[cfg(all(target_os = "wasi", target_env = "p1"))]
#[repr(C)]
pub struct WasmOutput {
pub len: usize,
pub buf: *mut u8,
}
#[cfg(all(target_os = "wasi", target_env = "p1"))]
#[no_mangle]
pub extern "C" fn wasm_main() -> *mut WasmOutput {
let res = inner().unwrap();
let mut s = serde_json::to_string_pretty(&res).unwrap();
let len = s.len();
let ptr = s.as_mut_ptr();
std::mem::forget(s);
let output = Box::new(WasmOutput { len, buf: ptr });
Box::into_raw(output)
}
その案においてJS側は以下のようになりました。
wasi.startはmainが空とはいえ初期化のために呼び出す必要があります。
あとはpointerからbufferへのpointerとlengthを取得し、それをJSONに変換することで結果をJS側で受け取っています。そのために#[repr(C)]をWasmOutputにつけています。
wasi.start(instance);
const res = instance.exports.wasm_main();
const view = new DataView(memory.buffer, res);
const len = view.getUint32(0, true);
const bufPtr = view.getUint32(4, true);
const stringData = new Uint8Array(memory.buffer, bufPtr, len);
const decoder = new TextDecoder('utf-8');
const string = decoder.decode(stringData);
instance.exports.free_wasm_output(m);
const report = JSON.parse(string);
このあたりのコードは以下です。
あとはこれらをreg-actionsから呼び出せば完了です。
まとめ
駆け足でしたが、古いJavaScriptのコードをlibwebpやwasi-threadsを利用しつつRustとwasmに移植した実例を紹介させていただきました。参考になれば幸いです。
余談ですが、wasi-threadsはComponent Modelに対応できていません。その対応としてwasi-threadsのフィードバックを得てShared-Everything Threads ProposalというProposalが検討されている認識です。
なので、そのあたりが固まったら、将来的にはwasi-threadsからShared-Everything Threadsへの変更を行う可能性がありそうです。その際はまた記事にできればと思います。
そのあたりは以下に書いたので合わせて御覧ください。
また、セキュリティやパフォーマンスの観点からどう変わったかは別途まとめて記事にしたいなと考えています。
FRAIM株式会社 開発チームのブログです。 AI 搭載クラウドドキュメントワークスペース「LAWGUE（ローグ）」を開発・提供しています。 利用技術はRust, Python, React, AWS, Kubernetes などです。
Discussion