WASM 多线程与 SharedArrayBuffer

// 普通内存：只能被当前实例 + 宿主 JS 访问
const mem = new WebAssembly.Memory({ initial: 10 });

// 共享内存：可跨 worker/实例共享
const sharedMem = new WebAssembly.Memory({ initial: 10, maximum: 20, shared: true });
// shared: true 时，maximum 必填；底层是 SharedArrayBuffer
console.log(sharedMem.buffer instanceof SharedArrayBuffer); // true

// 通过 postMessage 把内存"传"给 worker（零拷贝，引用共享）
worker.postMessage({ mem: sharedMem }, [sharedMem.buffer]);

(module
  (memory $m 10 20 shared)   ;; 初始 10 页，最大 20 页，共享
  ;; 注意：shared 内存的 memory.grow 会被广播给所有线程
  (export "memory" (memory $m)))

Cross-Origin-Opener-Policy: same-origin
Cross-Origin-Embedder-Policy: require-corp

// 隔离成功时，crossOriginIsolated === true
if (crossOriginIsolated) {
  console.log("可以安全使用 SharedArrayBuffer + Atomics");
} else {
  console.warn("需要设置 COOP/COEP 头");
}

;; 对 addr 做一次原子的 +1（计数器的标准写法）
(func $atomic_inc (param $addr i32)
  (local $expected i32)
  (loop $retry
    local.get $addr
    local.get $expected
    i32.atomic.rmw.cmpxchg     ;; CAS：若 [addr]==expected 则写 expected+1
    local.get $expected
    i32.add
    local.set $expected
    local.get $addr
    i32.load
    local.get $expected
    i32.ne
    br_if $retry))

# .cargo/config.toml（wasm32-unknown-unknown）
[target.wasm32-unknown-unknown]
rustflags = [
  "-C", "target-feature=+atomics,+bulk-memory,+mutable-globals",
  # 线程库需要 bulk-memory（memcpy/memset 的 memory.copy/fill）
]

RUSTFLAGS="-C target-feature=+atomics,+bulk-memory" \
cargo build --target wasm32-unknown-unknown --release
# 产物会 import env.memory（shared），且带 '__wbindgen_thread_xform__' 等符号

use std::thread;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};

#[no_mangle]
pub fn parallel_sum(data: *const u32, len: usize, nthreads: usize) -> u32 {
    let data = unsafe { std::slice::from_raw_parts(data, len) };
    let counter = Arc::new(AtomicU32::new(0));

    let handles: Vec<_> = (0..nthreads)
        .map(|t| {
            let data = data.to_vec(); // 跨线程闭包需要 'static
            let counter = counter.clone();
            thread::spawn(move || {
                let chunk = len / nthreads;
                let start = t * chunk;
                let end = if t == nthreads - 1 { len } else { start + chunk };
                let sum: u32 = data[start..end].iter().sum();
                counter.fetch_add(sum, Ordering::SeqCst);
            })
        })
        .collect();

    for h in handles { h.join().unwrap(); }
    counter.load(Ordering::SeqCst)
}

[dependencies]
wasm-bindgen-rayon = "1.0"
rayon = "1.8"

use wasm_bindgen_rayon::init_thread_pool;
use rayon::prelude::*;

#[wasm_bindgen]
pub async fn run(n: usize) -> usize {
    // 只需一步初始化：自动创建 worker、共享内存、启动线程池
    init_thread_pool().await;

    (0..n).into_par_iter().map(|i| i * i).sum()
}

import init, { run } from "./pkg/hello.js";
await init();            // 加载主模块
const result = await run(10_000_000);
console.log(result);

// 基于 i32.atomic.rmw.cmpxchg 的自旋锁
pub struct SpinLock {
    flag: AtomicBool, // 底层是 shared memory 中的 0/1
}

impl SpinLock {
    pub fn lock(&self) {
        while self.flag.compare_exchange(
            false, true, Ordering::SeqCst, Ordering::SeqCst
        ).is_err() {
            // 忙等：用 std::hint::spin_loop() 触发 pause 指令
            std::hint::spin_loop();
        }
    }

    pub fn unlock(&self) {
        self.flag.store(false, Ordering::SeqCst);
    }
}

(func $spin_lock (param $addr i32)
  (local $ok i32)
  (block $done
    (loop $wait
      (local.set $ok
        (i32.atomic.rmw.cmpxchg $addr (i32.const 0) (i32.const 1)))
      (if (i32.eqz (local.get $ok))
        (br $done))      ;; 拿到锁
      (br $wait))))      ;; 继续自旋

;; 线程 A：等待 addr 的值从 0 变为非 0（超时 1 秒）
(i32.atomic.wait $addr (i32.const 0) (i64.const 1_000_000_000))
;; 返回：0=被唤醒，1=值不匹配（立即返回），2=超时

;; 线程 B：把 addr 写为 1 并唤醒等待者
(i32.store $addr (i32.const 1))
(atomic.notify $addr (i32.const 1))   ;; 唤醒 1 个等待者

use std::sync::atomic::{AtomicU32, Ordering};

// 消费者等待：值变化则唤醒
let data_ready = AtomicU32::new(0);

// 生产者
data_ready.store(1, Ordering::SeqCst);
// notify 映射到 wasm 的 atomic.notify（std 内部已实现）

// 消费者
while data_ready.load(Ordering::SeqCst) == 0 {
    // 内部调用 i32.atomic.wait（阻塞，不烧 CPU）
    std::thread::park();
}

// main.js —— 线程池调度器
const workers = [];
const poolSize = navigator.hardwareConcurrency || 4;

const sharedMemory = new WebAssembly.Memory({
  initial: 100, maximum: 500, shared: true,
});

const moduleBytes = await fetch("app.wasm").then(r => r.arrayBuffer());
const { module, instance } = await WebAssembly.instantiate(moduleBytes, {
  env: { memory: sharedMemory },
});

// 每个 worker 共享同一份模块 + 内存
for (let i = 0; i < poolSize; i++) {
  const w = new Worker(new URL("./worker.js", import.meta.url), { type: "module" });
  w.postMessage({ module, memory: sharedMemory }, [moduleBytes, sharedMemory.buffer]);
  workers.push(w);
}

// worker.js —— 每个线程一个 WASM 实例
let instance;

self.onmessage = async ({ data }) => {
  if (!instance) {
    // 实例化：共享模块 + 共享内存 → 每个 worker 得到独立栈、共享堆
    instance = await WebAssembly.instantiate(data.module, {
      env: { memory: data.memory },
    });
  }
  // 分配任务区间，直接读写共享内存
  const start = data.start, end = data.end;
  const result = instance.exports.work(start, end);
  self.postMessage({ result });
};

// worker 侧：每个线程把自己的结果原子累加进共享内存
#[no_mangle]
pub extern "C" fn work(start: *mut u32, len: usize, counter: *mut AtomicU32) {
    let data = unsafe { std::slice::from_raw_parts(start, len) };
    let sum: u32 = data.iter().sum();
    unsafe {
        // fetch_add：多线程并发累加不丢失
        (*counter).fetch_add(sum, Ordering::SeqCst);
    }
}