diff --git a/src/component.rs b/src/component.rs index 3f3fc28..832b937 100644 --- a/src/component.rs +++ b/src/component.rs @@ -70,6 +70,7 @@ pub(crate) async fn wizen( config.wasm_bulk_memory(true); config.wasm_component_model(true); config.wasm_component_model_async(true); + crate::memory::configure(&mut config); let engine = Engine::new(&config)?; diff --git a/src/main.rs b/src/main.rs index 3b507e8..f1c66ca 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,6 +17,7 @@ mod filter; mod image; mod intrinsics; mod liveness; +mod memory; mod module; mod state; mod stats; @@ -62,6 +63,13 @@ pub enum Command { #[structopt(long = "allow-wasi-http")] allow_wasi_http: bool, + /// Virtual memory to reserve for each linear memory while wizening, + /// e.g. `64M` or `1G`. Wasmtime reserves 4GiB by default. Can also be + /// set with the WEVAL_MEMORY_RESERVATION environment variable; the flag + /// takes precedence. + #[structopt(long = "memory-reservation")] + memory_reservation: Option, + /// Cache file to use. #[structopt(long = "cache")] cache: Option, @@ -99,29 +107,33 @@ fn main() -> anyhow::Result<()> { allow_wasip2, allow_wasip3, allow_wasi_http, + memory_reservation, cache, cache_ro, show_stats, output_ir, verbose, - } => weval( - input_module, - output_module, - wizen, - preopens, - init_func, - keep_init_func, - cache, - cache_ro, - show_stats, - output_ir, - verbose, - WasiOpts { - p2: allow_wasip2, - p3: allow_wasip3, - http: allow_wasi_http, - }, - ), + } => { + memory::init(memory_reservation)?; + weval( + input_module, + output_module, + wizen, + preopens, + init_func, + keep_init_func, + cache, + cache_ro, + show_stats, + output_ir, + verbose, + WasiOpts { + p2: allow_wasip2, + p3: allow_wasip3, + http: allow_wasi_http, + }, + ) + } } } diff --git a/src/memory.rs b/src/memory.rs new file mode 100644 index 0000000..91434ed --- /dev/null +++ b/src/memory.rs @@ -0,0 +1,78 @@ +//! Linear memory configuration for the engine used while wizening. +//! +//! By default Wasmtime reserves 4GiB of virtual address space for every linear +//! memory on 64-bit hosts. This is primarily useful for configuring the virtual +//! address space when running in a resource constrained environment. The +//! reservation can be overridden with `--memory-reservation` or the +//! `WEVAL_MEMORY_RESERVATION` environment variable. + +use std::sync::OnceLock; + +use anyhow::{bail, Context}; +use wasmtime::Config; + +/// Environment variable consulted when `--memory-reservation` is not given. +pub(crate) const ENV_VAR: &str = "WEVAL_MEMORY_RESERVATION"; + +static RESERVATION: OnceLock> = OnceLock::new(); + +/// Parse a size such as `64M`, `64MiB`, `1G`, `4096` (bytes) into bytes. +pub(crate) fn parse_size(input: &str) -> anyhow::Result { + let s = input.trim().to_ascii_lowercase(); + let s = s.strip_suffix('b').unwrap_or(&s); + let s = s.strip_suffix('i').unwrap_or(s); + let (digits, shift) = match s.chars().last() { + Some('k') => (&s[..s.len() - 1], 10), + Some('m') => (&s[..s.len() - 1], 20), + Some('g') => (&s[..s.len() - 1], 30), + _ => (s, 0), + }; + let value: u64 = digits + .trim() + .parse() + .with_context(|| format!("invalid memory size: {input:?}"))?; + match value.checked_mul(1u64 << shift) { + Some(bytes) => Ok(bytes), + None => bail!("memory size overflows u64: {input:?}"), + } +} + +/// Record the reservation to use. The flag value wins over the environment. +pub(crate) fn init(flag: Option) -> anyhow::Result<()> { + let value = match flag { + Some(s) => Some(parse_size(&s)?), + None => match std::env::var(ENV_VAR) { + Ok(s) if !s.trim().is_empty() => { + Some(parse_size(&s).with_context(|| format!("while reading {ENV_VAR}"))?) + } + _ => None, + }, + }; + let _ = RESERVATION.set(value); + Ok(()) +} + +/// Apply the configured reservation, if any, to a Wasmtime `Config`. +pub(crate) fn configure(config: &mut Config) { + if let Some(Some(bytes)) = RESERVATION.get() { + config.memory_reservation(*bytes); + config.memory_reservation_for_growth(0); + config.memory_guard_size(0); + config.memory_may_move(true); + } +} + +#[cfg(test)] +mod tests { + use super::parse_size; + + #[test] + fn parses_sizes() { + assert_eq!(parse_size("4096").unwrap(), 4096); + assert_eq!(parse_size("64M").unwrap(), 64 << 20); + assert_eq!(parse_size("64MiB").unwrap(), 64 << 20); + assert_eq!(parse_size("1g").unwrap(), 1 << 30); + assert_eq!(parse_size("512k").unwrap(), 512 << 10); + assert!(parse_size("abc").is_err()); + } +} diff --git a/src/module.rs b/src/module.rs index b7470f0..7649278 100644 --- a/src/module.rs +++ b/src/module.rs @@ -24,6 +24,7 @@ pub(crate) async fn wizen( ) -> anyhow::Result> { let mut config = Config::new(); config.wasm_bulk_memory(true); + crate::memory::configure(&mut config); let engine = Engine::new(&config)?; let mut wasi_ctx = wasmtime_wasi::WasiCtxBuilder::new();