Skip to main content

aqua_registry/
cache.rs

1use crate::{AquaRegistryError, CompiledRegistry, ParsedRegistry, Result};
2use blake3::Hasher as Blake3Hasher;
3use siphasher::sip::SipHasher;
4use std::fs;
5use std::hash::{Hash, Hasher};
6use std::io::Write;
7use std::path::{Path, PathBuf};
8use std::time::{Duration, SystemTime};
9
10const COMPILED_REGISTRY_CACHE_VERSION: &str = "v9";
11
12#[derive(Debug, Clone)]
13pub struct RegistryCache {
14    root: PathBuf,
15}
16
17impl RegistryCache {
18    pub fn new(root: impl Into<PathBuf>) -> Self {
19        Self { root: root.into() }
20    }
21
22    pub fn source_path(&self, registry_url: &str) -> PathBuf {
23        self.root
24            .join("sources")
25            .join(format!("{}.yaml", registry_url_hash(registry_url)))
26    }
27
28    pub fn read_source(&self, registry_url: &str) -> Result<Option<String>> {
29        let path = self.source_path(registry_url);
30        read_optional_to_string(&path)
31    }
32
33    pub fn read_fresh_source(
34        &self,
35        registry_url: &str,
36        max_age: Duration,
37    ) -> Result<Option<String>> {
38        let path = self.source_path(registry_url);
39        if !path_is_fresh(&path, max_age)? {
40            return Ok(None);
41        }
42        read_optional_to_string(&path)
43    }
44
45    pub fn write_source(&self, registry_url: &str, source: &str) -> Result<()> {
46        let path = self.source_path(registry_url);
47        let Some(parent) = path.parent() else {
48            return Err(AquaRegistryError::RegistryNotAvailable(format!(
49                "cached aqua registry source path has no parent: {}",
50                path.display()
51            )));
52        };
53        fs::create_dir_all(parent)?;
54
55        let mut tmp = tempfile::NamedTempFile::with_prefix_in("registry-source.", parent)?;
56        tmp.write_all(source.as_bytes())?;
57        tmp.persist(&path).map_err(|err| err.error)?;
58        Ok(())
59    }
60
61    pub fn source_hash(source: &str) -> String {
62        source_hash(source)
63    }
64
65    pub fn compiled_dir(&self, registry_url: &str, source_hash: &str) -> PathBuf {
66        self.root
67            .join("compiled")
68            .join(registry_url_hash(registry_url))
69            .join(COMPILED_REGISTRY_CACHE_VERSION)
70            .join(source_hash)
71    }
72
73    pub fn load_compiled(&self, registry_url: &str, source_hash: &str) -> Result<CompiledRegistry> {
74        CompiledRegistry::load(self.compiled_dir(registry_url, source_hash))
75    }
76
77    pub fn write_compiled(
78        &self,
79        registry_url: &str,
80        source_hash: &str,
81        registry: &ParsedRegistry,
82    ) -> Result<CompiledRegistry> {
83        let compiled_dir = self.compiled_dir(registry_url, source_hash);
84        if let Ok(existing) = CompiledRegistry::load(&compiled_dir) {
85            self.prune_stale_compiled(registry_url, source_hash);
86            return Ok(existing);
87        }
88
89        let Some(parent) = compiled_dir.parent() else {
90            return Err(AquaRegistryError::RegistryNotAvailable(format!(
91                "compiled aqua registry cache path has no parent: {}",
92                compiled_dir.display()
93            )));
94        };
95        fs::create_dir_all(parent)?;
96
97        let tmp_dir = tempfile::Builder::new()
98            .prefix(&format!("{source_hash}.tmp-"))
99            .tempdir_in(parent)?;
100        let tmp_path = tmp_dir.path().to_path_buf();
101
102        registry.write_compiled_cache(&tmp_path)?;
103        let tmp_path = tmp_dir.keep();
104
105        if let Ok(existing) = CompiledRegistry::load(&compiled_dir) {
106            cleanup_tmp_dir_for_existing_compiled_cache(&tmp_path, &compiled_dir)?;
107            self.prune_stale_compiled(registry_url, source_hash);
108            return Ok(existing);
109        }
110
111        if compiled_dir.exists() {
112            remove_dir_all_if_exists(&compiled_dir)?;
113        }
114
115        if let Err(err) = fs::rename(&tmp_path, &compiled_dir) {
116            if let Ok(existing) = CompiledRegistry::load(&compiled_dir) {
117                cleanup_tmp_dir_for_existing_compiled_cache(&tmp_path, &compiled_dir)?;
118                self.prune_stale_compiled(registry_url, source_hash);
119                return Ok(existing);
120            }
121            let _ = remove_dir_all_if_exists(&tmp_path);
122            return Err(err.into());
123        }
124
125        let compiled = CompiledRegistry::load(&compiled_dir)?;
126        self.prune_stale_compiled(registry_url, source_hash);
127        Ok(compiled)
128    }
129
130    pub fn prune_stale_compiled(&self, registry_url: &str, source_hash: &str) {
131        let current_dir = self.compiled_dir(registry_url, source_hash);
132        prune_stale_compiled_registries(&current_dir);
133    }
134}
135
136fn registry_url_hash(registry_url: &str) -> String {
137    hash_to_str(&registry_url)
138}
139
140fn source_hash(source: &str) -> String {
141    let mut hasher = Blake3Hasher::new();
142    hasher.update(source.as_bytes());
143    hasher.finalize().to_hex().to_string()
144}
145
146fn hash_to_str<T: Hash>(t: &T) -> String {
147    let mut s = SipHasher::new();
148    t.hash(&mut s);
149    format!("{:x}", s.finish())
150}
151
152fn read_optional_to_string(path: &Path) -> Result<Option<String>> {
153    match fs::read_to_string(path) {
154        Ok(source) => Ok(Some(source)),
155        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
156        Err(err) => Err(err.into()),
157    }
158}
159
160fn path_is_fresh(path: &Path, max_age: Duration) -> Result<bool> {
161    let Some(age) = path_age(path)? else {
162        return Ok(false);
163    };
164    Ok(age < max_age)
165}
166
167fn path_age(path: &Path) -> Result<Option<Duration>> {
168    let metadata = match fs::metadata(path) {
169        Ok(metadata) => metadata,
170        Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
171        Err(err) => return Err(err.into()),
172    };
173    let modified = metadata.modified()?;
174    Ok(Some(
175        SystemTime::now()
176            .duration_since(modified)
177            .unwrap_or_default(),
178    ))
179}
180
181fn prune_stale_compiled_registries(current_dir: &Path) {
182    let Some(parent) = current_dir.parent() else {
183        return;
184    };
185    let Ok(entries) = fs::read_dir(parent) else {
186        return;
187    };
188
189    for entry in entries.flatten() {
190        let path = entry.path();
191        if path == current_dir {
192            continue;
193        }
194        if entry.file_type().is_ok_and(|file_type| file_type.is_dir())
195            && is_compiled_source_hash_dir(&path)
196            && let Err(err) = fs::remove_dir_all(&path)
197        {
198            log::debug!(
199                "failed to prune stale compiled aqua registry cache {}: {err}",
200                path.display()
201            );
202        }
203    }
204}
205
206fn is_compiled_source_hash_dir(path: &Path) -> bool {
207    path.file_name()
208        .and_then(|name| name.to_str())
209        .is_some_and(|name| name.len() == 64 && name.bytes().all(|b| b.is_ascii_hexdigit()))
210}
211
212fn cleanup_tmp_dir_for_existing_compiled_cache(tmp_dir: &Path, compiled_dir: &Path) -> Result<()> {
213    match fs::remove_dir_all(tmp_dir) {
214        Ok(()) => Ok(()),
215        Err(err)
216            if err.kind() == std::io::ErrorKind::NotFound
217                && CompiledRegistry::load(compiled_dir).is_ok() =>
218        {
219            Ok(())
220        }
221        Err(err) => Err(err.into()),
222    }
223}
224
225fn remove_dir_all_if_exists(path: &Path) -> std::io::Result<()> {
226    match fs::remove_dir_all(path) {
227        Ok(()) => Ok(()),
228        Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()),
229        Err(err) => Err(err),
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    fn registry_source(package_id: &str) -> String {
238        format!("packages:\n  - name: {package_id}\n    url: https://example.com/tool\n")
239    }
240
241    #[test]
242    fn source_cache_reads_fresh_sources_and_skips_stale_sources() {
243        let temp = tempfile::tempdir().unwrap();
244        let cache = RegistryCache::new(temp.path());
245        let registry_url = "https://example.com/aqua-registry";
246
247        cache.write_source(registry_url, "packages: []").unwrap();
248
249        assert_eq!(
250            cache
251                .read_fresh_source(registry_url, Duration::from_secs(60))
252                .unwrap()
253                .as_deref(),
254            Some("packages: []")
255        );
256        assert!(
257            cache
258                .read_fresh_source(registry_url, Duration::ZERO)
259                .unwrap()
260                .is_none()
261        );
262    }
263
264    #[test]
265    fn source_cache_writes_atomically_and_overwrites_existing_source() {
266        let temp = tempfile::tempdir().unwrap();
267        let cache = RegistryCache::new(temp.path());
268        let registry_url = "https://example.com/aqua-registry";
269
270        cache.write_source(registry_url, "first").unwrap();
271        cache.write_source(registry_url, "second").unwrap();
272
273        assert_eq!(
274            cache.read_source(registry_url).unwrap().as_deref(),
275            Some("second")
276        );
277        assert!(cache.source_path(registry_url).is_file());
278    }
279
280    #[test]
281    fn compiled_cache_is_scoped_by_registry_url() {
282        let cache = RegistryCache::new("/cache");
283        let source_hash = RegistryCache::source_hash("packages: []");
284        let first = cache.compiled_dir("https://example.com/one", &source_hash);
285        let second = cache.compiled_dir("https://example.com/two", &source_hash);
286
287        assert_ne!(first.parent(), second.parent());
288        assert_eq!(
289            first.file_name().and_then(|name| name.to_str()),
290            Some(source_hash.as_str())
291        );
292    }
293
294    #[test]
295    fn compiled_cache_writes_loads_and_prunes_stale_source_hash_siblings() {
296        let temp = tempfile::tempdir().unwrap();
297        let cache = RegistryCache::new(temp.path());
298        let registry_url = "https://example.com/aqua-registry";
299        let first_source = registry_source("example/first");
300        let second_source = registry_source("example/second");
301        let first_hash = RegistryCache::source_hash(&first_source);
302        let second_hash = RegistryCache::source_hash(&second_source);
303        let first_registry = ParsedRegistry::parse_yaml(&first_source).unwrap();
304        let second_registry = ParsedRegistry::parse_yaml(&second_source).unwrap();
305
306        cache
307            .write_compiled(registry_url, &first_hash, &first_registry)
308            .unwrap();
309        let first_dir = cache.compiled_dir(registry_url, &first_hash);
310        assert!(first_dir.is_dir());
311
312        cache
313            .write_compiled(registry_url, &second_hash, &second_registry)
314            .unwrap();
315        let second_dir = cache.compiled_dir(registry_url, &second_hash);
316        let loaded = cache.load_compiled(registry_url, &second_hash).unwrap();
317
318        assert!(second_dir.is_dir());
319        assert!(!first_dir.exists());
320        assert!(loaded.package("example/second").is_ok());
321    }
322
323    #[test]
324    fn compiled_cache_prune_skips_temp_directories() {
325        let temp = tempfile::tempdir().unwrap();
326        let cache = RegistryCache::new(temp.path());
327        let registry_url = "https://example.com/aqua-registry";
328        let current_hash = RegistryCache::source_hash(&registry_source("example/current"));
329        let stale_hash = RegistryCache::source_hash(&registry_source("example/stale"));
330        let current_dir = cache.compiled_dir(registry_url, &current_hash);
331        let stale_dir = cache.compiled_dir(registry_url, &stale_hash);
332        let tmp_dir = current_dir
333            .parent()
334            .unwrap()
335            .join(format!("{current_hash}.tmp-in-progress"));
336
337        fs::create_dir_all(&current_dir).unwrap();
338        fs::create_dir_all(&stale_dir).unwrap();
339        fs::create_dir_all(&tmp_dir).unwrap();
340
341        cache.prune_stale_compiled(registry_url, &current_hash);
342
343        assert!(current_dir.is_dir());
344        assert!(!stale_dir.exists());
345        assert!(tmp_dir.is_dir());
346    }
347
348    #[test]
349    fn compiled_temp_cleanup_treats_missing_temp_as_success_when_final_cache_exists() {
350        let temp = tempfile::tempdir().unwrap();
351        let cache = RegistryCache::new(temp.path());
352        let registry_url = "https://example.com/aqua-registry";
353        let source = registry_source("example/tool");
354        let source_hash = RegistryCache::source_hash(&source);
355        let registry = ParsedRegistry::parse_yaml(&source).unwrap();
356        let compiled_dir = cache.compiled_dir(registry_url, &source_hash);
357        let missing_tmp_dir = compiled_dir.with_file_name(format!("{source_hash}.tmp-missing"));
358
359        registry.write_compiled_cache(&compiled_dir).unwrap();
360
361        cleanup_tmp_dir_for_existing_compiled_cache(&missing_tmp_dir, &compiled_dir).unwrap();
362    }
363
364    #[test]
365    fn registry_url_hash_matches_existing_cache_layout() {
366        assert_eq!(registry_url_hash("foo"), "e1b19adfb2e348a2");
367    }
368}