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