mithril_cardano_node_internal_database/digesters/cache/
json_provider.rs

1use async_trait::async_trait;
2use std::{
3    collections::BTreeMap,
4    path::{Path, PathBuf},
5};
6use tokio::{
7    fs,
8    fs::File,
9    io::{AsyncReadExt, AsyncWriteExt},
10};
11
12use mithril_common::entities::{HexEncodedDigest, ImmutableFileName};
13
14use crate::digesters::{
15    cache::provider::{ImmutableDigesterCacheGetError, ImmutableDigesterCacheStoreError},
16    cache::CacheProviderResult,
17    cache::ImmutableFileDigestCacheProvider,
18};
19use crate::entities::ImmutableFile;
20
21type InnerStructure = BTreeMap<ImmutableFileName, HexEncodedDigest>;
22
23/// A in memory [ImmutableFileDigestCacheProvider].
24pub struct JsonImmutableFileDigestCacheProvider {
25    filepath: PathBuf,
26}
27
28impl JsonImmutableFileDigestCacheProvider {
29    /// [JsonImmutableFileDigestCacheProvider] factory
30    pub fn new(filepath: &Path) -> Self {
31        Self {
32            filepath: filepath.to_path_buf(),
33        }
34    }
35
36    #[cfg(test)]
37    /// [Test Only] Build a new [JsonImmutableFileDigestCacheProvider] that contains the given values.
38    pub async fn from(filepath: &Path, values: InnerStructure) -> Self {
39        let provider = Self::new(filepath);
40        provider.write_data(values).await.unwrap();
41        provider
42    }
43
44    async fn write_data(
45        &self,
46        values: InnerStructure,
47    ) -> Result<(), ImmutableDigesterCacheStoreError> {
48        let mut file = File::create(&self.filepath).await?;
49        file.write_all(serde_json::to_string_pretty(&values)?.as_bytes())
50            .await?;
51
52        Ok(())
53    }
54
55    async fn read_data(&self) -> Result<InnerStructure, ImmutableDigesterCacheGetError> {
56        match self.filepath.exists() {
57            true => {
58                let mut file = File::open(&self.filepath).await?;
59                let mut json_string = String::new();
60                file.read_to_string(&mut json_string).await?;
61                let values: InnerStructure = serde_json::from_str(&json_string)?;
62                Ok(values)
63            }
64            false => Ok(BTreeMap::new()),
65        }
66    }
67}
68
69#[async_trait]
70impl ImmutableFileDigestCacheProvider for JsonImmutableFileDigestCacheProvider {
71    async fn store(
72        &self,
73        digest_per_filenames: Vec<(ImmutableFileName, HexEncodedDigest)>,
74    ) -> CacheProviderResult<()> {
75        let mut data = self.read_data().await?;
76        for (filename, digest) in digest_per_filenames {
77            data.insert(filename, digest);
78        }
79        self.write_data(data).await?;
80
81        Ok(())
82    }
83
84    async fn get(
85        &self,
86        immutables: Vec<ImmutableFile>,
87    ) -> CacheProviderResult<BTreeMap<ImmutableFile, Option<HexEncodedDigest>>> {
88        let values = self.read_data().await?;
89        let mut result = BTreeMap::new();
90
91        for immutable in immutables {
92            let value = values.get(&immutable.filename).map(|f| f.to_owned());
93            result.insert(immutable, value);
94        }
95
96        Ok(result)
97    }
98
99    async fn reset(&self) -> CacheProviderResult<()> {
100        fs::remove_file(&self.filepath)
101            .await
102            .map_err(ImmutableDigesterCacheStoreError::from)?;
103
104        Ok(())
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use std::{collections::BTreeMap, path::PathBuf};
111
112    use mithril_common::test_utils::TempDir;
113
114    use crate::digesters::cache::{
115        ImmutableFileDigestCacheProvider, JsonImmutableFileDigestCacheProvider,
116    };
117    use crate::test::fake_data;
118
119    fn get_test_dir(subdir_name: &str) -> PathBuf {
120        TempDir::create("json_digester_cache_provider", subdir_name)
121    }
122
123    #[tokio::test]
124    async fn can_store_values() {
125        let file = get_test_dir("can_store_values").join("immutable-cache-store.json");
126        let provider = JsonImmutableFileDigestCacheProvider::new(&file);
127        let values_to_store = vec![
128            ("0.chunk".to_string(), "digest 0".to_string()),
129            ("1.chunk".to_string(), "digest 1".to_string()),
130        ];
131        let expected: BTreeMap<_, _> = BTreeMap::from([
132            (
133                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
134                Some("digest 0".to_string()),
135            ),
136            (
137                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
138                Some("digest 1".to_string()),
139            ),
140        ]);
141        let immutables = expected.keys().cloned().collect();
142
143        provider
144            .store(values_to_store)
145            .await
146            .expect("Cache write should not fail");
147        let result = provider
148            .get(immutables)
149            .await
150            .expect("Cache read should not fail");
151
152        assert_eq!(expected, result);
153    }
154
155    #[tokio::test]
156    async fn returns_only_asked_immutables_cache() {
157        let file =
158            get_test_dir("returns_only_asked_immutables_cache").join("immutable-cache-store.json");
159        let provider = JsonImmutableFileDigestCacheProvider::from(
160            &file,
161            BTreeMap::from([
162                ("0.chunk".to_string(), "digest 0".to_string()),
163                ("1.chunk".to_string(), "digest 1".to_string()),
164            ]),
165        )
166        .await;
167        let expected: BTreeMap<_, _> = BTreeMap::from([(
168            fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
169            Some("digest 0".to_string()),
170        )]);
171        let immutables = expected.keys().cloned().collect();
172
173        let result = provider
174            .get(immutables)
175            .await
176            .expect("Cache read should not fail");
177
178        assert_eq!(expected, result);
179    }
180
181    #[tokio::test]
182    async fn returns_none_for_uncached_asked_immutables() {
183        let file = get_test_dir("returns_none_for_uncached_asked_immutables")
184            .join("immutable-cache-store.json");
185        let provider = JsonImmutableFileDigestCacheProvider::from(
186            &file,
187            BTreeMap::from([("0.chunk".to_string(), "digest 0".to_string())]),
188        )
189        .await;
190        let expected: BTreeMap<_, _> = BTreeMap::from([(
191            fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
192            None,
193        )]);
194        let immutables = expected.keys().cloned().collect();
195
196        let result = provider
197            .get(immutables)
198            .await
199            .expect("Cache read should not fail");
200
201        assert_eq!(expected, result);
202    }
203
204    #[tokio::test]
205    async fn store_erase_existing_values() {
206        let file = get_test_dir("store_erase_existing_values").join("immutable-cache-store.json");
207        let provider = JsonImmutableFileDigestCacheProvider::from(
208            &file,
209            BTreeMap::from([
210                ("0.chunk".to_string(), "to erase".to_string()),
211                ("1.chunk".to_string(), "keep me".to_string()),
212                ("2.chunk".to_string(), "keep me too".to_string()),
213            ]),
214        )
215        .await;
216        let values_to_store = vec![
217            ("0.chunk".to_string(), "updated".to_string()),
218            ("1.chunk".to_string(), "keep me".to_string()),
219        ];
220        let expected: BTreeMap<_, _> = BTreeMap::from([
221            (
222                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
223                Some("updated".to_string()),
224            ),
225            (
226                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
227                Some("keep me".to_string()),
228            ),
229            (
230                fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
231                Some("keep me too".to_string()),
232            ),
233            (
234                fake_data::immutable_file(PathBuf::default(), 3, "3.chunk"),
235                None,
236            ),
237        ]);
238        let immutables = expected.keys().cloned().collect();
239
240        provider
241            .store(values_to_store)
242            .await
243            .expect("Cache write should not fail");
244        let result = provider
245            .get(immutables)
246            .await
247            .expect("Cache read should not fail");
248
249        assert_eq!(expected, result);
250    }
251
252    #[tokio::test]
253    async fn reset_clear_existing_values() {
254        let file = get_test_dir("reset_clear_existing_values").join("immutable-cache-store.json");
255        let provider = JsonImmutableFileDigestCacheProvider::new(&file);
256        let values_to_store = vec![
257            ("0.chunk".to_string(), "digest 0".to_string()),
258            ("1.chunk".to_string(), "digest 1".to_string()),
259        ];
260        let expected: BTreeMap<_, _> = BTreeMap::from([
261            (
262                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
263                Some("digest 0".to_string()),
264            ),
265            (
266                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
267                Some("digest 1".to_string()),
268            ),
269        ]);
270        let immutables = expected.keys().cloned().collect();
271
272        provider
273            .store(values_to_store)
274            .await
275            .expect("Cache write should not fail");
276        provider.reset().await.expect("reset should not fails");
277
278        let result: BTreeMap<_, _> = provider
279            .get(immutables)
280            .await
281            .expect("Cache read should not fail");
282
283        assert!(result.into_iter().all(|(_, cache)| cache.is_none()));
284    }
285}