Add jspi-hooks pass

[guybedford]

This implements a --jspi-hooks pass that wraps the JSPI entry/exit/suspend/resume boundaries of a module with calls to corresponding lifecycle hooks the module itself provides.

This is needed in Binaryen itself because it cannot be done at another layer of the system - suspension incurs a microtask, so that any JS wrappers around the suspend JS function will not be correct, and any exit wrapper on the WebAssembly.promising itself will also not have the correct timing.

Only by directly integrating the synchronous timing of the hooks into the Wasm can we get sound tracking of JSPI context switching.

The use cases are various here:

• Supporting stack save/restore for reentrant JSPI
• Checking if a suspension happened or not for a given fiber
• Giving JSPI "fibers" a unique identity
• Any async context associations associated with that ID
• Thread-local context switching if appropriate

The alternative to this approach would be to more heavily define JSPI instrumentation:

• Assuming the fiber identity model and creating our own ID system
• Assuming the stack switching model and handling that automatically instead of through hooks.

Given where JSPI is in its ecosystem adoption as an emerging convention, a general hooks-based approach feels like the best way to start on these problems for now.

The hooks are four function exports which must be provided by the binary for the transform to work:

__jspi_enter:   [] -> [i64 token]
__jspi_exit:    [i64 token, i32 error] -> []
__jspi_suspend: [] -> [i64 token]
__jspi_resume:  [i64 token, i32 error] -> []

They are called on every JSPI enter/exit/suspend/resume operation:

• token: opaque i64 - the wrapper passes the value the before hook returned to the matching after hook (enter→exit, suspend→resume) through a wasm local, which JSPI preserves across the suspension; the pass never interprets it. This is what supports passing e.g. fiber IDs / stack identifiers across the pairs.
• error: 1 on the after hooks when the wrapped call threw, else 0.

Then similar to the lines of asyncify, the PR here implements for Binaryen that contract as a transform pass:

• --pass-arg=jspi-exports@<patterns>: each matching export is retargeted to a wrapper calling __jspi_enter and __jspi_exit around it.
• --pass-arg=jspi-imports@<patterns>: each matching import is moved to a new import and the original function object becomes a wrapper calling __jspi_suspend and __jspi_resume and passing the token between them, so every existing use (calls, ref.func, element segments, exports) reaches the wrapper with no reference rewriting.
• --pass-arg=jspi-dyncalls (plus jspi-dyncall-sigs@ii,vi for extra signatures) also exports a wrapped __jspi_dyncall_<sig>(fptr, ...) trampoline around call_indirect per function signature in the table, so hosts that make function pointers promising (Emscripten's promising dynCall, embind async()) do not bypass the hooks. <sig> is the getSig() alphabet.

When instrumented:

• The wrappers are locals only and the pass enables EH and reference types so unwinds are fully supported.
• Every JSPI boundary crossing gets two extra direct calls, performance impact is expected to be negligible.
• Both error paths go through the hooks: a try_table (catch_all_ref) calls the after hooks with error=1, then throw_refs the exception unchanged, whatever its tag, so the pass needs no knowledge of any tag and no imports. Modules already using legacy EH get the legacy try/catch_all/rethrow form, since engines reject mixing the two; a module that already mixes them is rejected with a pointer to --translate-to-exnref.
• Multivalue results are supported (tuple local); pattern lists take *, comma/newline separators and @file response files like asyncify.

The actual behaviors of the JSPI system are otherwise entirely left to the hook runtime implementations, with Emscripten expected to support a C runtime implementation for interfacing with JSPI hooks and also using it to build REENTRANT_JSPI (PR which is based on this).

Due to the generality though, other runtimes could implement their own custom systems similarly.

Tests: lit coverage of import wrapping through direct calls, ref.func, element segments and exported imports; export retargeting with internal callers untouched; hook exclusion; features; response files; both EH forms; multivalue; trampolines with explicit signatures; the Fatal cases; --roundtrip/-O2 stability; and an engine-level test/lit/d8 run of the wrapped module under real JSPI (V8) checking the full event trace, exception identity through the rethrow, SuspendError reaching the resume hook, id forwarding across a real suspension, nested promising entry from a plain import, and a trampoline call.

Made with AI assistance under my review

[tlively]

Alternatively, we could also have 4 separate functions, but I wanted to try and keep the contract minimal here.

Using separate functions seems like it would be better for follow-up optimizations and seems like it might be simpler in the end than using opaque operation codes.

---

Stepping back, is the idea that this would be used for "advanced" JSPI usage like for fibers and multi-stack reentrancy? It seems that these use cases all boil down to some version of green threads. Will "plain" JSPI usage continue to not need this pass?

Would it simplify things at all if we can assume that there is only a single, blessed JSPI suspension import? All other "suspending" imports could simply take the returned promise as an externref and then pass it to the single suspend import to do the actual suspension.

cc @sbc100.

[guybedford]

Thanks for the feedback @tlively I've updated the PR to use 4 separate hook functions instead.

Stepping back, is the idea that this would be used for "advanced" JSPI usage like for fibers and multi-stack reentrancy? It seems that these use cases all boil down to some version of green threads. Will "plain" JSPI usage continue to not need this pass?

Exactly, in particular for Cloudflare we need to support reentrant JSPI with multiple Tokio runtimes via JSPI-scoped thread locals in the same Wasm instance so that durable objects can serve multiple IO-segmented requests. So yes exactly supporting green thread primitives or otherwise via these hooks.

Would it simplify things at all if we can assume that there is only a single, blessed JSPI suspension import? All other "suspending" imports could simply take the returned promise as an externref and then pass it to the single suspend import to do the actual suspension.

It might reduce the import interface at the cost of the double call and indirection. But imports are already a short list, and the problem is exports listing is the long one that would still be required so it wouldn't really fix the whole problem I don't think.

[tlively]

JSPI-scoped thread locals

Are these expected to literally be thread_local variables, or accessed via some new API? The former would require leaning all the way into each async task being modeled as a separate pthread, I think. That would be an interesting direction, but would require a significant amount of up-front design. For example, how would that interact with normal worker-backed threads? What changes would be necessary to support inter-task synchronization via mutexes and similar on the same browser thread?

[guybedford]

No plans to literally remap all thread locals. I suspect we'd want to introduce a new Rust macro utility like jspi_local! that would be enabled by these hooks that would then support arbitrary JSPI boundaries. This would be very useful for library and user code.

[sbc100]

How does this relate to the co-operative threading ABI that WASI is proposing: https://github.com/WebAssembly/wasi-sdk/blob/main/CoopThreading.md. Could JSPI piggyback on this instead of creating a separate standard?

[guybedford]

How does this relate to the co-operative threading ABI that WASI is proposing: https://github.com/WebAssembly/wasi-sdk/blob/main/CoopThreading.md. Could JSPI piggyback on this instead of creating a separate standard?

Very interesting, I hadn't seen that, I did an AI guided deep-dive on the topic and got it to the following response which seems to match my own intuition in that Jco might benefit from the same hooks:

_The coop-threading ABI already has the same shape on the export side: wit-component wraps every lifted export/callback/dtor with an in-wasm __wasm_task_hook(kind) call, and wasi-libc's hook allocates the task's stack and TLS and returns the SP to install — that's __jspi_enter/_jspi_exit. What it doesn't have is a suspend/resume hook around blocking calls, because the engine holds the SP/TLS in per-thread context slots that survive a block. JSPI has no such slots; the SP/TLS cells are wasm globals, so the switch has to happen in-fiber at the suspend/resume boundary with the identity carried in a wasm local. That's the whole delta: the import half plus the token (and an error flag, since JS/C++ exceptions cross the wrapper).

So I don't see this as a separate standard — the hook events and the runtime policy (stack+TLS per task, TLS kept across suspension, freed at exit) are the same, and I'd want the kinds aligned so one libc implementation serves both. The codegen half of the coop ABI (libcall-thread-context) doesn't help under JSPI: its purpose is to let the engine own the cell, and here the global already is the cell.

The reverse also holds: a component hosted on JSPI needs the import-half hook too. After a resume the first thing a p3 task does is context.get 0 for its stack pointer, and a JS-side "current task" register can't be updated soundly at resume time (two fibers' continuations can run before either resume job). jco already runs core modules through wasm-opt (its asyncify mode did exactly this), so it could apply the same pass — or wit-component's fixup could grow the import half natively, which is probably the right long-term home for the CM. Either way the hook contract is what's worth sharing; the instrumenter can differ per toolchain.

Of course, having feedback from @vados-cosmonic would help here to figure this out further too.

[tlively]

I was able to convince myself that we need cleanup hooks on at least the export side of things. When Wasm returns from a JSPI'd export, the export promise is resolved with its return value. But then any cleanup logic chained onto that promise is pushed to the end of the microtask queue, so it could be preempted by another task already on the microtask queue. Such a task could call into Wasm and observe the not-yet-cleaned-up global state from the completed task. There's no other way to clean up after an export returns besides by chaining the cleanup logic onto its promise.

If jco plans to lower the component model's thread.index builtin to a Wasm global, it seems it would have a similar problem. cc @alexcrichton regarding that.

It should still be possible to use Wasm wrapper functions around the original exports, but those would have to be generated offline, in which case they are no different than the hooks approach, or would have to be generated dynamically based on the types of the exports, which sounds unpleasant and would not scale to Wasm GC.

[vados-cosmonic]

Hi @tlively so our co-op threads implementation (thread.index) is actually still under development (@TartanLlama is actually taking point here), so we still have a choice on that. Right now most of the async task/state machinery is mostly not registered as Wasm globals but rather in the host bindgen machinery (if I'm understanding what you're noting correctly).

I'm not sure that we need this hooks pass to implement co-op threads on the Component Model side, because of the use of context builtins in particular... I'm not sure this is relevant there, but I'm still digesting this so I could be wrong.

[guybedford]

@vados-cosmonic the thinking here is that under JSPI, the Wasm stack and the shadow LLVM stack need to be kept in sync since JSPI will just switch the Wasm stack not the shadow LLVM stack. And the way to ensure they are correct is to always instrument all functions that are directly exposed in JSPI exports or imports. That instrumentation can either be on the bindings side for the system, or equivalently via something like these hooks. But because this stack alignment problem is unique to JS, I would be surprised if it just worked with reentrant JSPI without some extra machinery being needed, unless the component model bindings already do a full stack save and restore of the linear memory shadow stack around a JSPI suspension.

[vados-cosmonic]

Thanks for the explanation Guy, that makes it a bit easier to work out the concern here -- specifically dealing with the guest shadow stack, one thing missing from the discussion here is the work that TartanLlama & others already added to LLVM to enable upstream support for this:

llvm/llvm-project@577e9a7
llvm/llvm-project#175800

This is what I was referring to by the note about context.get/set, I had to look it up since I'm certainly not an expert in the wasi-libc implementation but the previously mutable single stack pointer becomes calls to __wasm_get_stack_pointer and __wasm_set_stack_pointer calls (these get mapped to context.get/set) and uses "TLS" (that Jco will manage post transpilation in the shims) to set state back up.

For example in the new Rust target wasm32-wasip3 which makes use of the new LLVM release:

The linear memory shadow stack pointer is stored in a component model task context slot instead of a WebAssembly global. Additionally the base pointer of TLS is managed differently than other targets. These changes are made to enable cooperative multithreading on this target.

IF I'm understanding right, as long as the version of LLVM is recent enough, we don't need the extra hooks here on the CM side (i.e. a recent-enough LLVM will be using the wasm builtins not __stack_pointer), and we actually have support right now in C/C++, Rust, and Python for this paradigm (we're going through the testing now and there's certainly more testing to be done but the JS host is the straggler here), so the other host & guest implementations could serve instructive for anyone interested.

However, in the case of Wasm w/ mutable stack pointers, IIRC the rewrite approach is what is currently done by wit-component BUT it isn't enough for the JSPI context as you were saying -- one reason being it only does it on exports right now, but not imports (and any import could be JSPI wrapped).

There's one new option though I think isn't represented here, which is adding the relevant configuration & filling out rewriting in wit-component to do this kind of work, but that doesn't help toolchains that don't include wit-component. Similarly, you can't take advantage of Jco's stack pointer management via context.get/set if you're not using Jco to do the transpilation, so obviously other ecosystems would not benefit.

So AFAICT we do need something like jspi-hooks in the mutable stack pointer case.

But because this stack alignment problem is unique to JS, I would be surprised if it just worked with reentrant JSPI without some extra machinery being needed, unless the component model bindings already do a full stack save and restore of the linear memory shadow stack around a JSPI suspension.

So essentially, under P3 yes, this is what happens and it's up to the guest toolchain. BUT for the the mutable stack pointer case, yeah it's not going to work out of the box.