wasm: poll VMTraps at loop back-edges so pure-Wasm loops can be terminated - #372
wasm: poll VMTraps at loop back-edges so pure-Wasm loops can be terminated#372robobun wants to merge 1 commit into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (10)
WalkthroughThe change adds VM-trap polling at WebAssembly loop back-edges for IPInt, BBQJIT, and OMGJIT. Trap handling resumes execution when possible and propagates termination through existing exception paths. ChangesWebAssembly loop VM-trap handling
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
Comment |
Preview Builds
|
There was a problem hiding this comment.
Thanks for addressing the comment-wording nit in bff0b81. I didn't find functional issues, but this touches hot-path codegen across all three Wasm tiers (IPInt dispatch-slot sizing, BBQ late-path + probe, OMG B3 patchpoint effects/clobbers), so it warrants a human reviewer familiar with the Wasm JIT.
What was reviewed:
- IPInt
_looppoll: confirmed the addedbpaeq+jmpmirrors thecheckStackOverflowpattern and the out-of-slot trampoline restores PC/MC before advancing. ipint_extern_handle_vm_traps_at_loop/operationWasmHandleTrapsAtLoop: match the existingcheck_stack_and_vm_trapstrap-servicing shape.- OMG: the slow-path patchpoint is now confined to a Rare block so the hot path carries no clobbers;
exitsSideways+reads=toplooks right for the throw edge.
Extended reasoning...
Overview
The PR adds a trap-aware stack-limit poll at every Wasm loop head so that VM::notifyNeedTermination() (e.g. worker.terminate(), watchdog) can preempt pure-Wasm loops that never re-enter JS. It spans all three execution tiers: IPInt (offlineasm _loop opcode + new out-of-slot trampoline + C++ slow path), BBQ (addLoop late-path with jit.probe), and OMG (addLoop B3 Load/Branch into a Rare block containing a PatchpointValue that probes and either throws or resumes). Supporting changes expose Mirror::offsetOfTrapAwareSoftStackLimit() and add operationWasmHandleTrapsAtLoop(Probe::Context&).
Since my earlier inline nit, commit bff0b81 reworded the ipint_loop_check_vm_traps comment, moved the OMG clobber set into a dedicated Rare block (so the hot path is a plain B3 Load+Above+Branch with no clobbers), flipped BBQ to use sp as the left comparand, and pointed the BBQ throw edge at the shared recordJumpToThrowException sink.
Security risks
None identified. The change adds a read of an existing per-instance field and a rarely-taken slow path that calls handleTrapsIfNeeded() — the same routine already reachable from the IPInt prologue. No new attacker-controlled inputs, no auth/crypto/permission surfaces.
Level of scrutiny
High. This is hot-path JIT code generation that runs on every Wasm loop back-edge across three tiers, each with tier-specific correctness constraints:
- IPInt: fixed-size aligned dispatch slots (
alignIPInt) — the addedbpaeq+jmpmust not overflow the_loopslot on any target. The out-of-slotop(...)trampoline sidesteps this, but the in-slot delta still needs a size check per architecture. - BBQ: the late-path uses
jit.probeto preserve live state, then readsnonPreservedNonArgumentGPR0(asserted== wasmScratchGPR) after the probe restores everything else. Register discipline here interacts with BBQ's own scratch-register conventions. - OMG: B3 effects modeling (
exitsSideways,reads = top, clobber set ofmacroClobberedGPRs+nonPreservedNonArgumentGPR0) must be sufficient for the probe +emitExceptionChecksequence. The generator lambda capturesthis(theOMGIRGenerator), which relies on the generator running before the IR generator is destroyed.
Any of these being subtly wrong would produce miscompiles or register corruption that only manifest under specific register-pressure / trap-timing conditions — exactly the kind of thing that needs eyes from someone who owns these tiers.
Other factors
- The bug-hunting pass found no functional issues; the one prior finding was a comment-wording nit, now addressed.
- The PR description includes a concrete before/after repro and notes hot-path cost matches JS
op_check_traps. - No test is added in-tree; verification is via the repro script and preview builds.
- The commit message on bff0b81 references "review" feedback beyond my nit (OMG rare-block confinement, sp-as-left, shared throw sink), suggesting there is already an ongoing human review conversation that should continue.
Given the breadth (three JIT tiers), the subtlety of B3 effects/clobber declarations, and the per-back-edge performance implications, this should be signed off by a human familiar with the Wasm JIT rather than auto-approved.
…nated
A Worker running a WebAssembly function whose body is a tight loop with no
JS re-entry (e.g. loop { br 0 }) could not be preempted by
worker.terminate(): VM::notifyNeedTermination() poisons
m_trapAwareSoftStackLimit via StackManager::requestStop(), but the only
Wasm-side consumer was the IPInt function prologue, so a pure-Wasm loop
never observed the request.
Add a trap-aware stack-limit poll at each Wasm loop head:
- IPInt: compare sp against m_trapAwareSoftStackLimit after ipintLoopOSR
(sp as first comparand: ARM64 cannot encode sp as Rm in SUBS), slow path
ipint_extern_handle_vm_traps_at_loop services the trap and resumes or
throws ExceptionType::Termination.
- BBQ: branchPtr(Below, sp, limit) late path runs
operationWasmHandleTrapsAtLoop under jit.probe so all live state is
preserved for the resume case; termination routes through
recordJumpToThrowException so multiple loops share one sink.
- OMG: lower the poll as plain B3 IR (Load + Above + Branch) into a
FrequencyClass::Rare block, matching FTL compileCheckTraps; fp stands in
for sp since B3 has no SP value. The probe patchpoint's clobber set is
confined to the rare block so the hot loop keeps its register budget.
Non-termination async traps (NeedStopTheWorld / NeedWatchdogCheck /
NeedDebuggerBreak) are serviced and the loop resumes with all live state
intact; termination unwinds via the existing Wasm exception machinery.
Rebased onto 78d45d3 (resolves drift from the upstream merge and
worker-lifetime changes).
bff0b81 to
a50d1e5
Compare
|
Rebased onto current main (78d45d3); the previous base predated the 3722912 upstream merge and the worker-lifetime changes on the Bun side. The diff applied without conflicts, squashed to one commit (a50d1e5) since the intermediate revisions were review iterations. Validated against oven-sh/bun main (9008ae7ab) using this branch as the local WebKit (linux x64, debug+ASAN):
|
| Effects effects = Effects::none(); | ||
| effects.reads = B3::HeapRange::top(); | ||
| effects.exitsSideways = true; | ||
| handle->effects = effects; |
There was a problem hiding this comment.
🔴 The slow-path patchpoint's effects leave writes empty (only reads/exitsSideways are set), so B3 CSE forwards an outer loop's limit Load to an inner loop's — the inner back-edge then tests a value read once and never re-loaded, and a nested pure-Wasm loop in OMG never observes requestStop(). Add effects.writes = HeapRange::top(); (matching Effects::forCall(), which FTL's lazySlowPath() keeps by default).
Extended reasoning...
What the bug is
In OMGIRGenerator::addLoop(), the slow-path patchpoint's effects are built as:
Effects effects = Effects::none();
effects.reads = B3::HeapRange::top();
effects.exitsSideways = true;
handle->effects = effects;This leaves effects.writes at its default HeapRange() (the empty range). The m_trapAwareSoftStackLimit load is a plain MemoryValue (no fence), so B3's eliminateCommonSubexpressions is free to forward one loop's Load to another's whenever no block on any path between them has a writes set that overlaps HeapRange::top().
The code path that triggers it
B3EliminateCommonSubexpressions.cpp:findMemoryValue() walks the CFG backward from a Load, bailing only on data.writes.overlaps(range) (line 1028). It never consults exitsSideways, fence (of other blocks), or reads. Per-block data.writes is populated at line ~349 via if (HeapRange writes = effects.writes) clobber(...); an empty HeapRange is falsy, so the slow-path block contributes nothing and does not stop the walk.
Both loads share the same ptr key: instanceValue() returns the single cached m_instanceValue, and both are Load pointerType() at offsetOfTrapAwareSoftStackLimit(), so the Load-case filter (offset + opcode + type) matches.
Why existing code doesn't prevent it
- LICM doesn't hoist the load (
hoistLoopInvariantValuesbails on control-dependent loads when the loop has side exits), but that's a different pass — CSE has no such guard. exitsSidewayson the patchpoint prevents the patchpoint itself from being moved/eliminated, but CSE's predecessor walk for other loads only inspectsdata.writes.- FTL
compileCheckTraps— which the comment cites as the model — avoids this becauselazySlowPath()creates aPatchpointValuewith the default constructor effects (B3PatchpointValue.cpp:54→Effects::forCall()), which haswrites = HeapRange::top(). That write barrier on the FTL slow-path block is exactly what makes CSE bail. The OMG code explicitly discards it.
Step-by-step proof
Take (loop $outer (loop $inner (br $inner))). addLoop produces (per loop) body → Load, Above, Branch → {slowPath (Rare), continuation}; slowPath → patchpoint(writes=∅), Jump → continuation. br 0 targets the Loop's special = body, so the inner back-edge is inner_cont → inner_body.
CSE processes Load₂ in inner_body:
- No local match;
m_data.writesis empty → walk predecessors {outer_cont, inner_cont}. - outer_cont: no writes, no match → push preds {outer_body, outer_slow}.
- inner_cont: no writes, no match → push preds {inner_body, inner_slow}.
- outer_body:
memoryValuesAtTailcontains Load₁ (same ptr/offset/type) →matches = {Load₁}; continue. - outer_slow: patchpoint has
writes = ∅→ does not bail; preds already visited. - inner_body:
memoryValuesAtTailcontains Load₂, but thematch != m_valueguard (line 1022) skips it; no writes → continue. - inner_slow:
writes = ∅→ does not bail. - Worklist exhausts with
matches = {Load₁}(never reached the root).replaceMemoryValuesees a single match,RELEASE_ASSERT(outer_body dominates inner_body)passes (every path to inner_body goes through outer_body), and rewrites Load₂ → Identity(Load₁).
The inner loop's Branch now tests Above(Load₁, fp). The inner back-edge cycle inner_body ↔ inner_cont never re-executes outer_body, so Load₁ is read exactly once. A later requestStop() poisoning m_trapAwareSoftStackLimit is never observed and the OMG-compiled function spins forever — the exact hang this PR fixes, reintroduced for nested pure-Wasm loops (local.set etc. lower to B3 ops with writesLocalState only, so realistic pure-compute inner loops qualify). OMG runs at optLevel = 2 (Options::wasmOMGOptimizationLevel() default), which enables eliminateCommonSubexpressions in B3Generate.cpp, so this fires in the default configuration; the PR's single-loop repro wouldn't catch it.
Fix
Add one line so CSE bails when it walks through the slow-path block via the back-edge:
Effects effects = Effects::none();
effects.reads = B3::HeapRange::top();
effects.writes = B3::HeapRange::top(); // block CSE across the probe (matches Effects::forCall())
effects.exitsSideways = true;(or equivalently start from Effects::forCall()). This is confined to the Rare slow-path block, so the hot path still carries no clobber set as intended.
Bump WEBKIT_VERSION to the oven-sh/WebKit#372 preview build (rebased onto fork main 78d45d3184 as a50d1e5072): JSC now polls the trap-aware stack limit at every Wasm loop back-edge (IPInt/BBQ/OMG), so a Worker spinning in a pure-Wasm loop (no JS re-entry, e.g. loop { br 0 }) observes VM::notifyNeedTermination() instead of never reaching a trap check. worker-terminate-wasm-loop.test.ts covers the three loop shapes (JS loop, Wasm loop calling a JS import, pure Wasm loop) and asserts terminate() resolves with exit code 1 for each. Before merging, oven-sh/WebKit#372 must land and WEBKIT_VERSION must be repinned to the merged 40-hex sha; the preview tag is CI-only.
A Worker running a WebAssembly function whose body is a tight loop with no JS re-entry (e.g.
(loop (br 0))) could not be preempted byVM::notifyNeedTermination(): theworker.terminate()promise never settled and the thread spun at 100% CPU until the process died. JSfor(;;){}and Wasm loops that call an imported JS function per iteration were already preemptible because they reach a JS-side trap check.Cause
notifyNeedTermination()firesVMTraps::NeedTermination, which callsStackManager::requestStop()to setm_trapAwareSoftStackLimiton every registeredMirror(including eachJSWebAssemblyInstance'sm_stackMirror) toUINTPTR_MAX. The only Wasm-side consumer of that field was the IPInt function prologue (InPlaceInterpreter.asmcheckStackOverflow); BBQ/OMG prologues read the non-trap-awarem_softStackLimit, and no tier reads it at loop back-edges. TheVMTraps::SignalSenderpath only patches JS DFG/FTL CodeBlocks (tryInstallTrapBreakpointsreturns early on a Wasm PC), so a pure-Wasm loop never observed the termination request.Fix
Add a trap-aware stack-limit poll at each Wasm loop head:
_loopop): onebpbeqagainstMirror::m_trapAwareSoftStackLimitin the dispatch slot; an out-of-slot trampolineipint_loop_check_vm_trapsservices the trap via a newipint_extern_handle_vm_traps_at_loopslow path (handleTrapsIfNeeded()thenIPINT_THROW(Termination)or resume).addLoop):branchPtragainstoffsetOfTrapAwareSoftStackLimit(); the late path runshandleTrapsIfNeededunderjit.probeso all live state is preserved, resumes on a non-termination async trap (NeedStopTheWorld / NeedWatchdogCheck / NeedDebuggerBreak), andemitThrowException(Termination)otherwise.addLoop): a B3PatchpointValuewithexitsSidewaysemits the same branch + probe + late-path throw/resume, declaringmacroClobberedGPRs+nonPreservedNonArgumentGPR0.Expose
Mirror::offsetOfTrapAwareSoftStackLimit()andJSWebAssemblyInstance::offsetOfTrapAwareSoftStackLimit(); addoperationWasmHandleTrapsAtLoop(Probe::Context&)for the JIT tiers.Hot-path cost is one load + one predicted-not-taken branch per back-edge, matching the JS
op_check_trapscost.Repro
Before:
HUNG >15s(worker thread remains Running at ~100% CPU). After:terminated(code 1)in ~140ms. Node.js terminates all three shapes (js / wasm-with-JS-call / pure-wasm) in ≤4ms.Branch is based on 5491700 (Bun's current
WEBKIT_VERSION) for a minimal-diff preview build; happy to rebase ontomain.