atagen

roundabout

public
Star0
Watch0

Commit 6f5b84fad3ef

tarball

atagen <boss@atagen.co>2026-09-27 03:59 UTC

18 files changed+3772-0
roundabout: admit builds by self-reported size so parallel agents stop ooming the box

A per-user daemon, spawned on first use and gone once idle, grants leases for as long as the requesting connection stays open, so a crashed agent can never strand a lock. Admission is fifo with easy backfill. Heavies never overlap, disk is never promised twice, and small jobs only jump the queue when they will be gone before the head could start. Build artifacts are tracked per project and gc'd once idle or over their caps.
Added file: .gitignore+2 additions
@@ -0,0 +1,2 @@
1+/target
2+/result*
Added file: Cargo.lock+198 additions
@@ -0,0 +1,198 @@
1+# This file is automatically @generated by Cargo.
2+# It is not intended for manual editing.
3+version = 4
4+
5+[[package]]
6+name = "itoa"
7+version = "1.0.18"
8+source = "registry+https://github.com/rust-lang/crates.io-index"
9+checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
10+
11+[[package]]
12+name = "libc"
13+version = "0.2.189"
14+source = "registry+https://github.com/rust-lang/crates.io-index"
15+checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
16+
17+[[package]]
18+name = "memchr"
19+version = "2.8.3"
20+source = "registry+https://github.com/rust-lang/crates.io-index"
21+checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
22+
23+[[package]]
24+name = "pound"
25+version = "0.1.8"
26+source = "registry+https://github.com/rust-lang/crates.io-index"
27+checksum = "c86a0567fb5ffd1e933983f6baef4e557d0cc2bf3541b26c169bbb8099889487"
28+dependencies = [
29+ "pound-derive",
30+]
31+
32+[[package]]
33+name = "pound-derive"
34+version = "0.1.8"
35+source = "registry+https://github.com/rust-lang/crates.io-index"
36+checksum = "e1ea8a0689abf825dc184168cda08d14ff35a451cea2650464400932589cf4cf"
37+dependencies = [
38+ "proc-macro2",
39+ "quote",
40+ "venial",
41+]
42+
43+[[package]]
44+name = "proc-macro2"
45+version = "1.0.107"
46+source = "registry+https://github.com/rust-lang/crates.io-index"
47+checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
48+dependencies = [
49+ "unicode-ident",
50+]
51+
52+[[package]]
53+name = "quote"
54+version = "1.0.47"
55+source = "registry+https://github.com/rust-lang/crates.io-index"
56+checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
57+dependencies = [
58+ "proc-macro2",
59+]
60+
61+[[package]]
62+name = "roundabout"
63+version = "0.1.0"
64+dependencies = [
65+ "libc",
66+ "pound",
67+ "serde",
68+ "serde_json",
69+ "toml",
70+]
71+
72+[[package]]
73+name = "serde"
74+version = "1.0.229"
75+source = "registry+https://github.com/rust-lang/crates.io-index"
76+checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
77+dependencies = [
78+ "serde_core",
79+ "serde_derive",
80+]
81+
82+[[package]]
83+name = "serde_core"
84+version = "1.0.229"
85+source = "registry+https://github.com/rust-lang/crates.io-index"
86+checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
87+dependencies = [
88+ "serde_derive",
89+]
90+
91+[[package]]
92+name = "serde_derive"
93+version = "1.0.229"
94+source = "registry+https://github.com/rust-lang/crates.io-index"
95+checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
96+dependencies = [
97+ "proc-macro2",
98+ "quote",
99+ "syn",
100+]
101+
102+[[package]]
103+name = "serde_json"
104+version = "1.0.151"
105+source = "registry+https://github.com/rust-lang/crates.io-index"
106+checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14"
107+dependencies = [
108+ "itoa",
109+ "memchr",
110+ "serde",
111+ "serde_core",
112+ "zmij",
113+]
114+
115+[[package]]
116+name = "serde_spanned"
117+version = "1.1.1"
118+source = "registry+https://github.com/rust-lang/crates.io-index"
119+checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26"
120+dependencies = [
121+ "serde_core",
122+]
123+
124+[[package]]
125+name = "syn"
126+version = "3.0.6"
127+source = "registry+https://github.com/rust-lang/crates.io-index"
128+checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee"
129+dependencies = [
130+ "proc-macro2",
131+ "quote",
132+ "unicode-ident",
133+]
134+
135+[[package]]
136+name = "toml"
137+version = "0.9.12+spec-1.1.0"
138+source = "registry+https://github.com/rust-lang/crates.io-index"
139+checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863"
140+dependencies = [
141+ "serde_core",
142+ "serde_spanned",
143+ "toml_datetime",
144+ "toml_parser",
145+ "winnow 0.7.15",
146+]
147+
148+[[package]]
149+name = "toml_datetime"
150+version = "0.7.5+spec-1.1.0"
151+source = "registry+https://github.com/rust-lang/crates.io-index"
152+checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347"
153+dependencies = [
154+ "serde_core",
155+]
156+
157+[[package]]
158+name = "toml_parser"
159+version = "1.1.3+spec-1.1.0"
160+source = "registry+https://github.com/rust-lang/crates.io-index"
161+checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56"
162+dependencies = [
163+ "winnow 1.0.4",
164+]
165+
166+[[package]]
167+name = "unicode-ident"
168+version = "1.0.26"
169+source = "registry+https://github.com/rust-lang/crates.io-index"
170+checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954"
171+
172+[[package]]
173+name = "venial"
174+version = "0.6.1"
175+source = "registry+https://github.com/rust-lang/crates.io-index"
176+checksum = "9a42528baceab6c7784446df2a10f4185078c39bf73dc614f154353f1a6b1229"
177+dependencies = [
178+ "proc-macro2",
179+ "quote",
180+]
181+
182+[[package]]
183+name = "winnow"
184+version = "0.7.15"
185+source = "registry+https://github.com/rust-lang/crates.io-index"
186+checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945"
187+
188+[[package]]
189+name = "winnow"
190+version = "1.0.4"
191+source = "registry+https://github.com/rust-lang/crates.io-index"
192+checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81"
193+
194+[[package]]
195+name = "zmij"
196+version = "1.0.23"
197+source = "registry+https://github.com/rust-lang/crates.io-index"
198+checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
Added file: Cargo.toml+38 additions
@@ -0,0 +1,38 @@
1+[package]
2+name = "roundabout"
3+version = "0.1.0"
4+edition = "2024"
5+authors = [ "atagen" ]
6+description = "build admission for agents that can't share"
7+license = "EUPL-1.2"
8+rust-version = "1.91"
9+
10+[[bin]]
11+name = "rabt"
12+path = "src/main.rs"
13+
14+[dependencies]
15+libc = "0.2"
16+pound = "0.1"
17+serde = { version = "1", features = [ "derive" ] }
18+serde_json = "1"
19+toml = { version = "0.9", default-features = false, features = [ "parse", "serde" ] }
20+
21+[lints.clippy]
22+nursery = { level = "warn", priority = -1 }
23+pedantic = { level = "warn", priority = -1 }
24+
25+cast_possible_truncation = "allow"
26+cast_precision_loss = "allow"
27+cast_sign_loss = "allow"
28+missing_errors_doc = "allow"
29+module_name_repetitions = "allow"
30+option_if_let_else = "allow"
31+too_many_lines = "allow"
32+
33+[profile.release]
34+codegen-units = 1
35+lto = "fat"
36+opt-level = "s"
37+panic = "abort"
38+strip = true
Added file: src/artifacts.rs+341 additions
@@ -0,0 +1,341 @@
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+}
Added file: src/cli.rs+326 additions
@@ -0,0 +1,326 @@
1+use std::{io, os::unix::net::UnixStream, path::PathBuf};
2+
3+use pound::Parse;
4+
5+use crate::{
6+ artifacts, client,
7+ config::Config,
8+ daemon, history, hook, paths,
9+ proto::{self, Event, Reader, Request},
10+ run,
11+ units::{Amount, Bytes, Dur},
12+};
13+
14+/// one build at a time, please
15+#[derive(Parse)]
16+#[pound(name = "rabt", version = env!("CARGO_PKG_VERSION"))]
17+pub enum Rabt {
18+ /// wait your turn, then run the thing
19+ Run {
20+ /// preset from config, s m l xl out of the box
21+ #[pound(short, long)]
22+ size: Option<String>,
23+ #[pound(long)]
24+ ram: Option<Bytes>,
25+ #[pound(long)]
26+ disk: Option<Bytes>,
27+ /// count or percentage of the machine
28+ #[pound(long)]
29+ cores: Option<Amount<u32>>,
30+ /// how long you reckon it'll take
31+ #[pound(short, long)]
32+ time: Option<Dur>,
33+ /// kill it past this, exit 124
34+ #[pound(long)]
35+ deadline: Option<Dur>,
36+ /// stop queueing past this, exit 75
37+ #[pound(long)]
38+ max_wait: Option<Dur>,
39+ #[pound(long, hidden, env = "CLAUDE_CODE_SESSION_ID")]
40+ session: Option<String>,
41+ #[pound(trailing)]
42+ cmd: Vec<String>,
43+ },
44+ /// who's running, who's waiting
45+ Status {
46+ #[pound(long)]
47+ json: bool,
48+ },
49+ /// a size guess from past runs
50+ Suggest {
51+ #[pound(trailing)]
52+ cmd: Vec<String>,
53+ },
54+ /// stop a lease's build
55+ Kill { lease: u64 },
56+ /// buy an overdue build more time
57+ Extend { lease: u64, by: Dur },
58+ /// exit when a lease goes overdue or finishes
59+ Watch { lease: u64 },
60+ /// build output and how much of it there is
61+ Artifacts {
62+ #[pound(subcommand)]
63+ action: Artifacts,
64+ },
65+ /// claude code hook entrypoint, json on stdin
66+ Hook {
67+ /// pre-tool-use post-tool-use prompt-submit stop subagent-stop
68+ event: String,
69+ },
70+ #[pound(hidden)]
71+ Daemon,
72+}
73+
74+#[derive(Parse)]
75+pub enum Artifacts {
76+ /// sizes against the caps, exits 1 when over
77+ Check {
78+ /// re-measure instead of trusting the cache
79+ #[pound(long)]
80+ fresh: bool,
81+ },
82+ /// everything tracked, biggest first
83+ List,
84+ /// delete idle artifacts
85+ Gc {
86+ /// keep going, oldest first, until under the caps
87+ #[pound(long)]
88+ cap: bool,
89+ #[pound(short = 'n', long)]
90+ dry_run: bool,
91+ },
92+ /// never gc this
93+ Pin {
94+ path: PathBuf,
95+ },
96+ Unpin {
97+ path: PathBuf,
98+ },
99+}
100+
101+fn artifacts(cfg: &Config, action: Artifacts) -> i32 {
102+ let cwd = std::env::current_dir().unwrap_or_else(|_| "/".into());
103+ let root = paths::project_root(&cwd);
104+ match action {
105+ Artifacts::Check { fresh } => {
106+ let r = artifacts::check(cfg, Some(&root), fresh);
107+ if let Some((root, size)) = &r.project {
108+ println!(
109+ "{}: {size} of {}",
110+ root.display(),
111+ cfg.artifacts.project_cap
112+ );
113+ }
114+ println!("everything: {} of {}", r.global, cfg.artifacts.global_cap);
115+ if r.over() {
116+ println!("over. `rabt artifacts gc --cap`");
117+ 1
118+ } else {
119+ 0
120+ }
121+ }
122+ Artifacts::List => {
123+ let r = artifacts::check(cfg, Some(&root), false);
124+ let mut rows: Vec<_> = r.entries.iter().collect();
125+ rows.sort_by_key(|(_, e)| std::cmp::Reverse(e.size));
126+ for (path, e) in rows {
127+ println!(
128+ "{:>10} {}{}",
129+ e.size.to_string(),
130+ path.display(),
131+ if e.pinned { " (pinned)" } else { "" }
132+ );
133+ }
134+ 0
135+ }
136+ Artifacts::Gc { cap, dry_run } => {
137+ let g = artifacts::gc(cfg, cap, dry_run);
138+ for p in &g.removed {
139+ println!(
140+ "{}{}",
141+ if dry_run { "would remove " } else { "removed " },
142+ p.display()
143+ );
144+ }
145+ for f in &g.failed {
146+ eprintln!("rabt: {f}");
147+ }
148+ println!(
149+ "{} {}",
150+ if dry_run { "would free" } else { "freed" },
151+ g.freed
152+ );
153+ i32::from(!g.failed.is_empty())
154+ }
155+ Artifacts::Pin { path } | Artifacts::Unpin { path } if !path.exists() => {
156+ eprintln!("rabt: {} doesn't exist", path.display());
157+ 1
158+ }
159+ Artifacts::Pin { path } => report(artifacts::pin(&path, true)),
160+ Artifacts::Unpin { path } => report(artifacts::pin(&path, false)),
161+ }
162+}
163+
164+fn report(r: Result<(), String>) -> i32 {
165+ r.map_or_else(
166+ |e| {
167+ eprintln!("rabt: {e}");
168+ 1
169+ },
170+ |()| 0,
171+ )
172+}
173+
174+fn ask(req: &Request) -> io::Result<Option<Event>> {
175+ let Some(mut s) = client::connect() else {
176+ return Ok(None);
177+ };
178+ proto::send(&mut s, req)?;
179+ Reader::new(&s)?.recv()
180+}
181+
182+fn reply(ev: io::Result<Option<Event>>) -> i32 {
183+ match ev {
184+ Ok(Some(Event::Ok)) => 0,
185+ Ok(Some(Event::Err { msg })) => {
186+ eprintln!("rabt: {msg}");
187+ 1
188+ }
189+ Ok(None) => {
190+ eprintln!("rabt: nothing running");
191+ 1
192+ }
193+ Ok(Some(_)) => 1,
194+ Err(e) => {
195+ eprintln!("rabt: {e}");
196+ 1
197+ }
198+ }
199+}
200+
201+fn status(json: bool) -> i32 {
202+ let Ok(Some(Event::Status(st))) = ask(&Request::Status) else {
203+ if json {
204+ println!("null");
205+ } else {
206+ println!("idle");
207+ }
208+ return 0;
209+ };
210+ if json {
211+ println!("{}", serde_json::to_string(&st).unwrap_or_default());
212+ return 0;
213+ }
214+ println!(
215+ "ram {} free of {}, cores {} free of {}",
216+ st.ram_free, st.ram_budget, st.cores_free, st.cores
217+ );
218+ for l in &st.leases {
219+ let state = if l.running {
220+ format!(
221+ "running {}{}",
222+ Dur::secs(l.elapsed),
223+ if l.overdue { " OVERDUE" } else { "" }
224+ )
225+ } else {
226+ format!(
227+ "#{} {}",
228+ l.position,
229+ l.waiting.map(|r| r.to_string()).unwrap_or_default()
230+ )
231+ };
232+ println!(
233+ "{:>4} {}{} {} ram {} cores {} rss {} {}",
234+ l.id,
235+ if l.heavy { "H " } else { "" },
236+ state,
237+ Dur(l.need.time),
238+ l.need.ram,
239+ l.need.cores,
240+ l.rss,
241+ l.cmd.join(" ")
242+ );
243+ }
244+ 0
245+}
246+
247+fn suggest(cfg: &Config, cmd: &[String]) -> i32 {
248+ let cwd = std::env::current_dir().unwrap_or_else(|_| "/".into());
249+ let key = history::key(&paths::project_root(&cwd), cmd);
250+ let Some(rec) = history::lookup(&key) else {
251+ println!("never seen it, pick a --size");
252+ return 1;
253+ };
254+ let base = match run::base(cfg) {
255+ Ok(b) => b,
256+ Err(e) => {
257+ eprintln!("rabt: {e}");
258+ return 2;
259+ }
260+ };
261+ let n = history::suggest(&rec, base, cmd);
262+ println!(
263+ "--ram {} --cores {} --time {} --disk {}",
264+ n.ram,
265+ n.cores,
266+ Dur(n.time),
267+ n.disk
268+ );
269+ 0
270+}
271+
272+fn watch(lease: u64) -> i32 {
273+ let Some(mut s): Option<UnixStream> = client::connect() else {
274+ return 0;
275+ };
276+ if proto::send(&mut s, &Request::Watch { lease }).is_err() {
277+ return 0;
278+ }
279+ if matches!(
280+ Reader::new(&s).and_then(|mut r| r.recv::<Event>()),
281+ Ok(Some(Event::Fired { overdue: true }))
282+ ) {
283+ println!("lease {lease} is overdue, `rabt kill {lease}` or `rabt extend {lease} <dur>`");
284+ 3
285+ } else {
286+ println!("lease {lease} is done");
287+ 0
288+ }
289+}
290+
291+pub fn dispatch(cfg: &Config, cmd: Rabt) -> i32 {
292+ match cmd {
293+ Rabt::Run {
294+ size,
295+ ram,
296+ disk,
297+ cores,
298+ time,
299+ deadline,
300+ max_wait,
301+ session,
302+ cmd,
303+ } => run::run(
304+ cfg,
305+ &run::Args {
306+ size,
307+ ram,
308+ disk,
309+ cores,
310+ time,
311+ deadline,
312+ max_wait,
313+ session,
314+ cmd,
315+ },
316+ ),
317+ Rabt::Status { json } => status(json),
318+ Rabt::Suggest { cmd } => suggest(cfg, &cmd),
319+ Rabt::Kill { lease } => reply(ask(&Request::Kill { lease })),
320+ Rabt::Extend { lease, by } => reply(ask(&Request::Extend { lease, by })),
321+ Rabt::Watch { lease } => watch(lease),
322+ Rabt::Artifacts { action } => artifacts(cfg, action),
323+ Rabt::Hook { event } => hook::main(cfg, &event),
324+ Rabt::Daemon => daemon::main(Config::load().unwrap_or_default()),
325+ }
326+}
Added file: src/client.rs+92 additions
@@ -0,0 +1,92 @@
1+use std::{
2+ fs::{self, File, OpenOptions},
3+ io,
4+ os::unix::{fs::OpenOptionsExt, net::UnixStream, process::CommandExt},
5+ process::{Command, Stdio},
6+ thread,
7+ time::Duration,
8+};
9+
10+use crate::paths;
11+
12+pub struct Flock(File);
13+
14+impl Flock {
15+ pub fn exclusive(file: File) -> io::Result<Self> {
16+ use std::os::fd::AsRawFd;
17+ if unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX) } != 0 {
18+ return Err(io::Error::last_os_error());
19+ }
20+ Ok(Self(file))
21+ }
22+
23+ pub fn try_exclusive(file: File) -> Option<Self> {
24+ use std::os::fd::AsRawFd;
25+ (unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } == 0)
26+ .then_some(Self(file))
27+ }
28+}
29+
30+impl Drop for Flock {
31+ fn drop(&mut self) {
32+ use std::os::fd::AsRawFd;
33+ unsafe { libc::flock(self.0.as_raw_fd(), libc::LOCK_UN) };
34+ }
35+}
36+
37+pub fn spawn_lock() -> io::Result<Flock> {
38+ let file = OpenOptions::new()
39+ .create(true)
40+ .truncate(false)
41+ .write(true)
42+ .mode(0o600)
43+ .open(paths::spawn_lock())?;
44+ Flock::exclusive(file)
45+}
46+
47+pub fn connect() -> Option<UnixStream> {
48+ UnixStream::connect(paths::socket()).ok()
49+}
50+
51+pub fn connect_or_spawn() -> io::Result<UnixStream> {
52+ if let Some(s) = connect() {
53+ return Ok(s);
54+ }
55+ let _lock = spawn_lock()?;
56+ if let Some(s) = connect() {
57+ return Ok(s);
58+ }
59+ let _ = fs::remove_file(paths::socket());
60+ let log = OpenOptions::new()
61+ .create(true)
62+ .append(true)
63+ .open(paths::state_dir().join("daemon.log"))?;
64+ let mut cmd = Command::new(std::env::current_exe()?);
65+ cmd.arg("daemon")
66+ .stdin(Stdio::null())
67+ .stdout(Stdio::null())
68+ .stderr(log);
69+ unsafe {
70+ cmd.pre_exec(|| {
71+ libc::setsid();
72+ Ok(())
73+ });
74+ }
75+ let mut child = cmd.spawn()?;
76+ for _ in 0..200 {
77+ if let Some(s) = connect() {
78+ let _ = child.wait();
79+ return Ok(s);
80+ }
81+ if let Ok(Some(status)) = child.try_wait()
82+ && !status.success()
83+ {
84+ return Err(io::Error::other(format!(
85+ "daemon died on startup ({status}), see {}",
86+ paths::state_dir().join("daemon.log").display()
87+ )));
88+ }
89+ thread::sleep(Duration::from_millis(10));
90+ }
91+ Err(io::Error::other("daemon never showed up"))
92+}
Added file: src/config.rs+380 additions
@@ -0,0 +1,380 @@
1+use std::{collections::BTreeMap, fs, path::PathBuf};
2+
3+use serde::Deserialize;
4+
5+use crate::{
6+ paths,
7+ units::{Amount, Bytes, Dur},
8+};
9+
10+#[derive(Debug, Deserialize)]
11+#[serde(default, deny_unknown_fields)]
12+pub struct Config {
13+ pub budget: Budget,
14+ pub heavy: Heavy,
15+ pub sizes: BTreeMap<String, Preset>,
16+ pub queue: Queue,
17+ pub overrun: Overrun,
18+ pub enforce: Enforce,
19+ pub artifacts: Artifacts,
20+ pub nix: Nix,
21+ pub detect: Detect,
22+}
23+
24+#[derive(Debug, Deserialize)]
25+#[serde(default, deny_unknown_fields)]
26+pub struct Budget {
27+ pub ram: Amount<Bytes>,
28+ pub ram_reserve: Bytes,
29+ pub cores: Amount<u32>,
30+ pub disk_reserve: Bytes,
31+ /// every lease's disk estimate is charged to all of these as well as its
32+ /// cwd. pessimistic when they share a fs, which is the point
33+ pub disk_paths: Vec<PathBuf>,
34+ /// memory psi `some avg10`. only gates heavies
35+ pub psi_max: Option<f64>,
36+}
37+
38+#[derive(Debug, Deserialize)]
39+#[serde(default, deny_unknown_fields)]
40+pub struct Heavy {
41+ pub ram: Bytes,
42+ pub cores: Amount<u32>,
43+ pub disk: Bytes,
44+}
45+
46+#[derive(Clone, Copy, Debug, Deserialize)]
47+#[serde(deny_unknown_fields)]
48+pub struct Preset {
49+ pub ram: Bytes,
50+ pub disk: Bytes,
51+ pub cores: Amount<u32>,
52+ pub time: Dur,
53+}
54+
55+#[derive(Debug, Deserialize)]
56+#[serde(default, deny_unknown_fields)]
57+pub struct Queue {
58+ pub linger: Dur,
59+ pub backfill_max_time: Dur,
60+ pub max_wait: Option<Dur>,
61+}
62+
63+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
64+#[serde(rename_all = "lowercase")]
65+pub enum RamOverrun {
66+ Warn,
67+ Throttle,
68+}
69+
70+#[derive(Debug, Deserialize)]
71+#[serde(default, deny_unknown_fields)]
72+pub struct Overrun {
73+ pub ram: RamOverrun,
74+ /// multiple of the ram estimate
75+ pub ram_slack: f64,
76+ /// multiple of the time estimate
77+ pub time_grace: f64,
78+ pub remind_every: Dur,
79+}
80+
81+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
82+#[serde(rename_all = "lowercase")]
83+pub enum Unwrapped {
84+ Deny,
85+ Wrap,
86+ Allow,
87+}
88+
89+#[derive(Debug, Deserialize)]
90+#[serde(default, deny_unknown_fields)]
91+pub struct Enforce {
92+ pub unwrapped: Unwrapped,
93+ pub default_size: String,
94+}
95+
96+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
97+#[serde(rename_all = "lowercase")]
98+pub enum OverCap {
99+ Warn,
100+ Block,
101+}
102+
103+#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
104+#[serde(rename_all = "lowercase")]
105+pub enum StopPolicy {
106+ Block,
107+ Nudge,
108+ Warn,
109+ Off,
110+}
111+
112+#[derive(Clone, Debug, Deserialize)]
113+#[serde(deny_unknown_fields)]
114+pub struct Rule {
115+ /// trailing `*` for a prefix match, that's all the glob you get
116+ pub name: String,
117+ /// sibling that must exist, so we don't go deleting someone's `build/`
118+ /// because it had a suggestive name
119+ pub marker: Option<String>,
120+ pub clean: Option<Vec<String>>,
121+}
122+
123+#[derive(Debug, Deserialize)]
124+#[serde(default, deny_unknown_fields)]
125+pub struct Artifacts {
126+ pub project_cap: Bytes,
127+ pub global_cap: Bytes,
128+ pub idle_after: Dur,
129+ pub size_ttl: Dur,
130+ pub over_cap: OverCap,
131+ pub stop: StopPolicy,
132+ pub rules: Vec<Rule>,
133+ pub extra_rules: Vec<Rule>,
134+ pub nix_gc: bool,
135+ pub nix_min_free: Bytes,
136+}
137+
138+#[derive(Debug, Deserialize)]
139+#[serde(default, deny_unknown_fields)]
140+pub struct Nix {
141+ pub inject: bool,
142+ pub cores_per_job: u32,
143+}
144+
145+#[derive(Debug, Deserialize)]
146+#[serde(default, deny_unknown_fields)]
147+pub struct Detect {
148+ pub commands: Vec<String>,
149+ pub extra: Vec<String>,
150+}
151+
152+impl Default for Config {
153+ fn default() -> Self {
154+ let preset = |ram, disk, cores: &str, time| Preset {
155+ ram,
156+ disk,
157+ cores: cores.parse().unwrap_or(Amount::Abs(1)),
158+ time: Dur::secs(time),
159+ };
160+ let sizes = BTreeMap::from([
161+ (
162+ "s".into(),
163+ preset(Bytes::mib(512), Bytes::mib(256), "1", 60),
164+ ),
165+ ("m".into(), preset(Bytes::gib(2), Bytes::gib(1), "4", 300)),
166+ ("l".into(), preset(Bytes::gib(6), Bytes::gib(5), "50%", 900)),
167+ (
168+ "xl".into(),
169+ preset(Bytes::gib(12), Bytes::gib(20), "100%", 2700),
170+ ),
171+ ]);
172+ Self {
173+ budget: Budget::default(),
174+ heavy: Heavy::default(),
175+ sizes,
176+ queue: Queue::default(),
177+ overrun: Overrun::default(),
178+ enforce: Enforce::default(),
179+ artifacts: Artifacts::default(),
180+ nix: Nix::default(),
181+ detect: Detect::default(),
182+ }
183+ }
184+}
185+
186+impl Default for Budget {
187+ fn default() -> Self {
188+ Self {
189+ ram: Amount::Percent(75.0),
190+ ram_reserve: Bytes::gib(4),
191+ cores: Amount::Percent(100.0),
192+ disk_reserve: Bytes::gib(20),
193+ disk_paths: vec!["/nix/store".into()],
194+ psi_max: Some(40.0),
195+ }
196+ }
197+}
198+
199+impl Default for Heavy {
200+ fn default() -> Self {
201+ Self {
202+ ram: Bytes::gib(4),
203+ cores: Amount::Percent(50.0),
204+ disk: Bytes::gib(5),
205+ }
206+ }
207+}
208+
209+impl Default for Queue {
210+ fn default() -> Self {
211+ Self {
212+ linger: Dur::secs(5),
213+ backfill_max_time: Dur::secs(120),
214+ max_wait: None,
215+ }
216+ }
217+}
218+
219+impl Default for Overrun {
220+ fn default() -> Self {
221+ Self {
222+ ram: RamOverrun::Warn,
223+ ram_slack: 1.25,
224+ time_grace: 1.5,
225+ remind_every: Dur::secs(300),
226+ }
227+ }
228+}
229+
230+impl Default for Enforce {
231+ fn default() -> Self {
232+ Self {
233+ unwrapped: Unwrapped::Deny,
234+ default_size: "m".into(),
235+ }
236+ }
237+}
238+
239+impl Default for Artifacts {
240+ fn default() -> Self {
241+ let rule = |name: &str, marker: Option<&str>| Rule {
242+ name: name.into(),
243+ marker: marker.map(Into::into),
244+ clean: None,
245+ };
246+ Self {
247+ project_cap: Bytes::gib(10),
248+ global_cap: Bytes::gib(60),
249+ idle_after: Dur::secs(86_400),
250+ size_ttl: Dur::secs(600),
251+ over_cap: OverCap::Warn,
252+ stop: StopPolicy::Nudge,
253+ rules: vec![
254+ rule("target", Some("Cargo.toml")),
255+ rule("result*", None),
256+ rule("node_modules", Some("package.json")),
257+ rule("zig-out", Some("build.zig")),
258+ rule(".zig-cache", Some("build.zig")),
259+ rule("_build", Some("dune-project")),
260+ rule("dist-newstyle", Some("cabal.project")),
261+ ],
262+ extra_rules: Vec::new(),
263+ nix_gc: false,
264+ nix_min_free: Bytes::gib(30),
265+ }
266+ }
267+}
268+
269+impl Default for Nix {
270+ fn default() -> Self {
271+ Self {
272+ inject: true,
273+ cores_per_job: 4,
274+ }
275+ }
276+}
277+
278+impl Default for Detect {
279+ fn default() -> Self {
280+ let commands = [
281+ "cargo build",
282+ "cargo b",
283+ "cargo test",
284+ "cargo t",
285+ "cargo check",
286+ "cargo c",
287+ "cargo clippy",
288+ "cargo run",
289+ "cargo r",
290+ "cargo nextest",
291+ "cargo bench",
292+ "cargo install",
293+ "cargo doc",
294+ "nix build",
295+ "nix flake check",
296+ "nix develop",
297+ "nix run",
298+ "nix shell",
299+ "nix-build",
300+ "nix-shell",
301+ "nixos-rebuild",
302+ "home-manager switch",
303+ "home-manager build",
304+ "make",
305+ "gmake",
306+ "cmake --build",
307+ "ninja",
308+ "meson compile",
309+ "go build",
310+ "go test",
311+ "zig build",
312+ "dune build",
313+ "cabal build",
314+ "stack build",
315+ "npm run build",
316+ "npm test",
317+ "npm ci",
318+ "npm install",
319+ "pnpm build",
320+ "pnpm install",
321+ "yarn build",
322+ "bun run build",
323+ "docker build",
324+ "podman build",
325+ "mvn",
326+ "gradle",
327+ "./gradlew",
328+ "bazel build",
329+ "bazel test",
330+ ];
331+ Self {
332+ commands: commands.into_iter().map(Into::into).collect(),
333+ extra: Vec::new(),
334+ }
335+ }
336+}
337+
338+impl Config {
339+ pub fn load() -> Result<Self, String> {
340+ let path = paths::config_file();
341+ match fs::read_to_string(&path) {
342+ Ok(text) => Self::parse(&text).map_err(|e| format!("{}: {e}", path.display())),
343+ Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
344+ Err(e) => Err(format!("{}: {e}", path.display())),
345+ }
346+ }
347+
348+ pub fn parse(text: &str) -> Result<Self, toml::de::Error> {
349+ let mut config: Self = toml::from_str(text)?;
350+ for (name, preset) in Self::default().sizes {
351+ config.sizes.entry(name).or_insert(preset);
352+ }
353+ Ok(config)
354+ }
355+
356+ pub fn preset(&self, name: &str) -> Result<Preset, String> {
357+ self.sizes.get(name).copied().ok_or_else(|| {
358+ let known: Vec<_> = self.sizes.keys().map(String::as_str).collect();
359+ format!(
360+ "unknown size '{name}', known sizes are {}",
361+ known.join(", ")
362+ )
363+ })
364+ }
365+
366+ pub fn artifact_rules(&self) -> impl Iterator<Item = &Rule> {
367+ self.artifacts
368+ .rules
369+ .iter()
370+ .chain(&self.artifacts.extra_rules)
371+ }
372+
373+ pub fn detect_patterns(&self) -> impl Iterator<Item = &str> {
374+ self.detect
375+ .commands
376+ .iter()
377+ .chain(&self.detect.extra)
378+ .map(String::as_str)
379+ }
380+}
Added file: src/daemon.rs+526 additions
@@ -0,0 +1,526 @@
1+use std::{
2+ collections::BTreeMap,
3+ fs, io,
4+ os::unix::net::{UnixListener, UnixStream},
5+ path::PathBuf,
6+ process,
7+ sync::{Arc, Mutex, MutexGuard},
8+ thread,
9+ time::{Duration, Instant},
10+};
11+
12+use crate::{
13+ client,
14+ config::Config,
15+ paths,
16+ proto::{self, Acquire, Event, LeaseId, LeaseInfo, Reader, Request, Status},
17+ sched::{self, Capacity, Disk, FsId, Need, Queued, Reason, Running, Verdict},
18+ sys,
19+ units::Bytes,
20+};
21+
22+struct Run {
23+ started: Instant,
24+ rss: Bytes,
25+ reminded: Option<Instant>,
26+}
27+
28+struct Lease {
29+ req: Acquire,
30+ heavy: bool,
31+ fs: Vec<FsId>,
32+ queued_at: Instant,
33+ run: Option<Run>,
34+ out: UnixStream,
35+ waiting: Option<Reason>,
36+ last_pos: usize,
37+ watchers: Vec<UnixStream>,
38+}
39+
40+impl Lease {
41+ fn elapsed(&self) -> Duration {
42+ self.run
43+ .as_ref()
44+ .map_or(self.queued_at, |r| r.started)
45+ .elapsed()
46+ }
47+
48+ fn overdue(&self, grace: f64) -> bool {
49+ self.run.is_some() && self.elapsed() > self.req.need.time.mul_f64(grace)
50+ }
51+}
52+
53+struct State {
54+ cfg: Config,
55+ next: LeaseId,
56+ leases: BTreeMap<LeaseId, Lease>,
57+ cap: Capacity,
58+ baselines: BTreeMap<FsId, Bytes>,
59+ conns: usize,
60+ idle_since: Instant,
61+}
62+
63+type Shared = Arc<Mutex<State>>;
64+
65+fn lock(shared: &Shared) -> MutexGuard<'_, State> {
66+ shared
67+ .lock()
68+ .unwrap_or_else(std::sync::PoisonError::into_inner)
69+}
70+
71+impl State {
72+ fn new(cfg: Config) -> Self {
73+ let (total, live) = sys::memory();
74+ let cap = Capacity {
75+ ram_budget: cfg.budget.ram.resolve(total),
76+ ram_live: live,
77+ ram_reserve: cfg.budget.ram_reserve,
78+ cores: cfg.budget.cores.resolve(sys::cores()),
79+ disks: BTreeMap::new(),
80+ pressure: false,
81+ };
82+ Self {
83+ cfg,
84+ next: 1,
85+ leases: BTreeMap::new(),
86+ cap,
87+ baselines: BTreeMap::new(),
88+ conns: 0,
89+ idle_since: Instant::now(),
90+ }
91+ }
92+
93+ fn heavy(&self, need: &Need) -> bool {
94+ let h = &self.cfg.heavy;
95+ need.ram >= h.ram || need.disk >= h.disk || need.cores >= h.cores.resolve(sys::cores())
96+ }
97+
98+ /// capped at a tenth of the fs, or a 16G tmpfs /tmp with a 20G reserve
99+ /// admits nothing ever
100+ fn reserve(&self, size: Bytes) -> Bytes {
101+ self.cfg.budget.disk_reserve.min(Bytes(size.0 / 10))
102+ }
103+
104+ fn fs_for(&mut self, cwd: &PathBuf) -> Vec<FsId> {
105+ let mut ids = Vec::new();
106+ for path in std::iter::once(cwd).chain(&self.cfg.budget.disk_paths) {
107+ if let Some(fs) = sys::fs(path) {
108+ if !ids.contains(&fs.id) {
109+ ids.push(fs.id);
110+ }
111+ let reserve = self.reserve(fs.size);
112+ self.cap.disks.entry(fs.id).or_insert(Disk {
113+ free: fs.free,
114+ reserve,
115+ consumed: Bytes::ZERO,
116+ });
117+ }
118+ }
119+ ids
120+ }
121+
122+ fn refresh(&mut self) {
123+ let (_, live) = sys::memory();
124+ self.cap.ram_live = live;
125+ self.cap.pressure = self
126+ .cfg
127+ .budget
128+ .psi_max
129+ .is_some_and(|max| sys::pressure().is_some_and(|p| p > max));
130+
131+ let mut paths: Vec<PathBuf> = self.leases.values().map(|l| l.req.cwd.clone()).collect();
132+ paths.extend(self.cfg.budget.disk_paths.iter().cloned());
133+ let mut disks = BTreeMap::new();
134+ for path in paths {
135+ if let Some(fs) = sys::fs(&path) {
136+ disks.insert(fs.id, (fs.free, self.reserve(fs.size)));
137+ }
138+ }
139+
140+ let busy: Vec<FsId> = self
141+ .leases
142+ .values()
143+ .filter(|l| l.run.is_some())
144+ .flat_map(|l| l.fs.iter().copied())
145+ .collect();
146+ self.baselines.retain(|id, _| busy.contains(id));
147+ self.cap.disks = disks
148+ .into_iter()
149+ .map(|(id, (free, reserve))| {
150+ let consumed = if busy.contains(&id) {
151+ self.baselines
152+ .entry(id)
153+ .or_insert(free)
154+ .saturating_sub(free)
155+ } else {
156+ Bytes::ZERO
157+ };
158+ (
159+ id,
160+ Disk {
161+ free,
162+ reserve,
163+ consumed,
164+ },
165+ )
166+ })
167+ .collect();
168+ }
169+
170+ fn running(&self) -> Vec<Running> {
171+ self.leases
172+ .values()
173+ .filter_map(|l| {
174+ let run = l.run.as_ref()?;
175+ Some(Running {
176+ need: l.req.need,
177+ heavy: l.heavy,
178+ fs: l.fs.clone(),
179+ rss: run.rss,
180+ remaining: l.req.need.time.saturating_sub(run.started.elapsed()),
181+ })
182+ })
183+ .collect()
184+ }
185+
186+ fn reschedule(&mut self) {
187+ let queue: Vec<Queued> = self
188+ .leases
189+ .iter()
190+ .filter(|(_, l)| l.run.is_none())
191+ .map(|(id, l)| Queued {
192+ id: *id,
193+ need: l.req.need,
194+ heavy: l.heavy,
195+ fs: l.fs.clone(),
196+ })
197+ .collect();
198+ if queue.is_empty() {
199+ return;
200+ }
201+ let verdicts = sched::plan(
202+ &queue,
203+ &self.running(),
204+ &self.cap,
205+ self.cfg.queue.backfill_max_time.0,
206+ );
207+ let cores = self.cap.cores;
208+ let mut position = 0;
209+ for (id, verdict) in verdicts {
210+ let Some(lease) = self.leases.get_mut(&id) else {
211+ continue;
212+ };
213+ match verdict {
214+ Verdict::Admit => {
215+ lease.run = Some(Run {
216+ started: Instant::now(),
217+ rss: Bytes::ZERO,
218+ reminded: None,
219+ });
220+ let _ = proto::send(
221+ &mut lease.out,
222+ &Event::Granted {
223+ lease: id,
224+ cores: lease.req.need.cores.clamp(1, cores.max(1)),
225+ heavy: lease.heavy,
226+ },
227+ );
228+ }
229+ Verdict::Wait(reason) => {
230+ position += 1;
231+ lease.waiting = Some(reason);
232+ if lease.last_pos != position {
233+ lease.last_pos = position;
234+ let _ = proto::send(
235+ &mut lease.out,
236+ &Event::Queued {
237+ lease: id,
238+ position,
239+ reason: Some(reason),
240+ },
241+ );
242+ }
243+ }
244+ }
245+ }
246+ }
247+
248+ fn remove(&mut self, id: LeaseId) {
249+ if let Some(lease) = self.leases.remove(&id) {
250+ for mut w in lease.watchers {
251+ let _ = proto::send(&mut w, &Event::Fired { overdue: false });
252+ }
253+ }
254+ if self.leases.is_empty() {
255+ self.idle_since = Instant::now();
256+ }
257+ self.refresh();
258+ self.reschedule();
259+ }
260+
261+ fn fire_overdue(&mut self) {
262+ let grace = self.cfg.overrun.time_grace;
263+ for lease in self.leases.values_mut() {
264+ if !lease.watchers.is_empty() && lease.overdue(grace) {
265+ for mut w in lease.watchers.drain(..) {
266+ let _ = proto::send(&mut w, &Event::Fired { overdue: true });
267+ }
268+ }
269+ }
270+ }
271+
272+ fn info(&self) -> Vec<LeaseInfo> {
273+ let grace = self.cfg.overrun.time_grace;
274+ let mut position = 0;
275+ self.leases
276+ .iter()
277+ .map(|(id, l)| {
278+ if l.run.is_none() {
279+ position += 1;
280+ }
281+ LeaseInfo {
282+ id: *id,
283+ cmd: l.req.cmd.clone(),
284+ cwd: l.req.cwd.clone(),
285+ session: l.req.session.clone(),
286+ need: l.req.need,
287+ heavy: l.heavy,
288+ running: l.run.is_some(),
289+ elapsed: l.elapsed().as_secs(),
290+ rss: l.run.as_ref().map_or(Bytes::ZERO, |r| r.rss),
291+ overdue: l.overdue(grace),
292+ waiting: l.waiting.filter(|_| l.run.is_none()),
293+ position: if l.run.is_none() { position } else { 0 },
294+ }
295+ })
296+ .collect()
297+ }
298+
299+ fn status(&self) -> Status {
300+ let (ram_free, cores_free) = sched::headroom(&self.running(), &self.cap);
301+ Status {
302+ leases: self.info(),
303+ ram_free,
304+ ram_budget: self.cap.ram_budget,
305+ cores_free,
306+ cores: self.cap.cores,
307+ }
308+ }
309+}
310+
311+fn serve(shared: &Shared, mut stream: UnixStream) -> io::Result<()> {
312+ let mut reader = Reader::new(&stream)?;
313+ let Some(req) = reader.recv::<Request>()? else {
314+ return Ok(());
315+ };
316+ match req {
317+ Request::Acquire(req) => {
318+ let id = {
319+ let mut st = lock(shared);
320+ let id = st.next;
321+ st.next += 1;
322+ let fs = st.fs_for(&req.cwd);
323+ let heavy = st.heavy(&req.need);
324+ st.leases.insert(
325+ id,
326+ Lease {
327+ req,
328+ heavy,
329+ fs,
330+ queued_at: Instant::now(),
331+ run: None,
332+ out: stream.try_clone()?,
333+ waiting: None,
334+ last_pos: 0,
335+ watchers: Vec::new(),
336+ },
337+ );
338+ st.reschedule();
339+ id
340+ };
341+ hold(shared, id, &mut reader);
342+ lock(shared).remove(id);
343+ Ok(())
344+ }
345+ Request::Status => {
346+ let status = lock(shared).status();
347+ proto::send(&mut stream, &Event::Status(status))
348+ }
349+ Request::Kill { lease } => {
350+ let mut st = lock(shared);
351+ let reply = match st.leases.get_mut(&lease) {
352+ Some(l) => {
353+ let _ = proto::send(&mut l.out, &Event::Kill);
354+ Event::Ok
355+ }
356+ None => Event::Err {
357+ msg: format!("no lease {lease}"),
358+ },
359+ };
360+ drop(st);
361+ proto::send(&mut stream, &reply)
362+ }
363+ Request::Extend { lease, by } => {
364+ let mut st = lock(shared);
365+ let reply = match st.leases.get_mut(&lease) {
366+ Some(l) => {
367+ l.req.need.time += by.0;
368+ if let Some(run) = &mut l.run {
369+ run.reminded = None;
370+ }
371+ Event::Ok
372+ }
373+ None => Event::Err {
374+ msg: format!("no lease {lease}"),
375+ },
376+ };
377+ drop(st);
378+ proto::send(&mut stream, &reply)
379+ }
380+ Request::Watch { lease } => {
381+ let mut st = lock(shared);
382+ let grace = st.cfg.overrun.time_grace;
383+ match st.leases.get_mut(&lease) {
384+ Some(l) if l.overdue(grace) => {
385+ proto::send(&mut stream, &Event::Fired { overdue: true })
386+ }
387+ Some(l) => {
388+ l.watchers.push(stream);
389+ Ok(())
390+ }
391+ None => proto::send(&mut stream, &Event::Fired { overdue: false }),
392+ }
393+ }
394+ Request::Overdue { session } => {
395+ let mut st = lock(shared);
396+ let grace = st.cfg.overrun.time_grace;
397+ let every = st.cfg.overrun.remind_every.0;
398+ let mut due = Vec::new();
399+ for (id, l) in &mut st.leases {
400+ if l.req.session.as_deref() != Some(&session) || !l.overdue(grace) {
401+ continue;
402+ }
403+ let Some(run) = &mut l.run else { continue };
404+ if run.reminded.is_some_and(|t| t.elapsed() < every) {
405+ continue;
406+ }
407+ run.reminded = Some(Instant::now());
408+ due.push(*id);
409+ }
410+ let leases = st
411+ .info()
412+ .into_iter()
413+ .filter(|i| due.contains(&i.id))
414+ .collect();
415+ drop(st);
416+ proto::send(&mut stream, &Event::Overdue { leases })
417+ }
418+ Request::Started { .. } | Request::Usage { .. } | Request::Done => proto::send(
419+ &mut stream,
420+ &Event::Err {
421+ msg: "no lease on this connection".into(),
422+ },
423+ ),
424+ }
425+}
426+
427+/// a client that vanishes without saying done got sigkilled or similar, and
428+/// pdeathsig only reaches its direct child. take the rest of the group down
429+/// so a headless build doesn't keep eating ram off the books
430+fn hold(shared: &Shared, id: LeaseId, reader: &mut Reader) {
431+ let mut pgid = None;
432+ let clean = loop {
433+ match reader.recv::<Request>() {
434+ Ok(Some(Request::Usage { rss })) => {
435+ if let Some(run) = lock(shared)
436+ .leases
437+ .get_mut(&id)
438+ .and_then(|l| l.run.as_mut())
439+ {
440+ run.rss = rss;
441+ }
442+ }
443+ Ok(Some(Request::Started { pid })) => pgid = Some(pid),
444+ Ok(Some(Request::Done)) => break true,
445+ Ok(Some(_)) => {}
446+ Ok(None) | Err(_) => break false,
447+ }
448+ };
449+ if !clean && let Some(pgid) = pgid.filter(|p| *p > 1) {
450+ unsafe { libc::kill(-pgid, libc::SIGTERM) };
451+ }
452+}
453+
454+fn tick(shared: &Shared) {
455+ loop {
456+ thread::sleep(Duration::from_secs(1));
457+ let mut st = lock(shared);
458+ st.refresh();
459+ st.reschedule();
460+ st.fire_overdue();
461+ if st.conns == 0 && st.leases.is_empty() && st.idle_since.elapsed() >= st.cfg.queue.linger.0
462+ {
463+ // a spawner holds this while it clears a stale socket, so taking it
464+ // here stops us unlinking our successor's socket on the way out
465+ if let Some(_guard) = try_spawn_lock() {
466+ let _ = fs::remove_file(paths::socket());
467+ process::exit(0);
468+ }
469+ }
470+ }
471+}
472+
473+fn try_spawn_lock() -> Option<client::Flock> {
474+ let file = fs::OpenOptions::new()
475+ .write(true)
476+ .open(paths::spawn_lock())
477+ .ok()?;
478+ client::Flock::try_exclusive(file)
479+}
480+
481+pub fn main(cfg: Config) -> ! {
482+ // fork once more so the client that spawned us can reap its child now
483+ // instead of carrying a zombie for the length of its build
484+ match unsafe { libc::fork() } {
485+ -1 => {
486+ eprintln!("roundabout: fork failed: {}", io::Error::last_os_error());
487+ process::exit(1);
488+ }
489+ 0 => {}
490+ _ => process::exit(0),
491+ }
492+
493+ let listener = match UnixListener::bind(paths::socket()) {
494+ Ok(l) => l,
495+ Err(e) if e.kind() == io::ErrorKind::AddrInUse && client::connect().is_some() => {
496+ process::exit(0)
497+ }
498+ Err(e) => {
499+ eprintln!("roundabout: bind: {e}");
500+ process::exit(1);
501+ }
502+ };
503+
504+ let shared: Shared = Arc::new(Mutex::new(State::new(cfg)));
505+ {
506+ let shared = shared.clone();
507+ thread::spawn(move || tick(&shared));
508+ }
509+
510+ for stream in listener.incoming() {
511+ let Ok(stream) = stream else { continue };
512+ lock(&shared).conns += 1;
513+ let shared = shared.clone();
514+ thread::spawn(move || {
515+ if let Err(e) = serve(&shared, stream) {
516+ eprintln!("roundabout: {e}");
517+ }
518+ let mut st = lock(&shared);
519+ st.conns -= 1;
520+ if st.leases.is_empty() {
521+ st.idle_since = Instant::now();
522+ }
523+ });
524+ }
525+ process::exit(0)
526+}
Added file: src/history.rs+89 additions
@@ -0,0 +1,89 @@
1+use std::{
2+ collections::BTreeMap,
3+ path::Path,
4+ time::{Duration, SystemTime, UNIX_EPOCH},
5+};
6+
7+use serde::{Deserialize, Serialize};
8+
9+use crate::{sched::Need, store, units::Bytes};
10+
11+const KEEP: usize = 1000;
12+
13+#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
14+pub struct Record {
15+ pub runs: u32,
16+ pub peak_rss: Bytes,
17+ /// ewma, so one freak run doesn't define the command forever
18+ pub wall: u64,
19+ pub cpu: u64,
20+ pub cores: u32,
21+ pub last: u64,
22+}
23+
24+type History = BTreeMap<String, Record>;
25+
26+pub fn key(root: &Path, cmd: &[String]) -> String {
27+ let mut k = format!("{}\t{}", root.display(), cmd.join(" "));
28+ k.truncate(k.floor_char_boundary(400));
29+ k
30+}
31+
32+fn now() -> u64 {
33+ SystemTime::now()
34+ .duration_since(UNIX_EPOCH)
35+ .map_or(0, |d| d.as_secs())
36+}
37+
38+pub fn record(key: String, peak_rss: Bytes, wall: Duration, cpu: Duration, cores: u32) {
39+ let _ = store::update("history", |h: &mut History| {
40+ let wall = wall.as_secs();
41+ h.entry(key)
42+ .and_modify(|r| {
43+ r.runs += 1;
44+ r.peak_rss = r.peak_rss.max(peak_rss);
45+ r.wall = (r.wall * 7 + wall * 3) / 10;
46+ r.cpu = cpu.as_secs();
47+ r.cores = cores;
48+ r.last = now();
49+ })
50+ .or_insert_with(|| Record {
51+ runs: 1,
52+ peak_rss,
53+ wall,
54+ cpu: cpu.as_secs(),
55+ cores,
56+ last: now(),
57+ });
58+ if h.len() > KEEP {
59+ let mut ages: Vec<_> = h.iter().map(|(k, r)| (r.last, k.clone())).collect();
60+ ages.sort();
61+ for (_, k) in ages.into_iter().take(h.len() - KEEP) {
62+ h.remove(&k);
63+ }
64+ }
65+ });
66+}
67+
68+pub fn lookup(key: &str) -> Option<Record> {
69+ store::load::<History>("history").get(key).copied()
70+}
71+
72+/// padded, since the agent that under-asks is the one that ooms the video.
73+/// nix builds happen in nix-daemon, so their rss never shows up in ours and
74+/// the base estimate stands
75+pub fn suggest(rec: &Record, base: Need, cmd: &[String]) -> Need {
76+ let via_daemon = cmd
77+ .iter()
78+ .any(|t| t.contains("nix") || t.contains("home-manager"));
79+ Need {
80+ ram: if via_daemon {
81+ base.ram
82+ } else {
83+ Bytes(rec.peak_rss.0 + rec.peak_rss.0 / 4).max(Bytes::mib(256))
84+ },
85+ cores: rec.cores.max(1),
86+ time: Duration::from_secs((rec.wall * 3 / 2).max(30)),
87+ ..base
88+ }
89+}
Added file: src/hook.rs+319 additions
@@ -0,0 +1,319 @@
1+use std::{
2+ io::{self, Read},
3+ path::{Path, PathBuf},
4+};
5+
6+use serde::Deserialize;
7+use serde_json::{Value, json};
8+
9+use crate::{
10+ artifacts, client,
11+ config::{Config, StopPolicy, Unwrapped},
12+ history, paths,
13+ proto::{self, Event, LeaseInfo, Reader, Request},
14+ units::Dur,
15+};
16+
17+#[derive(Deserialize)]
18+struct Input {
19+ #[serde(default)]
20+ session_id: String,
21+ #[serde(default)]
22+ cwd: Option<PathBuf>,
23+ #[serde(default)]
24+ tool_name: Option<String>,
25+ #[serde(default, rename = "tool_input")]
26+ tool_args: Option<Value>,
27+ #[serde(default)]
28+ stop_hook_active: bool,
29+}
30+
31+fn read_input() -> Option<Input> {
32+ let mut buf = String::new();
33+ io::stdin().read_to_string(&mut buf).ok()?;
34+ serde_json::from_str(&buf).ok()
35+}
36+
37+fn emit(v: &Value) {
38+ println!("{v}");
39+}
40+
41+/// quote-aware split on `&& || ; | & \n`. subshells, heredocs and friends get
42+/// treated as words, which errs toward seeing a build that isn't there
43+fn segments(cmd: &str) -> Vec<Vec<String>> {
44+ let mut out = Vec::new();
45+ let mut words = Vec::new();
46+ let mut word = String::new();
47+ let mut quote: Option<char> = None;
48+ let mut chars = cmd.chars().peekable();
49+ let flush_word = |word: &mut String, words: &mut Vec<String>| {
50+ if !word.is_empty() {
51+ words.push(std::mem::take(word));
52+ }
53+ };
54+ while let Some(c) = chars.next() {
55+ match (quote, c) {
56+ (Some(q), c) if c == q => quote = None,
57+ (Some('"') | None, '\\') => {
58+ if let Some(n) = chars.next() {
59+ word.push(n);
60+ }
61+ }
62+ (None, '\'' | '"') => quote = Some(c),
63+ (None, '&') if word.ends_with('>') || chars.peek() == Some(&'>') => word.push(c),
64+ (None, '&' | '|' | ';' | '\n' | '(' | ')') => {
65+ flush_word(&mut word, &mut words);
66+ if !words.is_empty() {
67+ out.push(std::mem::take(&mut words));
68+ }
69+ }
70+ (None, c) if c.is_whitespace() => flush_word(&mut word, &mut words),
71+ (_, c) => word.push(c),
72+ }
73+ }
74+ flush_word(&mut word, &mut words);
75+ if !words.is_empty() {
76+ out.push(words);
77+ }
78+ out
79+}
80+
81+/// peels env assignments and wrappers that don't change what runs
82+fn head(seg: &[String]) -> &[String] {
83+ let mut s = seg;
84+ loop {
85+ match s.first().map(String::as_str) {
86+ Some(w) if w.contains('=') && !w.starts_with('-') && !w.starts_with('=') => s = &s[1..],
87+ Some("time" | "nice" | "nohup" | "command" | "exec" | "env" | "sudo" | "doas") => {
88+ s = &s[1..];
89+ }
90+ Some("timeout") => {
91+ s = &s[1..];
92+ while s.first().is_some_and(|w| w.starts_with('-')) {
93+ s = &s[1..];
94+ }
95+ s = s.get(1..).unwrap_or_default();
96+ }
97+ _ => return s,
98+ }
99+ }
100+}
101+
102+fn is_rabt(word: &str) -> bool {
103+ word == "rabt" || word.ends_with("/rabt")
104+}
105+
106+fn matches_pattern(seg: &[String], pattern: &str) -> bool {
107+ let want: Vec<&str> = pattern.split_whitespace().collect();
108+ let first_ok = |w: &str, p: &str| w == p || w.rsplit('/').next() == Some(p);
109+ !want.is_empty()
110+ && seg.len() >= want.len()
111+ && first_ok(&seg[0], want[0])
112+ && seg[1..want.len()]
113+ .iter()
114+ .zip(&want[1..])
115+ .all(|(w, p)| w == p)
116+}
117+
118+enum Scan {
119+ Clean,
120+ Unwrapped(Vec<String>),
121+}
122+
123+fn scan(cfg: &Config, cmd: &str) -> Scan {
124+ let mut found = Vec::new();
125+ for seg in segments(cmd) {
126+ let seg = head(&seg);
127+ if seg.first().is_some_and(|w| is_rabt(w)) {
128+ continue;
129+ }
130+ if cfg.detect_patterns().any(|p| matches_pattern(seg, p)) {
131+ let words: Vec<&str> = seg
132+ .iter()
133+ .map(String::as_str)
134+ .take_while(|w| !w.contains('>') && !w.contains('<'))
135+ .collect();
136+ found.push(words.join(" "));
137+ }
138+ }
139+ if found.is_empty() {
140+ Scan::Clean
141+ } else {
142+ Scan::Unwrapped(found)
143+ }
144+}
145+
146+fn guess(cfg: &Config, cwd: &Path, cmd: &[String]) -> String {
147+ let key = history::key(&paths::project_root(cwd), cmd);
148+ history::lookup(&key).map_or_else(
149+ || format!("--size {}", cfg.enforce.default_size),
150+ |rec| {
151+ crate::run::base(cfg).map_or_else(
152+ |_| format!("--size {}", cfg.enforce.default_size),
153+ |b| {
154+ let n = history::suggest(&rec, b, cmd);
155+ format!("--ram {} --cores {} --time {}", n.ram, n.cores, Dur(n.time))
156+ },
157+ )
158+ },
159+ )
160+}
161+
162+fn shell_quote(s: &str) -> String {
163+ format!("'{}'", s.replace('\'', r"'\''"))
164+}
165+
166+fn pre_tool_use(cfg: &Config, input: Input) {
167+ if input.tool_name.as_deref() != Some("Bash") {
168+ return;
169+ }
170+ let Some(mut tool_input) = input.tool_args else {
171+ return;
172+ };
173+ let Some(cmd) = tool_input
174+ .get("command")
175+ .and_then(Value::as_str)
176+ .map(str::to_owned)
177+ else {
178+ return;
179+ };
180+ let cwd = input.cwd.unwrap_or_else(|| "/".into());
181+
182+ match (scan(cfg, &cmd), cfg.enforce.unwrapped) {
183+ (Scan::Clean, _) | (_, Unwrapped::Allow) => {}
184+ (Scan::Unwrapped(found), Unwrapped::Deny) => {
185+ let first: Vec<String> = found[0].split_whitespace().map(Into::into).collect();
186+ let hint = guess(cfg, &cwd, &first);
187+ let reason = format!(
188+ "builds go through roundabout here so parallel agents don't oom the box. wrap each build \
189+ in `rabt run <size> -- <cmd>`, e.g. `rabt run {hint} -- {}`. sizes are {}, or \
190+ give --ram/--disk/--cores/--time yourself. for anything long use \
191+ run_in_background and start `rabt watch <lease>` in the background too so you \
192+ hear if it drags. see the roundabout skill",
193+ found[0],
194+ cfg.sizes.keys().cloned().collect::<Vec<_>>().join("/"),
195+ );
196+ emit(&json!({
197+ "hookSpecificOutput": {
198+ "hookEventName": "PreToolUse",
199+ "permissionDecision": "deny",
200+ "permissionDecisionReason": reason,
201+ }
202+ }));
203+ }
204+ (Scan::Unwrapped(_), Unwrapped::Wrap) => {
205+ let hint = guess(cfg, &cwd, std::slice::from_ref(&cmd));
206+ tool_input["command"] =
207+ Value::String(format!("rabt run {hint} -- bash -c {}", shell_quote(&cmd)));
208+ // no permissionDecision on purpose. allow would skip the user's
209+ // prompt for whatever the agent wrote, rewriting shouldn't
210+ emit(&json!({
211+ "hookSpecificOutput": {
212+ "hookEventName": "PreToolUse",
213+ "updatedInput": tool_input,
214+ }
215+ }));
216+ }
217+ }
218+}
219+
220+fn overdue(session: &str) -> Vec<LeaseInfo> {
221+ if session.is_empty() {
222+ return Vec::new();
223+ }
224+ let Some(mut s) = client::connect() else {
225+ return Vec::new();
226+ };
227+ if proto::send(
228+ &mut s,
229+ &Request::Overdue {
230+ session: session.to_owned(),
231+ },
232+ )
233+ .is_err()
234+ {
235+ return Vec::new();
236+ }
237+ match Reader::new(&s).and_then(|mut r| r.recv::<Event>()) {
238+ Ok(Some(Event::Overdue { leases })) => leases,
239+ _ => Vec::new(),
240+ }
241+}
242+
243+/// same source `rabt run` reads. after `--continue` the env var can keep the
244+/// startup id while the json moves on, and we need both ends to agree
245+fn session(input: &Input) -> String {
246+ std::env::var("CLAUDE_CODE_SESSION_ID")
247+ .ok()
248+ .filter(|s| !s.is_empty())
249+ .unwrap_or_else(|| input.session_id.clone())
250+}
251+
252+fn remind(event: &str, input: &Input) {
253+ let late = overdue(&session(input));
254+ if late.is_empty() {
255+ return;
256+ }
257+ let lines: Vec<String> = late
258+ .iter()
259+ .map(|l| {
260+ format!(
261+ "roundabout lease {} (`{}`) has run {} against a {} estimate. `rabt kill {}` to stop it, \
262+ `rabt extend {} <dur>` if it's worth the wait",
263+ l.id,
264+ l.cmd.join(" "),
265+ Dur::secs(l.elapsed),
266+ Dur(l.need.time),
267+ l.id,
268+ l.id
269+ )
270+ })
271+ .collect();
272+ emit(&json!({
273+ "hookSpecificOutput": {
274+ "hookEventName": event,
275+ "additionalContext": lines.join("\n"),
276+ }
277+ }));
278+}
279+
280+fn stop(cfg: &Config, event: &str, input: &Input) {
281+ if input.stop_hook_active || cfg.artifacts.stop == StopPolicy::Off {
282+ return;
283+ }
284+ let root = input.cwd.as_deref().map(paths::project_root);
285+ let report = artifacts::check(cfg, root.as_deref(), false);
286+ if !report.over() {
287+ return;
288+ }
289+ let msg = format!(
290+ "{}. run `rabt artifacts gc --cap`, or `rabt artifacts list` and clear out what you \
291+ don't need",
292+ report.describe(cfg)
293+ );
294+ let out = match cfg.artifacts.stop {
295+ StopPolicy::Block => json!({ "decision": "block", "reason": msg }),
296+ StopPolicy::Nudge => json!({
297+ "hookSpecificOutput": { "hookEventName": event, "additionalContext": msg }
298+ }),
299+ StopPolicy::Warn => json!({ "systemMessage": msg }),
300+ StopPolicy::Off => return,
301+ };
302+ emit(&out);
303+}
304+
305+/// hooks never fail loudly. a broken roundabout shouldn't wedge the agent
306+pub fn main(cfg: &Config, event: &str) -> i32 {
307+ let Some(input) = read_input() else {
308+ return 0;
309+ };
310+ match event {
311+ "pre-tool-use" => pre_tool_use(cfg, input),
312+ "post-tool-use" => remind("PostToolUse", &input),
313+ "prompt-submit" => remind("UserPromptSubmit", &input),
314+ "stop" => stop(cfg, "Stop", &input),
315+ "subagent-stop" => stop(cfg, "SubagentStop", &input),
316+ other => eprintln!("rabt: unknown hook {other}"),
317+ }
318+ 0
319+}
Added file: src/main.rs+27 additions
@@ -0,0 +1,27 @@
1+mod artifacts;
2+mod cli;
3+mod client;
4+mod config;
5+mod daemon;
6+mod history;
7+mod hook;
8+mod paths;
9+mod proto;
10+mod run;
11+mod sched;
12+mod store;
13+mod sys;
14+mod units;
15+
16+use pound::Parse;
17+
18+fn main() {
19+ let cfg = match config::Config::load() {
20+ Ok(c) => c,
21+ Err(e) => {
22+ eprintln!("rabt: bad config, {e}");
23+ std::process::exit(2);
24+ }
25+ };
26+ std::process::exit(cli::dispatch(&cfg, cli::Rabt::parse()));
27+}
Added file: src/paths.rs+83 additions
@@ -0,0 +1,83 @@
1+use std::{
2+ env, fs,
3+ os::unix::fs::{DirBuilderExt, PermissionsExt},
4+ path::{Path, PathBuf},
5+};
6+
7+fn home() -> PathBuf {
8+ env::var_os("HOME").map_or_else(|| PathBuf::from("/"), PathBuf::from)
9+}
10+
11+fn xdg(var: &str, fallback: &str) -> PathBuf {
12+ env::var_os(var)
13+ .map(PathBuf::from)
14+ .filter(|p| p.is_absolute())
15+ .unwrap_or_else(|| home().join(fallback))
16+}
17+
18+pub fn config_file() -> PathBuf {
19+ env::var_os("ROUNDABOUT_CONFIG").map_or_else(
20+ || xdg("XDG_CONFIG_HOME", ".config").join("roundabout/config.toml"),
21+ PathBuf::from,
22+ )
23+}
24+
25+pub fn state_dir() -> PathBuf {
26+ let dir = env::var_os("ROUNDABOUT_STATE_DIR").map_or_else(
27+ || xdg("XDG_STATE_HOME", ".local/state").join("roundabout"),
28+ PathBuf::from,
29+ );
30+ let _ = fs::create_dir_all(&dir);
31+ dir
32+}
33+
34+pub fn runtime_dir() -> PathBuf {
35+ if let Some(dir) = env::var_os("ROUNDABOUT_RUNTIME_DIR") {
36+ let dir = PathBuf::from(dir);
37+ let _ = fs::DirBuilder::new()
38+ .recursive(true)
39+ .mode(0o700)
40+ .create(&dir);
41+ return dir;
42+ }
43+ if let Some(dir) = env::var_os("XDG_RUNTIME_DIR").map(PathBuf::from)
44+ && dir.is_absolute()
45+ && dir.is_dir()
46+ {
47+ return dir;
48+ }
49+ let uid = unsafe { libc::getuid() };
50+ let dir = env::temp_dir().join(format!("roundabout-{uid}"));
51+ let _ = fs::DirBuilder::new().mode(0o700).create(&dir);
52+ let _ = fs::set_permissions(&dir, fs::Permissions::from_mode(0o700));
53+ dir
54+}
55+
56+pub fn socket() -> PathBuf {
57+ runtime_dir().join("roundabout.sock")
58+}
59+
60+pub fn spawn_lock() -> PathBuf {
61+ runtime_dir().join("roundabout.lock")
62+}
63+
64+pub fn project_root(cwd: &Path) -> PathBuf {
65+ const MARKERS: [&str; 6] = [
66+ ".jj",
67+ ".git",
68+ ".hg",
69+ "flake.nix",
70+ "Cargo.toml",
71+ "package.json",
72+ ];
73+ let mut best = None;
74+ for dir in cwd.ancestors() {
75+ if dir.join(".jj").exists() || dir.join(".git").exists() || dir.join(".hg").exists() {
76+ return dir.to_path_buf();
77+ }
78+ if best.is_none() && MARKERS.iter().any(|m| dir.join(m).exists()) {
79+ best = Some(dir.to_path_buf());
80+ }
81+ }
82+ best.unwrap_or_else(|| cwd.to_path_buf())
83+}
Added file: src/proto.rs+115 additions
@@ -0,0 +1,115 @@
1+use std::{
2+ io::{self, BufRead, BufReader, Write},
3+ os::unix::net::UnixStream,
4+ path::PathBuf,
5+};
6+
7+use serde::{Deserialize, Serialize, de::DeserializeOwned};
8+
9+use crate::{
10+ sched::{Need, Reason},
11+ units::{Bytes, Dur},
12+};
13+
14+pub type LeaseId = u64;
15+
16+#[derive(Clone, Debug, Serialize, Deserialize)]
17+pub struct Acquire {
18+ pub need: Need,
19+ pub cwd: PathBuf,
20+ pub cmd: Vec<String>,
21+ pub session: Option<String>,
22+ pub deadline: Option<Dur>,
23+}
24+
25+#[derive(Debug, Serialize, Deserialize)]
26+#[serde(tag = "op", rename_all = "snake_case")]
27+pub enum Request {
28+ Acquire(Acquire),
29+ Started { pid: i32 },
30+ Usage { rss: Bytes },
31+ Done,
32+ Status,
33+ Kill { lease: LeaseId },
34+ Extend { lease: LeaseId, by: Dur },
35+ Watch { lease: LeaseId },
36+ Overdue { session: String },
37+}
38+
39+#[derive(Clone, Debug, Serialize, Deserialize)]
40+pub struct LeaseInfo {
41+ pub id: LeaseId,
42+ pub cmd: Vec<String>,
43+ pub cwd: PathBuf,
44+ pub session: Option<String>,
45+ pub need: Need,
46+ pub heavy: bool,
47+ pub running: bool,
48+ /// seconds since queued or started, whichever it's doing
49+ pub elapsed: u64,
50+ pub rss: Bytes,
51+ pub overdue: bool,
52+ pub waiting: Option<Reason>,
53+ pub position: usize,
54+}
55+
56+#[derive(Debug, Serialize, Deserialize)]
57+pub struct Status {
58+ pub leases: Vec<LeaseInfo>,
59+ pub ram_free: Bytes,
60+ pub ram_budget: Bytes,
61+ pub cores_free: u32,
62+ pub cores: u32,
63+}
64+
65+#[derive(Debug, Serialize, Deserialize)]
66+#[serde(tag = "ev", rename_all = "snake_case")]
67+pub enum Event {
68+ Queued {
69+ lease: LeaseId,
70+ position: usize,
71+ reason: Option<Reason>,
72+ },
73+ Granted {
74+ lease: LeaseId,
75+ cores: u32,
76+ heavy: bool,
77+ },
78+ Kill,
79+ Status(Status),
80+ Overdue {
81+ leases: Vec<LeaseInfo>,
82+ },
83+ Fired {
84+ overdue: bool,
85+ },
86+ Ok,
87+ Err {
88+ msg: String,
89+ },
90+}
91+
92+pub fn send<T: Serialize>(stream: &mut UnixStream, msg: &T) -> io::Result<()> {
93+ let mut line = serde_json::to_vec(msg)?;
94+ line.push(b'\n');
95+ stream.write_all(&line)
96+}
97+
98+pub struct Reader(BufReader<UnixStream>);
99+
100+impl Reader {
101+ pub fn new(stream: &UnixStream) -> io::Result<Self> {
102+ Ok(Self(BufReader::new(stream.try_clone()?)))
103+ }
104+
105+ /// `Ok(None)` on eof
106+ pub fn recv<T: DeserializeOwned>(&mut self) -> io::Result<Option<T>> {
107+ let mut line = String::new();
108+ if self.0.read_line(&mut line)? == 0 {
109+ return Ok(None);
110+ }
111+ serde_json::from_str(&line)
112+ .map(Some)
113+ .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
114+ }
115+}
Added file: src/run.rs+384 additions
@@ -0,0 +1,384 @@
1+use std::{
2+ env,
3+ fmt::Write as _,
4+ os::unix::{net::UnixStream, process::CommandExt},
5+ process::{Command, ExitStatus},
6+ sync::{
7+ Arc,
8+ atomic::{AtomicBool, AtomicI32, Ordering},
9+ },
10+ thread,
11+ time::{Duration, Instant},
12+};
13+
14+use crate::{
15+ artifacts, client,
16+ config::{Config, OverCap, RamOverrun},
17+ history, paths,
18+ proto::{self, Acquire, Event, Reader, Request},
19+ sched::Need,
20+ sys,
21+ units::{Amount, Bytes, Dur},
22+};
23+
24+pub const EXIT_TIMEOUT: i32 = 124;
25+pub const EXIT_GAVE_UP: i32 = 75;
26+
27+pub struct Args {
28+ pub size: Option<String>,
29+ pub ram: Option<Bytes>,
30+ pub disk: Option<Bytes>,
31+ pub cores: Option<Amount<u32>>,
32+ pub time: Option<Dur>,
33+ pub deadline: Option<Dur>,
34+ pub max_wait: Option<Dur>,
35+ pub session: Option<String>,
36+ pub cmd: Vec<String>,
37+}
38+
39+static PGID: AtomicI32 = AtomicI32::new(0);
40+
41+extern "C" fn forward(sig: libc::c_int) {
42+ let pgid = PGID.load(Ordering::Relaxed);
43+ if pgid > 1 {
44+ unsafe { libc::kill(-pgid, sig) };
45+ }
46+}
47+
48+fn say(msg: &str) {
49+ eprintln!("rabt: {msg}");
50+}
51+
52+fn need(cfg: &Config, a: &Args) -> Result<Need, String> {
53+ let base = cfg.preset(a.size.as_deref().unwrap_or(&cfg.enforce.default_size))?;
54+ Ok(Need {
55+ ram: a.ram.unwrap_or(base.ram),
56+ disk: a.disk.unwrap_or(base.disk),
57+ cores: a.cores.unwrap_or(base.cores).resolve(sys::cores()),
58+ time: a.time.unwrap_or(base.time).0,
59+ })
60+}
61+
62+pub fn base(cfg: &Config) -> Result<Need, String> {
63+ let p = cfg.preset(&cfg.enforce.default_size)?;
64+ Ok(Need {
65+ ram: p.ram,
66+ disk: p.disk,
67+ cores: p.cores.resolve(sys::cores()),
68+ time: p.time.0,
69+ })
70+}
71+
72+struct Grant {
73+ stream: UnixStream,
74+ lease: u64,
75+ cores: u32,
76+}
77+
78+fn acquire(req: &Acquire, max_wait: Option<Dur>) -> Result<Grant, i32> {
79+ let started = Instant::now();
80+ // a daemon on its way out can accept us then hang up, so go round again
81+ for _ in 0..3 {
82+ let mut stream = client::connect_or_spawn().map_err(|e| {
83+ say(&format!("can't reach the daemon: {e}"));
84+ 1
85+ })?;
86+ if proto::send(&mut stream, &Request::Acquire(req.clone())).is_err() {
87+ continue;
88+ }
89+ let _ = stream.set_read_timeout(Some(Duration::from_secs(1)));
90+ let mut reader = Reader::new(&stream).map_err(|_| 1)?;
91+ loop {
92+ match reader.recv::<Event>() {
93+ Ok(Some(Event::Granted {
94+ lease,
95+ cores,
96+ heavy,
97+ })) => {
98+ say(&format!(
99+ "lease {lease} go, cores {cores}{}",
100+ if heavy { ", heavy" } else { "" }
101+ ));
102+ let _ = stream.set_read_timeout(None);
103+ return Ok(Grant {
104+ stream,
105+ lease,
106+ cores,
107+ });
108+ }
109+ Ok(Some(Event::Queued {
110+ lease,
111+ position,
112+ reason,
113+ })) => {
114+ let why = reason.map(|r| format!(", {r}")).unwrap_or_default();
115+ say(&format!("lease {lease} queued at #{position}{why}"));
116+ }
117+ Ok(Some(Event::Kill)) => {
118+ say("killed while queued");
119+ return Err(143);
120+ }
121+ Ok(Some(_)) => {}
122+ Err(e)
123+ if matches!(
124+ e.kind(),
125+ std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
126+ ) =>
127+ {
128+ if max_wait.is_some_and(|m| started.elapsed() > m.0) {
129+ say("gave up waiting");
130+ return Err(EXIT_GAVE_UP);
131+ }
132+ }
133+ Ok(None) | Err(_) => break,
134+ }
135+ }
136+ }
137+ say("daemon keeps hanging up");
138+ Err(1)
139+}
140+
141+fn has_systemd_run() -> bool {
142+ env::var_os("PATH")
143+ .is_some_and(|p| env::split_paths(&p).any(|d| d.join("systemd-run").exists()))
144+}
145+
146+fn command(cfg: &Config, a: &Args, need: &Need, grant: &Grant) -> Command {
147+ let mut argv = a.cmd.clone();
148+ if cfg.overrun.ram == RamOverrun::Throttle && cfg!(target_os = "linux") && has_systemd_run() {
149+ let high = Bytes((need.ram.0 as f64 * cfg.overrun.ram_slack) as u64);
150+ let mut wrapped: Vec<String> = [
151+ "systemd-run",
152+ "--user",
153+ "--scope",
154+ "--quiet",
155+ "--collect",
156+ "-p",
157+ ]
158+ .map(Into::into)
159+ .into();
160+ wrapped.push(format!("MemoryHigh={}", high.0));
161+ wrapped.push("--".into());
162+ wrapped.append(&mut argv);
163+ argv = wrapped;
164+ }
165+
166+ let n = grant.cores.to_string();
167+ let mut cmd = Command::new(&argv[0]);
168+ cmd.args(&argv[1..])
169+ .process_group(0)
170+ .env("ROUNDABOUT_LEASE", grant.lease.to_string())
171+ .env("CARGO_BUILD_JOBS", &n)
172+ .env("CMAKE_BUILD_PARALLEL_LEVEL", &n);
173+
174+ let make = env::var("MAKEFLAGS").unwrap_or_default();
175+ if !make.contains("-j") {
176+ cmd.env("MAKEFLAGS", format!("-j{n} {make}").trim_end());
177+ }
178+
179+ if cfg.nix.inject {
180+ let per = cfg.nix.cores_per_job.max(1);
181+ let jobs = (grant.cores / per).max(1);
182+ let cores = grant.cores.min(per);
183+ let mut nix = env::var("NIX_CONFIG").unwrap_or_default();
184+ let _ = write!(nix, "\ncores = {cores}\nmax-jobs = {jobs}\n");
185+ cmd.env("NIX_CONFIG", nix);
186+ }
187+
188+ #[cfg(target_os = "linux")]
189+ unsafe {
190+ cmd.pre_exec(|| {
191+ libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM);
192+ Ok(())
193+ });
194+ }
195+ cmd
196+}
197+
198+fn cpu_time() -> Duration {
199+ let mut ru = std::mem::MaybeUninit::<libc::rusage>::uninit();
200+ unsafe {
201+ if libc::getrusage(libc::RUSAGE_CHILDREN, ru.as_mut_ptr()) != 0 {
202+ return Duration::ZERO;
203+ }
204+ let ru = ru.assume_init();
205+ let tv = |t: libc::timeval| {
206+ Duration::from_secs(t.tv_sec as u64) + Duration::from_micros(t.tv_usec as u64)
207+ };
208+ tv(ru.ru_utime) + tv(ru.ru_stime)
209+ }
210+}
211+
212+fn code(status: ExitStatus) -> i32 {
213+ use std::os::unix::process::ExitStatusExt;
214+ status
215+ .code()
216+ .or_else(|| status.signal().map(|s| 128 + s))
217+ .unwrap_or(1)
218+}
219+
220+pub fn run(cfg: &Config, a: &Args) -> i32 {
221+ if a.cmd.is_empty() {
222+ say("nothing to run, put the command after --");
223+ return 2;
224+ }
225+ if a.size.is_none() && a.ram.is_none() {
226+ let sizes: Vec<_> = cfg.sizes.keys().map(String::as_str).collect();
227+ say(&format!(
228+ "how big is it? pass --size {} or --ram and friends. `rabt suggest -- <cmd>` has a guess",
229+ sizes.join("|")
230+ ));
231+ return 2;
232+ }
233+ let need = match need(cfg, a) {
234+ Ok(n) => n,
235+ Err(e) => {
236+ say(&e);
237+ return 2;
238+ }
239+ };
240+ let cwd = env::current_dir().unwrap_or_else(|_| "/".into());
241+ let root = paths::project_root(&cwd);
242+
243+ let report = artifacts::check(cfg, Some(&root), false);
244+ if report.over() {
245+ let msg = report.describe(cfg);
246+ if cfg.artifacts.over_cap == OverCap::Block {
247+ say(&format!("{msg}. clean up first, `rabt artifacts gc --cap`"));
248+ return EXIT_GAVE_UP;
249+ }
250+ say(&format!(
251+ "{msg}. `rabt artifacts gc --cap` when you get a sec"
252+ ));
253+ }
254+
255+ let req = Acquire {
256+ need,
257+ cwd,
258+ cmd: a.cmd.clone(),
259+ session: a.session.clone(),
260+ deadline: a.deadline,
261+ };
262+
263+ let grant = match acquire(&req, a.max_wait.or(cfg.queue.max_wait)) {
264+ Ok(g) => g,
265+ Err(code) => return code,
266+ };
267+
268+ let mut child = match command(cfg, a, &need, &grant).spawn() {
269+ Ok(c) => c,
270+ Err(e) => {
271+ say(&format!("couldn't start {}: {e}", a.cmd[0]));
272+ return 127;
273+ }
274+ };
275+ let pgid = i32::try_from(child.id()).expect("pid fits in pid_t");
276+ PGID.store(pgid, Ordering::Relaxed);
277+ for sig in [libc::SIGINT, libc::SIGTERM, libc::SIGHUP] {
278+ unsafe { libc::signal(sig, forward as *const () as libc::sighandler_t) };
279+ }
280+
281+ let mut out = grant.stream.try_clone().ok();
282+ if let Some(s) = &mut out {
283+ let _ = proto::send(s, &Request::Started { pid: pgid });
284+ }
285+
286+ let killed = Arc::new(AtomicBool::new(false));
287+ if let Ok(mut reader) = Reader::new(&grant.stream) {
288+ let killed = killed.clone();
289+ thread::spawn(move || {
290+ while let Ok(Some(ev)) = reader.recv::<Event>() {
291+ if matches!(ev, Event::Kill) {
292+ killed.store(true, Ordering::Relaxed);
293+ }
294+ }
295+ });
296+ }
297+
298+ let started = Instant::now();
299+ let grace = need.time.mul_f64(cfg.overrun.time_grace);
300+ let loud_ram = Bytes((need.ram.0 as f64 * cfg.overrun.ram_slack) as u64);
301+ let mut peak = Bytes::ZERO;
302+ let mut last_sample: Option<Instant> = None;
303+ let mut term_sent: Option<Instant> = None;
304+ let mut timed_out = false;
305+ let mut warned_ram = false;
306+ let mut warned_time = false;
307+
308+ let status = loop {
309+ match child.try_wait() {
310+ Ok(Some(status)) => break status,
311+ Ok(None) => {}
312+ Err(e) => {
313+ say(&format!("lost track of the child: {e}"));
314+ return 1;
315+ }
316+ }
317+ let elapsed = started.elapsed();
318+
319+ if term_sent.is_none() {
320+ if killed.load(Ordering::Relaxed) {
321+ say("killed on request");
322+ term_sent = Some(Instant::now());
323+ } else if a.deadline.is_some_and(|d| elapsed > d.0) {
324+ say("deadline hit, killing it");
325+ timed_out = true;
326+ term_sent = Some(Instant::now());
327+ }
328+ if term_sent.is_some() {
329+ unsafe { libc::kill(-pgid, libc::SIGTERM) };
330+ }
331+ } else if term_sent.is_some_and(|t| t.elapsed() > Duration::from_secs(10)) {
332+ unsafe { libc::kill(-pgid, libc::SIGKILL) };
333+ }
334+
335+ if last_sample.is_none_or(|t| t.elapsed() >= Duration::from_secs(2)) {
336+ last_sample = Some(Instant::now());
337+ let rss = sys::group_rss(pgid);
338+ peak = peak.max(rss);
339+ if let Some(s) = &mut out {
340+ let _ = proto::send(s, &Request::Usage { rss });
341+ }
342+ if !warned_ram && rss > loud_ram {
343+ warned_ram = true;
344+ say(&format!("using {rss} against a {} estimate", need.ram));
345+ }
346+ }
347+ if !warned_time && elapsed > grace {
348+ warned_time = true;
349+ say(&format!(
350+ "{} past its {} estimate and still going",
351+ Dur(elapsed),
352+ Dur(need.time)
353+ ));
354+ }
355+ thread::sleep(Duration::from_millis(250));
356+ };
357+
358+ let wall = started.elapsed();
359+ if let Some(s) = &mut out {
360+ let _ = proto::send(s, &Request::Done);
361+ }
362+ history::record(
363+ history::key(&root, &a.cmd),
364+ peak,
365+ wall,
366+ cpu_time(),
367+ grant.cores,
368+ );
369+ artifacts::register(cfg, &root, true);
370+ if peak > loud_ram || wall > grace {
371+ say(&format!(
372+ "took {} and peaked at {}, asked for {} and {}. ask for more next time",
373+ Dur(wall),
374+ peak,
375+ Dur(need.time),
376+ need.ram
377+ ));
378+ }
379+ if timed_out {
380+ EXIT_TIMEOUT
381+ } else {
382+ code(status)
383+ }
384+}
Added file: src/sched.rs+353 additions
@@ -0,0 +1,353 @@
1+//! fifo with easy backfill. the first job that can't start gets a reservation,
2+//! and anyone behind it only cuts in if they'll be gone before it starts or fit
3+//! in the leftovers once it does
4+
5+use std::{collections::BTreeMap, time::Duration};
6+
7+use serde::{Deserialize, Serialize};
8+
9+use crate::units::Bytes;
10+
11+pub type JobId = u64;
12+pub type FsId = u64;
13+
14+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
15+pub struct Need {
16+ pub ram: Bytes,
17+ pub disk: Bytes,
18+ pub cores: u32,
19+ pub time: Duration,
20+}
21+
22+#[derive(Clone, Debug)]
23+pub struct Queued {
24+ pub id: JobId,
25+ pub need: Need,
26+ pub heavy: bool,
27+ pub fs: Vec<FsId>,
28+}
29+
30+#[derive(Clone, Debug)]
31+pub struct Running {
32+ pub need: Need,
33+ pub heavy: bool,
34+ pub fs: Vec<FsId>,
35+ pub rss: Bytes,
36+ /// zero once overdue
37+ pub remaining: Duration,
38+}
39+
40+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41+pub struct Disk {
42+ pub free: Bytes,
43+ pub reserve: Bytes,
44+ /// written since the fs last went from idle to busy. free space already
45+ /// reflects this, so it's netted off promises to avoid counting it twice
46+ pub consumed: Bytes,
47+}
48+
49+#[derive(Clone, Debug)]
50+pub struct Capacity {
51+ pub ram_budget: Bytes,
52+ pub ram_live: Bytes,
53+ pub ram_reserve: Bytes,
54+ pub cores: u32,
55+ pub disks: BTreeMap<FsId, Disk>,
56+ pub pressure: bool,
57+}
58+
59+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
60+#[serde(rename_all = "snake_case", tag = "reason", content = "fs")]
61+pub enum Reason {
62+ HeavyRunning,
63+ Pressure,
64+ Ram,
65+ Cores,
66+ Disk(FsId),
67+ Reserved,
68+}
69+
70+impl std::fmt::Display for Reason {
71+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72+ f.write_str(match self {
73+ Self::HeavyRunning => "another heavy build is running",
74+ Self::Pressure => "memory pressure is high",
75+ Self::Ram => "not enough free RAM",
76+ Self::Cores => "not enough free cores",
77+ Self::Disk(_) => "not enough free disk",
78+ Self::Reserved => "waiting behind a reserved build",
79+ })
80+ }
81+}
82+
83+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
84+pub enum Verdict {
85+ Admit,
86+ Wait(Reason),
87+}
88+
89+#[derive(Clone, Debug)]
90+struct Free {
91+ ram: Bytes,
92+ cores: u32,
93+ disk: BTreeMap<FsId, Bytes>,
94+ heavy_slot: bool,
95+}
96+
97+impl Free {
98+ fn shortfall(&self, need: &Need, fs: &[FsId], heavy: bool) -> Option<Reason> {
99+ if heavy && !self.heavy_slot {
100+ return Some(Reason::HeavyRunning);
101+ }
102+ if need.ram > self.ram {
103+ return Some(Reason::Ram);
104+ }
105+ if need.cores > self.cores {
106+ return Some(Reason::Cores);
107+ }
108+ fs.iter()
109+ .find(|id| self.disk.get(id).is_some_and(|free| need.disk > *free))
110+ .map(|id| Reason::Disk(*id))
111+ }
112+
113+ fn take(&mut self, need: &Need, fs: &[FsId], heavy: bool) {
114+ self.ram = self.ram.saturating_sub(need.ram);
115+ self.cores = self.cores.saturating_sub(need.cores);
116+ for id in fs {
117+ if let Some(free) = self.disk.get_mut(id) {
118+ *free = free.saturating_sub(need.disk);
119+ }
120+ }
121+ if heavy {
122+ self.heavy_slot = false;
123+ }
124+ }
125+
126+ /// no disk here. builds leave their output behind when they finish
127+ fn give_back(&mut self, ram: Bytes, cores: u32, heavy: bool) {
128+ self.ram += ram;
129+ self.cores += cores;
130+ if heavy {
131+ self.heavy_slot = true;
132+ }
133+ }
134+}
135+
136+fn clamp(need: Need, cap: &Capacity) -> Need {
137+ Need {
138+ ram: need.ram.min(cap.ram_budget),
139+ cores: need.cores.clamp(1, cap.cores.max(1)),
140+ ..need
141+ }
142+}
143+
144+fn free_now(running: &[Running], cap: &Capacity) -> Free {
145+ let charged: Bytes = running.iter().map(|r| r.need.ram.max(r.rss)).sum();
146+ let unrealised: Bytes = running
147+ .iter()
148+ .map(|r| r.need.ram.saturating_sub(r.rss))
149+ .sum();
150+ let by_budget = cap.ram_budget.saturating_sub(charged);
151+ let by_live = cap
152+ .ram_live
153+ .saturating_sub(cap.ram_reserve)
154+ .saturating_sub(unrealised);
155+
156+ let disk = cap
157+ .disks
158+ .iter()
159+ .map(|(id, d)| {
160+ let promised: Bytes = running
161+ .iter()
162+ .filter(|r| r.fs.contains(id))
163+ .map(|r| r.need.disk)
164+ .sum();
165+ let outstanding = promised.saturating_sub(d.consumed);
166+ (
167+ *id,
168+ d.free.saturating_sub(d.reserve).saturating_sub(outstanding),
169+ )
170+ })
171+ .collect();
172+
173+ Free {
174+ ram: by_budget.min(by_live),
175+ cores: cap
176+ .cores
177+ .saturating_sub(running.iter().map(|r| r.need.cores).sum()),
178+ disk,
179+ heavy_slot: !running.iter().any(|r| r.heavy),
180+ }
181+}
182+
183+pub fn headroom(running: &[Running], cap: &Capacity) -> (Bytes, u32) {
184+ let free = free_now(running, cap);
185+ (free.ram, free.cores)
186+}
187+
188+struct Reservation {
189+ /// `None` when nothing finishing would free enough, e.g. disk or someone
190+ /// else's ram. backfill falls back to the time cap then
191+ shadow: Option<Duration>,
192+ extra: Option<Free>,
193+}
194+
195+fn reserve(
196+ job: &Queued,
197+ need: &Need,
198+ free: &Free,
199+ ends: &[(Duration, Bytes, u32, bool)],
200+) -> Reservation {
201+ let mut sim = free.clone();
202+ let mut ends = ends.to_vec();
203+ ends.sort_by_key(|e| e.0);
204+ for (at, ram, cores, heavy) in ends {
205+ sim.give_back(ram, cores, heavy);
206+ if sim.shortfall(need, &job.fs, job.heavy).is_none() {
207+ sim.take(need, &job.fs, job.heavy);
208+ return Reservation {
209+ shadow: Some(at),
210+ extra: Some(sim),
211+ };
212+ }
213+ }
214+ Reservation {
215+ shadow: None,
216+ extra: None,
217+ }
218+}
219+
220+pub fn plan(
221+ queue: &[Queued],
222+ running: &[Running],
223+ cap: &Capacity,
224+ backfill_max: Duration,
225+) -> Vec<(JobId, Verdict)> {
226+ let mut free = free_now(running, cap);
227+ let mut ends: Vec<_> = running
228+ .iter()
229+ .map(|r| (r.remaining, r.need.ram.max(r.rss), r.need.cores, r.heavy))
230+ .collect();
231+ let mut reservation: Option<Reservation> = None;
232+ let mut out = Vec::with_capacity(queue.len());
233+
234+ for job in queue {
235+ let need = clamp(job.need, cap);
236+
237+ if job.heavy && cap.pressure {
238+ out.push((job.id, Verdict::Wait(Reason::Pressure)));
239+ continue;
240+ }
241+
242+ if let Some(why) = free.shortfall(&need, &job.fs, job.heavy) {
243+ if reservation.is_none() {
244+ reservation = Some(reserve(job, &need, &free, &ends));
245+ }
246+ out.push((job.id, Verdict::Wait(why)));
247+ continue;
248+ }
249+
250+ if let Some(res) = &mut reservation {
251+ let finishes_first = match res.shadow {
252+ Some(shadow) => need.time <= shadow,
253+ None => need.time <= backfill_max,
254+ };
255+ let fits_extra = res
256+ .extra
257+ .as_ref()
258+ .is_some_and(|x| x.shortfall(&need, &job.fs, job.heavy).is_none());
259+ if !finishes_first && !fits_extra {
260+ out.push((job.id, Verdict::Wait(Reason::Reserved)));
261+ continue;
262+ }
263+ if !finishes_first && let Some(extra) = &mut res.extra {
264+ extra.take(&need, &job.fs, job.heavy);
265+ }
266+ }
267+
268+ free.take(&need, &job.fs, job.heavy);
269+ ends.push((need.time, need.ram, need.cores, job.heavy));
270+ out.push((job.id, Verdict::Admit));
271+ }
272+ out
273+}
274+
275+#[cfg(test)]
276+mod tests {
277+ use super::*;
278+
279+ const FS: FsId = 1;
280+
281+ fn need(ram_gib: u64, cores: u32, disk_gib: u64, mins: u64) -> Need {
282+ Need {
283+ ram: Bytes::gib(ram_gib),
284+ disk: Bytes::gib(disk_gib),
285+ cores,
286+ time: Duration::from_secs(mins * 60),
287+ }
288+ }
289+
290+ fn q(id: JobId, need: Need, heavy: bool) -> Queued {
291+ Queued {
292+ id,
293+ need,
294+ heavy,
295+ fs: vec![FS],
296+ }
297+ }
298+
299+ fn cap() -> Capacity {
300+ Capacity {
301+ ram_budget: Bytes::gib(24),
302+ ram_live: Bytes::gib(28),
303+ ram_reserve: Bytes::gib(4),
304+ cores: 16,
305+ disks: BTreeMap::from([(
306+ FS,
307+ Disk {
308+ free: Bytes::gib(100),
309+ reserve: Bytes::gib(20),
310+ consumed: Bytes::ZERO,
311+ },
312+ )]),
313+ pressure: false,
314+ }
315+ }
316+
317+ fn verdicts(plan: &[(JobId, Verdict)]) -> Vec<Verdict> {
318+ plan.iter().map(|(_, v)| *v).collect()
319+ }
320+
321+ const BF: Duration = Duration::from_secs(120);
322+
323+ #[test]
324+ fn two_heavies_never_collide() {
325+ let out = plan(
326+ &[q(1, need(8, 8, 5, 10), true), q(2, need(8, 4, 5, 10), true)],
327+ &[],
328+ &cap(),
329+ BF,
330+ );
331+ assert_eq!(
332+ verdicts(&out),
333+ [Verdict::Admit, Verdict::Wait(Reason::HeavyRunning)]
334+ );
335+ }
336+
337+ #[test]
338+ fn disk_never_double_booked() {
339+ let out = plan(
340+ &[
341+ q(1, need(1, 1, 50, 10), false),
342+ q(2, need(1, 1, 50, 10), false),
343+ ],
344+ &[],
345+ &cap(),
346+ BF,
347+ );
348+ assert_eq!(
349+ verdicts(&out),
350+ [Verdict::Admit, Verdict::Wait(Reason::Disk(FS))]
351+ );
352+ }
353+}
Added file: src/store.rs+35 additions
@@ -0,0 +1,35 @@
1+use std::{fs, io, path::PathBuf};
2+
3+use serde::{Serialize, de::DeserializeOwned};
4+
5+use crate::{client::Flock, paths};
6+
7+fn path(name: &str) -> PathBuf {
8+ paths::state_dir().join(format!("{name}.json"))
9+}
10+
11+pub fn load<T: DeserializeOwned + Default>(name: &str) -> T {
12+ fs::read(path(name))
13+ .ok()
14+ .and_then(|b| serde_json::from_slice(&b).ok())
15+ .unwrap_or_default()
16+}
17+
18+/// readers skip the lock, the rename keeps them from seeing half a file
19+pub fn update<T, R>(name: &str, f: impl FnOnce(&mut T) -> R) -> io::Result<R>
20+where
21+ T: Serialize + DeserializeOwned + Default,
22+{
23+ let lock = fs::OpenOptions::new()
24+ .create(true)
25+ .truncate(false)
26+ .write(true)
27+ .open(paths::state_dir().join(format!("{name}.lock")))?;
28+ let _guard = Flock::exclusive(lock)?;
29+ let mut value = load(name);
30+ let out = f(&mut value);
31+ let tmp = path(&format!("{name}.tmp"));
32+ fs::write(&tmp, serde_json::to_vec(&value)?)?;
33+ fs::rename(tmp, path(name))?;
34+ Ok(out)
35+}
Added file: src/sys.rs+163 additions
@@ -0,0 +1,163 @@
1+use std::{ffi::CString, os::unix::ffi::OsStrExt, path::Path};
2+
3+use crate::{sched::FsId, units::Bytes};
4+
5+pub fn cores() -> u32 {
6+ std::thread::available_parallelism().map_or(1, |n| n.get() as u32)
7+}
8+
9+pub struct FsStat {
10+ pub id: FsId,
11+ pub free: Bytes,
12+ pub size: Bytes,
13+}
14+
15+pub fn fs(path: &Path) -> Option<FsStat> {
16+ let c = CString::new(path.as_os_str().as_bytes()).ok()?;
17+ let mut vfs = std::mem::MaybeUninit::<libc::statvfs>::uninit();
18+ let mut st = std::mem::MaybeUninit::<libc::stat>::uninit();
19+ unsafe {
20+ if libc::statvfs(c.as_ptr(), vfs.as_mut_ptr()) != 0
21+ || libc::stat(c.as_ptr(), st.as_mut_ptr()) != 0
22+ {
23+ return None;
24+ }
25+ let vfs = vfs.assume_init();
26+ let st = st.assume_init();
27+ #[allow(clippy::unnecessary_cast)]
28+ let block = vfs.f_frsize as u64;
29+ #[allow(clippy::unnecessary_cast)]
30+ Some(FsStat {
31+ id: st.st_dev as u64,
32+ free: Bytes(vfs.f_bavail as u64 * block),
33+ size: Bytes(vfs.f_blocks as u64 * block),
34+ })
35+ }
36+}
37+
38+pub fn page_size() -> u64 {
39+ unsafe { libc::sysconf(libc::_SC_PAGESIZE) as u64 }
40+}
41+
42+#[cfg(target_os = "linux")]
43+mod imp {
44+ use std::fs;
45+
46+ use crate::units::Bytes;
47+
48+ pub fn memory() -> (Bytes, Bytes) {
49+ let text = fs::read_to_string("/proc/meminfo").unwrap_or_default();
50+ let field = |name: &str| {
51+ text.lines()
52+ .find_map(|l| l.strip_prefix(name))
53+ .and_then(|rest| {
54+ rest.trim()
55+ .trim_end_matches("kB")
56+ .trim()
57+ .parse::<u64>()
58+ .ok()
59+ })
60+ .map_or(Bytes::ZERO, |kib| Bytes(kib * 1024))
61+ };
62+ (field("MemTotal:"), field("MemAvailable:"))
63+ }
64+
65+ pub fn pressure() -> Option<f64> {
66+ let text = fs::read_to_string("/proc/pressure/memory").ok()?;
67+ text.lines()
68+ .find(|l| l.starts_with("some"))?
69+ .split_whitespace()
70+ .find_map(|kv| kv.strip_prefix("avg10="))?
71+ .parse()
72+ .ok()
73+ }
74+
75+ pub fn group_rss(pgid: i32) -> Bytes {
76+ let page = super::page_size();
77+ let Ok(dir) = fs::read_dir("/proc") else {
78+ return Bytes::ZERO;
79+ };
80+ let mut total = 0;
81+ for entry in dir.flatten() {
82+ let name = entry.file_name();
83+ if !name.as_encoded_bytes().iter().all(u8::is_ascii_digit) {
84+ continue;
85+ }
86+ let Ok(stat) = fs::read_to_string(entry.path().join("stat")) else {
87+ continue;
88+ };
89+ // comm can hold spaces and parens, fields only make sense after the last ')'
90+ let Some(tail) = stat.rsplit_once(')').map(|(_, t)| t) else {
91+ continue;
92+ };
93+ let fields: Vec<&str> = tail.split_whitespace().collect();
94+ let pgrp = fields.get(2).and_then(|f| f.parse::<i32>().ok());
95+ let rss = fields.get(21).and_then(|f| f.parse::<u64>().ok());
96+ if pgrp == Some(pgid) {
97+ total += rss.unwrap_or(0) * page;
98+ }
99+ }
100+ Bytes(total)
101+ }
102+}
103+
104+#[cfg(target_os = "macos")]
105+mod imp {
106+ use std::process::Command;
107+
108+ use crate::units::Bytes;
109+
110+ pub fn memory() -> (Bytes, Bytes) {
111+ let mut total: u64 = 0;
112+ let mut len = std::mem::size_of::<u64>();
113+ unsafe {
114+ libc::sysctlbyname(
115+ c"hw.memsize".as_ptr(),
116+ (&raw mut total).cast(),
117+ &raw mut len,
118+ std::ptr::null_mut(),
119+ 0,
120+ );
121+ }
122+ // vm_stat over mach calls, libc's mach bindings are deprecated and this
123+ // runs once a second at most
124+ let out = Command::new("vm_stat")
125+ .output()
126+ .map(|o| o.stdout)
127+ .unwrap_or_default();
128+ let text = String::from_utf8_lossy(&out);
129+ let pages = |name: &str| {
130+ text.lines()
131+ .find_map(|l| l.strip_prefix(name))
132+ .and_then(|r| r.trim().trim_end_matches('.').parse::<u64>().ok())
133+ .unwrap_or(0)
134+ };
135+ let avail = (pages("Pages free:") + pages("Pages inactive:") + pages("Pages purgeable:"))
136+ * super::page_size();
137+ (Bytes(total), Bytes(avail))
138+ }
139+
140+ pub const fn pressure() -> Option<f64> {
141+ None
142+ }
143+
144+ pub fn group_rss(pgid: i32) -> Bytes {
145+ let out = Command::new("ps")
146+ .args(["-A", "-o", "pgid=,rss="])
147+ .output()
148+ .map(|o| o.stdout)
149+ .unwrap_or_default();
150+ let kib: u64 = String::from_utf8_lossy(&out)
151+ .lines()
152+ .filter_map(|l| {
153+ let mut f = l.split_whitespace();
154+ let g = f.next()?.parse::<i32>().ok()?;
155+ let rss = f.next()?.parse::<u64>().ok()?;
156+ (g == pgid).then_some(rss)
157+ })
158+ .sum();
159+ Bytes(kib * 1024)
160+ }
161+}
162+
163+pub use imp::{group_rss, memory, pressure};
Added file: src/units.rs+301 additions
@@ -0,0 +1,301 @@
1+use std::{fmt, str::FromStr, time::Duration};
2+
3+use serde::{Deserialize, Deserializer, Serialize, Serializer, de};
4+
5+const KIB: u64 = 1024;
6+
7+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
8+pub struct Bytes(pub u64);
9+
10+impl Bytes {
11+ pub const ZERO: Self = Self(0);
12+
13+ pub const fn gib(n: u64) -> Self {
14+ Self(n * KIB * KIB * KIB)
15+ }
16+
17+ pub const fn mib(n: u64) -> Self {
18+ Self(n * KIB * KIB)
19+ }
20+
21+ pub const fn saturating_sub(self, other: Self) -> Self {
22+ Self(self.0.saturating_sub(other.0))
23+ }
24+}
25+
26+impl std::ops::Add for Bytes {
27+ type Output = Self;
28+
29+ fn add(self, other: Self) -> Self {
30+ Self(self.0.saturating_add(other.0))
31+ }
32+}
33+
34+impl std::ops::AddAssign for Bytes {
35+ fn add_assign(&mut self, other: Self) {
36+ *self = *self + other;
37+ }
38+}
39+
40+impl std::iter::Sum for Bytes {
41+ fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
42+ iter.fold(Self::ZERO, |a, b| a + b)
43+ }
44+}
45+
46+fn split_number(s: &str) -> Result<(f64, &str), &'static str> {
47+ let s = s.trim();
48+ let end = s
49+ .find(|c: char| !(c.is_ascii_digit() || c == '.'))
50+ .unwrap_or(s.len());
51+ let (num, suffix) = s.split_at(end);
52+ let num = num.parse::<f64>().map_err(|_| "expected a number")?;
53+ if !num.is_finite() || num < 0.0 {
54+ return Err("expected a non-negative number");
55+ }
56+ Ok((num, suffix.trim()))
57+}
58+
59+impl FromStr for Bytes {
60+ type Err = &'static str;
61+
62+ fn from_str(s: &str) -> Result<Self, Self::Err> {
63+ let (num, suffix) = split_number(s)?;
64+ let scale = match suffix.to_ascii_lowercase().as_str() {
65+ "" | "b" => 1,
66+ "k" | "kb" | "kib" => KIB,
67+ "m" | "mb" | "mib" => KIB.pow(2),
68+ "g" | "gb" | "gib" => KIB.pow(3),
69+ "t" | "tb" | "tib" => KIB.pow(4),
70+ _ => return Err("expected a size like 512M or 4GiB"),
71+ };
72+ Ok(Self((num * scale as f64) as u64))
73+ }
74+}
75+
76+impl fmt::Display for Bytes {
77+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78+ const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
79+ let mut value = self.0 as f64;
80+ let mut unit = 0;
81+ while value >= 1024.0 && unit < UNITS.len() - 1 {
82+ value /= 1024.0;
83+ unit += 1;
84+ }
85+ if unit == 0 {
86+ write!(f, "{}B", self.0)
87+ } else {
88+ write!(f, "{value:.1}{}", UNITS[unit])
89+ }
90+ }
91+}
92+
93+impl pound::FromArg for Bytes {
94+ fn from_arg(s: &str) -> Result<Self, pound::ValueError> {
95+ s.parse().map_err(|e| pound::ValueError::new(s, e))
96+ }
97+}
98+
99+impl Serialize for Bytes {
100+ fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
101+ s.serialize_u64(self.0)
102+ }
103+}
104+
105+impl<'de> Deserialize<'de> for Bytes {
106+ fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
107+ struct V;
108+ impl de::Visitor<'_> for V {
109+ type Value = Bytes;
110+
111+ fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112+ f.write_str("a byte count or a size string like \"4GiB\"")
113+ }
114+
115+ fn visit_u64<E: de::Error>(self, v: u64) -> Result<Bytes, E> {
116+ Ok(Bytes(v))
117+ }
118+
119+ fn visit_i64<E: de::Error>(self, v: i64) -> Result<Bytes, E> {
120+ u64::try_from(v)
121+ .map(Bytes)
122+ .map_err(|_| E::custom("negative size"))
123+ }
124+
125+ fn visit_str<E: de::Error>(self, v: &str) -> Result<Bytes, E> {
126+ v.parse().map_err(E::custom)
127+ }
128+ }
129+ d.deserialize_any(V)
130+ }
131+}
132+
133+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
134+pub struct Dur(pub Duration);
135+
136+impl Dur {
137+ pub const fn secs(n: u64) -> Self {
138+ Self(Duration::from_secs(n))
139+ }
140+}
141+
142+impl FromStr for Dur {
143+ type Err = &'static str;
144+
145+ fn from_str(s: &str) -> Result<Self, Self::Err> {
146+ let mut rest = s.trim();
147+ if rest.is_empty() {
148+ return Err("expected a duration like 90s or 5m");
149+ }
150+ let mut total = 0.0;
151+ while !rest.is_empty() {
152+ let (num, tail) = split_number(rest)?;
153+ let unit_len = tail
154+ .find(|c: char| c.is_ascii_digit() || c == '.')
155+ .unwrap_or(tail.len());
156+ let (unit, next) = tail.split_at(unit_len);
157+ let scale = match unit.trim() {
158+ "" | "s" | "sec" | "secs" => 1.0,
159+ "ms" => 0.001,
160+ "m" | "min" | "mins" => 60.0,
161+ "h" | "hr" | "hrs" => 3600.0,
162+ "d" | "day" | "days" => 86400.0,
163+ _ => return Err("expected a duration like 90s or 5m"),
164+ };
165+ total = num.mul_add(scale, total);
166+ rest = next.trim();
167+ }
168+ Ok(Self(Duration::from_secs_f64(total)))
169+ }
170+}
171+
172+impl fmt::Display for Dur {
173+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174+ let s = self.0.as_secs();
175+ match s {
176+ 0..60 => write!(f, "{s}s"),
177+ 60..3600 if s.is_multiple_of(60) => write!(f, "{}m", s / 60),
178+ 60..3600 => write!(f, "{}m{}s", s / 60, s % 60),
179+ _ if s.is_multiple_of(3600) => write!(f, "{}h", s / 3600),
180+ _ => write!(f, "{}h{}m", s / 3600, (s % 3600) / 60),
181+ }
182+ }
183+}
184+
185+impl pound::FromArg for Dur {
186+ fn from_arg(s: &str) -> Result<Self, pound::ValueError> {
187+ s.parse().map_err(|e| pound::ValueError::new(s, e))
188+ }
189+}
190+
191+impl Serialize for Dur {
192+ fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
193+ s.serialize_u64(self.0.as_secs())
194+ }
195+}
196+
197+impl<'de> Deserialize<'de> for Dur {
198+ fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
199+ struct V;
200+ impl de::Visitor<'_> for V {
201+ type Value = Dur;
202+
203+ fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
204+ f.write_str("seconds or a duration string like \"5m\"")
205+ }
206+
207+ fn visit_u64<E: de::Error>(self, v: u64) -> Result<Dur, E> {
208+ Ok(Dur::secs(v))
209+ }
210+
211+ fn visit_i64<E: de::Error>(self, v: i64) -> Result<Dur, E> {
212+ u64::try_from(v)
213+ .map(Dur::secs)
214+ .map_err(|_| E::custom("negative duration"))
215+ }
216+
217+ fn visit_str<E: de::Error>(self, v: &str) -> Result<Dur, E> {
218+ v.parse().map_err(E::custom)
219+ }
220+ }
221+ d.deserialize_any(V)
222+ }
223+}
224+
225+#[derive(Clone, Copy, Debug, PartialEq)]
226+pub enum Amount<T> {
227+ Abs(T),
228+ Percent(f64),
229+}
230+
231+impl Amount<Bytes> {
232+ pub fn resolve(self, total: Bytes) -> Bytes {
233+ match self {
234+ Self::Abs(b) => b,
235+ Self::Percent(p) => Bytes((total.0 as f64 * p / 100.0) as u64),
236+ }
237+ }
238+}
239+
240+impl Amount<u32> {
241+ pub fn resolve(self, total: u32) -> u32 {
242+ match self {
243+ Self::Abs(n) => n,
244+ Self::Percent(p) => ((f64::from(total) * p / 100.0).round() as u32).max(1),
245+ }
246+ }
247+}
248+
249+impl<T: FromStr> FromStr for Amount<T> {
250+ type Err = &'static str;
251+
252+ fn from_str(s: &str) -> Result<Self, Self::Err> {
253+ if let Some(p) = s.trim().strip_suffix('%') {
254+ let p = p
255+ .trim()
256+ .parse::<f64>()
257+ .map_err(|_| "expected a percentage")?;
258+ if !(0.0..=100.0).contains(&p) {
259+ return Err("percentage must be between 0 and 100");
260+ }
261+ Ok(Self::Percent(p))
262+ } else {
263+ s.trim()
264+ .parse()
265+ .map(Self::Abs)
266+ .map_err(|_| "expected an amount or a percentage")
267+ }
268+ }
269+}
270+
271+impl<T: FromStr> pound::FromArg for Amount<T> {
272+ fn from_arg(s: &str) -> Result<Self, pound::ValueError> {
273+ s.parse().map_err(|e| pound::ValueError::new(s, e))
274+ }
275+}
276+
277+impl<'de, T: FromStr> Deserialize<'de> for Amount<T> {
278+ fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
279+ struct V<T>(std::marker::PhantomData<T>);
280+ impl<T: FromStr> de::Visitor<'_> for V<T> {
281+ type Value = Amount<T>;
282+
283+ fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
284+ f.write_str("an amount or a percentage string like \"75%\"")
285+ }
286+
287+ fn visit_u64<E: de::Error>(self, v: u64) -> Result<Amount<T>, E> {
288+ self.visit_str(&v.to_string())
289+ }
290+
291+ fn visit_i64<E: de::Error>(self, v: i64) -> Result<Amount<T>, E> {
292+ self.visit_str(&v.to_string())
293+ }
294+
295+ fn visit_str<E: de::Error>(self, v: &str) -> Result<Amount<T>, E> {
296+ v.parse().map_err(E::custom)
297+ }
298+ }
299+ d.deserialize_any(V(std::marker::PhantomData))
300+ }
301+}