kobako 0.28.0 → 0.30.0

data/crates/kobako-wasmtime/src/instance_pre.rs

@@ -1,13 +1,10 @@
 //! Per-path cache of pre-instantiated wasmtime artifacts.
 //!
-//! The `Linker` wiring (the WASI preview1 import set plus the
-//! `__kobako_upcall` host import) and its type-check against the
-//! compiled Module are identical for every `Kobako::Runtime` on the
-//! same Guest Binary — both host closures read all their state from
-//! the `Invocation` inside the calling Store, never from the Runtime.
-//! Caching the resolved `InstancePre` per path — together with the ABI
-//! version check it guards (`crate::abi`) — takes both off the
-//! `Driver::new` hot path.
+//! The host imports (`crate::imports`) and their type-check against the
+//! compiled Module are identical for every `Kobako::Runtime` on the same
+//! Guest Binary. Caching the resolved `InstancePre` per path — together
+//! with the ABI version check it guards (`crate::abi`) — takes both off
+//! the `Driver::new` hot path.
 //!
 //! Concurrency: the same as `crate::cache` — the Mutex covers lookup and
 //! insert only, and a concurrent miss builds an equivalent `InstancePre`.
@@ -15,19 +12,24 @@
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use std::sync::{Mutex, OnceLock};
+use std::time::Duration;
 
-use wasmtime::{Caller, InstancePre, Linker};
-use wasmtime_wasi::p1;
+use wasmtime::InstancePre;
 
-use crate::cache::{cached_module, shared_engine};
+use crate::cache::load_module;
 use crate::invocation::Invocation;
-use crate::{abi, dispatch, trap};
+use crate::{abi, imports, trap};
 use kobako_runtime::error::SetupError;
 
 static INSTANCE_PRE_CACHE: OnceLock<Mutex<HashMap<PathBuf, InstancePre<Invocation>>>> =
     OnceLock::new();
 
-pub(crate) fn cached_instance_pre(path: &Path) -> Result<InstancePre<Invocation>, SetupError> {
+/// The probe runs only on the first load of a path, so only that load's
+/// `setup_timeout` bounds it; a path already verified is not set up again.
+pub(crate) fn load_instance_pre(
+    path: &Path,
+    setup_timeout: Option<Duration>,
+) -> Result<InstancePre<Invocation>, SetupError> {
     let cache = INSTANCE_PRE_CACHE.get_or_init(|| Mutex::new(HashMap::new()));
 
     if let Some(pre) = cache
@@ -39,47 +41,17 @@ pub(crate) fn cached_instance_pre(path: &Path) -> Result<InstancePre<Invocation>
         return Ok(pre);
     }
 
-    let module = cached_module(path)?;
-    let linker = build_linker()?;
+    let module = load_module(path)?;
+    let linker = imports::linker()?;
     let pre = linker
         .instantiate_pre(&module)
         .map_err(trap::instantiate_err)?;
     // Verify before the insert, so a rejected artifact is never
     // remembered as a usable template.
-    abi::verify(&pre)?;
+    abi::verify(&pre, setup_timeout)?;
     cache
         .lock()
         .expect("instance_pre cache mutex poisoned")
         .insert(path.to_path_buf(), pre.clone());
     Ok(pre)
 }
-
-fn build_linker() -> Result<Linker<Invocation>, SetupError> {
-    let mut linker: Linker<Invocation> = Linker::new(shared_engine()?);
-
-    // Wire the wasmtime-wasi preview1 WASI imports. Routes guest fd 1/2
-    // to the MemoryOutputPipes set up before each run via
-    // `Driver::invoke`. The closure pulls a `&mut WasiP1Ctx` out of
-    // Invocation; the panic semantics live inside `Invocation::wasi_mut`
-    // so the wiring stays honest about its precondition.
-    p1::add_to_linker_sync(&mut linker, |state: &mut Invocation| state.wasi_mut())
-        .map_err(|e| SetupError::Dead(format!("failed to set up the WASI runtime: {e}")))?;
-
-    // `__kobako_upcall`, `(ptr: i32, len: i32) -> i64` per
-    // docs/wire/abi.md § ABI Signatures; `dispatch::handle` owns what it
-    // answers.
-    linker
-        .func_wrap(
-            "env",
-            "__kobako_upcall",
-            |mut caller: Caller<'_, Invocation>,
-             req_ptr: i32,
-             req_len: i32|
-             -> wasmtime::Result<i64> {
-                dispatch::handle(&mut caller, req_ptr, req_len)
-            },
-        )
-        .map_err(|e| SetupError::Dead(format!("failed to set up the host callback bridge: {e}")))?;
-
-    Ok(linker)
-}

data/crates/kobako-wasmtime/src/invocation.rs

@@ -4,9 +4,8 @@
 //! Owned as the data of each per-invocation `wasmtime::Store`
 //! and threaded through every host import —
 //! the `__kobako_upcall` dispatcher reads the bound dispatch handler,
-//! while `Driver::invoke` installs the invocation's WASI
-//! context + pipes (via `io::install_wasi`) before the guest
-//! export call.
+//! and the WASI imports read the invocation's `Wasi` context, which the
+//! Store is created with.
 //!
 //! The slot also carries the per-invocation wall-clock deadline and the
 //! memory cap `MemoryLimiter`, both read from the wasmtime
@@ -16,53 +15,56 @@
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
-use wasmtime_wasi::p1::WasiP1Ctx;
-use wasmtime_wasi::p2::pipe::MemoryOutputPipe;
-
 use kobako_runtime::dispatch::DispatchHandler;
+use kobako_runtime::endpoint::{Calls, MemoryCap, ABANDONED};
 
 use crate::exports::Exports;
 use crate::limiter::MemoryLimiter;
+use crate::wasi::Wasi;
 
 /// Per-invocation host state, owned by the invocation's own Store.
 /// Threaded through every host import callback.
 ///
-/// All field access is mediated by methods on this type — the WASI ctx
-/// is rebuilt fresh before each invocation via
-/// `Invocation::install_wasi`, the dispatch handler is set once via
-/// `Invocation::bind_on_dispatch`, and captured stdout/stderr bytes
-/// are read after the invocation via `Invocation::stdout_bytes` /
-/// `Invocation::stderr_bytes`. The fields are private so the mutation
-/// surface stays narrow.
+/// All field access is mediated by methods on this type — the `Wasi`
+/// context arrives with the invocation, the dispatch handler is set once
+/// via `Invocation::bind_on_dispatch`, and the captures are read back
+/// through `Invocation::wasi` after the run. The fields are private so
+/// the mutation surface stays narrow.
 pub(crate) struct Invocation {
-    wasi: Option<WasiP1Ctx>,
-    stdout_pipe: Option<MemoryOutputPipe>,
-    stderr_pipe: Option<MemoryOutputPipe>,
-    on_dispatch: Option<Arc<dyn DispatchHandler>>,
+    wasi: Wasi,
+    calls: Calls<wasmtime::Error>,
     deadline: Option<Instant>,
     limiter: MemoryLimiter,
     wall_entry: Option<Instant>,
     wall_time: Duration,
-    reentry_trap: Option<wasmtime::Error>,
     exports: Option<Exports>,
+    hostcall_fuel: usize,
 }
 
 impl Invocation {
-    pub(crate) fn new(memory_limit: Option<usize>) -> Self {
+    pub(crate) fn new(memory_limit: Option<usize>, wasi: Wasi) -> Self {
         Self {
-            wasi: None,
-            stdout_pipe: None,
-            stderr_pipe: None,
-            on_dispatch: None,
+            wasi,
+            calls: Calls::default(),
             deadline: None,
-            limiter: MemoryLimiter::new(memory_limit),
+            limiter: MemoryLimiter(MemoryCap::new(memory_limit)),
             wall_entry: None,
             wall_time: Duration::ZERO,
-            reentry_trap: None,
             exports: None,
+            hostcall_fuel: 0,
         }
     }
 
+    /// How much of the guest's memory the call under way may have the host
+    /// copy or walk, as wasmtime hands each host call.
+    pub(crate) fn hostcall_fuel(&self) -> usize {
+        self.hostcall_fuel
+    }
+
+    pub(crate) fn set_hostcall_fuel(&mut self, fuel: usize) {
+        self.hostcall_fuel = fuel;
+    }
+
     pub(crate) fn set_exports(&mut self, exports: Exports) {
         self.exports = Some(exports);
     }
@@ -73,61 +75,28 @@ impl Invocation {
         self.exports.clone()
     }
 
-    /// The first trap wins: the upcall import ends the invocation with it,
-    /// so the trap keeps its own kind.
-    pub(crate) fn record_reentry_trap(&mut self, trap: wasmtime::Error) {
-        self.reentry_trap.get_or_insert(trap);
-    }
-
-    /// Once a callback into the guest has trapped, the guest is past
-    /// resuming.
-    pub(crate) fn reentry_trapped(&self) -> bool {
-        self.reentry_trap.is_some()
-    }
-
-    pub(crate) fn take_reentry_trap(&mut self) -> Option<wasmtime::Error> {
-        self.reentry_trap.take()
-    }
-
-    pub(crate) fn install_wasi(
-        &mut self,
-        wasi: WasiP1Ctx,
-        stdout: MemoryOutputPipe,
-        stderr: MemoryOutputPipe,
-    ) {
-        self.wasi = Some(wasi);
-        self.stdout_pipe = Some(stdout);
-        self.stderr_pipe = Some(stderr);
+    /// The state of the Calls open on this instance.
+    pub(crate) fn calls_mut(&mut self) -> &mut Calls<wasmtime::Error> {
+        &mut self.calls
     }
 
     pub(crate) fn bind_on_dispatch(&mut self, handler: Arc<dyn DispatchHandler>) {
-        self.on_dispatch = Some(handler);
-    }
-
-    pub(crate) fn stdout_bytes(&self) -> Vec<u8> {
-        self.stdout_pipe
-            .as_ref()
-            .map(|p| p.contents().to_vec())
-            .unwrap_or_default()
+        self.calls.bind(handler);
     }
 
-    pub(crate) fn stderr_bytes(&self) -> Vec<u8> {
-        self.stderr_pipe
-            .as_ref()
-            .map(|p| p.contents().to_vec())
-            .unwrap_or_default()
+    /// The trap ending an invocation its host gave up on.
+    pub(crate) fn abandonment(&self) -> Option<wasmtime::Error> {
+        self.calls
+            .abandoned()
+            .then(|| wasmtime::Error::msg(ABANDONED))
     }
 
-    /// A clone, so the borrow on the `Caller` is released and the
-    /// dispatcher can re-borrow it to write the Reply.
-    pub(crate) fn on_dispatch(&self) -> Option<Arc<dyn DispatchHandler>> {
-        self.on_dispatch.clone()
+    pub(crate) fn wasi(&self) -> &Wasi {
+        &self.wasi
     }
 
-    pub(crate) fn wasi_mut(&mut self) -> &mut WasiP1Ctx {
-        self.wasi
-            .as_mut()
-            .expect("WASI context not initialised — the driver must install it before any WASI use")
+    pub(crate) fn wasi_mut(&mut self) -> &mut Wasi {
+        &mut self.wasi
     }
 
     pub(crate) fn set_deadline(&mut self, deadline: Option<Instant>) {
@@ -146,13 +115,13 @@ impl Invocation {
     /// image's initial allocation and the high-water mark of earlier
     /// invocations do not consume the budget.
     pub(crate) fn arm_memory_cap(&mut self, baseline: usize) {
-        self.limiter.activate(baseline);
+        self.limiter.0.activate(baseline);
     }
 
     /// Disarmed once the export returns, so post-run host bookkeeping such
     /// as reading the answer back is not charged to the script.
     pub(crate) fn disarm_memory_cap(&mut self) {
-        self.limiter.deactivate();
+        self.limiter.0.deactivate();
     }
 
     /// Stamped right before the guest export call, so the bracket matches
@@ -172,6 +141,6 @@ impl Invocation {
     }
 
     pub(crate) fn memory_peak(&self) -> usize {
-        self.limiter.peak()
+        self.limiter.0.peak()
     }
 }

data/crates/kobako-wasmtime/src/lib.rs

@@ -13,17 +13,18 @@
 mod abi;
 mod ambient;
 mod cache;
-mod capture;
 mod config;
-mod dispatch;
+mod descriptors;
 mod driver;
 mod exports;
-mod guest_mem;
+mod guest;
+mod imports;
 mod instance_pre;
 mod invocation;
-mod io;
 mod limiter;
+mod preview1;
 mod trap;
+mod wasi;
 
 pub use config::Config;
 pub use driver::Driver;

data/crates/kobako-wasmtime/src/limiter.rs

@@ -1,65 +1,11 @@
-//! The per-invocation memory cap: only the `memory.grow` delta past the
-//! linear-memory size captured at invocation entry is charged, so the
-//! image's initial allocation is outside the budget.
+//! Hands wasmtime's growth hook to the runtime's per-invocation memory cap.
 
+use kobako_runtime::endpoint::MemoryCap;
 use wasmtime::ResourceLimiter;
 
-/// Resource limiter that enforces the per-invocation `memory_limit`
-/// cap.
-///
-/// `max_memory` is the byte cap on per-invocation growth (`None` disables
-/// the cap). `baseline` is the linear-memory size captured at invocation
-/// entry by `MemoryLimiter::activate`; the limiter charges only the
-/// `memory.grow` delta past `baseline` against `max_memory`, so the
-/// mruby image's initial allocation and any high-water mark left by
-/// prior invocations on the same Sandbox do not consume the budget.
-/// `cap_active` gates whether the cap is enforced — wasmtime's
-/// `ResourceLimiter` also fires for the module's declared initial
-/// allocation at instantiation time, but the cap stays dormant until
-/// `MemoryLimiter::activate` flips the flag for one `Driver` invoke
-/// call. When `cap_active` is `false`, the limiter always allows
-/// growth.
-///
-/// When `memory.grow` would push the per-invocation delta past
-/// `max_memory`, the limiter returns `MemoryLimitTrap` from
-/// `memory_growing`; wasmtime turns that into the trap surfaced to the
-/// host as a guest invocation failure.
-#[derive(Debug, Clone, Copy)]
-pub(crate) struct MemoryLimiter {
-    max_memory: Option<usize>,
-    baseline: usize,
-    cap_active: bool,
-    peak: usize,
-}
-
-impl MemoryLimiter {
-    pub(crate) fn new(max_memory: Option<usize>) -> Self {
-        Self {
-            max_memory,
-            baseline: 0,
-            cap_active: false,
-            peak: 0,
-        }
-    }
-
-    /// The cap stays dormant until armed, because the module's declared
-    /// initial memory is allocated during instantiation and must pass.
-    pub(crate) fn activate(&mut self, baseline: usize) {
-        self.baseline = baseline;
-        self.cap_active = true;
-        self.peak = 0;
-    }
-
-    pub(crate) fn deactivate(&mut self) {
-        self.cap_active = false;
-    }
-
-    /// A grow the cap rejects never updates the peak, so the reported value
-    /// never exceeds `memory_limit`.
-    pub(crate) fn peak(&self) -> usize {
-        self.peak
-    }
-}
+/// wasmtime asks this as `memory.grow` runs; an `Err` becomes the trap
+/// surfaced to the host as a guest invocation failure.
+pub(crate) struct MemoryLimiter(pub(crate) MemoryCap);
 
 impl ResourceLimiter for MemoryLimiter {
     fn memory_growing(
@@ -68,18 +14,7 @@ impl ResourceLimiter for MemoryLimiter {
         desired: usize,
         _maximum: Option<usize>,
     ) -> wasmtime::Result<bool> {
-        if !self.cap_active {
-            return Ok(true);
-        }
-        let delta = desired.saturating_sub(self.baseline);
-        if let Some(limit) = self.max_memory {
-            if delta > limit {
-                return Err(wasmtime::Error::new(MemoryLimitTrap { desired, limit }));
-            }
-        }
-        if delta > self.peak {
-            self.peak = delta;
-        }
+        self.0.charge(desired).map_err(wasmtime::Error::new)?;
         Ok(true)
     }
 
@@ -92,162 +27,3 @@ impl ResourceLimiter for MemoryLimiter {
         Ok(true)
     }
 }
-
-/// Marker error returned from `MemoryLimiter::memory_growing` when the
-/// per-invocation memory cap is exceeded. Downcast from the wasmtime
-/// trap error to classify the failure as a memory-limit `Trap`. Callers
-/// use the `Display` impl below — no field is read directly — so the
-/// inner state stays private.
-#[derive(Debug)]
-pub(crate) struct MemoryLimitTrap {
-    desired: usize,
-    limit: usize,
-}
-
-impl MemoryLimitTrap {
-    #[cfg(test)]
-    pub(crate) fn new(desired: usize, limit: usize) -> Self {
-        Self { desired, limit }
-    }
-}
-
-impl std::fmt::Display for MemoryLimitTrap {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "memory usage exceeded memory_limit: \
-             requested={} bytes, limit={} bytes",
-            self.desired, self.limit
-        )
-    }
-}
-
-impl std::error::Error for MemoryLimitTrap {}
-
-#[cfg(test)]
-mod tests {
-    //! Unit tests for `MemoryLimiter` — the per-invocation memory
-    //! delta cap. The Ruby-facing E2E suite exercises the full path
-    //! through wasmtime; these tests pin the pure delta arithmetic so
-    //! a regression that breaks the baseline accounting (e.g. dropping
-    //! the `baseline` subtraction, or letting `activate` carry stale
-    //! state across invocations) is caught without spinning up a
-    //! Store.
-    use super::{MemoryLimitTrap, MemoryLimiter};
-    use wasmtime::ResourceLimiter;
-
-    fn assert_growing(limiter: &mut MemoryLimiter, desired: usize) {
-        assert!(
-            limiter.memory_growing(0, desired, None).unwrap(),
-            "expected memory_growing({desired}) to allow growth"
-        );
-    }
-
-    fn assert_trapping(limiter: &mut MemoryLimiter, desired: usize) {
-        let err = limiter
-            .memory_growing(0, desired, None)
-            .expect_err("expected memory_growing to trap");
-        assert!(
-            err.downcast_ref::<MemoryLimitTrap>().is_some(),
-            "expected MemoryLimitTrap, got {err:?}"
-        );
-    }
-
-    // @behavior S-113
-    #[test]
-    fn dormant_limiter_allows_any_growth() {
-        let mut limiter = MemoryLimiter::new(Some(1 << 20));
-        // Without `activate`, the cap is dormant — the module's
-        // declared initial allocation must pass through unconditionally.
-        assert_growing(&mut limiter, 100 << 20);
-    }
-
-    // @behavior S-114
-    #[test]
-    fn delta_below_cap_passes_after_activate() {
-        let mut limiter = MemoryLimiter::new(Some(1 << 20));
-        limiter.activate(2 << 20);
-        // baseline=2 MiB, desired=2.5 MiB → delta=0.5 MiB ≤ 1 MiB cap.
-        assert_growing(&mut limiter, (2 << 20) + (1 << 19));
-    }
-
-    // @behavior S-121
-    #[test]
-    fn delta_past_cap_traps_with_memory_limit_trap() {
-        let mut limiter = MemoryLimiter::new(Some(1 << 20));
-        limiter.activate(2 << 20);
-        // baseline=2 MiB, desired=4 MiB → delta=2 MiB > 1 MiB cap.
-        assert_trapping(&mut limiter, 4 << 20);
-    }
-
-    // @behavior S-008
-    #[test]
-    fn activate_resets_baseline_on_each_invocation() {
-        let mut limiter = MemoryLimiter::new(Some(1 << 20));
-        limiter.activate(2 << 20);
-        assert_growing(&mut limiter, (2 << 20) + (1 << 20));
-        // Second invocation: linear memory has grown to 3 MiB. Re-arming
-        // must re-anchor the baseline so the next 1 MiB of growth fits
-        // the per-invocation budget rather than being charged against
-        // the prior invocation's residue.
-        limiter.activate(3 << 20);
-        assert_growing(&mut limiter, (3 << 20) + (1 << 20));
-    }
-
-    // @behavior S-009
-    #[test]
-    fn disabled_cap_ignores_delta_size() {
-        let mut limiter = MemoryLimiter::new(None);
-        limiter.activate(0);
-        assert_growing(&mut limiter, 100 << 20);
-    }
-
-    // @behavior S-116
-    #[test]
-    fn peak_starts_at_zero_before_any_grow() {
-        let limiter = MemoryLimiter::new(Some(1 << 20));
-        assert_eq!(limiter.peak(), 0);
-    }
-
-    // @behavior S-115
-    #[test]
-    fn peak_tracks_high_water_of_delta_past_baseline() {
-        let mut limiter = MemoryLimiter::new(Some(1 << 20));
-        limiter.activate(2 << 20);
-        assert_growing(&mut limiter, (2 << 20) + (1 << 18)); // delta=256 KiB
-        assert_growing(&mut limiter, (2 << 20) + (1 << 19)); // delta=512 KiB (new peak)
-        assert_growing(&mut limiter, (2 << 20) + (1 << 17)); // delta=128 KiB (below peak)
-        assert_eq!(limiter.peak(), 1 << 19);
-    }
-
-    // @behavior S-062
-    #[test]
-    fn trap_does_not_update_peak() {
-        let mut limiter = MemoryLimiter::new(Some(1 << 20));
-        limiter.activate(2 << 20);
-        assert_growing(&mut limiter, (2 << 20) + (1 << 19)); // delta=512 KiB
-        assert_trapping(&mut limiter, (2 << 20) + (2 << 20)); // would be 2 MiB > 1 MiB cap
-                                                              // Peak reflects the last accepted grow, not the rejected desired.
-        assert_eq!(limiter.peak(), 1 << 19);
-    }
-
-    // @behavior S-117
-    #[test]
-    fn activate_resets_peak_for_new_invocation() {
-        let mut limiter = MemoryLimiter::new(Some(1 << 20));
-        limiter.activate(2 << 20);
-        assert_growing(&mut limiter, (2 << 20) + (1 << 19));
-        assert_eq!(limiter.peak(), 1 << 19);
-        limiter.activate(3 << 20);
-        assert_eq!(limiter.peak(), 0);
-    }
-
-    // @behavior S-118
-    #[test]
-    fn disabled_cap_still_tracks_peak() {
-        let mut limiter = MemoryLimiter::new(None);
-        limiter.activate(1 << 20);
-        assert_growing(&mut limiter, (1 << 20) + (4 << 20));
-        assert_eq!(limiter.peak(), 4 << 20);
-    }
-}