From 2885886747b9ca2500fe5e7d0092e31b9010f918 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rafael=20Fern=C3=A1ndez=20L=C3=B3pez?= Date: Thu, 28 Sep 2023 14:59:12 +0200 Subject: [PATCH] feat: introduce snapshot2 wasi context building --- crates/api-manage/src/handlers/v0/workers.rs | 2 +- crates/router/src/route.rs | 8 +- crates/runtimes/src/lib.rs | 1 + crates/runtimes/src/modules/external.rs | 34 +- crates/runtimes/src/modules/javascript.rs | 26 +- crates/runtimes/src/runtime.rs | 9 +- crates/server/src/handlers/worker.rs | 125 +++---- crates/worker/src/lib.rs | 368 ++++++++++++------- crates/worker/src/stdio.rs | 41 ++- 9 files changed, 390 insertions(+), 224 deletions(-) diff --git a/crates/api-manage/src/handlers/v0/workers.rs b/crates/api-manage/src/handlers/v0/workers.rs index ed247ca1..5cdaca7c 100644 --- a/crates/api-manage/src/handlers/v0/workers.rs +++ b/crates/api-manage/src/handlers/v0/workers.rs @@ -44,7 +44,7 @@ pub async fn handle_api_worker(routes: Data, path: Path) -> Http .map(|r| workers.get(&r.worker).expect("unexpected missing worker")); if let Some(worker) = worker { - HttpResponse::Ok().json(WorkerConfig::from(worker)) + HttpResponse::Ok().json(WorkerConfig::from(worker.as_ref())) } else { HttpResponse::NotFound().json("{}") } diff --git a/crates/router/src/route.rs b/crates/router/src/route.rs index 55724621..f4230b70 100644 --- a/crates/router/src/route.rs +++ b/crates/router/src/route.rs @@ -16,7 +16,7 @@ use std::{ collections::HashMap, ffi::OsStr, path::{Component, Path, PathBuf}, - sync::RwLock, + sync::{Arc, RwLock}, }; use wws_config::Config as ProjectConfig; use wws_worker::Worker; @@ -60,16 +60,16 @@ pub struct Route { /// other crates. #[derive(Default)] pub struct WorkerSet { - workers: HashMap, + workers: HashMap>, } impl WorkerSet { - pub fn get(&self, worker_id: &str) -> Option<&Worker> { + pub fn get(&self, worker_id: &str) -> Option<&Arc> { self.workers.get(worker_id) } pub fn register(&mut self, worker_id: String, worker: Worker) { - self.workers.insert(worker_id, worker); + self.workers.insert(worker_id, Arc::new(worker)); } } diff --git a/crates/runtimes/src/lib.rs b/crates/runtimes/src/lib.rs index 687bb4c1..967ea156 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 a41f85ef..19bdbf2c 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 9c45895d..adec59ca 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 f621f1a2..2988a3c3 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 87f19451..8666cb2a 100644 --- a/crates/server/src/handlers/worker.rs +++ b/crates/server/src/handlers/worker.rs @@ -38,11 +38,10 @@ pub async fn handle_worker(req: HttpRequest, body: Bytes) -> HttpResponse { let data_connectors = req .app_data::>>() .expect("error fetching data connectors"); - let result: HttpResponse; // First, we need to identify the best suited route let selected_route = app_data.routes.retrieve_best_route(req.path()); - if let Some(route) = selected_route { + let worker = if let Some(route) = selected_route { // First, check if there's an existing static file. Static assets have more priority // than dynamic routes. However, I cannot set the static assets as the first service // as it's captures everything. @@ -56,74 +55,76 @@ pub async fn handle_worker(req: HttpRequest, body: Bytes) -> HttpResponse { .read() .expect("error locking worker lock for reading"); - let worker = workers - .get(&route.worker) - .expect("unexpected missing worker"); + Some( + workers + .get(&route.worker) + .expect("unexpected missing worker") + .clone(), + ) + } else { + None + }; - // Let's continue - let body_str = String::from_utf8(body.to_vec()).unwrap_or_else(|_| String::from("")); + if worker.is_none() { + return handle_not_found(&req).await; + }; + let worker = worker.unwrap(); - // Init from configuration - let vars = &worker.config.vars; - let kv_namespace = worker.config.data_kv_namespace(); + let body_str = String::from_utf8(body.to_vec()).unwrap_or_else(|_| String::from("")); - let store = match &kv_namespace { - Some(namespace) => { - let connector = data_connectors - .read() - .expect("error locking data connectors lock for reading"); - let kv_store = connector.kv.find_store(namespace); + // Init from configuration + let vars = &worker.config.vars; + let kv_namespace = worker.config.data_kv_namespace(); - kv_store.map(|store| store.clone()) - } - None => None, - }; - - let (handler_result, handler_success) = match worker.run(&req, &body_str, store, vars) { - Ok(output) => (output, true), - Err(err) => (WasmOutput::failed(err), false), - }; - - let mut builder = HttpResponse::build( - StatusCode::from_u16(handler_result.status).unwrap_or(StatusCode::OK), - ); - // Default content type - builder.insert_header(("Content-Type", "text/html")); - - // Check if cors config has any origins to register - if let Some(origins) = app_data.cors_origins.as_ref() { - // Check if worker has overridden the header, if not - if !handler_result.headers.contains_key(CORS_HEADER) { - // insert those origins in 'Access-Control-Allow-Origin' header - let header_value = origins.join(","); - builder.insert_header((CORS_HEADER, header_value)); - } - } + let store = match &kv_namespace { + Some(namespace) => { + let connector = data_connectors + .read() + .expect("error locking data connectors lock for reading"); + let kv_store = connector.kv.find_store(namespace); - for (key, val) in handler_result.headers.iter() { - // Note that QuickJS is replacing the "-" character - // with "_" on property keys. Here, we rollback it - builder.insert_header((key.replace('_', "-").as_str(), val.as_str())); + kv_store.map(|store| store.clone()) } - - // Write to the state if required - if handler_success && kv_namespace.is_some() { - data_connectors - .write() - .expect("error locking data connectors lock for writing") - .kv - .replace_store(&kv_namespace.unwrap(), &handler_result.kv) + None => None, + }; + + let (handler_result, handler_success) = match worker.run(&req, &body_str, store, vars).await { + Ok(output) => (output, true), + Err(err) => (WasmOutput::failed(err), false), + }; + + let mut builder = + HttpResponse::build(StatusCode::from_u16(handler_result.status).unwrap_or(StatusCode::OK)); + // Default content type + builder.insert_header(("Content-Type", "text/html")); + + // Check if cors config has any origins to register + if let Some(origins) = app_data.cors_origins.as_ref() { + // Check if worker has overridden the header, if not + if !handler_result.headers.contains_key(CORS_HEADER) { + // insert those origins in 'Access-Control-Allow-Origin' header + let header_value = origins.join(","); + builder.insert_header((CORS_HEADER, header_value)); } + } - result = match handler_result.body() { - Ok(res) => builder.body(res), - Err(_) => { - HttpResponse::ServiceUnavailable().body("There was an error running the worker") - } - } - } else { - result = handle_not_found(&req).await; + for (key, val) in handler_result.headers.iter() { + // Note that QuickJS is replacing the "-" character + // with "_" on property keys. Here, we rollback it + builder.insert_header((key.replace('_', "-").as_str(), val.as_str())); } - result + // Write to the state if required + if handler_success && kv_namespace.is_some() { + data_connectors + .write() + .expect("error locking data connectors lock for writing") + .kv + .replace_store(&kv_namespace.unwrap(), &handler_result.kv) + } + + match handler_result.body() { + Ok(res) => builder.body(res), + Err(_) => HttpResponse::ServiceUnavailable().body("There was an error running the worker"), + } } diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 5a6ef2f4..3e45c788 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 5bd35c50..4b36fa1a 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() } }