Merge pull request #107 from bytecodealliance/dicej/fiber-switch-save-restore

fix save/restore of PC/FP/SP/limit on fiber switch

Author: dicej

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -1398,7 +1398,7 @@ impl ComponentInstance {
                                     break;
                                 } else {
                                     unsafe {
-                                        suspend_fiber(fiber.suspend, fiber.stack_limit, None)?;
+                                        suspend_fiber(fiber.suspend, None)?;
                                     };
                                 }
                             }
@@ -3310,7 +3310,6 @@ pub(crate) struct AsyncCx {
         Option<*mut dyn VMStore>,
         (Option<*mut dyn VMStore>, Result<()>),
     >,
-    current_stack_limit: *mut usize,
     current_poll_cx: *mut PollContext,
     track_pkey_context_switch: bool,
 }
@@ -3327,7 +3326,6 @@ impl AsyncCx {
         } else {
             Some(Self {
                 current_suspend: store.concurrent_async_state().current_suspend.get(),
-                current_stack_limit: store.vm_store_context().stack_limit.get(),
                 current_poll_cx,
                 track_pkey_context_switch: store.has_pkey(),
             })
@@ -3365,7 +3363,7 @@ impl AsyncCx {
         } else {
             ProtectionMask::all()
         };
-        let store = suspend_fiber(self.current_suspend, self.current_stack_limit, store);
+        let store = suspend_fiber(self.current_suspend, store);
         if self.track_pkey_context_switch {
             mpk::allow(previous_mask);
         }
@@ -3550,7 +3548,6 @@ struct StoreFiber<'a> {
         Option<*mut dyn VMStore>,
         (Option<*mut dyn VMStore>, Result<()>),
     >,
-    stack_limit: *mut usize,
     instance: Option<RuntimeComponentInstanceIndex>,
 }
 
@@ -3607,7 +3604,6 @@ unsafe fn make_fiber<'a>(
             .concurrent_async_state()
             .current_suspend
             .get(),
-        stack_limit: (*store).vm_store_context().stack_limit.get(),
         instance,
     })
 }
@@ -3631,7 +3627,6 @@ unsafe fn resume_fiber_raw<'a>(
     }
 
     let _reset_suspend = Reset((*fiber).suspend, *(*fiber).suspend);
-    let _reset_stack_limit = Reset((*fiber).stack_limit, *(*fiber).stack_limit);
     let state = Some((*fiber).state.take().unwrap().push());
     let restore = Restore { fiber, state };
     (*restore.fiber)
@@ -3667,11 +3662,9 @@ unsafe fn suspend_fiber(
         Option<*mut dyn VMStore>,
         (Option<*mut dyn VMStore>, Result<()>),
     >,
-    stack_limit: *mut usize,
     store: Option<*mut dyn VMStore>,
 ) -> Result<Option<*mut dyn VMStore>> {
     let _reset_suspend = Reset(suspend, *suspend);
-    let _reset_stack_limit = Reset(stack_limit, *stack_limit);
     assert!(!(*suspend).is_null());
     let (store, result) = (**suspend).suspend(store);
     result?;

crates/wasmtime/src/runtime/vm/traphandlers.rs

@@ -439,17 +439,19 @@ mod call_thread_state {
         #[cfg(all(has_native_signals, unix))]
         pub(crate) async_guard_range: Range<*mut u8>,
 
-        // The values of `VMStoreContext::last_wasm_{exit_{pc,fp},entry_sp}` for
+        // The values of
+        // `VMStoreContext::{stack_limit,last_wasm_{exit_{pc,fp},entry_sp}}` for
         // the *previous* `CallThreadState` for this same store/limits. Our
         // *current* last wasm PC/FP/SP are saved in `self.vm_store_context`. We
         // save a copy of the old registers here because the `VMStoreContext`
         // typically doesn't change across nested calls into Wasm (i.e. they are
         // typically calls back into the same store and `self.vm_store_context
         // == self.prev.vm_store_context`) and we must to maintain the list of
         // contiguous-Wasm-frames stack regions for backtracing purposes.
-        old_last_wasm_exit_fp: Cell<usize>,
-        old_last_wasm_exit_pc: Cell<usize>,
-        old_last_wasm_entry_fp: Cell<usize>,
+        pub(super) old_stack_limit: Cell<usize>,
+        pub(super) old_last_wasm_exit_fp: Cell<usize>,
+        pub(super) old_last_wasm_exit_pc: Cell<usize>,
+        pub(super) old_last_wasm_entry_fp: Cell<usize>,
     }
 
     impl Drop for CallThreadState {
@@ -460,6 +462,7 @@ mod call_thread_state {
 
             unsafe {
                 let cx = self.vm_store_context.as_ref();
+                *cx.stack_limit.get() = self.old_stack_limit.get();
                 *cx.last_wasm_exit_fp.get() = self.old_last_wasm_exit_fp.get();
                 *cx.last_wasm_exit_pc.get() = self.old_last_wasm_exit_pc.get();
                 *cx.last_wasm_entry_fp.get() = self.old_last_wasm_entry_fp.get();
@@ -500,6 +503,7 @@ mod call_thread_state {
                 #[cfg(all(has_native_signals, unix))]
                 async_guard_range,
                 prev: Cell::new(ptr::null()),
+                old_stack_limit: Cell::new(unsafe { *vm_store_context.as_ref().stack_limit.get() }),
                 old_last_wasm_exit_fp: Cell::new(unsafe {
                     *vm_store_context.as_ref().last_wasm_exit_fp.get()
                 }),
@@ -825,6 +829,21 @@ pub(crate) mod tls {
 
     pub use raw::initialize as tls_eager_initialize;
 
+    #[cfg(feature = "async")]
+    impl CallThreadState {
+        unsafe fn swap(&self) {
+            unsafe fn swap(a: &core::cell::UnsafeCell<usize>, b: &core::cell::Cell<usize>) {
+                *a.get() = b.replace(*a.get());
+            }
+
+            let cx = self.vm_store_context.as_ref();
+            swap(&cx.stack_limit, &self.old_stack_limit);
+            swap(&cx.last_wasm_exit_fp, &self.old_last_wasm_exit_fp);
+            swap(&cx.last_wasm_exit_pc, &self.old_last_wasm_exit_pc);
+            swap(&cx.last_wasm_entry_fp, &self.old_last_wasm_entry_fp);
+        }
+    }
+
     /// Opaque state used to persist the state of the `CallThreadState`
     /// activations associated with a fiber stack that's used as part of an
     /// async wasm call.
@@ -873,6 +892,13 @@ pub(crate) mod tls {
             // list as stored in the state of this current thread.
             let ret = PreviousAsyncWasmCallState { state: raw::get() };
             let mut ptr = self.state;
+
+            if let Some(state) = ptr.as_ref() {
+                // Swap the current PC/FP/SP/etc. with those in the oldest
+                // activation so we can restore them later in
+                // `PreviousAsyncWasmCallState::restore`
+                state.swap();
+            }
             while let Some(state) = ptr.as_ref() {
                 ptr = state.prev.replace(core::ptr::null_mut());
                 state.push();
@@ -932,14 +958,22 @@ pub(crate) mod tls {
                 // this loop is finished.
                 let ptr = raw::get();
                 if ptr == thread_head {
+                    if let Some(state) = ret.state.as_ref() {
+                        // Swap the current PC/FP/SP/etc. with those in the
+                        // oldest activation, which reverses what we did in
+                        // `AsyncWasmCallState::push`.
+                        state.swap();
+                    }
+
                     break ret;
                 }
 
                 // Pop this activation from the current thread's TLS state, and
                 // then afterwards push it onto our own linked list within this
-                // `AsyncWasmCallState`. Note that the linked list in `AsyncWasmCallState` is stored
-                // in reverse order so a subsequent `push` later on pushes
-                // everything in the right order.
+                // `AsyncWasmCallState`. Note that the linked list in
+                // `AsyncWasmCallState` is stored in reverse order so a
+                // subsequent `push` later on pushes everything in the right
+                // order.
                 (*ptr).pop();
                 if let Some(state) = ret.state.as_ref() {
                     (*ptr).prev.set(state);