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();
268 let pathsets = fs::read_dir(root.join("pathsets")).into_iter().flatten();
269 let projections = fs::read_dir(root.join("projections")).into_iter().flatten();
270 for item in fs::read_dir(root.join("entries"))?
271 .chain(warm)
272 .chain(pathsets)
273 .chain(projections)
274 {
275 let item = item?;
276 let metadata = item.metadata()?;
277 let blobs = fs::read(item.path())
279 .ok()
280 .and_then(|bytes| crate::record::StoredRecord::read(&bytes).ok())
281 .map(|record| record.blobs())
282 .unwrap_or_default();
283 for blob in &blobs {
284 *references.entry(blob.clone()).or_default() += 1;
285 }
286 entries.push(Entry {
287 path: item.path(),
288 used: metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH),
289 size: metadata.len(),
290 blobs,
291 });
292 }
293 for blob in pinned(root) {
295 *references.entry(blob).or_default() += 1;
296 }
297 let mut blob_sizes = BTreeMap::new();
298 let mut pruned = Pruned::default();
299 for item in fs::read_dir(root.join("blobs"))? {
300 let item = item?;
301 let metadata = item.metadata()?;
302 let name = item.file_name().to_string_lossy().into_owned();
303 if !references.contains_key(&name) && age(&metadata) > IN_FLIGHT {
304 if fs::remove_file(item.path()).is_ok() {
305 pruned.blobs += 1;
306 pruned.freed += metadata.len();
307 }
308 continue;
309 }
310 blob_sizes.insert(name, metadata.len());
311 }
312 let mut size: u64 =
313 entries.iter().map(|entry| entry.size).sum::<u64>() + blob_sizes.values().sum::<u64>();
314 entries.sort_by_key(|entry| entry.used);
315 let target = if size > limit {
317 limit - limit / 10
318 } else {
319 limit
320 };
321 for entry in entries {
322 let unused = now.duration_since(entry.used).is_ok_and(|age| age > UNUSED);
323 if size <= target && !unused {
324 break;
325 }
326 if fs::remove_file(&entry.path).is_err() {
327 continue;
328 }
329 pruned.entries += 1;
330 pruned.freed += entry.size;
331 size -= entry.size;
332 for blob in entry.blobs {
333 let count = references.get_mut(&blob).expect("counted above");
334 *count -= 1;
335 if *count == 0
336 && let Some(blob_size) = blob_sizes.remove(&blob)
337 && fs::remove_file(root.join("blobs").join(&blob)).is_ok()
338 {
339 pruned.blobs += 1;
340 pruned.freed += blob_size;
341 size -= blob_size;
342 }
343 }
344 }
345 drop(gate);
346 for item in fs::read_dir(root.join("tmp"))
348 .into_iter()
349 .flatten()
350 .flatten()
351 {
352 if item
353 .metadata()
354 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
355 {
356 let _ = fs::remove_dir_all(item.path()).or_else(|_| fs::remove_file(item.path()));
357 }
358 }
359 for item in fs::read_dir(root.join("locks"))?.flatten() {
361 let name = item.file_name().to_string_lossy().into_owned();
362 if name.starts_with('.') || root.join("entries").join(format!("{name}.json")).exists() {
363 continue;
364 }
365 if item
366 .metadata()
367 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
368 && let Ok(file) = File::open(item.path())
369 && file.try_lock_exclusive().is_ok()
370 {
371 let _ = fs::remove_file(item.path());
372 }
373 }
374 pruned.size = if automatic { cached_size(root)? } else { size };
375 Ok(pruned)
376}
377
378const STAMP_SETTLE: Duration = Duration::from_secs(2);
380
381#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
382struct DirectoryStamp {
383 modified: SystemTime,
384 len: u64,
385 #[cfg(unix)]
386 identity: (u64, u64, i64, i64),
387}
388
389#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
390struct SizeMemo {
391 measured: SystemTime,
392 stamps: Vec<Option<DirectoryStamp>>,
393 size: u64,
394}
395
396fn directory_stamps(root: &Path) -> Result<Vec<Option<DirectoryStamp>>> {
397 ["entries", "blobs", "warm", "pathsets", "projections"]
398 .iter()
399 .map(|name| {
400 let metadata = match fs::metadata(root.join(name)) {
401 Ok(metadata) => metadata,
402 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
403 Err(error) => return Err(error.into()),
404 };
405 #[cfg(unix)]
406 use std::os::unix::fs::MetadataExt;
407 Ok(Some(DirectoryStamp {
408 modified: metadata.modified()?,
409 len: metadata.len(),
410 #[cfg(unix)]
411 identity: (
412 metadata.dev(),
413 metadata.ino(),
414 metadata.ctime(),
415 metadata.ctime_nsec(),
416 ),
417 }))
418 })
419 .collect()
420}
421
422fn settled(stamps: &[Option<DirectoryStamp>], now: SystemTime) -> bool {
423 stamps.iter().flatten().all(|stamp| {
424 let old_enough = |time| {
425 now.duration_since(time)
426 .is_ok_and(|age| age >= STAMP_SETTLE)
427 };
428 if !old_enough(stamp.modified) {
429 return false;
430 }
431 #[cfg(unix)]
432 {
433 let (_, _, seconds, nanos) = stamp.identity;
434 if seconds < 0 || !(0..1_000_000_000).contains(&nanos) {
435 return false;
436 }
437 let changed = SystemTime::UNIX_EPOCH + Duration::new(seconds as u64, nanos as u32);
438 if !old_enough(changed) {
439 return false;
440 }
441 }
442 true
443 })
444}
445
446fn cached_size(root: &Path) -> Result<u64> {
447 let memo_path = root.join("locks").join(".size");
448 let now = SystemTime::now();
449 let before = directory_stamps(root)?;
450 if settled(&before, now)
451 && let Some(memo) = fs::read(&memo_path)
452 .ok()
453 .and_then(|bytes| serde_json::from_slice::<SizeMemo>(&bytes).ok())
454 && now
455 .duration_since(memo.measured)
456 .is_ok_and(|age| age < IN_FLIGHT)
457 && memo.stamps == before
458 {
459 return Ok(memo.size);
460 }
461 let size = ["entries", "blobs", "warm", "pathsets", "projections"]
462 .iter()
463 .zip(&before)
464 .filter(|(_, stamp)| stamp.is_some())
465 .try_fold(0u64, |total, (name, _)| {
466 directory_size(&root.join(name)).map(|size| total.saturating_add(size))
467 })?;
468 let after = directory_stamps(root)?;
469 if before == after && settled(&after, now) {
470 let memo = SizeMemo {
471 measured: now,
472 stamps: after,
473 size,
474 };
475 let mut temporary = tempfile::NamedTempFile::new_in(root.join("locks"))?;
478 serde_json::to_writer(temporary.as_file_mut(), &memo)?;
479 temporary.persist(&memo_path).map_err(|error| error.error)?;
480 } else {
481 let _ = fs::remove_file(memo_path);
482 }
483 Ok(size)
484}
485
486fn directory_size(directory: &Path) -> Result<u64> {
487 let mut size = 0;
488 for item in fs::read_dir(directory)? {
489 size += item?.metadata()?.len();
490 }
491 Ok(size)
492}
493
494#[cfg(test)]
495mod tests {
496 use super::*;
497
498 fn fixture() -> tempfile::TempDir {
499 let root = tempfile::tempdir().unwrap();
500 for name in ["entries", "blobs", "locks"] {
501 fs::create_dir(root.path().join(name)).unwrap();
502 }
503 fs::write(root.path().join("entries/a.json"), b"bad record").unwrap();
504 fs::write(root.path().join("blobs/a"), [0; 100]).unwrap();
505 fs::write(root.path().join("locks/.prune"), []).unwrap();
506 root
507 }
508
509 fn cited() -> (tempfile::TempDir, String) {
511 let root = fixture();
512 let blob = "b".repeat(64);
513 fs::write(root.path().join("blobs").join(&blob), [0; 100]).unwrap();
514 fs::write(
515 root.path().join("entries/a.json"),
516 serde_json::to_vec(&serde_json::json!({"log": blob})).unwrap(),
517 )
518 .unwrap();
519 let old = SystemTime::now() - Duration::from_secs(2 * 60 * 60);
520 for path in [format!("blobs/{blob}"), "entries/a.json".into()] {
521 File::options()
522 .write(true)
523 .open(root.path().join(path))
524 .unwrap()
525 .set_modified(old)
526 .unwrap();
527 }
528 (root, blob)
529 }
530
531 fn prune_away(root: &Path, limit: u64, blob: &str) {
535 let deadline = std::time::Instant::now() + Duration::from_secs(5);
536 while root.join("blobs").join(blob).exists() && std::time::Instant::now() < deadline {
537 prune(root, limit).unwrap();
538 std::thread::sleep(Duration::from_millis(10));
539 }
540 assert!(!root.join("blobs").join(blob).exists());
541 }
542
543 #[test]
544 fn a_pinned_blob_survives_evicting_the_last_entry_that_cites_it() {
545 let (root, blob) = cited();
546 let root = root.path();
547 let pin = pin(root, &BTreeSet::from([blob.clone()])).unwrap();
548 assert_eq!(prune(root, 0).unwrap().entries, 1);
549 assert!(root.join("blobs").join(&blob).exists());
550 drop(pin);
551 prune_away(root, 0, &blob);
552 }
553
554 #[test]
555 fn a_pinned_old_blob_no_entry_cites_survives() {
556 let (root, blob) = cited();
557 let root = root.path();
558 fs::remove_file(root.join("entries/a.json")).unwrap();
559 let pin = pin(root, &BTreeSet::from([blob.clone()])).unwrap();
560 prune(root, 1 << 30).unwrap();
561 assert!(root.join("blobs").join(&blob).exists());
562 drop(pin);
563 assert!(
564 fs::read_dir(pins(root))
565 .unwrap()
566 .all(|item| { item.unwrap().file_name() == ".gate" })
567 );
568 prune_away(root, 1 << 30, &blob);
569 }
570
571 #[test]
572 fn an_abandoned_pin_neither_keeps_blobs_nor_stays() {
573 let (root, blob) = cited();
574 let root = root.path();
575 fs::create_dir_all(pins(root)).unwrap();
576 fs::write(pins(root).join("gone.lock"), []).unwrap();
577 fs::write(pins(root).join("gone.blobs"), format!("{blob}\n")).unwrap();
578 prune(root, 0).unwrap();
579 assert!(!root.join("blobs").join(&blob).exists());
580 assert!(!pins(root).join("gone.lock").exists());
581 assert!(!pins(root).join("gone.blobs").exists());
582 }
583
584 #[test]
585 fn a_pin_waits_for_a_prune_that_is_deleting() {
586 let (root, blob) = cited();
587 let root = root.path().to_owned();
588 let gate = pin_gate(&root).unwrap();
589 gate.lock_exclusive().unwrap();
590 let (sender, receiver) = std::sync::mpsc::channel();
591 let pinning = std::thread::spawn({
592 let root = root.clone();
593 move || {
594 let pin = pin(&root, &BTreeSet::from([blob])).unwrap();
595 sender.send(()).unwrap();
596 pin
597 }
598 });
599 assert!(receiver.recv_timeout(Duration::from_millis(200)).is_err());
600 drop(gate);
601 receiver.recv_timeout(Duration::from_secs(5)).unwrap();
602 drop(pinning.join().unwrap());
603 }
604
605 #[test]
606 fn automatic_size_memo_reuses_and_reconciles() {
607 let root = fixture();
608 let root = root.path();
609 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
611 assert!(!root.join("locks/.size").exists());
612 std::thread::sleep(STAMP_SETTLE + Duration::from_millis(100));
613 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
614 assert!(root.join("locks/.size").exists());
615
616 fs::write(root.join("blobs/a"), [0; 200]).unwrap();
619 touch(&root.join("entries/a.json"));
620 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
621 assert_eq!(prune(root, 1_000).unwrap().size, 210);
622 assert!(!root.join("locks/.size").exists());
623
624 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
625 fs::write(root.join("locks/.size"), b"corrupt").unwrap();
626 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
627 let mut memo: SizeMemo =
628 serde_json::from_slice(&fs::read(root.join("locks/.size")).unwrap()).unwrap();
629 memo.measured = SystemTime::UNIX_EPOCH;
630 memo.size = 0;
631 fs::write(root.join("locks/.size"), serde_json::to_vec(&memo).unwrap()).unwrap();
632 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
633
634 let temp = root.join("entries/new.tmp");
636 fs::write(&temp, [0; 30]).unwrap();
637 fs::rename(temp, root.join("entries/new.json")).unwrap();
638 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 240);
639 fs::create_dir(root.join("warm")).unwrap();
640 fs::write(root.join("warm/a.json"), [0; 40]).unwrap();
641 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 280);
642 std::thread::sleep(STAMP_SETTLE + Duration::from_millis(100));
643 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 280);
644 assert!(root.join("locks/.size").exists());
645 let result = prune_automatic(root, 250).unwrap();
647 assert!(result.entries > 0);
648 assert!(result.size <= 250);
649 }
650
651 #[test]
652 fn old_prune_marker_forces_full_reconciliation() {
653 let root = fixture();
654 let root = root.path();
655 let marker = File::options()
656 .write(true)
657 .open(root.join("locks/.prune"))
658 .unwrap();
659 marker.set_modified(SystemTime::UNIX_EPOCH).unwrap();
660 let result = prune_automatic(root, 100).unwrap();
661 assert_eq!(result.entries, 1);
662 assert_eq!(result.size, 100);
663 }
664
665 #[test]
666 fn entries_unused_for_a_week_go_and_eviction_stops_a_tenth_below_the_limit() {
667 let root = tempfile::tempdir().unwrap();
668 let root = root.path();
669 for name in ["entries", "blobs", "locks"] {
670 fs::create_dir(root.join(name)).unwrap();
671 }
672 let now = SystemTime::now();
673 for (name, age_days, size) in [
674 ("stale", 8, 10),
675 ("old", 3, 15),
676 ("middle", 2, 100),
677 ("new", 1, 180),
678 ] {
679 let path = root.join("entries").join(format!("{name}.json"));
680 fs::write(&path, vec![0; size]).unwrap();
681 File::options()
682 .write(true)
683 .open(path)
684 .unwrap()
685 .set_modified(now - Duration::from_secs(age_days * 24 * 60 * 60))
686 .unwrap();
687 }
688 let left = || {
689 let mut names: Vec<String> = fs::read_dir(root.join("entries"))
690 .unwrap()
691 .map(|item| item.unwrap().file_name().to_string_lossy().into_owned())
692 .collect();
693 names.sort();
694 names
695 };
696 let prune = |limit: u64| {
699 let deadline = std::time::Instant::now() + Duration::from_secs(5);
700 loop {
701 let pruned = prune(root, limit).unwrap();
702 if pruned != Pruned::default() || std::time::Instant::now() > deadline {
703 return pruned;
704 }
705 std::thread::sleep(Duration::from_millis(10));
706 }
707 };
708 assert_eq!(prune(u64::MAX).entries, 1);
710 assert_eq!(left(), ["middle.json", "new.json", "old.json"]);
711 assert_eq!(prune(290).entries, 2);
714 assert_eq!(left(), ["new.json"]);
715 }
716
717 #[test]
718 fn parses_sizes_like_nx() {
719 assert_eq!(parse_size("1024").unwrap(), 1024);
720 assert_eq!(parse_size("1KB").unwrap(), 1024);
721 assert_eq!(parse_size("1.5 MB").unwrap(), 1_572_864);
722 assert_eq!(parse_size("2GB").unwrap(), 2 << 30);
723 assert_eq!(parse_size("0").unwrap(), 0);
724 for bad in ["", "1TB", "1.2.3MB", "MB", "-1"] {
725 assert!(parse_size(bad).is_err(), "{bad}");
726 }
727 }
728}