atagen

roundabout

public
Star0
Watch0
7770b8fcfa8b72ae913db787fd5ef641f31bd302
1
ssh://brewery@pls.lobotomise.me:33/atagen/roundabout.git
src/artifacts.rsRust357 lines10.2 KB
1
use std::{
2
    collections::{BTreeMap, BTreeSet},
3
    fs,
4
    os::unix::fs::MetadataExt,
5
    path::{Path, PathBuf},
6
    process::Command,
7
    time::{SystemTime, UNIX_EPOCH},
8
};
9
10
use serde::{Deserialize, Serialize};
11
12
use crate::{
13
    client,
14
    config::{Config, Rule},
15
    paths,
16
    proto::{self, Event, Reader, Request},
17
    store, sys,
18
    units::Bytes,
19
};
20
21
#[derive(Clone, Debug, Serialize, Deserialize)]
22
pub struct Entry {
23
    pub root: PathBuf,
24
    pub rule: String,
25
    pub last_used: u64,
26
    pub pinned: bool,
27
    pub size: Bytes,
28
    pub sized_at: u64,
29
}
30
31
type Registry = BTreeMap<PathBuf, Entry>;
32
33
fn now() -> u64 {
34
    SystemTime::now()
35
        .duration_since(UNIX_EPOCH)
36
        .map_or(0, |d| d.as_secs())
37
}
38
39
fn matches(rule: &Rule, name: &str) -> bool {
40
    match rule.name.strip_suffix('*') {
41
        Some(prefix) => name.starts_with(prefix),
42
        None => name == rule.name,
43
    }
44
}
45
46
fn discover<'a>(cfg: &'a Config, root: &Path) -> Vec<(PathBuf, &'a Rule)> {
47
    let Ok(dir) = fs::read_dir(root) else {
48
        return Vec::new();
49
    };
50
    let names: Vec<String> = dir
51
        .flatten()
52
        .map(|e| e.file_name().to_string_lossy().into_owned())
53
        .collect();
54
    names
55
        .iter()
56
        .filter_map(|name| {
57
            let rule = cfg.artifact_rules().find(|r| {
58
                matches(r, name) && r.marker.as_ref().is_none_or(|m| root.join(m).exists())
59
            })?;
60
            Some((root.join(name), rule))
61
        })
62
        .collect()
63
}
64
65
/// symlinks are followed at the top only, so a nix `result` counts its store
66
/// path. hardlinks count once, the store is full of them
67
fn size_of(path: &Path) -> Bytes {
68
    let mut seen = BTreeSet::new();
69
    let mut total = 0;
70
    let mut stack = vec![fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())];
71
    while let Some(p) = stack.pop() {
72
        let Ok(meta) = fs::symlink_metadata(&p) else {
73
            continue;
74
        };
75
        if meta.nlink() > 1 && !seen.insert((meta.dev(), meta.ino())) {
76
            continue;
77
        }
78
        total += meta.blocks() * 512;
79
        if meta.is_dir()
80
            && let Ok(dir) = fs::read_dir(&p)
81
        {
82
            stack.extend(dir.flatten().map(|e| e.path()));
83
        }
84
    }
85
    Bytes(total)
86
}
87
88
/// only builds bump `last_used`. checking sizes isn't using them
89
pub fn register(cfg: &Config, root: &Path, used: bool) {
90
    let found = discover(cfg, root);
91
    if found.is_empty() {
92
        return;
93
    }
94
    let _ = store::update("artifacts", |reg: &mut Registry| {
95
        for (path, rule) in found {
96
            reg.entry(path)
97
                .and_modify(|e| {
98
                    if used {
99
                        e.last_used = now();
100
                    }
101
                })
102
                .or_insert_with(|| Entry {
103
                    root: root.to_path_buf(),
104
                    rule: rule.name.clone(),
105
                    last_used: now(),
106
                    pinned: false,
107
                    size: Bytes::ZERO,
108
                    sized_at: 0,
109
                });
110
        }
111
    });
112
}
113
114
pub fn pin(path: &Path, pinned: bool) -> Result<(), String> {
115
    let path = fs::canonicalize(path).map_err(|e| format!("{}: {e}", path.display()))?;
116
    let root = paths::project_root(path.parent().unwrap_or(&path));
117
    store::update("artifacts", |reg: &mut Registry| {
118
        reg.entry(path.clone())
119
            .or_insert_with(|| Entry {
120
                root: root.clone(),
121
                rule: "manual".into(),
122
                last_used: now(),
123
                pinned,
124
                size: Bytes::ZERO,
125
                sized_at: 0,
126
            })
127
            .pinned = pinned;
128
    })
129
    .map_err(|e| e.to_string())
130
}
131
132
fn active_roots() -> BTreeSet<PathBuf> {
133
    let Some(mut s) = client::connect() else {
134
        return BTreeSet::new();
135
    };
136
    if proto::send(&mut s, &Request::Status).is_err() {
137
        return BTreeSet::new();
138
    }
139
    match Reader::new(&s).and_then(|mut r| r.recv::<Event>()) {
140
        Ok(Some(Event::Status(st))) => st
141
            .leases
142
            .iter()
143
            .map(|l| paths::project_root(&l.cwd))
144
            .collect(),
145
        _ => BTreeSet::new(),
146
    }
147
}
148
149
fn refresh(cfg: &Config, root: Option<&Path>, max_age: u64) -> Registry {
150
    if let Some(root) = root {
151
        register(cfg, root, false);
152
    }
153
    store::update("artifacts", |reg: &mut Registry| {
154
        reg.retain(|p, _| fs::symlink_metadata(p).is_ok());
155
        for (path, e) in reg.iter_mut() {
156
            if now().saturating_sub(e.sized_at) >= max_age {
157
                e.size = size_of(path);
158
                e.sized_at = now();
159
            }
160
        }
161
        reg.clone()
162
    })
163
    .unwrap_or_default()
164
}
165
166
pub struct Report {
167
    pub project: Option<(PathBuf, Bytes)>,
168
    pub global: Bytes,
169
    pub over_project: bool,
170
    pub over_global: bool,
171
    pub entries: Registry,
172
}
173
174
impl Report {
175
    pub const fn over(&self) -> bool {
176
        self.over_project || self.over_global
177
    }
178
179
    pub fn describe(&self, cfg: &Config) -> String {
180
        let mut out = Vec::new();
181
        if let Some((root, size)) = &self.project
182
            && self.over_project
183
        {
184
            out.push(format!(
185
                "{} holds {size} of build artifacts, cap is {}",
186
                root.display(),
187
                cfg.artifacts.project_cap
188
            ));
189
        }
190
        if self.over_global {
191
            out.push(format!(
192
                "artifacts total {} across projects, cap is {}",
193
                self.global, cfg.artifacts.global_cap
194
            ));
195
        }
196
        out.join(". ")
197
    }
198
}
199
200
const RECHECK_SECS: u64 = 60;
201
202
/// the cache can predate someone's `rm`, so being over earns a second look
203
/// before anyone gets nagged
204
pub fn check(cfg: &Config, root: Option<&Path>, fresh: bool) -> Report {
205
    let ttl = if fresh {
206
        0
207
    } else {
208
        cfg.artifacts.size_ttl.0.as_secs()
209
    };
210
    let report = measure(cfg, root, ttl);
211
    if report.over() && ttl > RECHECK_SECS {
212
        return measure(cfg, root, RECHECK_SECS);
213
    }
214
    report
215
}
216
217
fn measure(cfg: &Config, root: Option<&Path>, max_age: u64) -> Report {
218
    let entries = refresh(cfg, root, max_age);
219
    let global = entries.values().map(|e| e.size).sum();
220
    let project = root.map(|r| {
221
        let size = entries
222
            .values()
223
            .filter(|e| e.root == r)
224
            .map(|e| e.size)
225
            .sum();
226
        (r.to_path_buf(), size)
227
    });
228
    Report {
229
        over_project: project
230
            .as_ref()
231
            .is_some_and(|(_, s)| *s > cfg.artifacts.project_cap),
232
        over_global: global > cfg.artifacts.global_cap,
233
        project,
234
        global,
235
        entries,
236
    }
237
}
238
239
fn remove(cfg: &Config, path: &Path, e: &Entry) -> Result<(), String> {
240
    let rule = cfg.artifact_rules().find(|r| r.name == e.rule);
241
    if let Some(clean) = rule.and_then(|r| r.clean.as_ref())
242
        && let Some((prog, args)) = clean.split_first()
243
    {
244
        let ok = Command::new(prog)
245
            .args(args)
246
            .current_dir(&e.root)
247
            .status()
248
            .is_ok_and(|s| s.success());
249
        return if ok {
250
            Ok(())
251
        } else {
252
            Err(format!("`{}` failed", clean.join(" ")))
253
        };
254
    }
255
    let meta = fs::symlink_metadata(path).map_err(|err| err.to_string())?;
256
    if meta.is_dir() {
257
        fs::remove_dir_all(path)
258
    } else {
259
        fs::remove_file(path)
260
    }
261
    .map_err(|err| format!("{}: {err}", path.display()))
262
}
263
264
pub struct Gc {
265
    pub freed: Bytes,
266
    pub removed: Vec<PathBuf>,
267
    pub failed: Vec<String>,
268
}
269
270
pub fn gc(cfg: &Config, to_cap: bool, dry_run: bool) -> Gc {
271
    let entries = refresh(cfg, None, cfg.artifacts.size_ttl.0.as_secs());
272
    let active = active_roots();
273
    let idle_after = cfg.artifacts.idle_after.0.as_secs();
274
    let movable = |e: &Entry| !e.pinned && !active.contains(&e.root);
275
276
    let mut doomed: Vec<PathBuf> = entries
277
        .iter()
278
        .filter(|(_, e)| movable(e) && now().saturating_sub(e.last_used) > idle_after)
279
        .map(|(p, _)| p.clone())
280
        .collect();
281
282
    if to_cap {
283
        let mut global: Bytes = entries
284
            .iter()
285
            .filter(|(p, _)| !doomed.contains(p))
286
            .map(|(_, e)| e.size)
287
            .sum();
288
        let mut per_root: BTreeMap<&Path, Bytes> = BTreeMap::new();
289
        for (p, e) in &entries {
290
            if !doomed.contains(p) {
291
                *per_root.entry(&e.root).or_default() += e.size;
292
            }
293
        }
294
        let mut lru: Vec<(&PathBuf, &Entry)> = entries
295
            .iter()
296
            .filter(|(p, e)| movable(e) && !doomed.contains(p))
297
            .collect();
298
        lru.sort_by_key(|(_, e)| e.last_used);
299
        for (p, e) in lru {
300
            let root_size = per_root.get(e.root.as_path()).copied().unwrap_or_default();
301
            if global <= cfg.artifacts.global_cap && root_size <= cfg.artifacts.project_cap {
302
                continue;
303
            }
304
            global = global.saturating_sub(e.size);
305
            per_root.insert(&e.root, root_size.saturating_sub(e.size));
306
            doomed.push(p.clone());
307
        }
308
    }
309
310
    let mut out = Gc {
311
        freed: Bytes::ZERO,
312
        removed: Vec::new(),
313
        failed: Vec::new(),
314
    };
315
    for path in doomed {
316
        let e = &entries[&path];
317
        if !dry_run && let Err(err) = remove(cfg, &path, e) {
318
            out.failed.push(err);
319
            continue;
320
        }
321
        out.freed += e.size;
322
        out.removed.push(path);
323
    }
324
    if !dry_run {
325
        let removed = out.removed.clone();
326
        let _ = store::update("artifacts", |reg: &mut Registry| {
327
            for p in &removed {
328
                reg.remove(p);
329
            }
330
        });
331
        nix_gc(cfg, &mut out);
332
    }
333
    out
334
}
335
336
fn nix_gc(cfg: &Config, out: &mut Gc) {
337
    if !cfg.artifacts.nix_gc {
338
        return;
339
    }
340
    let Some(store) = sys::fs(Path::new("/nix/store")) else {
341
        return;
342
    };
343
    let want = cfg.artifacts.nix_min_free.saturating_sub(store.free);
344
    if want == Bytes::ZERO {
345
        return;
346
    }
347
    let ok = Command::new("nix")
348
        .args(["store", "gc", "--max", &want.0.to_string()])
349
        .status()
350
        .is_ok_and(|s| s.success());
351
    if ok {
352
        let after = sys::fs(Path::new("/nix/store")).map_or(store.free, |s| s.free);
353
        out.freed += after.saturating_sub(store.free);
354
    } else {
355
        out.failed.push("nix store gc failed".into());
356
    }
357
}