diff --git a/src/commands/run.rs b/src/commands/run.rs index 2af92f6ddf24..b45f48a5ae6b 100644 --- a/src/commands/run.rs +++ b/src/commands/run.rs @@ -5,27 +5,23 @@ allow(irrefutable_let_patterns, unreachable_patterns) )] +use crate::common::{Profile, RunCommon, RunTarget}; + use anyhow::{anyhow, bail, Context as _, Error, Result}; use clap::Parser; use std::fs::File; use std::io::Write; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::sync::Arc; use std::thread; -use std::time::Duration; use wasmtime::{ - AsContextMut, Engine, Func, GuestProfiler, Module, Precompiled, Store, StoreLimits, - StoreLimitsBuilder, UpdateDeadline, Val, ValType, + AsContextMut, Engine, Func, GuestProfiler, Module, Store, StoreLimits, UpdateDeadline, Val, + ValType, }; -use wasmtime_cli_flags::opt::WasmtimeOptionValue; -use wasmtime_cli_flags::CommonOptions; use wasmtime_wasi::maybe_exit_on_error; use wasmtime_wasi::preview2; use wasmtime_wasi::sync::{ambient_authority, Dir, TcpListener, WasiCtxBuilder}; -#[cfg(feature = "component-model")] -use wasmtime::component::Component; - #[cfg(feature = "wasi-nn")] use wasmtime_wasi_nn::WasiNnCtx; @@ -61,43 +57,12 @@ fn parse_preloads(s: &str) -> Result<(String, PathBuf)> { Ok((parts[0].into(), parts[1].into())) } -fn parse_profile(s: &str) -> Result { - let parts = s.split(',').collect::>(); - match &parts[..] { - ["perfmap"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::PerfMap)), - ["jitdump"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::JitDump)), - ["vtune"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::VTune)), - ["guest"] => Ok(Profile::Guest { - path: "wasmtime-guest-profile.json".to_string(), - interval: Duration::from_millis(10), - }), - ["guest", path] => Ok(Profile::Guest { - path: path.to_string(), - interval: Duration::from_millis(10), - }), - ["guest", path, dur] => Ok(Profile::Guest { - path: path.to_string(), - interval: WasmtimeOptionValue::parse(Some(dur))?, - }), - _ => bail!("unknown profiling strategy: {s}"), - } -} - /// Runs a WebAssembly module #[derive(Parser)] #[structopt(name = "run")] pub struct RunCommand { #[clap(flatten)] - common: CommonOptions, - - /// Allow executing precompiled WebAssembly modules as `*.cwasm` files. - /// - /// Note that this option is not safe to pass if the module being passed in - /// is arbitrary user input. Only `wasmtime`-precompiled modules generated - /// via the `wasmtime compile` command or equivalent should be passed as an - /// argument with this option specified. - #[clap(long = "allow-precompiled")] - allow_precompiled: bool, + run: RunCommon, /// Grant access of a host directory to a guest. /// @@ -131,28 +96,6 @@ pub struct RunCommand { )] preloads: Vec<(String, PathBuf)>, - /// Profiling strategy (valid options are: perfmap, jitdump, vtune, guest) - /// - /// The perfmap, jitdump, and vtune profiling strategies integrate Wasmtime - /// with external profilers such as `perf`. The guest profiling strategy - /// enables in-process sampling and will write the captured profile to - /// `wasmtime-guest-profile.json` by default which can be viewed at - /// https://profiler.firefox.com/. - /// - /// The `guest` option can be additionally configured as: - /// - /// --profile=guest[,path[,interval]] - /// - /// where `path` is where to write the profile and `interval` is the - /// duration between samples. When used with `--wasm-timeout` the timeout - /// will be rounded up to the nearest multiple of this interval. - #[clap( - long, - value_name = "STRATEGY", - value_parser = parse_profile, - )] - profile: Option, - /// The WebAssembly module to run and arguments to pass to it. /// /// Arguments passed to the wasm module will be configured as WASI CLI @@ -162,53 +105,23 @@ pub struct RunCommand { module_and_args: Vec, } -#[derive(Clone)] -enum Profile { - Native(wasmtime::ProfilingStrategy), - Guest { path: String, interval: Duration }, -} - enum CliLinker { Core(wasmtime::Linker), #[cfg(feature = "component-model")] Component(wasmtime::component::Linker), } -enum CliModule { - Core(wasmtime::Module), - #[cfg(feature = "component-model")] - Component(Component), -} - -impl CliModule { - fn unwrap_core(&self) -> &Module { - match self { - CliModule::Core(module) => module, - #[cfg(feature = "component-model")] - CliModule::Component(_) => panic!("expected a core wasm module, not a component"), - } - } - - #[cfg(feature = "component-model")] - fn unwrap_component(&self) -> &Component { - match self { - CliModule::Component(c) => c, - CliModule::Core(_) => panic!("expected a component, not a core wasm module"), - } - } -} - impl RunCommand { /// Executes the command. pub fn execute(mut self) -> Result<()> { - self.common.init_logging(); + self.run.common.init_logging(); - let mut config = self.common.config(None)?; + let mut config = self.run.common.config(None)?; - if self.common.wasm.timeout.is_some() { + if self.run.common.wasm.timeout.is_some() { config.epoch_interruption(true); } - match self.profile { + match self.run.profile { Some(Profile::Native(s)) => { config.profiler(s); } @@ -222,23 +135,23 @@ impl RunCommand { let engine = Engine::new(&config)?; // Read the wasm module binary either as `*.wat` or a raw binary. - let main = self.load_module(&engine, &self.module_and_args[0])?; + let main = self.run.load_module(&engine, &self.module_and_args[0])?; // Validate coredump-on-trap argument - if let Some(path) = &self.common.debug.coredump { + if let Some(path) = &self.run.common.debug.coredump { if path.contains("%") { bail!("the coredump-on-trap path does not support patterns yet.") } } let mut linker = match &main { - CliModule::Core(_) => CliLinker::Core(wasmtime::Linker::new(&engine)), + RunTarget::Core(_) => CliLinker::Core(wasmtime::Linker::new(&engine)), #[cfg(feature = "component-model")] - CliModule::Component(_) => { + RunTarget::Component(_) => { CliLinker::Component(wasmtime::component::Linker::new(&engine)) } }; - if let Some(enable) = self.common.wasm.unknown_exports_allow { + if let Some(enable) = self.run.common.wasm.unknown_exports_allow { match &mut linker { CliLinker::Core(l) => { l.allow_unknown_exports(enable); @@ -254,45 +167,26 @@ impl RunCommand { let mut store = Store::new(&engine, host); self.populate_with_wasi(&mut linker, &mut store, &main)?; - let mut limits = StoreLimitsBuilder::new(); - if let Some(max) = self.common.wasm.max_memory_size { - limits = limits.memory_size(max); - } - if let Some(max) = self.common.wasm.max_table_elements { - limits = limits.table_elements(max); - } - if let Some(max) = self.common.wasm.max_instances { - limits = limits.instances(max); - } - if let Some(max) = self.common.wasm.max_tables { - limits = limits.tables(max); - } - if let Some(max) = self.common.wasm.max_memories { - limits = limits.memories(max); - } - if let Some(enable) = self.common.wasm.trap_on_grow_failure { - limits = limits.trap_on_grow_failure(enable); - } - store.data_mut().limits = limits.build(); + store.data_mut().limits = self.run.store_limits(); store.limiter(|t| &mut t.limits); // If fuel has been configured, we want to add the configured // fuel amount to this store. - if let Some(fuel) = self.common.wasm.fuel { + if let Some(fuel) = self.run.common.wasm.fuel { store.add_fuel(fuel)?; } // Load the preload wasm modules. let mut modules = Vec::new(); - if let CliModule::Core(m) = &main { + if let RunTarget::Core(m) = &main { modules.push((String::new(), m.clone())); } for (name, path) in self.preloads.iter() { // Read the wasm module binary either as `*.wat` or a raw binary - let module = match self.load_module(&engine, path)? { - CliModule::Core(m) => m, + let module = match self.run.load_module(&engine, path)? { + RunTarget::Core(m) => m, #[cfg(feature = "component-model")] - CliModule::Component(_) => bail!("components cannot be loaded with `--preload`"), + RunTarget::Component(_) => bail!("components cannot be loaded with `--preload`"), }; modules.push((name.clone(), module.clone())); @@ -350,7 +244,7 @@ impl RunCommand { fn compute_preopen_sockets(&self) -> Result> { let mut listeners = vec![]; - for address in &self.common.wasi.tcplisten { + for address in &self.run.common.wasi.tcplisten { let stdlistener = std::net::TcpListener::bind(address) .with_context(|| format!("failed to bind to address '{}'", address))?; @@ -387,7 +281,7 @@ impl RunCommand { store: &mut Store, modules: Vec<(String, Module)>, ) -> Box)> { - if let Some(Profile::Guest { path, interval }) = &self.profile { + if let Some(Profile::Guest { path, interval }) = &self.run.profile { let module_name = self.module_and_args[0].to_str().unwrap_or("
"); let interval = *interval; store.data_mut().guest_profiler = @@ -406,7 +300,7 @@ impl RunCommand { store.as_context_mut().data_mut().guest_profiler = Some(profiler); } - if let Some(timeout) = self.common.wasm.timeout { + if let Some(timeout) = self.run.common.wasm.timeout { let mut timeout = (timeout.as_secs_f64() / interval.as_secs_f64()).ceil() as u64; assert!(timeout > 0); store.epoch_deadline_callback(move |mut store| { @@ -448,7 +342,7 @@ impl RunCommand { }); } - if let Some(timeout) = self.common.wasm.timeout { + if let Some(timeout) = self.run.common.wasm.timeout { store.set_epoch_deadline(1); let engine = store.engine().clone(); thread::spawn(move || { @@ -464,12 +358,12 @@ impl RunCommand { &self, store: &mut Store, linker: &mut CliLinker, - module: &CliModule, + module: &RunTarget, modules: Vec<(String, Module)>, ) -> Result<()> { // The main module might be allowed to have unknown imports, which // should be defined as traps: - if self.common.wasm.unknown_imports_trap == Some(true) { + if self.run.common.wasm.unknown_imports_trap == Some(true) { match linker { CliLinker::Core(linker) => { linker.define_unknown_imports_as_traps(module.unwrap_core())?; @@ -479,7 +373,7 @@ impl RunCommand { } // ...or as default values. - if self.common.wasm.unknown_imports_default == Some(true) { + if self.run.common.wasm.unknown_imports_default == Some(true) { match linker { CliLinker::Core(linker) => { linker.define_unknown_imports_as_default_values(module.unwrap_core())?; @@ -620,7 +514,7 @@ impl RunCommand { } fn handle_core_dump(&self, store: &mut Store, err: Error) -> Error { - let coredump_path = match &self.common.debug.coredump { + let coredump_path = match &self.run.common.debug.coredump { Some(path) => path, None => return err, }; @@ -639,130 +533,17 @@ impl RunCommand { } } - fn load_module(&self, engine: &Engine, path: &Path) -> Result { - let path = match path.to_str() { - #[cfg(unix)] - Some("-") => "/dev/stdin".as_ref(), - _ => path, - }; - - // First attempt to load the module as an mmap. If this succeeds then - // detection can be done with the contents of the mmap and if a - // precompiled module is detected then `deserialize_file` can be used - // which is a slightly more optimal version than `deserialize` since we - // can leave most of the bytes on disk until they're referenced. - // - // If the mmap fails, for example if stdin is a pipe, then fall back to - // `std::fs::read` to load the contents. At that point precompiled - // modules must go through the `deserialize` functions. - // - // Note that this has the unfortunate side effect for precompiled - // modules on disk that they're opened once to detect what they are and - // then again internally in Wasmtime as part of the `deserialize_file` - // API. Currently there's no way to pass the `MmapVec` here through to - // Wasmtime itself (that'd require making `wasmtime-runtime` a public - // dependency or `MmapVec` a public type, both of which aren't ready to - // happen at this time). It's hoped though that opening a file twice - // isn't too bad in the grand scheme of things with respect to the CLI. - match wasmtime_runtime::MmapVec::from_file(path) { - Ok(map) => self.load_module_contents( - engine, - path, - &map, - || unsafe { Module::deserialize_file(engine, path) }, - #[cfg(feature = "component-model")] - || unsafe { Component::deserialize_file(engine, path) }, - ), - Err(_) => { - let bytes = std::fs::read(path) - .with_context(|| format!("failed to read file: {}", path.display()))?; - self.load_module_contents( - engine, - path, - &bytes, - || unsafe { Module::deserialize(engine, &bytes) }, - #[cfg(feature = "component-model")] - || unsafe { Component::deserialize(engine, &bytes) }, - ) - } - } - } - - fn load_module_contents( - &self, - engine: &Engine, - path: &Path, - bytes: &[u8], - deserialize_module: impl FnOnce() -> Result, - #[cfg(feature = "component-model")] deserialize_component: impl FnOnce() -> Result, - ) -> Result { - Ok(match engine.detect_precompiled(bytes) { - Some(Precompiled::Module) => { - self.ensure_allow_precompiled()?; - CliModule::Core(deserialize_module()?) - } - #[cfg(feature = "component-model")] - Some(Precompiled::Component) => { - self.ensure_allow_precompiled()?; - self.ensure_allow_components()?; - CliModule::Component(deserialize_component()?) - } - #[cfg(not(feature = "component-model"))] - Some(Precompiled::Component) => { - bail!("support for components was not enabled at compile time"); - } - None => { - // Parse the text format here specifically to add the `path` to - // the error message if there's a syntax error. - let wasm = wat::parse_bytes(bytes).map_err(|mut e| { - e.set_path(path); - e - })?; - if wasmparser::Parser::is_component(&wasm) { - #[cfg(feature = "component-model")] - { - self.ensure_allow_components()?; - CliModule::Component(Component::new(engine, &wasm)?) - } - #[cfg(not(feature = "component-model"))] - { - bail!("support for components was not enabled at compile time"); - } - } else { - CliModule::Core(Module::new(engine, &wasm)?) - } - } - }) - } - - fn ensure_allow_precompiled(&self) -> Result<()> { - if self.allow_precompiled { - Ok(()) - } else { - bail!("running a precompiled module requires the `--allow-precompiled` flag") - } - } - - #[cfg(feature = "component-model")] - fn ensure_allow_components(&self) -> Result<()> { - if self.common.wasm.component_model != Some(true) { - bail!("cannot execute a component without `--wasm component-model`"); - } - - Ok(()) - } - /// Populates the given `Linker` with WASI APIs. fn populate_with_wasi( &self, linker: &mut CliLinker, store: &mut Store, - module: &CliModule, + module: &RunTarget, ) -> Result<()> { - if self.common.wasi.common != Some(false) { + if self.run.common.wasi.common != Some(false) { match linker { CliLinker::Core(linker) => { - if self.common.wasi.preview2 == Some(true) { + if self.run.common.wasi.preview2 == Some(true) { preview2::preview1::add_to_linker_sync(linker)?; self.set_preview2_ctx(store)?; } else { @@ -780,7 +561,7 @@ impl RunCommand { } } - if self.common.wasi.nn == Some(true) { + if self.run.common.wasi.nn == Some(true) { #[cfg(not(feature = "wasi-nn"))] { bail!("Cannot enable wasi-nn when the binary is not compiled with this feature."); @@ -810,6 +591,7 @@ impl RunCommand { } } let graphs = self + .run .common .wasi .nn_graph @@ -821,7 +603,7 @@ impl RunCommand { } } - if self.common.wasi.threads == Some(true) { + if self.run.common.wasi.threads == Some(true) { #[cfg(not(feature = "wasi-threads"))] { // Silence the unused warning for `module` as it is only used in the @@ -849,7 +631,7 @@ impl RunCommand { } } - if self.common.wasi.http == Some(true) { + if self.run.common.wasi.http == Some(true) { #[cfg(not(all(feature = "wasi-http", feature = "component-model")))] { bail!("Cannot enable wasi-http when the binary is not compiled with this feature."); @@ -887,7 +669,7 @@ impl RunCommand { let mut num_fd: usize = 3; - if self.common.wasi.listenfd == Some(true) { + if self.run.common.wasi.listenfd == Some(true) { num_fd = ctx_set_listenfd(num_fd, &mut builder)?; } @@ -917,7 +699,7 @@ impl RunCommand { builder.env(key, &value); } - if self.common.wasi.listenfd == Some(true) { + if self.run.common.wasi.listenfd == Some(true) { bail!("components do not support --listenfd"); } for _ in self.compute_preopen_sockets()? { @@ -933,10 +715,10 @@ impl RunCommand { ); } - if self.common.wasi.inherit_network == Some(true) { + if self.run.common.wasi.inherit_network == Some(true) { builder.inherit_network(ambient_authority()); } - if let Some(enable) = self.common.wasi.allow_ip_name_lookup { + if let Some(enable) = self.run.common.wasi.allow_ip_name_lookup { builder.allow_ip_name_lookup(enable); } diff --git a/src/commands/serve.rs b/src/commands/serve.rs index bfca73c92cd2..6993217839a7 100644 --- a/src/commands/serve.rs +++ b/src/commands/serve.rs @@ -1,26 +1,31 @@ -use anyhow::Result; +use crate::common::{Profile, RunCommon, RunTarget}; +use anyhow::{bail, Result}; use clap::Parser; -use std::{path::PathBuf, pin::Pin, sync::Arc}; -use wasmtime::component::{Component, InstancePre, Linker}; -use wasmtime::{Engine, Store}; -use wasmtime_cli_flags::CommonOptions; +use std::{ + path::PathBuf, + pin::Pin, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, +}; +use wasmtime::component::{InstancePre, Linker}; +use wasmtime::{Engine, Store, StoreLimits}; use wasmtime_wasi::preview2::{Table, WasiCtx, WasiCtxBuilder, WasiView}; use wasmtime_wasi_http::{body::HyperOutgoingBody, WasiHttpCtx, WasiHttpView}; +#[cfg(feature = "wasi-nn")] +use wasmtime_wasi_nn::WasiNnCtx; + struct Host { table: Table, ctx: WasiCtx, http: WasiHttpCtx, -} -impl Host { - fn new() -> Self { - Host { - table: Table::new(), - ctx: WasiCtxBuilder::new().build(), - http: WasiHttpCtx, - } - } + limits: StoreLimits, + + #[cfg(feature = "wasi-nn")] + nn: Option, } impl WasiView for Host { @@ -51,16 +56,21 @@ impl WasiHttpView for Host { } } +const DEFAULT_ADDR: std::net::SocketAddr = std::net::SocketAddr::new( + std::net::IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)), + 8080, +); + /// Runs a WebAssembly module #[derive(Parser)] -#[structopt(name = "run")] +#[structopt(name = "serve")] pub struct ServeCommand { #[clap(flatten)] - common: CommonOptions, + run: RunCommon, - /// Socket address for the web server to bind to. Defaults to 0.0.0.0:8080. - #[clap(long = "addr", value_name = "SOCKADDR")] - addr: Option, + /// Socket address for the web server to bind to. + #[clap(long = "addr", value_name = "SOCKADDR", default_value_t = DEFAULT_ADDR )] + addr: std::net::SocketAddr, /// The WebAssembly component to run. #[clap(value_name = "WASM", required = true)] @@ -68,12 +78,43 @@ pub struct ServeCommand { } impl ServeCommand { - fn addr(&self) -> std::net::SocketAddr { - self.addr.unwrap_or("0.0.0.0:8080".parse().unwrap()) - } - /// Start a server to run the given wasi-http proxy component pub fn execute(mut self) -> Result<()> { + self.run.common.init_logging(); + + // We force cli errors before starting to listen for connections so tha we don't + // accidentally delay them to the first request. + if self.run.common.wasi.nn == Some(true) { + #[cfg(not(feature = "wasi-nn"))] + { + bail!("Cannot enable wasi-nn when the binary is not compiled with this feature."); + } + } + + if let Some(Profile::Guest { .. }) = &self.run.profile { + bail!("Cannot use the guest profiler with components"); + } + + if self.run.common.wasi.nn == Some(true) { + #[cfg(not(feature = "wasi-nn"))] + { + bail!("Cannot enable wasi-nn when the binary is not compiled with this feature."); + } + } + + if self.run.common.wasi.threads == Some(true) { + bail!("wasi-threads does not support components yet") + } + + // The serve command requires both wasi-http and the component model, so we enable those by + // default here. + if self.run.common.wasi.http.replace(true) == Some(false) { + bail!("wasi-http is required for the serve command, and must not be disabled"); + } + if self.run.common.wasm.component_model.replace(true) == Some(false) { + bail!("components are required for the serve command, and must not be disabled"); + } + let runtime = tokio::runtime::Builder::new_multi_thread() .enable_time() .enable_io() @@ -94,38 +135,124 @@ impl ServeCommand { Ok(()) } + fn new_store(&self, engine: &Engine) -> Result> { + let mut builder = WasiCtxBuilder::new(); + + // TODO: connect stdio to logging infrastructure + + let mut host = Host { + table: Table::new(), + ctx: builder.build(), + http: WasiHttpCtx, + + limits: StoreLimits::default(), + + #[cfg(feature = "wasi-nn")] + nn: None, + }; + + if self.run.common.wasi.nn == Some(true) { + #[cfg(feature = "wasi-nn")] + { + let graphs = self + .run + .common + .wasi + .nn_graph + .iter() + .map(|g| (g.format.clone(), g.dir.clone())) + .collect::>(); + let (backends, registry) = wasmtime_wasi_nn::preload(&graphs)?; + host.nn.replace(WasiNnCtx::new(backends, registry)); + } + } + + let mut store = Store::new(engine, host); + + if self.run.common.wasm.timeout.is_some() { + store.set_epoch_deadline(1); + } + + store.data_mut().limits = self.run.store_limits(); + store.limiter(|t| &mut t.limits); + + // If fuel has been configured, we want to add the configured + // fuel amount to this store. + if let Some(fuel) = self.run.common.wasm.fuel { + store.add_fuel(fuel)?; + } + + Ok(store) + } + fn add_to_linker(&self, linker: &mut Linker) -> Result<()> { + // wasi-http and the component model are implicitly enabled for `wasmtime serve`, so we + // don't test for `self.run.common.wasi.common` or `self.run.common.wasi.http` in this + // function. + wasmtime_wasi_http::proxy::add_to_linker(linker)?; + + if self.run.common.wasi.nn == Some(true) { + #[cfg(feature = "wasi-nn")] + { + wasmtime_wasi_nn::wit::ML::add_to_linker(linker, |host| host.nn.as_mut().unwrap())?; + } + } + Ok(()) } - async fn serve(&mut self) -> Result<()> { + async fn serve(mut self) -> Result<()> { use hyper::server::conn::http1; - let mut config = self.common.config(None)?; + let mut config = self.run.common.config(None)?; config.wasm_component_model(true); config.async_support(true); - let engine = Arc::new(Engine::new(&config)?); + if self.run.common.wasm.timeout.is_some() { + config.epoch_interruption(true); + } + + match self.run.profile { + Some(Profile::Native(s)) => { + config.profiler(s); + } + + // We bail early in `execute` if the guest profiler is configured. + Some(Profile::Guest { .. }) => unreachable!(), + + None => {} + } + + let engine = Engine::new(&config)?; let mut linker = Linker::new(&engine); self.add_to_linker(&mut linker)?; - let component = Component::from_file(&engine, &self.component)?; + let component = match self.run.load_module(&engine, &self.component)? { + RunTarget::Core(_) => bail!("The serve command currently requires a component"), + RunTarget::Component(c) => c, + }; + + let instance = linker.instantiate_pre(&component)?; - let instance = Arc::new(linker.instantiate_pre(&component)?); + let listener = tokio::net::TcpListener::bind(self.addr).await?; - let listener = tokio::net::TcpListener::bind(self.addr()).await?; + let _epoch_thread = if let Some(timeout) = self.run.common.wasm.timeout { + Some(EpochThread::spawn(timeout, engine.clone())) + } else { + None + }; + + let handler = ProxyHandler::new(self, engine, instance); loop { let (stream, _) = listener.accept().await?; - let engine = Arc::clone(&engine); - let instance = Arc::clone(&instance); + let h = handler.clone(); tokio::task::spawn(async move { - let handler = ProxyHandler::new(engine, instance); if let Err(e) = http1::Builder::new() .keep_alive(true) - .serve_connection(stream, handler) + .serve_connection(stream, h) .await { eprintln!("error: {e:?}"); @@ -135,18 +262,54 @@ impl ServeCommand { } } -#[derive(Clone)] -struct ProxyHandler { - engine: Arc, - instance_pre: Arc>, +struct EpochThread { + shutdown: Arc, + handle: Option>, } +impl EpochThread { + fn spawn(timeout: std::time::Duration, engine: Engine) -> Self { + let shutdown = Arc::new(AtomicBool::new(false)); + let handle = { + let shutdown = Arc::clone(&shutdown); + let handle = std::thread::spawn(move || { + while !shutdown.load(Ordering::Relaxed) { + std::thread::sleep(timeout); + engine.increment_epoch(); + } + }); + Some(handle) + }; + + EpochThread { shutdown, handle } + } +} + +impl Drop for EpochThread { + fn drop(&mut self) { + if let Some(handle) = self.handle.take() { + self.shutdown.store(true, Ordering::Relaxed); + handle.join().unwrap(); + } + } +} + +struct ProxyHandlerInner { + cmd: ServeCommand, + engine: Engine, + instance_pre: InstancePre, +} + +#[derive(Clone)] +struct ProxyHandler(Arc); + impl ProxyHandler { - fn new(engine: Arc, instance_pre: Arc>) -> Self { - Self { + fn new(cmd: ServeCommand, engine: Engine, instance_pre: InstancePre) -> Self { + Self(Arc::new(ProxyHandlerInner { + cmd, engine, instance_pre, - } + })) } } @@ -166,7 +329,7 @@ impl hyper::service::Service for ProxyHandler { // TODO: need to track the join handle, but don't want to block the response on it tokio::task::spawn(async move { - let mut store = Store::new(&handler.engine, Host::new()); + let mut store = handler.0.cmd.new_store(&handler.0.engine)?; let req = store.data_mut().new_incoming_request( req.map(|body| body.map_err(|e| anyhow::anyhow!(e)).boxed()), @@ -176,7 +339,7 @@ impl hyper::service::Service for ProxyHandler { let (proxy, _inst) = wasmtime_wasi_http::proxy::Proxy::instantiate_pre( &mut store, - &handler.instance_pre, + &handler.0.instance_pre, ) .await?; diff --git a/src/common.rs b/src/common.rs new file mode 100644 index 000000000000..92708335d1b6 --- /dev/null +++ b/src/common.rs @@ -0,0 +1,243 @@ +//! Common functionality shared between command implementations. + +use anyhow::{bail, Context, Result}; +use clap::Parser; +use std::{path::Path, time::Duration}; +use wasmtime::{Engine, Module, Precompiled, StoreLimits, StoreLimitsBuilder}; +use wasmtime_cli_flags::{opt::WasmtimeOptionValue, CommonOptions}; + +#[cfg(feature = "component-model")] +use wasmtime::component::Component; + +pub enum RunTarget { + Core(Module), + + #[cfg(feature = "component-model")] + Component(Component), +} + +impl RunTarget { + pub fn unwrap_core(&self) -> &Module { + match self { + RunTarget::Core(module) => module, + #[cfg(feature = "component-model")] + RunTarget::Component(_) => panic!("expected a core wasm module, not a component"), + } + } + + #[cfg(feature = "component-model")] + pub fn unwrap_component(&self) -> &Component { + match self { + RunTarget::Component(c) => c, + RunTarget::Core(_) => panic!("expected a component, not a core wasm module"), + } + } +} + +/// Common command line arguments for run commands. +#[derive(Parser)] +pub struct RunCommon { + #[clap(flatten)] + pub common: CommonOptions, + + /// Allow executing precompiled WebAssembly modules as `*.cwasm` files. + /// + /// Note that this option is not safe to pass if the module being passed in + /// is arbitrary user input. Only `wasmtime`-precompiled modules generated + /// via the `wasmtime compile` command or equivalent should be passed as an + /// argument with this option specified. + #[clap(long = "allow-precompiled")] + pub allow_precompiled: bool, + + /// Profiling strategy (valid options are: perfmap, jitdump, vtune, guest) + /// + /// The perfmap, jitdump, and vtune profiling strategies integrate Wasmtime + /// with external profilers such as `perf`. The guest profiling strategy + /// enables in-process sampling and will write the captured profile to + /// `wasmtime-guest-profile.json` by default which can be viewed at + /// https://profiler.firefox.com/. + /// + /// The `guest` option can be additionally configured as: + /// + /// --profile=guest[,path[,interval]] + /// + /// where `path` is where to write the profile and `interval` is the + /// duration between samples. When used with `--wasm-timeout` the timeout + /// will be rounded up to the nearest multiple of this interval. + #[clap( + long, + value_name = "STRATEGY", + value_parser = Profile::parse, + )] + pub profile: Option, +} + +impl RunCommon { + pub fn store_limits(&self) -> StoreLimits { + let mut limits = StoreLimitsBuilder::new(); + if let Some(max) = self.common.wasm.max_memory_size { + limits = limits.memory_size(max); + } + if let Some(max) = self.common.wasm.max_table_elements { + limits = limits.table_elements(max); + } + if let Some(max) = self.common.wasm.max_instances { + limits = limits.instances(max); + } + if let Some(max) = self.common.wasm.max_tables { + limits = limits.tables(max); + } + if let Some(max) = self.common.wasm.max_memories { + limits = limits.memories(max); + } + if let Some(enable) = self.common.wasm.trap_on_grow_failure { + limits = limits.trap_on_grow_failure(enable); + } + + limits.build() + } + + pub fn ensure_allow_precompiled(&self) -> Result<()> { + if self.allow_precompiled { + Ok(()) + } else { + bail!("running a precompiled module requires the `--allow-precompiled` flag") + } + } + + #[cfg(feature = "component-model")] + fn ensure_allow_components(&self) -> Result<()> { + if self.common.wasm.component_model != Some(true) { + bail!("cannot execute a component without `--wasm component-model`"); + } + + Ok(()) + } + + pub fn load_module(&self, engine: &Engine, path: &Path) -> Result { + let path = match path.to_str() { + #[cfg(unix)] + Some("-") => "/dev/stdin".as_ref(), + _ => path, + }; + + // First attempt to load the module as an mmap. If this succeeds then + // detection can be done with the contents of the mmap and if a + // precompiled module is detected then `deserialize_file` can be used + // which is a slightly more optimal version than `deserialize` since we + // can leave most of the bytes on disk until they're referenced. + // + // If the mmap fails, for example if stdin is a pipe, then fall back to + // `std::fs::read` to load the contents. At that point precompiled + // modules must go through the `deserialize` functions. + // + // Note that this has the unfortunate side effect for precompiled + // modules on disk that they're opened once to detect what they are and + // then again internally in Wasmtime as part of the `deserialize_file` + // API. Currently there's no way to pass the `MmapVec` here through to + // Wasmtime itself (that'd require making `wasmtime-runtime` a public + // dependency or `MmapVec` a public type, both of which aren't ready to + // happen at this time). It's hoped though that opening a file twice + // isn't too bad in the grand scheme of things with respect to the CLI. + match wasmtime_runtime::MmapVec::from_file(path) { + Ok(map) => self.load_module_contents( + engine, + path, + &map, + || unsafe { Module::deserialize_file(engine, path) }, + #[cfg(feature = "component-model")] + || unsafe { Component::deserialize_file(engine, path) }, + ), + Err(_) => { + let bytes = std::fs::read(path) + .with_context(|| format!("failed to read file: {}", path.display()))?; + self.load_module_contents( + engine, + path, + &bytes, + || unsafe { Module::deserialize(engine, &bytes) }, + #[cfg(feature = "component-model")] + || unsafe { Component::deserialize(engine, &bytes) }, + ) + } + } + } + + pub fn load_module_contents( + &self, + engine: &Engine, + path: &Path, + bytes: &[u8], + deserialize_module: impl FnOnce() -> Result, + #[cfg(feature = "component-model")] deserialize_component: impl FnOnce() -> Result, + ) -> Result { + Ok(match engine.detect_precompiled(bytes) { + Some(Precompiled::Module) => { + self.ensure_allow_precompiled()?; + RunTarget::Core(deserialize_module()?) + } + #[cfg(feature = "component-model")] + Some(Precompiled::Component) => { + self.ensure_allow_precompiled()?; + self.ensure_allow_components()?; + RunTarget::Component(deserialize_component()?) + } + #[cfg(not(feature = "component-model"))] + Some(Precompiled::Component) => { + bail!("support for components was not enabled at compile time"); + } + None => { + // Parse the text format here specifically to add the `path` to + // the error message if there's a syntax error. + let wasm = wat::parse_bytes(bytes).map_err(|mut e| { + e.set_path(path); + e + })?; + if wasmparser::Parser::is_component(&wasm) { + #[cfg(feature = "component-model")] + { + self.ensure_allow_components()?; + RunTarget::Component(Component::new(engine, &wasm)?) + } + #[cfg(not(feature = "component-model"))] + { + bail!("support for components was not enabled at compile time"); + } + } else { + RunTarget::Core(Module::new(engine, &wasm)?) + } + } + }) + } +} + +#[derive(Clone)] +pub enum Profile { + Native(wasmtime::ProfilingStrategy), + Guest { path: String, interval: Duration }, +} + +impl Profile { + /// Parse the `profile` argument to either the `run` or `serve` commands. + pub fn parse(s: &str) -> Result { + let parts = s.split(',').collect::>(); + match &parts[..] { + ["perfmap"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::PerfMap)), + ["jitdump"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::JitDump)), + ["vtune"] => Ok(Profile::Native(wasmtime::ProfilingStrategy::VTune)), + ["guest"] => Ok(Profile::Guest { + path: "wasmtime-guest-profile.json".to_string(), + interval: Duration::from_millis(10), + }), + ["guest", path] => Ok(Profile::Guest { + path: path.to_string(), + interval: Duration::from_millis(10), + }), + ["guest", path, dur] => Ok(Profile::Guest { + path: path.to_string(), + interval: WasmtimeOptionValue::parse(Some(dur))?, + }), + _ => bail!("unknown profiling strategy: {s}"), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 5cf21f6fe768..11473981a467 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,3 +11,5 @@ #![warn(unused_import_braces)] pub mod commands; + +pub(crate) mod common;