Remove StoreOpaque::async_yield_impl (#11482)

alexcrichton · web-flow · commit 6e21cf1f47d8 · 2025-08-21T00:52:31.000Z
* Remove `StoreOpaque::async_yield_impl`

Now that we have all this fancy support for natively running `async`
things this commit refactors fuel/epochs to use it. This simplifies the
`VMStore` trait, removes a usage of `block_on`, and helps keep the
`async` boundary close to the libcall entrypoint rather than further
down the stack. This all in turn enables using rustc to check our
stack-locals for non-`Send` values instead of pinky promising that we're
doing the right thing everywhere.

* Clean up some code movement and comments

* Fix a merge conflict
diff --git a/crates/wasmtime/src/runtime/fiber.rs b/crates/wasmtime/src/runtime/fiber.rs
@@ -343,35 +343,7 @@ impl<T> StoreContextMut<'_, T> {
     }
 }
 
-impl<T> crate::store::StoreInner<T> {
-    /// Blocks on the future computed by `f`.
-    ///
-    /// # Panics
-    ///
-    /// Panics if this is invoked outside the context of a fiber.
-    pub(crate) fn block_on<R>(
-        &mut self,
-        f: impl FnOnce(StoreContextMut<'_, T>) -> Pin<Box<dyn Future<Output = R> + Send + '_>>,
-    ) -> Result<R> {
-        BlockingContext::with(self, |store, cx| {
-            cx.block_on(f(StoreContextMut(store)).as_mut())
-        })
-    }
-}
-
 impl StoreOpaque {
-    /// Blocks on the future computed by `f`.
-    ///
-    /// # Panics
-    ///
-    /// Panics if this is invoked outside the context of a fiber.
-    pub(crate) fn block_on<R>(
-        &mut self,
-        f: impl FnOnce(&mut Self) -> Pin<Box<dyn Future<Output = R> + Send + '_>>,
-    ) -> Result<R> {
-        BlockingContext::with(self, |store, cx| cx.block_on(f(store).as_mut()))
-    }
-
     /// Creates a `BlockingContext` suitable for blocking on futures or
     /// suspending the current fiber.
     ///
diff --git a/crates/wasmtime/src/runtime/store.rs b/crates/wasmtime/src/runtime/store.rs
@@ -96,7 +96,7 @@ use crate::runtime::vm::{
     SignalHandler, StoreBox, Unwind, VMContext, VMFuncRef, VMGcRef, VMStore, VMStoreContext,
 };
 use crate::trampoline::VMHostGlobalContext;
-use crate::{Engine, Module, Trap, Val, ValRaw, module::ModuleRegistry};
+use crate::{Engine, Module, Val, ValRaw, module::ModuleRegistry};
 use crate::{Global, Instance, Memory, Table, Uninhabited};
 use alloc::sync::Arc;
 use core::fmt;
@@ -324,6 +324,9 @@ enum CallHookInner<T: 'static> {
 /// the deadline for a Store during execution of a function using that store.
 #[non_exhaustive]
 pub enum UpdateDeadline {
+    /// Halt execution of WebAssembly, don't update the epoch deadline, and
+    /// raise a trap.
+    Interrupt,
     /// Extend the deadline by the specified number of ticks.
     Continue(u64),
     /// Extend the deadline by the specified number of ticks after yielding to
@@ -429,7 +432,7 @@ pub struct StoreOpaque {
     // together. Then when we run out of gas, we inject the yield amount from the reserve
     // until the reserve is empty.
     fuel_reserve: u64,
-    fuel_yield_interval: Option<NonZeroU64>,
+    pub(crate) fuel_yield_interval: Option<NonZeroU64>,
     /// Indexed data within this `Store`, used to store information about
     /// globals, functions, memories, etc.
     store_data: StoreData,
@@ -1893,7 +1896,7 @@ impl StoreOpaque {
         Ok(get_fuel(injected_fuel, self.fuel_reserve))
     }
 
-    fn refuel(&mut self) -> bool {
+    pub(crate) fn refuel(&mut self) -> bool {
         let injected_fuel = unsafe { &mut *self.vm_store_context.fuel_consumed.get() };
         refuel(
             injected_fuel,
@@ -2277,6 +2280,22 @@ at https://bytecodealliance.org/security.
 
         Ok(id)
     }
+
+    #[cfg(target_has_atomic = "64")]
+    pub(crate) fn set_epoch_deadline(&mut self, delta: u64) {
+        // Set a new deadline based on the "epoch deadline delta".
+        //
+        // Also, note that when this update is performed while Wasm is
+        // on the stack, the Wasm will reload the new value once we
+        // return into it.
+        let current_epoch = self.engine().current_epoch();
+        let epoch_deadline = self.vm_store_context.epoch_deadline.get_mut();
+        *epoch_deadline = current_epoch + delta;
+    }
+
+    pub(crate) fn get_epoch_deadline(&mut self) -> u64 {
+        *self.vm_store_context.epoch_deadline.get_mut()
+    }
 }
 
 /// Helper parameter to [`StoreOpaque::allocate_instance`].
@@ -2327,67 +2346,19 @@ unsafe impl<T> VMStore for StoreInner<T> {
         )
     }
 
-    fn out_of_gas(&mut self) -> Result<()> {
-        if !self.refuel() {
-            return Err(Trap::OutOfFuel.into());
-        }
-        #[cfg(feature = "async")]
-        if self.fuel_yield_interval.is_some() {
-            self.async_yield_impl()?;
-        }
-        Ok(())
-    }
-
     #[cfg(target_has_atomic = "64")]
-    fn new_epoch(&mut self) -> Result<u64, anyhow::Error> {
+    fn new_epoch_updated_deadline(&mut self) -> Result<UpdateDeadline> {
         // Temporarily take the configured behavior to avoid mutably borrowing
         // multiple times.
         let mut behavior = self.epoch_deadline_behavior.take();
-        let delta_result = match &mut behavior {
-            None => Err(Trap::Interrupt.into()),
-            Some(callback) => callback((&mut *self).as_context_mut()).and_then(|update| {
-                let delta = match update {
-                    UpdateDeadline::Continue(delta) => delta,
-                    #[cfg(feature = "async")]
-                    UpdateDeadline::Yield(delta) => {
-                        assert!(
-                            self.async_support(),
-                            "cannot use `UpdateDeadline::Yield` without enabling async support in the config"
-                        );
-                        // Do the async yield. May return a trap if future was
-                        // canceled while we're yielded.
-                        self.async_yield_impl()?;
-                        delta
-                    }
-                    #[cfg(feature = "async")]
-                    UpdateDeadline::YieldCustom(delta, future) => {
-                        assert!(
-                            self.async_support(),
-                            "cannot use `UpdateDeadline::YieldCustom` without enabling async support in the config"
-                        );
-
-                        // When control returns, we have a `Result<()>` passed
-                        // in from the host fiber. If this finished successfully then
-                        // we were resumed normally via a `poll`, so keep going.  If
-                        // the future was dropped while we were yielded, then we need
-                        // to clean up this fiber. Do so by raising a trap which will
-                        // abort all wasm and get caught on the other side to clean
-                        // things up.
-                        self.block_on(|_| future)?;
-                        delta
-                    }
-                };
-
-                // Set a new deadline and return the new epoch deadline so
-                // the Wasm code doesn't have to reload it.
-                self.set_epoch_deadline(delta);
-                Ok(self.get_epoch_deadline())
-            })
+        let update = match &mut behavior {
+            Some(callback) => callback((&mut *self).as_context_mut()),
+            None => Ok(UpdateDeadline::Interrupt),
         };
 
         // Put back the original behavior which was replaced by `take`.
         self.epoch_deadline_behavior = behavior;
-        delta_result
+        update
     }
 
     #[cfg(feature = "component-model")]
@@ -2397,18 +2368,6 @@ unsafe impl<T> VMStore for StoreInner<T> {
 }
 
 impl<T> StoreInner<T> {
-    #[cfg(target_has_atomic = "64")]
-    pub(crate) fn set_epoch_deadline(&mut self, delta: u64) {
-        // Set a new deadline based on the "epoch deadline delta".
-        //
-        // Also, note that when this update is performed while Wasm is
-        // on the stack, the Wasm will reload the new value once we
-        // return into it.
-        let current_epoch = self.engine().current_epoch();
-        let epoch_deadline = self.vm_store_context.epoch_deadline.get_mut();
-        *epoch_deadline = current_epoch + delta;
-    }
-
     #[cfg(target_has_atomic = "64")]
     fn epoch_deadline_trap(&mut self) {
         self.epoch_deadline_behavior = None;
@@ -2421,10 +2380,6 @@ impl<T> StoreInner<T> {
     ) {
         self.epoch_deadline_behavior = Some(callback);
     }
-
-    fn get_epoch_deadline(&mut self) -> u64 {
-        *self.vm_store_context.epoch_deadline.get_mut()
-    }
 }
 
 impl<T: Default> Default for Store<T> {
diff --git a/crates/wasmtime/src/runtime/store/async_.rs b/crates/wasmtime/src/runtime/store/async_.rs
@@ -169,21 +169,6 @@ impl<T> StoreInner<T> {
 
 #[doc(hidden)]
 impl StoreOpaque {
-    /// Yields execution to the caller on out-of-gas or epoch interruption.
-    ///
-    /// This only works on async futures and stores, and assumes that we're
-    /// executing on a fiber. This will yield execution back to the caller once.
-    pub fn async_yield_impl(&mut self) -> Result<()> {
-        // When control returns, we have a `Result<()>` passed
-        // in from the host fiber. If this finished successfully then
-        // we were resumed normally via a `poll`, so keep going.  If
-        // the future was dropped while we were yielded, then we need
-        // to clean up this fiber. Do so by raising a trap which will
-        // abort all wasm and get caught on the other side to clean
-        // things up.
-        self.block_on(|_| Box::pin(crate::runtime::vm::Yield::new()))
-    }
-
     pub(crate) fn allocate_fiber_stack(&mut self) -> Result<wasmtime_fiber::FiberStack> {
         if let Some(stack) = self.async_state.last_fiber_stack().take() {
             return Ok(stack);
diff --git a/crates/wasmtime/src/runtime/vm.rs b/crates/wasmtime/src/runtime/vm.rs
@@ -208,16 +208,11 @@ pub unsafe trait VMStore: 'static {
         &mut self,
     ) -> (Option<StoreResourceLimiter<'_>>, &mut StoreOpaque);
 
-    /// Callback invoked whenever fuel runs out by a wasm instance. If an error
-    /// is returned that's raised as a trap. Otherwise wasm execution will
-    /// continue as normal.
-    fn out_of_gas(&mut self) -> Result<(), Error>;
-
     /// Callback invoked whenever an instance observes a new epoch
     /// number. Cannot fail; cooperative epoch-based yielding is
     /// completely semantically transparent. Returns the new deadline.
     #[cfg(target_has_atomic = "64")]
-    fn new_epoch(&mut self) -> Result<u64, Error>;
+    fn new_epoch_updated_deadline(&mut self) -> Result<crate::UpdateDeadline>;
 
     /// Metadata required for resources for the component model.
     #[cfg(feature = "component-model")]
diff --git a/crates/wasmtime/src/runtime/vm/libcalls.rs b/crates/wasmtime/src/runtime/vm/libcalls.rs
@@ -1211,13 +1211,61 @@ fn memory_atomic_wait64(
 
 // Hook for when an instance runs out of fuel.
 fn out_of_gas(store: &mut dyn VMStore, _instance: Pin<&mut Instance>) -> Result<()> {
-    store.out_of_gas()
+    block_on!(store, async |store| {
+        if !store.refuel() {
+            return Err(Trap::OutOfFuel.into());
+        }
+        #[cfg(feature = "async")]
+        if store.fuel_yield_interval.is_some() {
+            crate::runtime::vm::Yield::new().await;
+        }
+        Ok(())
+    })?
 }
 
 // Hook for when an instance observes that the epoch has changed.
 #[cfg(target_has_atomic = "64")]
 fn new_epoch(store: &mut dyn VMStore, _instance: Pin<&mut Instance>) -> Result<NextEpoch> {
-    store.new_epoch().map(NextEpoch)
+    use crate::UpdateDeadline;
+
+    let update_deadline = store.new_epoch_updated_deadline()?;
+    block_on!(store, async move |store| {
+        let delta = match update_deadline {
+            UpdateDeadline::Interrupt => return Err(Trap::Interrupt.into()),
+            UpdateDeadline::Continue(delta) => delta,
+
+            // Note that custom assertions for `async_support` are needed below
+            // as otherwise if these are used in an
+            // `async_support`-disabled-build it'll trip the `assert_ready` part
+            // of `block_on!` above. The assertion here provides a more direct
+            // error message as to what's going on.
+            #[cfg(feature = "async")]
+            UpdateDeadline::Yield(delta) => {
+                assert!(
+                    store.async_support(),
+                    "cannot use `UpdateDeadline::Yield` without enabling \
+                     async support in the config"
+                );
+                crate::runtime::vm::Yield::new().await;
+                delta
+            }
+            #[cfg(feature = "async")]
+            UpdateDeadline::YieldCustom(delta, future) => {
+                assert!(
+                    store.async_support(),
+                    "cannot use `UpdateDeadline::YieldCustom` without enabling \
+                     async support in the config"
+                );
+                future.await;
+                delta
+            }
+        };
+
+        // Set a new deadline and return the new epoch deadline so
+        // the Wasm code doesn't have to reload it.
+        store.set_epoch_deadline(delta);
+        Ok(NextEpoch(store.get_epoch_deadline()))
+    })?
 }
 
 struct NextEpoch(u64);
