diff --git a/crates/runtimes/src/lib.rs b/crates/runtimes/src/lib.rs index 687bb4c..967ea15 100644 --- a/crates/runtimes/src/lib.rs +++ b/crates/runtimes/src/lib.rs @@ -6,6 +6,7 @@ use errors::Result; mod modules; mod runtime; +pub use runtime::CtxBuilder; use modules::{external::ExternalRuntime, javascript::JavaScriptRuntime, native::NativeRuntime}; use std::path::Path; diff --git a/crates/runtimes/src/modules/external.rs b/crates/runtimes/src/modules/external.rs index a41f85e..19bdbf2 100644 --- a/crates/runtimes/src/modules/external.rs +++ b/crates/runtimes/src/modules/external.rs @@ -3,12 +3,12 @@ use crate::errors::{self, Result}; -use crate::runtime::Runtime; +use crate::runtime::{CtxBuilder, Runtime}; use std::{ fs, path::{Path, PathBuf}, }; -use wasmtime_wasi::{ambient_authority, Dir, WasiCtxBuilder}; +use wasmtime_wasi::{ambient_authority, preview2, Dir}; use wws_project::metadata::Runtime as RuntimeMetadata; use wws_store::Store; @@ -90,14 +90,28 @@ impl Runtime for ExternalRuntime { /// Mount the source code in the WASI context so it can be /// processed by the engine - fn prepare_wasi_ctx(&self, builder: &mut WasiCtxBuilder) -> Result<()> { - builder - .preopened_dir( - Dir::open_ambient_dir(&self.store.folder, ambient_authority())?, - "/src", - )? - .args(&self.metadata.args) - .map_err(|_| errors::RuntimeError::WasiContextError)?; + fn prepare_wasi_ctx(&self, builder: &mut CtxBuilder) -> Result<()> { + match builder { + CtxBuilder::Preview1(ref mut builder) => { + builder + .preopened_dir( + Dir::open_ambient_dir(&self.store.folder, ambient_authority())?, + "/src", + )? + .args(&self.metadata.args) + .map_err(|_| errors::RuntimeError::WasiContextError)?; + } + CtxBuilder::Preview2(ref mut builder) => { + builder + .preopened_dir( + Dir::open_ambient_dir(&self.store.folder, ambient_authority())?, + preview2::DirPerms::all(), + preview2::FilePerms::all(), + "/src", + ) + .args(&self.metadata.args); + } + } Ok(()) } diff --git a/crates/runtimes/src/modules/javascript.rs b/crates/runtimes/src/modules/javascript.rs index 9c45895..adec59c 100644 --- a/crates/runtimes/src/modules/javascript.rs +++ b/crates/runtimes/src/modules/javascript.rs @@ -2,10 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 use crate::errors::Result; -use crate::runtime::Runtime; +use crate::runtime::{CtxBuilder, Runtime}; use std::path::{Path, PathBuf}; -use wasmtime_wasi::{ambient_authority, Dir, WasiCtxBuilder}; +use wasmtime_wasi::{ambient_authority, preview2, Dir}; use wws_store::Store; static JS_ENGINE_WASM: &[u8] = @@ -46,11 +46,23 @@ impl Runtime for JavaScriptRuntime { /// Mount the source code in the WASI context so it can be /// processed by the engine - fn prepare_wasi_ctx(&self, builder: &mut WasiCtxBuilder) -> Result<()> { - builder.preopened_dir( - Dir::open_ambient_dir(&self.store.folder, ambient_authority())?, - "/src", - )?; + fn prepare_wasi_ctx(&self, builder: &mut CtxBuilder) -> Result<()> { + match builder { + CtxBuilder::Preview1(ref mut builder) => { + builder.preopened_dir( + Dir::open_ambient_dir(&self.store.folder, ambient_authority())?, + "/src", + )?; + } + CtxBuilder::Preview2(ref mut builder) => { + builder.preopened_dir( + Dir::open_ambient_dir(&self.store.folder, ambient_authority())?, + preview2::DirPerms::all(), + preview2::FilePerms::all(), + "/src", + ); + } + } Ok(()) } diff --git a/crates/runtimes/src/runtime.rs b/crates/runtimes/src/runtime.rs index f621f1a..2988a3c 100644 --- a/crates/runtimes/src/runtime.rs +++ b/crates/runtimes/src/runtime.rs @@ -3,7 +3,12 @@ use crate::errors::Result; -use wasmtime_wasi::WasiCtxBuilder; +use wasmtime_wasi::{preview2, WasiCtxBuilder}; + +pub enum CtxBuilder { + Preview1(WasiCtxBuilder), + Preview2(preview2::WasiCtxBuilder), +} /// Define the behavior a Runtime must have. This includes methods /// to initialize the environment for the given runtime as well as @@ -21,7 +26,7 @@ pub trait Runtime { /// WASI context builder. This allow runtimes to mount /// specific lib folders, source code and adding /// environment variables. - fn prepare_wasi_ctx(&self, _builder: &mut WasiCtxBuilder) -> Result<()> { + fn prepare_wasi_ctx(&self, _builder: &mut CtxBuilder) -> Result<()> { Ok(()) } diff --git a/crates/server/src/handlers/worker.rs b/crates/server/src/handlers/worker.rs index 87f1945..5344d41 100644 --- a/crates/server/src/handlers/worker.rs +++ b/crates/server/src/handlers/worker.rs @@ -79,7 +79,8 @@ pub async fn handle_worker(req: HttpRequest, body: Bytes) -> HttpResponse { None => None, }; - let (handler_result, handler_success) = match worker.run(&req, &body_str, store, vars) { + let (handler_result, handler_success) = match worker.run(&req, &body_str, store, vars).await + { Ok(output) => (output, true), Err(err) => (WasmOutput::failed(err), false), }; diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 5a6ef2f..3e45c78 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -19,11 +19,14 @@ use std::path::PathBuf; use std::sync::Arc; use std::{collections::HashMap, path::Path}; use stdio::Stdio; -use wasmtime::{component::Component, Engine, Linker, Module, Store}; +use wasmtime::{ + component::{self, Component}, + Config as WasmtimeConfig, Engine, Linker, Module, Store, +}; use wasmtime_wasi::{ambient_authority, preview2, Dir, WasiCtxBuilder}; use wasmtime_wasi_nn::{InMemoryRegistry, Registry, WasiNnCtx}; use wws_config::Config as ProjectConfig; -use wws_runtimes::{init_runtime, Runtime}; +use wws_runtimes::{init_runtime, CtxBuilder, Runtime}; pub enum ModuleOrComponent { Module(Module), @@ -38,16 +41,17 @@ pub struct Worker { pub id: String, /// Wasmtime engine to run this worker engine: Engine, - /// Wasm Module or component - module_or_component: ModuleOrComponent, /// Worker runtime runtime: Box, + /// Wasm Module or component + module_or_component: ModuleOrComponent, /// Current config pub config: Config, /// The worker filepath path: PathBuf, } +#[derive(Default)] struct Host { pub wasi_preview1_ctx: Option, pub wasi_preview2_ctx: Option>, @@ -62,7 +66,7 @@ struct Host { wasi_preview2_adapter: Arc, pub wasi_nn: Option>, - pub http: HttpBindings, + pub http: Option, } impl preview2::WasiView for Host { @@ -116,7 +120,14 @@ impl Worker { } } - let engine = Engine::default(); + let engine = Engine::new( + WasmtimeConfig::default() + .async_support(true) + .wasm_component_model(true), + ) + .map_err(|err| errors::WorkerError::ConfigureRuntimeError { + error: format!("error creating engine ({err})"), + })?; let runtime = init_runtime(project_root, path, project_config)?; let bytes = runtime.module_bytes()?; let module_or_component = if let Ok(module) = Module::from_binary(&engine, &bytes) { @@ -133,14 +144,86 @@ impl Worker { Ok(Self { id, engine, - module_or_component, runtime, + module_or_component, config, path: path.to_path_buf(), }) } - pub fn run( + pub fn prepare_wasi_context( + &self, + environment_variables: &[(String, String)], + wasi_builder: &mut CtxBuilder, + ) -> Result<()> { + match wasi_builder { + CtxBuilder::Preview1(wasi_builder) => { + // Set up environment variables + wasi_builder.envs(environment_variables).map_err(|error| { + errors::WorkerError::ConfigureRuntimeError { + error: format!("error configuring runtime: {error}"), + } + })?; + + // Setup pre-opens + if let Some(folders) = self.config.folders.as_ref() { + for folder in folders { + if let Some(base) = &self.path.parent() { + let dir = + Dir::open_ambient_dir(base.join(&folder.from), ambient_authority()) + .map_err(|error| { + errors::WorkerError::ConfigureRuntimeError { + error: format!( + "error setting up pre-opened folders: {error}" + ), + } + })?; + wasi_builder + .preopened_dir(dir, &folder.to) + .map_err(|error| errors::WorkerError::ConfigureRuntimeError { + error: format!("error setting up pre-opened folders: {error}"), + })?; + } else { + return Err(errors::WorkerError::FailedToInitialize); + } + } + } + } + CtxBuilder::Preview2(wasi_builder) => { + // Set up environment variables + wasi_builder.envs(environment_variables); + + // Setup pre-opens + if let Some(folders) = self.config.folders.as_ref() { + for folder in folders { + if let Some(base) = &self.path.parent() { + let dir = + Dir::open_ambient_dir(base.join(&folder.from), ambient_authority()) + .map_err(|error| { + errors::WorkerError::ConfigureRuntimeError { + error: format!( + "error setting up pre-opened folders: {error}" + ), + } + })?; + wasi_builder.preopened_dir( + dir, + preview2::DirPerms::all(), + preview2::FilePerms::all(), + &folder.to, + ); + } else { + return Err(errors::WorkerError::FailedToInitialize); + } + } + } + } + } + + Ok(()) + } + + pub async fn run( &self, request: &HttpRequest, body: &str, @@ -150,145 +233,178 @@ impl Worker { let input = serde_json::to_string(&WasmInput::new(request, body, kv)).unwrap(); let mut linker = Linker::new(&self.engine); + let mut component_linker = component::Linker::new(&self.engine); - http_add_to_linker(&mut linker, |host: &mut Host| &mut host.http).map_err(|error| { - errors::WorkerError::ConfigureRuntimeError { - error: format!("error adding HTTP bindings to linker ({error})"), - } - })?; - wasmtime_wasi::add_to_linker(&mut linker, |host| host.wasi_preview1_ctx.as_mut().unwrap()) + if let ModuleOrComponent::Module(_) = &self.module_or_component { + wasmtime_wasi::add_to_linker(&mut linker, |host: &mut Host| { + host.wasi_preview1_ctx.as_mut().unwrap() + }) .map_err(|error| errors::WorkerError::ConfigureRuntimeError { - error: format!("error adding WASI to linker ({error})"), + error: format!("error adding WASI preview1 to linker ({error})"), })?; - // I have to use `String` as it's required by WasiCtxBuilder - let tuple_vars: Vec<(String, String)> = - vars.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); - - // Create the initial WASI context - let mut wasi_builder = WasiCtxBuilder::new(); - wasi_builder.envs(&tuple_vars).map_err(|error| { - errors::WorkerError::ConfigureRuntimeError { - error: format!("error configuring runtime: {error}"), - } - })?; - - // Configure the stdio - let stdio = Stdio::new(&input); - stdio.configure_wasi_ctx(&mut wasi_builder); - - // Mount folders from the configuration - if let Some(folders) = self.config.folders.as_ref() { - for folder in folders { - if let Some(base) = &self.path.parent() { - let dir = Dir::open_ambient_dir(base.join(&folder.from), ambient_authority()) - .map_err(|error| errors::WorkerError::ConfigureRuntimeError { - error: format!("error setting up pre-opened folders: {error}"), - })?; - wasi_builder - .preopened_dir(dir, &folder.to) - .map_err(|error| errors::WorkerError::ConfigureRuntimeError { - error: format!("error setting up pre-opened folders: {error}"), - })?; - } else { - return Err(errors::WorkerError::FailedToInitialize); + http_add_to_linker(&mut linker, |host: &mut Host| host.http.as_mut().unwrap()) + .map_err(|error| errors::WorkerError::ConfigureRuntimeError { + error: format!("error adding HTTP bindings to linker ({error})"), + })?; + } else { + preview2::command::add_to_linker(&mut component_linker).map_err(|error| { + errors::WorkerError::ConfigureRuntimeError { + error: format!("error adding WASI preview2 to linker ({error})"), } - } - } - - // WASI-NN - let allowed_backends = &self.config.features.wasi_nn.allowed_backends; - let preload_models = &self.config.features.wasi_nn.preload_models; - - let wasi_nn = if !preload_models.is_empty() { - // Preload the models on the host. - let graphs = preload_models - .iter() - .map(|m| m.build_graph_data(&self.path)) - .collect::>(); - let (backends, registry) = wasmtime_wasi_nn::preload(&graphs).map_err(|_| { - errors::WorkerError::RuntimeError( - wws_runtimes::errors::RuntimeError::WasiContextError, - ) })?; + } - Some(Arc::new(WasiNnCtx::new(backends, registry))) - } else if !allowed_backends.is_empty() { - let registry = Registry::from(InMemoryRegistry::new()); - let mut backends = Vec::new(); - - // Load the given backends: - for b in allowed_backends.iter() { - if let Some(backend) = b.to_backend() { - backends.push(backend); - } - } + let environment_variables: Vec<(String, String)> = + vars.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); - Some(Arc::new(WasiNnCtx::new(backends, registry))) + let mut wasi_builder = if let ModuleOrComponent::Module(_) = &self.module_or_component { + CtxBuilder::Preview1(WasiCtxBuilder::new()) } else { - None + CtxBuilder::Preview2(preview2::WasiCtxBuilder::new()) }; + self.prepare_wasi_context(&environment_variables, &mut wasi_builder)?; - // Load the Wasi NN linker - if wasi_nn.is_some() { - wasmtime_wasi_nn::witx::add_to_linker(&mut linker, |host: &mut Host| { - Arc::get_mut(host.wasi_nn.as_mut().unwrap()) - .expect("wasi-nn is not implemented with multi-threading support") - }) - .map_err(|_| { - errors::WorkerError::RuntimeError( - wws_runtimes::errors::RuntimeError::WasiContextError, - ) - })?; - } + let stdio = Stdio::new(&input); + let mut wasi_builder = stdio.configure_wasi_ctx(wasi_builder); - // Pass to the runtime to add any WASI specific requirement self.runtime.prepare_wasi_ctx(&mut wasi_builder)?; - let wasi = wasi_builder.build(); - let host = Host { - wasi_preview1_ctx: Some(wasi), - wasi_preview2_ctx: None, - wasi_preview2_table: Arc::new(preview2::Table::default()), - wasi_preview2_adapter: Arc::new(preview2::preview1::WasiPreview1Adapter::default()), - wasi_nn, - http: HttpBindings { - http_config: self.config.features.http_requests.clone(), + let host = match wasi_builder { + CtxBuilder::Preview1(mut wasi_builder) => Host { + wasi_preview1_ctx: Some(wasi_builder.build()), + wasi_nn: None, + http: Some(HttpBindings { + http_config: self.config.features.http_requests.clone(), + }), + ..Host::default() }, + CtxBuilder::Preview2(mut wasi_builder) => { + let mut table = preview2::Table::default(); + Host { + wasi_preview2_ctx: Some(Arc::new(wasi_builder.build(&mut table).map_err( + |error| errors::WorkerError::ConfigureRuntimeError { + error: format!("error configuring WASI preview 2: {error}"), + }, + )?)), + wasi_preview2_table: Arc::new(table), + wasi_preview2_adapter: Arc::new( + preview2::preview1::WasiPreview1Adapter::default(), + ), + wasi_nn: None, + http: Some(HttpBindings { + http_config: self.config.features.http_requests.clone(), + }), + ..Host::default() + } + } }; - let mut store = Store::new(&self.engine, host); - let module = match &self.module_or_component { - ModuleOrComponent::Module(module) => module, - ModuleOrComponent::Component(_) => unimplemented!(), - }; + // Setup wasi-nn + { + let allowed_backends = &self.config.features.wasi_nn.allowed_backends; + let preload_models = &self.config.features.wasi_nn.preload_models; + let wasi_nn = if !preload_models.is_empty() { + // Preload the models on the host. + let graphs = preload_models + .iter() + .map(|m| m.build_graph_data(&self.path)) + .collect::>(); + let (backends, registry) = wasmtime_wasi_nn::preload(&graphs).map_err(|_| { + errors::WorkerError::RuntimeError( + wws_runtimes::errors::RuntimeError::WasiContextError, + ) + })?; + + Some(Arc::new(WasiNnCtx::new(backends, registry))) + } else if !allowed_backends.is_empty() { + let registry = Registry::from(InMemoryRegistry::new()); + let mut backends = Vec::new(); + + // Load the given backends: + for b in allowed_backends.iter() { + if let Some(backend) = b.to_backend() { + backends.push(backend); + } + } - linker.module(&mut store, "", module).map_err(|error| { - errors::WorkerError::ConfigureRuntimeError { - error: format!("error retrieving module from linker: {error}"), + Some(Arc::new(WasiNnCtx::new(backends, registry))) + } else { + None + }; + + // Load the Wasi NN linker + if wasi_nn.is_some() { + wasmtime_wasi_nn::witx::add_to_linker(&mut linker, |host: &mut Host| { + Arc::get_mut(host.wasi_nn.as_mut().unwrap()) + .expect("wasi-nn is not implemented with multi-threading support") + }) + .map_err(|_| { + errors::WorkerError::RuntimeError( + wws_runtimes::errors::RuntimeError::WasiContextError, + ) + })?; } - })?; - linker - .get_default(&mut store, "") - .map_err(|error| errors::WorkerError::ConfigureRuntimeError { - error: format!("error getting default export from module: {error}"), - })? - .typed::<(), ()>(&store) - .map_err(|error| errors::WorkerError::ConfigureRuntimeError { - error: format!("error getting default typed export from module: {error}"), - })? - .call(&mut store, ()) - .map_err(|error| errors::WorkerError::ConfigureRuntimeError { - error: format!("error calling module default export: {error}"), - })?; + } + + let contents = { + let mut store = Store::new(&self.engine, host); + match &self.module_or_component { + ModuleOrComponent::Module(module) => { + linker + .module_async(&mut store, "", module) + .await + .map_err(|error| errors::WorkerError::ConfigureRuntimeError { + error: format!("error retrieving module from linker: {error}"), + })?; + + linker + .get_default(&mut store, "") + .map_err(|error| errors::WorkerError::ConfigureRuntimeError { + error: format!("error getting default export from module: {error}"), + })? + .typed::<(), ()>(&store) + .map_err(|error| errors::WorkerError::ConfigureRuntimeError { + error: format!( + "error getting default typed export from module: {error}" + ), + })? + .call_async(&mut store, ()) + .await + .map_err(|error| errors::WorkerError::ConfigureRuntimeError { + error: format!("error calling module default export: {error}"), + })?; - drop(store); + drop(store); - let contents: Vec = stdio - .stdout - .try_into_inner() - .unwrap_or_default() - .into_inner(); + stdio + .stdout + .try_into_inner() + .unwrap_or_default() + .into_inner() + } + ModuleOrComponent::Component(component) => { + let (command, _instance) = preview2::command::Command::instantiate_async( + &mut store, + component, + &component_linker, + ) + .await + .unwrap(); + let _ = command + .wasi_cli_run() + .call_run(&mut store) + .await + .map_err(|error| errors::WorkerError::ConfigureRuntimeError { + error: format!("error calling component cli::run: {error}"), + })?; + + drop(store); + + stdio.stdout_preview2.contents().to_vec() + } + } + }; // Build the output let output: WasmOutput = serde_json::from_slice(&contents).map_err(|error| { diff --git a/crates/worker/src/stdio.rs b/crates/worker/src/stdio.rs index 5bd35c5..4b36fa1 100644 --- a/crates/worker/src/stdio.rs +++ b/crates/worker/src/stdio.rs @@ -1,6 +1,9 @@ use std::io::Cursor; use wasi_common::pipe::{ReadPipe, WritePipe}; -use wasmtime_wasi::WasiCtxBuilder; +use wasmtime_wasi::preview2; +use wws_runtimes::CtxBuilder; + +const MAX_OUTPUT_BYTES: usize = 10240; /// A library to configure the stdio of the WASI context. /// Note that currently, wws relies on stdin and stdout @@ -9,28 +12,42 @@ use wasmtime_wasi::WasiCtxBuilder; /// The stdin/stdout approach will change in the future with /// a more performant and appropiate approach. pub struct Stdio { - /// Defines the stdin ReadPipe to send the data to the module - pub stdin: ReadPipe>, - /// Defines the stdout to extract the data from the module + /// Defines the stdin ReadPipe to send data to the module + pub stdin: Vec, + /// Defines the stdout to extract data from the module pub stdout: WritePipe>>, + /// Defines the stdout to extract data from the module + pub stdout_preview2: preview2::pipe::MemoryOutputPipe, } impl Stdio { /// Initialize the stdio. The stdin will contain the input data. pub fn new(input: &str) -> Self { Self { - stdin: ReadPipe::from(input), + stdin: Vec::from(input), stdout: WritePipe::new_in_memory(), + stdout_preview2: preview2::pipe::MemoryOutputPipe::new(MAX_OUTPUT_BYTES), } } - pub fn configure_wasi_ctx<'a>( - &self, - builder: &'a mut WasiCtxBuilder, - ) -> &'a mut WasiCtxBuilder { + pub fn configure_wasi_ctx(&self, mut builder: CtxBuilder) -> CtxBuilder { + match builder { + CtxBuilder::Preview1(ref mut wasi_builder) => { + wasi_builder + .stdin(Box::new(ReadPipe::from(self.stdin.clone()).clone())) + .stdout(Box::new(self.stdout.clone())) + .inherit_stderr(); + } + CtxBuilder::Preview2(ref mut wasi_builder) => { + wasi_builder + .stdin( + preview2::pipe::MemoryInputPipe::new(self.stdin.clone().into()), + preview2::IsATTY::No, + ) + .stdout(self.stdout_preview2.clone(), preview2::IsATTY::No) + .inherit_stderr(); + } + } builder - .stdin(Box::new(self.stdin.clone())) - .stdout(Box::new(self.stdout.clone())) - .inherit_stderr() } }