feat: JSPI support for wasm32-unknown-emscripten

.github/workflows/ci.yml

@@ -1202,12 +1202,19 @@ jobs:
         env:
           RUSTFLAGS: ""
 
-      - name: Test tokio for emscripten
+      - name: Test tokio for emscripten (JSPI)
         run: cargo test -p tokio --target wasm32-unknown-emscripten --features "rt,time,sync,macros,io-util,test-util" --tests
         working-directory: tokio
         env:
           CARGO_TARGET_WASM32_UNKNOWN_EMSCRIPTEN_RUNNER: node
-          RUSTFLAGS: "-C link-args=-sALLOW_MEMORY_GROWTH=1 -C link-args=-sEXIT_RUNTIME=1 -C link-args=-sSTACK_SIZE=1048576"
+          RUSTFLAGS: "-C link-args=-sALLOW_MEMORY_GROWTH=1 -C link-args=-sEXIT_RUNTIME=1 -C link-args=-sJSPI -C link-args=-sSTACK_SIZE=1048576"
+
+      - name: Test tokio for emscripten (non-JSPI)
+        run: cargo test -p tokio --target wasm32-unknown-emscripten --features "rt,time,sync,macros,io-util,test-util" --test rt_emscripten_block_on
+        working-directory: tokio
+        env:
+          CARGO_TARGET_WASM32_UNKNOWN_EMSCRIPTEN_RUNNER: node
+          RUSTFLAGS: "-C link-args=-sALLOW_MEMORY_GROWTH=1 -C link-args=-sEXIT_RUNTIME=1"
 
   check-external-types:
     name: check-external-types (${{ matrix.os }})

spellcheck.dic

@@ -1,4 +1,4 @@
-327
+330
 &
 +
 <
@@ -162,6 +162,8 @@ IP
 IPv4
 IPv6
 iteratively
+JS
+JSPI
 Kotlin's
 latencies
 Lauck
@@ -266,6 +268,7 @@ subfield
 suboptimal
 subprocess
 superset
+suspendable
 symlink
 symlinks
 sys

tokio-macros/src/entry.rs

@@ -521,12 +521,57 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
 
     };
 
+    let output_type = match &input.sig.output {
+        // For functions with no return value syn doesn't print anything,
+        // but that doesn't work as `Output` for our boxed `Future`, so
+        // default to `()` (the same type as the function output).
+        syn::ReturnType::Default => quote! { () },
+        syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
+    };
+    let raw_body = input.body();
+    // JSPI is enabled for the module, but suspending imports are valid only
+    // from a promising-wrapped activation. The guard marks libtest's known
+    // promising activation so its `block_on` may suspend; other entries fail
+    // at the park leaf instead of trapping in the engine.
+    let emscripten_test_block = quote_spanned! {last_stmt_end_span=>
+        #[allow(clippy::expect_used, clippy::diverging_sub_expression, clippy::needless_return, clippy::unwrap_in_result)]
+        {
+            let _suspend_guard = if #crate_path::runtime::jspi_enabled() {
+                ::core::option::Option::Some(#crate_path::runtime::SuspendGuard::new())
+            } else {
+                ::core::option::Option::None
+            };
+
+            let body = async #raw_body;
+            #crate_path::pin!(body);
+            let body: ::core::pin::Pin<&mut dyn ::core::future::Future<Output = #output_type>> = body;
+
+            // Inner block: keep the `use` from shadowing names in the
+            // user body, which shares this scope.
+            let rt = {
+                #use_builder
+
+                #rt
+                    .enable_all()
+                    .#build
+                    .expect("Failed building the Runtime")
+            };
+            return rt.block_on(body);
+        }
+    };
+
     let last_block = match config.flavor {
         RuntimeFlavor::Threaded => quote! {
             #crate_path::__tokio_unsupported_multi_thread_on_emscripten! {
                 #native_last_block
             }
         },
+        _ if is_test => quote! {
+            #[cfg(not(target_os = "emscripten"))]
+            #native_last_block
+            #[cfg(target_os = "emscripten")]
+            #emscripten_test_block
+        },
         _ => native_last_block,
     };
 
@@ -541,18 +586,15 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
     //
     // We don't do this for the main function as it should only be used once so
     // there will be no benefit.
-    let output_type = match &input.sig.output {
-        // For functions with no return value syn doesn't print anything,
-        // but that doesn't work as `Output` for our boxed `Future`, so
-        // default to `()` (the same type as the function output).
-        syn::ReturnType::Default => quote! { () },
-        syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
-    };
-
     let body = if is_test {
+        // On Emscripten the body construction lives inside the capture-free
+        // closure (see `emscripten_test_block`).
         quote! {
+            #[cfg(not(target_os = "emscripten"))]
             let body = async #body;
+            #[cfg(not(target_os = "emscripten"))]
             #crate_path::pin!(body);
+            #[cfg(not(target_os = "emscripten"))]
             let body: ::core::pin::Pin<&mut dyn ::core::future::Future<Output = #output_type>> = body;
         }
     } else {

tokio/src/lib.rs

@@ -451,6 +451,11 @@
 //! `test-util` features. The `net`, `process`, `signal`, and `rt-multi-thread`
 //! features are not supported.
 //!
+//! With [JSPI], `#[tokio::test]` can suspend on the host event loop while
+//! waiting for timers. Without JSPI, a wait that cannot progress panics.
+//!
+//! [JSPI]: https://github.com/WebAssembly/js-promise-integration
+//!
 //! ## Unstable `WASM` support
 //!
 //! Tokio also has unstable support for some additional `WASM` features. This
@@ -489,6 +494,16 @@ compile_error!("Only features sync,macros,io-util,rt,time are supported on wasm.
 #[cfg(all(target_os = "emscripten", feature = "process"))]
 compile_error!("The `process` feature is not supported on wasm32-unknown-emscripten.");
 
+// We currently only support the single-threaded Emscripten runtime, with
+// support for JSPI.
+// A `-pthread` (atomics) build breaks its assumptions, until the kernel is made
+// thread-aware.
+#[cfg(all(target_os = "emscripten", feature = "rt", target_feature = "atomics"))]
+compile_error!(
+    "Tokio's `wasm32-unknown-emscripten` runtime is single-threaded and does \
+ not support `-pthread` (atomics) builds."
+);
+
 #[cfg(all(not(tokio_unstable), feature = "io-uring"))]
 compile_error!("The `io-uring` feature requires `--cfg tokio_unstable`.");
 

tokio/src/runtime/jspi.rs

@@ -0,0 +1,101 @@
+//! Minimal JSPI primitives for `wasm32-unknown-emscripten`.
+//!
+//! Tokio's suspension model is strictly non-reentrant: `#[tokio::test]` claims
+//! suspension for its whole promising activation with a [`SuspendGuard`],
+//! and [`sleep`] — the one suspending import the runtime issues — parks
+//! that activation on a host timer. With never more than one suspension in
+//! flight, no spill-stack save/restore is needed: any Wasm re-entered
+//! while the activation is parked completes before the resume, leaving the
+//! spill stack above the suspended frame untouched.
+
+use std::cell::Cell;
+use std::time::Duration;
+
+thread_local! {
+    static SUSPENDABLE: Cell<bool> = const { Cell::new(false) };
+}
+
+/// Marks the `#[tokio::test]` promising activation as suspendable for the
+/// body's extent. JSPI availability is module-wide, but only an activation
+/// entered through `WebAssembly.promising` may call a suspending import; this
+/// guard records that per-activation capability. Internal to the test
+/// expansion, not a user convention.
+#[derive(Debug)]
+pub struct SuspendGuard(());
+
+impl SuspendGuard {
+    /// Marks the current activation suspendable until drop.
+    #[allow(clippy::new_without_default)]
+    pub fn new() -> SuspendGuard {
+        SUSPENDABLE.set(true);
+        SuspendGuard(())
+    }
+}
+
+impl Drop for SuspendGuard {
+    fn drop(&mut self) {
+        SUSPENDABLE.set(false);
+    }
+}
+
+/// Whether the park leaf may suspend: a [`SuspendGuard`] is live.
+pub(crate) fn can_suspend() -> bool {
+    SUSPENDABLE.get()
+}
+
+// Emscripten EM_JS convention: `__em_js__<name>` data exports carry JS
+// bodies into the objects, and `__asyncjs__` names get
+// `WebAssembly.Suspending` treatment under `-sJSPI`. The static must be
+// referenced from linked code (`anchor`) so its archive member is pulled
+// in.
+const TOKIO_JSPI_SLEEP: &str = "(ms)<::>{ return Asyncify.handleAsync(async () => { await new Promise((r) => setTimeout(r, ms)); }); }";
+
+const fn em_js<const N: usize>(s: &str) -> [u8; N] {
+    // NUL-terminated: N == s.len() + 1
+    let mut a = [0u8; N];
+    let b = s.as_bytes();
+    let mut i = 0;
+    while i < b.len() {
+        a[i] = b[i];
+        i += 1;
+    }
+    a
+}
+
+#[allow(non_upper_case_globals)]
+#[no_mangle]
+#[used]
+static __em_js____asyncjs__tokio_jspi_sleep: [u8; TOKIO_JSPI_SLEEP.len() + 1] =
+    em_js(TOKIO_JSPI_SLEEP);
+
+unsafe extern "C" {
+    /// Reports the `ASYNCIFY` build mode: 0 = none, 1 = legacy `Asyncify`,
+    /// 2 = JSPI. Only mode 2 supports Tokio's JSPI import.
+    safe fn emscripten_has_asyncify() -> i32;
+}
+
+// Suspending import: parks on a host timeout. Unit return, never rejects,
+// `Asyncify.handleAsync` keeps the runtime alive across the suspension.
+#[link(wasm_import_module = "env")]
+unsafe extern "C-unwind" {
+    #[link_name = "__asyncjs__tokio_jspi_sleep"]
+    safe fn tokio_jspi_sleep_import(ms: f64);
+}
+
+#[inline(never)]
+fn ensure_jspi_sleep_linked() {
+    // `#[used]` retains the data in its object; this reference also causes
+    // the archive member containing the EM_JS body to be linked.
+    std::hint::black_box(__em_js____asyncjs__tokio_jspi_sleep.as_ptr());
+}
+
+/// Whether JSPI suspension is available: linked with `-sJSPI`.
+pub fn jspi_enabled() -> bool {
+    emscripten_has_asyncify() == 2
+}
+
+/// Suspend the owning activation for `dur` on a host timer.
+pub(crate) fn sleep(dur: Duration) {
+    ensure_jspi_sleep_linked();
+    tokio_jspi_sleep_import(dur.as_secs_f64() * 1000.0);
+}

tokio/src/runtime/mod.rs

@@ -416,6 +416,16 @@ mod tests;
 
 pub(crate) mod context;
 
+#[cfg(all(target_os = "emscripten", feature = "rt"))]
+pub(crate) mod jspi;
+
+// Used by the `#[tokio::test]` expansion on Emscripten: the guard claims
+// JSPI suspension for the test's promising activation so the body's
+// `block_on` can suspend on the host loop. Not public API.
+#[cfg(all(target_os = "emscripten", feature = "rt", feature = "test-util"))]
+#[doc(hidden)]
+pub use jspi::{jspi_enabled, SuspendGuard};
+
 pub(crate) mod park;
 
 pub(crate) mod driver;
@@ -550,6 +560,8 @@ cfg_rt! {
     }
 
     cfg_fs! {
+        // Emscripten's `fs` uses the inline blocking shim, not the pool.
+        #[cfg_attr(target_os = "emscripten", allow(unused_imports))]
         pub(crate) use blocking::spawn_mandatory_blocking;
     }