atagen

roundabout

public
Star0
Watch0
fdad58f8133a77a20c44d0e1472cb73a5128ef2c
1
ssh://brewery@pls.lobotomise.me:33/atagen/roundabout.git
src/artifacts.rsRust371 lines10.5 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 const fn remedy(&self) -> &'static str {
180
        if self.over_global {
181
            "roundabout artifacts gc --cap"
182
        } else {
183
            "roundabout artifacts gc --project --cap"
184
        }
185
    }
186
187
    pub fn describe(&self, cfg: &Config) -> String {
188
        let mut out = Vec::new();
189
        if let Some((root, size)) = &self.project
190
            && self.over_project
191
        {
192
            out.push(format!(
193
                "{} holds {size} of build artifacts, cap is {}",
194
                root.display(),
195
                cfg.artifacts.project_cap
196
            ));
197
        }
198
        if self.over_global {
199
            out.push(format!(
200
                "artifacts total {} across projects, cap is {}",
201
                self.global, cfg.artifacts.global_cap
202
            ));
203
        }
204
        out.join(". ")
205
    }
206
}
207
208
const RECHECK_SECS: u64 = 60;
209
210
/// the cache can predate someone's `rm`, so being over earns a second look
211
/// before anyone gets nagged
212
pub fn check(cfg: &Config, root: Option<&Path>, fresh: bool) -> Report {
213
    let ttl = if fresh {
214
        0
215
    } else {
216
        cfg.artifacts.size_ttl.0.as_secs()
217
    };
218
    let report = measure(cfg, root, ttl);
219
    if report.over() && ttl > RECHECK_SECS {
220
        return measure(cfg, root, RECHECK_SECS);
221
    }
222
    report
223
}
224
225
fn measure(cfg: &Config, root: Option<&Path>, max_age: u64) -> Report {
226
    let entries = refresh(cfg, root, max_age);
227
    let global = entries.values().map(|e| e.size).sum();
228
    let project = root.map(|r| {
229
        let size = entries
230
            .values()
231
            .filter(|e| e.root == r)
232
            .map(|e| e.size)
233
            .sum();
234
        (r.to_path_buf(), size)
235
    });
236
    Report {
237
        over_project: project
238
            .as_ref()
239
            .is_some_and(|(_, s)| *s > cfg.artifacts.project_cap),
240
        over_global: global > cfg.artifacts.global_cap,
241
        project,
242
        global,
243
        entries,
244
    }
245
}
246
247
fn remove(cfg: &Config, path: &Path, e: &Entry) -> Result<(), String> {
248
    let rule = cfg.artifact_rules().find(|r| r.name == e.rule);
249
    if let Some(clean) = rule.and_then(|r| r.clean.as_ref())
250
        && let Some((prog, args)) = clean.split_first()
251
    {
252
        let ok = Command::new(prog)
253
            .args(args)
254
            .current_dir(&e.root)
255
            .status()
256
            .is_ok_and(|s| s.success());
257
        return if ok {
258
            Ok(())
259
        } else {
260
            Err(format!("`{}` failed", clean.join(" ")))
261
        };
262
    }
263
    let meta = fs::symlink_metadata(path).map_err(|err| err.to_string())?;
264
    if meta.is_dir() {
265
        fs::remove_dir_all(path)
266
    } else {
267
        fs::remove_file(path)
268
    }
269
    .map_err(|err| format!("{}: {err}", path.display()))
270
}
271
272
pub struct Gc {
273
    pub freed: Bytes,
274
    pub removed: Vec<PathBuf>,
275
    pub failed: Vec<String>,
276
}
277
278
pub fn gc(cfg: &Config, project: Option<&Path>, to_cap: bool, dry_run: bool) -> Gc {
279
    let max_age = if to_cap {
280
        RECHECK_SECS
281
    } else {
282
        cfg.artifacts.size_ttl.0.as_secs()
283
    };
284
    let entries = refresh(cfg, project, max_age);
285
    let active = active_roots();
286
    let idle_after = cfg.artifacts.idle_after.0.as_secs();
287
    let movable =
288
        |e: &Entry| !e.pinned && !active.contains(&e.root) && project.is_none_or(|r| e.root == r);
289
290
    let mut doomed: Vec<PathBuf> = entries
291
        .iter()
292
        .filter(|(_, e)| movable(e) && now().saturating_sub(e.last_used) > idle_after)
293
        .map(|(p, _)| p.clone())
294
        .collect();
295
296
    if to_cap {
297
        let mut global: Bytes = entries
298
            .iter()
299
            .filter(|(p, _)| !doomed.contains(p))
300
            .map(|(_, e)| e.size)
301
            .sum();
302
        let mut per_root: BTreeMap<&Path, Bytes> = BTreeMap::new();
303
        for (p, e) in &entries {
304
            if !doomed.contains(p) {
305
                *per_root.entry(&e.root).or_default() += e.size;
306
            }
307
        }
308
        let mut lru: Vec<(&PathBuf, &Entry)> = entries
309
            .iter()
310
            .filter(|(p, e)| movable(e) && !doomed.contains(p))
311
            .collect();
312
        lru.sort_by_key(|(_, e)| e.last_used);
313
        for (p, e) in lru {
314
            let root_size = per_root.get(e.root.as_path()).copied().unwrap_or_default();
315
            if global <= cfg.artifacts.global_cap && root_size <= cfg.artifacts.project_cap {
316
                continue;
317
            }
318
            global = global.saturating_sub(e.size);
319
            per_root.insert(&e.root, root_size.saturating_sub(e.size));
320
            doomed.push(p.clone());
321
        }
322
    }
323
324
    let mut out = Gc {
325
        freed: Bytes::ZERO,
326
        removed: Vec::new(),
327
        failed: Vec::new(),
328
    };
329
    for path in doomed {
330
        let e = &entries[&path];
331
        if !dry_run && let Err(err) = remove(cfg, &path, e) {
332
            out.failed.push(err);
333
            continue;
334
        }
335
        out.freed += e.size;
336
        out.removed.push(path);
337
    }
338
    if !dry_run {
339
        let removed = out.removed.clone();
340
        let _ = store::update("artifacts", |reg: &mut Registry| {
341
            for p in &removed {
342
                reg.remove(p);
343
            }
344
        });
345
        nix_gc(cfg, &mut out);
346
    }
347
    out
348
}
349
350
fn nix_gc(cfg: &Config, out: &mut Gc) {
351
    if !cfg.artifacts.nix_gc {
352
        return;
353
    }
354
    let Some(store) = sys::fs(Path::new("/nix/store")) else {
355
        return;
356
    };
357
    let want = cfg.artifacts.nix_min_free.saturating_sub(store.free);
358
    if want == Bytes::ZERO {
359
        return;
360
    }
361
    let ok = Command::new("nix")
362
        .args(["store", "gc", "--max", &want.0.to_string()])
363
        .status()
364
        .is_ok_and(|s| s.success());
365
    if ok {
366
        let after = sys::fs(Path::new("/nix/store")).map_or(store.free, |s| s.free);
367
        out.freed += after.saturating_sub(store.free);
368
    } else {
369
        out.failed.push("nix store gc failed".into());
370
    }
371
}