cuttlefish_host/
module_cache.rs

1//! A content-hash-keyed cache of compiled `wasmtime::Module`s.
2//!
3//! `wasmtime::Module::new` measurably costs ~1.5s for a ~3.3MB module (the
4//! shared Rhai interpreter this feature adds) vs. ~125ms for a small
5//! example block — over 10x — and today nothing in this codebase caches a
6//! compiled module at all: `Guest::new` recompiles from scratch on every
7//! single job run. Every `Script`-kind node, across every spec and every
8//! job ever run against it, shares byte-identical `module_bytes` (the one
9//! embedded interpreter), so this cache turns an otherwise-repeated ~1.5s
10//! tax into a one-time cost per process lifetime.
11
12use std::collections::HashMap;
13use std::sync::{Arc, Mutex};
14
15/// **Load-bearing invariant this cache does not check**: every `compile`
16/// call through one `ModuleCache` must pass the *same* `wasmtime::Engine`
17/// every time. A `Module` is only valid for the `Engine` that compiled it —
18/// mixing engines through one cache would silently return a `Module`
19/// compiled for the wrong `Engine`. This codebase constructs exactly one
20/// `Engine` per process (`cuttlefishd`'s and `cuttlefish build`'s `main`),
21/// so a `ModuleCache` constructed once alongside it, and never shared
22/// across processes, upholds this automatically. A test that constructs
23/// its own throwaway `Engine` should also construct its own throwaway
24/// `ModuleCache` — never reuse one across two different `Engine`s.
25pub struct ModuleCache {
26    modules: Mutex<HashMap<String, Arc<wasmtime::Module>>>,
27}
28
29impl ModuleCache {
30    /// A fresh, empty cache.
31    pub fn new() -> Self {
32        Self {
33            modules: Mutex::new(HashMap::new()),
34        }
35    }
36
37    /// Compile `module_bytes` against `engine`, reusing a cached
38    /// compilation if these exact bytes were compiled before through this
39    /// same cache.
40    pub fn compile(
41        &self,
42        engine: &wasmtime::Engine,
43        module_bytes: &[u8],
44    ) -> anyhow::Result<Arc<wasmtime::Module>> {
45        use sha2::{Digest, Sha256};
46        let key = crate::hex::encode(Sha256::digest(module_bytes));
47
48        if let Some(cached) = self.modules.lock().unwrap().get(&key) {
49            return Ok(cached.clone());
50        }
51
52        let module = Arc::new(wasmtime::Module::new(engine, module_bytes)?);
53        self.modules.lock().unwrap().insert(key, module.clone());
54        Ok(module)
55    }
56}
57
58impl Default for ModuleCache {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    fn trivial_wasm() -> Vec<u8> {
69        wat::parse_str("(module)").unwrap()
70    }
71
72    #[test]
73    fn compiling_the_same_bytes_twice_returns_the_same_arc() {
74        let engine = wasmtime::Engine::default();
75        let cache = ModuleCache::new();
76        let bytes = trivial_wasm();
77
78        let a = cache.compile(&engine, &bytes).unwrap();
79        let b = cache.compile(&engine, &bytes).unwrap();
80        assert!(
81            Arc::ptr_eq(&a, &b),
82            "second call should hit the cache, not recompile"
83        );
84    }
85
86    #[test]
87    fn compiling_different_bytes_returns_different_modules() {
88        let engine = wasmtime::Engine::default();
89        let cache = ModuleCache::new();
90        let a = cache.compile(&engine, &trivial_wasm()).unwrap();
91        let other_wasm = wat::parse_str("(module (func))").unwrap();
92        let b = cache.compile(&engine, &other_wasm).unwrap();
93        assert!(!Arc::ptr_eq(&a, &b));
94    }
95}
cuttlefish_host/module_cache.rs

cuttlefish_host/
module_cache.rs
