[wasm][stack-switching] Propagate exceptions across stack segments

Exceptions should propagate inside the logical stack, which can consist
of multiple wasm stack segments. When the outermost frame of the current
segment is reached, pick up the parent stack and continue the search
from there, and update the state to reflect the implicit stack switch.

Drive-by: cleanups.

R=ahaas@chromium.org
CC=​fgm@chromium.org

Bug: v8:12191, v8:12960
Change-Id: Ia5cb39a6ae197fb68e635f986952419dc43c7b98
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3695376Reviewed-by: Jakob Kummerow <jkummerow@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/main@{#81208}

src/builtins/x64/builtins-x64.cc

@@ -2927,7 +2927,8 @@ void PrepareForBuiltinCall(MacroAssembler* masm, MemOperand GCScanSlotPlace,
                            Register current_int_param_slot,
                            Register current_float_param_slot,
                            Register valuetypes_array_ptr,
-                           Register wasm_instance, Register function_data) {
+                           Register wasm_instance, Register function_data,
+                           Register original_fp) {
   // Pushes and puts the values in order onto the stack before builtin calls for
   // the GenericJSToWasmWrapper.
   __ Move(GCScanSlotPlace, GCScanSlotCount);
@@ -2938,6 +2939,10 @@ void PrepareForBuiltinCall(MacroAssembler* masm, MemOperand GCScanSlotPlace,
   __ pushq(valuetypes_array_ptr);
   __ pushq(wasm_instance);
   __ pushq(function_data);
+  if (original_fp != rbp) {
+    // For stack-switching only: keep the pointer to the parent stack alive.
+    __ pushq(original_fp);
+  }
   // We had to prepare the parameters for the Call: we have to put the context
   // into rsi.
   __ LoadAnyTaggedField(
@@ -2951,9 +2956,13 @@ void RestoreAfterBuiltinCall(MacroAssembler* masm, Register function_data,
                              Register valuetypes_array_ptr,
                              Register current_float_param_slot,
                              Register current_int_param_slot,
-                             Register param_limit, Register current_param) {
+                             Register param_limit, Register current_param,
+                             Register original_fp) {
   // Pop and load values from the stack in order into the registers after
   // builtin calls for the GenericJSToWasmWrapper.
+  if (original_fp != rbp) {
+    __ popq(original_fp);
+  }
   __ popq(function_data);
   __ popq(wasm_instance);
   __ popq(valuetypes_array_ptr);
@@ -3094,14 +3103,14 @@ void RestoreParentSuspender(MacroAssembler* masm) {
   Register state = rbx;
   __ LoadTaggedSignedField(
       state, FieldOperand(suspender, WasmSuspenderObject::kStateOffset));
-  __ SmiCompare(state, Smi::FromInt(WasmSuspenderObject::Inactive));
+  __ SmiCompare(state, Smi::FromInt(WasmSuspenderObject::kInactive));
   __ j(equal, &parent_inactive, Label::kNear);
   __ Trap();
   __ bind(&parent_inactive);
 #endif
   __ StoreTaggedSignedField(
       FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
-      Smi::FromInt(WasmSuspenderObject::State::Active));
+      Smi::FromInt(WasmSuspenderObject::kActive));
   __ bind(&undefined);
   __ movq(masm->RootAsOperand(RootIndex::kActiveSuspender), suspender);
 }
@@ -3796,7 +3805,8 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   PrepareForBuiltinCall(masm, MemOperand(rbp, kGCScanSlotCountOffset),
                         kBuiltinCallGCScanSlotCount, current_param, param_limit,
                         current_int_param_slot, current_float_param_slot,
-                        valuetypes_array_ptr, wasm_instance, function_data);
+                        valuetypes_array_ptr, wasm_instance, function_data,
+                        original_fp);
 
   Label param_kWasmI32_not_smi;
   Label param_kWasmI64;
@@ -3827,7 +3837,8 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   __ movq(MemOperand(rbp, kHasRefTypesOffset), Immediate(1));
   RestoreAfterBuiltinCall(masm, function_data, wasm_instance,
                           valuetypes_array_ptr, current_float_param_slot,
-                          current_int_param_slot, param_limit, current_param);
+                          current_int_param_slot, param_limit, current_param,
+                          original_fp);
   __ jmp(&param_conversion_done);
 
   __ int3();
@@ -3839,7 +3850,8 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   __ AssertZeroExtended(param);
   RestoreAfterBuiltinCall(masm, function_data, wasm_instance,
                           valuetypes_array_ptr, current_float_param_slot,
-                          current_int_param_slot, param_limit, current_param);
+                          current_int_param_slot, param_limit, current_param,
+                          original_fp);
   __ movq(MemOperand(current_int_param_slot, 0), param);
   __ subq(current_int_param_slot, Immediate(kSystemPointerSize));
   __ jmp(&param_conversion_done);
@@ -3848,7 +3860,8 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
   __ Call(BUILTIN_CODE(masm->isolate(), BigIntToI64), RelocInfo::CODE_TARGET);
   RestoreAfterBuiltinCall(masm, function_data, wasm_instance,
                           valuetypes_array_ptr, current_float_param_slot,
-                          current_int_param_slot, param_limit, current_param);
+                          current_int_param_slot, param_limit, current_param,
+                          original_fp);
   __ movq(MemOperand(current_int_param_slot, 0), param);
   __ subq(current_int_param_slot, Immediate(kSystemPointerSize));
   __ jmp(&param_conversion_done);
@@ -3858,7 +3871,8 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
           RelocInfo::CODE_TARGET);
   RestoreAfterBuiltinCall(masm, function_data, wasm_instance,
                           valuetypes_array_ptr, current_float_param_slot,
-                          current_int_param_slot, param_limit, current_param);
+                          current_int_param_slot, param_limit, current_param,
+                          original_fp);
   // Clear higher bits.
   __ Xorpd(xmm1, xmm1);
   // Truncate float64 to float32.
@@ -3872,7 +3886,8 @@ void GenericJSToWasmWrapperHelper(MacroAssembler* masm, bool stack_switch) {
           RelocInfo::CODE_TARGET);
   RestoreAfterBuiltinCall(masm, function_data, wasm_instance,
                           valuetypes_array_ptr, current_float_param_slot,
-                          current_int_param_slot, param_limit, current_param);
+                          current_int_param_slot, param_limit, current_param,
+                          original_fp);
   __ Movsd(MemOperand(current_float_param_slot, 0), xmm0);
   __ subq(current_float_param_slot, Immediate(kSystemPointerSize));
   __ jmp(&param_conversion_done);
@@ -4034,7 +4049,7 @@ void Builtins::Generate_WasmSuspend(MacroAssembler* masm) {
   FillJumpBuffer(masm, jmpbuf, &resume);
   __ StoreTaggedSignedField(
       FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
-      Smi::FromInt(WasmSuspenderObject::Suspended));
+      Smi::FromInt(WasmSuspenderObject::kSuspended));
   jmpbuf = no_reg;
   // live: [rax, rbx, rcx]
 
@@ -4144,7 +4159,7 @@ void Builtins::Generate_WasmResume(MacroAssembler* masm) {
   Register state = rdx;
   __ LoadTaggedSignedField(
       state, FieldOperand(suspender, WasmSuspenderObject::kStateOffset));
-  __ SmiCompare(state, Smi::FromInt(WasmSuspenderObject::Suspended));
+  __ SmiCompare(state, Smi::FromInt(WasmSuspenderObject::kSuspended));
   __ j(equal, &suspender_is_suspended);
   __ Trap();  // TODO(thibaudm): Throw a wasm trap.
   closure = no_reg;
@@ -4173,7 +4188,7 @@ void Builtins::Generate_WasmResume(MacroAssembler* masm) {
   // -------------------------------------------
   __ StoreTaggedSignedField(
       FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
-      Smi::FromInt(WasmSuspenderObject::Active));
+      Smi::FromInt(WasmSuspenderObject::kActive));
   Register target_continuation = rdi;
   __ LoadAnyTaggedField(
       target_continuation,

src/execution/frames.h

@@ -1288,12 +1288,14 @@ class StackFrameIterator : public StackFrameIteratorBase {
   V8_EXPORT_PRIVATE void Advance();
   StackFrame* Reframe();
 
- private:
-  // Go back to the first frame.
-  void Reset(ThreadLocalTop* top);
 #if V8_ENABLE_WEBASSEMBLY
+  // Go to the first frame of this stack.
   void Reset(ThreadLocalTop* top, wasm::StackMemory* stack);
 #endif
+
+ private:
+  // Go back to the first frame.
+  void Reset(ThreadLocalTop* top);
 };
 
 // Iterator that supports iterating through all JavaScript frames.

src/execution/isolate.cc

@@ -1860,6 +1860,7 @@ class SetThreadInWasmFlagScope {
 Object Isolate::UnwindAndFindHandler() {
   // TODO(v8:12676): Fix gcmole failures in this function.
   DisableGCMole no_gcmole;
+  DisallowGarbageCollection no_gc;
 #if V8_ENABLE_WEBASSEMBLY
   // Create the {SetThreadInWasmFlagScope} first in this function so that its
   // destructor gets called after all the other destructors. It is important
@@ -1900,9 +1901,55 @@ Object Isolate::UnwindAndFindHandler() {
   bool catchable_by_js = is_catchable_by_javascript(exception);
   int visited_frames = 0;
 
+#if V8_ENABLE_WEBASSEMBLY
+  // Iterate the chain of stack segments for wasm stack switching.
+  WasmContinuationObject current_stack;
+  if (FLAG_experimental_wasm_stack_switching) {
+    current_stack =
+        WasmContinuationObject::cast(root(RootIndex::kActiveContinuation));
+  }
+#endif
+
   // Compute handler and stack unwinding information by performing a full walk
   // over the stack and dispatching according to the frame type.
   for (StackFrameIterator iter(this);; iter.Advance(), visited_frames++) {
+#if V8_ENABLE_WEBASSEMBLY
+    if (FLAG_experimental_wasm_stack_switching && iter.done()) {
+      // We reached the end of the current stack segment. Follow the linked-list
+      // of stacks to find the next frame, and perform the implicit stack
+      // switch.
+      auto stack = Managed<wasm::StackMemory>::cast(current_stack.stack());
+      // Mark this stack as empty.
+      stack.get()->jmpbuf()->sp = 0x0;
+      HeapObject parent = current_stack.parent();
+      DCHECK(!parent.IsUndefined());
+      current_stack = WasmContinuationObject::cast(parent);
+      wasm::StackMemory* parent_stack =
+          Managed<wasm::StackMemory>::cast(current_stack.stack()).get().get();
+      iter.Reset(thread_local_top(), parent_stack);
+
+      // Update the continuation and suspender state.
+      roots_table().slot(RootIndex::kActiveContinuation).store(current_stack);
+      WasmSuspenderObject suspender =
+          WasmSuspenderObject::cast(root(RootIndex::kActiveSuspender));
+      if (!suspender.parent().IsUndefined()) {
+        suspender.set_state(WasmSuspenderObject::State::kInactive);
+        auto parent_suspender = WasmSuspenderObject::cast(suspender.parent());
+        parent_suspender.set_state(WasmSuspenderObject::State::kActive);
+        // For now, assume that a suspender contains a single continuation.
+        // TODO(thibaudm): When core stack-switching is added, only update the
+        // suspender when we exit its outermost stack.
+        DCHECK_EQ(current_stack, parent_suspender.continuation());
+      }
+      roots_table().slot(RootIndex::kActiveSuspender).store(suspender.parent());
+      if (FLAG_trace_wasm_stack_switching) {
+        PrintF("Switch to stack #%d (unwind)\n", parent_stack->id());
+      }
+      uintptr_t limit =
+          reinterpret_cast<uintptr_t>(parent_stack->jmpbuf()->stack_limit);
+      stack_guard()->SetStackLimit(limit);
+    }
+#endif
     // Handler must exist.
     DCHECK(!iter.done());
 

src/runtime/runtime-wasm.cc

@@ -781,7 +781,7 @@ namespace {
 void SyncStackLimit(Isolate* isolate) {
   DisallowGarbageCollection no_gc;
   auto continuation = WasmContinuationObject::cast(
-      *isolate->roots_table().slot(RootIndex::kActiveContinuation));
+      isolate->root(RootIndex::kActiveContinuation));
   auto stack = Managed<wasm::StackMemory>::cast(continuation.stack()).get();
   if (FLAG_trace_wasm_stack_switching) {
     PrintF("Switch to stack #%d\n", stack->id());
@@ -799,10 +799,9 @@ RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
   Handle<WasmSuspenderObject> suspender = args.at<WasmSuspenderObject>(0);
 
   // Update the continuation state.
-  auto parent =
-      handle(WasmContinuationObject::cast(
-                 *isolate->roots_table().slot(RootIndex::kActiveContinuation)),
-             isolate);
+  auto parent = handle(WasmContinuationObject::cast(
+                           isolate->root(RootIndex::kActiveContinuation)),
+                       isolate);
   Handle<WasmContinuationObject> target =
       WasmContinuationObject::New(isolate, parent);
   auto target_stack =
@@ -816,9 +815,9 @@ RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
   suspender->set_parent(HeapObject::cast(*active_suspender_slot));
   if (!(*active_suspender_slot).IsUndefined()) {
     WasmSuspenderObject::cast(*active_suspender_slot)
-        .set_state(WasmSuspenderObject::Inactive);
+        .set_state(WasmSuspenderObject::kInactive);
   }
-  suspender->set_state(WasmSuspenderObject::State::Active);
+  suspender->set_state(WasmSuspenderObject::kActive);
   suspender->set_continuation(*target);
   active_suspender_slot.store(*suspender);
 

src/wasm/wasm-objects.cc

@@ -1804,7 +1804,7 @@ Handle<WasmSuspenderObject> WasmSuspenderObject::New(Isolate* isolate) {
   // which it will wrap the imports/exports, allocate in old space too.
   auto suspender = Handle<WasmSuspenderObject>::cast(
       isolate->factory()->NewJSObject(suspender_cons, AllocationType::kOld));
-  suspender->set_state(Inactive);
+  suspender->set_state(kInactive);
   // Instantiate the callable object which resumes this Suspender. This will be
   // used implicitly as the onFulfilled callback of the returned JS promise.
   Handle<WasmOnFulfilledData> function_data =

src/wasm/wasm-objects.h

@@ -1033,7 +1033,7 @@ class WasmContinuationObject
 class WasmSuspenderObject
     : public TorqueGeneratedWasmSuspenderObject<WasmSuspenderObject, JSObject> {
  public:
-  enum State : int { Inactive = 0, Active, Suspended };
+  enum State : int { kInactive = 0, kActive, kSuspended };
   static Handle<WasmSuspenderObject> New(Isolate* isolate);
   // TODO(thibaudm): returnPromiseOnSuspend & suspendOnReturnedPromise.
   DECL_PRINTER(WasmSuspenderObject)

test/mjsunit/wasm/stack-switching.js

@@ -249,3 +249,49 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   let combined_promise = wrapped_export();
   combined_promise.then(v => assertEquals(0.5, v));
 })();
+
+// Throw an exception after the initial prompt.
+(function TestStackSwitchException1() {
+  print(arguments.callee.name);
+  let builder = new WasmModuleBuilder();
+  let tag = builder.addTag(kSig_v_v);
+  builder.addFunction("throw", kSig_i_v)
+      .addBody([kExprThrow, tag]).exportFunc();
+  let instance = builder.instantiate();
+  let suspender = new WebAssembly.Suspender();
+  let wrapper = suspender.returnPromiseOnSuspend(instance.exports.throw);
+  try {
+    wrapper();
+    assertUnreachable();
+  } catch (e) {
+    assertTrue(e instanceof WebAssembly.Exception);
+  }
+})();
+
+// Throw an exception after the first resume event, which propagates to the
+// promise wrapper.
+(function TestStackSwitchException2() {
+  print(arguments.callee.name);
+  let tag = new WebAssembly.Tag({parameters: []});
+  let builder = new WasmModuleBuilder();
+  import_index = builder.addImport('m', 'import', kSig_i_v);
+  tag_index = builder.addImportedTag('m', 'tag', kSig_v_v);
+  builder.addFunction("test", kSig_i_v)
+      .addBody([
+          kExprCallFunction, import_index,
+          kExprThrow, tag_index
+      ]).exportFunc();
+  let suspender = new WebAssembly.Suspender();
+  function js_import() {
+    return Promise.resolve(42);
+  };
+  let wasm_js_import = new WebAssembly.Function(
+      {parameters: [], results: ['externref']}, js_import);
+  let suspending_wasm_js_import =
+      suspender.suspendOnReturnedPromise(wasm_js_import);
+
+  let instance = builder.instantiate({m: {import: suspending_wasm_js_import, tag: tag}});
+  let wrapped_export = suspender.returnPromiseOnSuspend(instance.exports.test);
+  let combined_promise = wrapped_export();
+  assertThrowsAsync(combined_promise, WebAssembly.Exception);
+})();