Commit 6f5b84fad3ef
tarball18 files changed+3772-0
@@ -0,0 +1,2 @@
1
+/target2
+/result*@@ -0,0 +1,198 @@
1
+# This file is automatically @generated by Cargo.2
+# It is not intended for manual editing.3
+version = 44
+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"@@ -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 = 135
+lto = "fat"36
+opt-level = "s"37
+panic = "abort"38
+strip = true@@ -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> = dir51
+ .flatten()52
+ .map(|e| e.file_name().to_string_lossy().into_owned())53
+ .collect();54
+ names55
+ .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 store66
+/// path. hardlinks count once, the store is full of them67
+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 them89
+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))) => st141
+ .leases142
+ .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_global178
+ }179
+180
+ pub fn describe(&self, cfg: &Config) -> String {181
+ let mut out = Vec::new();182
+ if let Some((root, size)) = &self.project183
+ && self.over_project184
+ {185
+ out.push(format!(186
+ "{} holds {size} of build artifacts, cap is {}",187
+ root.display(),188
+ cfg.artifacts.project_cap189
+ ));190
+ }191
+ if self.over_global {192
+ out.push(format!(193
+ "artifacts total {} across projects, cap is {}",194
+ self.global, cfg.artifacts.global_cap195
+ ));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 = entries206
+ .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: project214
+ .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> = entries261
+ .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 = entries268
+ .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)> = entries279
+ .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
+ out318
+}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
+}@@ -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, please15
+#[derive(Parse)]16
+#[pound(name = "rabt", version = env!("CARGO_PKG_VERSION"))]17
+pub enum Rabt {18
+ /// wait your turn, then run the thing19
+ Run {20
+ /// preset from config, s m l xl out of the box21
+ #[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 machine28
+ #[pound(long)]29
+ cores: Option<Amount<u32>>,30
+ /// how long you reckon it'll take31
+ #[pound(short, long)]32
+ time: Option<Dur>,33
+ /// kill it past this, exit 12434
+ #[pound(long)]35
+ deadline: Option<Dur>,36
+ /// stop queueing past this, exit 7537
+ #[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 waiting45
+ Status {46
+ #[pound(long)]47
+ json: bool,48
+ },49
+ /// a size guess from past runs50
+ Suggest {51
+ #[pound(trailing)]52
+ cmd: Vec<String>,53
+ },54
+ /// stop a lease's build55
+ Kill { lease: u64 },56
+ /// buy an overdue build more time57
+ Extend { lease: u64, by: Dur },58
+ /// exit when a lease goes overdue or finishes59
+ Watch { lease: u64 },60
+ /// build output and how much of it there is61
+ Artifacts {62
+ #[pound(subcommand)]63
+ action: Artifacts,64
+ },65
+ /// claude code hook entrypoint, json on stdin66
+ Hook {67
+ /// pre-tool-use post-tool-use prompt-submit stop subagent-stop68
+ 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 over77
+ Check {78
+ /// re-measure instead of trusting the cache79
+ #[pound(long)]80
+ fresh: bool,81
+ },82
+ /// everything tracked, biggest first83
+ List,84
+ /// delete idle artifacts85
+ Gc {86
+ /// keep going, oldest first, until under the caps87
+ #[pound(long)]88
+ cap: bool,89
+ #[pound(short = 'n', long)]90
+ dry_run: bool,91
+ },92
+ /// never gc this93
+ 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_cap112
+ );113
+ }114
+ println!("everything: {} of {}", r.global, cfg.artifacts.global_cap);115
+ if r.over() {116
+ println!("over. `rabt artifacts gc --cap`");117
+ 1118
+ } else {119
+ 0120
+ }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
+ 0135
+ }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.freed152
+ );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
+ 1158
+ }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
+ 1169
+ },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
+ 1188
+ }189
+ Ok(None) => {190
+ eprintln!("rabt: nothing running");191
+ 1192
+ }193
+ Ok(Some(_)) => 1,194
+ Err(e) => {195
+ eprintln!("rabt: {e}");196
+ 1197
+ }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.cores217
+ );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
+ 0245
+}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.disk268
+ );269
+ 0270
+}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
+ 3285
+ } else {286
+ println!("lease {lease} is done");287
+ 0288
+ }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
+}@@ -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
+}@@ -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 its32
+ /// cwd. pessimistic when they share a fs, which is the point33
+ pub disk_paths: Vec<PathBuf>,34
+ /// memory psi `some avg10`. only gates heavies35
+ 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 estimate75
+ pub ram_slack: f64,76
+ /// multiple of the time estimate77
+ 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 get116
+ pub name: String,117
+ /// sibling that must exist, so we don't go deleting someone's `build/`118
+ /// because it had a suggestive name119
+ 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.artifacts368
+ .rules369
+ .iter()370
+ .chain(&self.artifacts.extra_rules)371
+ }372
+373
+ pub fn detect_patterns(&self) -> impl Iterator<Item = &str> {374
+ self.detect375
+ .commands376
+ .iter()377
+ .chain(&self.detect.extra)378
+ .map(String::as_str)379
+ }380
+}@@ -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.run43
+ .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
+ shared67
+ .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 reserve99
+ /// admits nothing ever100
+ 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
+ ids120
+ }121
+122
+ fn refresh(&mut self) {123
+ let (_, live) = sys::memory();124
+ self.cap.ram_live = live;125
+ self.cap.pressure = self126
+ .cfg127
+ .budget128
+ .psi_max129
+ .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> = self141
+ .leases142
+ .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 = disks148
+ .into_iter()149
+ .map(|(id, (free, reserve))| {150
+ let consumed = if busy.contains(&id) {151
+ self.baselines152
+ .entry(id)153
+ .or_insert(free)154
+ .saturating_sub(free)155
+ } else {156
+ Bytes::ZERO157
+ };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.leases172
+ .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> = self188
+ .leases189
+ .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.leases276
+ .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
+ id340
+ };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::Ok355
+ }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::Ok372
+ }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 = st411
+ .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, and428
+/// pdeathsig only reaches its direct child. take the rest of the group down429
+/// so a headless build doesn't keep eating ram off the books430
+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
+ .leases437
+ .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.0462
+ {463
+ // a spawner holds this while it clears a stale socket, so taking it464
+ // here stops us unlinking our successor's socket on the way out465
+ 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 now483
+ // instead of carrying a zombie for the length of its build484
+ 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
+}@@ -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 forever18
+ 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
+ k30
+}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 and74
+/// the base estimate stands75
+pub fn suggest(rec: &Record, base: Need, cmd: &[String]) -> Need {76
+ let via_daemon = cmd77
+ .iter()78
+ .any(|t| t.contains("nix") || t.contains("home-manager"));79
+ Need {80
+ ram: if via_daemon {81
+ base.ram82
+ } 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
+ ..base88
+ }89
+}@@ -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 get42
+/// treated as words, which errs toward seeing a build that isn't there43
+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
+ out79
+}80
+81
+/// peels env assignments and wrappers that don't change what runs82
+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> = seg132
+ .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::Clean141
+ } 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_input174
+ .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's209
+ // prompt for whatever the agent wrote, rewriting shouldn't210
+ 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 the244
+/// startup id while the json moves on, and we need both ends to agree245
+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> = late258
+ .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.id269
+ )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 agent306
+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
+ 0319
+}@@ -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
+}@@ -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
+ dir32
+}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
+ dir54
+}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
+}@@ -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 doing49
+ 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 eof106
+ 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
+}@@ -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 again81
+ 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
+ 185
+ })?;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::TimedOut126
+ ) =>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
+ cmd196
+}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
+ status215
+ .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.ram377
+ ));378
+ }379
+ if timed_out {380
+ EXIT_TIMEOUT381
+ } else {382
+ code(status)383
+ }384
+}@@ -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 fit3
+//! in the leftovers once it does4
+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 overdue37
+ 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 already45
+ /// reflects this, so it's netted off promises to avoid counting it twice46
+ 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 finish127
+ 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
+ ..need141
+ }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 = running147
+ .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 = cap152
+ .ram_live153
+ .saturating_sub(cap.ram_reserve)154
+ .saturating_sub(unrealised);155
+156
+ let disk = cap157
+ .disks158
+ .iter()159
+ .map(|(id, d)| {160
+ let promised: Bytes = running161
+ .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: cap176
+ .cores177
+ .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 someone190
+ /// else's ram. backfill falls back to the time cap then191
+ 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<_> = running228
+ .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 = res256
+ .extra257
+ .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
+ out273
+}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
+}@@ -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 file19
+pub fn update<T, R>(name: &str, f: impl FnOnce(&mut T) -> R) -> io::Result<R>20
+where21
+ 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
+}@@ -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()) != 021
+ || libc::stat(c.as_ptr(), st.as_mut_ptr()) != 022
+ {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 this123
+ // runs once a second at most124
+ 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
+ None142
+ }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};@@ -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 = s49
+ .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 = tail154
+ .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 = p255
+ .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
+}