Skip to main content

wasmer_wasix/runtime/module_cache/
shared.rs

1use dashmap::DashMap;
2use wasmer::{Engine, Module};
3
4use crate::runtime::module_cache::{CacheError, ModuleCache};
5use wasmer_types::ModuleHash;
6
7/// A [`ModuleCache`] based on a <code>[DashMap]</code> keyed by module hash, engine ID, and format.
8#[derive(Debug, Default, Clone)]
9pub struct SharedCache {
10    modules: DashMap<SharedCacheKey, Module>,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14struct SharedCacheKey {
15    module_hash: ModuleHash,
16    deterministic_id: String,
17    artifact_format: String,
18}
19
20impl SharedCache {
21    pub fn new() -> SharedCache {
22        SharedCache::default()
23    }
24
25    fn cache_key(key: ModuleHash, engine: &Engine) -> SharedCacheKey {
26        SharedCacheKey {
27            module_hash: key,
28            deterministic_id: engine.deterministic_id(),
29            artifact_format: engine.artifact_format(),
30        }
31    }
32}
33
34#[async_trait::async_trait]
35impl ModuleCache for SharedCache {
36    #[tracing::instrument(level = "debug", skip_all, fields(%key))]
37    async fn load(&self, key: ModuleHash, engine: &Engine) -> Result<Module, CacheError> {
38        let key = Self::cache_key(key, engine);
39
40        match self.modules.get(&key) {
41            Some(m) => {
42                tracing::debug!("Cache hit!");
43                Ok(m.value().clone())
44            }
45
46            None => Err(CacheError::NotFound),
47        }
48    }
49
50    async fn contains(&self, key: ModuleHash, engine: &Engine) -> Result<bool, CacheError> {
51        let key = Self::cache_key(key, engine);
52        Ok(self.modules.contains_key(&key))
53    }
54
55    #[tracing::instrument(level = "debug", skip_all, fields(%key))]
56    async fn save(
57        &self,
58        key: ModuleHash,
59        engine: &Engine,
60        module: &Module,
61    ) -> Result<(), CacheError> {
62        let key = Self::cache_key(key, engine);
63        self.modules.insert(key, module.clone());
64
65        Ok(())
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    const ADD_WAT: &[u8] = br#"(
74        module
75            (func
76                (export "add")
77                (param $x i64)
78                (param $y i64)
79                (result i64)
80                (i64.add (local.get $x) (local.get $y)))
81        )"#;
82
83    #[tokio::test]
84    async fn round_trip_via_cache() {
85        let engine = Engine::default();
86        let module = Module::new(&engine, ADD_WAT).unwrap();
87        let cache = SharedCache::default();
88        let key = ModuleHash::from_bytes([0; _]);
89
90        cache.save(key, &engine, &module).await.unwrap();
91        let round_tripped = cache.load(key, &engine).await.unwrap();
92
93        let exports: Vec<_> = round_tripped
94            .exports()
95            .map(|export| export.name().to_string())
96            .collect();
97        assert_eq!(exports, ["add"]);
98    }
99}