[wasm] Set up Suspender chain for stack switching

- Add an ActiveSuspender root, similar to the ActiveContinuation root.
- Add the missing "parent" field to the Suspender, which points to the
outer Suspender when they are nested, and update that field when
entering a new Suspender.
- Add the missing "state" field and update it when the state of the
Suspender changes.

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

Bug: v8:12191
Change-Id: I7a95f44f81390a347c6ef252ec6184fb4f0b0455
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3345003Reviewed-by: Andreas Haas <ahaas@chromium.org>
Reviewed-by: Dominik Inführ <dinfuehr@chromium.org>
Reviewed-by: Jakob Kummerow <jkummerow@chromium.org>
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/main@{#78582}

src/builtins/x64/builtins-x64.cc

@@ -3732,17 +3732,23 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   // -------------------------------------------
   // Allocate a new continuation.
   // -------------------------------------------
+  Register suspender = rbx;
+  __ LoadAnyTaggedField(
+      suspender,
+      FieldOperand(function_data, WasmExportedFunctionData::kSuspenderOffset));
   MemOperand GCScanSlotPlace =
       MemOperand(rbp, BuiltinWasmWrapperConstants::kGCScanSlotCountOffset);
-  __ Move(GCScanSlotPlace, 2);
+  __ Move(GCScanSlotPlace, 3);
   __ Push(wasm_instance);
   __ Push(function_data);
+  __ Push(suspender);  // Argument.
   __ Move(kContextRegister, Smi::zero());
   __ CallRuntime(Runtime::kWasmAllocateContinuation);
   __ Pop(function_data);
   __ Pop(wasm_instance);
   STATIC_ASSERT(kReturnRegister0 == rax);
   Register target_continuation = rax;
+  suspender = no_reg;
   // live: [rsi, rdi, rax]
 
   // -------------------------------------------
@@ -3831,6 +3837,31 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   active_continuation = no_reg;
   foreign_jmpbuf = no_reg;
   wasm_instance = no_reg;
+  // live: [rsi]
+
+  // -------------------------------------------
+  // Restore parent suspender.
+  // -------------------------------------------
+  suspender = rax;
+  __ LoadRoot(suspender, RootIndex::kActiveSuspender);
+  __ LoadAnyTaggedField(
+      suspender, FieldOperand(suspender, WasmSuspenderObject::kParentOffset));
+  __ CompareRoot(suspender, RootIndex::kUndefinedValue);
+  Label undefined;
+  __ j(equal, &undefined, Label::kNear);
+#ifdef DEBUG
+  // Check that the parent suspender is inactive.
+  Label parent_inactive;
+  __ cmpq(FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
+          Immediate(WasmSuspenderObject::Inactive));
+  __ j(equal, &parent_inactive, Label::kNear);
+  __ Trap();
+  __ bind(&parent_inactive);
+#endif
+  __ movq(FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
+          Immediate(WasmSuspenderObject::State::Active));
+  __ bind(&undefined);
+  __ movq(masm->RootAsOperand(RootIndex::kActiveSuspender), suspender);
 
   // -------------------------------------------
   // Epilogue.

src/diagnostics/objects-printer.cc

@@ -1929,6 +1929,7 @@ void WasmExportedFunctionData::WasmExportedFunctionDataPrint(std::ostream& os) {
   os << "\n - function_index: " << function_index();
   os << "\n - signature: " << Brief(signature());
   os << "\n - wrapper_budget: " << wrapper_budget();
+  os << "\n - suspender: " << suspender();
   os << "\n";
 }
 

src/heap/factory.cc

@@ -1600,6 +1600,7 @@ Handle<WasmExportedFunctionData> Factory::NewWasmExportedFunctionData(
   result.set_c_wrapper_code(*BUILTIN_CODE(isolate(), Illegal),
                             SKIP_WRITE_BARRIER);
   result.set_packed_args_size(0);
+  result.set_suspender(*undefined_value());
   return handle(result, isolate());
 }
 

src/heap/setup-heap-internal.cc

@@ -811,6 +811,7 @@ void Heap::CreateInitialObjects() {
   set_shared_wasm_memories(roots.empty_weak_array_list());
 #ifdef V8_ENABLE_WEBASSEMBLY
   set_active_continuation(roots.undefined_value());
+  set_active_suspender(roots.undefined_value());
 #endif  // V8_ENABLE_WEBASSEMBLY
 
   set_script_list(roots.empty_weak_array_list());

src/roots/roots.h

@@ -318,7 +318,8 @@ class Symbol;
     PendingOptimizeForTestBytecode)                                        \
   V(ArrayList, basic_block_profiling_data, BasicBlockProfilingData)        \
   V(WeakArrayList, shared_wasm_memories, SharedWasmMemories)               \
-  IF_WASM(V, HeapObject, active_continuation, ActiveContinuation)
+  IF_WASM(V, HeapObject, active_continuation, ActiveContinuation)          \
+  IF_WASM(V, HeapObject, active_suspender, ActiveSuspender)
 
 // Entries in this list are limited to Smis and are not visited during GC.
 #define SMI_ROOT_LIST(V)                                                       \

src/runtime/runtime-wasm.cc

@@ -714,19 +714,35 @@ void SyncStackLimit(Isolate* isolate) {
 }  // namespace
 
 // Allocate a new continuation, and prepare for stack switching by updating the
-// active continuation and setting the stack limit.
+// active continuation, active suspender and stack limit.
 RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
   CHECK(FLAG_experimental_wasm_stack_switching);
   HandleScope scope(isolate);
+  CONVERT_ARG_HANDLE_CHECKED(WasmSuspenderObject, suspender, 0);
+
+  // Update the continuation state.
   auto parent =
       handle(WasmContinuationObject::cast(
                  *isolate->roots_table().slot(RootIndex::kActiveContinuation)),
              isolate);
-  auto target = WasmContinuationObject::New(isolate, *parent);
+  Handle<WasmContinuationObject> target =
+      WasmContinuationObject::New(isolate, *parent);
   auto target_stack =
       Managed<wasm::StackMemory>::cast(target->stack()).get().get();
   isolate->wasm_stacks()->Add(target_stack);
   isolate->roots_table().slot(RootIndex::kActiveContinuation).store(*target);
+
+  // Update the suspender state.
+  FullObjectSlot active_suspender_slot =
+      isolate->roots_table().slot(RootIndex::kActiveSuspender);
+  suspender->set_parent(HeapObject::cast(*active_suspender_slot));
+  if (!(*active_suspender_slot).IsUndefined()) {
+    WasmSuspenderObject::cast(*active_suspender_slot)
+        .set_state(WasmSuspenderObject::Inactive);
+  }
+  suspender->set_state(WasmSuspenderObject::State::Active);
+  active_suspender_slot.store(*suspender);
+
   SyncStackLimit(isolate);
   return *target;
 }

src/runtime/runtime.h

@@ -596,7 +596,7 @@ namespace internal {
   F(WasmDebugBreak, 0, 1)             \
   F(WasmAllocateRtt, 3, 1)            \
   F(WasmArrayCopy, 5, 1)              \
-  F(WasmAllocateContinuation, 0, 1)   \
+  F(WasmAllocateContinuation, 1, 1)   \
   F(WasmSyncStackLimit, 0, 1)
 
 #define FOR_EACH_INTRINSIC_WASM_TEST(F, I) \

src/wasm/wasm-js.cc

@@ -2584,10 +2584,15 @@ void WebAssemblySuspenderReturnPromiseOnSuspend(
       i::WasmInstanceObject::cast(data.internal().ref()), i_isolate);
   i::Handle<i::CodeT> wrapper =
       BUILTIN_CODE(i_isolate, WasmReturnPromiseOnSuspend);
-  i::Handle<i::JSObject> result =
+  // Upcast to JSFunction to re-use the existing ToLocal helper below.
+  i::Handle<i::JSFunction> result =
       i::Handle<i::WasmExternalFunction>::cast(i::WasmExportedFunction::New(
           i_isolate, instance, index,
           static_cast<int>(data.sig()->parameter_count()), wrapper));
+  EXTRACT_THIS(suspender, WasmSuspenderObject);
+  auto function_data = i::WasmExportedFunctionData::cast(
+      result->shared().function_data(kAcquireLoad));
+  function_data.set_suspender(*suspender);
   args.GetReturnValue().Set(Utils::ToLocal(result));
 }
 

src/wasm/wasm-objects.cc

@@ -1790,6 +1790,8 @@ Handle<WasmSuspenderObject> WasmSuspenderObject::New(Isolate* isolate) {
   auto suspender = Handle<WasmSuspenderObject>::cast(
       isolate->factory()->NewJSObject(suspender_cons, AllocationType::kOld));
   suspender->set_continuation(ReadOnlyRoots(isolate).undefined_value());
+  suspender->set_parent(ReadOnlyRoots(isolate).undefined_value());
+  suspender->set_state(Inactive);
   return suspender;
 }
 

src/wasm/wasm-objects.h

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

src/wasm/wasm-objects.tq

@@ -67,6 +67,9 @@ extern class WasmExportedFunctionData extends WasmFunctionData {
   @if(V8_EXTERNAL_CODE_SPACE) c_wrapper_code: CodeDataContainer;
   @ifnot(V8_EXTERNAL_CODE_SPACE) c_wrapper_code: Code;
   packed_args_size: Smi;
+  // Functions returned by suspender.returnPromiseOnSuspend() have this field
+  // set to the host suspender object.
+  suspender: WasmSuspenderObject|Undefined;
 }
 
 extern class WasmJSFunctionData extends WasmFunctionData {
@@ -98,6 +101,8 @@ extern class WasmContinuationObject extends Struct {
 
 extern class WasmSuspenderObject extends JSObject {
   continuation: WasmContinuationObject|Undefined;
+  parent: WasmSuspenderObject|Undefined;
+  state: Smi;  // 0: Inactive, 1: Active, 2: Suspended.
 }
 
 extern class WasmExceptionTag extends Struct {