core: Replace epoch deadline with yield

With async yielding, instance execution can be implemented by dropping
the async call future, with e.g. tokio::time::timeout.

Signed-off-by: Lann Martin <[email protected]>

Author: lann

crates/core/src/lib.rs

@@ -39,7 +39,7 @@ pub use io::OutputBuffer;
 pub use store::{Store, StoreBuilder, Wasi, WasiVersion};
 
 /// The default [`EngineBuilder::epoch_tick_interval`].
-pub const DEFAULT_EPOCH_TICK_INTERVAL: Duration = Duration::from_millis(10);
+pub const DEFAULT_EPOCH_TICK_INTERVAL: Duration = Duration::from_millis(1);
 
 const MB: u64 = 1 << 20;
 const GB: u64 = 1 << 30;
@@ -276,10 +276,11 @@ impl<T: Send + Sync> EngineBuilder<T> {
             .add_host_component(&mut self.linker, host_component)
     }
 
-    /// Sets the epoch tick internal for the built [`Engine`].
+    /// Sets the epoch tick interval for the built [`Engine`].
     ///
-    /// This is used by [`Store::set_deadline`] to calculate the number of
-    /// "ticks" for epoch interruption, and by the default epoch ticker thread.
+    /// This determines how often the engine's "epoch" will be incremented,
+    /// which determines the resolution of interrupt-based features like
+    /// [`Store::yield_interval`].
     /// The default is [`DEFAULT_EPOCH_TICK_INTERVAL`].
     ///
     /// See [`EngineBuilder::epoch_ticker_thread`] and
@@ -292,8 +293,8 @@ impl<T: Send + Sync> EngineBuilder<T> {
     /// [`Engine`] is built.
     ///
     /// Enabled by default; if disabled, the user must arrange to call
-    /// `engine.as_ref().increment_epoch()` every `epoch_tick_interval` or
-    /// interrupt-based features like `Store::set_deadline` will not work.
+    /// `engine.as_ref().increment_epoch()` periodically or interrupt-based
+    /// yielding will not work.
     pub fn epoch_ticker_thread(&mut self, enable: bool) {
         self.epoch_ticker_thread = enable;
     }

crates/core/src/store.rs

@@ -1,9 +1,9 @@
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, ensure, Result};
 use std::{
     io::{Read, Write},
     path::{Path, PathBuf},
     sync::Mutex,
-    time::{Duration, Instant},
+    time::Duration,
 };
 use system_interface::io::ReadReady;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -67,35 +67,13 @@ pub enum WasiVersion {
 /// A `Store` can be built with a [`StoreBuilder`].
 pub struct Store<T> {
     inner: wasmtime::Store<Data<T>>,
-    epoch_tick_interval: Duration,
 }
 
 impl<T> Store<T> {
     /// Returns a mutable reference to the [`HostComponentsData`] of this [`Store`].
     pub fn host_components_data(&mut self) -> &mut HostComponentsData {
         &mut self.inner.data_mut().host_components_data
     }
-
-    /// Sets the execution deadline.
-    ///
-    /// This is a rough deadline; an instance will trap some time after this
-    /// deadline, determined by [`EngineBuilder::epoch_tick_interval`] and
-    /// details of the system's thread scheduler.
-    ///
-    /// See [`wasmtime::Store::set_epoch_deadline`](https://docs.rs/wasmtime/latest/wasmtime/struct.Store.html#method.set_epoch_deadline).
-    pub fn set_deadline(&mut self, deadline: Instant) {
-        let now = Instant::now();
-        let duration = deadline - now;
-        let ticks = if duration.is_zero() {
-            tracing::warn!("Execution deadline set in past: {deadline:?} < {now:?}");
-            0
-        } else {
-            let ticks = duration.as_micros() / self.epoch_tick_interval.as_micros();
-            let ticks = ticks.min(u64::MAX as u128) as u64;
-            ticks + 1 // Add one to allow for current partially-completed tick
-        };
-        self.inner.set_epoch_deadline(ticks);
-    }
 }
 
 impl<T> AsRef<wasmtime::Store<Data<T>>> for Store<T> {
@@ -130,6 +108,7 @@ impl<T> wasmtime::AsContextMut for Store<T> {
 pub struct StoreBuilder {
     engine: wasmtime::Engine,
     epoch_tick_interval: Duration,
+    yield_interval: Duration,
     wasi: std::result::Result<WasiCtxBuilder, String>,
     host_components_data: HostComponentsData,
     store_limits: StoreLimitsAsync,
@@ -146,6 +125,7 @@ impl StoreBuilder {
         Self {
             engine,
             epoch_tick_interval,
+            yield_interval: epoch_tick_interval,
             wasi: Ok(wasi.into()),
             host_components_data: host_components.new_data(),
             store_limits: StoreLimitsAsync::default(),
@@ -160,6 +140,20 @@ impl StoreBuilder {
         self.store_limits = StoreLimitsAsync::new(Some(max_memory_size), None);
     }
 
+    /// Sets the execution yield interval.
+    ///
+    /// A CPU-bound running instance will be forced to yield approximately
+    /// every interval, which gives the host thread an opportunity to cancel
+    /// the instance or schedule other work on the thread.
+    ///
+    /// The exact interval of yielding is determined by [`EngineBuilder::epoch_tick_interval`]
+    /// and details of the task scheduler.
+    ///
+    /// The interval defaults to the epoch tick interval.
+    pub fn yield_interval(&mut self, interval: Duration) {
+        self.yield_interval = interval;
+    }
+
     /// Inherit stdin from the host process.
     pub fn inherit_stdin(&mut self) {
         self.with_wasi(|wasi| match wasi {
@@ -386,16 +380,26 @@ impl StoreBuilder {
 
         inner.limiter_async(move |data| &mut data.store_limits);
 
-        // With epoch interruption enabled, there must be _some_ deadline set
-        // or execution will trap immediately. Since this is a delta, we need
-        // to avoid overflow so we'll use 2^63 which is still "practically
-        // forever" for any plausible tick interval.
-        inner.set_epoch_deadline(u64::MAX / 2);
+        ensure!(
+            !self.epoch_tick_interval.is_zero(),
+            "epoch_tick_interval may not be zero"
+        );
+        let delta = self.yield_interval.as_nanos() / self.epoch_tick_interval.as_nanos();
+        let delta = if delta == 0 {
+            tracing::warn!(
+                "Yield interval {interval:?} too small to resolve; clamping to tick interval {tick:?}",
+                interval = self.yield_interval,
+                tick = self.epoch_tick_interval);
+            1
+        } else if delta > u64::MAX as u128 {
+            tracing::warn!("Yield interval too large; yielding effectively disabled");
+            u64::MAX
+        } else {
+            delta as u64
+        };
+        inner.epoch_deadline_async_yield_and_update(delta);
 
-        Ok(Store {
-            inner,
-            epoch_tick_interval: self.epoch_tick_interval,
-        })
+        Ok(Store { inner })
     }
 
     /// Builds a [`Store`] from this builder with `Default` host state data.

crates/core/tests/integration_test.rs

@@ -1,8 +1,4 @@
-use std::{
-    io::Cursor,
-    path::PathBuf,
-    time::{Duration, Instant},
-};
+use std::{io::Cursor, path::PathBuf, time::Duration};
 
 use anyhow::Context;
 use spin_core::{
@@ -102,33 +98,11 @@ async fn test_max_memory_size_violated() {
 }
 
 #[tokio::test(flavor = "multi_thread")]
-async fn test_set_deadline_obeyed() {
-    run_core_wasi_test_engine(
-        &test_engine(),
-        ["sleep", "20"],
-        |_| {},
-        |store| {
-            store.set_deadline(Instant::now() + Duration::from_millis(1000));
-        },
-    )
-    .await
-    .unwrap();
-}
-
-#[tokio::test(flavor = "multi_thread")]
-async fn test_set_deadline_violated() {
-    let err = run_core_wasi_test_engine(
-        &test_engine(),
-        ["sleep", "100"],
-        |_| {},
-        |store| {
-            store.set_deadline(Instant::now() + Duration::from_millis(10));
-        },
-    )
-    .await
-    .unwrap_err();
-    let trap = err.downcast::<Trap>().expect("trap");
-    assert_eq!(trap, Trap::Interrupt);
+async fn test_yield_interval_timeout() {
+    let forever = u64::MAX.to_string();
+    let fut = run_core_wasi_test(["sleep", &forever], |_| {});
+    let res = tokio::time::timeout(Duration::from_micros(1), fut).await;
+    assert!(res.is_err());
 }
 
 #[tokio::test(flavor = "multi_thread")]