zen_engine/loader/
cached.rs

1use ahash::{HashMap, HashMapExt};
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6
7use crate::loader::{DecisionLoader, DynamicLoader, LoaderResponse};
8use crate::model::DecisionContent;
9
10#[derive(Debug)]
11pub struct CachedLoader {
12    loader: DynamicLoader,
13    cache: Mutex<HashMap<String, Arc<DecisionContent>>>,
14}
15
16impl From<DynamicLoader> for CachedLoader {
17    fn from(value: DynamicLoader) -> Self {
18        Self {
19            loader: value,
20            cache: Mutex::new(HashMap::new()),
21        }
22    }
23}
24
25impl DecisionLoader for CachedLoader {
26    fn load<'a>(
27        &'a self,
28        key: &'a str,
29    ) -> Pin<Box<dyn Future<Output = LoaderResponse> + 'a + Send>> {
30        Box::pin(async move {
31            let mut cache = self.cache.lock().await;
32            if let Some(content) = cache.get(key) {
33                return Ok(content.clone());
34            }
35
36            let decision_content = self.loader.load(key).await?;
37            cache.insert(key.to_string(), decision_content.clone());
38            Ok(decision_content)
39        })
40    }
41
42    fn keys(&self) -> Option<Vec<Arc<str>>> {
43        self.loader.keys()
44    }
45
46    fn load_sync(&self, key: &str) -> Option<LoaderResponse> {
47        let Ok(mut cache) = self.cache.try_lock() else {
48            return self.loader.load_sync(key);
49        };
50        if let Some(content) = cache.get(key) {
51            return Some(Ok(content.clone()));
52        }
53        let response = self.loader.load_sync(key)?;
54        if let Ok(content) = &response {
55            cache.insert(key.to_string(), content.clone());
56        }
57        Some(response)
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64    use crate::loader::MemoryLoader;
65    use crate::model::DecisionContent;
66    use std::sync::atomic::{AtomicUsize, Ordering};
67
68    #[derive(Debug, Default)]
69    struct CountingLoader {
70        inner: MemoryLoader,
71        sync_loads: AtomicUsize,
72    }
73
74    impl DecisionLoader for CountingLoader {
75        fn load<'a>(
76            &'a self,
77            key: &'a str,
78        ) -> Pin<Box<dyn Future<Output = LoaderResponse> + 'a + Send>> {
79            self.inner.load(key)
80        }
81
82        fn load_sync(&self, key: &str) -> Option<LoaderResponse> {
83            self.sync_loads.fetch_add(1, Ordering::SeqCst);
84            self.inner.load_sync(key)
85        }
86    }
87
88    #[test]
89    fn load_sync_uses_cache_and_hits_inner_once() {
90        let counting = Arc::new(CountingLoader::default());
91        counting.inner.add("graph.json", DecisionContent::default());
92        let cached = CachedLoader::from(counting.clone() as DynamicLoader);
93
94        let first = cached.load_sync("graph.json").unwrap().unwrap();
95        let second = cached.load_sync("graph.json").unwrap().unwrap();
96
97        assert!(Arc::ptr_eq(&first, &second));
98        assert_eq!(counting.sync_loads.load(Ordering::SeqCst), 1);
99    }
100
101    #[test]
102    fn delegates_keys_and_load_sync_to_inner_loader() {
103        let memory_loader = MemoryLoader::default();
104        memory_loader.add("graph.json", DecisionContent::default());
105
106        let cached = CachedLoader::from(Arc::new(memory_loader) as DynamicLoader);
107
108        let keys = cached.keys().unwrap();
109        assert_eq!(keys, vec![Arc::from("graph.json")]);
110
111        let content = cached.load_sync("graph.json").unwrap().unwrap();
112        assert!(content.as_graph().is_some());
113
114        assert!(cached.load_sync("missing.json").unwrap().is_err());
115    }
116}
zen_engine/loader/cached.rs

zen_engine/loader/
cached.rs
