diff --git a/book/src/SUMMARY.md b/book/src/SUMMARY.md index e3e716ae..e9c4af9d 100644 --- a/book/src/SUMMARY.md +++ b/book/src/SUMMARY.md @@ -24,6 +24,7 @@ - [MCP Server](./guide/mcp-server.md) - [gRPC Server](./guide/grpc-server.md) - [Pre-compiling Runtimes](./guide/precompile.md) +- [Performance Tuning](./guide/performance.md) - [CLI](./guide/cli.md) # API Reference diff --git a/book/src/guide/performance.md b/book/src/guide/performance.md new file mode 100644 index 00000000..a20e85ac --- /dev/null +++ b/book/src/guide/performance.md @@ -0,0 +1,82 @@ +# Performance Tuning + +Most of the cost of a short `Sandbox::execute()` is not running Python. It is +creating the WebAssembly instance the code runs on and tearing it down again: +mapping the pre-initialized heap image, building a function reference for +every entry in the runtime's ~6k-entry function table, and, afterwards, +resetting the memory slot. Running `pass` on a fresh instance takes about +1.3 ms on a fast workstation; the Python part of that is roughly 0.2 ms. + +Eryx hides that cost in two ways, both on by default. This page explains what +they do and the environment variables that tune them. All of these are read +once, when the process-wide wasmtime engine is created, so set them in the +environment of the process that embeds eryx (for `pyeryx`, the Python process). + +## Warm instances + +After every stateless execution, a background task instantiates a replacement +store for the same runtime and parks it in a process-wide pool. The next +`execute()` takes it instead of instantiating, and the used store is dropped on +the background task too. Every execution still runs on an instance that has +never run user code, so isolation is unchanged; only the timing moves. + +The pool is keyed by runtime component rather than by `Sandbox`, so the common +pattern of building a short-lived sandbox per request (for example +`SandboxFactory.create_sandbox()` from Python) benefits from it: the first +request in a process pays for instantiation, later ones do not. + +| Variable | Default | Effect | +|----------|---------|--------| +| `ERYX_WARM_INSTANCES` | `1` | Instances to keep ready per runtime. `0` disables the pool. | + +One ready instance is enough for a caller that executes serially. If several +tasks execute concurrently, raise it towards that concurrency; each ready +instance costs about the resident size of the pre-initialized heap (tens of +megabytes, mostly shared copy-on-write pages) plus wasmtime's per-instance +metadata. + +The pool is only used on a multi-threaded Tokio runtime, where the background +work actually runs in parallel. On a current-thread runtime it would only add to +the next request's latency, so it stays off. `PythonExecutor::warm_instances_ready()` +reports how many instances are waiting for that executor's runtime. + +## Instance allocation + +Wasmtime can allocate each instance's linear memory, tables, and async stacks +on demand (a fresh `mmap` per instance) or from a pool of pre-reserved slots. +Eryx uses the pooling allocator: a slot that has just been released is reused +by the next instance, so its pages are still mapped, and on Linux only the +pages the previous instance dirtied are reset. + +| Variable | Default | Effect | +|----------|---------|--------| +| `ERYX_ALLOCATOR` | `pooling` | `pooling` or `on-demand`. | +| `ERYX_POOL_INSTANCES` | `1000` | Maximum instances alive at once. Every running `execute()` and every live session holds one; instantiation fails once the pool is full. | +| `ERYX_POOL_KEEP_RESIDENT_MB` | `64` | How much of each linear memory to keep mapped between uses. | +| `ERYX_POOL_PAGEMAP_SCAN` | `1` | `1`: reset only dirty pages (Linux 6.7+, via `PAGEMAP_SCAN`). `0`: `memcpy` the whole keep-resident budget instead, which trades a larger copy for fewer page faults on the next execution. | + +The pool reserves virtual address space for every slot up front — several +terabytes with the defaults, which is normal for wasmtime deployments but can +be refused by a host with a low `ulimit -v`. If the pool cannot be created, eryx +logs a warning and falls back to on-demand allocation; execution behaves +identically either way. + +If a process holds more than `ERYX_POOL_INSTANCES` sessions open at once, +raise the limit or switch to `on-demand`. Note that the choice of allocator does +not affect precompiled `.cwasm` artifacts; only compilation settings do. + +## Measuring + +`crates/eryx/examples/profile_stateless.rs` times the stateless path and is a +convenient target for `perf` or `samply`: + +```bash +cargo build --example profile_stateless --features embedded --release +# 2000 executions of `pass`, with a 2 ms gap between them so the background +# replenishment gets the same chance it has between real requests +./target/release/examples/profile_stateless 2000 pass 2000 +ERYX_WARM_INSTANCES=0 ./target/release/examples/profile_stateless 2000 pass 2000 +``` + +The criterion benchmarks (`cargo bench --package eryx --features embedded`) +cover the same paths with callbacks registered. diff --git a/book/src/guide/resource-limits.md b/book/src/guide/resource-limits.md index 141c177e..98ed37c9 100644 --- a/book/src/guide/resource-limits.md +++ b/book/src/guide/resource-limits.md @@ -260,6 +260,10 @@ Fuel limits provide fine-grained control over execution by limiting the number o - Billing based on actual computation performed - Preventing CPU-intensive attacks +Fuel counts the instructions your code executes; instantiating the sandbox's +WebAssembly instance is not charged, so the number does not depend on whether +the execution ran on a fresh or a pre-instantiated instance. + ```rust # extern crate eryx; diff --git a/book/src/guide/sandboxes.md b/book/src/guide/sandboxes.md index fba8aaf2..618fb0c7 100644 --- a/book/src/guide/sandboxes.md +++ b/book/src/guide/sandboxes.md @@ -230,7 +230,7 @@ If you need state to persist across executions, use a [Session](./sessions.md). ## SandboxFactory for Fast Creation -When creating many sandboxes, use `SandboxFactory` to pre-initialize Python and packages once, then quickly instantiate sandboxes from that snapshot: +When creating many sandboxes, use `SandboxFactory` to pre-initialize Python and packages once, then quickly instantiate sandboxes from that snapshot. Executions on short-lived sandboxes also pick up a pre-instantiated WebAssembly instance from a process-wide pool; see [Performance Tuning](./performance.md) for how that works and how to size it. ```rust diff --git a/crates/eryx-python/README.md b/crates/eryx-python/README.md index ff8aebf5..72709811 100644 --- a/crates/eryx-python/README.md +++ b/crates/eryx-python/README.md @@ -85,6 +85,15 @@ print(f"Execution took {(time.perf_counter() - start) * 1000:.1f}ms") For repeated sandbox creation with custom packages, see [`SandboxFactory`](#sandboxfactory) below. +Each `execute()` on a fresh sandbox also has to instantiate the WebAssembly +runtime. Eryx hides most of that by keeping a pre-instantiated instance ready +in the background and handing it to the next execution; the environment +variables that tune this (`ERYX_WARM_INSTANCES`, `ERYX_ALLOCATOR`, ...) are +documented in the +[Performance Tuning](https://docs.eryx.run/guide/performance.html) +guide. Set them in the environment of the Python process before importing +`eryx`. + ## API Reference **Core Classes:** diff --git a/crates/eryx/Cargo.toml b/crates/eryx/Cargo.toml index 924759d9..1f0bb7c0 100644 --- a/crates/eryx/Cargo.toml +++ b/crates/eryx/Cargo.toml @@ -48,6 +48,10 @@ name = "precompile" name = "profile_execution" required-features = ["embedded"] +[[example]] +name = "profile_stateless" +required-features = ["embedded"] + [[example]] name = "resource_limits" required-features = ["embedded"] @@ -115,7 +119,7 @@ tar.workspace = true target-lexicon = "0.13" tempfile.workspace = true thiserror.workspace = true -tokio = { workspace = true, features = ["sync", "net", "io-util", "time"] } +tokio = { workspace = true, features = ["sync", "net", "io-util", "time", "rt"] } tokio-rustls.workspace = true tokio-util.workspace = true tracing.workspace = true diff --git a/crates/eryx/benches/execution.rs b/crates/eryx/benches/execution.rs index 648afcd8..a1ddd891 100644 --- a/crates/eryx/benches/execution.rs +++ b/crates/eryx/benches/execution.rs @@ -158,8 +158,10 @@ fn bench_sandbox_creation(c: &mut Criterion) { /// Benchmark stateless execution via `Sandbox::execute()`. /// -/// Each call creates a fresh WASM instance and initializes Python from scratch. -/// This is ~500ms per execution due to Python interpreter initialization. +/// Each call runs on a fresh WASM instance. Python itself is already +/// initialized in the pre-initialized snapshot, so the per-call cost is +/// instantiation (or taking a warm instance from the pool, since this runtime +/// is multi-threaded), the per-execute callback setup, and the code itself. /// /// Use this when you need complete isolation between executions. fn bench_stateless_execution(c: &mut Criterion) { @@ -168,7 +170,7 @@ fn bench_stateless_execution(c: &mut Criterion) { let mut group = c.benchmark_group("stateless_execution"); - // Stateless execution is slow (~500ms), so reduce sample size + // Stateless execution is the slowest path, so reduce sample size group.sample_size(10); group.measurement_time(Duration::from_secs(10)); diff --git a/crates/eryx/examples/profile_stateless.rs b/crates/eryx/examples/profile_stateless.rs new file mode 100644 index 00000000..0d732921 --- /dev/null +++ b/crates/eryx/examples/profile_stateless.rs @@ -0,0 +1,70 @@ +//! Profiling harness for stateless (fresh-instance) execution overhead. +//! +//! Each iteration goes through the full `Sandbox::execute()` path: a new +//! `Store`, `instantiate_async` from the cached `InstancePre`, one `execute` +//! export call, and teardown. This is the path a `SandboxFactory` render +//! takes, so it is the one to profile for per-request latency. +//! +//! Run with samply: +//! cargo build --example profile_stateless --features embedded --release +//! samply record ./target/release/examples/profile_stateless +//! +//! Or count page faults per execution: +//! perf stat -e page-faults ./target/release/examples/profile_stateless 1000 +//! +//! Arguments: `[iterations] [code] [gap_us]`. `gap_us` sleeps between +//! executions so background work (warm-instance replenishment) gets the same +//! chance it has between real requests; only the executions are timed. + +use std::time::{Duration, Instant}; + +use eryx::Sandbox; + +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let iterations: u32 = args.next().and_then(|s| s.parse().ok()).unwrap_or(2000); + let code = args.next().unwrap_or_else(|| "pass".to_string()); + let gap = Duration::from_micros(args.next().and_then(|s| s.parse().ok()).unwrap_or(0)); + + let rt = tokio::runtime::Runtime::new()?; + + rt.block_on(async { + eprintln!("Creating sandbox..."); + let sandbox = Sandbox::embedded().build()?; + let executor = sandbox.executor(); + + eprintln!("Warming up (10 iterations)..."); + for _ in 0..10 { + sandbox.execute(&code).await?; + std::thread::sleep(gap); + } + + eprintln!("Profiling {iterations} iterations of {code:?} with a {gap:?} gap..."); + let mut executing = Duration::ZERO; + let mut warm_hits = 0u32; + + for _ in 0..iterations { + if executor.warm_instances_ready() > 0 { + warm_hits += 1; + } + let start = Instant::now(); + sandbox.execute(&code).await?; + executing += start.elapsed(); + std::thread::sleep(gap); + } + + eprintln!("\nResults:"); + eprintln!(" Time executing: {executing:?}"); + eprintln!(" Iterations: {iterations}"); + eprintln!(" Warm instance available: {warm_hits}"); + eprintln!(" Average: {:?} per execution", executing / iterations); + eprintln!( + " Throughput: {:.0} executions/sec", + iterations as f64 / executing.as_secs_f64() + ); + + Ok::<_, Box>(()) + })?; + + Ok(()) +} diff --git a/crates/eryx/src/lib.rs b/crates/eryx/src/lib.rs index 90a7c5ed..7bab4a83 100644 --- a/crates/eryx/src/lib.rs +++ b/crates/eryx/src/lib.rs @@ -68,6 +68,7 @@ mod schema; pub mod secrets; pub mod session; mod trace; +mod warm; mod wasm; /// Pre-initialization support for capturing Python memory state. diff --git a/crates/eryx/src/warm.rs b/crates/eryx/src/warm.rs new file mode 100644 index 00000000..2fefd323 --- /dev/null +++ b/crates/eryx/src/warm.rs @@ -0,0 +1,191 @@ +//! Pre-instantiated ("warm") stores for stateless execution. +//! +//! Instantiating the runtime component costs far more than running a short +//! script on it. Wasmtime has to build a `VMFuncRef` for every entry of the +//! dynamically linked modules' function tables (their element segments are +//! offset by an imported `__table_base`, so they cannot be initialised +//! lazily), and dropping the store afterwards resets the linear-memory slot. +//! For `pass` that is roughly 80% of the wall time. +//! +//! [`WarmPool`] keeps instantiated stores ready so [`crate::Sandbox::execute`] +//! can pick one up immediately. The used store is handed to a background task +//! that drops it and instantiates a replacement, moving both costs off the +//! request path. Isolation is unchanged: every execution still runs on an +//! instance that has never run user code. +//! +//! The pool is process-global and keyed by component, so short-lived +//! `Sandbox` values (the `SandboxFactory` pattern) share it. It is only active +//! on a multi-threaded Tokio runtime, where the background work actually runs +//! in parallel; on a current-thread runtime it would just be added to the next +//! request's latency. + +use std::collections::HashMap; +use std::sync::{Mutex, OnceLock}; + +use tokio::runtime::{Handle, RuntimeFlavor}; +use wasmtime::Store; + +use crate::wasm::{ExecutorState, Sandbox, SandboxPre}; + +/// An instantiated store that has not run user code yet. +pub(crate) struct WarmInstance { + pub(crate) store: Store, + pub(crate) bindings: Sandbox, +} + +#[derive(Default)] +struct Slot { + ready: Vec, + /// Replacements a background task has committed to instantiate. + pending: usize, +} + +/// Process-global pool of warm instances, one slot per component. +pub(crate) struct WarmPool { + /// Instances to keep ready per component; `0` disables the pool. + target: usize, + slots: Mutex>, +} + +impl WarmPool { + /// Default number of ready instances per component. + /// + /// One is enough to hide instantiation from a caller that executes + /// serially, which is the `SandboxFactory` render pattern; concurrent + /// callers should raise `ERYX_WARM_INSTANCES` towards their concurrency. + const DEFAULT_TARGET: usize = 1; + + pub(crate) fn global() -> &'static Self { + static POOL: OnceLock = OnceLock::new(); + POOL.get_or_init(|| { + let target = std::env::var("ERYX_WARM_INSTANCES") + .ok() + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(Self::DEFAULT_TARGET); + WarmPool { + target, + slots: Mutex::new(HashMap::new()), + } + }) + } + + /// Whether this call site can use the pool. + fn active(&self) -> bool { + self.target > 0 + && Handle::try_current() + .is_ok_and(|handle| handle.runtime_flavor() == RuntimeFlavor::MultiThread) + } + + /// Take a ready instance for `pre` whose instantiation fits in `memory_limit`. + /// + /// A limit below the snapshot's baseline has to fail the way a fresh + /// instantiation does, so such executions are left to the cold path. + pub(crate) fn take( + &self, + pre: &SandboxPre, + memory_limit: Option, + ) -> Option { + if !self.active() { + return None; + } + let mut slots = self.slots.lock().ok()?; + let ready = &mut slots.get_mut(&key(pre))?.ready; + let baseline = ready + .last()? + .store + .data() + .memory_tracker + .peak_memory_bytes(); + if memory_limit.is_some_and(|limit| baseline > limit) { + return None; + } + ready.pop() + } + + /// Number of ready instances for `pre`. + pub(crate) fn ready(&self, pre: &SandboxPre) -> usize { + self.slots + .lock() + .ok() + .and_then(|slots| slots.get(&key(pre)).map(|slot| slot.ready.len())) + .unwrap_or(0) + } + + /// Dispose of a store that has run user code and top the pool back up. + /// + /// Both happen on a background task when the pool is active; otherwise + /// the store is dropped here. + pub(crate) fn retire( + &'static self, + pre: &SandboxPre, + mut store: Store, + ) { + if !self.active() { + drop(store); + return; + } + + // The caller's handler tasks run until every sender of their channel + // is gone, so release those now rather than when the background task + // gets around to dropping the store. + store.data_mut().disconnect(); + + let key = key(pre); + let replenish = self.reserve(key); + let pre = pre.clone(); + tokio::spawn(async move { + // Dropping first returns the memory slot to wasmtime's pool, so the + // replacement below can reuse it while it is still cache-hot. + drop(store); + if !replenish { + return; + } + match instantiate_warm(&pre).await { + Ok(instance) => self.fulfil(key, Some(instance)), + Err(e) => { + tracing::warn!(error = %e, "failed to pre-instantiate a warm instance"); + self.fulfil(key, None); + } + } + }); + } + + /// Commit to instantiating one more instance for `key` if the slot is + /// below target, counting instances already in flight. + fn reserve(&self, key: usize) -> bool { + let Ok(mut slots) = self.slots.lock() else { + return false; + }; + let slot = slots.entry(key).or_default(); + if slot.ready.len() + slot.pending >= self.target { + return false; + } + slot.pending += 1; + true + } + + fn fulfil(&self, key: usize, instance: Option) { + let Ok(mut slots) = self.slots.lock() else { + return; + }; + let slot = slots.entry(key).or_default(); + slot.pending = slot.pending.saturating_sub(1); + if let Some(instance) = instance { + slot.ready.push(instance); + } + } +} + +/// Identify the component behind `pre`. +/// +/// The compiled image's address is unique for as long as the component is +/// loaded, and a warm store keeps its component loaded. +fn key(pre: &SandboxPre) -> usize { + pre.instance_pre().component().image_range().start as usize +} + +async fn instantiate_warm(pre: &SandboxPre) -> Result { + let (store, bindings) = + crate::wasm::instantiate_store(pre, ExecutorState::placeholder(), u64::MAX).await?; + Ok(WarmInstance { store, bindings }) +} diff --git a/crates/eryx/src/wasm.rs b/crates/eryx/src/wasm.rs index 4818c5c8..1c86eddf 100644 --- a/crates/eryx/src/wasm.rs +++ b/crates/eryx/src/wasm.rs @@ -34,12 +34,16 @@ use crate::cache::{CacheKey, InstancePreCache}; use tokio::sync::{mpsc, oneshot}; use tokio_util::sync::CancellationToken; use wasmtime::component::{Accessor, Component, HasSelf, Linker, ResourceTable}; -use wasmtime::{AsContextMut, Config, Engine, ResourceLimiter, Store, UpdateDeadline}; +use wasmtime::{ + AsContextMut, Config, Enabled, Engine, InstanceAllocationStrategy, PoolingAllocationConfig, + ResourceLimiter, Store, UpdateDeadline, +}; use wasmtime_wasi::{FsPerms, WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView}; use crate::callback::Callback; use crate::error::Error; use crate::trace::TraceEvent; +use crate::warm::{WarmInstance, WarmPool}; /// Interval between increments of the process-wide epoch ticker. pub(crate) const EPOCH_TICK_INTERVAL: Duration = Duration::from_millis(10); @@ -489,6 +493,132 @@ impl std::str::FromStr for CpuFeatureLevel { } } +/// How the shared engine allocates instance resources (linear memories, +/// tables, and async stacks). +/// +/// Read from the environment rather than a builder because the engine is +/// process-global ([`PythonExecutor::shared_engine`]) and is created before +/// any sandbox configuration exists. +/// +/// - `ERYX_ALLOCATOR`: `pooling` (default) or `on-demand`. +/// - `ERYX_POOL_INSTANCES`: maximum concurrently live instances with the +/// pooling allocator (default 1000). Every live `Sandbox::execute()` and +/// every live session holds one; instantiation fails once the pool is full. +/// - `ERYX_POOL_KEEP_RESIDENT_MB`: how much of each linear memory to keep +/// mapped between uses (default 64). Dirty pages within this budget are +/// restored with `memcpy` instead of being unmapped, so the next instance +/// in the slot takes no page faults for them. +/// - `ERYX_POOL_PAGEMAP_SCAN`: `1` (default) lets wasmtime find the dirty +/// pages with Linux's `PAGEMAP_SCAN` ioctl (6.7+) and reset only those; +/// `0` resets the whole keep-resident budget with `memcpy` instead, which +/// trades a larger copy for fewer page faults on the next execution. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct AllocatorSettings { + pooling: bool, + pool_instances: u32, + keep_resident_bytes: usize, + pagemap_scan: bool, +} + +impl Default for AllocatorSettings { + fn default() -> Self { + Self { + pooling: true, + pool_instances: 1000, + keep_resident_bytes: 64 << 20, + pagemap_scan: true, + } + } +} + +impl AllocatorSettings { + /// Upper bound on funcref table growth per instance. + /// + /// The runtime's function table starts at ~6k entries and grows when + /// native extensions are `dlopen`ed, so this sits well above wasmtime's + /// 20k default while still bounding the per-slot address reservation. + const TABLE_ELEMENTS: usize = 1 << 20; + + pub(crate) fn from_env() -> std::result::Result { + Self::from_lookup(|name| std::env::var(name).ok()) + } + + /// Parse the settings from `lookup`, a view of the environment. + fn from_lookup(lookup: impl Fn(&str) -> Option) -> std::result::Result { + let mut settings = Self::default(); + + match lookup("ERYX_ALLOCATOR") { + Some(v) if v.eq_ignore_ascii_case("pooling") => settings.pooling = true, + Some(v) if v.eq_ignore_ascii_case("on-demand") => settings.pooling = false, + Some(v) => { + return Err(Error::WasmEngine(format!( + "Unknown allocator '{v}' in ERYX_ALLOCATOR. Valid values: pooling, on-demand" + ))); + } + None => {} + } + + if let Some(v) = lookup("ERYX_POOL_INSTANCES") { + settings.pool_instances = + v.trim() + .parse::() + .ok() + .filter(|n| *n > 0) + .ok_or_else(|| { + Error::WasmEngine(format!( + "ERYX_POOL_INSTANCES must be a positive integer, got '{v}'" + )) + })?; + } + + if let Some(v) = lookup("ERYX_POOL_KEEP_RESIDENT_MB") { + let mib = v.trim().parse::().map_err(|_| { + Error::WasmEngine(format!( + "ERYX_POOL_KEEP_RESIDENT_MB must be a non-negative integer, got '{v}'" + )) + })?; + settings.keep_resident_bytes = mib << 20; + } + + match lookup("ERYX_POOL_PAGEMAP_SCAN") { + Some(v) if v.trim() == "0" => settings.pagemap_scan = false, + Some(v) if v.trim() == "1" => settings.pagemap_scan = true, + Some(v) => { + return Err(Error::WasmEngine(format!( + "ERYX_POOL_PAGEMAP_SCAN must be 0 or 1, got '{v}'" + ))); + } + None => {} + } + + Ok(settings) + } + + /// The pooling configuration to use, or `None` for on-demand allocation. + pub(crate) fn pooling_config(&self) -> Option { + if !self.pooling { + return None; + } + + let mut pooling = PoolingAllocationConfig::new(); + pooling + .total_component_instances(self.pool_instances) + .total_core_instances(self.pool_instances) + .total_memories(self.pool_instances) + .total_tables(self.pool_instances) + .total_stacks(self.pool_instances) + .table_elements(Self::TABLE_ELEMENTS) + .linear_memory_keep_resident(self.keep_resident_bytes) + .table_keep_resident(Self::TABLE_ELEMENTS * std::mem::size_of::()) + .pagemap_scan(if self.pagemap_scan { + Enabled::Auto + } else { + Enabled::No + }); + Some(pooling) + } +} + /// Tracks memory usage during WASM execution. /// /// This struct implements `ResourceLimiter` to intercept memory growth @@ -522,6 +652,11 @@ impl MemoryTracker { pub fn reset(&self) { self.peak_memory_bytes.store(0, Ordering::Relaxed); } + + /// Record memory that was already in use before this tracker was installed. + pub(crate) fn observe(&self, bytes: u64) { + self.peak_memory_bytes.fetch_max(bytes, Ordering::Relaxed); + } } impl ResourceLimiter for MemoryTracker { @@ -642,6 +777,73 @@ impl std::fmt::Debug for ExecutorState { } } +impl ExecutorState { + /// State for instantiating a store ahead of any execution. + /// + /// Instantiation runs no guest code beyond table initialisation, so the + /// guest never observes this state; [`PythonExecutor::execute_internal`] + /// replaces it wholesale with the real per-execution state before the + /// first export call. + pub(crate) fn placeholder() -> Self { + Self { + wasi: WasiCtxBuilder::new().build(), + table: ResourceTable::new(), + callback_tx: None, + trace_tx: None, + callbacks: Vec::new(), + memory_tracker: MemoryTracker::new(None), + net_tx: None, + output_tx: None, + #[cfg(feature = "vfs")] + hybrid_vfs_ctx: None, + suspended: None, + reuse_empty_callbacks: true, + } + } + + /// Drop every channel to the host so their receivers see end-of-stream. + pub(crate) fn disconnect(&mut self) { + self.callback_tx = None; + self.trace_tx = None; + self.net_tx = None; + self.output_tx = None; + } +} + +/// Create a store around `state` and instantiate `pre` into it, leaving +/// `initial_fuel` in the tank for user code. +/// +/// The epoch deadline is left effectively unbounded so instantiation cannot be +/// interrupted; callers arm the real deadline before running user code. +/// Instantiation burns fuel of its own (the modules' table-initialising start +/// functions), so it runs on an unlimited tank and `initial_fuel` is set +/// afterwards: `fuel_consumed` then measures only the user's code and is the +/// same whether or not the store came from the warm pool. +pub(crate) async fn instantiate_store( + pre: &SandboxPre, + state: ExecutorState, + initial_fuel: u64, +) -> std::result::Result<(Store, Sandbox), Error> { + let set_fuel = |store: &mut Store, fuel: u64| { + store + .set_fuel(fuel) + .map_err(|e| Error::Initialization(format!("Failed to set fuel: {e}"))) + }; + + let mut store = Store::new(pre.engine(), state); + store.limiter(|state| &mut state.memory_tracker); + store.set_epoch_deadline(u64::MAX / 2); + set_fuel(&mut store, u64::MAX)?; + + let bindings = pre + .instantiate_async(&mut store) + .await + .map_err(Error::WasmComponent)?; + + set_fuel(&mut store, initial_fuel)?; + Ok((store, bindings)) +} + impl WasiView for ExecutorState { fn ctx(&mut self) -> WasiCtxView<'_> { WasiCtxView { @@ -1336,6 +1538,19 @@ impl PythonExecutor { &self.instance_pre } + /// Number of pre-instantiated stores currently waiting for this + /// executor's component. + /// + /// Stateless executions take one of these instead of instantiating, and a + /// background task replaces it afterwards. The pool is shared by every + /// executor loading the same component and sized by `ERYX_WARM_INSTANCES` + /// (default 1, `0` disables it); it is only used on a multi-threaded Tokio + /// runtime. + #[must_use] + pub fn warm_instances_ready(&self) -> usize { + WarmPool::global().ready(&self.instance_pre) + } + /// Get the Python stdlib path if configured. #[must_use] pub fn python_stdlib_path(&self) -> Option<&PathBuf> { @@ -1965,36 +2180,7 @@ impl PythonExecutor { /// list of `flag=value` pairs. Example: /// `ERYX_CRANELIFT_FLAGS=has_avx512f=false,has_avx512bw=false` fn create_engine_with_target(target: Option<&str>) -> std::result::Result { - let mut config = Config::new(); - config.wasm_component_model(true); - // Enable component model async for the `invoke` callback function. - // The invoke function is async because Python code awaits on it. - // TCP/TLS functions are sync `func` in WIT but use fiber-based async - // on the host (via `async` bindgen flag) - they appear blocking to guest. - config.wasm_component_model_async(true); - - // Enable epoch-based interruption for execution timeouts. - // This allows us to interrupt WASM execution even in tight loops - // that don't yield to the async runtime (e.g., `while True: pass`). - config.epoch_interruption(true); - - // Enable fuel consumption for instruction tracking and limiting. - // Fuel provides fine-grained, deterministic execution bounds at the - // instruction level. Even when no limit is set, fuel consumption is - // tracked and reported for billing/metering purposes. - config.consume_fuel(true); - - // Enable copy-on-write heap images for faster instantiation - // This defers memory initialization from instantiation time to first write - config.memory_init_cow(true); - - // Optimize for smaller generated code (slight runtime perf tradeoff) - // This reduces .cwasm file sizes and memory footprint - config.cranelift_opt_level(wasmtime::OptLevel::SpeedAndSize); - - // Reduce async stack size from default 2 MiB to 512 KiB - // Python scripts don't need deep call stacks - config.async_stack_size(512 * 1024); + let mut config = Self::base_engine_config(); // Configure target triple for cross-compilation or portable builds. // Check explicit parameter first, then environment variable, then use native. @@ -2013,7 +2199,7 @@ impl PythonExecutor { #[cfg(any(feature = "embedded", feature = "preinit"))] Self::apply_cpu_feature_flags(&mut config)?; - Engine::new(&config).map_err(|e| Error::WasmEngine(e.to_string())) + Self::build_engine(config) } /// Create a configured wasmtime engine with explicit CPU feature control. @@ -2025,20 +2211,81 @@ impl PythonExecutor { target: Option<&str>, cpu_features: CpuFeatureLevel, ) -> std::result::Result { + let mut config = Self::base_engine_config(); + + // Target triple and CPU features are set together: pinning a triple is + // also what stops Cranelift inferring features from the build machine. + Self::apply_cpu_feature_level(&mut config, cpu_features, target)?; + + Self::build_engine(config) + } + + /// The engine options shared by every engine eryx creates. + /// + /// Everything here except the allocation strategy affects the generated + /// code, so a change to this function must be paired with a bump of + /// [`ENGINE_CONFIG_VERSION`](Self::ENGINE_CONFIG_VERSION). + fn base_engine_config() -> Config { let mut config = Config::new(); config.wasm_component_model(true); + // Enable component model async for the `invoke` callback function. + // The invoke function is async because Python code awaits on it. + // TCP/TLS functions are sync `func` in WIT but use fiber-based async + // on the host (via `async` bindgen flag) - they appear blocking to guest. config.wasm_component_model_async(true); + + // Enable epoch-based interruption for execution timeouts. + // This allows us to interrupt WASM execution even in tight loops + // that don't yield to the async runtime (e.g., `while True: pass`). config.epoch_interruption(true); + + // Enable fuel consumption for instruction tracking and limiting. + // Fuel provides fine-grained, deterministic execution bounds at the + // instruction level. Even when no limit is set, fuel consumption is + // tracked and reported for billing/metering purposes. config.consume_fuel(true); + + // Enable copy-on-write heap images for faster instantiation + // This defers memory initialization from instantiation time to first write config.memory_init_cow(true); + + // Optimize for smaller generated code (slight runtime perf tradeoff) + // This reduces .cwasm file sizes and memory footprint config.cranelift_opt_level(wasmtime::OptLevel::SpeedAndSize); + + // Reduce async stack size from default 2 MiB to 512 KiB + // Python scripts don't need deep call stacks config.async_stack_size(512 * 1024); - // Target triple and CPU features are set together: pinning a triple is - // also what stops Cranelift inferring features from the build machine. - Self::apply_cpu_feature_level(&mut config, cpu_features, target)?; + config + } - Engine::new(&config).map_err(|e| Error::WasmEngine(e.to_string())) + /// Build the engine, falling back to on-demand allocation if the pool + /// cannot be reserved. + /// + /// The pooling allocator reserves virtual address space for every slot up + /// front (terabytes with the default limits), which a host with a low + /// `ulimit -v` or an exotic kernel can refuse. Execution works identically + /// on either allocator, so that refusal only costs instantiation latency. + fn build_engine(mut config: Config) -> std::result::Result { + let allocator = AllocatorSettings::from_env()?; + let Some(pooling) = allocator.pooling_config() else { + return Engine::new(&config).map_err(|e| Error::WasmEngine(e.to_string())); + }; + + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pooling)); + match Engine::new(&config) { + Ok(engine) => Ok(engine), + Err(pooling_err) => { + tracing::warn!( + error = %pooling_err, + "failed to create wasmtime engine with the pooling allocator; \ + falling back to on-demand allocation" + ); + config.allocation_strategy(InstanceAllocationStrategy::OnDemand); + Engine::new(&config).map_err(|e| Error::WasmEngine(e.to_string())) + } + } } /// Pin the compilation target and enable the requested x86-64 psABI level. @@ -2451,30 +2698,73 @@ impl PythonExecutor { reuse_empty_callbacks: true, }; - // Create store for this execution - let mut store = Store::new(&self.engine, state); - - // Register the memory tracker as a resource limiter - store.limiter(|state| &mut state.memory_tracker); - - // Set a high epoch deadline for instantiation - we don't want to timeout during - // Python initialization, only during user code execution. - store.set_epoch_deadline(u64::MAX / 2); - // Set up fuel for tracking/limiting. We use u64::MAX for tracking-only mode // when no explicit limit is set. Fuel is consumed per WASM instruction. let initial_fuel = fuel_limit.unwrap_or(u64::MAX); - store - .set_fuel(initial_fuel) - .map_err(|e| Error::Initialization(format!("Failed to set fuel: {e}")))?; - // Instantiate from the pre-compiled template (includes Python initialization) - let bindings = self - .instance_pre - .instantiate_async(&mut store) - .await - .map_err(Error::WasmComponent)?; + let (mut store, bindings) = self + .acquire_store(state, memory_limit, initial_fuel) + .await?; + + let result = self + .run_on_store( + &mut store, + &bindings, + code, + execution_timeout, + cancellation_token, + fuel_limit, + initial_fuel, + ) + .await; + + // Teardown (resetting the memory slot) and re-instantiation happen off + // the request path when the warm pool is active. + WarmPool::global().retire(&self.instance_pre, store); + result + } + + /// Take a warm store for this component, or instantiate a fresh one. + /// + /// A warm store was instantiated with a placeholder state; the real + /// per-execution `state` replaces it here, before any export is called. + async fn acquire_store( + &self, + state: ExecutorState, + memory_limit: Option, + initial_fuel: u64, + ) -> std::result::Result<(Store, Sandbox), Error> { + if let Some(WarmInstance { + mut store, + bindings, + }) = WarmPool::global().take(&self.instance_pre, memory_limit) + { + // The placeholder's tracker saw the memory instantiation mapped; + // carry that into this execution's peak. + let baseline = store.data().memory_tracker.peak_memory_bytes(); + state.memory_tracker.observe(baseline); + *store.data_mut() = state; + store + .set_fuel(initial_fuel) + .map_err(|e| Error::Initialization(format!("Failed to set fuel: {e}")))?; + return Ok((store, bindings)); + } + instantiate_store(&self.instance_pre, state, initial_fuel).await + } + + /// Run `code` on an instantiated store. + #[allow(clippy::too_many_arguments)] + async fn run_on_store( + &self, + store: &mut Store, + bindings: &Sandbox, + code: &str, + execution_timeout: Option, + cancellation_token: Option, + fuel_limit: Option, + initial_fuel: u64, + ) -> std::result::Result { // Configure the guest's result-capture variable name. The guest defaults to // "result", so only call the export when a non-default name is configured. if self.result_variable != "result" { @@ -2527,7 +2817,7 @@ impl PythonExecutor { } }); } else { - arm_epoch_deadline(&mut store, execution_timeout); + arm_epoch_deadline(store, execution_timeout); } // Call the async execute export @@ -2825,6 +3115,69 @@ mod tests { } } + fn allocator_settings(vars: &[(&str, &str)]) -> std::result::Result { + AllocatorSettings::from_lookup(|name| { + vars.iter() + .find(|(k, _)| *k == name) + .map(|(_, v)| (*v).to_string()) + }) + } + + #[test] + fn allocator_defaults_to_pooling() { + let settings = allocator_settings(&[]).unwrap(); + assert_eq!(settings, AllocatorSettings::default()); + assert!(settings.pooling_config().is_some()); + } + + #[test] + fn allocator_can_be_switched_to_on_demand() { + let settings = allocator_settings(&[("ERYX_ALLOCATOR", "On-Demand")]).unwrap(); + assert!(settings.pooling_config().is_none()); + } + + #[test] + fn allocator_pool_limits_are_applied() { + let settings = allocator_settings(&[ + ("ERYX_POOL_INSTANCES", " 42 "), + ("ERYX_POOL_KEEP_RESIDENT_MB", "8"), + ("ERYX_POOL_PAGEMAP_SCAN", "0"), + ]) + .unwrap(); + let pooling = settings.pooling_config().unwrap(); + assert_eq!(pooling.get_total_memories(), 42); + assert_eq!(pooling.get_total_component_instances(), 42); + assert_eq!(pooling.get_total_stacks(), 42); + assert_eq!(pooling.get_memory_keep_resident(), 8 << 20); + assert!(matches!(pooling.get_pagemap_scan(), Enabled::No)); + } + + #[test] + fn allocator_rejects_invalid_values() { + for vars in [ + [("ERYX_ALLOCATOR", "mmap")], + [("ERYX_POOL_INSTANCES", "0")], + [("ERYX_POOL_INSTANCES", "lots")], + [("ERYX_POOL_KEEP_RESIDENT_MB", "-1")], + [("ERYX_POOL_PAGEMAP_SCAN", "yes")], + ] { + let err = allocator_settings(&vars).expect_err("invalid setting must be rejected"); + assert!(matches!(err, Error::WasmEngine(_)), "{vars:?}: {err:?}"); + } + } + + /// The pooling configuration must be accepted by wasmtime, and instantiation + /// must work under it: the pool imposes per-instance limits (table size, + /// instance metadata size) that the on-demand allocator does not. + #[test] + #[cfg(feature = "embedded")] + fn pooling_engine_is_constructible() { + let mut config = PythonExecutor::base_engine_config(); + let pooling = AllocatorSettings::default().pooling_config().unwrap(); + config.allocation_strategy(InstanceAllocationStrategy::Pooling(pooling)); + Engine::new(&config).expect("pooling configuration must be valid"); + } + /// Test that all CPU feature level presets use valid Cranelift flags. /// /// This test ensures that if someone adds a new flag name that doesn't exist diff --git a/crates/eryx/tests/warm_instances.rs b/crates/eryx/tests/warm_instances.rs new file mode 100644 index 00000000..81bef81a --- /dev/null +++ b/crates/eryx/tests/warm_instances.rs @@ -0,0 +1,226 @@ +//! Coverage for the warm-instance pool behind stateless execution. +//! +//! After a `Sandbox::execute()` on a multi-threaded runtime, a background task +//! pre-instantiates a replacement store so the next execution skips +//! instantiation. These tests check that the pool fills, that executions on a +//! pre-instantiated store see exactly the per-execution state they would on a +//! fresh one, and that the pool stays out of the way where it cannot help. +#![cfg(feature = "embedded")] +#![allow(clippy::expect_used, clippy::unwrap_used)] + +use std::future::Future; +use std::pin::Pin; +use std::time::Duration; + +use eryx::{CallbackError, JsonSchema, ResourceLimits, Sandbox, TypedCallback}; +use serde::Deserialize; +use serde_json::Value; + +/// Wait for the background task to top the pool up, returning the ready count. +async fn wait_for_warm(sandbox: &Sandbox) -> usize { + let executor = sandbox.executor(); + for _ in 0..500 { + let ready = executor.warm_instances_ready(); + if ready > 0 { + return ready; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + executor.warm_instances_ready() +} + +/// Build a sandbox and run one execution so that a warm instance is waiting. +async fn sandbox_with_warm_instance() -> Sandbox { + let sandbox = Sandbox::embedded().build().unwrap(); + sandbox.execute("pass").await.unwrap(); + assert_eq!(wait_for_warm(&sandbox).await, 1); + sandbox +} + +#[derive(Deserialize, JsonSchema)] +struct EchoArgs { + /// Data to echo back + data: Value, +} + +struct EchoCallback; + +impl TypedCallback for EchoCallback { + type Args = EchoArgs; + + fn name(&self) -> &str { + "echo" + } + + fn description(&self) -> &str { + "Echoes the input data back" + } + + fn invoke_typed( + &self, + args: EchoArgs, + ) -> Pin> + Send + '_>> { + Box::pin(async move { Ok(args.data) }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_replacement_is_prepared_after_each_execution() { + let sandbox = Sandbox::embedded().build().unwrap(); + assert_eq!(sandbox.executor().warm_instances_ready(), 0); + + let cold = sandbox.execute("print('cold')").await.unwrap(); + assert_eq!(cold.stdout, "cold"); + assert_eq!(wait_for_warm(&sandbox).await, 1); + + let warm = sandbox + .execute("import json\nprint(json.dumps({'warm': True}))") + .await + .unwrap(); + assert_eq!(warm.stdout, "{\"warm\": true}"); + assert_eq!(wait_for_warm(&sandbox).await, 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn warm_instances_are_shared_between_sandboxes_of_the_same_runtime() { + let _first = sandbox_with_warm_instance().await; + + // A brand-new sandbox (the SandboxFactory pattern) finds the instance + // the first one left behind. + let second = Sandbox::embedded().build().unwrap(); + assert_eq!(second.executor().warm_instances_ready(), 1); + let output = second.execute("print('shared')").await.unwrap(); + assert_eq!(output.stdout, "shared"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn executions_on_warm_instances_are_isolated() { + let sandbox = sandbox_with_warm_instance().await; + + sandbox + .execute("leaked = 'state'\nprint('one')") + .await + .unwrap(); + wait_for_warm(&sandbox).await; + + let output = sandbox + .execute("print('leaked' in globals())") + .await + .unwrap(); + assert_eq!(output.stdout, "False"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn per_execution_configuration_applies_to_a_warm_instance() { + // The instance waiting in the pool was instantiated by a sandbox with no + // callbacks and the default result variable. + let _plain = sandbox_with_warm_instance().await; + + let configured = Sandbox::embedded() + .with_callback(EchoCallback) + .with_result_variable("outcome") + .build() + .unwrap(); + assert_eq!(configured.executor().warm_instances_ready(), 1); + + let output = configured + .execute("outcome = await echo(data={'n': 1})\nprint(outcome)") + .await + .unwrap(); + assert_eq!(output.stdout, "{'n': 1}"); + assert_eq!(output.result.as_deref(), Some(r#"{"n": 1}"#)); +} + +#[cfg(feature = "vfs")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn the_virtual_filesystem_is_mounted_on_a_warm_instance() { + let sandbox = sandbox_with_warm_instance().await; + + let output = sandbox + .execute( + "with open('/data/note.txt', 'w') as f:\n f.write('hi')\nprint(open('/data/note.txt').read())", + ) + .await + .unwrap(); + assert_eq!(output.stdout, "hi"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_memory_limit_below_the_snapshot_baseline_is_still_rejected() { + let limits = ResourceLimits::unlimited().with_max_memory_bytes(1024 * 1024); + let limited = Sandbox::embedded() + .with_resource_limits(limits) + .build() + .unwrap(); + + let cold_error = limited + .execute("pass") + .await + .expect_err("1 MiB cannot hold the snapshot"); + assert!( + cold_error.to_string().to_lowercase().contains("memory"), + "unexpected error: {cold_error:?}" + ); + + // With an instance ready, the limited execution must not silently run on + // it: the pool leaves it alone and the same error comes back. + let _plain = sandbox_with_warm_instance().await; + let warm_error = limited + .execute("pass") + .await + .expect_err("1 MiB cannot hold the snapshot"); + assert!( + warm_error.to_string().to_lowercase().contains("memory"), + "unexpected error: {warm_error:?}" + ); + assert_eq!(limited.executor().warm_instances_ready(), 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fuel_limits_apply_to_a_warm_instance() { + let _plain = sandbox_with_warm_instance().await; + + let limits = ResourceLimits::unlimited().with_max_fuel(10_000); + let limited = Sandbox::embedded() + .with_resource_limits(limits) + .build() + .unwrap(); + assert_eq!(limited.executor().warm_instances_ready(), 1); + + let error = limited + .execute("total = sum(range(1_000_000))") + .await + .expect_err("the loop needs far more than 10k fuel"); + assert!( + matches!(error, eryx::Error::FuelExhausted { .. }), + "unexpected error: {error:?}" + ); +} + +#[tokio::test] +async fn the_pool_is_inactive_on_a_current_thread_runtime() { + let sandbox = Sandbox::embedded().build().unwrap(); + sandbox.execute("pass").await.unwrap(); + sandbox.execute("pass").await.unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!(sandbox.executor().warm_instances_ready(), 0); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fuel_consumed_is_the_same_on_cold_and_warm_instances() { + let sandbox = Sandbox::embedded().build().unwrap(); + let code = "x = sum(range(100))"; + + // The first execution instantiates, the rest take a warm instance; + // instantiation must not be charged to either. + let mut fuel = Vec::new(); + for _ in 0..3 { + fuel.push(sandbox.execute(code).await.unwrap().stats.fuel_consumed); + wait_for_warm(&sandbox).await; + } + assert!(fuel[0].is_some()); + assert!( + fuel.iter().all(|f| *f == fuel[0]), + "fuel should not depend on the instance's origin: {fuel:?}" + ); +} diff --git a/mise.toml b/mise.toml index 5b341218..05d45d02 100644 --- a/mise.toml +++ b/mise.toml @@ -1,6 +1,6 @@ [tools] # Profile "default" includes rustc, rust-std, cargo, rust-docs, rustfmt, and clippy -rust = { version = "1.98", profile = "default" } +rust = { version = "1.98.1", profile = "default" } [tools."cargo:cargo-nextest"] version = "0.9.143"