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::CacheProviderResult,
16    cache::ImmutableFileDigestCacheProvider,
17    cache::provider::{ImmutableDigesterCacheGetError, ImmutableDigesterCacheStoreError},
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.get(immutables).await.expect("Cache read should not fail");
148
149        assert_eq!(expected, result);
150    }
151
152    #[tokio::test]
153    async fn returns_only_asked_immutables_cache() {
154        let file =
155            get_test_dir("returns_only_asked_immutables_cache").join("immutable-cache-store.json");
156        let provider = JsonImmutableFileDigestCacheProvider::from(
157            &file,
158            BTreeMap::from([
159                ("0.chunk".to_string(), "digest 0".to_string()),
160                ("1.chunk".to_string(), "digest 1".to_string()),
161            ]),
162        )
163        .await;
164        let expected: BTreeMap<_, _> = BTreeMap::from([(
165            fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
166            Some("digest 0".to_string()),
167        )]);
168        let immutables = expected.keys().cloned().collect();
169
170        let result = provider.get(immutables).await.expect("Cache read should not fail");
171
172        assert_eq!(expected, result);
173    }
174
175    #[tokio::test]
176    async fn returns_none_for_uncached_asked_immutables() {
177        let file = get_test_dir("returns_none_for_uncached_asked_immutables")
178            .join("immutable-cache-store.json");
179        let provider = JsonImmutableFileDigestCacheProvider::from(
180            &file,
181            BTreeMap::from([("0.chunk".to_string(), "digest 0".to_string())]),
182        )
183        .await;
184        let expected: BTreeMap<_, _> = BTreeMap::from([(
185            fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
186            None,
187        )]);
188        let immutables = expected.keys().cloned().collect();
189
190        let result = provider.get(immutables).await.expect("Cache read should not fail");
191
192        assert_eq!(expected, result);
193    }
194
195    #[tokio::test]
196    async fn store_erase_existing_values() {
197        let file = get_test_dir("store_erase_existing_values").join("immutable-cache-store.json");
198        let provider = JsonImmutableFileDigestCacheProvider::from(
199            &file,
200            BTreeMap::from([
201                ("0.chunk".to_string(), "to erase".to_string()),
202                ("1.chunk".to_string(), "keep me".to_string()),
203                ("2.chunk".to_string(), "keep me too".to_string()),
204            ]),
205        )
206        .await;
207        let values_to_store = vec![
208            ("0.chunk".to_string(), "updated".to_string()),
209            ("1.chunk".to_string(), "keep me".to_string()),
210        ];
211        let expected: BTreeMap<_, _> = BTreeMap::from([
212            (
213                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
214                Some("updated".to_string()),
215            ),
216            (
217                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
218                Some("keep me".to_string()),
219            ),
220            (
221                fake_data::immutable_file(PathBuf::default(), 2, "2.chunk"),
222                Some("keep me too".to_string()),
223            ),
224            (
225                fake_data::immutable_file(PathBuf::default(), 3, "3.chunk"),
226                None,
227            ),
228        ]);
229        let immutables = expected.keys().cloned().collect();
230
231        provider
232            .store(values_to_store)
233            .await
234            .expect("Cache write should not fail");
235        let result = provider.get(immutables).await.expect("Cache read should not fail");
236
237        assert_eq!(expected, result);
238    }
239
240    #[tokio::test]
241    async fn reset_clear_existing_values() {
242        let file = get_test_dir("reset_clear_existing_values").join("immutable-cache-store.json");
243        let provider = JsonImmutableFileDigestCacheProvider::new(&file);
244        let values_to_store = vec![
245            ("0.chunk".to_string(), "digest 0".to_string()),
246            ("1.chunk".to_string(), "digest 1".to_string()),
247        ];
248        let expected: BTreeMap<_, _> = BTreeMap::from([
249            (
250                fake_data::immutable_file(PathBuf::default(), 0, "0.chunk"),
251                Some("digest 0".to_string()),
252            ),
253            (
254                fake_data::immutable_file(PathBuf::default(), 1, "1.chunk"),
255                Some("digest 1".to_string()),
256            ),
257        ]);
258        let immutables = expected.keys().cloned().collect();
259
260        provider
261            .store(values_to_store)
262            .await
263            .expect("Cache write should not fail");
264        provider.reset().await.expect("reset should not fails");
265
266        let result: BTreeMap<_, _> =
267            provider.get(immutables).await.expect("Cache read should not fail");
268
269        assert!(result.into_iter().all(|(_, cache)| cache.is_none()));
270    }
271}