1use fs2::FileExt;
9use serde::{Serialize, de::DeserializeOwned};
10use std::io::{Read, Write};
11#[cfg(unix)]
12use std::os::unix::fs::PermissionsExt;
13use std::path::PathBuf;
14use std::sync::atomic::{AtomicU64, Ordering};
15use tempfile::NamedTempFile;
16use tracing::{error, warn};
17
18const DISK_CACHE_DEGRADED_THRESHOLD: u64 = 100;
20
21pub struct DiskCache {
27 base: PathBuf,
28 disabled: bool,
29 write_failures: AtomicU64,
31 total_write_failures: AtomicU64,
33}
34
35impl DiskCache {
36 pub fn drain_write_failures(&self) -> u64 {
39 self.write_failures.swap(0, Ordering::Relaxed)
40 }
41
42 pub fn is_degraded(&self) -> bool {
45 self.total_write_failures.load(Ordering::Relaxed) >= DISK_CACHE_DEGRADED_THRESHOLD
46 }
47}
48
49impl DiskCache {
50 pub fn new(base: PathBuf, disabled: bool) -> Self {
53 if disabled {
54 return Self {
55 base,
56 disabled: true,
57 write_failures: AtomicU64::new(0),
58 total_write_failures: AtomicU64::new(0),
59 };
60 }
61 if let Err(e) = std::fs::create_dir_all(&base) {
62 warn!(path = %base.display(), error = %e, "disk cache disabled: failed to create cache directory");
63 return Self {
64 base,
65 disabled: true,
66 write_failures: AtomicU64::new(0),
67 total_write_failures: AtomicU64::new(0),
68 };
69 }
70 #[cfg(unix)]
71 if let Err(e) = std::fs::set_permissions(&base, std::fs::Permissions::from_mode(0o700)) {
72 warn!(path = %base.display(), error = %e, "disk cache: failed to set directory permissions to 0700");
73 }
74 #[cfg(not(unix))]
75 let _ = &base; Self {
77 base,
78 disabled: false,
79 write_failures: AtomicU64::new(0),
80 total_write_failures: AtomicU64::new(0),
81 }
82 }
83
84 pub fn entry_path(&self, tool: &str, key: &blake3::Hash) -> PathBuf {
85 let hex = format!("{}", key);
86 self.base
87 .join(tool)
88 .join(&hex[..2])
89 .join(format!("{}.json.snap", hex))
90 }
91
92 pub fn get<T: DeserializeOwned>(&self, tool: &str, key: &blake3::Hash) -> Option<T> {
95 if self.disabled {
96 return None;
97 }
98 let path = self.entry_path(tool, key);
99 let dir = path.parent()?;
100
101 let _lock = lock_shard_shared(dir)?;
103
104 let compressed = std::fs::read(&path).ok()?;
105 let mut decompressed_data = Vec::new();
106 snap::read::FrameDecoder::new(&compressed[..])
107 .read_to_end(&mut decompressed_data)
108 .ok()?;
109 serde_json::from_slice(&decompressed_data).ok()
110 }
111
112 fn serialize_entry<T: Serialize>(value: &T) -> Option<Vec<u8>> {
114 let json = serde_json::to_vec(value).ok()?;
115 let mut compressed = Vec::new();
116 snap::write::FrameEncoder::new(&mut compressed)
117 .write_all(&json)
118 .ok()?;
119 Some(compressed)
120 }
121
122 fn write_entry_atomically(
126 dir: &std::path::Path,
127 path: &std::path::Path,
128 compressed: &[u8],
129 ) -> Result<(), std::io::Error> {
130 use std::io::Write;
131 let _lock = lock_shard_exclusive(dir)?;
133 let mut tmp = NamedTempFile::new_in(dir)?;
134 tmp.write_all(compressed)?;
135 tmp.persist(path).map(|_| ()).map_err(|e| e.error)
136 }
137
138 pub fn put<T: Serialize>(&self, tool: &str, key: &blake3::Hash, value: &T) {
140 if self.disabled {
141 return;
142 }
143 let path = self.entry_path(tool, key);
144 let dir = match path.parent() {
145 Some(d) => d.to_path_buf(),
146 None => return,
147 };
148 if let Err(e) = std::fs::create_dir_all(&dir) {
149 warn!(tool, error = %e, "disk cache: failed to create cache directory");
150 self.record_write_failure();
151 return;
152 }
153 let compressed = match Self::serialize_entry(value) {
154 Some(c) => c,
155 None => return,
156 };
157 if Self::write_entry_atomically(&dir, &path, &compressed)
158 .ok()
159 .is_none()
160 {
161 self.record_write_failure();
162 }
163 }
164
165 fn record_write_failure(&self) {
169 self.write_failures.fetch_add(1, Ordering::Relaxed);
170 let total = self.total_write_failures.fetch_add(1, Ordering::Relaxed) + 1;
171 if total == DISK_CACHE_DEGRADED_THRESHOLD {
172 error!(
173 path = %self.base.display(),
174 total,
175 threshold = DISK_CACHE_DEGRADED_THRESHOLD,
176 "disk cache is degraded: consecutive write failures have reached the alert threshold; \
177 check disk space and permissions at the cache directory"
178 );
179 }
180 }
181
182 pub fn evict_stale(&self, retention_days: u64) {
184 if self.disabled {
185 return;
186 }
187 let cutoff = std::time::SystemTime::now()
188 .checked_sub(std::time::Duration::from_secs(retention_days * 86_400))
189 .unwrap_or(std::time::UNIX_EPOCH);
190 let _ = evict_dir_recursive(&self.base, cutoff);
191 }
192}
193
194fn evict_dir_recursive(
195 dir: &std::path::Path,
196 cutoff: std::time::SystemTime,
197) -> std::io::Result<()> {
198 for entry in std::fs::read_dir(dir)? {
199 let entry = entry?;
200 let meta = entry.metadata()?;
201 let path = entry.path();
202 if meta.is_dir() {
203 let _ = evict_dir_recursive(&path, cutoff);
204 } else if meta.is_file()
205 && let Ok(mtime) = meta.modified()
206 && mtime < cutoff
207 && let Err(e) = std::fs::remove_file(&path)
208 {
209 warn!(path = %path.display(), error = %e, "disk cache: failed to evict stale cache file");
210 }
211 }
212 Ok(())
213}
214
215fn lock_shard_shared(shard_dir: &std::path::Path) -> Option<ShardLockGuard> {
220 let lock_path = shard_dir.join(".lock");
221 let file = std::fs::OpenOptions::new()
222 .create(true)
223 .write(true)
224 .truncate(false)
225 .open(&lock_path)
226 .ok()?;
227 match file.lock_shared() {
228 Ok(()) => Some(ShardLockGuard(file)),
229 Err(e) => {
230 warn!(
231 error = %e, lock_path = %lock_path.display(),
232 "disk cache: failed to acquire shared lock on shard; proceeding without lock"
233 );
234 None
235 }
236 }
237}
238
239fn lock_shard_exclusive(shard_dir: &std::path::Path) -> Result<ShardLockGuard, std::io::Error> {
244 let lock_path = shard_dir.join(".lock");
245 let file = std::fs::OpenOptions::new()
246 .create(true)
247 .write(true)
248 .truncate(false)
249 .open(&lock_path)?;
250 file.lock_exclusive()?;
251 Ok(ShardLockGuard(file))
252}
253
254struct ShardLockGuard(
257 #[expect(dead_code)]
260 std::fs::File,
261);
262
263#[cfg(test)]
264mod disk_cache_tests {
265 use super::*;
266 use std::io::Read;
267 use tempfile::TempDir;
268
269 #[test]
270 fn test_disk_cache_roundtrip() {
271 let dir = TempDir::new().unwrap();
272 let cache = DiskCache::new(dir.path().to_path_buf(), false);
273 let key = blake3::hash(b"test-key");
274 let value = serde_json::json!({"result": "success", "count": 42});
275 cache.put("analyze_file", &key, &value);
276 let retrieved: Option<serde_json::Value> = cache.get("analyze_file", &key);
277 assert_eq!(retrieved, Some(value));
278 }
279
280 #[test]
281 fn test_disk_cache_permissions() {
282 #[cfg(unix)]
283 {
284 use std::os::unix::fs::PermissionsExt;
285 let dir = TempDir::new().unwrap();
286 let cache_dir = dir.path().join("analysis-cache");
287 let _cache = DiskCache::new(cache_dir.clone(), false);
288 let meta = std::fs::metadata(&cache_dir).unwrap();
289 let mode = meta.permissions().mode() & 0o777;
290 assert_eq!(mode, 0o700, "cache dir must be mode 0700");
291 }
292 }
293
294 #[test]
295 fn test_disk_cache_corrupt_entry_returns_none() {
296 let dir = TempDir::new().unwrap();
297 let cache = DiskCache::new(dir.path().to_path_buf(), false);
298 let key = blake3::hash(b"corrupt-key");
299 let path = cache.entry_path("analyze_file", &key);
300 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
301 std::fs::write(&path, b"not valid snappy data").unwrap();
302 let result: Option<serde_json::Value> = cache.get("analyze_file", &key);
303 assert!(result.is_none(), "corrupt entry must return None");
304 }
305
306 #[test]
307 fn test_disk_cache_disabled_on_dir_creation_failure() {
308 let dir = TempDir::new().unwrap();
309 let blocked = dir.path().join("blocked");
312 std::fs::write(&blocked, b"").unwrap();
313 let cache = DiskCache::new(blocked, false);
314 let key = blake3::hash(b"should-not-exist");
316 cache.put("analyze_file", &key, &serde_json::json!({"x": 1}));
317 let result: Option<serde_json::Value> = cache.get("analyze_file", &key);
318 assert!(
319 result.is_none(),
320 "cache must be disabled after dir creation failure"
321 );
322 assert!(
323 cache.disabled,
324 "disabled flag must be true after dir creation failure"
325 );
326 }
327
328 #[test]
329 fn test_concurrent_get_put_same_shard() {
330 let dir = TempDir::new().unwrap();
333 let cache = std::sync::Arc::new(DiskCache::new(dir.path().to_path_buf(), false));
334 let key = blake3::hash(b"concurrent-test-key");
335 let value = serde_json::json!({"result": "from put thread", "n": 42});
336
337 cache.put("analyze_file", &key, &value);
339
340 let cache_get = cache.clone();
341 let cache_put = cache.clone();
342 let key_put = key;
343 let key_get = key;
344 let value_put = serde_json::json!({"result": "from put thread", "n": 100});
345
346 std::thread::scope(|scope| {
347 scope.spawn(|| {
348 cache_put.put("analyze_file", &key_put, &value_put);
350 });
351 scope.spawn(|| {
352 let _: Option<serde_json::Value> = cache_get.get("analyze_file", &key_get);
354 });
355 });
356
357 let result: Option<serde_json::Value> = cache.get("analyze_file", &key);
359 assert!(
360 result.is_some(),
361 "entry must still be retrievable after concurrent access"
362 );
363 }
364
365 #[test]
366 fn test_concurrent_puts_same_shard() {
367 let dir = TempDir::new().unwrap();
370 let cache = std::sync::Arc::new(DiskCache::new(dir.path().to_path_buf(), false));
371 let key = blake3::hash(b"concurrent-put-key");
372 let value_a = serde_json::json!({"writer": "A", "data": "hello from A"});
373 let value_b = serde_json::json!({"writer": "B", "data": "hello from B"});
374
375 let cache_a = cache.clone();
376 let cache_b = cache.clone();
377 let key_a = key;
378 let key_b = key;
379
380 std::thread::scope(|scope| {
381 scope.spawn(|| {
382 cache_a.put("analyze_file", &key_a, &value_a);
383 });
384 scope.spawn(|| {
385 cache_b.put("analyze_file", &key_b, &value_b);
386 });
387 });
388
389 let result: Option<serde_json::Value> = cache.get("analyze_file", &key);
391 assert!(
392 result.is_some(),
393 "entry must be retrievable after concurrent puts"
394 );
395 let v = result.unwrap();
397 let writer = v.get("writer").and_then(|w| w.as_str());
398 assert!(
399 writer == Some("A") || writer == Some("B"),
400 "entry must contain data from one of the concurrent writers, got {writer:?}"
401 );
402 }
403}