mithril_cardano_node_internal_database/digesters/cache/
memory_provider.rs1use 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
12pub struct MemoryImmutableFileDigestCacheProvider {
14 store: RwLock<HashMap<ImmutableFileName, HexEncodedDigest>>,
15}
16
17impl MemoryImmutableFileDigestCacheProvider {
18 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}