Bug 1435360 - Baldr: implement wasm interrupt in terms of TlsData branch and stack overflow check (r=bbouvier)

--HG--
extra : rebase_source : 2384561de07e6849cb83a00713ba57e71b3715f6

js/public/MemoryMetrics.h

@@ -516,7 +516,8 @@ struct RuntimeSizes
     macro(_, MallocHeap, sharedIntlData) \
     macro(_, MallocHeap, uncompressedSourceCache) \
     macro(_, MallocHeap, scriptData) \
-    macro(_, MallocHeap, tracelogger)
+    macro(_, MallocHeap, tracelogger) \
+    macro(_, MallocHeap, wasmRuntime)
 
     RuntimeSizes()
       : FOR_EACH_SIZE(ZERO_SIZE)

js/src/jit/CodeGenerator.cpp

@@ -12710,6 +12710,14 @@ CodeGenerator::visitInterruptCheck(LInterruptCheck* lir)
     masm.bind(ool->rejoin());
 }
 
+void
+CodeGenerator::visitWasmInterruptCheck(LWasmInterruptCheck* lir)
+{
+    MOZ_ASSERT(gen->compilingWasm());
+
+    masm.wasmInterruptCheck(ToRegister(lir->tlsPtr()), lir->mir()->bytecodeOffset());
+}
+
 void
 CodeGenerator::visitWasmTrap(LWasmTrap* lir)
 {

js/src/jit/CodeGenerator.h

@@ -465,6 +465,7 @@ class CodeGenerator final : public CodeGeneratorSpecific
 
     void visitInterruptCheck(LInterruptCheck* lir);
     void visitOutOfLineInterruptCheckImplicit(OutOfLineInterruptCheckImplicit* ins);
+    void visitWasmInterruptCheck(LWasmInterruptCheck* lir);
     void visitWasmTrap(LWasmTrap* lir);
     void visitWasmLoadTls(LWasmLoadTls* ins);
     void visitWasmBoundsCheck(LWasmBoundsCheck* ins);

js/src/jit/Lowering.cpp

@@ -2737,6 +2737,13 @@ LIRGenerator::visitInterruptCheck(MInterruptCheck* ins)
     assignSafepoint(lir, ins);
 }
 
+void
+LIRGenerator::visitWasmInterruptCheck(MWasmInterruptCheck* ins)
+{
+    auto* lir = new(alloc()) LWasmInterruptCheck(useRegisterAtStart(ins->tlsPtr()));
+    add(lir, ins);
+}
+
 void
 LIRGenerator::visitWasmTrap(MWasmTrap* ins)
 {

js/src/jit/Lowering.h

@@ -211,6 +211,7 @@ class LIRGenerator : public LIRGeneratorSpecific
     void visitHomeObject(MHomeObject* ins) override;
     void visitHomeObjectSuperBase(MHomeObjectSuperBase* ins) override;
     void visitInterruptCheck(MInterruptCheck* ins) override;
+    void visitWasmInterruptCheck(MWasmInterruptCheck* ins) override;
     void visitWasmTrap(MWasmTrap* ins) override;
     void visitWasmReinterpret(MWasmReinterpret* ins) override;
     void visitStoreSlot(MStoreSlot* ins) override;

js/src/jit/MIR.h

@@ -8303,6 +8303,33 @@ class MInterruptCheck : public MNullaryInstruction
     }
 };
 
+// Check whether we need to fire the interrupt handler (in wasm code).
+class MWasmInterruptCheck
+  : public MUnaryInstruction,
+    public NoTypePolicy::Data
+{
+    wasm::BytecodeOffset bytecodeOffset_;
+
+    MWasmInterruptCheck(MDefinition* tlsPointer, wasm::BytecodeOffset bytecodeOffset)
+      : MUnaryInstruction(classOpcode, tlsPointer),
+        bytecodeOffset_(bytecodeOffset)
+    {
+        setGuard();
+    }
+
+  public:
+    INSTRUCTION_HEADER(WasmInterruptCheck)
+    TRIVIAL_NEW_WRAPPERS
+    NAMED_OPERANDS((0, tlsPtr))
+
+    AliasSet getAliasSet() const override {
+        return AliasSet::None();
+    }
+    wasm::BytecodeOffset bytecodeOffset() const {
+        return bytecodeOffset_;
+    }
+};
+
 // Directly jumps to the indicated trap, leaving Wasm code and reporting a
 // runtime error.
 

js/src/jit/MOpcodes.h

@@ -277,6 +277,7 @@ namespace jit {
     _(InstanceOf)                                                           \
     _(InstanceOfCache)                                                      \
     _(InterruptCheck)                                                       \
+    _(WasmInterruptCheck)                                                   \
     _(GetDOMProperty)                                                       \
     _(GetDOMMember)                                                         \
     _(SetDOMProperty)                                                       \

js/src/jit/MacroAssembler.cpp

@@ -3357,7 +3357,19 @@ MacroAssembler::maybeBranchTestType(MIRType type, MDefinition* maybeDef, Registe
 void
 MacroAssembler::wasmTrap(wasm::Trap trap, wasm::BytecodeOffset bytecodeOffset)
 {
-    append(trap, wasm::TrapSite(wasmTrapInstruction().offset(), bytecodeOffset));
+    uint32_t trapOffset = wasmTrapInstruction().offset();
+    MOZ_ASSERT_IF(!oom(), currentOffset() - trapOffset == WasmTrapInstructionLength);
+
+    append(trap, wasm::TrapSite(trapOffset, bytecodeOffset));
+}
+
+void
+MacroAssembler::wasmInterruptCheck(Register tls, wasm::BytecodeOffset bytecodeOffset)
+{
+    Label ok;
+    branch32(Assembler::Equal, Address(tls, offsetof(wasm::TlsData, interrupt)), Imm32(0), &ok);
+    wasmTrap(wasm::Trap::CheckInterrupt, bytecodeOffset);
+    bind(&ok);
 }
 
 void

js/src/jit/MacroAssembler.h

@@ -1493,6 +1493,7 @@ class MacroAssembler : public MacroAssemblerSpecific
     CodeOffset wasmTrapInstruction() PER_SHARED_ARCH;
 
     void wasmTrap(wasm::Trap trap, wasm::BytecodeOffset bytecodeOffset);
+    void wasmInterruptCheck(Register tls, wasm::BytecodeOffset bytecodeOffset);
 
     // Emit a bounds check against the wasm heap limit, jumping to 'label' if
     // 'cond' holds. Required when WASM_HUGE_MEMORY is not defined. If

js/src/jit/arm/Assembler-arm.h

@@ -128,6 +128,7 @@ static constexpr FloatRegister ABINonArgDoubleReg { FloatRegisters::d8, VFPRegis
 // Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = r4;
 static constexpr Register ABINonArgReturnReg1 = r5;
+static constexpr Register ABINonVolatileReg = r6;
 
 // This register is guaranteed to be clobberable during the prologue and
 // epilogue of an ABI call which must preserve both ABI argument, return
@@ -249,6 +250,7 @@ static_assert(JitStackAlignment % SimdMemoryAlignment == 0,
   "spilled values.  Thus it should be larger than the alignment for SIMD accesses.");
 
 static const uint32_t WasmStackAlignment = SimdMemoryAlignment;
+static const uint32_t WasmTrapInstructionLength = 4;
 
 // Does this architecture support SIMD conversions between Uint32x4 and Float32x4?
 static constexpr bool SupportsUint32x4FloatConversions = false;

js/src/jit/arm64/Assembler-arm64.h

@@ -172,6 +172,7 @@ static_assert(CodeAlignment % SimdMemoryAlignment == 0,
   "alignment for SIMD constants.");
 
 static const uint32_t WasmStackAlignment = SimdMemoryAlignment;
+static const uint32_t WasmTrapInstructionLength = 4;
 
 // Does this architecture support SIMD conversions between Uint32x4 and Float32x4?
 static constexpr bool SupportsUint32x4FloatConversions = false;
@@ -457,6 +458,7 @@ static constexpr FloatRegister ABINonArgDoubleReg = { FloatRegisters::s16, Float
 // Note: these three registers are all guaranteed to be different
 static constexpr Register ABINonArgReturnReg0 = r8;
 static constexpr Register ABINonArgReturnReg1 = r9;
+static constexpr Register ABINonVolatileReg { Registers::x19 };
 
 // This register is guaranteed to be clobberable during the prologue and
 // epilogue of an ABI call which must preserve both ABI argument, return

js/src/jit/arm64/vixl/MozSimulator-vixl.cpp

@@ -267,7 +267,7 @@ Simulator::handle_wasm_seg_fault(uintptr_t addr, unsigned numBytes)
 
     const js::wasm::MemoryAccess* memoryAccess = instance->code().lookupMemoryAccess(pc);
     if (!memoryAccess) {
-        act->startWasmTrap(wasm::Trap::OutOfBounds, 0, registerState());
+        act->startWasmTrap(js::wasm::Trap::OutOfBounds, 0, registerState());
         if (!instance->code().containsCodePC(pc))
             MOZ_CRASH("Cannot map PC to trap handler");
         set_pc((Instruction*)moduleSegment->outOfBoundsCode());

js/src/jit/mips32/Assembler-mips32.h

@@ -135,6 +135,7 @@ static_assert(JitStackAlignment % sizeof(Value) == 0 && JitStackValueAlignment >
 // TODO Copy the static_asserts from x64/x86 assembler files.
 static constexpr uint32_t SimdMemoryAlignment = 8;
 static constexpr uint32_t WasmStackAlignment = SimdMemoryAlignment;
+static const uint32_t WasmTrapInstructionLength = 4;
 
 // Does this architecture support SIMD conversions between Uint32x4 and Float32x4?
 static constexpr bool SupportsUint32x4FloatConversions = false;

js/src/jit/mips64/Assembler-mips64.h

@@ -146,6 +146,7 @@ static_assert(JitStackAlignment % sizeof(Value) == 0 && JitStackValueAlignment >
 static constexpr uint32_t SimdMemoryAlignment = 16;
 
 static constexpr uint32_t WasmStackAlignment = SimdMemoryAlignment;
+static const uint32_t WasmTrapInstructionLength = 4;
 
 // Does this architecture support SIMD conversions between Uint32x4 and Float32x4?
 static constexpr bool SupportsUint32x4FloatConversions = false;

js/src/jit/none/Architecture-none.h

@@ -19,6 +19,7 @@ namespace jit {
 static const bool SupportsSimd = false;
 static const uint32_t SimdMemoryAlignment = 4; // Make it 4 to avoid a bunch of div-by-zero warnings
 static const uint32_t WasmStackAlignment = 8;
+static const uint32_t WasmTrapInstructionLength = 0;
 
 // Does this architecture support SIMD conversions between Uint32x4 and Float32x4?
 static constexpr bool SupportsUint32x4FloatConversions = false;

js/src/jit/none/MacroAssembler-none.h

@@ -81,6 +81,7 @@ static constexpr Register ABINonArgReg1 { Registers::invalid_reg };
 static constexpr Register ABINonArgReg2 { Registers::invalid_reg };
 static constexpr Register ABINonArgReturnReg0 { Registers::invalid_reg };
 static constexpr Register ABINonArgReturnReg1 { Registers::invalid_reg };
+static constexpr Register ABINonVolatileReg { Registers::invalid_reg };
 static constexpr Register ABINonArgReturnVolatileReg { Registers::invalid_reg };
 
 static constexpr FloatRegister ABINonArgDoubleReg = { FloatRegisters::invalid_reg };

js/src/jit/shared/LIR-shared.h

@@ -1663,6 +1663,24 @@ class LInterruptCheck : public LInstructionHelper<0, 0, 1>
     }
 };
 
+class LWasmInterruptCheck : public LInstructionHelper<0, 1, 0>
+{
+  public:
+    LIR_HEADER(WasmInterruptCheck)
+
+    explicit LWasmInterruptCheck(const LAllocation& tlsData)
+      : LInstructionHelper(classOpcode)
+    {
+        setOperand(0, tlsData);
+    }
+    MWasmInterruptCheck* mir() const {
+        return mir_->toWasmInterruptCheck();
+    }
+    const LAllocation* tlsPtr() {
+        return getOperand(0);
+    }
+};
+
 class LDefVar : public LCallInstructionHelper<0, 1, 0>
 {
   public:

js/src/jit/shared/LOpcodes-shared.h

@@ -389,6 +389,7 @@
     _(InstanceOfV)                  \
     _(InstanceOfCache)              \
     _(InterruptCheck)               \
+    _(WasmInterruptCheck)           \
     _(Rotate)                       \
     _(RotateI64)                    \
     _(GetDOMProperty)               \

js/src/jit/x64/Assembler-x64.h

@@ -250,6 +250,7 @@ static_assert(JitStackAlignment % SimdMemoryAlignment == 0,
   "spilled values.  Thus it should be larger than the alignment for SIMD accesses.");
 
 static const uint32_t WasmStackAlignment = SimdMemoryAlignment;
+static const uint32_t WasmTrapInstructionLength = 2;
 
 static const Scale ScalePointer = TimesEight;
 

js/src/jit/x86/Assembler-x86.h

@@ -167,6 +167,7 @@ static_assert(JitStackAlignment % SimdMemoryAlignment == 0,
   "spilled values.  Thus it should be larger than the alignment for SIMD accesses.");
 
 static const uint32_t WasmStackAlignment = SimdMemoryAlignment;
+static const uint32_t WasmTrapInstructionLength = 2;
 
 struct ImmTag : public Imm32
 {

js/src/vm/JSCompartment.cpp

@@ -73,7 +73,7 @@ JSCompartment::JSCompartment(Zone* zone, const JS::CompartmentOptions& options =
     objectMetadataTable(nullptr),
     innerViews(zone),
     lazyArrayBuffers(nullptr),
-    wasm(zone),
+    wasm(zone->runtimeFromActiveCooperatingThread()),
     nonSyntacticLexicalEnvironments_(nullptr),
     gcIncomingGrayPointers(nullptr),
     debugModeBits(0),

js/src/vm/MutexIDs.h

@@ -36,7 +36,6 @@
   _(FutexThread,                 500) \
   _(GeckoProfilerStrings,        500) \
   _(ProtectedRegionTree,         500) \
-  _(WasmSigIdSet,                500) \
   _(ShellOffThreadState,         500) \
   _(SimulatorCacheLock,          500) \
   _(Arm64SimulatorLock,          500) \
@@ -49,14 +48,16 @@
   _(ProcessExecutableRegion,     500) \
   _(OffThreadPromiseState,       500) \
   _(BufferStreamState,           500) \
+  _(SharedArrayGrow,             500) \
+  _(RuntimeScriptData,           500) \
+  _(WasmSigIdSet,                500) \
   _(WasmCodeProfilingLabels,     500) \
   _(WasmModuleTieringLock,       500) \
   _(WasmCompileTaskState,        500) \
   _(WasmCodeStreamEnd,           500) \
   _(WasmTailBytesPtr,            500) \
   _(WasmStreamStatus,            500) \
-  _(SharedArrayGrow,             500) \
+  _(WasmRuntimeInstances,        500) \
-  _(RuntimeScriptData,           500) \
                                       \
   _(ThreadId,                    600) \
   _(WasmCodeSegmentMap,          600) \

js/src/vm/Runtime.cpp

@@ -177,7 +177,8 @@ JSRuntime::JSRuntime(JSRuntime* parentRuntime)
     lastAnimationTime(0),
     performanceMonitoring_(),
     stackFormat_(parentRuntime ? js::StackFormat::Default
-                               : js::StackFormat::SpiderMonkey)
+                               : js::StackFormat::SpiderMonkey),
+    wasmInstances(mutexid::WasmRuntimeInstances)
 {
     liveRuntimesCount++;
 
@@ -193,6 +194,8 @@ JSRuntime::~JSRuntime()
 
     DebugOnly<size_t> oldCount = liveRuntimesCount--;
     MOZ_ASSERT(oldCount > 0);
+
+    MOZ_ASSERT(wasmInstances.lock()->empty());
 }
 
 bool
@@ -509,6 +512,8 @@ JSRuntime::addSizeOfIncludingThis(mozilla::MallocSizeOf mallocSizeOf, JS::Runtim
         jitRuntime_->execAlloc().addSizeOfCode(&rtSizes->code);
         jitRuntime_->backedgeExecAlloc().addSizeOfCode(&rtSizes->code);
     }
+
+    rtSizes->wasmRuntime += wasmInstances.lock()->sizeOfExcludingThis(mallocSizeOf);
 }
 
 static bool
@@ -599,6 +604,7 @@ JSContext::requestInterrupt(InterruptMode mode)
             fx.wake(FutexThread::WakeForJSInterrupt);
         fx.unlock();
         jit::InterruptRunningCode(this);
+        wasm::InterruptRunningCode(this);
     }
 }
 

js/src/vm/Runtime.h

@@ -1003,6 +1003,11 @@ struct JSRuntime : public js::MallocProvider<JSRuntime>
     // purposes. Wasm code can't trap reentrantly.
     js::ActiveThreadData<mozilla::Maybe<js::wasm::TrapData>> wasmTrapData;
 
+    // List of all the live wasm::Instances in the runtime. Equal to the union
+    // of all instances registered in all JSCompartments. Accessed from watchdog
+    // threads for purposes of wasm::InterruptRunningCode().
+    js::ExclusiveData<js::wasm::InstanceVector> wasmInstances;
+
   public:
 #if defined(NIGHTLY_BUILD)
     // Support for informing the embedding of any error thrown.

js/src/vm/Stack.cpp

@@ -1762,7 +1762,13 @@ jit::JitActivation::startWasmTrap(wasm::Trap trap, uint32_t bytecodeOffset,
     if (unwound)
         bytecodeOffset = code.lookupCallSite(pc)->lineOrBytecode();
 
-    cx_->runtime()->wasmTrapData.ref().emplace(pc, trap, bytecodeOffset);
+    wasm::TrapData trapData;
+    trapData.resumePC = ((uint8_t*)state.pc) + jit::WasmTrapInstructionLength;
+    trapData.unwoundPC = pc;
+    trapData.trap = trap;
+    trapData.bytecodeOffset = bytecodeOffset;
+
+    cx_->runtime()->wasmTrapData = Some(trapData);
     setWasmExitFP(fp);
 }
 
@@ -1789,17 +1795,22 @@ jit::JitActivation::isWasmTrapping() const
     if (act != this)
         return false;
 
-    DebugOnly<const wasm::Frame*> fp = wasmExitFP();
+    MOZ_ASSERT(wasmExitFP()->instance()->code().containsCodePC(rt->wasmTrapData->unwoundPC));
-    DebugOnly<void*> unwindPC = rt->wasmTrapData->pc;
-    MOZ_ASSERT(fp->instance()->code().containsCodePC(unwindPC));
     return true;
 }
 
 void*
-jit::JitActivation::wasmTrapPC() const
+jit::JitActivation::wasmTrapResumePC() const
 {
     MOZ_ASSERT(isWasmTrapping());
-    return cx_->runtime()->wasmTrapData->pc;
+    return cx_->runtime()->wasmTrapData->resumePC;
+}
+
+void*
+jit::JitActivation::wasmTrapUnwoundPC() const
+{
+    MOZ_ASSERT(isWasmTrapping());
+    return cx_->runtime()->wasmTrapData->unwoundPC;
 }
 
 uint32_t

js/src/vm/Stack.h

@@ -1713,7 +1713,8 @@ class JitActivation : public Activation
     void startWasmTrap(wasm::Trap trap, uint32_t bytecodeOffset, const wasm::RegisterState& state);
     void finishWasmTrap();
     bool isWasmTrapping() const;
-    void* wasmTrapPC() const;
+    void* wasmTrapResumePC() const;
+    void* wasmTrapUnwoundPC() const;
     uint32_t wasmTrapBytecodeOffset() const;
 };
 

js/src/wasm/WasmBaselineCompile.cpp

@@ -3371,9 +3371,10 @@ class BaseCompiler final : public BaseCompilerInterface
         masm.loadConstantDouble(d, dest);
     }
 
-    void addInterruptCheck()
+    void addInterruptCheck() {
-    {
+        ScratchI32 tmp(*this);
-        // TODO
+        masm.loadWasmTlsRegFromFrame(tmp);
+        masm.wasmInterruptCheck(tmp, bytecodeOffset());
     }
 
     void jumpTable(const LabelVector& labels, Label* theTable) {

js/src/wasm/WasmBuiltins.cpp

@@ -211,61 +211,94 @@ WasmHandleThrow()
     return HandleThrow(cx, iter);
 }
 
-static void
+// Unconditionally returns nullptr per calling convention of OnTrap().
-WasmOldReportTrap(int32_t trapIndex)
+static void*
+ReportError(JSContext* cx, unsigned errorNumber)
 {
-    JSContext* cx = TlsContext.get();
+    JS_ReportErrorNumberUTF8(cx, GetErrorMessage, nullptr, errorNumber);
+    return nullptr;
+};
 
-    MOZ_ASSERT(trapIndex < int32_t(Trap::Limit) && trapIndex >= 0);
+// Has the same return-value convention as OnTrap().
-    Trap trap = Trap(trapIndex);
+static void*
+CheckInterrupt(JSContext* cx, JitActivation* activation)
+{
+    ResetInterruptState(cx);
+
+    if (!CheckForInterrupt(cx))
+        return nullptr;
+
+    void* resumePC = activation->wasmTrapResumePC();
+    activation->finishWasmTrap();
+    return resumePC;
+}
+
+// The calling convention between this function and its caller in the stub
+// generated by GenerateTrapExit() is:
+//   - return nullptr if the stub should jump to the throw stub to unwind
+//     the activation;
+//   - return the (non-null) resumePC that should be jumped if execution should
+//     resume after the trap.
+static void*
+OnTrap(Trap trap)
+{
+    JitActivation* activation = CallingActivation();
+    JSContext* cx = activation->cx();
 
-    unsigned errorNumber;
     switch (trap) {
       case Trap::Unreachable:
-        errorNumber = JSMSG_WASM_UNREACHABLE;
+        return ReportError(cx, JSMSG_WASM_UNREACHABLE);
-        break;
       case Trap::IntegerOverflow:
-        errorNumber = JSMSG_WASM_INTEGER_OVERFLOW;
+        return ReportError(cx, JSMSG_WASM_INTEGER_OVERFLOW);
-        break;
       case Trap::InvalidConversionToInteger:
-        errorNumber = JSMSG_WASM_INVALID_CONVERSION;
+        return ReportError(cx, JSMSG_WASM_INVALID_CONVERSION);
-        break;
       case Trap::IntegerDivideByZero:
-        errorNumber = JSMSG_WASM_INT_DIVIDE_BY_ZERO;
+        return ReportError(cx, JSMSG_WASM_INT_DIVIDE_BY_ZERO);
-        break;
       case Trap::IndirectCallToNull:
-        errorNumber = JSMSG_WASM_IND_CALL_TO_NULL;
+        return ReportError(cx, JSMSG_WASM_IND_CALL_TO_NULL);
-        break;
       case Trap::IndirectCallBadSig:
-        errorNumber = JSMSG_WASM_IND_CALL_BAD_SIG;
+        return ReportError(cx, JSMSG_WASM_IND_CALL_BAD_SIG);
-        break;
       case Trap::ImpreciseSimdConversion:
-        errorNumber = JSMSG_SIMD_FAILED_CONVERSION;
+        return ReportError(cx, JSMSG_SIMD_FAILED_CONVERSION);
-        break;
       case Trap::OutOfBounds:
-        errorNumber = JSMSG_WASM_OUT_OF_BOUNDS;
+        return ReportError(cx, JSMSG_WASM_OUT_OF_BOUNDS);
-        break;
       case Trap::UnalignedAccess:
-        errorNumber = JSMSG_WASM_UNALIGNED_ACCESS;
+        return ReportError(cx, JSMSG_WASM_UNALIGNED_ACCESS);
-        break;
+      case Trap::CheckInterrupt:
+        return CheckInterrupt(cx, activation);
       case Trap::StackOverflow:
-        errorNumber = JSMSG_OVER_RECURSED;
+        // TlsData::setInterrupt() causes a fake stack overflow. Since
-        break;
+        // TlsData::setInterrupt() is called racily, it's possible for a real
+        // stack overflow to trap, followed by a racy call to setInterrupt().
+        // Thus, we must check for a real stack overflow first before we
+        // CheckInterrupt() and possibly resume execution.
+        if (!CheckRecursionLimit(cx))
+            return nullptr;
+        if (activation->wasmExitFP()->tls->isInterrupted())
+            return CheckInterrupt(cx, activation);
+        return ReportError(cx, JSMSG_OVER_RECURSED);
       case Trap::ThrowReported:
         // Error was already reported under another name.
-        return;
+        return nullptr;
-      default:
+      case Trap::Limit:
-        MOZ_CRASH("unexpected trap");
+        break;
     }
 
-    JS_ReportErrorNumberUTF8(cx, GetErrorMessage, nullptr, errorNumber);
+    MOZ_CRASH("unexpected trap");
 }
 
 static void
-WasmReportTrap()
+WasmOldReportTrap(int32_t trapIndex)
 {
-    Trap trap = TlsContext.get()->runtime()->wasmTrapData->trap;
+    MOZ_ASSERT(trapIndex < int32_t(Trap::Limit) && trapIndex >= 0);
-    WasmOldReportTrap(int32_t(trap));
+    DebugOnly<void*> resumePC = OnTrap(Trap(trapIndex));
+    MOZ_ASSERT(!resumePC);
+}
+
+static void*
+WasmOnTrap()
+{
+    return OnTrap(TlsContext.get()->runtime()->wasmTrapData->trap);
 }
 
 static void
@@ -494,9 +527,9 @@ wasm::AddressOf(SymbolicAddress imm, ABIFunctionType* abiType)
       case SymbolicAddress::HandleThrow:
         *abiType = Args_General0;
         return FuncCast(WasmHandleThrow, *abiType);
-      case SymbolicAddress::ReportTrap:
+      case SymbolicAddress::OnTrap:
         *abiType = Args_General0;
-        return FuncCast(WasmReportTrap, *abiType);
+        return FuncCast(WasmOnTrap, *abiType);
       case SymbolicAddress::OldReportTrap:
         *abiType = Args_General1;
         return FuncCast(WasmOldReportTrap, *abiType);
@@ -668,7 +701,7 @@ wasm::NeedsBuiltinThunk(SymbolicAddress sym)
     switch (sym) {
       case SymbolicAddress::HandleDebugTrap:          // GenerateDebugTrapStub
       case SymbolicAddress::HandleThrow:              // GenerateThrowStub
-      case SymbolicAddress::ReportTrap:               // GenerateTrapExit
+      case SymbolicAddress::OnTrap:                   // GenerateTrapExit
       case SymbolicAddress::OldReportTrap:            // GenerateOldTrapExit
       case SymbolicAddress::ReportOutOfBounds:        // GenerateOutOfBoundsExit
       case SymbolicAddress::ReportUnalignedAccess:    // GenerateUnalignedExit

js/src/wasm/WasmCompartment.cpp

@@ -26,7 +26,8 @@
 using namespace js;
 using namespace wasm;
 
-Compartment::Compartment(Zone* zone)
+Compartment::Compartment(JSRuntime* rt)
+  : runtime_(rt)
 {}
 
 Compartment::~Compartment()
@@ -62,6 +63,8 @@ struct InstanceComparator
 bool
 Compartment::registerInstance(JSContext* cx, HandleWasmInstanceObject instanceObj)
 {
+    MOZ_ASSERT(runtime_ == cx->runtime());
+
     Instance& instance = instanceObj->instance();
     MOZ_ASSERT(this == &instance.compartment()->wasm);
 
@@ -70,15 +73,27 @@ Compartment::registerInstance(JSContext* cx, HandleWasmInstanceObject instanceOb
     if (instance.debugEnabled() && instance.compartment()->debuggerObservesAllExecution())
         instance.ensureEnterFrameTrapsState(cx, true);
 
-    size_t index;
+    {
-    if (BinarySearchIf(instances_, 0, instances_.length(), InstanceComparator(instance), &index))
+        if (!instances_.reserve(instances_.length() + 1))
-        MOZ_CRASH("duplicate registration");
+            return false;
 
-    if (!instances_.insert(instances_.begin() + index, &instance)) {
+        auto runtimeInstances = cx->runtime()->wasmInstances.lock();
-        ReportOutOfMemory(cx);
+        if (!runtimeInstances->reserve(runtimeInstances->length() + 1))
-        return false;
+            return false;
+
+        // To avoid implementing rollback, do not fail after mutations start.
+
+        InstanceComparator cmp(instance);
+        size_t index;
+
+        MOZ_ALWAYS_FALSE(BinarySearchIf(instances_, 0, instances_.length(), cmp, &index));
+        MOZ_ALWAYS_TRUE(instances_.insert(instances_.begin() + index, &instance));
+
+        MOZ_ALWAYS_FALSE(BinarySearchIf(runtimeInstances.get(), 0, runtimeInstances->length(), cmp, &index));
+        MOZ_ALWAYS_TRUE(runtimeInstances->insert(runtimeInstances->begin() + index, &instance));
     }
 
+    // Notify the debugger after wasmInstances is unlocked.
     Debugger::onNewWasmInstance(cx, instanceObj);
     return true;
 }
@@ -86,10 +101,15 @@ Compartment::registerInstance(JSContext* cx, HandleWasmInstanceObject instanceOb
 void
 Compartment::unregisterInstance(Instance& instance)
 {
+    InstanceComparator cmp(instance);
     size_t index;
-    if (!BinarySearchIf(instances_, 0, instances_.length(), InstanceComparator(instance), &index))
+
-        return;
+    if (BinarySearchIf(instances_, 0, instances_.length(), cmp, &index))
-    instances_.erase(instances_.begin() + index);
+        instances_.erase(instances_.begin() + index);
+
+    auto runtimeInstances = runtime_->wasmInstances.lock();
+    if (BinarySearchIf(runtimeInstances.get(), 0, runtimeInstances->length(), cmp, &index))
+        runtimeInstances->erase(runtimeInstances->begin() + index);
 }
 
 void
@@ -104,3 +124,19 @@ Compartment::addSizeOfExcludingThis(MallocSizeOf mallocSizeOf, size_t* compartme
 {
     *compartmentTables += instances_.sizeOfExcludingThis(mallocSizeOf);
 }
+
+void
+wasm::InterruptRunningCode(JSContext* cx)
+{
+    auto runtimeInstances = cx->runtime()->wasmInstances.lock();
+    for (Instance* instance : runtimeInstances.get())
+        instance->tlsData()->setInterrupt();
+}
+
+void
+wasm::ResetInterruptState(JSContext* cx)
+{
+    auto runtimeInstances = cx->runtime()->wasmInstances.lock();
+    for (Instance* instance : runtimeInstances.get())
+        instance->tlsData()->resetInterrupt(cx);
+}

js/src/wasm/WasmCompartment.h

@@ -24,8 +24,6 @@
 namespace js {
 namespace wasm {
 
-typedef Vector<Instance*, 0, SystemAllocPolicy> InstanceVector;
-
 // wasm::Compartment lives in JSCompartment and contains the wasm-related
 // per-compartment state. wasm::Compartment tracks every live instance in the
 // compartment and must be notified, via registerInstance(), of any new
@@ -33,10 +31,11 @@ typedef Vector<Instance*, 0, SystemAllocPolicy> InstanceVector;
 
 class Compartment
 {
+    JSRuntime* runtime_;
     InstanceVector instances_;
 
   public:
-    explicit Compartment(Zone* zone);
+    explicit Compartment(JSRuntime* rt);
     ~Compartment();
 
     // Before a WasmInstanceObject can be considered fully constructed and
@@ -64,6 +63,19 @@ class Compartment
     void addSizeOfExcludingThis(mozilla::MallocSizeOf mallocSizeOf, size_t* compartmentTables);
 };
 
+// Interrupt all running wasm Instances that have been registered with
+// wasm::Compartments in the given JSContext.
+
+extern void
+InterruptRunningCode(JSContext* cx);
+
+// After a wasm Instance sees an interrupt request and calls
+// CheckForInterrupt(), it should call RunningCodeInterrupted() to clear the
+// interrupt request for all wasm Instances to avoid spurious trapping.
+
+void
+ResetInterruptState(JSContext* cx);
+
 } // namespace wasm
 } // namespace js
 

js/src/wasm/WasmFrameIter.cpp

@@ -50,9 +50,9 @@ WasmFrameIter::WasmFrameIter(JitActivation* activation, wasm::Frame* fp)
 
     if (activation->isWasmTrapping()) {
         code_ = &fp_->tls->instance->code();
-        MOZ_ASSERT(code_ == LookupCode(activation->wasmTrapPC()));
+        MOZ_ASSERT(code_ == LookupCode(activation->wasmTrapUnwoundPC()));
 
-        codeRange_ = code_->lookupFuncRange(activation->wasmTrapPC());
+        codeRange_ = code_->lookupFuncRange(activation->wasmTrapUnwoundPC());
         MOZ_ASSERT(codeRange_);
 
         lineOrBytecode_ = activation->wasmTrapBytecodeOffset();
@@ -1074,7 +1074,7 @@ ThunkedNativeToDescription(SymbolicAddress func)
     switch (func) {
       case SymbolicAddress::HandleDebugTrap:
       case SymbolicAddress::HandleThrow:
-      case SymbolicAddress::ReportTrap:
+      case SymbolicAddress::OnTrap:
       case SymbolicAddress::OldReportTrap:
       case SymbolicAddress::ReportOutOfBounds:
       case SymbolicAddress::ReportUnalignedAccess:

js/src/wasm/WasmInstance.cpp

@@ -406,7 +406,7 @@ Instance::Instance(JSContext* cx,
 #endif
     tlsData()->instance = this;
     tlsData()->cx = cx;
-    tlsData()->stackLimit = cx->stackLimitForJitCode(JS::StackForUntrustedScript);
+    tlsData()->resetInterrupt(cx);
     tlsData()->jumpTable = code_->tieringJumpTable();
 
     Tier callerTier = code_->bestTier();

js/src/wasm/WasmIonCompile.cpp

@@ -1033,7 +1033,9 @@ class FunctionCompiler
 
     void addInterruptCheck()
     {
-        // TODO
+        if (inDeadCode())
+            return;
+        curBlock_->add(MWasmInterruptCheck::New(alloc(), tlsPointer_, bytecodeOffset()));
     }
 
     MDefinition* extractSimdElement(unsigned lane, MDefinition* base, MIRType type, SimdSign sign)

js/src/wasm/WasmSignalHandlers.h

@@ -71,13 +71,10 @@ class MachExceptionHandler
 
 struct TrapData
 {
-    void* pc;
+    void* resumePC;
+    void* unwoundPC;
     Trap trap;
     uint32_t bytecodeOffset;
-
-    TrapData(void* pc, Trap trap, uint32_t bytecodeOffset)
-      : pc(pc), trap(trap), bytecodeOffset(bytecodeOffset)
-    {}
 };
 
 } // namespace wasm

js/src/wasm/WasmStubs.cpp

@@ -1358,6 +1358,26 @@ wasm::GenerateBuiltinThunk(MacroAssembler& masm, ABIFunctionType abiType, ExitRe
     return FinishOffsets(masm, offsets);
 }
 
+#if defined(JS_CODEGEN_ARM)
+static const LiveRegisterSet RegsToPreserve(
+    GeneralRegisterSet(Registers::AllMask & ~((uint32_t(1) << Registers::sp) |
+                                              (uint32_t(1) << Registers::pc))),
+    FloatRegisterSet(FloatRegisters::AllDoubleMask));
+static_assert(!SupportsSimd, "high lanes of SIMD registers need to be saved too.");
+#elif defined(JS_CODEGEN_MIPS32) || defined(JS_CODEGEN_MIPS64)
+static const LiveRegisterSet RegsToPreserve(
+    GeneralRegisterSet(Registers::AllMask & ~((uint32_t(1) << Registers::k0) |
+                                              (uint32_t(1) << Registers::k1) |
+                                              (uint32_t(1) << Registers::sp) |
+                                              (uint32_t(1) << Registers::zero))),
+    FloatRegisterSet(FloatRegisters::AllDoubleMask));
+static_assert(!SupportsSimd, "high lanes of SIMD registers need to be saved too.");
+#else
+static const LiveRegisterSet RegsToPreserve(
+    GeneralRegisterSet(Registers::AllMask & ~(uint32_t(1) << Registers::StackPointer)),
+    FloatRegisterSet(FloatRegisters::AllMask));
+#endif
+
 // Generate a stub which calls WasmReportTrap() and can be executed by having
 // the signal handler redirect PC from any trapping instruction.
 static bool
@@ -1367,16 +1387,37 @@ GenerateTrapExit(MacroAssembler& masm, Label* throwLabel, Offsets* offsets)
 
     offsets->begin = masm.currentOffset();
 
+    // Traps can only happen at well-defined program points. However, since
+    // traps may resume and the optimal assumption for the surrounding code is
+    // that registers are not clobbered, we need to preserve all registers in
+    // the trap exit. One simplifying assumption is that flags may be clobbered.
+    // Push a dummy word to use as return address below.
+    masm.push(ImmWord(0));
+    masm.setFramePushed(0);
+    masm.PushRegsInMask(RegsToPreserve);
+
     // We know that StackPointer is word-aligned, but not necessarily
     // stack-aligned, so we need to align it dynamically.
+    Register preAlignStackPointer = ABINonVolatileReg;
+    masm.moveStackPtrTo(preAlignStackPointer);
     masm.andToStackPtr(Imm32(~(ABIStackAlignment - 1)));
     if (ShadowStackSpace)
         masm.subFromStackPtr(Imm32(ShadowStackSpace));
 
     masm.assertStackAlignment(ABIStackAlignment);
-    masm.call(SymbolicAddress::ReportTrap);
+    masm.call(SymbolicAddress::OnTrap);
 
-    masm.jump(throwLabel);
+    // OnTrap returns null if control should transfer to the throw stub.
+    masm.branchTestPtr(Assembler::Zero, ReturnReg, ReturnReg, throwLabel);
+
+    // Otherwise, the return value is the TrapData::resumePC we must jump to.
+    // We must restore register state before jumping, which will clobber
+    // ReturnReg, so store ReturnReg in the above-reserved stack slot which we
+    // use to jump to via ret.
+    masm.moveToStackPtr(preAlignStackPointer);
+    masm.storePtr(ReturnReg, Address(masm.getStackPointer(), masm.framePushed()));
+    masm.PopRegsInMask(RegsToPreserve);
+    masm.ret();
 
     return FinishOffsets(masm, offsets);
 }
@@ -1460,26 +1501,6 @@ GenerateUnalignedExit(MacroAssembler& masm, Label* throwLabel, Offsets* offsets)
                                            offsets);
 }
 
-#if defined(JS_CODEGEN_ARM)
-static const LiveRegisterSet AllRegsExceptPCSP(
-    GeneralRegisterSet(Registers::AllMask & ~((uint32_t(1) << Registers::sp) |
-                                              (uint32_t(1) << Registers::pc))),
-    FloatRegisterSet(FloatRegisters::AllDoubleMask));
-static_assert(!SupportsSimd, "high lanes of SIMD registers need to be saved too.");
-#elif defined(JS_CODEGEN_MIPS32) || defined(JS_CODEGEN_MIPS64)
-static const LiveRegisterSet AllUserRegsExceptSP(
-    GeneralRegisterSet(Registers::AllMask & ~((uint32_t(1) << Registers::k0) |
-                                              (uint32_t(1) << Registers::k1) |
-                                              (uint32_t(1) << Registers::sp) |
-                                              (uint32_t(1) << Registers::zero))),
-    FloatRegisterSet(FloatRegisters::AllDoubleMask));
-static_assert(!SupportsSimd, "high lanes of SIMD registers need to be saved too.");
-#else
-static const LiveRegisterSet AllRegsExceptSP(
-    GeneralRegisterSet(Registers::AllMask & ~(uint32_t(1) << Registers::StackPointer)),
-    FloatRegisterSet(FloatRegisters::AllMask));
-#endif
-
 // Generate a stub that restores the stack pointer to what it was on entry to
 // the wasm activation, sets the return register to 'false' and then executes a
 // return which will return from this wasm activation to the caller. This stub
@@ -1638,6 +1659,7 @@ wasm::GenerateStubs(const ModuleEnvironment& env, const FuncImportVector& import
           case Trap::IndirectCallBadSig:
           case Trap::ImpreciseSimdConversion:
           case Trap::StackOverflow:
+          case Trap::CheckInterrupt:
           case Trap::ThrowReported:
             break;
           // The TODO list of "old" traps to convert to new traps:

js/src/wasm/WasmTypes.cpp

@@ -911,3 +911,23 @@ wasm::CreateTlsData(uint32_t globalDataLength)
 
     return UniqueTlsData(tlsData);
 }
+
+void
+TlsData::setInterrupt()
+{
+    interrupt = true;
+    stackLimit = UINTPTR_MAX;
+}
+
+bool
+TlsData::isInterrupted() const
+{
+    return interrupt || stackLimit == UINTPTR_MAX;
+}
+
+void
+TlsData::resetInterrupt(JSContext* cx)
+{
+    interrupt = false;
+    stackLimit = cx->stackLimitForJitCode(JS::StackForUntrustedScript);
+}

js/src/wasm/WasmTypes.h

@@ -85,12 +85,6 @@ using mozilla::PodEqual;
 using mozilla::Some;
 using mozilla::Unused;
 
-typedef Vector<uint32_t, 0, SystemAllocPolicy> Uint32Vector;
-typedef Vector<uint8_t, 0, SystemAllocPolicy> Bytes;
-typedef UniquePtr<Bytes> UniqueBytes;
-typedef UniquePtr<const Bytes> UniqueConstBytes;
-typedef Vector<char, 0, SystemAllocPolicy> UTF8Bytes;
-
 typedef int8_t I8x16[16];
 typedef int16_t I16x8[8];
 typedef int32_t I32x4[4];
@@ -104,6 +98,13 @@ class Module;
 class Instance;
 class Table;
 
+typedef Vector<uint32_t, 0, SystemAllocPolicy> Uint32Vector;
+typedef Vector<uint8_t, 0, SystemAllocPolicy> Bytes;
+typedef UniquePtr<Bytes> UniqueBytes;
+typedef UniquePtr<const Bytes> UniqueConstBytes;
+typedef Vector<char, 0, SystemAllocPolicy> UTF8Bytes;
+typedef Vector<Instance*, 0, SystemAllocPolicy> InstanceVector;
+
 // To call Vector::podResizeToFit, a type must specialize mozilla::IsPod
 // which is pretty verbose to do within js::wasm, so factor that process out
 // into a macro.
@@ -928,6 +929,10 @@ enum class Trap
     // the same over-recursed error as JS.
     StackOverflow,
 
+    // The wasm execution has potentially run too long and the engine must call
+    // CheckForInterrupt(). This trap is resumable.
+    CheckInterrupt,
+
     // Signal an error that was reported in C++ code.
     ThrowReported,
 
@@ -1374,7 +1379,7 @@ enum class SymbolicAddress
     ATan2D,
     HandleDebugTrap,
     HandleThrow,
-    ReportTrap,
+    OnTrap,
     OldReportTrap,
     ReportOutOfBounds,
     ReportUnalignedAccess,
@@ -1527,9 +1532,19 @@ struct TlsData
     // The containing JSContext.
     JSContext* cx;
 
-    // The native stack limit which is checked by prologues. Shortcut for
+    // Usually equal to cx->stackLimitForJitCode(JS::StackForUntrustedScript),
-    // cx->stackLimitForJitCode(JS::StackForUntrustedScript).
+    // but can be racily set to trigger immediate trap as an opportunity to
-    uintptr_t stackLimit;
+    // CheckForInterrupt without an additional branch.
+    Atomic<uintptr_t, mozilla::Relaxed> stackLimit;
+
+    // Set to 1 when wasm should call CheckForInterrupt.
+    Atomic<uint32_t, mozilla::Relaxed> interrupt;
+
+    // Methods to set, test and clear the above two fields. Both interrupt
+    // fields are Relaxed and so no consistency/ordering can be assumed.
+    void setInterrupt();
+    bool isInterrupted() const;
+    void resetInterrupt(JSContext* cx);
 
     // Pointer that should be freed (due to padding before the TlsData).
     void* allocatedBase;

js/xpconnect/src/XPCJSRuntime.cpp

@@ -2393,6 +2393,11 @@ JSReporter::CollectReports(WindowPaths* windowPaths,
         KIND_OTHER, rtStats.runtime.tracelogger,
         "The memory used for the tracelogger, including the graph and events.");
 
+    // Report the numbers for memory used by wasm Runtime state.
+    REPORT_BYTES(NS_LITERAL_CSTRING("wasm-runtime"),
+        KIND_OTHER, rtStats.runtime.wasmRuntime,
+        "The memory used for wasm runtime bookkeeping.");
+
     // Report the numbers for memory outside of compartments.
 
     REPORT_BYTES(NS_LITERAL_CSTRING("js-main-runtime/gc-heap/unused-chunks"),