mithril_common/digesters/cache/
json_provider.rs

1use crate::{
2    digesters::{
3        cache::provider::{ImmutableDigesterCacheGetError, ImmutableDigesterCacheStoreError},
4        cache::CacheProviderResult,
5        cache::ImmutableFileDigestCacheProvider,
6        ImmutableFile,
7    },
8    entities::{HexEncodedDigest, ImmutableFileName},
9};
10
11use async_trait::async_trait;
12use std::{
13    collections::BTreeMap,
14    path::{Path, PathBuf},
15};
16#[cfg(feature = "fs")]
17use tokio::{
18    fs,
19    fs::File,
20    io::{AsyncReadExt, AsyncWriteExt},
21};
22
23type InnerStructure = BTreeMap<ImmutableFileName, HexEncodedDigest>;
24
25/// A in memory [ImmutableFileDigestCacheProvider].
26pub struct JsonImmutableFileDigestCacheProvider {
27    filepath: PathBuf,
28}
29
30impl JsonImmutableFileDigestCacheProvider {
31    /// [JsonImmutableFileDigestCacheProvider] factory
32    pub fn new(filepath: &Path) -> Self {
33        Self {
34            filepath: filepath.to_path_buf(),
35        }
36    }
37
38    #[cfg(test)]
39    /// [Test Only] Build a new [JsonImmutableFileDigestCacheProvider] that contains the given values.
40    pub async fn from(filepath: &Path, values: InnerStructure) -> Self {
41        let provider = Self::new(filepath);
42        provider.write_data(values).await.unwrap();
43        provider
44    }
45
46    async fn write_data(
47        &self,
48        values: InnerStructure,
49    ) -> Result<(), ImmutableDigesterCacheStoreError> {
50        let mut file = File::create(&self.filepath).await?;
51        file.write_all(serde_json::to_string_pretty(&values)?.as_bytes())
52            .await?;
53
54        Ok(())
55    }
56
57    async fn read_data(&self) -> Result<InnerStructure, ImmutableDigesterCacheGetError> {
58        match self.filepath.exists() {
59            true => {
60                let mut file = File::open(&self.filepath).await?;
61                let mut json_string = String::new();
62                file.read_to_string(&mut json_string).await?;
63                let values: InnerStructure = serde_json::from_str(&json_string)?;
64                Ok(values)
65            }
66            false => Ok(BTreeMap::new()),
67        }
68    }
69}
70
71#[async_trait]
72impl ImmutableFileDigestCacheProvider for JsonImmutableFileDigestCacheProvider {
73    async fn store(
74        &self,
75        digest_per_filenames: Vec<(ImmutableFileName, HexEncodedDigest)>,
76    ) -> CacheProviderResult<()> {
77        let mut data = self.read_data().await?;
78        for (filename, digest) in digest_per_filenames {
79            data.insert(filename, digest);
80        }
81        self.write_data(data).await?;
82
83        Ok(())
84    }
85
86    async fn get(
87        &self,
88        immutables: Vec<ImmutableFile>,
89    ) -> CacheProviderResult<BTreeMap<ImmutableFile, Option<HexEncodedDigest>>> {
90        let values = self.read_data().await?;
91        let mut result = BTreeMap::new();
92
93        for immutable in immutables {
94            let value = values.get(&immutable.filename).map(|f| f.to_owned());
95            result.insert(immutable, value);
96        }
97
98        Ok(result)
99    }
100
101    async fn reset(&self) -> CacheProviderResult<()> {
102        fs::remove_file(&self.filepath)
103            .await
104            .map_err(ImmutableDigesterCacheStoreError::from)?;
105
106        Ok(())
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use crate::digesters::cache::{
113        ImmutableFileDigestCacheProvider, JsonImmutableFileDigestCacheProvider,
114    };
115    use crate::digesters::ImmutableFile;
116    use crate::test_utils::TempDir;
117    use std::{collections::BTreeMap, path::PathBuf};
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                ImmutableFile::dummy(PathBuf::default(), 0, "0.chunk"),
134                Some("digest 0".to_string()),
135            ),
136            (
137                ImmutableFile::dummy(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            ImmutableFile::dummy(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<_, _> =
191            BTreeMap::from([(ImmutableFile::dummy(PathBuf::default(), 2, "2.chunk"), None)]);
192        let immutables = expected.keys().cloned().collect();
193
194        let result = provider
195            .get(immutables)
196            .await
197            .expect("Cache read should not fail");
198
199        assert_eq!(expected, result);
200    }
201
202    #[tokio::test]
203    async fn store_erase_existing_values() {
204        let file = get_test_dir("store_erase_existing_values").join("immutable-cache-store.json");
205        let provider = JsonImmutableFileDigestCacheProvider::from(
206            &file,
207            BTreeMap::from([
208                ("0.chunk".to_string(), "to erase".to_string()),
209                ("1.chunk".to_string(), "keep me".to_string()),
210                ("2.chunk".to_string(), "keep me too".to_string()),
211            ]),
212        )
213        .await;
214        let values_to_store = vec![
215            ("0.chunk".to_string(), "updated".to_string()),
216            ("1.chunk".to_string(), "keep me".to_string()),
217        ];
218        let expected: BTreeMap<_, _> = BTreeMap::from([
219            (
220                ImmutableFile::dummy(PathBuf::default(), 0, "0.chunk"),
221                Some("updated".to_string()),
222            ),
223            (
224                ImmutableFile::dummy(PathBuf::default(), 1, "1.chunk"),
225                Some("keep me".to_string()),
226            ),
227            (
228                ImmutableFile::dummy(PathBuf::default(), 2, "2.chunk"),
229                Some("keep me too".to_string()),
230            ),
231            (ImmutableFile::dummy(PathBuf::default(), 3, "3.chunk"), None),
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
240            .get(immutables)
241            .await
242            .expect("Cache read should not fail");
243
244        assert_eq!(expected, result);
245    }
246
247    #[tokio::test]
248    async fn reset_clear_existing_values() {
249        let file = get_test_dir("reset_clear_existing_values").join("immutable-cache-store.json");
250        let provider = JsonImmutableFileDigestCacheProvider::new(&file);
251        let values_to_store = vec![
252            ("0.chunk".to_string(), "digest 0".to_string()),
253            ("1.chunk".to_string(), "digest 1".to_string()),
254        ];
255        let expected: BTreeMap<_, _> = BTreeMap::from([
256            (
257                ImmutableFile::dummy(PathBuf::default(), 0, "0.chunk"),
258                Some("digest 0".to_string()),
259            ),
260            (
261                ImmutableFile::dummy(PathBuf::default(), 1, "1.chunk"),
262                Some("digest 1".to_string()),
263            ),
264        ]);
265        let immutables = expected.keys().cloned().collect();
266
267        provider
268            .store(values_to_store)
269            .await
270            .expect("Cache write should not fail");
271        provider.reset().await.expect("reset should not fails");
272
273        let result: BTreeMap<_, _> = provider
274            .get(immutables)
275            .await
276            .expect("Cache read should not fail");
277
278        assert!(result.into_iter().all(|(_, cache)| cache.is_none()));
279    }
280}
mithril_common/digesters/cache/json_provider.rs

mithril_common/digesters/cache/
json_provider.rs
