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_json::Value;
17
18const IN_FLIGHT: Duration = Duration::from_secs(60 * 60);
20
21#[derive(Debug, Default, PartialEq, Eq)]
22pub struct Pruned {
23 pub size: u64,
25 pub entries: usize,
26 pub blobs: usize,
27 pub freed: u64,
28}
29
30pub fn parse_size(text: &str) -> Result<u64> {
33 let text = text.trim();
34 let (number, unit) = match text.find(|c: char| !(c.is_ascii_digit() || c == '.')) {
35 Some(at) => (&text[..at], text[at..].trim_start()),
36 None => (text, ""),
37 };
38 let factor: u64 = match unit {
39 "" | "B" => 1,
40 "KB" => 1 << 10,
41 "MB" => 1 << 20,
42 "GB" => 1 << 30,
43 _ => bail!("invalid cache size {text:?}: use a number with an optional KB, MB or GB"),
44 };
45 if number.matches('.').count() > 1 || number.is_empty() {
46 bail!("invalid cache size {text:?}");
47 }
48 let number: f64 = number
49 .parse()
50 .with_context(|| format!("invalid cache size {text:?}"))?;
51 Ok((number * factor as f64) as u64)
52}
53
54pub fn max_size(configured: Option<&Value>, cache: &Path) -> Result<Option<u64>> {
58 let configured = match std::env::var("NX_MAX_CACHE_SIZE") {
59 Ok(value) => Some(value),
60 Err(_) => configured.map(|value| match value {
61 Value::String(text) => text.clone(),
62 other => other.to_string(),
63 }),
64 };
65 let limit = match configured {
66 Some(text) => parse_size(&text)?,
67 None => {
68 let existing = cache
69 .ancestors()
70 .find(|path| path.exists())
71 .context("cache path has no existing ancestor")?;
72 fs2::total_space(existing)? / 10
73 }
74 };
75 Ok((limit > 0).then_some(limit))
76}
77
78pub(crate) fn touch(manifest: &Path) {
80 if let Ok(file) = File::options().write(true).open(manifest) {
81 let _ = file.set_modified(SystemTime::now());
82 }
83}
84
85struct Entry {
86 path: std::path::PathBuf,
87 used: SystemTime,
88 size: u64,
89 blobs: BTreeSet<String>,
90}
91
92pub fn prune(root: &Path, limit: u64) -> Result<Pruned> {
100 if !root.join("entries").is_dir() {
101 return Ok(Pruned::default());
102 }
103 fs::create_dir_all(root.join("locks"))?;
104 let marker = root.join("locks").join(".prune");
105 let recent = fs::metadata(&marker)
106 .ok()
107 .and_then(|metadata| metadata.modified().ok())
108 .and_then(|modified| SystemTime::now().duration_since(modified).ok())
109 .is_some_and(|age| age < IN_FLIGHT);
110 if recent {
111 let warm = root.join("warm");
112 let size = directory_size(&root.join("entries"))?
113 + directory_size(&root.join("blobs"))?
114 + if warm.is_dir() {
115 directory_size(&warm)?
116 } else {
117 0
118 };
119 if size <= limit {
120 return Ok(Pruned {
121 size,
122 ..Pruned::default()
123 });
124 }
125 }
126 let guard = File::options()
127 .create(true)
128 .truncate(false)
129 .write(true)
130 .open(&marker)?;
131 if guard.try_lock_exclusive().is_err() {
132 return Ok(Pruned::default());
133 }
134 let _ = guard.set_modified(SystemTime::now());
135 let now = SystemTime::now();
136 let age = |metadata: &fs::Metadata| {
137 metadata
138 .modified()
139 .ok()
140 .and_then(|modified| now.duration_since(modified).ok())
141 .unwrap_or_default()
142 };
143 let mut entries = Vec::new();
144 let mut references: BTreeMap<String, usize> = BTreeMap::new();
145 let warm = fs::read_dir(root.join("warm")).into_iter().flatten();
147 for item in fs::read_dir(root.join("entries"))?.chain(warm) {
148 let item = item?;
149 let metadata = item.metadata()?;
150 let blobs = fs::read(item.path())
152 .ok()
153 .and_then(|bytes| crate::record::StoredRecord::read(&bytes).ok())
154 .map(|record| record.blobs())
155 .unwrap_or_default();
156 for blob in &blobs {
157 *references.entry(blob.clone()).or_default() += 1;
158 }
159 entries.push(Entry {
160 path: item.path(),
161 used: metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH),
162 size: metadata.len(),
163 blobs,
164 });
165 }
166 let mut blob_sizes = BTreeMap::new();
167 let mut pruned = Pruned::default();
168 for item in fs::read_dir(root.join("blobs"))? {
169 let item = item?;
170 let metadata = item.metadata()?;
171 let name = item.file_name().to_string_lossy().into_owned();
172 if !references.contains_key(&name) && age(&metadata) > IN_FLIGHT {
173 if fs::remove_file(item.path()).is_ok() {
174 pruned.blobs += 1;
175 pruned.freed += metadata.len();
176 }
177 continue;
178 }
179 blob_sizes.insert(name, metadata.len());
180 }
181 let mut size: u64 =
182 entries.iter().map(|entry| entry.size).sum::<u64>() + blob_sizes.values().sum::<u64>();
183 entries.sort_by_key(|entry| entry.used);
184 for entry in entries {
185 if size <= limit {
186 break;
187 }
188 if fs::remove_file(&entry.path).is_err() {
189 continue;
190 }
191 pruned.entries += 1;
192 pruned.freed += entry.size;
193 size -= entry.size;
194 for blob in entry.blobs {
195 let count = references.get_mut(&blob).expect("counted above");
196 *count -= 1;
197 if *count == 0
198 && let Some(blob_size) = blob_sizes.remove(&blob)
199 && fs::remove_file(root.join("blobs").join(&blob)).is_ok()
200 {
201 pruned.blobs += 1;
202 pruned.freed += blob_size;
203 size -= blob_size;
204 }
205 }
206 }
207 for item in fs::read_dir(root.join("tmp"))
209 .into_iter()
210 .flatten()
211 .flatten()
212 {
213 if item
214 .metadata()
215 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
216 {
217 let _ = fs::remove_dir_all(item.path()).or_else(|_| fs::remove_file(item.path()));
218 }
219 }
220 for item in fs::read_dir(root.join("locks"))?.flatten() {
222 let name = item.file_name().to_string_lossy().into_owned();
223 if name.starts_with('.') || root.join("entries").join(format!("{name}.json")).exists() {
224 continue;
225 }
226 if item
227 .metadata()
228 .is_ok_and(|metadata| age(&metadata) > IN_FLIGHT)
229 && let Ok(file) = File::open(item.path())
230 && file.try_lock_exclusive().is_ok()
231 {
232 let _ = fs::remove_file(item.path());
233 }
234 }
235 pruned.size = size;
236 Ok(pruned)
237}
238
239fn directory_size(directory: &Path) -> Result<u64> {
240 let mut size = 0;
241 for item in fs::read_dir(directory)? {
242 size += item?.metadata()?.len();
243 }
244 Ok(size)
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250
251 #[test]
252 fn parses_sizes_like_nx() {
253 assert_eq!(parse_size("1024").unwrap(), 1024);
254 assert_eq!(parse_size("1KB").unwrap(), 1024);
255 assert_eq!(parse_size("1.5 MB").unwrap(), 1_572_864);
256 assert_eq!(parse_size("2GB").unwrap(), 2 << 30);
257 assert_eq!(parse_size("0").unwrap(), 0);
258 for bad in ["", "1TB", "1.2.3MB", "MB", "-1"] {
259 assert!(parse_size(bad).is_err(), "{bad}");
260 }
261 }
262}