mithril_cardano_node_internal_database/digesters/cache/
memory_provider.rs

1use async_trait::async_trait;
2use std::collections::{BTreeMap, HashMap};
3use tokio::sync::RwLock;
4
5use mithril_common::entities::{HexEncodedDigest, ImmutableFileName};
6
7use crate::entities::ImmutableFile;
8use crate::{
9    digesters::cache::CacheProviderResult, digesters::cache::ImmutableFileDigestCacheProvider,
10};
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 std::collections::{BTreeMap, HashMap};
73    use std::path::PathBuf;
74
75    use crate::digesters::cache::{
76        ImmutableFileDigestCacheProvider, MemoryImmutableFileDigestCacheProvider,
77    };
78    use crate::test::fake_data;
79
80    #[tokio::test]
81    async fn can_store_values() {
82        let provider = MemoryImmutableFileDigestCacheProvider::default();
83        let values_to_store = vec![
84            ("0.chunk".to_string(), "digest 0".to_string()),
85            ("1.chunk".to_string(), "digest 1".to_string()),
86        ];
87        let expected: BTreeMap<_, _> = BTreeMap::from([
88            (
89                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
90                Some("digest 0".to_string()),
91            ),
92            (
93                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
94                Some("digest 1".to_string()),
95            ),
96        ]);
97        let immutables = expected.keys().cloned().collect();
98
99        provider
100            .store(values_to_store)
101            .await
102            .expect("Cache write should not fail");
103        let result = provider
104            .get(immutables)
105            .await
106            .expect("Cache read should not fail");
107
108        assert_eq!(expected, result);
109    }
110
111    #[tokio::test]
112    async fn returns_only_asked_immutables_cache() {
113        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([
114            ("0.chunk".to_string(), "digest 0".to_string()),
115            ("1.chunk".to_string(), "digest 1".to_string()),
116        ]));
117        let expected: BTreeMap<_, _> = BTreeMap::from([(
118            fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
119            Some("digest 0".to_string()),
120        )]);
121        let immutables = expected.keys().cloned().collect();
122
123        let result = provider
124            .get(immutables)
125            .await
126            .expect("Cache read should not fail");
127
128        assert_eq!(expected, result);
129    }
130
131    #[tokio::test]
132    async fn returns_none_for_uncached_asked_immutables() {
133        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([(
134            "0.chunk".to_string(),
135            "digest 0".to_string(),
136        )]));
137        let expected: BTreeMap<_, _> = BTreeMap::from([(
138            fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
139            None,
140        )]);
141        let immutables = expected.keys().cloned().collect();
142
143        let result = provider
144            .get(immutables)
145            .await
146            .expect("Cache read should not fail");
147
148        assert_eq!(expected, result);
149    }
150
151    #[tokio::test]
152    async fn store_erase_existing_values() {
153        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([
154            ("0.chunk".to_string(), "to erase".to_string()),
155            ("1.chunk".to_string(), "keep me".to_string()),
156            ("2.chunk".to_string(), "keep me too".to_string()),
157        ]));
158        let values_to_store = vec![
159            ("0.chunk".to_string(), "updated".to_string()),
160            ("1.chunk".to_string(), "keep me".to_string()),
161        ];
162        let expected: BTreeMap<_, _> = BTreeMap::from([
163            (
164                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
165                Some("updated".to_string()),
166            ),
167            (
168                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
169                Some("keep me".to_string()),
170            ),
171            (
172                fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
173                Some("keep me too".to_string()),
174            ),
175            (
176                fake_data::immutable_file(PathBuf::default(), 3, "3.chunk"),
177                None,
178            ),
179        ]);
180        let immutables = expected.keys().cloned().collect();
181
182        provider
183            .store(values_to_store)
184            .await
185            .expect("Cache write should not fail");
186        let result = provider
187            .get(immutables)
188            .await
189            .expect("Cache read should not fail");
190
191        assert_eq!(expected, result);
192    }
193
194    #[tokio::test]
195    async fn reset_clear_existing_values() {
196        let provider = MemoryImmutableFileDigestCacheProvider::default();
197        let values_to_store = vec![
198            ("0.chunk".to_string(), "digest 0".to_string()),
199            ("1.chunk".to_string(), "digest 1".to_string()),
200        ];
201        let expected: BTreeMap<_, _> = BTreeMap::from([
202            (
203                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
204                Some("digest 0".to_string()),
205            ),
206            (
207                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
208                Some("digest 1".to_string()),
209            ),
210        ]);
211        let immutables = expected.keys().cloned().collect();
212
213        provider
214            .store(values_to_store)
215            .await
216            .expect("Cache write should not fail");
217        provider.reset().await.expect("reset should not fails");
218
219        let result: BTreeMap<_, _> = provider
220            .get(immutables)
221            .await
222            .expect("Cache read should not fail");
223
224        assert!(result.into_iter().all(|(_, cache)| cache.is_none()));
225    }
226}