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