Reuse the compiled provider across runs

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<FunctionRunResult> {
         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<FunctionRunResult> {
     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,

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<T>(
         &mut self,
-        engine: &Engine,
         linker: &mut Linker<T>,
         store: &mut Store<T>,
     ) -> 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<T>(
     module: &ValidatedModule,
-    engine: &Engine,
     linker: &mut Linker<T>,
     mut store: &mut Store<T>,
 ) -> Option<Instance> {
     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() {

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<dyn Read + Sync + Send + 'static> = 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<dyn BufRead> = 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(),
         });
 

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<Provider>,
 }
 
 impl ValidatedModule {
-    pub(crate) fn new(module: Module) -> Result<Self> {
+    pub fn new(module: Module, engine: &Engine) -> Result<Self> {
         // 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(())
     }
 }

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::<serde_json::Value>)
+            .collect::<Result<Vec<_>, _>>()?;
+        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!());