Compile the standard provider once per module

`run` compiled the embedded standard provider (Javy plugin or Shopify
Function provider) from its bytes on every call. Library callers that run
the same Function many times paid that cost for every run.

Make `ValidatedModule` public and compile the provider when it is built.
`FunctionRunParams` now borrows a `ValidatedModule`, and `run` takes the
engine from the module, so callers build the module once and reuse it.

Author: dnagoda

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,8 +31,9 @@ pub struct FunctionRunParams<'a> {
     pub export: &'a str,
     pub profile_opts: Option<&'a ProfileOpts>,
     pub scale_factor: f64,
-    pub module: Module,
-    pub engine: Engine,
+    /// The validated Function module. Reuse it across runs to compile the
+    /// Function and its provider once.
+    pub module: &'a ValidatedModule,
 }
 
 const STARTING_FUEL: u64 = u64::MAX;
@@ -95,15 +96,15 @@ pub fn run(params: FunctionRunParams) -> Result<FunctionRunResult> {
         export,
         profile_opts,
         scale_factor,
-        engine,
         module,
     } = params;
 
-    let mut io_handler = IOHandler::new(ValidatedModule::new(module)?, input.clone());
+    let engine = module.inner().engine();
+    let mut io_handler = IOHandler::new(module, input.clone());
 
     let mut error_logs: String = String::new();
 
-    let mut linker = Linker::new(&engine);
+    let mut linker = Linker::new(engine);
     let wasi = io_handler.wasi();
     if wasi.is_some() {
         wasmtime_wasi::preview1::add_to_linker_sync(&mut linker, |ctx: &mut FunctionContext| {
@@ -113,10 +114,10 @@ pub fn run(params: FunctionRunParams) -> Result<FunctionRunResult> {
     }
 
     let function_context = FunctionContext::new(wasi);
-    let mut store = Store::new(&engine, function_context);
+    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,8 +236,7 @@ mod tests {
             function_path: Path::new("tests/fixtures/build/js_function.wasm").to_path_buf(),
             input,
             export: DEFAULT_EXPORT,
-            module,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -260,8 +260,7 @@ mod tests {
             function_path: Path::new("tests/fixtures/build/js_function_v2.wasm").to_path_buf(),
             input,
             export: DEFAULT_EXPORT,
-            module,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -285,8 +284,7 @@ mod tests {
             function_path: Path::new("tests/fixtures/build/js_function_v3.wasm").to_path_buf(),
             input,
             export: DEFAULT_EXPORT,
-            module,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -295,6 +293,37 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_validated_module_can_be_reused_across_runs() -> Result<()> {
+        let path = Path::new("tests/fixtures/build/js_function_javy_plugin_v3.wasm");
+        let engine = new_engine()?;
+        let module = ValidatedModule::new(Module::from_file(&engine, path)?)?;
+
+        for _ in 0..3 {
+            let input = BytesContainer::new(
+                BytesContainerType::Input,
+                Codec::Messagepack,
+                serde_json::to_vec(&json!({ "hello": "world" }))?,
+            )?;
+            let function_run_result = run(FunctionRunParams {
+                function_path: path.to_path_buf(),
+                input,
+                export: DEFAULT_EXPORT,
+                module: &module,
+                scale_factor: 1.0,
+                profile_opts: None,
+            })?;
+
+            assert!(function_run_result.success);
+            assert_eq!(
+                function_run_result.output.json_value,
+                Some(json!({ "hello": "world output" }))
+            );
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn test_js_functions_javy_v1() -> Result<()> {
         let engine = new_engine()?;
@@ -311,8 +340,7 @@ mod tests {
                 .to_path_buf(),
             input,
             export: DEFAULT_EXPORT,
-            module,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -329,8 +357,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,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -347,8 +374,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,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -368,8 +394,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,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -390,8 +415,7 @@ mod tests {
             function_path: Path::new("tests/fixtures/build/log_truncation_function.wasm")
                 .to_path_buf(),
             export: DEFAULT_EXPORT,
-            module,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -416,8 +440,7 @@ mod tests {
             function_path: file_path.to_path_buf(),
             input: json_input(&serde_json::to_vec(&json!({ "code": 0 }))?)?,
             export: DEFAULT_EXPORT,
-            module,
-            engine,
+            module: &ValidatedModule::new(module)?,
             scale_factor: 1.0,
             profile_opts: None,
         })?;
@@ -445,8 +468,7 @@ mod tests {
             function_path: trampolined_path.to_path_buf(),
             input: input_bytes.unwrap(),
             export: DEFAULT_EXPORT,
-            module,
-            engine,
+            module: &ValidatedModule::new(module)?,
             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::{
     p2::pipe::{MemoryInputPipe, MemoryOutputPipe},
     preview1::WasiP1Ctx,
@@ -25,14 +25,14 @@ enum IOStrategy {
     Memory(Option<Instance>),
 }
 
-pub(crate) struct IOHandler {
+pub(crate) struct IOHandler<'a> {
     strategy: IOStrategy,
-    module: ValidatedModule,
+    module: &'a ValidatedModule,
     input: BytesContainer,
 }
 
-impl IOHandler {
-    pub(crate) fn new(module: ValidatedModule, input: BytesContainer) -> Self {
+impl<'a> IOHandler<'a> {
+    pub(crate) fn new(module: &'a ValidatedModule, input: BytesContainer) -> Self {
         Self {
             strategy: if module.uses_mem_io() {
                 IOStrategy::Memory(None)
@@ -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;
@@ -144,6 +144,8 @@ fn main() -> Result<()> {
         Codec::Json
     };
 
+    let module = ValidatedModule::new(module)?;
+
     let input = BytesContainer::new(BytesContainerType::Input, codec, buffer)?;
     let scale_factor = if let (Some(schema_string), Some(query_string), Some(json_value)) =
         (schema_string, query_string, input.json_value.clone())
@@ -167,8 +169,7 @@ fn main() -> Result<()> {
         export: opts.export.as_ref(),
         profile_opts: profile_opts.as_ref(),
         scale_factor,
-        module,
-        engine,
+        module: &module,
     })?;
 
     if opts.json {

src/validated_module.rs

@@ -1,16 +1,14 @@
-use std::borrow::Cow;
-
-use anyhow::{bail, Result};
+use anyhow::{anyhow, bail, Result};
 use rust_embed::RustEmbed;
 use wasmtime::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,21 @@ impl Provider {
     }
 }
 
-#[derive(Debug)]
-pub(crate) struct ValidatedModule {
+/// A Function module that has been checked for a valid combination of imports,
+/// together with the compiled standard provider it imports, if any.
+///
+/// Building a `ValidatedModule` compiles the provider once. Reuse the same
+/// `ValidatedModule` for every run of the Function to avoid compiling the
+/// provider again for each input.
+#[derive(Debug, Clone)]
+pub struct ValidatedModule {
     module: Module,
     std_import: Option<Provider>,
 }
 
 impl ValidatedModule {
-    pub(crate) fn new(module: Module) -> Result<Self> {
+    /// Validates `module` and compiles its standard provider with the module's engine.
+    pub fn new(module: Module) -> 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 +61,20 @@ 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(),
+        // Compile the provider once so every run of this module can reuse it.
+        let std_import = imports
+            .iter()
+            .find_map(|import| {
+                StandardProviders::get(&format!("{import}.wasm")).map(|file| {
+                    Module::from_binary(module.engine(), &file.data)
+                        .map(|module| Provider {
+                            module,
+                            name: import.into(),
+                        })
+                        .map_err(|e| anyhow!("Failed to load provider {import}: {e}"))
+                })
             })
-        });
+            .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