Rust编译WASM共享内存时线程循环变量异常共享问题排查

问题概述 #

当在多个Web Worker中并行运行Rust编译的WASM代码时，原本属于线程局部的循环变量（chunk_start、chunk_end、n）出现跨线程共享覆盖的情况，导致输出完全混乱。串行运行时，每个Worker能正确输出对应区间的数值（0-99、100-199、200-299），但并行时输出无规律。

相关代码 #

Rust核心代码（run函数） #

#![no_main]
#![no_std]

use core::panic::PanicInfo;
use js::*;

mod js {
    #[link(wasm_import_module = "imports")]
    extern "C" {
        pub fn abort(msgPtr: usize, filePtr: usize, line: u32, column: u32) -> !;
        pub fn _log_num(number: usize);
    }
}

#[no_mangle]
pub unsafe extern "C" fn run(worker_id: i32) {
    let worker_index = worker_id as u32 - 1;
    let chunk_start = 100 * worker_index;
    let chunk_end = chunk_start + 100; //Total pixels may not divide evenly into number of worker cores.
    for n in chunk_start as usize..chunk_end as usize {
        _log_num(n);
    }
}

#[panic_handler]
unsafe fn panic(_: &PanicInfo) -> ! { abort(0, 0, 0, 0) }

Cargo编译配置（.cargo/config.toml） #

[target.wasm32-unknown-unknown]
rustflags = [
    "-C", "target-feature=+atomics,+mutable-globals,+bulk-memory",
    "-C", "link-args=--no-entry --shared-memory --import-memory --max-memory=2130706432",
]

Worker端JS代码 #

const wasmSource = fetch("sim.wasm") //kick off the request now, we're going to need it

//See message sending code for why we use multiple messages.
let messageArgQueue = [];
addEventListener("message", ({data}) => {
    messageArgQueue.push(data)
    if (messageArgQueue.length === 4) {
        self[messageArgQueue[0]].apply(0, messageArgQueue.slice(1))
    }
})

self.start = async (workerID, worldBackingBuffer, world) => {
    const wasm = await WebAssembly.instantiateStreaming(wasmSource, {
        env: { memory: worldBackingBuffer },
        imports: {
            abort: (messagePtr, locationPtr, row, column) => {
                throw new Error(`? (?:${row}:${column}, thread ${workerID})`)
            },
            _log_num: num => console.log(`thread ${workerID}: n is ${num}`),
        },
    })

    //Initialise thread-local storage, so we get separate stacks for our local variables.
    wasm.instance.exports.__wasm_init_tls(workerID-1)   

    //Loop, running the Rust logging loop when the "tick" advances.
    let lastProcessedTick = 0
    while (1) {
        Atomics.wait(world.globalTick, 0, lastProcessedTick)
        lastProcessedTick = world.globalTick[0]
        wasm.instance.exports.run(workerID)
    }
}

主线程JS代码 #

//Let's count to 300. We'll have three web workers, each taking ⅓rd of the task. 0-100, 100-200, 200-300...

//First, allocate some shared memory. (The original task wants to share some values around.)
const memory = new WebAssembly.Memory({
    initial: 23,
    maximum: 23,
    shared: true,
})

//Then, allocate the data views into the memory.
//This is shared memory which will get updated by the worker threads, off the main thread.
const world = {
    globalTick: new Int32Array(memory.buffer, 1200000, 1), //Current global tick. Increment to tell the workers to count up in scratchA!
}

//Load a core and send the "start" event to it.
const startAWorkerCore = coreIndex => {
    const worker = new Worker('worker/sim.mjs', {type:'module'})
    ;['start', coreIndex+1, memory, world].forEach(arg => worker.postMessage(arg)) //Marshal the "start" message across multiple postMessages because of the following bugs: 1. Must transfer memory BEFORE world. 2. Must transfer world BEFORE memory.
}

//Now, let's start some worker threads! They will work on different memory locations, so they don't conflict.
startAWorkerCore(0) //works fine
startAWorkerCore(1) //breaks counting - COMMENT THIS OUT TO FIX COUNTING
startAWorkerCore(2) //breaks counting - COMMENT THIS OUT TO FIX COUNTING


//Run the simulation thrice. Each thread should print a hundred numbers in order, thrice.
//For thread 1, it should print 0, then 1, then 2, etc. up to 99.
//Thread 2 should run from 100 to 199, and thread 3 200 to 299.
//But when they're run simultaneously, all three threads seem to use the same counter.
setTimeout(tick, 500)
setTimeout(tick, 700)
setTimeout(tick, 900)
function tick() {
    Atomics.add(world.globalTick, 0, 1)
    Atomics.notify(world.globalTick, 0)
}

问题根源 #

1. 共享内存滥用：所有Worker复用同一块共享内存作为WASM的模块内存，但WASM的栈、线程局部存储（TLS）等线程专属数据默认存储在模块内存中。多个线程共享同一块内存时，这些线程局部数据区域会重叠，导致变量互相覆盖。
2. TLS初始化错误：__wasm_init_tls函数需要为每个线程分配独立的内存区域存储TLS数据，在共享内存上调用该函数会导致所有线程的TLS数据存放在同一块内存，完全失去线程隔离的作用。
3. 实例内存混淆：每个WASM实例需要独立的内存维护自身执行环境，仅业务数据需要放在共享内存中，代码混淆了这两者的用途。

修复方案 #

1. 为每个Worker分配独立的WASM内存 #

Worker的WASM执行内存（栈、TLS）必须独立，仅把业务共享内存单独导入：

self.start = async (workerID, sharedMemory, world) => {
    // 创建当前Worker专属的WASM内存，用于栈和TLS
    const wasmMemory = new WebAssembly.Memory({
        initial: 23,
        maximum: 23,
    });

    const wasm = await WebAssembly.instantiateStreaming(wasmSource, {
        env: { 
            memory: wasmMemory,
            // 单独导入共享内存供业务逻辑使用
            shared_memory: sharedMemory
        },
        imports: {
            abort: (messagePtr, locationPtr, row, column) => {
                throw new Error(`错误 (${row}:${column}, 线程 ${workerID})`)
            },
            _log_num: num => console.log(`线程 ${workerID}: n 是 ${num}`),
        },
    })

    // 现在初始化TLS用的是独立内存，不会和其他线程冲突
    wasm.instance.exports.__wasm_init_tls(workerID-1)   

    let lastProcessedTick = 0
    while (1) {
        Atomics.wait(world.globalTick, 0, lastProcessedTick)
        lastProcessedTick = world.globalTick[0]
        wasm.instance.exports.run(workerID)
    }
}

2. 调整主线程的共享内存用途 #

共享内存仅用于存储业务需要共享的数据（比如globalTick），不需要分配过大空间：

// 共享内存仅用于业务共享数据
const sharedMemory = new WebAssembly.Memory({
    initial: 1,
    maximum: 1,
    shared: true,
})

const world = {
    globalTick: new Int32Array(sharedMemory.buffer, 0, 1),
}

const startAWorkerCore = coreIndex => {
    const worker = new Worker('worker/sim.mjs', {type:'module'})
    ;['start', coreIndex+1, sharedMemory, world].forEach(arg => worker.postMessage(arg))
}

3. （可选）Rust中访问共享内存 #

如果业务逻辑需要操作共享内存，在Rust中导入共享内存：

mod js {
    #[link(wasm_import_module = "imports")]
    extern "C" {
        pub fn abort(msgPtr: usize, filePtr: usize, line: u32, column: u32) -> !;
        pub fn _log_num(number: usize);
    }

    #[link(wasm_import_module = "env")]
    extern "C" {
        #[link_name = "shared_memory"]
        pub static SHARED_MEMORY: *mut u8;
    }
}

验证 #

修改后重新编译Rust代码，运行主线程代码，三个Worker会各自输出对应区间的数值，并行运行时不会再出现变量共享覆盖的问题。

内容的提问来源于stack exchange，提问作者DDR