atagen

roundabout

public
Star0
Watch0
56a853ebf2549fd63ad8a481423b088d23dd0dc1
1
ssh://brewery@pls.lobotomise.me:33/atagen/roundabout.git
src/artifacts.rsRust341 lines9.7 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>, fresh: bool) -> Registry {
150
    if let Some(root) = root {
151
        register(cfg, root, false);
152
    }
153
    let ttl = cfg.artifacts.size_ttl.0.as_secs();
154
    store::update("artifacts", |reg: &mut Registry| {
155
        reg.retain(|p, _| fs::symlink_metadata(p).is_ok());
156
        for (path, e) in reg.iter_mut() {
157
            if fresh || now().saturating_sub(e.sized_at) > ttl {
158
                e.size = size_of(path);
159
                e.sized_at = now();
160
            }
161
        }
162
        reg.clone()
163
    })
164
    .unwrap_or_default()
165
}
166
167
pub struct Report {
168
    pub project: Option<(PathBuf, Bytes)>,
169
    pub global: Bytes,
170
    pub over_project: bool,
171
    pub over_global: bool,
172
    pub entries: Registry,
173
}
174
175
impl Report {
176
    pub const fn over(&self) -> bool {
177
        self.over_project || self.over_global
178
    }
179
180
    pub fn describe(&self, cfg: &Config) -> String {
181
        let mut out = Vec::new();
182
        if let Some((root, size)) = &self.project
183
            && self.over_project
184
        {
185
            out.push(format!(
186
                "{} holds {size} of build artifacts, cap is {}",
187
                root.display(),
188
                cfg.artifacts.project_cap
189
            ));
190
        }
191
        if self.over_global {
192
            out.push(format!(
193
                "artifacts total {} across projects, cap is {}",
194
                self.global, cfg.artifacts.global_cap
195
            ));
196
        }
197
        out.join(". ")
198
    }
199
}
200
201
pub fn check(cfg: &Config, root: Option<&Path>, fresh: bool) -> Report {
202
    let entries = refresh(cfg, root, fresh);
203
    let global = entries.values().map(|e| e.size).sum();
204
    let project = root.map(|r| {
205
        let size = entries
206
            .values()
207
            .filter(|e| e.root == r)
208
            .map(|e| e.size)
209
            .sum();
210
        (r.to_path_buf(), size)
211
    });
212
    Report {
213
        over_project: project
214
            .as_ref()
215
            .is_some_and(|(_, s)| *s > cfg.artifacts.project_cap),
216
        over_global: global > cfg.artifacts.global_cap,
217
        project,
218
        global,
219
        entries,
220
    }
221
}
222
223
fn remove(cfg: &Config, path: &Path, e: &Entry) -> Result<(), String> {
224
    let rule = cfg.artifact_rules().find(|r| r.name == e.rule);
225
    if let Some(clean) = rule.and_then(|r| r.clean.as_ref())
226
        && let Some((prog, args)) = clean.split_first()
227
    {
228
        let ok = Command::new(prog)
229
            .args(args)
230
            .current_dir(&e.root)
231
            .status()
232
            .is_ok_and(|s| s.success());
233
        return if ok {
234
            Ok(())
235
        } else {
236
            Err(format!("`{}` failed", clean.join(" ")))
237
        };
238
    }
239
    let meta = fs::symlink_metadata(path).map_err(|err| err.to_string())?;
240
    if meta.is_dir() {
241
        fs::remove_dir_all(path)
242
    } else {
243
        fs::remove_file(path)
244
    }
245
    .map_err(|err| format!("{}: {err}", path.display()))
246
}
247
248
pub struct Gc {
249
    pub freed: Bytes,
250
    pub removed: Vec<PathBuf>,
251
    pub failed: Vec<String>,
252
}
253
254
pub fn gc(cfg: &Config, to_cap: bool, dry_run: bool) -> Gc {
255
    let entries = refresh(cfg, None, false);
256
    let active = active_roots();
257
    let idle_after = cfg.artifacts.idle_after.0.as_secs();
258
    let movable = |e: &Entry| !e.pinned && !active.contains(&e.root);
259
260
    let mut doomed: Vec<PathBuf> = entries
261
        .iter()
262
        .filter(|(_, e)| movable(e) && now().saturating_sub(e.last_used) > idle_after)
263
        .map(|(p, _)| p.clone())
264
        .collect();
265
266
    if to_cap {
267
        let mut global: Bytes = entries
268
            .iter()
269
            .filter(|(p, _)| !doomed.contains(p))
270
            .map(|(_, e)| e.size)
271
            .sum();
272
        let mut per_root: BTreeMap<&Path, Bytes> = BTreeMap::new();
273
        for (p, e) in &entries {
274
            if !doomed.contains(p) {
275
                *per_root.entry(&e.root).or_default() += e.size;
276
            }
277
        }
278
        let mut lru: Vec<(&PathBuf, &Entry)> = entries
279
            .iter()
280
            .filter(|(p, e)| movable(e) && !doomed.contains(p))
281
            .collect();
282
        lru.sort_by_key(|(_, e)| e.last_used);
283
        for (p, e) in lru {
284
            let root_size = per_root.get(e.root.as_path()).copied().unwrap_or_default();
285
            if global <= cfg.artifacts.global_cap && root_size <= cfg.artifacts.project_cap {
286
                continue;
287
            }
288
            global = global.saturating_sub(e.size);
289
            per_root.insert(&e.root, root_size.saturating_sub(e.size));
290
            doomed.push(p.clone());
291
        }
292
    }
293
294
    let mut out = Gc {
295
        freed: Bytes::ZERO,
296
        removed: Vec::new(),
297
        failed: Vec::new(),
298
    };
299
    for path in doomed {
300
        let e = &entries[&path];
301
        if !dry_run && let Err(err) = remove(cfg, &path, e) {
302
            out.failed.push(err);
303
            continue;
304
        }
305
        out.freed += e.size;
306
        out.removed.push(path);
307
    }
308
    if !dry_run {
309
        let removed = out.removed.clone();
310
        let _ = store::update("artifacts", |reg: &mut Registry| {
311
            for p in &removed {
312
                reg.remove(p);
313
            }
314
        });
315
        nix_gc(cfg, &mut out);
316
    }
317
    out
318
}
319
320
fn nix_gc(cfg: &Config, out: &mut Gc) {
321
    if !cfg.artifacts.nix_gc {
322
        return;
323
    }
324
    let Some(store) = sys::fs(Path::new("/nix/store")) else {
325
        return;
326
    };
327
    let want = cfg.artifacts.nix_min_free.saturating_sub(store.free);
328
    if want == Bytes::ZERO {
329
        return;
330
    }
331
    let ok = Command::new("nix")
332
        .args(["store", "gc", "--max", &want.0.to_string()])
333
        .status()
334
        .is_ok_and(|s| s.success());
335
    if ok {
336
        let after = sys::fs(Path::new("/nix/store")).map_or(store.free, |s| s.free);
337
        out.freed += after.saturating_sub(store.free);
338
    } else {
339
        out.failed.push("nix store gc failed".into());
340
    }
341
}