diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index a4015eacdd..9c1ed87cab 100644 --- a/crates/wasmtime/src/config.rs +++ b/crates/wasmtime/src/config.rs @@ -360,23 +360,27 @@ impl Config { /// /// The [`Store`](crate::Store) tracks the deadline, and controls /// what happens when the deadline is reached during - /// execution. Two behaviors are possible: + /// execution. Several behaviors are possible: /// /// - Trap if code is executing when the epoch deadline is /// met. See /// [`Store::epoch_deadline_trap`](crate::Store::epoch_deadline_trap). /// + /// - Call an arbitrary function. This function may chose to trap or + /// increment the epoch. See + /// [`Store::epoch_deadline_callback`](crate::Store::epoch_deadline_callback). + /// /// - Yield to the executor loop, then resume when the future is /// next polled. See /// [`Store::epoch_deadline_async_yield_and_update`](crate::Store::epoch_deadline_async_yield_and_update). /// - /// The first is the default; set the second for the timeslicing - /// behavior described above. + /// Trapping is the default. The yielding behaviour may be used for + /// the timeslicing behavior described above. /// - /// This feature is available with or without async - /// support. However, without async support, only the trapping - /// behavior is available. In this mode, epoch-based interruption - /// can serve as a simple external-interruption mechanism. + /// This feature is available with or without async support. + /// However, without async support, the timeslicing behaviour is + /// not available. This means epoch-based interruption can only + /// serve as a simple external-interruption mechanism. /// /// An initial deadline can be set before executing code by /// calling @@ -404,6 +408,7 @@ impl Config { /// - [`Engine::increment_epoch`](crate::Engine::increment_epoch) /// - [`Store::set_epoch_deadline`](crate::Store::set_epoch_deadline) /// - [`Store::epoch_deadline_trap`](crate::Store::epoch_deadline_trap) + /// - [`Store::epoch_deadline_callback`](crate::Store::epoch_deadline_callback) /// - [`Store::epoch_deadline_async_yield_and_update`](crate::Store::epoch_deadline_async_yield_and_update) pub fn epoch_interruption(&mut self, enable: bool) -> &mut Self { self.tunables.epoch_interruption = enable; diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index cb56e67822..430e201816 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -202,6 +202,7 @@ pub struct StoreInner { limiter: Option>, call_hook: Option>, + epoch_deadline_behavior: EpochDeadline, // for comments about `ManuallyDrop`, see `Store::into_data` data: ManuallyDrop, } @@ -296,7 +297,6 @@ pub struct StoreOpaque { #[cfg(feature = "async")] async_state: AsyncState, out_of_gas_behavior: OutOfGas, - epoch_deadline_behavior: EpochDeadline, /// Indexed data within this `Store`, used to store information about /// globals, functions, memories, etc. /// @@ -426,10 +426,11 @@ enum OutOfGas { /// What to do when the engine epoch reaches the deadline for a Store /// during execution of a function using that store. -#[derive(Copy, Clone)] -enum EpochDeadline { +enum EpochDeadline { /// Return early with a trap. 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")] @@ -491,7 +492,6 @@ impl Store { current_poll_cx: UnsafeCell::new(ptr::null_mut()), }, out_of_gas_behavior: OutOfGas::Trap, - epoch_deadline_behavior: EpochDeadline::Trap, store_data: ManuallyDrop::new(StoreData::new()), default_callee, hostcall_val_storage: Vec::new(), @@ -500,6 +500,7 @@ impl Store { }, limiter: None, call_hook: None, + epoch_deadline_behavior: EpochDeadline::Trap, data: ManuallyDrop::new(data), }); @@ -842,7 +843,8 @@ impl Store { /// /// This behavior is the default if the store is not otherwise /// configured via - /// [`epoch_deadline_trap()`](Store::epoch_deadline_trap) or + /// [`epoch_deadline_trap()`](Store::epoch_deadline_trap), + /// [`epoch_deadline_callback()`](Store::epoch_deadline_callback) or /// [`epoch_deadline_async_yield_and_update()`](Store::epoch_deadline_async_yield_and_update). /// /// This setting is intended to allow for coarse-grained @@ -857,6 +859,33 @@ impl Store { self.inner.epoch_deadline_trap(); } + /// Configures epoch-deadline expiration to invoke a custom callback + /// function. + /// + /// When epoch-interruption-instrumented code is executed on this + /// 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 setting is intended to allow for coarse-grained + /// interruption, but not a deterministic deadline of a fixed, + /// finite interval. For deterministic interruption, see the + /// "fuel" mechanism instead. + /// + /// See documentation on + /// [`Config::epoch_interruption()`](crate::Config::epoch_interruption) + /// for an introduction to epoch-based interruption. + pub fn epoch_deadline_callback( + &mut self, + callback: impl FnMut(&mut T) -> Result + Send + Sync + 'static, + ) { + self.inner.epoch_deadline_callback(Box::new(callback)); + } + #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))] /// Configures epoch-deadline expiration to yield to the async /// caller and the update the deadline. @@ -1351,22 +1380,6 @@ impl StoreOpaque { } } - fn epoch_deadline_trap(&mut self) { - self.epoch_deadline_behavior = EpochDeadline::Trap; - } - - fn epoch_deadline_async_yield_and_update(&mut self, delta: u64) { - assert!( - self.async_support(), - "cannot use `epoch_deadline_async_yield_and_update` without enabling async support in the config" - ); - #[cfg(feature = "async")] - { - self.epoch_deadline_behavior = EpochDeadline::YieldAndExtendDeadline { delta }; - } - drop(delta); // suppress warning in non-async build - } - #[inline] pub fn signal_handler(&self) -> Option<*const SignalHandler<'static>> { let handler = self.signal_handler.as_ref()?; @@ -1855,8 +1868,8 @@ unsafe impl wasmtime_runtime::Store for StoreInner { } fn new_epoch(&mut self) -> Result { - return match &self.epoch_deadline_behavior { - &EpochDeadline::Trap => { + return match &mut self.epoch_deadline_behavior { + EpochDeadline::Trap => { let trap = Trap::new_wasm( None, wasmtime_environ::TrapCode::Interrupt, @@ -1864,8 +1877,16 @@ unsafe impl wasmtime_runtime::Store for StoreInner { ); Err(anyhow::Error::from(trap)) } + EpochDeadline::Callback(callback) => { + let delta = callback(&mut self.data)?; + // 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 } => { + 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()?; @@ -1895,6 +1916,29 @@ impl StoreInner { *epoch_deadline = self.engine().current_epoch() + delta; } + fn epoch_deadline_trap(&mut self) { + self.epoch_deadline_behavior = EpochDeadline::Trap; + } + + fn epoch_deadline_callback( + &mut self, + callback: Box Result + Send + Sync>, + ) { + self.epoch_deadline_behavior = EpochDeadline::Callback(callback); + } + + fn epoch_deadline_async_yield_and_update(&mut self, delta: u64) { + assert!( + self.async_support(), + "cannot use `epoch_deadline_async_yield_and_update` without enabling async support in the config" + ); + #[cfg(feature = "async")] + { + self.epoch_deadline_behavior = EpochDeadline::YieldAndExtendDeadline { delta }; + } + drop(delta); // suppress warning in non-async build + } + fn get_epoch_deadline(&self) -> u64 { // Safety: this is safe because, as above, it is only invoked // from within `new_epoch` which is called from guest Wasm diff --git a/tests/all/epoch_interruption.rs b/tests/all/epoch_interruption.rs index 61d7c1b49e..c34a8b7f75 100644 --- a/tests/all/epoch_interruption.rs +++ b/tests/all/epoch_interruption.rs @@ -1,4 +1,5 @@ use crate::async_functions::{CountPending, PollOnce}; +use anyhow::{anyhow, Result}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use wasmtime::*; @@ -10,7 +11,7 @@ fn build_engine() -> Arc { Arc::new(Engine::new(&config).unwrap()) } -fn make_env(engine: &Engine) -> Linker<()> { +fn make_env(engine: &Engine) -> Linker { let mut linker = Linker::new(engine); let engine = engine.clone(); @@ -29,29 +30,38 @@ fn make_env(engine: &Engine) -> Linker<()> { linker } +enum InterruptMode { + Trap, + Callback(fn(&mut usize) -> Result), + Yield(u64), +} + /// Run a test with the given wasm, giving an initial deadline of /// `initial` ticks in the future, and either configuring the wasm to /// yield and set a deadline `delta` ticks in the future if `delta` is /// `Some(..)` or trapping if `delta` is `None`. /// -/// Returns `Some(yields)` if function completed normally, giving the -/// number of yields that occured, or `None` if a trap occurred. +/// Returns `Some((yields, store))` if function completed normally, giving +/// the number of yields that occurred, or `None` if a trap occurred. async fn run_and_count_yields_or_trap)>( wasm: &str, initial: u64, - delta: Option, + delta: InterruptMode, setup_func: F, -) -> Option { +) -> Option<(usize, usize)> { let engine = build_engine(); let linker = make_env(&engine); let module = Module::new(&engine, wasm).unwrap(); - let mut store = Store::new(&engine, ()); + let mut store = Store::new(&engine, 0); store.set_epoch_deadline(initial); match delta { - Some(delta) => { + InterruptMode::Yield(delta) => { store.epoch_deadline_async_yield_and_update(delta); } - None => { + InterruptMode::Callback(func) => { + store.epoch_deadline_callback(func); + } + InterruptMode::Trap => { store.epoch_deadline_trap(); } } @@ -63,14 +73,15 @@ async fn run_and_count_yields_or_trap)>( let f = instance.get_func(&mut store, "run").unwrap(); let (result, yields) = CountPending::new(Box::pin(f.call_async(&mut store, &[], &mut []))).await; - return result.ok().map(|_| yields); + let store = store.data(); + return result.ok().map(|_| (yields, *store)); } #[tokio::test] async fn epoch_yield_at_func_entry() { // Should yield at start of call to func $subfunc. assert_eq!( - Some(1), + Some((1, 0)), run_and_count_yields_or_trap( " (module @@ -81,7 +92,7 @@ async fn epoch_yield_at_func_entry() { (func $subfunc)) ", 1, - Some(1), + InterruptMode::Yield(1), |_| {}, ) .await @@ -92,7 +103,7 @@ async fn epoch_yield_at_func_entry() { async fn epoch_yield_at_loop_header() { // Should yield at top of loop, once per five iters. assert_eq!( - Some(2), + Some((2, 0)), run_and_count_yields_or_trap( " (module @@ -105,7 +116,7 @@ async fn epoch_yield_at_loop_header() { (br_if $l (local.tee $i (i32.sub (local.get $i) (i32.const 1))))))) ", 0, - Some(5), + InterruptMode::Yield(5), |_| {}, ) .await @@ -117,7 +128,7 @@ async fn epoch_yield_immediate() { // We should see one yield immediately when the initial deadline // is zero. assert_eq!( - Some(1), + Some((1, 0)), run_and_count_yields_or_trap( " (module @@ -125,7 +136,7 @@ async fn epoch_yield_immediate() { (func (export \"run\"))) ", 0, - Some(1), + InterruptMode::Yield(1), |_| {}, ) .await @@ -139,7 +150,7 @@ async fn epoch_yield_only_once() { // not yield again (the double-check block will reload the correct // epoch). assert_eq!( - Some(1), + Some((1, 0)), run_and_count_yields_or_trap( " (module @@ -155,7 +166,7 @@ async fn epoch_yield_only_once() { (call $bump))) ", 1, - Some(1), + InterruptMode::Yield(1), |_| {}, ) .await @@ -175,7 +186,7 @@ async fn epoch_interrupt_infinite_loop() { (br $l)))) ", 1, - None, + InterruptMode::Trap, |engine| { std::thread::spawn(move || { std::thread::sleep(std::time::Duration::from_millis(50)); @@ -297,7 +308,7 @@ async fn epoch_interrupt_function_entries() { (func $f9)) ", 1, - None, + InterruptMode::Trap, |engine| { std::thread::spawn(move || { std::thread::sleep(std::time::Duration::from_millis(50)); @@ -309,6 +320,51 @@ async fn epoch_interrupt_function_entries() { ); } +#[tokio::test] +async fn epoch_callback_continue() { + assert_eq!( + Some((0, 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(|s| { + *s += 1; + Ok(1) + }), + |_| {}, + ) + .await + ); +} + +#[tokio::test] +async fn epoch_callback_trap() { + assert_eq!( + None, + 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(|_| Err(anyhow!("Failing in callback"))), + |_| {}, + ) + .await + ); +} + #[tokio::test] async fn drop_future_on_epoch_yield() { let wasm = "