diff --git a/src/engine.rs b/src/engine.rs index aa026118..d362cebd 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -7,7 +7,7 @@ use wasmtime_wasi::I32Exit; use crate::function_run_result::FunctionRunResult; use crate::io::{IOHandler, OutputAndLogs}; -use crate::validated_module::ValidatedModule; +pub use crate::validated_module::ValidatedModule; use crate::{BytesContainer, BytesContainerType}; #[derive(Clone)] @@ -31,7 +31,7 @@ pub struct FunctionRunParams<'a> { pub export: &'a str, pub profile_opts: Option<&'a ProfileOpts>, pub scale_factor: f64, - pub module: Module, + pub module: ValidatedModule, pub engine: Engine, } @@ -99,7 +99,7 @@ pub fn run(params: FunctionRunParams) -> Result { module, } = params; - let mut io_handler = IOHandler::new(ValidatedModule::new(module)?, input.clone()); + let mut io_handler = IOHandler::new(module, input.clone()); let mut error_logs: String = String::new(); @@ -116,7 +116,7 @@ pub fn run(params: FunctionRunParams) -> Result { let mut store = Store::new(&engine, function_context); store.limiter(|s| &mut s.limiter); - io_handler.initialize(&engine, &mut linker, &mut store)?; + io_handler.initialize(&mut linker, &mut store)?; store.set_fuel(STARTING_FUEL)?; store.set_epoch_deadline(1); @@ -235,7 +235,7 @@ mod tests { function_path: Path::new("tests/fixtures/build/js_function.wasm").to_path_buf(), input, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -260,7 +260,7 @@ mod tests { function_path: Path::new("tests/fixtures/build/js_function_v2.wasm").to_path_buf(), input, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -285,7 +285,7 @@ mod tests { function_path: Path::new("tests/fixtures/build/js_function_v3.wasm").to_path_buf(), input, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -311,7 +311,7 @@ mod tests { .to_path_buf(), input, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -329,7 +329,7 @@ mod tests { function_path: Path::new("tests/fixtures/build/exit_code.wasm").to_path_buf(), input: json_input(&serde_json::to_vec(&json!({ "code": 0 }))?)?, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -347,7 +347,7 @@ mod tests { function_path: Path::new("tests/fixtures/build/exit_code.wasm").to_path_buf(), input: json_input(&serde_json::to_vec(&json!({ "code": 1 }))?)?, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -368,7 +368,7 @@ mod tests { function_path: Path::new("tests/fixtures/build/linear_memory.wasm").to_path_buf(), input: json_input(&serde_json::to_vec(&json!({}))?)?, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -390,7 +390,7 @@ mod tests { function_path: Path::new("tests/fixtures/build/log_truncation_function.wasm") .to_path_buf(), export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -416,7 +416,7 @@ mod tests { function_path: file_path.to_path_buf(), input: json_input(&serde_json::to_vec(&json!({ "code": 0 }))?)?, export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, @@ -445,7 +445,7 @@ mod tests { function_path: trampolined_path.to_path_buf(), input: input_bytes.unwrap(), export: DEFAULT_EXPORT, - module, + module: ValidatedModule::new(module, &engine)?, engine, scale_factor: 1.0, profile_opts: None, diff --git a/src/io.rs b/src/io.rs index 46f26e04..bebb59f6 100644 --- a/src/io.rs +++ b/src/io.rs @@ -1,5 +1,5 @@ use anyhow::{anyhow, Result}; -use wasmtime::{AsContext, AsContextMut, Engine, Instance, Linker, Module, Store}; +use wasmtime::{AsContext, AsContextMut, Instance, Linker, Module, Store}; use wasmtime_wasi::{ p1::WasiP1Ctx, p2::pipe::{MemoryInputPipe, MemoryOutputPipe}, @@ -68,14 +68,13 @@ impl IOHandler { pub(crate) fn initialize( &mut self, - engine: &Engine, linker: &mut Linker, store: &mut Store, ) -> Result<()> { store.set_epoch_deadline(1); // Need to make sure we don't timeout during initialization. let old_fuel = store.get_fuel()?; store.set_fuel(u64::MAX)?; // Make sure we have fuel for initialization. - let mem_io_instance = instantiate_imports(&self.module, engine, linker, store); + let mem_io_instance = instantiate_imports(&self.module, linker, store); if let IOStrategy::Memory(ref mut instance) = self.strategy { *instance = mem_io_instance; } @@ -157,18 +156,14 @@ impl IOHandler { fn instantiate_imports( module: &ValidatedModule, - engine: &Engine, linker: &mut Linker, mut store: &mut Store, ) -> Option { let mut mem_io_instance = None; if let Some(std_import) = module.std_import() { - let imported_module = Module::from_binary(engine, &std_import.bytes) - .unwrap_or_else(|_| panic!("Failed to load module {}", std_import.name)); - let imported_module_instance = linker - .instantiate(&mut store, &imported_module) + .instantiate(&mut store, &std_import.module) .expect("Failed to instantiate imported instance"); if std_import.is_mem_io_provider() { diff --git a/src/main.rs b/src/main.rs index cd8dea94..d026aded 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,7 +11,7 @@ use anyhow::{anyhow, Result}; use clap::Parser; use function_runner::{ bluejay_schema_analyzer::BluejaySchemaAnalyzer, - engine::{run, FunctionRunParams, ProfileOpts}, + engine::{run, FunctionRunParams, ProfileOpts, ValidatedModule}, }; use is_terminal::IsTerminal; @@ -134,17 +134,21 @@ fn main() -> Result<()> { Codec::Json }; + // Validate and compile the module (and its standard provider) once, then + // reuse it across every run so batch mode doesn't recompile per input. + let validated_module = ValidatedModule::new(module, &engine)?; + if opts.batch { - run_batch_mode(&opts, &engine, &module, codec) + run_batch_mode(&opts, &engine, &validated_module, codec) } else { - run_single_mode(&opts, &engine, &module, codec) + run_single_mode(&opts, &engine, &validated_module, codec) } } fn run_single_mode( opts: &Opts, engine: &wasmtime::Engine, - module: &Module, + validated_module: &ValidatedModule, codec: Codec, ) -> Result<()> { let mut input: Box = if let Some(ref input) = opts.input { @@ -188,7 +192,7 @@ fn run_single_mode( export: opts.export.as_ref(), profile_opts: profile_opts.as_ref(), scale_factor, - module: module.clone(), + module: validated_module.clone(), engine: engine.clone(), })?; @@ -212,7 +216,7 @@ fn run_single_mode( fn run_batch_mode( opts: &Opts, engine: &wasmtime::Engine, - module: &Module, + validated_module: &ValidatedModule, codec: Codec, ) -> Result<()> { let input_reader: Box = if let Some(ref input) = opts.input { @@ -319,7 +323,7 @@ fn run_batch_mode( export: opts.export.as_ref(), profile_opts, scale_factor, - module: module.clone(), + module: validated_module.clone(), engine: engine.clone(), }); diff --git a/src/validated_module.rs b/src/validated_module.rs index 2d4b40e9..d24e3127 100644 --- a/src/validated_module.rs +++ b/src/validated_module.rs @@ -1,16 +1,14 @@ -use std::borrow::Cow; - use anyhow::{bail, Result}; use rust_embed::RustEmbed; -use wasmtime::Module; +use wasmtime::{Engine, Module}; #[derive(RustEmbed)] #[folder = "providers/"] struct StandardProviders; -#[derive(Debug)] +#[derive(Debug, Clone)] pub(crate) struct Provider { - pub(crate) bytes: Cow<'static, [u8]>, + pub(crate) module: Module, pub(crate) name: String, } @@ -38,14 +36,14 @@ impl Provider { } } -#[derive(Debug)] -pub(crate) struct ValidatedModule { +#[derive(Debug, Clone)] +pub struct ValidatedModule { module: Module, std_import: Option, } impl ValidatedModule { - pub(crate) fn new(module: Module) -> Result { + pub fn new(module: Module, engine: &Engine) -> Result { // Need to track with deterministic order so don't use a hash let mut imports = vec![]; for import in module.imports().map(|i| i.module().to_string()) { @@ -56,12 +54,17 @@ impl ValidatedModule { let uses_wasi = imports.contains(&"wasi_snapshot_preview1".to_string()); - let std_import = imports.iter().find_map(|import| { - StandardProviders::get(&format!("{import}.wasm")).map(|file| Provider { - bytes: file.data, - name: import.into(), + let std_import = imports + .iter() + .find_map(|import| { + StandardProviders::get(&format!("{import}.wasm")).map(|file| { + Module::from_binary(engine, &file.data).map(|module| Provider { + module, + name: import.into(), + }) + }) }) - }); + .transpose()?; // If there are multiple standard imports or more than zero unknown imports, // the module will fail to instantiate because we only link the one @@ -105,8 +108,9 @@ mod tests { (import "wasi_snapshot_preview1" "fd_read" (func)) ) "#; - let module = Module::new(&Engine::default(), &wat)?; - ValidatedModule::new(module)?; + let engine = Engine::default(); + let module = Module::new(&engine, wat)?; + ValidatedModule::new(module, &engine)?; Ok(()) } @@ -118,8 +122,9 @@ mod tests { (import "shopify_function_v1" "shopify_function_input_get" (func)) ) "#; - let module = Module::new(&Engine::default(), &wat)?; - ValidatedModule::new(module)?; + let engine = Engine::default(); + let module = Module::new(&engine, wat)?; + ValidatedModule::new(module, &engine)?; Ok(()) } @@ -130,8 +135,9 @@ mod tests { (import "shopify_function_v2" "shopify_function_input_get" (func)) ) "#; - let module = Module::new(&Engine::default(), &wat)?; - ValidatedModule::new(module)?; + let engine = Engine::default(); + let module = Module::new(&engine, wat)?; + ValidatedModule::new(module, &engine)?; Ok(()) } @@ -143,8 +149,9 @@ mod tests { (import "shopify_function_v2" "shopify_function_input_get" (func)) ) "#; - let module = Module::new(&Engine::default(), &wat)?; - ValidatedModule::new(module).unwrap_err(); + let engine = Engine::default(); + let module = Module::new(&engine, wat)?; + ValidatedModule::new(module, &engine).unwrap_err(); Ok(()) } } diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index bf18c63b..e58de3bf 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -127,6 +127,38 @@ mod tests { Ok(()) } + #[test] + fn batch_javy_plugin_v3_uses_provider() -> Result<()> { + let mut cmd = Command::new(cargo_bin!()); + let input_file = temp_batch_input("{\"hello\":\"world\"}\n{\"hello\":\"world\"}\n")?; + + cmd.args([ + "--function", + "tests/fixtures/build/js_function_javy_plugin_v3.wasm", + ]) + .arg("--json") + .arg("--batch") + .arg("--input") + .arg(input_file.as_os_str()); + + let output = cmd.output()?; + + assert!(output.status.success()); + + let stdout = String::from_utf8(output.stdout)?; + let results = stdout + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + assert_eq!(results.len(), 2); + assert!(results.iter().all(|result| result["success"] == true)); + assert!(results + .iter() + .all(|result| result["output"].to_string().contains("world output"))); + + Ok(()) + } + #[test] fn run_no_opts() -> Result<()> { let mut cmd = Command::new(cargo_bin!());