[wasm] Update stack limits on stack switch

This is done in a separate runtime function call for now, so that we can
update the limit under the ExecutionAcess lock.

Also set the thread-in-wasm flag before calling the wasm function.

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

Bug: v8:12191
Change-Id: I914856bc261fa0f75e93620bc6597bd28bec0695
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3250902
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/main@{#77613}

src/builtins/x64/builtins-x64.cc

@@ -3642,11 +3642,9 @@ void Builtins::Generate_GenericJSToWasmWrapper(MacroAssembler* masm) {
 
 namespace {
 // Helper function for WasmReturnPromiseOnSuspend.
-void LoadJumpBuffer(MacroAssembler* masm, Register jmpbuf,
-                    Register stack_limit_address, Register stack_limit_tmp) {
+void LoadJumpBuffer(MacroAssembler* masm, Register jmpbuf) {
   __ movq(rsp, MemOperand(jmpbuf, wasm::kJmpBufSpOffset));
-  __ movq(stack_limit_tmp, MemOperand(jmpbuf, wasm::kJmpBufStackLimitOffset));
-  __ movq(MemOperand(stack_limit_address, 0), stack_limit_tmp);
+  // The stack limit is set separately under the ExecutionAccess lock.
   // TODO(thibaudm): Reload live registers.
 }
 }  // namespace
@@ -3712,7 +3710,7 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   Register stack_limit_address = rcx;
   __ movq(stack_limit_address,
           FieldOperand(wasm_instance,
-                       WasmInstanceObject::kStackLimitAddressOffset));
+                       WasmInstanceObject::kRealStackLimitAddressOffset));
   Register stack_limit = rdx;
   __ movq(stack_limit, MemOperand(stack_limit_address, 0));
   __ movq(MemOperand(jmpbuf, wasm::kJmpBufStackLimitOffset), stack_limit);
@@ -3720,7 +3718,7 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   foreign_jmpbuf = no_reg;
   stack_limit = no_reg;
   stack_limit_address = no_reg;
-  // live: [rsi, rdi, rax, rcx]
+  // live: [rsi, rdi, rax]
 
   // -------------------------------------------
   // Allocate a new continuation.
@@ -3749,16 +3747,11 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
       target_jmpbuf,
       FieldOperand(foreign_jmpbuf, Foreign::kForeignAddressOffset),
       kForeignForeignAddressTag, r8);
-  stack_limit_address = rcx;
-  __ movq(stack_limit_address,
-          FieldOperand(wasm_instance,
-                       WasmInstanceObject::kStackLimitAddressOffset));
   // Switch stack!
-  LoadJumpBuffer(masm, target_jmpbuf, stack_limit_address, rdx);
+  LoadJumpBuffer(masm, target_jmpbuf);
   __ movq(rbp, rsp);  // New stack, there is no frame yet.
   foreign_jmpbuf = no_reg;
   target_jmpbuf = no_reg;
-  stack_limit_address = no_reg;
   // live: [rsi, rdi]
 
   // -------------------------------------------
@@ -3766,6 +3759,12 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   // -------------------------------------------
   // TODO(thibaudm): Handle arguments.
   // TODO(thibaudm): Handle GC.
+  // Set thread_in_wasm_flag.
+  Register thread_in_wasm_flag_addr = rax;
+  __ movq(
+      thread_in_wasm_flag_addr,
+      MemOperand(kRootRegister, Isolate::thread_in_wasm_flag_address_offset()));
+  __ movl(MemOperand(thread_in_wasm_flag_addr, 0), Immediate(1));
   Register function_entry = function_data;
   __ LoadExternalPointerField(
       function_entry,
@@ -3775,6 +3774,12 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   __ Push(wasm_instance);
   __ call(function_entry);
   __ Pop(wasm_instance);
+  // Unset thread_in_wasm_flag.
+  __ movq(
+      thread_in_wasm_flag_addr,
+      MemOperand(kRootRegister, Isolate::thread_in_wasm_flag_address_offset()));
+  __ movl(MemOperand(thread_in_wasm_flag_addr, 0), Immediate(0));
+  thread_in_wasm_flag_addr = no_reg;
   function_entry = no_reg;
   function_data = no_reg;
   // live: [rsi]
@@ -3819,16 +3824,16 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   __ LoadExternalPointerField(
       jmpbuf, FieldOperand(foreign_jmpbuf, Foreign::kForeignAddressOffset),
       kForeignForeignAddressTag, r8);
-  stack_limit_address = rcx;
-  __ movq(stack_limit_address,
-          FieldOperand(wasm_instance,
-                       WasmInstanceObject::kStackLimitAddressOffset));
   // Switch stack!
-  LoadJumpBuffer(masm, jmpbuf, stack_limit_address, rdx);
+  LoadJumpBuffer(masm, jmpbuf);
   __ leaq(rbp, Operand(rsp, (kNumSpillSlots + 1) * kSystemPointerSize));
+  __ Push(wasm_instance);  // Spill.
+  __ Push(wasm_instance);  // First arg.
+  __ Move(kContextRegister, Smi::zero());
+  __ CallRuntime(Runtime::kWasmSyncStackLimit);
+  __ Pop(wasm_instance);
   parent = no_reg;
   active_continuation = no_reg;
-  stack_limit_address = no_reg;
   foreign_jmpbuf = no_reg;
   wasm_instance = no_reg;
 

src/runtime/runtime-wasm.cc

@@ -681,6 +681,25 @@ RUNTIME_FUNCTION(Runtime_WasmArrayCopy) {
   return ReadOnlyRoots(isolate).undefined_value();
 }
 
+namespace {
+// Synchronize the stack limit with the active continuation for stack-switching.
+// This can be done before or after changing the stack pointer itself, as long
+// as we update both before the next stack check.
+// {StackGuard::SetStackLimit} doesn't update the value of the jslimit if it
+// contains a sentinel value, and it is also thread-safe. So if an interrupt is
+// requested before, during or after this call, it will be preserved and handled
+// at the next stack check.
+void SyncStackLimit(Isolate* isolate, Handle<WasmInstanceObject> instance) {
+  auto jmpbuf = Managed<wasm::JumpBuffer>::cast(
+                    instance->active_continuation().managed_jmpbuf())
+                    .get();
+  uintptr_t limit = reinterpret_cast<uintptr_t>(jmpbuf->stack_limit);
+  isolate->stack_guard()->SetStackLimit(limit);
+}
+}  // namespace
+
+// Allocate a new continuation, and prepare for stack switching by updating the
+// active continuation and setting the stack limit.
 RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
   CHECK(FLAG_experimental_wasm_stack_switching);
   DCHECK_EQ(1, args.length());
@@ -689,8 +708,18 @@ RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
   auto parent = instance->active_continuation();
   auto target = WasmContinuationObject::New(isolate, parent);
   instance->set_active_continuation(*target);
+  SyncStackLimit(isolate, instance);
   return *target;
 }
 
+// Update the stack limit after a stack switch, and preserve pending interrupts.
+RUNTIME_FUNCTION(Runtime_WasmSyncStackLimit) {
+  CHECK(FLAG_experimental_wasm_stack_switching);
+  HandleScope scope(isolate);
+  CONVERT_ARG_HANDLE_CHECKED(WasmInstanceObject, instance, 0);
+  SyncStackLimit(isolate, instance);
+  return ReadOnlyRoots(isolate).undefined_value();
+}
+
 }  // namespace internal
 }  // namespace v8

src/runtime/runtime.h

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