feat: add support for epoch-based deadline extensions for host I/O

ruslanti · ruslanti · commit a8457cb155b6 · 2026-05-12T09:44:55.000+03:00
- Introduced `epoch_pause_ms` to track elapsed host I/O time and extend execution deadlines.
- Enhanced `StoreBuilder` with `epoch_pause_ms` and `max_external_duration_ms` configuration.
- Updated HTTP backends to deposit I/O time into shared counters and refund guest execution ticks.
- Integrated `epoch_deadline_callback` for dynamic deadline adjustments.
diff --git a/crates/http-backend/src/lib.rs b/crates/http-backend/src/lib.rs
@@ -6,8 +6,9 @@ use std::future::Future;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 use std::pin::Pin;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::task::{Context, Poll};
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use anyhow::{Error, Result, anyhow};
 use http::{HeaderMap, HeaderName, Uri, header, uri::Scheme};
@@ -64,6 +65,10 @@ pub struct Backend<C> {
     pub strategy: BackendStrategy,
     ext_http_stats: Option<Arc<dyn ExtRequestStats>>,
     hostname: Option<SmolStr>,
+    /// Counter shared with the wasmtime Store: each outbound HTTP call
+    /// deposits its elapsed wall-clock time (ms) here so the epoch
+    /// deadline callback can refund those ticks to the guest.
+    epoch_pause_ms: Arc<AtomicU64>,
 }
 
 pub struct Builder {
@@ -112,6 +117,7 @@ impl Builder {
             strategy: self.strategy,
             ext_http_stats: None,
             hostname: self.hostname.clone(),
+            epoch_pause_ms: Arc::new(AtomicU64::new(0)),
         }
     }
 }
@@ -140,6 +146,12 @@ impl<C> Backend<C> {
         self.ext_http_stats.replace(stats);
     }
 
+    /// Share the epoch-pause counter with the Store so that outbound HTTP
+    /// time can be excluded from the guest's execution-time budget.
+    pub fn set_epoch_pause_ms(&mut self, counter: Arc<AtomicU64>) {
+        self.epoch_pause_ms = counter;
+    }
+
     pub fn propagate_header_names(&self) -> HeaderNameList {
         self.propagate_header_names.clone()
     }
@@ -321,48 +333,57 @@ where
             .as_ref()
             .map(|s| ExtStatsTimer::new(s.clone()));
 
-        let res = self.client.request(request).await.map_err(|error| {
-            warn!(cause=?error, "sending request to backend");
-            HttpError::RequestError
-        })?;
-
-        let status = res.status().as_u16();
-        let (parts, body) = res.into_parts();
-        let headers = if !parts.headers.is_empty() {
-            Some(
-                parts
-                    .headers
-                    .iter()
-                    .filter_map(|(name, value)| match value.to_str() {
-                        Ok(value) => Some((name.to_string(), value.to_string())),
-                        Err(error) => {
-                            warn!(cause=?error, "invalid value: {:?}", value);
-                            None
-                        }
-                    })
-                    .collect::<Vec<(String, String)>>(),
-            )
-        } else {
-            None
-        };
-
-        let body_bytes = body
-            .collect()
-            .await
-            .map_err(|error| {
-                warn!(cause=?error, "receiving body from backend");
+        // Time the network I/O so we can refund the equivalent epoch ticks
+        // to the guest. Both the request send and the body read count.
+        let started = Instant::now();
+        let result = async {
+            let res = self.client.request(request).await.map_err(|error| {
+                warn!(cause=?error, "sending request to backend");
                 HttpError::RequestError
-            })?
-            .to_bytes();
-        let body = Some(body_bytes.to_vec());
-
-        trace!(?status, ?headers, len = body_bytes.len(), "reply");
+            })?;
+
+            let status = res.status().as_u16();
+            let (parts, body) = res.into_parts();
+            let headers = if !parts.headers.is_empty() {
+                Some(
+                    parts
+                        .headers
+                        .iter()
+                        .filter_map(|(name, value)| match value.to_str() {
+                            Ok(value) => Some((name.to_string(), value.to_string())),
+                            Err(error) => {
+                                warn!(cause=?error, "invalid value: {:?}", value);
+                                None
+                            }
+                        })
+                        .collect::<Vec<(String, String)>>(),
+                )
+            } else {
+                None
+            };
 
-        Ok(Response {
-            status,
-            headers,
-            body,
-        })
+            let body_bytes = body
+                .collect()
+                .await
+                .map_err(|error| {
+                    warn!(cause=?error, "receiving body from backend");
+                    HttpError::RequestError
+                })?
+                .to_bytes();
+            let body = Some(body_bytes.to_vec());
+
+            trace!(?status, ?headers, len = body_bytes.len(), "reply");
+
+            Ok::<_, HttpError>(Response {
+                status,
+                headers,
+                body,
+            })
+        }
+        .await;
+        self.epoch_pause_ms
+            .fetch_add(started.elapsed().as_millis() as u64, Ordering::Relaxed);
+        result
     }
 }
 
diff --git a/crates/http-service/src/executor/http.rs b/crates/http-service/src/executor/http.rs
@@ -11,6 +11,7 @@ use reactor::gcore::fastedge;
 use runtime::util::stats::{StatsTimer, StatsVisitor};
 use runtime::{InstancePre, store::StoreBuilder};
 use std::sync::Arc;
+use std::sync::atomic::AtomicU64;
 use std::time::Duration;
 use wasmtime_wasi_http::body::HyperOutgoingBody;
 
@@ -73,14 +74,21 @@ where
 
         let properties = executor::get_properties(&parts.headers);
 
-        let store_builder = self.store_builder.with_properties(properties);
+        // Shared counter so external HTTP time in `http_backend.send_request`
+        // refunds epoch ticks via the Store's `epoch_deadline_callback`.
+        let epoch_pause_ms = Arc::new(AtomicU64::new(0));
+        let store_builder = self
+            .store_builder
+            .with_properties(properties)
+            .epoch_pause_ms(epoch_pause_ms.clone());
         let mut http_backend = self.backend;
 
         http_backend
             .propagate_headers(parts.headers.clone())
             .context("propagate headers")?;
 
         http_backend.set_ext_http_stats(stats.clone());
+        http_backend.set_epoch_pause_ms(epoch_pause_ms);
 
         let propagate_header_names = http_backend.propagate_header_names();
         let backend_uri = http_backend.uri();
diff --git a/crates/http-service/src/executor/wasi_http.rs b/crates/http-service/src/executor/wasi_http.rs
@@ -1,4 +1,5 @@
 use std::sync::Arc;
+use std::sync::atomic::AtomicU64;
 use std::time::Duration;
 
 use crate::executor;
@@ -88,8 +89,16 @@ where
         let body = body.boxed();
 
         let properties = executor::get_properties(&parts.headers);
-        let store_builder = self.store_builder.with_properties(properties);
+        // Shared counter so wasi-http outbound calls (handled by
+        // `WasiHttpView::send_request` in the runtime) refund epoch ticks
+        // via the Store's `epoch_deadline_callback`.
+        let epoch_pause_ms = Arc::new(AtomicU64::new(0));
+        let store_builder = self
+            .store_builder
+            .with_properties(properties)
+            .epoch_pause_ms(epoch_pause_ms.clone());
         let mut http_backend = self.backend;
+        http_backend.set_epoch_pause_ms(epoch_pause_ms);
 
         http_backend
             .propagate_headers(parts.headers.clone())
diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs
@@ -2,6 +2,8 @@ use crate::app::KvStoreOption;
 use crate::store::HasStats;
 use http_backend::stats::ExtStatsTimer;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::time::Instant;
 use std::{fmt::Debug, ops::Deref};
 use utils::{Dictionary, Utils};
 use wasmtime_wasi::ResourceTable;
@@ -97,6 +99,11 @@ pub struct Data<T: 'static> {
     pub dictionary: Dictionary,
     pub utils: Utils,
     pub cache: cache::CacheImpl,
+    /// Milliseconds of host I/O that should not count against the epoch
+    /// deadline. The `epoch_deadline_callback` installed on the Store reads
+    /// and clears this counter, converting milliseconds into extra ticks to
+    /// extend the deadline.
+    pub epoch_pause_ms: Arc<AtomicU64>,
 }
 
 pub trait BackendRequest {
@@ -143,13 +150,16 @@ impl<T: Send + BackendRequest + HasStats> WasiHttpView for Data<T> {
         let request = Request::from_parts(head, body);
         // start external request stats timer
         let stats = self.inner.get_stats();
+        let epoch_pause_ms = self.epoch_pause_ms.clone();
 
         let handle = wasmtime_wasi::runtime::spawn(async move {
             let _stats_timer = ExtStatsTimer::new(stats); // keep timer alive until request is done
-            Ok(
+            let started = Instant::now();
+            let resp =
                 default_send_request_handler(request, OutgoingRequestConfig { use_tls, ..config })
-                    .await,
-            )
+                    .await;
+            epoch_pause_ms.fetch_add(started.elapsed().as_millis() as u64, Ordering::Relaxed);
+            Ok(resp)
         });
         Ok(HostFutureIncomingResponse::pending(handle))
     }
diff --git a/crates/runtime/src/store.rs b/crates/runtime/src/store.rs
@@ -6,18 +6,26 @@ use crate::{DEFAULT_EPOCH_TICK_INTERVAL, Data, Wasi, WasiVersion};
 use anyhow::Result;
 use secret::SecretStore;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::{
     collections::HashMap,
     fmt::{Debug, Formatter},
     ops::{Deref, DerefMut},
 };
 use tracing::{debug, instrument, trace};
 use utils::{Dictionary, Utils};
+use wasmtime::UpdateDeadline;
 use wasmtime::component::ResourceTable;
 use wasmtime_wasi::WasiCtxBuilder;
 use wasmtime_wasi_http::WasiHttpCtx;
 use wasmtime_wasi_nn::wit::WasiNnCtx;
 
+/// Default extra wall-clock budget (in ms) added to the tokio timeout that
+/// wraps wasm execution. Acts as a soft outer bound that must comfortably
+/// exceed the epoch budget once host I/O credit (e.g. external HTTP) extends
+/// the epoch deadline.
+pub const DEFAULT_MAX_EXTERNAL_DURATION_MS: u64 = 30_000;
+
 /// A `Store` holds the runtime state of a app instance.
 ///
 /// `Store` lives only for the lifetime of a single app invocation.
@@ -92,6 +100,8 @@ pub struct StoreBuilder {
     key_value_store: key_value_store::Builder,
     dictionary: Dictionary,
     cache_backend: Option<Arc<dyn cache::CacheBackend>>,
+    epoch_pause_ms: Option<Arc<AtomicU64>>,
+    max_external_duration_ms: u64,
 }
 
 impl StoreBuilder {
@@ -110,6 +120,8 @@ impl StoreBuilder {
             key_value_store: key_value_store::Builder::default(),
             dictionary: Default::default(),
             cache_backend: None,
+            epoch_pause_ms: None,
+            max_external_duration_ms: DEFAULT_MAX_EXTERNAL_DURATION_MS,
         }
     }
 
@@ -193,6 +205,28 @@ impl StoreBuilder {
         }
     }
 
+    /// Provide the shared counter that host functions use to "pause" the
+    /// epoch deadline. Each ms accumulated here grants the guest extra
+    /// ticks when the epoch deadline next fires. If unset, the store
+    /// falls back to a private counter (no host call can deposit credit).
+    pub fn epoch_pause_ms(self, counter: Arc<AtomicU64>) -> Self {
+        Self {
+            epoch_pause_ms: Some(counter),
+            ..self
+        }
+    }
+
+    /// Wall-clock slack (in ms) added to the tokio timeout that wraps
+    /// wasm execution. Must comfortably exceed the worst-case total time
+    /// spent in epoch-paused host calls so the tokio bound does not fire
+    /// before the epoch one.
+    pub fn max_external_duration_ms(self, ms: u64) -> Self {
+        Self {
+            max_external_duration_ms: ms,
+            ..self
+        }
+    }
+
     pub fn make_wasi_nn(&self) -> Result<WasiNnCtx> {
         // initialize application specific graph
         let backends: Vec<&str> = self
@@ -271,6 +305,9 @@ impl StoreBuilder {
             self.cache_backend
                 .unwrap_or_else(|| Arc::new(cache::NoCacheBackend)),
         );
+        let epoch_pause_ms = self
+            .epoch_pause_ms
+            .unwrap_or_else(|| Arc::new(AtomicU64::new(0)));
 
         let mut inner = wasmtime::Store::new(
             &self.engine,
@@ -279,7 +316,8 @@ impl StoreBuilder {
                 wasi,
                 wasi_nn,
                 store_limits: self.store_limits,
-                timeout: (self.max_duration + 1) * DEFAULT_EPOCH_TICK_INTERVAL,
+                timeout: (self.max_duration + 1) * DEFAULT_EPOCH_TICK_INTERVAL
+                    + self.max_external_duration_ms,
                 table,
                 logger,
                 http: WasiHttpCtx::new(),
@@ -288,12 +326,23 @@ impl StoreBuilder {
                 dictionary: self.dictionary,
                 utils,
                 cache: cache_impl,
+                epoch_pause_ms: epoch_pause_ms.clone(),
             },
         );
         inner.limiter(|state| &mut state.store_limits);
         // allow max number of epoch ticks (1 tick = 10 ms)
-        inner.set_epoch_deadline(self.max_duration); // allow max number of epoch ticks (1 tick = 10 ms)
-        inner.epoch_deadline_trap();
+        inner.set_epoch_deadline(self.max_duration);
+        // When the deadline fires, consume any host-call credit accumulated by
+        // `epoch_pause_ms` to extend it; if there is no credit, trap as before.
+        inner.epoch_deadline_callback(move |_ctx| {
+            let credit_ms = epoch_pause_ms.swap(0, Ordering::Relaxed);
+            if credit_ms == 0 {
+                Err(anyhow::Error::new(wasmtime::Trap::Interrupt))
+            } else {
+                let credit_ticks = (credit_ms / DEFAULT_EPOCH_TICK_INTERVAL).max(1);
+                Ok(UpdateDeadline::Continue(credit_ticks))
+            }
+        });
         Ok(Store { inner })
     }
 }
