From 7962a53cbf227b4ecc8b841194301045cc11e6f3 Mon Sep 17 00:00:00 2001 From: Dave Nagoda Date: Fri, 2 Oct 2026 11:20:11 -0700 Subject: [PATCH 1/2] Reuse compiled standard providers across runs `run` compiled the embedded standard provider (Javy plugin or Shopify Function provider) from its bytes on every call. Even with a warm Wasmtime cache, the Javy providers took about 8 ms per call, so callers that run the same Function many times paid that cost on every run. Keep the compiled providers for the most recent engine in a private cache, and reuse them while callers pass the same engine. A call with a different engine replaces the cache, so it holds one engine at most. The public API does not change. --- src/engine.rs | 43 +++++++++++++++ src/io.rs | 3 +- src/validated_module.rs | 114 ++++++++++++++++++++++++++++++++++++++-- 3 files changed, 156 insertions(+), 4 deletions(-) diff --git a/src/engine.rs b/src/engine.rs index 9185552a..6f120a93 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -321,6 +321,49 @@ mod tests { Ok(()) } + #[test] + fn test_repeated_runs_match_fresh_runs() -> Result<()> { + for path in [ + "tests/fixtures/build/js_function.wasm", + "tests/fixtures/build/js_function_javy_plugin_v3.wasm", + ] + .map(Path::new) + { + let run_with = |engine: &Engine| -> Result { + let module = Module::from_file(engine, path)?; + let codec = if uses_msgpack_provider(&module) { + Codec::Messagepack + } else { + Codec::Json + }; + let input = BytesContainer::new( + BytesContainerType::Input, + codec, + include_bytes!("../tests/fixtures/input/js_function_input.json").to_vec(), + )?; + let result = run(FunctionRunParams { + function_path: path.to_path_buf(), + input, + export: DEFAULT_EXPORT, + module, + engine: engine.clone(), + scale_factor: 1.0, + profile_opts: None, + })?; + assert!(result.success); + Ok(serde_json::to_value(result)?) + }; + + let engine = new_engine()?; + let first = run_with(&engine)?; + assert_eq!(first, run_with(&engine)?); + assert_eq!(first, run_with(&new_engine()?)?); + assert_eq!(first, run_with(&engine)?); + } + + Ok(()) + } + #[test] fn test_exit_code_zero() -> Result<()> { let engine = new_engine()?; diff --git a/src/io.rs b/src/io.rs index 596cb9cb..ce2ea671 100644 --- a/src/io.rs +++ b/src/io.rs @@ -164,7 +164,8 @@ fn instantiate_imports( let mut mem_io_instance = None; if let Some(std_import) = module.std_import() { - let imported_module = Module::from_binary(engine, &std_import.bytes) + let imported_module = std_import + .module(engine) .unwrap_or_else(|_| panic!("Failed to load module {}", std_import.name)); let imported_module_instance = linker diff --git a/src/validated_module.rs b/src/validated_module.rs index 2d4b40e9..f1b0c062 100644 --- a/src/validated_module.rs +++ b/src/validated_module.rs @@ -1,13 +1,52 @@ -use std::borrow::Cow; +use std::{ + borrow::Cow, + collections::HashMap, + sync::{Mutex, PoisonError}, +}; use anyhow::{bail, Result}; use rust_embed::RustEmbed; -use wasmtime::Module; +use wasmtime::{Engine, Module}; #[derive(RustEmbed)] #[folder = "providers/"] struct StandardProviders; +static COMPILED_PROVIDERS: Mutex = Mutex::new(ProviderCache::new()); + +struct CompiledProviders { + engine: Engine, + modules: HashMap, +} + +struct ProviderCache(Option); + +impl ProviderCache { + const fn new() -> Self { + Self(None) + } + + fn module(&mut self, engine: &Engine, provider: &Provider) -> Result { + let compiled = match &mut self.0 { + Some(compiled) if Engine::same(&compiled.engine, engine) => compiled, + cache => cache.insert(CompiledProviders { + engine: engine.clone(), + modules: HashMap::new(), + }), + }; + + if let Some(module) = compiled.modules.get(&provider.name) { + return Ok(module.clone()); + } + + let module = Module::from_binary(engine, &provider.bytes)?; + compiled + .modules + .insert(provider.name.clone(), module.clone()); + Ok(module) + } +} + #[derive(Debug)] pub(crate) struct Provider { pub(crate) bytes: Cow<'static, [u8]>, @@ -15,6 +54,13 @@ pub(crate) struct Provider { } impl Provider { + pub(crate) fn module(&self, engine: &Engine) -> Result { + COMPILED_PROVIDERS + .lock() + .unwrap_or_else(PoisonError::into_inner) + .module(engine, self) + } + pub(crate) fn is_mem_io_provider(&self) -> bool { let javy_plugin_version = self .name @@ -96,7 +142,69 @@ mod tests { use anyhow::Result; use wasmtime::{Engine, Module}; - use crate::validated_module::ValidatedModule; + use crate::validated_module::{Provider, ProviderCache, StandardProviders, ValidatedModule}; + + fn provider(name: &str) -> Provider { + Provider { + bytes: StandardProviders::get(&format!("{name}.wasm")) + .unwrap() + .data, + name: name.into(), + } + } + + #[test] + fn test_provider_cache_reuses_module_for_same_engine() -> Result<()> { + let engine = Engine::default(); + let provider = provider("shopify_function_v2"); + let mut cache = ProviderCache::new(); + + let first = cache.module(&engine, &provider)?; + let second = cache.module(&engine.clone(), &provider)?; + + assert_eq!(first.image_range(), second.image_range()); + Ok(()) + } + + #[test] + fn test_provider_cache_keeps_each_provider() -> Result<()> { + let engine = Engine::default(); + let v1 = provider("shopify_function_v1"); + let v2 = provider("shopify_function_v2"); + let mut cache = ProviderCache::new(); + + let first_v1 = cache.module(&engine, &v1)?; + let first_v2 = cache.module(&engine, &v2)?; + + assert_ne!(first_v1.image_range(), first_v2.image_range()); + assert_eq!( + first_v1.image_range(), + cache.module(&engine, &v1)?.image_range() + ); + assert_eq!( + first_v2.image_range(), + cache.module(&engine, &v2)?.image_range() + ); + Ok(()) + } + + #[test] + fn test_provider_cache_holds_only_most_recent_engine() -> Result<()> { + let first_engine = Engine::default(); + let second_engine = Engine::default(); + let provider = provider("shopify_function_v2"); + let mut cache = ProviderCache::new(); + + let first = cache.module(&first_engine, &provider)?; + let second = cache.module(&second_engine, &provider)?; + assert!(Engine::same(second.engine(), &second_engine)); + assert_ne!(first.image_range(), second.image_range()); + + let first_again = cache.module(&first_engine, &provider)?; + assert!(Engine::same(first_again.engine(), &first_engine)); + assert_ne!(first.image_range(), first_again.image_range()); + Ok(()) + } #[test] fn test_module_with_just_wasi() -> Result<()> { From ee0fd237f3e05435ef47d6c619e69658a25ce151 Mon Sep 17 00:00:00 2001 From: Dave Nagoda Date: Fri, 2 Oct 2026 17:40:44 -0700 Subject: [PATCH 2/2] Compile providers outside the provider cache lock The provider cache held its mutex while Module::from_binary compiled a cold provider, so every other lookup waited, including warm cache hits. Look up the provider under the lock, compile it without the lock, and take the lock again to insert it. If the cache moved to another engine during compilation, return the module without caching it, so the cache still holds one engine at most. --- src/validated_module.rs | 148 +++++++++++++++++++++++++++++++--------- 1 file changed, 115 insertions(+), 33 deletions(-) diff --git a/src/validated_module.rs b/src/validated_module.rs index f1b0c062..5224e1a6 100644 --- a/src/validated_module.rs +++ b/src/validated_module.rs @@ -26,25 +26,47 @@ impl ProviderCache { Self(None) } - fn module(&mut self, engine: &Engine, provider: &Provider) -> Result { - let compiled = match &mut self.0 { - Some(compiled) if Engine::same(&compiled.engine, engine) => compiled, - cache => cache.insert(CompiledProviders { - engine: engine.clone(), - modules: HashMap::new(), - }), - }; - - if let Some(module) = compiled.modules.get(&provider.name) { - return Ok(module.clone()); + fn get(&mut self, engine: &Engine, name: &str) -> Option { + match &mut self.0 { + Some(compiled) if Engine::same(&compiled.engine, engine) => { + compiled.modules.get(name).cloned() + } + cache => { + *cache = Some(CompiledProviders { + engine: engine.clone(), + modules: HashMap::new(), + }); + None + } + } + } + + fn insert(&mut self, engine: &Engine, name: &str, module: Module) -> Module { + match &mut self.0 { + Some(compiled) if Engine::same(&compiled.engine, engine) => compiled + .modules + .entry(name.to_string()) + .or_insert(module) + .clone(), + _ => module, } + } +} + +fn cached_module( + cache: &Mutex, + engine: &Engine, + name: &str, + compile: impl FnOnce() -> Result, +) -> Result { + let lock = || cache.lock().unwrap_or_else(PoisonError::into_inner); - let module = Module::from_binary(engine, &provider.bytes)?; - compiled - .modules - .insert(provider.name.clone(), module.clone()); - Ok(module) + if let Some(module) = lock().get(engine, name) { + return Ok(module); } + + let module = compile()?; + Ok(lock().insert(engine, name, module)) } #[derive(Debug)] @@ -55,10 +77,9 @@ pub(crate) struct Provider { impl Provider { pub(crate) fn module(&self, engine: &Engine) -> Result { - COMPILED_PROVIDERS - .lock() - .unwrap_or_else(PoisonError::into_inner) - .module(engine, self) + cached_module(&COMPILED_PROVIDERS, engine, &self.name, || { + Module::from_binary(engine, &self.bytes) + }) } pub(crate) fn is_mem_io_provider(&self) -> bool { @@ -139,10 +160,14 @@ impl ValidatedModule { #[cfg(test)] mod tests { + use std::sync::Mutex; + use anyhow::Result; use wasmtime::{Engine, Module}; - use crate::validated_module::{Provider, ProviderCache, StandardProviders, ValidatedModule}; + use crate::validated_module::{ + cached_module, Provider, ProviderCache, StandardProviders, ValidatedModule, + }; fn provider(name: &str) -> Provider { Provider { @@ -153,14 +178,24 @@ mod tests { } } + fn module( + cache: &Mutex, + engine: &Engine, + provider: &Provider, + ) -> Result { + cached_module(cache, engine, &provider.name, || { + Module::from_binary(engine, &provider.bytes) + }) + } + #[test] fn test_provider_cache_reuses_module_for_same_engine() -> Result<()> { let engine = Engine::default(); let provider = provider("shopify_function_v2"); - let mut cache = ProviderCache::new(); + let cache = Mutex::new(ProviderCache::new()); - let first = cache.module(&engine, &provider)?; - let second = cache.module(&engine.clone(), &provider)?; + let first = module(&cache, &engine, &provider)?; + let second = module(&cache, &engine.clone(), &provider)?; assert_eq!(first.image_range(), second.image_range()); Ok(()) @@ -171,19 +206,19 @@ mod tests { let engine = Engine::default(); let v1 = provider("shopify_function_v1"); let v2 = provider("shopify_function_v2"); - let mut cache = ProviderCache::new(); + let cache = Mutex::new(ProviderCache::new()); - let first_v1 = cache.module(&engine, &v1)?; - let first_v2 = cache.module(&engine, &v2)?; + let first_v1 = module(&cache, &engine, &v1)?; + let first_v2 = module(&cache, &engine, &v2)?; assert_ne!(first_v1.image_range(), first_v2.image_range()); assert_eq!( first_v1.image_range(), - cache.module(&engine, &v1)?.image_range() + module(&cache, &engine, &v1)?.image_range() ); assert_eq!( first_v2.image_range(), - cache.module(&engine, &v2)?.image_range() + module(&cache, &engine, &v2)?.image_range() ); Ok(()) } @@ -193,19 +228,66 @@ mod tests { let first_engine = Engine::default(); let second_engine = Engine::default(); let provider = provider("shopify_function_v2"); - let mut cache = ProviderCache::new(); + let cache = Mutex::new(ProviderCache::new()); - let first = cache.module(&first_engine, &provider)?; - let second = cache.module(&second_engine, &provider)?; + let first = module(&cache, &first_engine, &provider)?; + let second = module(&cache, &second_engine, &provider)?; assert!(Engine::same(second.engine(), &second_engine)); assert_ne!(first.image_range(), second.image_range()); - let first_again = cache.module(&first_engine, &provider)?; + let first_again = module(&cache, &first_engine, &provider)?; assert!(Engine::same(first_again.engine(), &first_engine)); assert_ne!(first.image_range(), first_again.image_range()); Ok(()) } + #[test] + fn test_provider_cache_serves_warm_lookups_during_compilation() -> Result<()> { + let engine = Engine::default(); + let v1 = provider("shopify_function_v1"); + let v2 = provider("shopify_function_v2"); + let cache = Mutex::new(ProviderCache::new()); + let warm_v1 = module(&cache, &engine, &v1)?; + + cached_module(&cache, &engine, &v2.name, || { + let during_compile = cache + .try_lock() + .expect("compilation holds the cache lock") + .get(&engine, &v1.name) + .expect("warm provider is cached"); + assert_eq!(warm_v1.image_range(), during_compile.image_range()); + Module::from_binary(&engine, &v2.bytes) + })?; + + Ok(()) + } + + #[test] + fn test_provider_cache_skips_insert_when_engine_changes_during_compilation() -> Result<()> { + let first_engine = Engine::default(); + let second_engine = Engine::default(); + let provider = provider("shopify_function_v2"); + let cache = Mutex::new(ProviderCache::new()); + let mut second = None; + + let first = cached_module(&cache, &first_engine, &provider.name, || { + second = Some(module(&cache, &second_engine, &provider)?); + Module::from_binary(&first_engine, &provider.bytes) + })?; + let second = second.unwrap(); + + assert!(Engine::same(first.engine(), &first_engine)); + assert_eq!( + second.image_range(), + module(&cache, &second_engine, &provider)?.image_range() + ); + assert_ne!( + first.image_range(), + module(&cache, &first_engine, &provider)?.image_range() + ); + Ok(()) + } + #[test] fn test_module_with_just_wasi() -> Result<()> { let wat = r#"