From a50d1e50720c914d9e5bd3422dcf59a19eab48cb Mon Sep 17 00:00:00 2001 From: robobun Date: Sun, 9 Aug 2026 13:31:39 +0000 Subject: [PATCH] wasm: poll VMTraps at loop back-edges so pure-Wasm loops can be terminated 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 78d45d3184 (resolves drift from the upstream merge and worker-lifetime changes). --- .../llint/InPlaceInterpreter.asm | 16 +++++++ .../llint/InPlaceInterpreter64.asm | 7 +++ Source/JavaScriptCore/runtime/StackManager.h | 5 +++ Source/JavaScriptCore/wasm/WasmBBQJIT.cpp | 18 ++++++++ .../wasm/WasmIPIntSlowPaths.cpp | 16 +++++++ .../JavaScriptCore/wasm/WasmIPIntSlowPaths.h | 1 + .../wasm/WasmOMGIRGenerator.cpp | 43 +++++++++++++++++++ Source/JavaScriptCore/wasm/WasmOperations.cpp | 14 ++++++ Source/JavaScriptCore/wasm/WasmOperations.h | 1 + .../wasm/js/JSWebAssemblyInstance.h | 1 + 10 files changed, 122 insertions(+) diff --git a/Source/JavaScriptCore/llint/InPlaceInterpreter.asm b/Source/JavaScriptCore/llint/InPlaceInterpreter.asm index 576d58beb2049..879a4590c4f6f 100644 --- a/Source/JavaScriptCore/llint/InPlaceInterpreter.asm +++ b/Source/JavaScriptCore/llint/InPlaceInterpreter.asm @@ -1173,6 +1173,22 @@ macro handleDebuggerTrapIfNeeded() addp 4 * MachineRegisterSize, sp end +# Slow path for ipintOp(_loop) when m_trapAwareSoftStackLimit has been poisoned by +# VMTraps::requestStop(). Services the pending async trap; on Termination the call +# throws (operationCallMayThrow unwinds), otherwise PC/MC (restored by +# operationCallMayThrow, still pointing at the loop opcode) are advanced past it and +# the next instruction is dispatched. +op(ipint_loop_check_vm_traps, macro () + operationCallMayThrow(macro() + move cfr, a1 + cCall2(_ipint_extern_handle_vm_traps_at_loop) + end) + loadb IPInt::InstructionLengthMetadata::length[MC], t0 + advancePCByReg(t0) + advanceMCByReg(constexpr (sizeof(IPInt::InstructionLengthMetadata))) + nextIPIntInstruction() +end) + op(wasm_ipint_check_debugger_hook_and_throw_trap, macro () handleDebuggerTrapIfNeeded() # r0 == 0 i.e. DebuggerTrapStatus::ResolvedByDebugger i.e. this was purely a debugger trap / breakpoint, diff --git a/Source/JavaScriptCore/llint/InPlaceInterpreter64.asm b/Source/JavaScriptCore/llint/InPlaceInterpreter64.asm index 518aba1fdc39d..432cb500670b9 100644 --- a/Source/JavaScriptCore/llint/InPlaceInterpreter64.asm +++ b/Source/JavaScriptCore/llint/InPlaceInterpreter64.asm @@ -272,6 +272,13 @@ ipintOp(_loop, macro() # loop # We already validateOpcodeConfig in ipintLoopOSR. ipintLoopOSR(1) + # VMTraps poll: requestStop() sets m_trapAwareSoftStackLimit = UINTPTR_MAX, so a pure + # Wasm loop (no calls) observes termination/watchdog requests at each back-edge instead of + # only at the next function prologue. sp as the first comparand (not second) because ARM64 + # cannot encode sp as Rm in SUBS. + bpaeq sp, JSWebAssemblyInstance::m_stackMirror + StackManager::Mirror::m_trapAwareSoftStackLimit[wasmInstance], .ipint_loop_no_trap + jmp _ipint_loop_check_vm_traps +.ipint_loop_no_trap: loadb IPInt::InstructionLengthMetadata::length[MC], t0 advancePCByReg(t0) advanceMCByReg(constexpr (sizeof(IPInt::InstructionLengthMetadata))) diff --git a/Source/JavaScriptCore/runtime/StackManager.h b/Source/JavaScriptCore/runtime/StackManager.h index 7cd701b836e70..c878330dc9953 100644 --- a/Source/JavaScriptCore/runtime/StackManager.h +++ b/Source/JavaScriptCore/runtime/StackManager.h @@ -49,6 +49,11 @@ class StackManager { return OBJECT_OFFSETOF(Mirror, m_softStackLimit); } + static constexpr ptrdiff_t offsetOfTrapAwareSoftStackLimit() + { + return OBJECT_OFFSETOF(Mirror, m_trapAwareSoftStackLimit); + } + private: Atomic m_trapAwareSoftStackLimit { nullptr }; void* m_softStackLimit { nullptr }; diff --git a/Source/JavaScriptCore/wasm/WasmBBQJIT.cpp b/Source/JavaScriptCore/wasm/WasmBBQJIT.cpp index c6e1574cbe431..904df6e70d539 100644 --- a/Source/JavaScriptCore/wasm/WasmBBQJIT.cpp +++ b/Source/JavaScriptCore/wasm/WasmBBQJIT.cpp @@ -3648,6 +3648,24 @@ void BBQJIT::emitLoopTierUpCheckAndOSREntryData(const ControlData& data, std::sp ASSERT(enclosingStack.size() >= args.size()); auto enclosingWithoutArgs = enclosingStack.first(enclosingStack.size() - args.size()); emitLoopTierUpCheckAndOSREntryData(result, enclosingWithoutArgs, loopIndex); + + // VMTraps poll: requestStop() (e.g. VM::notifyNeedTermination) poisons + // m_trapAwareSoftStackLimit so a pure-Wasm loop observes termination/watchdog requests at + // each back-edge instead of never. The slow path runs under jit.probe so all live state is + // preserved for the resume case (non-termination async traps like NeedStopTheWorld). + JIT_COMMENT(m_jit, "Loop VMTraps poll"); + static_assert(GPRInfo::nonPreservedNonArgumentGPR0 == wasmScratchGPR); + // sp as the left comparand: matches the offsetOfSoftStackLimit() checks above and avoids + // ARM64's extra mov when sp is Rm. + Jump needTrapHandling = m_jit.branchPtr(CCallHelpers::Below, MacroAssembler::stackPointerRegister, CCallHelpers::Address(GPRInfo::wasmContextInstancePointer, JSWebAssemblyInstance::offsetOfTrapAwareSoftStackLimit())); + MacroAssembler::Label trapResume = m_jit.label(); + addLatePath(origin(), [needTrapHandling, trapResume](BBQJIT& bbq, CCallHelpers& jit) { + needTrapHandling.link(&jit); + jit.probe(tagCFunction(operationWasmHandleTrapsAtLoop), nullptr); + bbq.recordJumpToThrowException(ExceptionType::Termination, jit.branchTestPtr(CCallHelpers::NonZero, GPRInfo::nonPreservedNonArgumentGPR0)); + jit.jump().linkTo(trapResume, &jit); + }); + return { }; } diff --git a/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.cpp b/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.cpp index 8c2b25daf4057..9578c160ef221 100644 --- a/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.cpp +++ b/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.cpp @@ -1467,6 +1467,22 @@ WASM_IPINT_EXTERN_CPP_DECL(check_stack_and_vm_traps, void* candidateNewStackPoin IPINT_THROW(Wasm::ExceptionType::StackOverflow); } +// Reached from the loop back-edge when m_trapAwareSoftStackLimit has been poisoned by +// VMTraps::requestStop(). Services whatever async trap is pending (NeedTermination / +// NeedWatchdogCheck / NeedStopTheWorld / NeedDebuggerBreak); resumes the loop when it +// wasn't a termination. Without this a pure-Wasm loop never observes VMTraps at all. +WASM_IPINT_EXTERN_CPP_DECL(handle_vm_traps_at_loop, CallFrame* callFrame) +{ + UNUSED_PARAM(callFrame); + VM& vm = instance->vm(); + if (vm.traps().handleTrapsIfNeeded()) { + if (vm.hasPendingTerminationException()) + IPINT_THROW(Wasm::ExceptionType::Termination); + ASSERT(!vm.exceptionForInspection()); + } + IPINT_END(); +} + #if ENABLE(WEBASSEMBLY_DEBUGGER) static UNUSED_FUNCTION void displayWasmDebugState(JSWebAssemblyInstance* instance, Wasm::IPIntCallee* callee, CallFrame* callFrame, IPIntStackEntry* sp) { diff --git a/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.h b/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.h index b0068816796fe..eb03ff2adb8cb 100644 --- a/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.h +++ b/Source/JavaScriptCore/wasm/WasmIPIntSlowPaths.h @@ -139,6 +139,7 @@ WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(memory_atomic_wait64, IPIntStackEntry*); WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(memory_atomic_notify, IPIntStackEntry*); WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(check_stack_and_vm_traps, void* candidateNewStackPointer, Wasm::IPIntCallee*, CallFrame*); +WASM_IPINT_EXTERN_CPP_HIDDEN_DECL(handle_vm_traps_at_loop, CallFrame*); WASM_IPINT_EXTERN_CPP_DECL(handle_debugger_trap_if_needed, CallFrame*, Register*); diff --git a/Source/JavaScriptCore/wasm/WasmOMGIRGenerator.cpp b/Source/JavaScriptCore/wasm/WasmOMGIRGenerator.cpp index 3b0e5bbe91fef..078a62f230ede 100644 --- a/Source/JavaScriptCore/wasm/WasmOMGIRGenerator.cpp +++ b/Source/JavaScriptCore/wasm/WasmOMGIRGenerator.cpp @@ -4619,6 +4619,49 @@ auto OMGIRGenerator::addLoop(BlockSignature&& signature, std::span sp so fp < limit only under the poisoned-limit, which is the branch we want. + { + BasicBlock* slowPath = m_proc.addBlock(); + BasicBlock* continuation = m_proc.addBlock(); + + Value* limit = m_currentBlock->appendNew(m_proc, Load, pointerType(), origin(), instanceValue(), safeCast(JSWebAssemblyInstance::offsetOfTrapAwareSoftStackLimit())); + m_currentBlock->appendNewControlValue(m_proc, B3::Branch, origin(), + m_currentBlock->appendNew(m_proc, Above, origin(), limit, framePointer()), + FrequentedBlock(slowPath, FrequencyClass::Rare), FrequentedBlock(continuation)); + slowPath->addPredecessor(m_currentBlock); + continuation->addPredecessor(m_currentBlock); + + m_currentBlock = slowPath; + // The probe preserves all live state for the resume case (non-termination async traps + // like NeedStopTheWorld); clobbers apply only to this rare block. + B3::PatchpointValue* handle = m_currentBlock->appendNew(m_proc, B3::Void, origin()); + Effects effects = Effects::none(); + effects.reads = B3::HeapRange::top(); + effects.exitsSideways = true; + handle->effects = effects; + RegisterSet clobbers = RegisterSet::macroClobberedGPRs(); + clobbers.add(GPRInfo::nonPreservedNonArgumentGPR0, IgnoreVectors); + handle->clobber(clobbers); + handle->append(instanceValue(), ValueRep::reg(GPRInfo::wasmContextInstancePointer)); + handle->setGenerator([this, origin = this->origin()](CCallHelpers& jit, const B3::StackmapGenerationParams&) { + AllowMacroScratchRegisterUsage allowScratch(jit); + jit.probe(tagCFunction(operationWasmHandleTrapsAtLoop), nullptr); + CCallHelpers::Jump resume = jit.branchTestPtr(CCallHelpers::Zero, GPRInfo::nonPreservedNonArgumentGPR0); + this->emitExceptionCheck(jit, origin, ExceptionType::Termination); + resume.link(&jit); + }); + m_currentBlock->appendNewControlValue(m_proc, Jump, origin(), continuation); + continuation->addPredecessor(m_currentBlock); + + m_currentBlock = continuation; + } + return { }; } diff --git a/Source/JavaScriptCore/wasm/WasmOperations.cpp b/Source/JavaScriptCore/wasm/WasmOperations.cpp index c2ee2bccd8c24..08f8e558d69cd 100644 --- a/Source/JavaScriptCore/wasm/WasmOperations.cpp +++ b/Source/JavaScriptCore/wasm/WasmOperations.cpp @@ -1171,6 +1171,20 @@ JSC_DEFINE_NOEXCEPT_JIT_OPERATION(operationWasmLoopOSREnterBBQJIT, void, (Probe: context.gpr(GPRInfo::nonPreservedNonArgumentGPR0) = std::bit_cast(callee.loopEntrypoints()[loopIndex].taggedPtr()); } +// Reached from BBQ/OMG loop heads when m_trapAwareSoftStackLimit has been poisoned by +// VMTraps::requestStop() (e.g. VM::notifyNeedTermination()). Services whatever async trap +// is pending and reports whether a TerminationException was raised so the caller can unwind; +// otherwise all registers are restored by the probe and the loop resumes. Without this a +// pure-Wasm loop that never calls into JS (e.g. `loop { br 0 }`) never observes VMTraps and +// cannot be preempted by worker.terminate() / the watchdog. +JSC_DEFINE_NOEXCEPT_JIT_OPERATION(operationWasmHandleTrapsAtLoop, void, (Probe::Context& context)) +{ + JSWebAssemblyInstance* instance = context.gpr(GPRInfo::wasmContextInstancePointer); + VM& vm = instance->vm(); + vm.traps().handleTrapsIfNeeded(); + context.gpr(GPRInfo::nonPreservedNonArgumentGPR0) = static_cast(vm.hasPendingTerminationException()); +} + #endif #if ENABLE(WEBASSEMBLY_BBQJIT) diff --git a/Source/JavaScriptCore/wasm/WasmOperations.h b/Source/JavaScriptCore/wasm/WasmOperations.h index 2668eea594414..3d5a3f7a12367 100644 --- a/Source/JavaScriptCore/wasm/WasmOperations.h +++ b/Source/JavaScriptCore/wasm/WasmOperations.h @@ -69,6 +69,7 @@ JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmTriggerTierUpNow, void, (CallFra #if ENABLE(WEBASSEMBLY_OMGJIT) || ENABLE(WEBASSEMBLY_BBQJIT) JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmTriggerOSREntryNow, void, (Probe::Context&)); JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmLoopOSREnterBBQJIT, void, (Probe::Context&)); +JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmHandleTrapsAtLoop, void, (Probe::Context&)); #endif #if ENABLE(WEBASSEMBLY_BBQJIT) JSC_DECLARE_NOEXCEPT_JIT_OPERATION(operationWasmMaterializeBaselineData, void, (CallFrame*, JSWebAssemblyInstance*)); diff --git a/Source/JavaScriptCore/wasm/js/JSWebAssemblyInstance.h b/Source/JavaScriptCore/wasm/js/JSWebAssemblyInstance.h index cfb1565b1d974..0b7b6e623fc5d 100644 --- a/Source/JavaScriptCore/wasm/js/JSWebAssemblyInstance.h +++ b/Source/JavaScriptCore/wasm/js/JSWebAssemblyInstance.h @@ -177,6 +177,7 @@ class JSWebAssemblyInstance final : public JSNonFinalObject { using FunctionWrapperMap = UncheckedKeyHashMap, IntHash, WTF::UnsignedWithZeroKeyHashTraits>; static constexpr ptrdiff_t offsetOfSoftStackLimit() { return OBJECT_OFFSETOF(JSWebAssemblyInstance, m_stackMirror) + StackManager::Mirror::offsetOfSoftStackLimit(); } + static constexpr ptrdiff_t offsetOfTrapAwareSoftStackLimit() { return OBJECT_OFFSETOF(JSWebAssemblyInstance, m_stackMirror) + StackManager::Mirror::offsetOfTrapAwareSoftStackLimit(); } Wasm::Module& module() const { return m_module.get(); } SourceTaintedOrigin taintedness() const { return m_sourceProvider->sourceTaintedOrigin(); }