1use std::collections::{BTreeMap, BTreeSet};
13use std::fs::{self, File};
14use std::io::{BufRead, Write};
15use std::path::{Path, PathBuf};
16use std::time::{Duration, SystemTime};
17
18use anyhow::{Context, Result, bail};
19use fs2::FileExt;
20use serde::{Deserialize, Serialize};
21use serde_json::Value;
22
23const IN_FLIGHT: Duration = Duration::from_secs(60 * 60);
25
26const UNUSED: Duration = Duration::from_secs(7 * 24 * 60 * 60);
29
30#[derive(Debug, Default, PartialEq, Eq)]
31pub struct Pruned {
32 pub size: u64,
34 pub entries: usize,
35 pub blobs: usize,
36 pub freed: u64,
37}
38
39pub fn parse_size(text: &str) -> Result<u64> {
42 let text = text.trim();
43 let (number, unit) = match text.find(|c: char| !(c.is_ascii_digit() || c == '.')) {
44 Some(at) => (&text[..at], text[at..].trim_start()),
45 None => (text, ""),
46 };
47 let factor: u64 = match unit {
48 "" | "B" => 1,
49 "KB" => 1 << 10,
50 "MB" => 1 << 20,
51 "GB" => 1 << 30,
52 _ => bail!("invalid cache size {text:?}: use a number with an optional KB, MB or GB"),
53 };
54 if number.matches('.').count() > 1 || number.is_empty() {
55 bail!("invalid cache size {text:?}");
56 }
57 let number: f64 = number
58 .parse()
59 .with_context(|| format!("invalid cache size {text:?}"))?;
60 Ok((number * factor as f64) as u64)
61}
62
63pub fn max_size(configured: Option<&Value>, cache: &Path) -> Result<Option<u64>> {
68 let configured =
69 match std::env::var("QK_MAX_CACHE_SIZE").or_else(|_| std::env::var("NX_MAX_CACHE_SIZE")) {
70 Ok(value) => Some(value),
71 Err(_) => configured.map(|value| match value {
72 Value::String(text) => text.clone(),
73 other => other.to_string(),
74 }),
75 };
76 let limit = match configured {
77 Some(text) => parse_size(&text)?,
78 None => cache
79 .ancestors()
80 .find(|path| path.exists())
81 .and_then(|existing| fs2::total_space(existing).ok())
82 .map_or(100 << 30, |total| total / 10),
83 };
84 Ok((limit > 0).then_some(limit))
85}
86
87fn pins(root: &Path) -> PathBuf {
89 root.join("locks").join(".pins")
90}
91
92fn pin_gate(root: &Path) -> Result<File> {
95 let directory = pins(root);
96 fs::create_dir_all(&directory)?;
97 Ok(File::options()
98 .create(true)
99 .truncate(false)
100 .write(true)
101 .open(directory.join(".gate"))?)
102}
103
104#[derive(Debug)]
109pub(crate) struct Pin {
110 lock: Option<File>,
111 path: PathBuf,
112}
113
114pub(crate) fn pin(root: &Path, blobs: &BTreeSet<String>) -> Result<Pin> {
118 let gate = pin_gate(root)?;
119 gate.lock_shared()?;
120 let directory = pins(root);
121 let name = format!(
122 "{}-{}",
123 std::process::id(),
124 &blake3::hash(
125 format!("{:?}{:?}", std::thread::current().id(), SystemTime::now()).as_bytes()
126 )
127 .to_hex()[..16]
128 );
129 let path = directory.join(name);
130 let lock = File::options()
131 .create_new(true)
132 .write(true)
133 .open(path.with_extension("lock"))?;
134 lock.lock_exclusive()?;
135 let pin = Pin {
136 lock: Some(lock),
137 path,
138 };
139 let mut listing = String::new();
140 for blob in blobs {
141 listing.push_str(blob);
142 listing.push('\n');
143 }
144 let mut file = File::create(pin.path.with_extension("blobs"))?;
145 file.write_all(listing.as_bytes())?;
146 drop(gate);
147 Ok(pin)
148}
149
150impl Drop for Pin {
151 fn drop(&mut self) {
152 let _ = fs::remove_file(self.path.with_extension("blobs"));
153 drop(self.lock.take());
155 let _ = fs::remove_file(self.path.with_extension("lock"));
156 }
157}
158
159fn pinned(root: &Path) -> BTreeSet<String> {
161 let mut held = BTreeSet::new();
162 for item in fs::read_dir(pins(root)).into_iter().flatten().flatten() {
163 let path = item.path();
164 if path.extension().is_none_or(|extension| extension != "lock") {
165 continue;
166 }
167 let Ok(lock) = File::open(&path) else {
168 continue;
169 };
170 if lock.try_lock_exclusive().is_ok() {
171 let _ = fs::remove_file(path.with_extension("blobs"));
172 let _ = fs::remove_file(&path);
173 continue;
174 }
175 if let Ok(listing) = File::open(path.with_extension("blobs")) {
176 held.extend(
177 std::io::BufReader::new(listing)
178 .lines()
179 .map_while(Result::ok),
180 );
181 }
182 }
183 held
184}
185
186pub(crate) fn touch(manifest: &Path) {
188 if let Ok(file) = File::options().write(true).open(manifest) {
189 let _ = file.set_modified(SystemTime::now());
190 }
191}
192
193struct Entry {
194 path: std::path::PathBuf,
195 used: SystemTime,
196 size: u64,
197 blobs: BTreeSet<String>,
198}
199
200pub fn prune(root: &Path, limit: u64) -> Result<Pruned> {
208 prune_inner(root, limit, false)
209}
210
211pub fn prune_automatic(root: &Path, limit: u64) -> Result<Pruned> {
216 prune_inner(root, limit, true)
217}
218
219fn prune_inner(root: &Path, limit: u64, automatic: bool) -> Result<Pruned> {
220 if !root.join("entries").is_dir() {
221 return Ok(Pruned::default());
222 }
223 fs::create_dir_all(root.join("locks"))?;
224 let marker = root.join("locks").join(".prune");
225 let recent = fs::metadata(&marker)
226 .ok()
227 .and_then(|metadata| metadata.modified().ok())
228 .and_then(|modified| SystemTime::now().duration_since(modified).ok())
229 .is_some_and(|age| age < IN_FLIGHT);
230 if automatic && recent {
231 let size = cached_size(root)?;
232 if size <= limit {
233 return Ok(Pruned {
234 size,
235 ..Pruned::default()
236 });
237 }
238 }
239 let guard = File::options()
240 .create(true)
241 .truncate(false)
242 .write(true)
243 .open(&marker)?;
244 if guard.try_lock_exclusive().is_err() {
245 return Ok(Pruned::default());
246 }
247 let _ = fs::remove_file(root.join("locks").join(".size"));
250 let _ = guard.set_modified(SystemTime::now());
251 let gate = pin_gate(root)?;
253 gate.lock_exclusive()?;
254 let now = SystemTime::now();
255 let age = |metadata: &fs::Metadata| {
256 metadata
257 .modified()
258 .ok()
259 .and_then(|modified| now.duration_since(modified).ok())
260 .unwrap_or_default()
261 };
262 let mut entries = Vec::new();
263 let mut references: BTreeMap<String, usize> = BTreeMap::new();
264 let warm = fs::read_dir(root.join("warm")).into_iter().flatten();
267 let pathsets = fs::read_dir(root.join("pathsets")).into_iter().flatten();
268 for item in fs::read_dir(root.join("entries"))?
269 .chain(warm)
270 .chain(pathsets)
271 {
272 let item = item?;
273 let metadata = item.metadata()?;
274 let blobs = fs::read(item.path())
276 .ok()
277 .and_then(|bytes| crate::record::StoredRecord::read(&bytes).ok())
278 .map(|record| record.blobs())
279 .unwrap_or_default();
280 for blob in &blobs {
281 *references.entry(blob.clone()).or_default() += 1;
282 }
283 entries.push(Entry {
284 path: item.path(),
285 used: metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH),
286 size: metadata.len(),
287 blobs,
288 });
289 }
290 for blob in pinned(root) {
292 *references.entry(blob).or_default() += 1;
293 }
294 let mut blob_sizes = BTreeMap::new();
295 let mut pruned = Pruned::default();
296 for item in fs::read_dir(root.join("blobs"))? {
297 let item = item?;
298 let metadata = item.metadata()?;
299 let name = item.file_name().to_string_lossy().into_owned();
300 if !references.contains_key(&name) && age(&metadata) > IN_FLIGHT {
301 if fs::remove_file(item.path()).is_ok() {
302 pruned.blobs += 1;
303 pruned.freed += metadata.len();
304 }
305 continue;
306 }
307 blob_sizes.insert(name, metadata.len());
308 }
309 let mut size: u64 =
310 entries.iter().map(|entry| entry.size).sum::<u64>() + blob_sizes.values().sum::<u64>();
311 entries.sort_by_key(|entry| entry.used);
312 let target = if size > limit {
314 limit - limit / 10
315 } else {
316 limit
317 };
318 for entry in entries {
319 let unused = now.duration_since(entry.used).is_ok_and(|age| age > UNUSED);
320 if size <= target && !unused {
321 break;
322 }
323 if fs::remove_file(&entry.path).is_err() {
324 continue;
325 }
326 pruned.entries += 1;
327 pruned.freed += entry.size;
328 size -= entry.size;
329 for blob in entry.blobs {
330 let count = references.get_mut(&blob).expect("counted above");
331 *count -= 1;
332 if *count == 0
333 && let Some(blob_size) = blob_sizes.remove(&blob)
334 && fs::remove_file(root.join("blobs").join(&blob)).is_ok()
335 {
336 pruned.blobs += 1;
337 pruned.freed += blob_size;
338 size -= blob_size;
339 }
340 }
341 }
342 drop(gate);
343 for item in fs::read_dir(root.join("tmp"))
345 .into_iter()
346 .flatten()
347 .flatten()
348 {
349 if item
350 .metadata()
351 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
352 {
353 let _ = fs::remove_dir_all(item.path()).or_else(|_| fs::remove_file(item.path()));
354 }
355 }
356 for item in fs::read_dir(root.join("locks"))?.flatten() {
358 let name = item.file_name().to_string_lossy().into_owned();
359 if name.starts_with('.') || root.join("entries").join(format!("{name}.json")).exists() {
360 continue;
361 }
362 if item
363 .metadata()
364 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
365 && let Ok(file) = File::open(item.path())
366 && file.try_lock_exclusive().is_ok()
367 {
368 let _ = fs::remove_file(item.path());
369 }
370 }
371 pruned.size = if automatic { cached_size(root)? } else { size };
372 Ok(pruned)
373}
374
375const STAMP_SETTLE: Duration = Duration::from_secs(2);
377
378#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
379struct DirectoryStamp {
380 modified: SystemTime,
381 len: u64,
382 #[cfg(unix)]
383 identity: (u64, u64, i64, i64),
384}
385
386#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
387struct SizeMemo {
388 measured: SystemTime,
389 stamps: Vec<Option<DirectoryStamp>>,
390 size: u64,
391}
392
393fn directory_stamps(root: &Path) -> Result<Vec<Option<DirectoryStamp>>> {
394 ["entries", "blobs", "warm", "pathsets"]
395 .iter()
396 .map(|name| {
397 let metadata = match fs::metadata(root.join(name)) {
398 Ok(metadata) => metadata,
399 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
400 Err(error) => return Err(error.into()),
401 };
402 #[cfg(unix)]
403 use std::os::unix::fs::MetadataExt;
404 Ok(Some(DirectoryStamp {
405 modified: metadata.modified()?,
406 len: metadata.len(),
407 #[cfg(unix)]
408 identity: (
409 metadata.dev(),
410 metadata.ino(),
411 metadata.ctime(),
412 metadata.ctime_nsec(),
413 ),
414 }))
415 })
416 .collect()
417}
418
419fn settled(stamps: &[Option<DirectoryStamp>], now: SystemTime) -> bool {
420 stamps.iter().flatten().all(|stamp| {
421 let old_enough = |time| {
422 now.duration_since(time)
423 .is_ok_and(|age| age >= STAMP_SETTLE)
424 };
425 if !old_enough(stamp.modified) {
426 return false;
427 }
428 #[cfg(unix)]
429 {
430 let (_, _, seconds, nanos) = stamp.identity;
431 if seconds < 0 || !(0..1_000_000_000).contains(&nanos) {
432 return false;
433 }
434 let changed = SystemTime::UNIX_EPOCH + Duration::new(seconds as u64, nanos as u32);
435 if !old_enough(changed) {
436 return false;
437 }
438 }
439 true
440 })
441}
442
443fn cached_size(root: &Path) -> Result<u64> {
444 let memo_path = root.join("locks").join(".size");
445 let now = SystemTime::now();
446 let before = directory_stamps(root)?;
447 if settled(&before, now)
448 && let Some(memo) = fs::read(&memo_path)
449 .ok()
450 .and_then(|bytes| serde_json::from_slice::<SizeMemo>(&bytes).ok())
451 && now
452 .duration_since(memo.measured)
453 .is_ok_and(|age| age < IN_FLIGHT)
454 && memo.stamps == before
455 {
456 return Ok(memo.size);
457 }
458 let size = ["entries", "blobs", "warm"]
459 .iter()
460 .zip(&before)
461 .filter(|(_, stamp)| stamp.is_some())
462 .try_fold(0u64, |total, (name, _)| {
463 directory_size(&root.join(name)).map(|size| total.saturating_add(size))
464 })?;
465 let after = directory_stamps(root)?;
466 if before == after && settled(&after, now) {
467 let memo = SizeMemo {
468 measured: now,
469 stamps: after,
470 size,
471 };
472 let mut temporary = tempfile::NamedTempFile::new_in(root.join("locks"))?;
475 serde_json::to_writer(temporary.as_file_mut(), &memo)?;
476 temporary.persist(&memo_path).map_err(|error| error.error)?;
477 } else {
478 let _ = fs::remove_file(memo_path);
479 }
480 Ok(size)
481}
482
483fn directory_size(directory: &Path) -> Result<u64> {
484 let mut size = 0;
485 for item in fs::read_dir(directory)? {
486 size += item?.metadata()?.len();
487 }
488 Ok(size)
489}
490
491#[cfg(test)]
492mod tests {
493 use super::*;
494
495 fn fixture() -> tempfile::TempDir {
496 let root = tempfile::tempdir().unwrap();
497 for name in ["entries", "blobs", "locks"] {
498 fs::create_dir(root.path().join(name)).unwrap();
499 }
500 fs::write(root.path().join("entries/a.json"), b"bad record").unwrap();
501 fs::write(root.path().join("blobs/a"), [0; 100]).unwrap();
502 fs::write(root.path().join("locks/.prune"), []).unwrap();
503 root
504 }
505
506 fn cited() -> (tempfile::TempDir, String) {
508 let root = fixture();
509 let blob = "b".repeat(64);
510 fs::write(root.path().join("blobs").join(&blob), [0; 100]).unwrap();
511 fs::write(
512 root.path().join("entries/a.json"),
513 serde_json::to_vec(&serde_json::json!({"log": blob})).unwrap(),
514 )
515 .unwrap();
516 let old = SystemTime::now() - Duration::from_secs(2 * 60 * 60);
517 for path in [format!("blobs/{blob}"), "entries/a.json".into()] {
518 File::options()
519 .write(true)
520 .open(root.path().join(path))
521 .unwrap()
522 .set_modified(old)
523 .unwrap();
524 }
525 (root, blob)
526 }
527
528 fn prune_away(root: &Path, limit: u64, blob: &str) {
532 let deadline = std::time::Instant::now() + Duration::from_secs(5);
533 while root.join("blobs").join(blob).exists() && std::time::Instant::now() < deadline {
534 prune(root, limit).unwrap();
535 std::thread::sleep(Duration::from_millis(10));
536 }
537 assert!(!root.join("blobs").join(blob).exists());
538 }
539
540 #[test]
541 fn a_pinned_blob_survives_evicting_the_last_entry_that_cites_it() {
542 let (root, blob) = cited();
543 let root = root.path();
544 let pin = pin(root, &BTreeSet::from([blob.clone()])).unwrap();
545 assert_eq!(prune(root, 0).unwrap().entries, 1);
546 assert!(root.join("blobs").join(&blob).exists());
547 drop(pin);
548 prune_away(root, 0, &blob);
549 }
550
551 #[test]
552 fn a_pinned_old_blob_no_entry_cites_survives() {
553 let (root, blob) = cited();
554 let root = root.path();
555 fs::remove_file(root.join("entries/a.json")).unwrap();
556 let pin = pin(root, &BTreeSet::from([blob.clone()])).unwrap();
557 prune(root, 1 << 30).unwrap();
558 assert!(root.join("blobs").join(&blob).exists());
559 drop(pin);
560 assert!(
561 fs::read_dir(pins(root))
562 .unwrap()
563 .all(|item| { item.unwrap().file_name() == ".gate" })
564 );
565 prune_away(root, 1 << 30, &blob);
566 }
567
568 #[test]
569 fn an_abandoned_pin_neither_keeps_blobs_nor_stays() {
570 let (root, blob) = cited();
571 let root = root.path();
572 fs::create_dir_all(pins(root)).unwrap();
573 fs::write(pins(root).join("gone.lock"), []).unwrap();
574 fs::write(pins(root).join("gone.blobs"), format!("{blob}\n")).unwrap();
575 prune(root, 0).unwrap();
576 assert!(!root.join("blobs").join(&blob).exists());
577 assert!(!pins(root).join("gone.lock").exists());
578 assert!(!pins(root).join("gone.blobs").exists());
579 }
580
581 #[test]
582 fn a_pin_waits_for_a_prune_that_is_deleting() {
583 let (root, blob) = cited();
584 let root = root.path().to_owned();
585 let gate = pin_gate(&root).unwrap();
586 gate.lock_exclusive().unwrap();
587 let (sender, receiver) = std::sync::mpsc::channel();
588 let pinning = std::thread::spawn({
589 let root = root.clone();
590 move || {
591 let pin = pin(&root, &BTreeSet::from([blob])).unwrap();
592 sender.send(()).unwrap();
593 pin
594 }
595 });
596 assert!(receiver.recv_timeout(Duration::from_millis(200)).is_err());
597 drop(gate);
598 receiver.recv_timeout(Duration::from_secs(5)).unwrap();
599 drop(pinning.join().unwrap());
600 }
601
602 #[test]
603 fn automatic_size_memo_reuses_and_reconciles() {
604 let root = fixture();
605 let root = root.path();
606 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
608 assert!(!root.join("locks/.size").exists());
609 std::thread::sleep(STAMP_SETTLE + Duration::from_millis(100));
610 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
611 assert!(root.join("locks/.size").exists());
612
613 fs::write(root.join("blobs/a"), [0; 200]).unwrap();
616 touch(&root.join("entries/a.json"));
617 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
618 assert_eq!(prune(root, 1_000).unwrap().size, 210);
619 assert!(!root.join("locks/.size").exists());
620
621 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
622 fs::write(root.join("locks/.size"), b"corrupt").unwrap();
623 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
624 let mut memo: SizeMemo =
625 serde_json::from_slice(&fs::read(root.join("locks/.size")).unwrap()).unwrap();
626 memo.measured = SystemTime::UNIX_EPOCH;
627 memo.size = 0;
628 fs::write(root.join("locks/.size"), serde_json::to_vec(&memo).unwrap()).unwrap();
629 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
630
631 let temp = root.join("entries/new.tmp");
633 fs::write(&temp, [0; 30]).unwrap();
634 fs::rename(temp, root.join("entries/new.json")).unwrap();
635 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 240);
636 fs::create_dir(root.join("warm")).unwrap();
637 fs::write(root.join("warm/a.json"), [0; 40]).unwrap();
638 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 280);
639 std::thread::sleep(STAMP_SETTLE + Duration::from_millis(100));
640 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 280);
641 assert!(root.join("locks/.size").exists());
642 let result = prune_automatic(root, 250).unwrap();
644 assert!(result.entries > 0);
645 assert!(result.size <= 250);
646 }
647
648 #[test]
649 fn old_prune_marker_forces_full_reconciliation() {
650 let root = fixture();
651 let root = root.path();
652 let marker = File::options()
653 .write(true)
654 .open(root.join("locks/.prune"))
655 .unwrap();
656 marker.set_modified(SystemTime::UNIX_EPOCH).unwrap();
657 let result = prune_automatic(root, 100).unwrap();
658 assert_eq!(result.entries, 1);
659 assert_eq!(result.size, 100);
660 }
661
662 #[test]
663 fn entries_unused_for_a_week_go_and_eviction_stops_a_tenth_below_the_limit() {
664 let root = tempfile::tempdir().unwrap();
665 let root = root.path();
666 for name in ["entries", "blobs", "locks"] {
667 fs::create_dir(root.join(name)).unwrap();
668 }
669 let now = SystemTime::now();
670 for (name, age_days, size) in [
671 ("stale", 8, 10),
672 ("old", 3, 15),
673 ("middle", 2, 100),
674 ("new", 1, 180),
675 ] {
676 let path = root.join("entries").join(format!("{name}.json"));
677 fs::write(&path, vec![0; size]).unwrap();
678 File::options()
679 .write(true)
680 .open(path)
681 .unwrap()
682 .set_modified(now - Duration::from_secs(age_days * 24 * 60 * 60))
683 .unwrap();
684 }
685 let left = || {
686 let mut names: Vec<String> = fs::read_dir(root.join("entries"))
687 .unwrap()
688 .map(|item| item.unwrap().file_name().to_string_lossy().into_owned())
689 .collect();
690 names.sort();
691 names
692 };
693 let prune = |limit: u64| {
696 let deadline = std::time::Instant::now() + Duration::from_secs(5);
697 loop {
698 let pruned = prune(root, limit).unwrap();
699 if pruned != Pruned::default() || std::time::Instant::now() > deadline {
700 return pruned;
701 }
702 std::thread::sleep(Duration::from_millis(10));
703 }
704 };
705 assert_eq!(prune(u64::MAX).entries, 1);
707 assert_eq!(left(), ["middle.json", "new.json", "old.json"]);
708 assert_eq!(prune(290).entries, 2);
711 assert_eq!(left(), ["new.json"]);
712 }
713
714 #[test]
715 fn parses_sizes_like_nx() {
716 assert_eq!(parse_size("1024").unwrap(), 1024);
717 assert_eq!(parse_size("1KB").unwrap(), 1024);
718 assert_eq!(parse_size("1.5 MB").unwrap(), 1_572_864);
719 assert_eq!(parse_size("2GB").unwrap(), 2 << 30);
720 assert_eq!(parse_size("0").unwrap(), 0);
721 for bad in ["", "1TB", "1.2.3MB", "MB", "-1"] {
722 assert!(parse_size(bad).is_err(), "{bad}");
723 }
724 }
725}