Zenn
🎷

Go to Wasm: 並行処理の移し替え奮闘記

2023/02/24に公開
2件
Go
Go
WebAssembly
WebAssembly
concurrency
tech

この記事は何

  • WebAssembly(以下、Wasm)の普及によりC/C++、Rustなどの処理を手軽にWebアプリに移植できて嬉しいですね。
  • 私はGo言語が好きなので、Goで作ったCLIをWasmに移植してみたのですが、CLIで実装したgoroutineによる並行処理ちゃんも、Wasm移植時は全力を発揮できず…。
  • 結局JavaScript + Web Workerで並行処理を書き直しました。このあたりのトリックを経験できて良かったです。

Go & Wasm

多くのドキュメント・記事があるのでここでは割愛します。個人的に、理解の手助けになった情報の紹介のみさせて頂きます。

GoでWasmをビルドターゲットにする場合、アプリとして作るか、関数として作るか、関数にするならJavaScriptのグローバルオブジェクトに入れ込むか、TinyGoを使ってexportするかなど、まだ選択肢がいろいろある(決まっていない)のが悩ましいですね。

CLI版の並行処理

本題ではないのですが、次節の内容を理解するための例として以下のCLIについて軽く説明します。この処理をそのままWasmに移植するといまいちだよね、というお話です。

https://github.com/tenkoh/go-pubmine

このツールは、昨今話題のNostrで使うおしゃれな公開鍵を生成するCLIです。公開鍵はランダム生成→発掘を繰り返して探しますが、かなりの試行回数が必要なため並行処理で実装しています。以下、そのあたりの抜粋です。

pubmine.go
type Generator struct {
    maxWorkers int64
    prefix     string
}

func (g *Generator) Mine(ctx context.Context) (*KeyPair, error) {
    //ワーカープールの作成
    sem := semaphore.NewWeighted(g.maxWorkers)
    ctx, cancel := context.WithCancel(ctx)
    //結果取得用のチャネル
    ckp := make(chan *KeyPair)
    go func() {
        for {
            //プールに空きができるまで待機
            if err := sem.Acquire(ctx, 1); err != nil {
                return
            }
            go func() {
                //公開鍵を生成して条件にあったら結果を送信
                defer sem.Release(1)
                kp, _ := genKeyPair()
                if strings.HasPrefix(kp.Public, g.prefix) {
                    ckp <- kp
                }
            }()
        }
    }()

    var kp *KeyPair
    var ok bool
    //期待した結果を1つ取得するか、contextを通じた中断が入ったら終了
    select {
    case kp, ok = <-ckp:
    case <-ctx.Done():
    }
    //お片付け
    cancel()
    for {
        if sem.TryAcquire(g.maxWorkers) {
            break
        }
    }
    close(ckp)
    //以下省略
}

処理の概要

  1. ワーカープールを作り、各ワーカーで公開鍵をランダム生成。(goroutineによる並行処理)
  2. 条件に合致する公開鍵が見つかったら処理終了。プールを片付ける。

goroutineとチャネルを使うことで簡単にこの処理が書けるのが良いですね。

JavaScript + Wasmで並行処理を書きたかったら

Wasmを取り巻く並行処理の現状

ご存知のようにJavaScriptはシングルスレッドを基本とした言語ですね。

https://developer.mozilla.org/en-US/docs/Web/API/HTML_DOM_API/Microtask_guide/In_depth

Promiseを使った非同期処理(結果的に並行処理)もありますが、いかんせんシングルスレッドですから、重い処理が流れると画面描画等がフリーズしたようになります。上記のようにWasmの中身を並行に実装しようがこの点からは逃れられません。(逆に1スレッドをひたすら使い倒す側になります)。これを回避するには、Web Workerを使ったマルチスレッドでの並列計算が選択肢となるようです。

https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API/Using_web_workers

Wasmでは重い処理を回すこともあるためか、並行処理のニーズは強く存在するようです。TensorFlowをSIMDとマルチスレッド処理で高速化した、というニュースも過去にはありますね。

https://developers-jp.googleblog.com/2020/09/simd-tensorflowjs-webassembly.html

並行処理を助けるため、Wasmでは”thread" proposalが出されています。この提案が通ったら気軽に並行処理が書けるのかな？とも思えますが、proposalの書き出しに「スレッドを生成・結合する責務は実装者にある」と書いてあるように、提供されるのは並行処理を手助けする新しい共有メモリとアトミック操作のようです。つまりGoのgoroutineで書いた並行処理を、あたかも魔法のように並列実行されるWasmバイナリとして出力するような仕組みではありません。

https://github.com/WebAssembly/threads/blob/main/proposals/threads/Overview.md

したがってWasmを使うと言っても、マルチスレッドを活用した並行処理は、従来のJavaScriptと同様にWeb Workerを使うことが基本になると解釈しました。上記proposalの例でも、メインスレッドとWorkerの間で共有メモリを使ってデータをやりとりしています。

現状の確認はこのあたりにして、そろそろWeb Workerを使って並行処理を移植してみましょう。

Web Workerを使った実装

早速ですが完成形です。

概略図

main
// initialize workers
const workers = [];
const threads = navigator.hardwareConcurrency;
const maxWorkers = threads > 1 ? threads - 1 : threads;
for(let i=0; i<maxWorkers; i++){
    const w = new Worker("mine.js");
    workers.push(w);
}

const mine = (prefix) => {
    //複数のワーカーから1つの結果を取り出すためにPromiseでラップする
    //ユーザーの操作のたびにこれを行う
    const jobs = [];
    for(const w of workers){
        const promise = new Promise((resolve) => {
            w.onmessage = (e) => {
                //エラー処理は省略
                resolve(e.data);
            }
        })
        jobs.push(promise);
        w.postMessage(prefix);
    }
    Promise.any(jobs).then((result)=>{
        // 何か結果を使った処理
        terminate();
    })
}

const terminate = () => {
    //処理の打ち切りと、連続実行に備えたワーカーの再生成
    for(let [i,w] of workers.entries()){
        w.terminate();
        w = null;
        const nw = new Worker("mine.js");
        workers[i] = nw;
    }
}
mine.js
//GoのWasm実行に必要なスクリプトの読み込み
importScripts("wasm_exec.js");
const go = new Go();
WebAssembly.instantiateStreaming(fetch("main.wasm"), go.importObject).then(
    (result) => {
        go.run(result.instance);
    },
);
onmessage = (e) => {
    //Wasmの中のmine(s string)関数を呼び出す
    const ret = mine(e.data);
    postMessage(ret);
};
main.go
func mine(this js.Value, args []js.Value) any {
    // もろもろのエラーハンドリングは記載省略。
    prefix := args[0].String()
    g, _ := pubmine.NewGenerator(prefix, maxWorkers)
    kp, _ := g.SimpleMine(context.Background())
    // エラーなく返せる型が決まっているので注意!
    return map[string]any{"public": kp.Public, "private": kp.Private}
}

func main() {
    js.Global().Set("mine", js.FuncOf(mine))
    select {}
}

かなり試行錯誤した結果つかんだポイントは以下の通りです。

  • Workerの生成・呼び出しの流れは基本通りでOK。ただしいくつか注意点がある。
  • Workerは直接DOMを操作できないため、メインスレッドに計算結果を渡し、メインスレッドがDOMを操作する必要がある。そのためGoのWasmをアプリではなく、関数を吐き出す方式にした。(グローバルオブジェクトに関数を生やした)。
    • Go側でpostMessageメソッドを取得すればどうにかなったかもしれないけど試していない
  • Wasmは各Workerの中でinstatntiateする。各Workerはメインスレッドと独立したスコープを持っため、メインスレッドでインスタンス化してもWorkerから触れないから。
  • 今回のユースケースのように『複数のWorkerを走らせて、どれか1つでも答えを得られたら処理を終える』場合には、WorkerのタスクをPromiseでラップした上で、Promise.any()を使う。処理を強制的に打ち切るにはWorker.terminate()が使えるが、Worker自体が終了してしまうため、何度も処理を繰り返すならその都度再生成が必要。
    • もっとうまい方法もあるようだが、そこまでこだわりたくなかった

未解決ポイントは以下の通りです。

  • Wasmの関数がデッドロックした。結果的には解消したものの真因は分からなかった
    • Wasm内でgoroutine -> channelを使った待ちを行うパターンがデッドロックした。おそらくこのパターンに合致するのかな？と思いますが、単純にドキュメント内容がつかみ切れていません…。
    • 上記パターンを除去すれば普通に動作しました。

なかなか骨が折れましたが、やりたい処理ができてHappyです。

おわりに

Wasmやってみたい！と勢いで着手しましたが、いろいろ学ぶことができて楽しかったですね。

JavaScriptのWorkerを使った並行処理にGoのgoroutineとの親近感が感じられたことが、面白い気づきでした。postMessageを使ったメインスレッドとWorker間のデータ授受は、チャネルを使った授受と同じような挙動ですものね。うん、書いていて面白い！

GitHubで編集を提案

Discussion

NoboNoboNoboNobo

JSのイベント駆動時、その呼び出された関数はできるだけ速やかにreturnしなければなりません。
そこで「待ち」の発生する処理を書いてはいけません。JSのイベントループが停止してしまいます。
JSでは巧妙に簡単に「待ち」の発生する処理を書けないようになっているので気づきにくいですが、Goでは簡単に書けちゃうという罠。
Goの記述的に「待ち」が発生するようなものはすべてgoroutineに載せて元のスタックは早々にreturnする必要があります。

tenkohtenkoh

コメントありがとうございます。
私の「Wasmの関数がデッドロックした。結果的には解消したものの真因は分からなかった」に対するご助言だと捉えたのですが、合っているでしょうか？？🤔

「Goの記述的に待ち」、というのはとても分かりやすいですね！なるほどです！