mithril_common/digesters/cache/
memory_provider.rs

1use crate::{
2    digesters::cache::CacheProviderResult,
3    digesters::cache::ImmutableFileDigestCacheProvider,
4    digesters::ImmutableFile,
5    entities::{HexEncodedDigest, ImmutableFileName},
6};
7
8use async_trait::async_trait;
9use std::collections::{BTreeMap, HashMap};
10use tokio::sync::RwLock;
11
12/// A in memory [ImmutableFileDigestCacheProvider].
13pub struct MemoryImmutableFileDigestCacheProvider {
14    store: RwLock<HashMap<ImmutableFileName, HexEncodedDigest>>,
15}
16
17impl MemoryImmutableFileDigestCacheProvider {
18    /// Build a new [MemoryImmutableFileDigestCacheProvider] that contains the given values.
19    pub fn from(values: HashMap<ImmutableFileName, HexEncodedDigest>) -> Self {
20        Self {
21            store: RwLock::new(values),
22        }
23    }
24}
25
26impl Default for MemoryImmutableFileDigestCacheProvider {
27    fn default() -> Self {
28        Self {
29            store: RwLock::new(HashMap::new()),
30        }
31    }
32}
33
34#[async_trait]
35impl ImmutableFileDigestCacheProvider for MemoryImmutableFileDigestCacheProvider {
36    async fn store(
37        &self,
38        digest_per_filenames: Vec<(ImmutableFileName, HexEncodedDigest)>,
39    ) -> CacheProviderResult<()> {
40        let mut store = self.store.write().await;
41        for (filename, digest) in digest_per_filenames {
42            store.insert(filename, digest);
43        }
44
45        Ok(())
46    }
47
48    async fn get(
49        &self,
50        immutables: Vec<ImmutableFile>,
51    ) -> CacheProviderResult<BTreeMap<ImmutableFile, Option<HexEncodedDigest>>> {
52        let store = self.store.read().await;
53        let mut result = BTreeMap::new();
54
55        for immutable in immutables {
56            let value = store.get(&immutable.filename).map(|f| f.to_owned());
57            result.insert(immutable, value);
58        }
59
60        Ok(result)
61    }
62
63    async fn reset(&self) -> CacheProviderResult<()> {
64        let mut store = self.store.write().await;
65        store.clear();
66        Ok(())
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use crate::{
73        digesters::cache::{
74            ImmutableFileDigestCacheProvider, MemoryImmutableFileDigestCacheProvider,
75        },
76        digesters::ImmutableFile,
77    };
78    use std::collections::{BTreeMap, HashMap};
79    use std::path::PathBuf;
80
81    #[tokio::test]
82    async fn can_store_values() {
83        let provider = MemoryImmutableFileDigestCacheProvider::default();
84        let values_to_store = vec![
85            ("0.chunk".to_string(), "digest 0".to_string()),
86            ("1.chunk".to_string(), "digest 1".to_string()),
87        ];
88        let expected: BTreeMap<_, _> = BTreeMap::from([
89            (
90                ImmutableFile::dummy(PathBuf::default(), 0, "0.chunk"),
91                Some("digest 0".to_string()),
92            ),
93            (
94                ImmutableFile::dummy(PathBuf::default(), 1, "1.chunk"),
95                Some("digest 1".to_string()),
96            ),
97        ]);
98        let immutables = expected.keys().cloned().collect();
99
100        provider
101            .store(values_to_store)
102            .await
103            .expect("Cache write should not fail");
104        let result = provider
105            .get(immutables)
106            .await
107            .expect("Cache read should not fail");
108
109        assert_eq!(expected, result);
110    }
111
112    #[tokio::test]
113    async fn returns_only_asked_immutables_cache() {
114        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([
115            ("0.chunk".to_string(), "digest 0".to_string()),
116            ("1.chunk".to_string(), "digest 1".to_string()),
117        ]));
118        let expected: BTreeMap<_, _> = BTreeMap::from([(
119            ImmutableFile::dummy(PathBuf::default(), 0, "0.chunk"),
120            Some("digest 0".to_string()),
121        )]);
122        let immutables = expected.keys().cloned().collect();
123
124        let result = provider
125            .get(immutables)
126            .await
127            .expect("Cache read should not fail");
128
129        assert_eq!(expected, result);
130    }
131
132    #[tokio::test]
133    async fn returns_none_for_uncached_asked_immutables() {
134        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([(
135            "0.chunk".to_string(),
136            "digest 0".to_string(),
137        )]));
138        let expected: BTreeMap<_, _> =
139            BTreeMap::from([(ImmutableFile::dummy(PathBuf::default(), 2, "2.chunk"), None)]);
140        let immutables = expected.keys().cloned().collect();
141
142        let result = provider
143            .get(immutables)
144            .await
145            .expect("Cache read should not fail");
146
147        assert_eq!(expected, result);
148    }
149
150    #[tokio::test]
151    async fn store_erase_existing_values() {
152        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([
153            ("0.chunk".to_string(), "to erase".to_string()),
154            ("1.chunk".to_string(), "keep me".to_string()),
155            ("2.chunk".to_string(), "keep me too".to_string()),
156        ]));
157        let values_to_store = vec![
158            ("0.chunk".to_string(), "updated".to_string()),
159            ("1.chunk".to_string(), "keep me".to_string()),
160        ];
161        let expected: BTreeMap<_, _> = BTreeMap::from([
162            (
163                ImmutableFile::dummy(PathBuf::default(), 0, "0.chunk"),
164                Some("updated".to_string()),
165            ),
166            (
167                ImmutableFile::dummy(PathBuf::default(), 1, "1.chunk"),
168                Some("keep me".to_string()),
169            ),
170            (
171                ImmutableFile::dummy(PathBuf::default(), 2, "2.chunk"),
172                Some("keep me too".to_string()),
173            ),
174            (ImmutableFile::dummy(PathBuf::default(), 3, "3.chunk"), None),
175        ]);
176        let immutables = expected.keys().cloned().collect();
177
178        provider
179            .store(values_to_store)
180            .await
181            .expect("Cache write should not fail");
182        let result = provider
183            .get(immutables)
184            .await
185            .expect("Cache read should not fail");
186
187        assert_eq!(expected, result);
188    }
189
190    #[tokio::test]
191    async fn reset_clear_existing_values() {
192        let provider = MemoryImmutableFileDigestCacheProvider::default();
193        let values_to_store = vec![
194            ("0.chunk".to_string(), "digest 0".to_string()),
195            ("1.chunk".to_string(), "digest 1".to_string()),
196        ];
197        let expected: BTreeMap<_, _> = BTreeMap::from([
198            (
199                ImmutableFile::dummy(PathBuf::default(), 0, "0.chunk"),
200                Some("digest 0".to_string()),
201            ),
202            (
203                ImmutableFile::dummy(PathBuf::default(), 1, "1.chunk"),
204                Some("digest 1".to_string()),
205            ),
206        ]);
207        let immutables = expected.keys().cloned().collect();
208
209        provider
210            .store(values_to_store)
211            .await
212            .expect("Cache write should not fail");
213        provider.reset().await.expect("reset should not fails");
214
215        let result: BTreeMap<_, _> = provider
216            .get(immutables)
217            .await
218            .expect("Cache read should not fail");
219
220        assert!(result.into_iter().all(|(_, cache)| cache.is_none()));
221    }
222}
mithril_common/digesters/cache/memory_provider.rs

mithril_common/digesters/cache/
memory_provider.rs
