From a8174bb1824a8fb1edab9bba98821ef155fcfe6a Mon Sep 17 00:00:00 2001 From: Dave Nagoda Date: Tue, 30 Jun 2026 16:38:23 -0700 Subject: [PATCH] Reuse the compiled provider across runs Batch mode reuses the loaded function module, but Javy/provider functions still recompiled the embedded standard provider for every input row, so provider setup dominated batch runtime. Store the compiled provider Module on Provider and compile it once when the ValidatedModule is built. FunctionRunParams now carries a ValidatedModule (built once in main), so single and batch runs reuse the already-compiled function and provider modules and per-row work is just instantiation. This drops the separate CompiledProvider type, the parallel *_with_compiled_provider entry points, and the per-row Module::from_binary fallback. Add batch coverage for a Javy v3 function to exercise the provider path. Measured locally: 5,000 js_function_javy_plugin_v3 rows complete in ~1.2s, versus ~52s when the provider was recompiled per input row. Verified with: - cargo test - cargo build --release Assisted-By: devx/c659e918-9568-4750-b122-e3890447348a --- src/engine.rs | 28 +++++++++++----------- src/io.rs | 11 +++------ src/main.rs | 18 ++++++++------ src/validated_module.rs | 49 ++++++++++++++++++++++---------------- tests/integration_tests.rs | 32 +++++++++++++++++++++++++ 5 files changed, 88 insertions(+), 50 deletions(-) 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!());