1use std::collections::{BTreeMap, BTreeSet};
10use std::fs::{self, File};
11use std::path::Path;
12use std::time::{Duration, SystemTime};
13
14use anyhow::{Context, Result, bail};
15use fs2::FileExt;
16use serde::{Deserialize, Serialize};
17use serde_json::Value;
18
19const IN_FLIGHT: Duration = Duration::from_secs(60 * 60);
21
22#[derive(Debug, Default, PartialEq, Eq)]
23pub struct Pruned {
24 pub size: u64,
26 pub entries: usize,
27 pub blobs: usize,
28 pub freed: u64,
29}
30
31pub fn parse_size(text: &str) -> Result<u64> {
34 let text = text.trim();
35 let (number, unit) = match text.find(|c: char| !(c.is_ascii_digit() || c == '.')) {
36 Some(at) => (&text[..at], text[at..].trim_start()),
37 None => (text, ""),
38 };
39 let factor: u64 = match unit {
40 "" | "B" => 1,
41 "KB" => 1 << 10,
42 "MB" => 1 << 20,
43 "GB" => 1 << 30,
44 _ => bail!("invalid cache size {text:?}: use a number with an optional KB, MB or GB"),
45 };
46 if number.matches('.').count() > 1 || number.is_empty() {
47 bail!("invalid cache size {text:?}");
48 }
49 let number: f64 = number
50 .parse()
51 .with_context(|| format!("invalid cache size {text:?}"))?;
52 Ok((number * factor as f64) as u64)
53}
54
55pub fn max_size(configured: Option<&Value>, cache: &Path) -> Result<Option<u64>> {
59 let configured = match std::env::var("NX_MAX_CACHE_SIZE") {
60 Ok(value) => Some(value),
61 Err(_) => configured.map(|value| match value {
62 Value::String(text) => text.clone(),
63 other => other.to_string(),
64 }),
65 };
66 let limit = match configured {
67 Some(text) => parse_size(&text)?,
68 None => {
69 let existing = cache
70 .ancestors()
71 .find(|path| path.exists())
72 .context("cache path has no existing ancestor")?;
73 fs2::total_space(existing)? / 10
74 }
75 };
76 Ok((limit > 0).then_some(limit))
77}
78
79pub(crate) fn touch(manifest: &Path) {
81 if let Ok(file) = File::options().write(true).open(manifest) {
82 let _ = file.set_modified(SystemTime::now());
83 }
84}
85
86struct Entry {
87 path: std::path::PathBuf,
88 used: SystemTime,
89 size: u64,
90 blobs: BTreeSet<String>,
91}
92
93pub fn prune(root: &Path, limit: u64) -> Result<Pruned> {
100 prune_inner(root, limit, false)
101}
102
103pub fn prune_automatic(root: &Path, limit: u64) -> Result<Pruned> {
108 prune_inner(root, limit, true)
109}
110
111fn prune_inner(root: &Path, limit: u64, automatic: bool) -> Result<Pruned> {
112 if !root.join("entries").is_dir() {
113 return Ok(Pruned::default());
114 }
115 fs::create_dir_all(root.join("locks"))?;
116 let marker = root.join("locks").join(".prune");
117 let recent = fs::metadata(&marker)
118 .ok()
119 .and_then(|metadata| metadata.modified().ok())
120 .and_then(|modified| SystemTime::now().duration_since(modified).ok())
121 .is_some_and(|age| age < IN_FLIGHT);
122 if automatic && recent {
123 let size = cached_size(root)?;
124 if size <= limit {
125 return Ok(Pruned {
126 size,
127 ..Pruned::default()
128 });
129 }
130 }
131 let guard = File::options()
132 .create(true)
133 .truncate(false)
134 .write(true)
135 .open(&marker)?;
136 if guard.try_lock_exclusive().is_err() {
137 return Ok(Pruned::default());
138 }
139 let _ = fs::remove_file(root.join("locks").join(".size"));
142 let _ = guard.set_modified(SystemTime::now());
143 let now = SystemTime::now();
144 let age = |metadata: &fs::Metadata| {
145 metadata
146 .modified()
147 .ok()
148 .and_then(|modified| now.duration_since(modified).ok())
149 .unwrap_or_default()
150 };
151 let mut entries = Vec::new();
152 let mut references: BTreeMap<String, usize> = BTreeMap::new();
153 let warm = fs::read_dir(root.join("warm")).into_iter().flatten();
155 for item in fs::read_dir(root.join("entries"))?.chain(warm) {
156 let item = item?;
157 let metadata = item.metadata()?;
158 let blobs = fs::read(item.path())
160 .ok()
161 .and_then(|bytes| crate::record::StoredRecord::read(&bytes).ok())
162 .map(|record| record.blobs())
163 .unwrap_or_default();
164 for blob in &blobs {
165 *references.entry(blob.clone()).or_default() += 1;
166 }
167 entries.push(Entry {
168 path: item.path(),
169 used: metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH),
170 size: metadata.len(),
171 blobs,
172 });
173 }
174 let mut blob_sizes = BTreeMap::new();
175 let mut pruned = Pruned::default();
176 for item in fs::read_dir(root.join("blobs"))? {
177 let item = item?;
178 let metadata = item.metadata()?;
179 let name = item.file_name().to_string_lossy().into_owned();
180 if !references.contains_key(&name) && age(&metadata) > IN_FLIGHT {
181 if fs::remove_file(item.path()).is_ok() {
182 pruned.blobs += 1;
183 pruned.freed += metadata.len();
184 }
185 continue;
186 }
187 blob_sizes.insert(name, metadata.len());
188 }
189 let mut size: u64 =
190 entries.iter().map(|entry| entry.size).sum::<u64>() + blob_sizes.values().sum::<u64>();
191 entries.sort_by_key(|entry| entry.used);
192 for entry in entries {
193 if size <= limit {
194 break;
195 }
196 if fs::remove_file(&entry.path).is_err() {
197 continue;
198 }
199 pruned.entries += 1;
200 pruned.freed += entry.size;
201 size -= entry.size;
202 for blob in entry.blobs {
203 let count = references.get_mut(&blob).expect("counted above");
204 *count -= 1;
205 if *count == 0
206 && let Some(blob_size) = blob_sizes.remove(&blob)
207 && fs::remove_file(root.join("blobs").join(&blob)).is_ok()
208 {
209 pruned.blobs += 1;
210 pruned.freed += blob_size;
211 size -= blob_size;
212 }
213 }
214 }
215 for item in fs::read_dir(root.join("tmp"))
217 .into_iter()
218 .flatten()
219 .flatten()
220 {
221 if item
222 .metadata()
223 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
224 {
225 let _ = fs::remove_dir_all(item.path()).or_else(|_| fs::remove_file(item.path()));
226 }
227 }
228 for item in fs::read_dir(root.join("locks"))?.flatten() {
230 let name = item.file_name().to_string_lossy().into_owned();
231 if name.starts_with('.') || root.join("entries").join(format!("{name}.json")).exists() {
232 continue;
233 }
234 if item
235 .metadata()
236 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
237 && let Ok(file) = File::open(item.path())
238 && file.try_lock_exclusive().is_ok()
239 {
240 let _ = fs::remove_file(item.path());
241 }
242 }
243 pruned.size = if automatic { cached_size(root)? } else { size };
244 Ok(pruned)
245}
246
247const STAMP_SETTLE: Duration = Duration::from_secs(2);
249
250#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
251struct DirectoryStamp {
252 modified: SystemTime,
253 len: u64,
254 #[cfg(unix)]
255 identity: (u64, u64, i64, i64),
256}
257
258#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
259struct SizeMemo {
260 measured: SystemTime,
261 stamps: Vec<Option<DirectoryStamp>>,
262 size: u64,
263}
264
265fn directory_stamps(root: &Path) -> Result<Vec<Option<DirectoryStamp>>> {
266 ["entries", "blobs", "warm"]
267 .iter()
268 .map(|name| {
269 let metadata = match fs::metadata(root.join(name)) {
270 Ok(metadata) => metadata,
271 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
272 Err(error) => return Err(error.into()),
273 };
274 #[cfg(unix)]
275 use std::os::unix::fs::MetadataExt;
276 Ok(Some(DirectoryStamp {
277 modified: metadata.modified()?,
278 len: metadata.len(),
279 #[cfg(unix)]
280 identity: (
281 metadata.dev(),
282 metadata.ino(),
283 metadata.ctime(),
284 metadata.ctime_nsec(),
285 ),
286 }))
287 })
288 .collect()
289}
290
291fn settled(stamps: &[Option<DirectoryStamp>], now: SystemTime) -> bool {
292 stamps.iter().flatten().all(|stamp| {
293 let old_enough = |time| {
294 now.duration_since(time)
295 .is_ok_and(|age| age >= STAMP_SETTLE)
296 };
297 if !old_enough(stamp.modified) {
298 return false;
299 }
300 #[cfg(unix)]
301 {
302 let (_, _, seconds, nanos) = stamp.identity;
303 if seconds < 0 || !(0..1_000_000_000).contains(&nanos) {
304 return false;
305 }
306 let changed = SystemTime::UNIX_EPOCH + Duration::new(seconds as u64, nanos as u32);
307 if !old_enough(changed) {
308 return false;
309 }
310 }
311 true
312 })
313}
314
315fn cached_size(root: &Path) -> Result<u64> {
316 let memo_path = root.join("locks").join(".size");
317 let now = SystemTime::now();
318 let before = directory_stamps(root)?;
319 if settled(&before, now)
320 && let Some(memo) = fs::read(&memo_path)
321 .ok()
322 .and_then(|bytes| serde_json::from_slice::<SizeMemo>(&bytes).ok())
323 && now
324 .duration_since(memo.measured)
325 .is_ok_and(|age| age < IN_FLIGHT)
326 && memo.stamps == before
327 {
328 return Ok(memo.size);
329 }
330 let size = ["entries", "blobs", "warm"]
331 .iter()
332 .zip(&before)
333 .filter(|(_, stamp)| stamp.is_some())
334 .try_fold(0u64, |total, (name, _)| {
335 directory_size(&root.join(name)).map(|size| total.saturating_add(size))
336 })?;
337 let after = directory_stamps(root)?;
338 if before == after && settled(&after, now) {
339 let memo = SizeMemo {
340 measured: now,
341 stamps: after,
342 size,
343 };
344 let mut temporary = tempfile::NamedTempFile::new_in(root.join("locks"))?;
347 serde_json::to_writer(temporary.as_file_mut(), &memo)?;
348 temporary.persist(&memo_path).map_err(|error| error.error)?;
349 } else {
350 let _ = fs::remove_file(memo_path);
351 }
352 Ok(size)
353}
354
355fn directory_size(directory: &Path) -> Result<u64> {
356 let mut size = 0;
357 for item in fs::read_dir(directory)? {
358 size += item?.metadata()?.len();
359 }
360 Ok(size)
361}
362
363#[cfg(test)]
364mod tests {
365 use super::*;
366
367 fn fixture() -> tempfile::TempDir {
368 let root = tempfile::tempdir().unwrap();
369 for name in ["entries", "blobs", "locks"] {
370 fs::create_dir(root.path().join(name)).unwrap();
371 }
372 fs::write(root.path().join("entries/a.json"), b"bad record").unwrap();
373 fs::write(root.path().join("blobs/a"), [0; 100]).unwrap();
374 fs::write(root.path().join("locks/.prune"), []).unwrap();
375 root
376 }
377
378 #[test]
379 fn automatic_size_memo_reuses_and_reconciles() {
380 let root = fixture();
381 let root = root.path();
382 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
384 assert!(!root.join("locks/.size").exists());
385 std::thread::sleep(STAMP_SETTLE + Duration::from_millis(100));
386 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
387 assert!(root.join("locks/.size").exists());
388
389 fs::write(root.join("blobs/a"), [0; 200]).unwrap();
392 touch(&root.join("entries/a.json"));
393 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 110);
394 assert_eq!(prune(root, 1_000).unwrap().size, 210);
395 assert!(!root.join("locks/.size").exists());
396
397 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
398 fs::write(root.join("locks/.size"), b"corrupt").unwrap();
399 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
400 let mut memo: SizeMemo =
401 serde_json::from_slice(&fs::read(root.join("locks/.size")).unwrap()).unwrap();
402 memo.measured = SystemTime::UNIX_EPOCH;
403 memo.size = 0;
404 fs::write(root.join("locks/.size"), serde_json::to_vec(&memo).unwrap()).unwrap();
405 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 210);
406
407 let temp = root.join("entries/new.tmp");
409 fs::write(&temp, [0; 30]).unwrap();
410 fs::rename(temp, root.join("entries/new.json")).unwrap();
411 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 240);
412 fs::create_dir(root.join("warm")).unwrap();
413 fs::write(root.join("warm/a.json"), [0; 40]).unwrap();
414 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 280);
415 std::thread::sleep(STAMP_SETTLE + Duration::from_millis(100));
416 assert_eq!(prune_automatic(root, 1_000).unwrap().size, 280);
417 assert!(root.join("locks/.size").exists());
418 let result = prune_automatic(root, 250).unwrap();
420 assert!(result.entries > 0);
421 assert!(result.size <= 250);
422 }
423
424 #[test]
425 fn old_prune_marker_forces_full_reconciliation() {
426 let root = fixture();
427 let root = root.path();
428 let marker = File::options()
429 .write(true)
430 .open(root.join("locks/.prune"))
431 .unwrap();
432 marker.set_modified(SystemTime::UNIX_EPOCH).unwrap();
433 let result = prune_automatic(root, 100).unwrap();
434 assert_eq!(result.entries, 1);
435 assert_eq!(result.size, 100);
436 }
437
438 #[test]
439 fn parses_sizes_like_nx() {
440 assert_eq!(parse_size("1024").unwrap(), 1024);
441 assert_eq!(parse_size("1KB").unwrap(), 1024);
442 assert_eq!(parse_size("1.5 MB").unwrap(), 1_572_864);
443 assert_eq!(parse_size("2GB").unwrap(), 2 << 30);
444 assert_eq!(parse_size("0").unwrap(), 0);
445 for bad in ["", "1TB", "1.2.3MB", "MB", "-1"] {
446 assert!(parse_size(bad).is_err(), "{bad}");
447 }
448 }
449}