diff --git a/RELEASES.md b/RELEASES.md index 64e1b2d11394..7f1fc7c6fa96 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -12,6 +12,10 @@ Unreleased. ### Changed +* An `epoch_deadline_callback` now returns an `UpdateDeadline` enum to allow + optionally yielding to the async executor after the callback runs. + [#6464](https://github.com/bytecodealliance/wasmtime/pull/6464) + * The "raw" representation of `funcref` and `externref` in the embedding API has been updated from a `usize` to a `*mut u8` to be compatible with Rust's proposed strict provenance rules. This change is additionally reflected into diff --git a/crates/c-api/src/store.rs b/crates/c-api/src/store.rs index 5aedaf371860..d56977619888 100644 --- a/crates/c-api/src/store.rs +++ b/crates/c-api/src/store.rs @@ -4,7 +4,7 @@ use std::ffi::c_void; use std::sync::Arc; use wasmtime::{ AsContext, AsContextMut, Store, StoreContext, StoreContextMut, StoreLimits, StoreLimitsBuilder, - Val, + UpdateDeadline, Val, }; /// This representation of a `Store` is used to implement the `wasm.h` API. @@ -143,7 +143,7 @@ pub extern "C" fn wasmtime_store_epoch_deadline_callback( Some(err) => Err(wasmtime::Error::from(>::into(*err))), - None => Ok(delta), + None => Ok(UpdateDeadline::Continue(delta)), } }); } diff --git a/crates/wasmtime/src/lib.rs b/crates/wasmtime/src/lib.rs index 2dfac02c25ec..37e58185229c 100644 --- a/crates/wasmtime/src/lib.rs +++ b/crates/wasmtime/src/lib.rs @@ -424,7 +424,9 @@ pub use crate::profiling::GuestProfiler; pub use crate::r#ref::ExternRef; #[cfg(feature = "async")] pub use crate::store::CallHookHandler; -pub use crate::store::{AsContext, AsContextMut, CallHook, Store, StoreContext, StoreContextMut}; +pub use crate::store::{ + AsContext, AsContextMut, CallHook, Store, StoreContext, StoreContextMut, UpdateDeadline, +}; pub use crate::trap::*; pub use crate::types::*; pub use crate::values::*; diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 14b3a3d638a9..8e915ec994da 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -204,7 +204,8 @@ pub struct StoreInner { limiter: Option>, call_hook: Option>, - epoch_deadline_behavior: EpochDeadline, + epoch_deadline_behavior: + Option) -> Result + Send + Sync>>, // for comments about `ManuallyDrop`, see `Store::into_data` data: ManuallyDrop, } @@ -230,6 +231,18 @@ enum CallHookInner { Async(Box + Send + Sync>), } +/// What to do after returning from a callback when the engine epoch reaches +/// the deadline for a Store during execution of a function using that store. +pub enum UpdateDeadline { + /// Extend the deadline by the specified number of ticks. + Continue(u64), + /// Extend the deadline by the specified number of ticks after yielding to + /// the async executor loop. This can only be used with an async [`Store`] + /// configured via [`Config::async_support`](crate::Config::async_support). + #[cfg(feature = "async")] + Yield(u64), +} + // Forward methods on `StoreOpaque` to also being on `StoreInner` impl Deref for StoreInner { type Target = StoreOpaque; @@ -422,21 +435,6 @@ enum OutOfGas { }, } -/// What to do when the engine epoch reaches the deadline for a Store -/// during execution of a function using that store. -#[derive(Default)] -enum EpochDeadline { - /// Return early with a trap. - #[default] - Trap, - /// Call a custom deadline handler. - Callback(Box) -> Result + Send + Sync>), - /// Extend the deadline by the specified number of ticks after - /// yielding to the async executor loop. - #[cfg(feature = "async")] - YieldAndExtendDeadline { delta: u64 }, -} - impl Store { /// Creates a new [`Store`] to be associated with the given [`Engine`] and /// `data` provided. @@ -480,7 +478,7 @@ impl Store { }, limiter: None, call_hook: None, - epoch_deadline_behavior: EpochDeadline::Trap, + epoch_deadline_behavior: None, data: ManuallyDrop::new(data), }); @@ -928,10 +926,17 @@ impl Store { /// store and the epoch deadline is reached before completion, the /// provided callback function is invoked. /// - /// This function should return a positive `delta`, which is used to - /// update the new epoch, setting it to the current epoch plus - /// `delta` ticks. Alternatively, the callback may return an error, - /// which will terminate execution. + /// This callback should either return an [`UpdateDeadline`], or + /// return an error, which will terminate execution with a trap. + /// + /// The [`UpdateDeadline`] is a positive number of ticks to + /// add to the epoch deadline, as well as indicating what + /// to do after the callback returns. If the [`Store`] is + /// configured with async support, then the callback may return + /// [`UpdateDeadline::Yield`] to yield to the async executor before + /// updating the epoch deadline. Alternatively, the callback may + /// return [`UpdateDeadline::Continue`] to update the epoch deadline + /// immediately. /// /// This setting is intended to allow for coarse-grained /// interruption, but not a deterministic deadline of a fixed, @@ -943,7 +948,7 @@ impl Store { /// for an introduction to epoch-based interruption. pub fn epoch_deadline_callback( &mut self, - callback: impl FnMut(StoreContextMut) -> Result + Send + Sync + 'static, + callback: impl FnMut(StoreContextMut) -> Result + Send + Sync + 'static, ) { self.inner.epoch_deadline_callback(Box::new(callback)); } @@ -1945,29 +1950,31 @@ unsafe impl wasmtime_runtime::Store for StoreInner { fn new_epoch(&mut self) -> Result { // Temporarily take the configured behavior to avoid mutably borrowing // multiple times. - let mut behavior = std::mem::take(&mut self.epoch_deadline_behavior); + let mut behavior = self.epoch_deadline_behavior.take(); let delta_result = match &mut behavior { - EpochDeadline::Trap => Err(Trap::Interrupt.into()), - EpochDeadline::Callback(callback) => { - let delta = callback((&mut *self).as_context_mut())?; + 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 + } + }; + // 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()) - } - #[cfg(feature = "async")] - EpochDeadline::YieldAndExtendDeadline { delta } => { - let delta = *delta; - // Do the async yield. May return a trap if future was - // canceled while we're yielded. - self.async_yield_impl()?; - // Set a new deadline. - self.set_epoch_deadline(delta); - - // Return the new epoch deadline so the Wasm code - // doesn't have to reload it. - Ok(self.get_epoch_deadline()) - } + }) }; // Put back the original behavior which was replaced by `take`. @@ -1992,14 +1999,14 @@ impl StoreInner { } fn epoch_deadline_trap(&mut self) { - self.epoch_deadline_behavior = EpochDeadline::Trap; + self.epoch_deadline_behavior = None; } fn epoch_deadline_callback( &mut self, - callback: Box) -> Result + Send + Sync>, + callback: Box) -> Result + Send + Sync>, ) { - self.epoch_deadline_behavior = EpochDeadline::Callback(callback); + self.epoch_deadline_behavior = Some(callback); } fn epoch_deadline_async_yield_and_update(&mut self, delta: u64) { @@ -2009,7 +2016,8 @@ impl StoreInner { ); #[cfg(feature = "async")] { - self.epoch_deadline_behavior = EpochDeadline::YieldAndExtendDeadline { delta }; + self.epoch_deadline_behavior = + Some(Box::new(move |_store| Ok(UpdateDeadline::Yield(delta)))); } let _ = delta; // suppress warning in non-async build } diff --git a/src/commands/run.rs b/src/commands/run.rs index 95c5fc43ef0b..18eec5a500be 100644 --- a/src/commands/run.rs +++ b/src/commands/run.rs @@ -11,7 +11,7 @@ use std::thread; use std::time::Duration; use wasmtime::{ AsContextMut, Engine, Func, GuestProfiler, Linker, Module, Store, StoreLimits, - StoreLimitsBuilder, Val, ValType, + StoreLimitsBuilder, UpdateDeadline, Val, ValType, }; use wasmtime_cli_flags::{CommonOptions, WasiModules}; use wasmtime_wasi::maybe_exit_on_error; @@ -467,12 +467,12 @@ impl RunCommand { if timeout == 0 { bail!("timeout exceeded"); } - Ok(1) + Ok(UpdateDeadline::Continue(1)) }); } else { store.epoch_deadline_callback(move |mut store| { sample(&mut store); - Ok(1) + Ok(UpdateDeadline::Continue(1)) }); } diff --git a/tests/all/epoch_interruption.rs b/tests/all/epoch_interruption.rs index 8e32881c2dc3..8e2424b280cf 100644 --- a/tests/all/epoch_interruption.rs +++ b/tests/all/epoch_interruption.rs @@ -34,7 +34,7 @@ fn make_env(engine: &Engine) -> Linker { enum InterruptMode { Trap, - Callback(fn(StoreContextMut) -> Result), + Callback(fn(StoreContextMut) -> Result), Yield(u64), } @@ -339,7 +339,32 @@ async fn epoch_callback_continue() { InterruptMode::Callback(|mut cx| { let s = cx.data_mut(); *s += 1; - Ok(1) + Ok(UpdateDeadline::Continue(1)) + }), + |_| {}, + ) + .await + ); +} + +#[tokio::test] +async fn epoch_callback_yield() { + assert_eq!( + Some((1, 1)), + run_and_count_yields_or_trap( + " + (module + (import \"\" \"bump_epoch\" (func $bump)) + (func (export \"run\") + call $bump ;; bump epoch + call $subfunc) ;; call func; will notice new epoch and yield + (func $subfunc)) + ", + 1, + InterruptMode::Callback(|mut cx| { + let s = cx.data_mut(); + *s += 1; + Ok(UpdateDeadline::Yield(1)) }), |_| {}, )