Hacky tweaks to get Bevy closer to Wasm threading support

crates/bevy_asset/src/asset_server.rs

@@ -449,7 +449,7 @@ impl AssetServer {
         let server = self.clone();
         let owned_path = asset_path.to_owned();
         IoTaskPool::get()
-            .spawn(async move {
+            .spawn_local(async move {
                 if let Err(err) = server.load_async(owned_path, force).await {
                     warn!("{}", err);
                 }

crates/bevy_tasks/Cargo.toml

@@ -17,6 +17,9 @@ once_cell = "1.7"
 concurrent-queue = "2.0.0"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
+wasm-bindgen = { version = "0.2" }
+web-sys = { version = "0.3", features = ["Window", "Worker", "WorkerOptions", "WorkerType", "Navigator"] }
+js-sys = "0.3"
 wasm-bindgen-futures = "0.4"
 
 [dev-dependencies]

crates/bevy_tasks/src/lib.rs

@@ -7,24 +7,29 @@ pub use slice::{ParallelSlice, ParallelSliceMut};
 mod task;
 pub use task::Task;
 
-#[cfg(not(target_arch = "wasm32"))]
+//#[cfg(not(target_arch = "wasm32"))]
 mod task_pool;
-#[cfg(not(target_arch = "wasm32"))]
+//#[cfg(not(target_arch = "wasm32"))]
 pub use task_pool::{Scope, TaskPool, TaskPoolBuilder};
 
+#[cfg(target_arch = "wasm32")]
+pub(crate) mod wasm_worker;
+
+/*
 #[cfg(target_arch = "wasm32")]
 mod single_threaded_task_pool;
 #[cfg(target_arch = "wasm32")]
 pub use single_threaded_task_pool::{Scope, TaskPool, TaskPoolBuilder, ThreadExecutor};
+*/
 
 mod usages;
-#[cfg(not(target_arch = "wasm32"))]
+//#[cfg(not(target_arch = "wasm32"))]
 pub use usages::tick_global_task_pools_on_main_thread;
 pub use usages::{AsyncComputeTaskPool, ComputeTaskPool, IoTaskPool};
 
-#[cfg(not(target_arch = "wasm32"))]
+//#[cfg(not(target_arch = "wasm32"))]
 mod thread_executor;
-#[cfg(not(target_arch = "wasm32"))]
+//#[cfg(not(target_arch = "wasm32"))]
 pub use thread_executor::{ThreadExecutor, ThreadExecutorTicker};
 
 mod iter;
@@ -49,7 +54,13 @@ use std::num::NonZeroUsize;
 ///
 /// This will always return at least 1.
 pub fn available_parallelism() -> usize {
-    std::thread::available_parallelism()
+    #[cfg(not(target_arch = "wasm32"))]
+    return std::thread::available_parallelism()
         .map(NonZeroUsize::get)
-        .unwrap_or(1)
+        .unwrap_or(1);
+    #[cfg(target_arch = "wasm32")]
+    return web_sys::window()
+        .expect("Missing Window")
+        .navigator()
+        .hardware_concurrency() as _;
 }

crates/bevy_tasks/src/task_pool.rs

@@ -105,7 +105,10 @@ pub struct TaskPool {
     executor: Arc<async_executor::Executor<'static>>,
 
     /// Inner state of the pool
+    #[cfg(not(target_arch = "wasm32"))]
     threads: Vec<JoinHandle<()>>,
+    #[cfg(target_arch = "wasm32")]
+    threads: Vec<crate::wasm_worker::WebWorkerJoinHandle>,
     shutdown_tx: async_channel::Sender<()>,
 }
 
@@ -133,6 +136,7 @@ impl TaskPool {
         let num_threads = builder
             .num_threads
             .unwrap_or_else(crate::available_parallelism);
+        println!("NUM THREADS: {:?}", num_threads);
 
         let threads = (0..num_threads)
             .map(|i| {
@@ -144,8 +148,14 @@ impl TaskPool {
                 } else {
                     format!("TaskPool ({i})")
                 };
+
+                #[cfg(not(target_arch = "wasm32"))]
                 let mut thread_builder = thread::Builder::new().name(thread_name);
 
+                #[cfg(target_arch = "wasm32")]
+                let mut thread_builder =
+                    crate::wasm_worker::WasmWorkerBuilder::new().name(thread_name);
+
                 if let Some(stack_size) = builder.stack_size {
                     thread_builder = thread_builder.stack_size(stack_size);
                 }

crates/bevy_tasks/src/usages.rs

@@ -116,7 +116,7 @@ impl Deref for IoTaskPool {
 /// # Warning
 ///
 /// This function *must* be called on the main thread, or the task pools will not be updated appropriately.
-#[cfg(not(target_arch = "wasm32"))]
+//#[cfg(not(target_arch = "wasm32"))]
 pub fn tick_global_task_pools_on_main_thread() {
     COMPUTE_TASK_POOL
         .get()

crates/bevy_tasks/src/wasm_worker.rs

@@ -0,0 +1,56 @@
+use wasm_bindgen::prelude::*;
+
+pub struct WasmWorkerBuilder {
+    stack_size: usize,
+}
+
+impl WasmWorkerBuilder {
+    pub fn new() -> Self {
+        Self { stack_size: 65536 }
+    }
+
+    pub fn name(self, _name: String) -> Self {
+        self
+    }
+
+    pub fn stack_size(self, size: usize) -> Self {
+        Self { stack_size: size }
+    }
+
+    pub fn spawn<F: FnOnce() + Send + 'static>(self, f: F) -> Result<WebWorkerJoinHandle, JsValue> {
+        let worker = web_sys::Worker::new_with_options(
+            "./worker.js",
+            web_sys::WorkerOptions::new().type_(web_sys::WorkerType::Module),
+        )?;
+        // Double-boxing because `dyn FnOnce` is unsized and so `Box<dyn FnOnce()>` is a fat pointer.
+        // But `Box<Box<dyn FnOnce()>>` is just a plain pointer, and since wasm has 32-bit pointers,
+        // we can cast it to a `u32` and back.
+        let ptr = Box::into_raw(Box::new(Box::new(f) as Box<dyn FnOnce()>));
+        let msg = js_sys::Array::new();
+
+        msg.push(&wasm_bindgen::module());
+        msg.push(&wasm_bindgen::memory());
+
+        // Send the address of the closure to execute.
+        msg.push(&JsValue::from(ptr as u32));
+        worker.post_message(&msg);
+        Ok(WebWorkerJoinHandle {})
+    }
+}
+
+#[derive(Debug)]
+pub struct WebWorkerJoinHandle {}
+
+impl WebWorkerJoinHandle {
+    pub fn join(self) -> Result<(), ()> {
+        Ok(())
+    }
+}
+
+#[wasm_bindgen]
+/// This function is here for `worker.js` to call.
+pub fn bevy_worker_entry_point(addr: u32) {
+    // Interpret the address we were given as a pointer to a closure to call.
+    let closure = unsafe { Box::from_raw(addr as *mut Box<dyn FnOnce()>) };
+    (*closure)();
+}

examples/wasm/worker.js

@@ -0,0 +1,11 @@
+import wasm_bindgen from "./target/wasm_example.js";
+
+self.onmessage = async event => {
+  const initialized = await wasm_bindgen(
+    event.data[0], // Module
+    event.data[1]  // Memoryx
+  );
+
+  console.log("INITIALIZED: ", initialized);
+  initialized.bevy_worker_entry_point(Number(event.data[2]));
+}

run.sh

@@ -0,0 +1,6 @@
+RUSTFLAGS='-C target-feature=+atomics,+bulk-memory,+mutable-globals' \
+    cargo build --example breakout --target wasm32-unknown-unknown -Z build-std=std,panic_abort
+wasm-bindgen --out-name wasm_example \
+  --out-dir examples/wasm/target \
+  --target web target/wasm32-unknown-unknown/debug/examples/breakout.wasm
+devserver --header Cross-Origin-Opener-Policy='same-origin' --header Cross-Origin-Embedder-Policy='require-corp' --path examples/wasm