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_json::Value;
17
18/// How long an unreferenced blob or scratch file may be part of a publish.
19const IN_FLIGHT: Duration = Duration::from_secs(60 * 60);
20
21#[derive(Debug, Default, PartialEq, Eq)]
22pub struct Pruned {
23    /// Bytes the cache holds afterwards.
24    pub size: u64,
25    pub entries: usize,
26    pub blobs: usize,
27    pub freed: u64,
28}
29
30/// Nx's `parseMaxCacheSize`: a number of bytes with an optional `KB`, `MB` or
31/// `GB` unit, in powers of 1024.
32pub 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
54/// The size limit: `NX_MAX_CACHE_SIZE`, then nx.json `maxCacheSize`, then, as
55/// in Nx, a tenth of the disk holding the cache. `None` when the limit is 0,
56/// which means unlimited.
57pub 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
78/// Refreshes an entry's last use.
79pub(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
92/// Evicts least recently used entries until the cache holds at most `limit`
93/// bytes, and removes what no entry needs. Another prune in progress makes
94/// this one return without doing anything.
95///
96/// Runs call this after every run, so a cache under its limit is only summed:
97/// the full pass, which reads every manifest, runs when the cache is over its
98/// limit or when the last full pass is more than an hour old.
99pub 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    // Warm state records cite blobs as entries do, and age out the same way.
146    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        // Unreadable manifests are misses already; they only take space.
151        let blobs = fs::read(item.path())
152            .ok()
153            .and_then(|bytes| serde_json::from_slice::<Value>(&bytes).ok())
154            .map(|manifest| manifest_blobs(&manifest))
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    // Scratch space of runs that ended without cleaning up.
208    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    // Locks of keys without an entry, unless another run holds them.
221    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/// Every blob a manifest cites: its log and its file outputs.
248pub(crate) fn manifest_blobs(manifest: &Value) -> BTreeSet<String> {
249    let mut blobs = BTreeSet::new();
250    if let Some(log) = manifest.get("log").and_then(Value::as_str) {
251        blobs.insert(log.to_owned());
252    }
253    let groups = manifest
254        .get("groups")
255        .and_then(Value::as_object)
256        .into_iter()
257        .flat_map(|groups| groups.values())
258        .filter_map(|group| group.get("artifacts").and_then(Value::as_object));
259    for artifact in manifest
260        .get("artifacts")
261        .and_then(Value::as_object)
262        .into_iter()
263        .chain(groups)
264        .flat_map(|artifacts| artifacts.values())
265    {
266        if let Some(blob) = artifact.get("blob").and_then(Value::as_str) {
267            blobs.insert(blob.to_owned());
268        }
269    }
270    blobs
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn parses_sizes_like_nx() {
279        assert_eq!(parse_size("1024").unwrap(), 1024);
280        assert_eq!(parse_size("1KB").unwrap(), 1024);
281        assert_eq!(parse_size("1.5 MB").unwrap(), 1_572_864);
282        assert_eq!(parse_size("2GB").unwrap(), 2 << 30);
283        assert_eq!(parse_size("0").unwrap(), 0);
284        for bad in ["", "1TB", "1.2.3MB", "MB", "-1"] {
285            assert!(parse_size(bad).is_err(), "{bad}");
286        }
287    }
288}
qk_cache/evict.rs

qk_cache/
evict.rs
