Skip to main content

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