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.get(immutables).await.expect("Cache read should not fail");
104
105        assert_eq!(expected, result);
106    }
107
108    #[tokio::test]
109    async fn returns_only_asked_immutables_cache() {
110        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([
111            ("0.chunk".to_string(), "digest 0".to_string()),
112            ("1.chunk".to_string(), "digest 1".to_string()),
113        ]));
114        let expected: BTreeMap<_, _> = BTreeMap::from([(
115            fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
116            Some("digest 0".to_string()),
117        )]);
118        let immutables = expected.keys().cloned().collect();
119
120        let result = provider.get(immutables).await.expect("Cache read should not fail");
121
122        assert_eq!(expected, result);
123    }
124
125    #[tokio::test]
126    async fn returns_none_for_uncached_asked_immutables() {
127        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([(
128            "0.chunk".to_string(),
129            "digest 0".to_string(),
130        )]));
131        let expected: BTreeMap<_, _> = BTreeMap::from([(
132            fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
133            None,
134        )]);
135        let immutables = expected.keys().cloned().collect();
136
137        let result = provider.get(immutables).await.expect("Cache read should not fail");
138
139        assert_eq!(expected, result);
140    }
141
142    #[tokio::test]
143    async fn store_erase_existing_values() {
144        let provider = MemoryImmutableFileDigestCacheProvider::from(HashMap::from([
145            ("0.chunk".to_string(), "to erase".to_string()),
146            ("1.chunk".to_string(), "keep me".to_string()),
147            ("2.chunk".to_string(), "keep me too".to_string()),
148        ]));
149        let values_to_store = vec![
150            ("0.chunk".to_string(), "updated".to_string()),
151            ("1.chunk".to_string(), "keep me".to_string()),
152        ];
153        let expected: BTreeMap<_, _> = BTreeMap::from([
154            (
155                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
156                Some("updated".to_string()),
157            ),
158            (
159                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
160                Some("keep me".to_string()),
161            ),
162            (
163                fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
164                Some("keep me too".to_string()),
165            ),
166            (
167                fake_data::immutable_file(PathBuf::default(), 3, "3.chunk"),
168                None,
169            ),
170        ]);
171        let immutables = expected.keys().cloned().collect();
172
173        provider
174            .store(values_to_store)
175            .await
176            .expect("Cache write should not fail");
177        let result = provider.get(immutables).await.expect("Cache read should not fail");
178
179        assert_eq!(expected, result);
180    }
181
182    #[tokio::test]
183    async fn reset_clear_existing_values() {
184        let provider = MemoryImmutableFileDigestCacheProvider::default();
185        let values_to_store = vec![
186            ("0.chunk".to_string(), "digest 0".to_string()),
187            ("1.chunk".to_string(), "digest 1".to_string()),
188        ];
189        let expected: BTreeMap<_, _> = BTreeMap::from([
190            (
191                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
192                Some("digest 0".to_string()),
193            ),
194            (
195                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
196                Some("digest 1".to_string()),
197            ),
198        ]);
199        let immutables = expected.keys().cloned().collect();
200
201        provider
202            .store(values_to_store)
203            .await
204            .expect("Cache write should not fail");
205        provider.reset().await.expect("reset should not fails");
206
207        let result: BTreeMap<_, _> =
208            provider.get(immutables).await.expect("Cache read should not fail");
209
210        assert!(result.into_iter().all(|(_, cache)| cache.is_none()));
211    }
212}