zen_engine/loader/
cached.rs

1use ahash::{HashMap, HashMapExt};
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::{Arc, RwLock};
5
6use crate::loader::{DecisionLoader, DynamicLoader, LoaderResponse};
7use crate::model::DecisionContent;
8
9#[derive(Debug)]
10pub struct CachedLoader {
11    loader: DynamicLoader,
12    cache: RwLock<HashMap<String, Arc<DecisionContent>>>,
13}
14
15impl From<DynamicLoader> for CachedLoader {
16    fn from(value: DynamicLoader) -> Self {
17        Self {
18            loader: value,
19            cache: RwLock::new(HashMap::new()),
20        }
21    }
22}
23
24fn compiled(content: Arc<DecisionContent>) -> Arc<DecisionContent> {
25    let DecisionContent::Graph(graph) = content.as_ref() else {
26        return content;
27    };
28    if graph.compiled_cache.is_some() {
29        return content;
30    }
31
32    let mut owned = (**graph).clone();
33    owned.compile();
34    Arc::new(DecisionContent::Graph(Arc::new(owned)))
35}
36
37async fn prepared(loader: &DynamicLoader, content: Arc<DecisionContent>) -> Arc<DecisionContent> {
38    let DecisionContent::Graph(graph) = content.as_ref() else {
39        return content;
40    };
41    if graph.compiled_cache.is_some() && graph.resolved_schemas.is_some() {
42        return content;
43    }
44
45    let mut owned = (**graph).clone();
46    owned.compile();
47    let _ = owned.resolve_schemas(loader).await;
48    Arc::new(DecisionContent::Graph(Arc::new(owned)))
49}
50
51impl DecisionLoader for CachedLoader {
52    fn load<'a>(
53        &'a self,
54        key: &'a str,
55    ) -> Pin<Box<dyn Future<Output = LoaderResponse> + 'a + Send>> {
56        Box::pin(async move {
57            let cached = self
58                .cache
59                .read()
60                .ok()
61                .and_then(|cache| cache.get(key).cloned());
62
63            let loaded = match &cached {
64                Some(content) => content.clone(),
65                None => self.loader.load(key).await?,
66            };
67
68            let decision_content = prepared(&self.loader, loaded).await;
69            let unchanged = cached
70                .as_ref()
71                .is_some_and(|content| Arc::ptr_eq(content, &decision_content));
72            if !unchanged {
73                if let Ok(mut cache) = self.cache.write() {
74                    cache.insert(key.to_string(), decision_content.clone());
75                }
76            }
77            Ok(decision_content)
78        })
79    }
80
81    fn keys(&self) -> Option<Vec<Arc<str>>> {
82        self.loader.keys()
83    }
84
85    fn load_sync(&self, key: &str) -> Option<LoaderResponse> {
86        if let Ok(cache) = self.cache.read() {
87            if let Some(content) = cache.get(key) {
88                return Some(Ok(content.clone()));
89            }
90        }
91
92        let response = self.loader.load_sync(key)?.map(compiled);
93        if let Ok(content) = &response {
94            if let Ok(mut cache) = self.cache.write() {
95                cache.insert(key.to_string(), content.clone());
96            }
97        }
98        Some(response)
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use crate::loader::MemoryLoader;
106    use crate::model::DecisionContent;
107    use std::sync::atomic::{AtomicUsize, Ordering};
108
109    #[derive(Debug, Default)]
110    struct CountingLoader {
111        inner: MemoryLoader,
112        sync_loads: AtomicUsize,
113    }
114
115    impl DecisionLoader for CountingLoader {
116        fn load<'a>(
117            &'a self,
118            key: &'a str,
119        ) -> Pin<Box<dyn Future<Output = LoaderResponse> + 'a + Send>> {
120            self.inner.load(key)
121        }
122
123        fn load_sync(&self, key: &str) -> Option<LoaderResponse> {
124            self.sync_loads.fetch_add(1, Ordering::SeqCst);
125            self.inner.load_sync(key)
126        }
127    }
128
129    #[test]
130    fn load_sync_uses_cache_and_hits_inner_once() {
131        let counting = Arc::new(CountingLoader::default());
132        counting.inner.add("graph.json", DecisionContent::default());
133        let cached = CachedLoader::from(counting.clone() as DynamicLoader);
134
135        let first = cached.load_sync("graph.json").unwrap().unwrap();
136        let second = cached.load_sync("graph.json").unwrap().unwrap();
137
138        assert!(Arc::ptr_eq(&first, &second));
139        assert_eq!(counting.sync_loads.load(Ordering::SeqCst), 1);
140    }
141
142    #[test]
143    fn delegates_keys_and_load_sync_to_inner_loader() {
144        let memory_loader = MemoryLoader::default();
145        memory_loader.add("graph.json", DecisionContent::default());
146
147        let cached = CachedLoader::from(Arc::new(memory_loader) as DynamicLoader);
148
149        let keys = cached.keys().unwrap();
150        assert_eq!(keys, vec![Arc::from("graph.json")]);
151
152        let content = cached.load_sync("graph.json").unwrap().unwrap();
153        assert!(content.as_graph().is_some());
154
155        assert!(cached.load_sync("missing.json").unwrap().is_err());
156    }
157}
zen_engine/loader/cached.rs

zen_engine/loader/
cached.rs
