Skip to main content

qk_cache/
evict.rs

1//! Keeps the cache under its size limit by evicting the least recently used
2//! entries, as Nx's `maxCacheSize` does.
3//!
4//! An entry's last use is its manifest's modification time, refreshed on every
5//! hit. Blobs are shared between entries, so a blob is deleted only once no
6//! remaining manifest cites it. A blob no manifest cites may belong to an entry
7//! another run is still publishing, so it is kept until it is an hour old.
8
9use 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
19/// How long an unreferenced blob or scratch file may be part of a publish.
20const IN_FLIGHT: Duration = Duration::from_secs(60 * 60);
21
22#[derive(Debug, Default, PartialEq, Eq)]
23pub struct Pruned {
24    /// Bytes the cache holds afterwards.
25    pub size: u64,
26    pub entries: usize,
27    pub blobs: usize,
28    pub freed: u64,
29}
30
31/// Nx's `parseMaxCacheSize`: a number of bytes with an optional `KB`, `MB` or
32/// `GB` unit, in powers of 1024.
33pub 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
55/// The size limit: `NX_MAX_CACHE_SIZE`, then nx.json `maxCacheSize`, then, as
56/// in Nx, a tenth of the disk holding the cache. `None` when the limit is 0,
57/// which means unlimited.
58pub 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
79/// Refreshes an entry's last use.
80pub(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
93/// Evicts least recently used entries until the cache holds at most `limit`
94/// bytes, and removes what no entry needs. Another prune in progress makes
95/// this one return without doing anything.
96///
97/// Explicit pruning always remeasures the cache, including files changed in
98/// place by external tools. Normal runs use [`prune_automatic`].
99pub fn prune(root: &Path, limit: u64) -> Result<Pruned> {
100    prune_inner(root, limit, false)
101}
102
103/// Enforces the limit after a run, reusing size accounting while the cache's
104/// directories are unchanged. Normal publication replaces records atomically
105/// and blobs are immutable, so directory changes invalidate the memo. A full
106/// reconciliation still runs at least once an hour.
107pub 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    // Never associate a computed post-eviction size with stamps taken only
140    // after eviction: a concurrent publisher could have added uncounted bytes.
141    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    // Warm state records cite blobs as entries do, and age out the same way.
154    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        // Unreadable manifests are misses already; they only take space.
159        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    // Scratch space of runs that ended without cleaning up.
216    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    // Locks of keys without an entry, unless another run holds them.
229    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
247// Recent timestamps may have coarse resolution on the backing filesystem.
248const 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        // Concurrent writers can replace this memo safely: readers still compare
345        // its stamps with the current directories. Partial writes are misses.
346        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        // No memo can be trusted while directory timestamps are recent.
383        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        // Unsupported external in-place changes intentionally demonstrate that
390        // the automatic path reads only directory stamps, rather than each file.
391        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        // An atomic publication invalidates accounting immediately.
408        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        // A lower limit must still trigger eviction with an unchanged memo.
419        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}