From f425f64276750b63cdf2605f538c5802a1c4b84c Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 10 Oct 2026 17:23:57 +0000 Subject: [PATCH 01/27] chore: project borg from Mono a5293e119fc2 --- .repository-projection.json | 13 + Cargo.lock | 2390 +++++++++++++++++++++++++++++++++++ Cargo.toml | 67 + LICENSE | 94 ++ README.md | 163 +++ src/allocator.rs | 1038 +++++++++++++++ src/api.rs | 482 +++++++ src/classes.rs | 297 +++++ src/config.rs | 196 +++ src/drone/cleanup.rs | 489 +++++++ src/drone/client.rs | 97 ++ src/drone/config.rs | 201 +++ src/drone/exec.rs | 42 + src/drone/install.rs | 956 ++++++++++++++ src/drone/mod.rs | 302 +++++ src/drone/pins.rs | 103 ++ src/drone/runner.rs | 430 +++++++ src/enroll.rs | 54 + src/github.rs | 367 ++++++ src/lib.rs | 21 + src/main.rs | 264 ++++ src/probe/mod.rs | 382 ++++++ src/probe/system.rs | 103 ++ src/queen.rs | 413 ++++++ src/secret.rs | 42 + src/store.rs | 522 ++++++++ tests/drone.rs | 291 +++++ tests/queen.rs | 314 +++++ tests/support/mod.rs | 345 +++++ 29 files changed, 10478 insertions(+) create mode 100644 .repository-projection.json create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 LICENSE create mode 100644 README.md create mode 100644 src/allocator.rs create mode 100644 src/api.rs create mode 100644 src/classes.rs create mode 100644 src/config.rs create mode 100644 src/drone/cleanup.rs create mode 100644 src/drone/client.rs create mode 100644 src/drone/config.rs create mode 100644 src/drone/exec.rs create mode 100644 src/drone/install.rs create mode 100644 src/drone/mod.rs create mode 100644 src/drone/pins.rs create mode 100644 src/drone/runner.rs create mode 100644 src/enroll.rs create mode 100644 src/github.rs create mode 100644 src/lib.rs create mode 100644 src/main.rs create mode 100644 src/probe/mod.rs create mode 100644 src/probe/system.rs create mode 100644 src/queen.rs create mode 100644 src/secret.rs create mode 100644 src/store.rs create mode 100644 tests/drone.rs create mode 100644 tests/queen.rs create mode 100644 tests/support/mod.rs diff --git a/.repository-projection.json b/.repository-projection.json new file mode 100644 index 0000000..fb2eacb --- /dev/null +++ b/.repository-projection.json @@ -0,0 +1,13 @@ +{ + "schemaVersion": 1, + "projection": "borg", + "projectionSchemaVersion": 1, + "sourceRepository": "dx-corp/mono", + "sourceSha": "a5293e119fc29a43fef8e2f3fb3e06f8db5a931b", + "destinationRepository": "dx-corp/borg", + "priorProjectedBase": "bb547084fc1e8c6827704c24b9165fa148c84945", + "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", + "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", + "contentDigest": "38a3bf905bd076b969e0df500b40705791f6e6b163d9fc992462efaed7943c75", + "publicationEligible": true +} diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..70bf517 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,2390 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "aws-lc-rs" +version = "1.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b281d307588d634de920874890732659e2e7672f72b5e10e81badc1a8a83621e" +dependencies = [ + "aws-lc-sys", + "untrusted 0.7.1", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bff6c3b54fad79a2e60b8102caf565819711497c1f5f092f49508e2f5c31b27" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", +] + +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "borg" +version = "0.1.0" +dependencies = [ + "anyhow", + "aws-lc-rs", + "axum", + "base64", + "clap", + "hex", + "jsonwebtoken", + "rand 0.8.6", + "reqwest", + "serde", + "serde_json", + "sha2", + "sqlx", + "tempfile", + "thiserror", + "time", + "tokio", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" + +[[package]] +name = "cc" +version = "1.2.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + +[[package]] +name = "crossbeam-queue" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dotenvy" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" + +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" +dependencies = [ + "serde", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flume" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da0e4dd2a88388a1f4ccc7c9ce104604dab68d9f408dc34cd45823d5a9069095" +dependencies = [ + "futures-core", + "futures-sink", + "spin", +] + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-intrusive" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d930c203dd0b6ff06e0201a4a2fe9149b43c684fd4420555b26d21b1a02956f" +dependencies = [ + "futures-core", + "lock_api", + "parking_lot", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-io", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 5.3.0", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "hashlink" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1" +dependencies = [ + "hashbrown 0.15.5", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "tokio", + "tokio-rustls", + "tower-service", + "webpki-roots 1.0.8", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "jsonwebtoken" +version = "10.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0529410abe238729a60b108898784df8984c87f6054c9c4fcacc47e4803c1ce1" +dependencies = [ + "aws-lc-rs", + "base64", + "getrandom 0.2.17", + "js-sys", + "pem", + "serde", + "serde_json", + "signature", + "simple_asn1", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libsqlite3-sys" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "mio" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64", + "serde_core", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +dependencies = [ + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.4", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.52.0", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "webpki-roots 1.0.8", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted 0.9.0", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted 0.9.0", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "rand_core 0.6.4", +] + +[[package]] +name = "simple_asn1" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d585997b0ac10be3c5ee635f1bab02d512760d14b7c468801ac8a01d9ae5f1d" +dependencies = [ + "num-bigint", + "num-traits", + "thiserror", + "time", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "spin" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" +dependencies = [ + "lock_api", +] + +[[package]] +name = "sqlx" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fefb893899429669dcdd979aff487bd78f4064e5e7907e4269081e0ef7d97dc" +dependencies = [ + "sqlx-core", + "sqlx-macros", + "sqlx-sqlite", +] + +[[package]] +name = "sqlx-core" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6" +dependencies = [ + "base64", + "bytes", + "crc", + "crossbeam-queue", + "either", + "event-listener", + "futures-core", + "futures-intrusive", + "futures-io", + "futures-util", + "hashbrown 0.15.5", + "hashlink", + "indexmap", + "log", + "memchr", + "once_cell", + "percent-encoding", + "rustls", + "serde", + "sha2", + "smallvec", + "thiserror", + "tokio", + "tokio-stream", + "tracing", + "url", + "webpki-roots 0.26.11", +] + +[[package]] +name = "sqlx-macros" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2d452988ccaacfbf5e0bdbc348fb91d7c8af5bee192173ac3636b5fb6e6715d" +dependencies = [ + "proc-macro2", + "quote", + "sqlx-core", + "sqlx-macros-core", + "syn", +] + +[[package]] +name = "sqlx-macros-core" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19a9c1841124ac5a61741f96e1d9e2ec77424bf323962dd894bdb93f37d5219b" +dependencies = [ + "dotenvy", + "either", + "heck", + "hex", + "once_cell", + "proc-macro2", + "quote", + "serde", + "serde_json", + "sha2", + "sqlx-core", + "sqlx-sqlite", + "syn", + "tokio", + "url", +] + +[[package]] +name = "sqlx-sqlite" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2d12fe70b2c1b4401038055f90f151b78208de1f9f89a7dbfd41587a10c3eea" +dependencies = [ + "atoi", + "flume", + "futures-channel", + "futures-core", + "futures-executor", + "futures-intrusive", + "futures-util", + "libsqlite3-sys", + "log", + "percent-encoding", + "serde", + "serde_urlencoded", + "sqlx-core", + "thiserror", + "tracing", + "url", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "time" +version = "0.3.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e48db7b415311b615f910b3dcaa4557bcd4bf1982379c95c223fd8c2a20e210" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "time-macros" +version = "0.2.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c431b87111666e491a90baa837f914fb45cd5dc3c268591b0220ff5057f2085f" +dependencies = [ + "num-conv", + "time-core", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "untrusted" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.8", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..d99b00f --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,67 @@ +[package] +name = "borg" +version = "0.1.0" +edition = "2024" +license = "BUSL-1.1" +rust-version = "1.95" +publish = false +repository = "https://github.com/dx-corp/borg" +description = "Self-organizing CI fleet agent: drones probe their host and run ephemeral runners that a queen leases to them." + +[lib] +name = "borg" +path = "src/lib.rs" + +[[bin]] +name = "borg" +path = "src/main.rs" +test = false + +[dependencies] +anyhow = "1.0" +axum = "0.8" +clap = { version = "4.5", features = ["derive", "env"] } +hex = "0.4" +jsonwebtoken = { version = "10.3", default-features = false, features = ["aws_lc_rs", "use_pem"] } +rand = "0.8" +reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "json"] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +sha2 = "0.10" +sqlx = { version = "0.8", default-features = false, features = ["runtime-tokio-rustls", "sqlite"] } +thiserror = "2.0" +time = { version = "0.3", features = ["formatting", "parsing"] } +tokio = { version = "1.48", features = ["macros", "net", "rt-multi-thread", "process", "signal", "sync", "time"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } + +[dev-dependencies] +aws-lc-rs = "1.18" +base64 = "0.22" +tempfile = "3.14" + +[lints.rust] +non_ascii_idents = "deny" +unsafe_code = "forbid" +unexpected_cfgs = { level = "warn", check-cfg = ["cfg(kani)"] } +unused_lifetimes = "warn" + +[lints.clippy] +await_holding_lock = "deny" +dbg_macro = "deny" +disallowed_methods = "deny" +empty_drop = "deny" +exit = "deny" +filetype_is_file = "deny" +fn_to_numeric_cast_any = "deny" +lossy_float_literal = "deny" +mem_forget = "deny" +mutex_atomic = "deny" +rc_buffer = "deny" +rest_pat_in_fully_bound_structs = "deny" +string_add = "deny" +string_lit_as_bytes = "deny" +todo = "deny" +unchecked_time_subtraction = "deny" +undocumented_unsafe_blocks = "deny" +verbose_file_reads = "deny" diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..acb5227 --- /dev/null +++ b/LICENSE @@ -0,0 +1,94 @@ +SPDX-License-Identifier: BUSL-1.1 + +Business Source License 1.1 + +License text copyright (c) 2017 MariaDB Corporation Ab, All Rights Reserved. +"Business Source License" is a trademark of MariaDB Corporation Ab. + +Licensor: + +EvalOps, Inc. + +Licensed Work: + +The Licensed Work is EvalOps Mono (Deixic), including all source code and associated files +made available in this repository, and is Copyright (c) 2026 EvalOps, Inc. + +Additional Use Grant: + +None + +Change Date: + +Four years from the date each version of the Licensed Work is first publicly +distributed. + +Change License: + +GPL Version 2.0 or any later version + +Terms + +The Licensor hereby grants you the right to copy, modify, create derivative +works, redistribute, and make non-production use of the Licensed Work. The +Licensor may make an Additional Use Grant, above, permitting limited production +use. + +Effective on the Change Date, or the fourth anniversary of the first publicly +available distribution of a specific version of the Licensed Work under this +License, whichever comes first, the Licensor hereby grants you rights under the +terms of the Change License, and the rights granted in the paragraph above +terminate. + +If your use of the Licensed Work does not comply with the requirements +currently in effect as described in this License, you must purchase a +commercial license from the Licensor, its affiliated entities, or authorized +resellers, or you must refrain from using the Licensed Work. + +All copies of the original and modified Licensed Work, and derivative works of +the Licensed Work, are subject to this License. This License applies separately +for each version of the Licensed Work and the Change Date may vary for each +version of the Licensed Work released by Licensor. + +You must conspicuously display this License on each original or modified copy +of the Licensed Work. If you receive the Licensed Work in original or modified +form from a third party, the terms and conditions set forth in this License +apply to your use of that work. + +Any use of the Licensed Work in violation of this License will automatically +terminate your rights under this License for the current and all other versions +of the Licensed Work. + +This License does not grant you any right in any trademark or logo of Licensor +or its affiliates (provided that you may use a trademark or logo of Licensor as +expressly required by this License). + +TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED ON AN +"AS IS" BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND CONDITIONS, EXPRESS +OR IMPLIED, INCLUDING (WITHOUT LIMITATION) WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT, AND TITLE. + +MariaDB hereby grants you permission to use this License's text to license your +works, and to refer to it using the trademark "Business Source License", as +long as you comply with the Covenants of Licensor below. + +Covenants of Licensor + +In consideration of the right to use this License's text and the "Business +Source License" name and trademark, Licensor covenants to MariaDB, and to all +other recipients of the licensed work to be provided by Licensor: + +1. To specify as the Change License the GPL Version 2.0 or any later version, + or a license that is compatible with GPL Version 2.0 or a later version, + where "compatible" means that software provided under the Change License can + be included in a program with software provided under GPL Version 2.0 or a + later version. Licensor may specify additional Change Licenses without + limitation. + +2. To either: (a) specify an additional grant of rights to use that does not + impose any additional restriction on the right granted in this License, as + the Additional Use Grant; or (b) insert the text "None". + +3. To specify a Change Date. + +4. Not to modify this License in any other way. diff --git a/README.md b/README.md new file mode 100644 index 0000000..2a558fb --- /dev/null +++ b/README.md @@ -0,0 +1,163 @@ +# borg + +borg is a self-organizing CI fleet agent for self-hosted GitHub Actions +runners. Every machine that joins the fleet runs a drone. A single queen +watches the queue, decides how much of each machine to lend to which kind of +job, and hands drones one-shot (ephemeral, JIT) runner configurations. + +The crate contains the pure core and the queen: + +| Module | Does | +| --- | --- | +| `probe` | Reports vCPU, memory, free disk, architecture, OS, host kind (WSL, LXC, VM, bare metal), and whether Docker, mise and git are usable. | +| `classes` | Defines the `light`, `medium` and `heavy` runner classes, their resource requirements, and the label scheme. All of it is configuration with defaults. | +| `allocator` | A pure function from queued demand, drone capacity, in-flight leases, warm floors and per-host caps to lease grants with exact label sets. | + +| `queen`, `github`, `store`, `api` | The coordinator: polls queued jobs, runs the allocator, mints single-use JIT runner configs, keeps lease bookkeeping in SQLite and serves the drone API plus `/status` and `/metrics`. | + +| `drone` | The agent on every machine: heartbeats to the queen, runs each lease as an ephemeral runner in a private directory, cleans up after every job, and installs itself. | + +## Build + +``` +cargo build --locked +cargo test --locked +``` + +## Probe + +``` +borg probe # pretty JSON +borg probe --compact # one line +``` + +## Labels + +A job asks for a class with a runner label: `` for jobs from +the default branch or a merge queue, and `-pr` for pull +requests. The default prefix is `borg-ci-`; set `label_prefix` in the classes +config to change it (for example `example-ci-`). A runner minted for main +demand carries only the bare label, so GitHub can hand it nothing but main +jobs. A runner minted for pull-request demand carries both labels. + +## Classes config + +`Classes::from_json` accepts partial JSON; omitted fields keep their defaults: + +```json +{ + "label_prefix": "example-ci-", + "heavy": { "mem_mib": 20480, "vcpu": 6, "docker": true } +} +``` + +## Queen + +``` +borg queen enroll --config queen.json host-a # prints a one-time token +borg queen run --config queen.json +``` + +Example `queen.json` (every field except `github`, `repos` and `db_path` is +optional): + +```json +{ + "github": { "org": "example-org", "app_id": "12345", "private_key_path": "/etc/borg/app.pem" }, + "repos": ["example-org/app"], + "db_path": "/var/lib/borg/queen.db", + "listen": "127.0.0.1:8080", + "classes": { "label_prefix": "example-ci-" }, + "starvation_after_secs": 120 +} +``` + +The GitHub App needs organization self-hosted runners: read and write, plus +read access to Actions on the listed repositories. Set +`BORG_APP_PRIVATE_KEY` instead of `private_key_path` to pass the PEM in the +environment. The queen speaks plain HTTP; bind it to a private interface or +put a TLS terminator in front. + +| Endpoint | Purpose | +| --- | --- | +| `GET /status` | Per-class demand, slots and leases; drones and heartbeat age; oldest queued age per label; `reserved_for`. | +| `GET /metrics` | The same numbers in Prometheus text format. | +| `POST /v1/heartbeat` | Drone API. Reports capacity; returns new leases until their JIT config is acknowledged. | +| `POST /v1/leases/{id}/result` | Drone API. Reports a lease's exit code, duration and peak RSS. | + +Drone requests carry `Authorization: Bearer `. The queen stores only +the SHA-256 of each token. + +### Starvation guard + +When demand for a class has waited longer than `starvation_after_secs`, the +allocator places it first. If it still does not fit, it holds back the drone +that can free up soonest (fewest in-flight leases, ties by name) and grants no +smaller lease there until the starving lease fits. Held drones are listed in +`/status` under `reserved_for`. + +## Drone + +``` +sudo borg drone install --queen https://queen.example.test:8443 --token --ca-cert queen-ca.pem +borg drone install ... --dry-run # list changes without applying them +borg drone run --config /etc/borg/drone.json +``` + +`install` is idempotent. It prints one line per action (`created`, `updated`, +`unchanged`, `ran`, ...) and restarts the service only if something changed. +It sets up: + +* **Linux and WSL:** a `borg` system user (added to the `docker` group when + that group exists) and a systemd unit. The unit sends `SIGTERM` to the drone + only (`KillMode=mixed`), so in-flight jobs finish while it drains. WSL needs + systemd enabled: put `[boot]` and `systemd=true` in `/etc/wsl.conf`, run + `wsl --shutdown` from Windows, and start the distribution again. +* **macOS:** a LaunchDaemon (`org.borg.drone`) that runs as `--run-as `, + which must already exist. +* The pinned actions-runner (version and SHA-256 are in `src/drone/pins.rs`), + downloaded and checksum-verified before it is unpacked, a copy of this + binary, `drone.json` and the token file (mode 0600). + +The queen URL must be `https://`; `ca_cert_path` makes that CA the only trust +root. A plain `http://` URL is refused unless `allow_insecure` is set. + +### Owner caps + +`--max-vcpu`, `--max-mem-mib` and `--max-leases` (config keys of the same +names) are sent with every heartbeat. The queen allocates within the capped +budget, so an owner can say "this PC gives at most 8 cores and 24 GiB". + +### What a lease does + +1. The pinned runner install is cloned into `/` (hard + links where possible) and `run.sh` starts in its own process group. The JIT + config is passed in `ACTIONS_RUNNER_INPUT_JITCONFIG`, never on the command + line, so it is not visible in `ps`; the runner masks it and removes it from its + environment. Jobs see `BORG_LEASE_ID`, `BORG_DOCKER_PREFIX`, + `COMPOSE_PROJECT_NAME` (both `borg-`) and `BORG_CACHE_ROOT`. +2. When the runner exits, or `max_job_secs` passes, the process group is + killed, and the work directory (including `_diag`) is deleted. +3. Docker containers, networks and volumes are removed only if `docker + inspect` shows they carry the label `borg.lease=`, the compose + project `borg-`, or a name that is exactly that prefix or starts + with it plus `-` or `_`. There is no global prune. Jobs that start their own + containers should name them with `$BORG_DOCKER_PREFIX` (docker compose does + this through `COMPOSE_PROJECT_NAME`). +4. Tool caches under `cache_root` are never touched. +5. The exit code, duration and peak RSS of the runner's process group (Linux) + are reported to the queen. + +On start the drone deletes leftover lease directories and borg-labelled docker +resources from a previous run. On `SIGTERM` it stops accepting leases and waits +up to `drain_timeout_secs` for running jobs before stopping them. + +## Contributing + +Issues are welcome here. Code changes land in the source-of-truth repository +and are projected into this one; pull requests opened directly against this +repository will be overwritten by the next projection. + +## License + +Business Source License 1.1 (BUSL-1.1). See [LICENSE](LICENSE). diff --git a/src/allocator.rs b/src/allocator.rs new file mode 100644 index 0000000..790f0c8 --- /dev/null +++ b/src/allocator.rs @@ -0,0 +1,1038 @@ +//! The lease allocator: a pure function from fleet state to lease grants. +//! +//! Inputs are queued demand per (class, scope), the drones' usable capacity, +//! leases already in flight, warm floors and per-host caps. The output says +//! which drone should run which runner, with the exact label set to mint it +//! with. Nothing here does I/O, reads a clock or draws a random number, so the +//! same input always yields the same grants. +//! +//! Rules, in priority order: +//! +//! 1. Main demand (bare label) is satisfied before pull-request demand, and +//! within a scope larger classes are placed first so they are not starved +//! of contiguous room by small ones. +//! 2. A main grant carries only the bare label; a pull-request grant carries +//! both labels (see [`Classes::runner_labels`]). +//! 3. Queued demand is first offset by leases that were granted but have not +//! yet taken a job ([`LeaseState::Pending`]); only the shortfall is minted. +//! 4. A warm floor keeps that many pending leases of a (class, scope) alive +//! even with nothing queued, so the first job never waits on a cold mint. +//! 5. A grant is only made if the drone still has the memory, vCPU and lease +//! slots for it after every in-flight lease and earlier grant. Hosts are +//! carved dynamically: the same drone can end up as one heavy, two medium +//! or several light slots depending on what is queued. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::classes::{Class, ClassSpec, Classes, Scope}; +use crate::probe::Capacity; + +/// Owner-set limits for one host ("this PC gives at most 8 cores / 24 GiB"). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct HostCaps { + pub max_vcpu: Option, + pub max_mem_mib: Option, + /// Maximum concurrent leases on the host, whatever their size. + pub max_leases: Option, + /// Memory kept back for the operating system and the drone itself. + pub reserve_mem_mib: u64, +} + +impl Default for HostCaps { + fn default() -> Self { + Self { + max_vcpu: None, + max_mem_mib: None, + max_leases: None, + reserve_mem_mib: 2048, + } + } +} + +/// One drone as the allocator sees it: budgets already net of caps. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct DroneState { + pub id: String, + pub vcpu: u32, + pub mem_mib: u64, + pub docker: bool, + /// False while the drone drains (SIGTERM) or yields to its owner. + pub accepting: bool, + pub max_leases: Option, +} + +impl DroneState { + /// Budget a probed host: apply the OS reserve and the owner's caps. + pub fn from_capacity(id: impl Into, capacity: &Capacity, caps: &HostCaps) -> Self { + let mem = capacity.mem_mib.saturating_sub(caps.reserve_mem_mib); + Self { + id: id.into(), + vcpu: caps + .max_vcpu + .map_or(capacity.vcpu, |m| m.min(capacity.vcpu)), + mem_mib: caps.max_mem_mib.map_or(mem, |m| m.min(mem)), + docker: capacity.docker_usable, + accepting: true, + max_leases: caps.max_leases, + } + } +} + +/// Whether a lease's runner has picked up a job yet. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum LeaseState { + /// Granted or minted, runner idle: can still absorb queued demand. + Pending, + /// Runner is executing a job: only consumes capacity. + Running, +} + +/// A lease already granted and not yet finished. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ActiveLease { + pub drone_id: String, + pub class: Class, + pub scope: Scope, + pub state: LeaseState, +} + +/// Queued jobs waiting for a runner of `(class, scope)`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct Demand { + pub class: Class, + pub scope: Scope, + pub count: u32, +} + +/// Keep at least `count` pending leases of `(class, scope)` alive. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct WarmFloor { + pub class: Class, + pub scope: Scope, + pub count: u32, +} + +/// Default floor: one idle main-only medium runner. +pub fn default_warm_floors() -> Vec { + vec![WarmFloor { + class: Class::Medium, + scope: Scope::Main, + count: 1, + }] +} + +/// Age of the oldest queued job for `(class, scope)`, from the queen's view of +/// the jobs' `created_at`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct QueueAge { + pub class: Class, + pub scope: Scope, + pub oldest_age_secs: u64, +} + +/// Everything the allocator needs. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AllocatorInput { + pub classes: Classes, + pub demand: Vec, + pub drones: Vec, + pub leases: Vec, + pub warm_floors: Vec, + pub queue_ages: Vec, + /// Demand that has waited at least this long is starving; `None` turns + /// the guard off. + pub starvation_after_secs: Option, +} + +/// An instruction to mint one runner and hand it to one drone. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Grant { + pub drone_id: String, + pub class: Class, + pub scope: Scope, + /// Exact label set to mint the runner with. + pub labels: Vec, + pub mem_mib: u64, + pub vcpu: u32, +} + +/// Demand that could not be placed this round. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct Unmet { + pub class: Class, + pub scope: Scope, + pub count: u32, +} + +/// A drone held back for starving demand: it takes no new lease except one of +/// this `(class, scope)` until that lease fits. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Reservation { + pub drone_id: String, + pub class: Class, + pub scope: Scope, +} + +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] +pub struct Allocation { + pub grants: Vec, + pub unmet: Vec, + pub reserved_for: Vec, +} + +/// How many leases of `spec` this drone could hold if it were empty. +pub fn max_slots(drone: &DroneState, spec: &ClassSpec) -> u32 { + if spec.docker && !drone.docker { + return 0; + } + let by_mem = drone.mem_mib / spec.mem_mib.max(1); + let by_cpu = u64::from(drone.vcpu / spec.vcpu.max(1)); + let by_cap = drone.max_leases.map_or(u64::MAX, u64::from); + u32::try_from(by_mem.min(by_cpu).min(by_cap)).unwrap_or(u32::MAX) +} + +/// Remaining room on a drone while allocating. +struct Free<'a> { + drone: &'a DroneState, + mem_mib: u64, + vcpu: u32, + leases_left: Option, + in_flight: u32, +} + +impl Free<'_> { + fn fits(&self, spec: &ClassSpec) -> bool { + self.drone.accepting + && (!spec.docker || self.drone.docker) + && self.mem_mib >= spec.mem_mib + && self.vcpu >= spec.vcpu + && self.leases_left != Some(0) + } + + /// Could this drone ever hold `spec` if it were empty? + fn could_ever_fit(&self, spec: &ClassSpec) -> bool { + self.drone.accepting + && (!spec.docker || self.drone.docker) + && self.drone.mem_mib >= spec.mem_mib + && self.drone.vcpu >= spec.vcpu + && self.drone.max_leases != Some(0) + } + + fn take(&mut self, spec: &ClassSpec) { + self.mem_mib = self.mem_mib.saturating_sub(spec.mem_mib); + self.vcpu = self.vcpu.saturating_sub(spec.vcpu); + self.leases_left = self.leases_left.map(|n| n.saturating_sub(1)); + } +} + +type Counts = BTreeMap<(Class, Scope), u32>; + +fn sum_counts>(items: I) -> Counts { + let mut map = BTreeMap::new(); + for (class, scope, count) in items { + let slot: &mut u32 = map.entry((class, scope)).or_default(); + *slot = slot.saturating_add(count); + } + map +} + +/// Place up to `wanted` runners of `(class, scope)` on the tightest-fitting +/// drones, skipping drones reserved for something else. Returns how many fit. +fn place( + input: &AllocatorInput, + free: &mut BTreeMap<&str, Free<'_>>, + reserved: &BTreeMap, + (scope, class): (Scope, Class), + wanted: u32, + out: &mut Vec, +) -> u32 { + let spec = input.classes.spec(class); + let mut placed = 0; + for _ in 0..wanted { + // Tightest fit, then id: keeps hosts packed and output stable. + let best = free + .iter_mut() + .filter(|(id, f)| { + f.fits(spec) + && reserved + .get(**id) + .is_none_or(|held| *held == (class, scope)) + }) + .min_by_key(|(id, f)| (f.mem_mib - spec.mem_mib, **id)) + .map(|(_, f)| f); + let Some(f) = best else { break }; + f.take(spec); + placed += 1; + out.push(Grant { + drone_id: f.drone.id.clone(), + class, + scope, + labels: input.classes.runner_labels(class, scope), + mem_mib: spec.mem_mib, + vcpu: spec.vcpu, + }); + } + placed +} + +/// Decide which runners to mint and where. Pure and deterministic. +pub fn allocate(input: &AllocatorInput) -> Allocation { + // Remaining room per drone after leases already in flight. Duplicate + // drone ids collapse to the first entry. + let mut free: BTreeMap<&str, Free<'_>> = BTreeMap::new(); + for drone in &input.drones { + free.entry(drone.id.as_str()).or_insert(Free { + drone, + mem_mib: drone.mem_mib, + vcpu: drone.vcpu, + leases_left: drone.max_leases, + in_flight: 0, + }); + } + for lease in &input.leases { + if let Some(f) = free.get_mut(lease.drone_id.as_str()) { + f.take(input.classes.spec(lease.class)); + f.in_flight += 1; + } + } + + let queued = sum_counts(input.demand.iter().map(|d| (d.class, d.scope, d.count))); + let floors = sum_counts( + input + .warm_floors + .iter() + .map(|w| (w.class, w.scope, w.count)), + ); + let pending = sum_counts( + input + .leases + .iter() + .filter(|l| l.state == LeaseState::Pending && free.contains_key(l.drone_id.as_str())) + .map(|l| (l.class, l.scope, 1)), + ); + let get = |m: &Counts, c, s| m.get(&(c, s)).copied().unwrap_or(0); + + // How many runners each (scope, class) still has to mint. + let mut need: BTreeMap<(Scope, Class), u32> = BTreeMap::new(); + for class in Class::ALL { + let pending_main = get(&pending, class, Scope::Main); + let pending_pr = get(&pending, class, Scope::Pr); + + // Main demand uses main-only pending runners first, then spare + // pull-request runners (they carry the bare label too). + let main_total = get(&queued, class, Scope::Main).max(get(&floors, class, Scope::Main)); + let after_main_pending = main_total.saturating_sub(pending_main); + let from_pr_pending = after_main_pending.min(pending_pr); + let main_need = after_main_pending - from_pr_pending; + let pending_pr_left = pending_pr - from_pr_pending; + + let pr_total = get(&queued, class, Scope::Pr).max(get(&floors, class, Scope::Pr)); + let pr_need = pr_total.saturating_sub(pending_pr_left); + + need.insert((Scope::Main, class), main_need); + need.insert((Scope::Pr, class), pr_need); + } + + let initial_need = need.clone(); + let mut out = Allocation::default(); + let mut placed: BTreeMap<(Scope, Class), u32> = BTreeMap::new(); + let mut reserved: BTreeMap = BTreeMap::new(); + + // Starvation guard. Demand that has waited too long is placed first; if it + // still cannot fit, hold back the drone that will free up soonest so the + // stream of smaller jobs cannot keep it full. + if let Some(after) = input.starvation_after_secs { + let aged = |class, scope| { + input + .queue_ages + .iter() + .any(|a| a.class == class && a.scope == scope && a.oldest_age_secs >= after) + }; + for scope in Scope::ALL { + for class in Class::ALL.into_iter().rev() { + let wanted = need.get(&(scope, class)).copied().unwrap_or(0); + if wanted == 0 || get(&queued, class, scope) == 0 || !aged(class, scope) { + continue; + } + let n = place( + input, + &mut free, + &reserved, + (scope, class), + wanted, + &mut out.grants, + ); + placed.insert((scope, class), n); + need.insert((scope, class), 0); + if n < wanted { + let spec = input.classes.spec(class); + let target = free + .iter() + .filter(|(id, f)| !reserved.contains_key(**id) && f.could_ever_fit(spec)) + .min_by_key(|(id, f)| (f.in_flight, **id)) + .map(|(id, _)| (*id).to_owned()); + if let Some(id) = target { + reserved.insert(id.clone(), (class, scope)); + out.reserved_for.push(Reservation { + drone_id: id, + class, + scope, + }); + } + } + } + } + } + + for scope in Scope::ALL { + for class in Class::ALL.into_iter().rev() { + let wanted = need.get(&(scope, class)).copied().unwrap_or(0); + let n = place( + input, + &mut free, + &reserved, + (scope, class), + wanted, + &mut out.grants, + ); + *placed.entry((scope, class)).or_default() += n; + } + } + + // Only real queued demand is reported as unmet, not unfilled floors. + for scope in Scope::ALL { + for class in Class::ALL.into_iter().rev() { + let got = placed.get(&(scope, class)).copied().unwrap_or(0); + let wanted = initial_need.get(&(scope, class)).copied().unwrap_or(0); + let unmet = get(&queued, class, scope).min(wanted).saturating_sub(got); + if unmet > 0 { + out.unmet.push(Unmet { + class, + scope, + count: unmet, + }); + } + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 28 GiB host with 16 vCPU and Docker, already net of the OS reserve. + fn host(id: &str) -> DroneState { + DroneState { + id: id.into(), + vcpu: 16, + mem_mib: 28 * 1024, + docker: true, + accepting: true, + max_leases: None, + } + } + + fn input(demand: &[(Class, Scope, u32)], drones: Vec) -> AllocatorInput { + AllocatorInput { + classes: Classes::default(), + demand: demand + .iter() + .map(|&(class, scope, count)| Demand { + class, + scope, + count, + }) + .collect(), + drones, + leases: vec![], + warm_floors: vec![], + queue_ages: vec![], + starvation_after_secs: None, + } + } + + fn count(a: &Allocation, class: Class, scope: Scope) -> usize { + a.grants + .iter() + .filter(|g| g.class == class && g.scope == scope) + .count() + } + + fn lease(drone: &str, class: Class, scope: Scope, state: LeaseState) -> ActiveLease { + ActiveLease { + drone_id: drone.into(), + class, + scope, + state, + } + } + + fn assert_within_budget(i: &AllocatorInput, a: &Allocation) { + for d in &i.drones { + let spec = |c| i.classes.spec(c); + let mut mem = 0; + let mut vcpu = 0; + let mut n = 0u32; + for l in i.leases.iter().filter(|l| l.drone_id == d.id) { + mem += spec(l.class).mem_mib; + vcpu += spec(l.class).vcpu; + n += 1; + } + for g in a.grants.iter().filter(|g| g.drone_id == d.id) { + mem += g.mem_mib; + vcpu += g.vcpu; + n += 1; + assert!(d.accepting, "grant to non-accepting drone"); + assert!( + d.docker || !spec(g.class).docker, + "docker-less drone got docker class" + ); + } + if a.grants.iter().any(|g| g.drone_id == d.id) { + assert!( + mem <= d.mem_mib, + "{} over-granted memory: {mem} > {}", + d.id, + d.mem_mib + ); + assert!( + vcpu <= d.vcpu, + "{} over-granted vcpu: {vcpu} > {}", + d.id, + d.vcpu + ); + if let Some(cap) = d.max_leases { + assert!(n <= cap, "{} over lease cap", d.id); + } + } + } + } + + #[test] + fn no_demand_no_floor_no_grants() { + let a = allocate(&input(&[], vec![host("a")])); + assert_eq!(a, Allocation::default()); + } + + #[test] + fn warm_floor_alone_grants_exactly_the_floor() { + let mut i = input(&[], vec![host("a")]); + i.warm_floors = default_warm_floors(); + let a = allocate(&i); + assert_eq!(a.grants.len(), 1); + assert_eq!( + (a.grants[0].class, a.grants[0].scope), + (Class::Medium, Scope::Main) + ); + assert_eq!(a.grants[0].labels, ["borg-ci-medium"]); + assert!(a.unmet.is_empty(), "an unfilled floor is not unmet demand"); + } + + #[test] + fn warm_floor_is_satisfied_by_an_existing_pending_lease() { + let mut i = input(&[], vec![host("a")]); + i.warm_floors = default_warm_floors(); + i.leases = vec![lease("a", Class::Medium, Scope::Main, LeaseState::Pending)]; + assert!(allocate(&i).grants.is_empty()); + } + + #[test] + fn running_lease_does_not_satisfy_the_floor() { + let mut i = input(&[], vec![host("a")]); + i.warm_floors = default_warm_floors(); + i.leases = vec![lease("a", Class::Medium, Scope::Main, LeaseState::Running)]; + assert_eq!(allocate(&i).grants.len(), 1); + } + + #[test] + fn floor_and_queued_demand_do_not_double_count() { + let mut i = input(&[(Class::Medium, Scope::Main, 1)], vec![host("a")]); + i.warm_floors = default_warm_floors(); + assert_eq!(allocate(&i).grants.len(), 1); + } + + #[test] + fn main_demand_is_placed_before_pr_demand_when_room_is_short() { + // Room for exactly one heavy; main must win over pr. + let i = input( + &[(Class::Heavy, Scope::Pr, 1), (Class::Heavy, Scope::Main, 1)], + vec![host("a")], + ); + let a = allocate(&i); + assert_eq!(a.grants.len(), 1); + assert_eq!(a.grants[0].scope, Scope::Main); + assert_eq!( + a.unmet, + [Unmet { + class: Class::Heavy, + scope: Scope::Pr, + count: 1 + }] + ); + } + + #[test] + fn small_main_beats_large_pr() { + // 24 GiB budget: main medium x2 fills it, leaving no room for pr heavy. + let mut h = host("a"); + h.mem_mib = 24 * 1024; + let i = input( + &[ + (Class::Heavy, Scope::Pr, 1), + (Class::Medium, Scope::Main, 2), + ], + vec![h], + ); + let a = allocate(&i); + assert_eq!(count(&a, Class::Medium, Scope::Main), 2); + assert_eq!(count(&a, Class::Heavy, Scope::Pr), 0); + } + + #[test] + fn main_grants_carry_only_the_bare_label() { + let a = allocate(&input(&[(Class::Light, Scope::Main, 1)], vec![host("a")])); + assert_eq!(a.grants[0].labels, ["borg-ci-light"]); + } + + #[test] + fn pr_grants_carry_both_labels() { + let a = allocate(&input(&[(Class::Light, Scope::Pr, 1)], vec![host("a")])); + assert_eq!(a.grants[0].labels, ["borg-ci-light", "borg-ci-light-pr"]); + } + + #[test] + fn labels_follow_the_configured_prefix() { + let mut i = input(&[(Class::Light, Scope::Pr, 1)], vec![host("a")]); + i.classes.label_prefix = "example-ci-".into(); + assert_eq!( + allocate(&i).grants[0].labels, + ["example-ci-light", "example-ci-light-pr"] + ); + } + + #[test] + fn carves_one_28gib_host_into_one_heavy() { + let i = input(&[(Class::Heavy, Scope::Main, 3)], vec![host("a")]); + let a = allocate(&i); + assert_eq!(count(&a, Class::Heavy, Scope::Main), 1); + assert_eq!( + a.unmet, + [Unmet { + class: Class::Heavy, + scope: Scope::Main, + count: 2 + }] + ); + assert_within_budget(&i, &a); + } + + #[test] + fn carves_one_28gib_host_into_two_medium() { + let i = input(&[(Class::Medium, Scope::Pr, 5)], vec![host("a")]); + let a = allocate(&i); + assert_eq!(count(&a, Class::Medium, Scope::Pr), 2); + assert_within_budget(&i, &a); + } + + #[test] + fn carves_one_28gib_host_into_eight_light() { + // Memory alone would allow 11; the 16 vCPU budget caps it at 8. + let i = input(&[(Class::Light, Scope::Pr, 20)], vec![host("a")]); + let a = allocate(&i); + assert_eq!(count(&a, Class::Light, Scope::Pr), 8); + assert_within_budget(&i, &a); + } + + #[test] + fn mixed_demand_fills_remaining_room_after_the_big_class() { + let i = input( + &[ + (Class::Light, Scope::Main, 4), + (Class::Heavy, Scope::Main, 1), + ], + vec![host("a")], + ); + let a = allocate(&i); + assert_eq!(count(&a, Class::Heavy, Scope::Main), 1); + // 28 - 24 = 4 GiB left: one light (2.5 GiB) fits. + assert_eq!(count(&a, Class::Light, Scope::Main), 1); + assert_within_budget(&i, &a); + } + + #[test] + fn in_flight_leases_reduce_available_room() { + let mut i = input(&[(Class::Medium, Scope::Main, 2)], vec![host("a")]); + i.leases = vec![lease("a", Class::Medium, Scope::Pr, LeaseState::Running)]; + let a = allocate(&i); + assert_eq!(a.grants.len(), 1); + assert_within_budget(&i, &a); + } + + #[test] + fn pending_leases_absorb_queued_demand_of_the_same_label() { + let mut i = input(&[(Class::Light, Scope::Main, 2)], vec![host("a")]); + i.leases = vec![lease("a", Class::Light, Scope::Main, LeaseState::Pending)]; + assert_eq!(allocate(&i).grants.len(), 1); + } + + #[test] + fn pending_main_lease_cannot_serve_pr_demand() { + let mut i = input(&[(Class::Light, Scope::Pr, 1)], vec![host("a")]); + i.leases = vec![lease("a", Class::Light, Scope::Main, LeaseState::Pending)]; + assert_eq!(allocate(&i).grants.len(), 1); + } + + #[test] + fn pending_pr_lease_can_serve_main_demand() { + let mut i = input(&[(Class::Light, Scope::Main, 1)], vec![host("a")]); + i.leases = vec![lease("a", Class::Light, Scope::Pr, LeaseState::Pending)]; + assert!(allocate(&i).grants.is_empty()); + } + + #[test] + fn a_pending_pr_lease_is_not_counted_twice() { + let mut i = input( + &[(Class::Light, Scope::Main, 1), (Class::Light, Scope::Pr, 1)], + vec![host("a")], + ); + i.leases = vec![lease("a", Class::Light, Scope::Pr, LeaseState::Pending)]; + let a = allocate(&i); + assert_eq!(a.grants.len(), 1); + assert_eq!(a.grants[0].scope, Scope::Pr); + } + + #[test] + fn host_lease_cap_limits_grants() { + let mut h = host("a"); + h.max_leases = Some(3); + let i = input(&[(Class::Light, Scope::Main, 8)], vec![h]); + let a = allocate(&i); + assert_eq!(a.grants.len(), 3); + assert_within_budget(&i, &a); + } + + #[test] + fn host_lease_cap_counts_in_flight_leases() { + let mut h = host("a"); + h.max_leases = Some(2); + let mut i = input(&[(Class::Light, Scope::Main, 8)], vec![h]); + i.leases = vec![lease("a", Class::Light, Scope::Pr, LeaseState::Running)]; + assert_eq!(allocate(&i).grants.len(), 1); + } + + #[test] + fn owner_caps_shrink_the_budget() { + let capacity = Capacity { + hostname: "h".into(), + os: "linux".into(), + arch: "x86_64".into(), + kind: crate::probe::HostKind::Wsl, + vcpu: 20, + mem_mib: 64 * 1024, + free_disk_mib: None, + docker_usable: true, + mise_version: None, + git_version: None, + }; + let caps = HostCaps { + max_vcpu: Some(8), + max_mem_mib: Some(24 * 1024), + ..HostCaps::default() + }; + let d = DroneState::from_capacity("pc", &capacity, &caps); + assert_eq!((d.vcpu, d.mem_mib), (8, 24 * 1024)); + + let i = input(&[(Class::Light, Scope::Main, 20)], vec![d]); + let a = allocate(&i); + assert_eq!(a.grants.len(), 4); // 8 vCPU / 2 + assert_within_budget(&i, &a); + } + + #[test] + fn reserve_is_subtracted_from_total_memory() { + let capacity = Capacity { + hostname: "h".into(), + os: "linux".into(), + arch: "x86_64".into(), + kind: crate::probe::HostKind::Bare, + vcpu: 16, + mem_mib: 28 * 1024, + free_disk_mib: None, + docker_usable: true, + mise_version: None, + git_version: None, + }; + let d = DroneState::from_capacity("a", &capacity, &HostCaps::default()); + assert_eq!(d.mem_mib, 26 * 1024); + // Still fits one heavy. + let a = allocate(&input(&[(Class::Heavy, Scope::Main, 1)], vec![d])); + assert_eq!(a.grants.len(), 1); + } + + #[test] + fn docker_classes_skip_hosts_without_docker() { + let mut nodocker = host("a"); + nodocker.docker = false; + let i = input( + &[ + (Class::Medium, Scope::Main, 1), + (Class::Light, Scope::Main, 1), + ], + vec![nodocker, host("b")], + ); + let a = allocate(&i); + let medium = a.grants.iter().find(|g| g.class == Class::Medium).unwrap(); + assert_eq!(medium.drone_id, "b"); + assert_within_budget(&i, &a); + } + + #[test] + fn draining_drones_get_nothing() { + let mut h = host("a"); + h.accepting = false; + let a = allocate(&input(&[(Class::Light, Scope::Main, 2)], vec![h])); + assert!(a.grants.is_empty()); + assert_eq!( + a.unmet, + [Unmet { + class: Class::Light, + scope: Scope::Main, + count: 2 + }] + ); + } + + #[test] + fn spreads_over_hosts_when_one_is_full_and_prefers_the_tightest_fit() { + let mut small = host("small"); + small.mem_mib = 13 * 1024; + let i = input(&[(Class::Medium, Scope::Main, 3)], vec![host("big"), small]); + let a = allocate(&i); + assert_eq!(a.grants.len(), 3); + // The first medium goes to the tighter 13 GiB host. + assert_eq!(a.grants[0].drone_id, "small"); + assert_within_budget(&i, &a); + } + + #[test] + fn leases_on_unknown_drones_are_ignored() { + let mut i = input(&[(Class::Light, Scope::Main, 1)], vec![host("a")]); + i.leases = vec![lease( + "gone", + Class::Heavy, + Scope::Main, + LeaseState::Pending, + )]; + assert_eq!(allocate(&i).grants.len(), 1); + } + + #[test] + fn duplicate_demand_entries_are_summed() { + let i = input( + &[ + (Class::Light, Scope::Main, 1), + (Class::Light, Scope::Main, 2), + ], + vec![host("a")], + ); + assert_eq!(allocate(&i).grants.len(), 3); + } + + #[test] + fn allocation_is_deterministic() { + let i = input( + &[ + (Class::Light, Scope::Pr, 5), + (Class::Medium, Scope::Main, 1), + ], + vec![host("b"), host("a")], + ); + assert_eq!(allocate(&i), allocate(&i)); + } + + fn aged(i: &mut AllocatorInput, class: Class, scope: Scope, secs: u64) { + i.queue_ages.push(QueueAge { + class, + scope, + oldest_age_secs: secs, + }); + i.starvation_after_secs = Some(120); + } + + /// Two 28 GiB drones, each too full for a heavy lease; `a` carries more + /// leases than `b`. + fn crowded() -> AllocatorInput { + let mut i = input( + &[ + (Class::Heavy, Scope::Pr, 1), + (Class::Light, Scope::Main, 10), + ], + vec![host("a"), host("b")], + ); + i.leases = vec![ + lease("a", Class::Medium, Scope::Pr, LeaseState::Running), + lease("a", Class::Light, Scope::Pr, LeaseState::Running), + lease("b", Class::Medium, Scope::Pr, LeaseState::Running), + ]; + i + } + + #[test] + fn starving_heavy_reserves_the_drone_with_fewest_leases() { + let mut i = crowded(); + aged(&mut i, Class::Heavy, Scope::Pr, 300); + let a = allocate(&i); + assert_eq!( + a.reserved_for, + [Reservation { + drone_id: "b".into(), + class: Class::Heavy, + scope: Scope::Pr + }] + ); + assert!( + a.grants.iter().all(|g| g.drone_id != "b"), + "reserved drone got a new lease" + ); + assert!( + count(&a, Class::Light, Scope::Main) > 0, + "other drone keeps working" + ); + assert!(a.unmet.contains(&Unmet { + class: Class::Heavy, + scope: Scope::Pr, + count: 1 + })); + assert_within_budget(&i, &a); + } + + #[test] + fn reservation_ties_break_by_drone_id() { + let mut i = crowded(); + i.leases = vec![ + lease("b", Class::Medium, Scope::Pr, LeaseState::Running), + lease("a", Class::Medium, Scope::Pr, LeaseState::Running), + ]; + aged(&mut i, Class::Heavy, Scope::Pr, 300); + assert_eq!(allocate(&i).reserved_for[0].drone_id, "a"); + } + + #[test] + fn young_demand_is_not_starving() { + let mut i = crowded(); + aged(&mut i, Class::Heavy, Scope::Pr, 119); + let a = allocate(&i); + assert!(a.reserved_for.is_empty()); + assert!(a.grants.iter().any(|g| g.drone_id == "b")); + } + + #[test] + fn guard_can_be_disabled() { + let mut i = crowded(); + aged(&mut i, Class::Heavy, Scope::Pr, 9999); + i.starvation_after_secs = None; + assert!(allocate(&i).reserved_for.is_empty()); + } + + #[test] + fn starving_demand_that_fits_is_granted_first_without_reservation() { + let mut i = input( + &[ + (Class::Light, Scope::Main, 20), + (Class::Heavy, Scope::Pr, 1), + ], + vec![host("a")], + ); + aged(&mut i, Class::Heavy, Scope::Pr, 500); + let a = allocate(&i); + assert_eq!(count(&a, Class::Heavy, Scope::Pr), 1); + assert!(a.reserved_for.is_empty()); + assert_within_budget(&i, &a); + } + + #[test] + fn drones_that_could_never_fit_are_not_reserved() { + let mut small = host("small"); + small.mem_mib = 8 * 1024; + let mut nodocker = host("nodocker"); + nodocker.docker = false; + let mut i = input(&[(Class::Heavy, Scope::Main, 1)], vec![small, nodocker]); + aged(&mut i, Class::Heavy, Scope::Main, 500); + let a = allocate(&i); + assert!(a.reserved_for.is_empty()); + assert!(a.grants.is_empty()); + } + + #[test] + fn starving_medium_also_reserves() { + let mut h = host("a"); + h.mem_mib = 14 * 1024; + let mut i = input( + &[ + (Class::Medium, Scope::Main, 1), + (Class::Light, Scope::Pr, 5), + ], + vec![h], + ); + i.leases = vec![lease("a", Class::Light, Scope::Pr, LeaseState::Running)]; + aged(&mut i, Class::Medium, Scope::Main, 200); + let a = allocate(&i); + assert_eq!(a.reserved_for.len(), 1); + assert!( + a.grants.is_empty(), + "nothing else may land on the reserved drone" + ); + } + + #[test] + fn never_over_grants_across_a_demand_sweep() { + let drones = vec![ + host("a"), + DroneState { + id: "b".into(), + vcpu: 4, + mem_mib: 8 * 1024, + docker: false, + accepting: true, + max_leases: Some(2), + }, + DroneState { + id: "c".into(), + vcpu: 8, + mem_mib: 14 * 1024, + docker: true, + accepting: true, + max_leases: None, + }, + ]; + for lm in 0..3 { + for mm in 0..3 { + for hm in 0..3 { + for lp in [0, 4] { + for hp in [0, 2] { + let mut i = input( + &[ + (Class::Light, Scope::Main, lm), + (Class::Medium, Scope::Main, mm), + (Class::Heavy, Scope::Main, hm), + (Class::Light, Scope::Pr, lp), + (Class::Heavy, Scope::Pr, hp), + ], + drones.clone(), + ); + i.warm_floors = default_warm_floors(); + i.leases = + vec![lease("c", Class::Light, Scope::Pr, LeaseState::Running)]; + assert_within_budget(&i, &allocate(&i)); + } + } + } + } + } + } +} diff --git a/src/api.rs b/src/api.rs new file mode 100644 index 0000000..1c831fd --- /dev/null +++ b/src/api.rs @@ -0,0 +1,482 @@ +//! HTTP surface of the queen. +//! +//! * `GET /status` and `GET /metrics` are read-only and unauthenticated; bind +//! them to a private interface. +//! * `POST /v1/heartbeat` and `POST /v1/leases/{id}/result` are the drone API, +//! authenticated with `Authorization: Bearer `. +//! +//! The queen speaks plain HTTP; terminate TLS in front of it. + +use std::collections::BTreeMap; +use std::fmt::Write as _; +use std::sync::Arc; + +use axum::extract::{Path, State}; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use serde::{Deserialize, Serialize}; +use time::OffsetDateTime; + +use crate::allocator::{DroneState, HostCaps, Reservation, Unmet}; +use crate::classes::{Class, Scope}; +use crate::enroll::hash_token; +use crate::probe::Capacity; +use crate::queen::{Queen, QueenError, QueueEntry}; +use crate::secret::Secret; +use crate::store::{LeaseResult, LeaseStatus, StoreError}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct HeartbeatRequest { + pub capacity: Capacity, + #[serde(default)] + pub caps: HostCaps, + /// False while the drone drains or yields to its owner. + #[serde(default = "yes")] + pub accepting: bool, + /// Lease ids the drone currently holds; acknowledges their JIT configs. + #[serde(default)] + pub active_leases: Vec, +} + +fn yes() -> bool { + true +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct LeaseGrant { + pub lease_id: String, + pub class: Class, + pub scope: Scope, + pub labels: Vec, + pub runner_name: String, + pub jit_config: Secret, + pub expires_at: i64, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct HeartbeatResponse { + pub leases: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ResultRequest { + pub exit_code: i32, + pub duration_secs: u64, + #[serde(default)] + pub peak_rss_mib: Option, +} + +#[derive(Debug)] +pub enum ApiError { + Unauthorized, + NotFound, + Internal(String), +} + +impl From for ApiError { + fn from(e: StoreError) -> Self { + Self::Internal(e.to_string()) + } +} + +impl From for ApiError { + fn from(e: QueenError) -> Self { + Self::Internal(e.to_string()) + } +} + +impl IntoResponse for ApiError { + fn into_response(self) -> Response { + match self { + Self::Unauthorized => StatusCode::UNAUTHORIZED.into_response(), + Self::NotFound => StatusCode::NOT_FOUND.into_response(), + Self::Internal(msg) => { + tracing::error!("api error: {msg}"); + StatusCode::INTERNAL_SERVER_ERROR.into_response() + } + } + } +} + +pub fn router(queen: Arc) -> Router { + Router::new() + .route("/status", get(status)) + .route("/metrics", get(metrics)) + .route("/v1/heartbeat", post(heartbeat)) + .route("/v1/leases/{id}/result", post(lease_result)) + .with_state(queen) +} + +/// The drone name behind the request's bearer token. +async fn authenticate(queen: &Queen, headers: &HeaderMap) -> Result { + let presented = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.strip_prefix("Bearer ")) + .ok_or(ApiError::Unauthorized)?; + queen + .store() + .drone_for_token_hash(&hash_token(presented)) + .await? + .ok_or(ApiError::Unauthorized) +} + +async fn heartbeat( + State(queen): State>, + headers: HeaderMap, + Json(req): Json, +) -> Result, ApiError> { + let drone = authenticate(&queen, &headers).await?; + let now = OffsetDateTime::now_utc().unix_timestamp(); + let store = queen.store(); + store + .record_heartbeat(&drone, &req.capacity, &req.caps, req.accepting, now) + .await?; + store.acknowledge(&drone, &req.active_leases).await?; + let leases = store + .undelivered(&drone) + .await? + .into_iter() + .filter_map(|l| { + Some(LeaseGrant { + lease_id: l.id, + class: l.class, + scope: l.scope, + labels: l.labels, + runner_name: l.runner_name, + jit_config: l.jit_config?, + expires_at: l.expires_at, + }) + }) + .collect(); + Ok(Json(HeartbeatResponse { leases })) +} + +async fn lease_result( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, + Json(req): Json, +) -> Result { + let drone = authenticate(&queen, &headers).await?; + let result = LeaseResult { + exit_code: req.exit_code, + duration_secs: req.duration_secs, + peak_rss_mib: req.peak_rss_mib, + }; + let now = OffsetDateTime::now_utc().unix_timestamp(); + if queen + .store() + .finish_lease(&id, &drone, &result, now) + .await? + { + Ok(StatusCode::NO_CONTENT) + } else { + Err(ApiError::NotFound) + } +} + +#[derive(Debug, Serialize)] +pub struct ClassStatus { + pub queued_main: u32, + pub queued_pr: u32, + pub oldest_age_secs_main: u64, + pub oldest_age_secs_pr: u64, + /// Leases of this class the fleet could hold if empty. + pub slots: u32, + pub leases_pending: u32, + pub leases_running: u32, +} + +#[derive(Debug, Serialize)] +pub struct DroneStatus { + pub name: String, + pub heartbeat_age_secs: Option, + pub fresh: bool, + pub accepting: bool, + pub hostname: Option, + pub kind: Option, + pub vcpu: Option, + pub mem_mib: Option, + pub active_leases: u32, +} + +#[derive(Debug, Serialize)] +pub struct Status { + pub now: i64, + pub last_tick_at: Option, + pub last_error: Option, + pub classes: BTreeMap, + pub queue: Vec, + pub drones: Vec, + pub unmet: Vec, + /// Drones held back for starving demand. + pub reserved_for: Vec, +} + +pub async fn build_status(queen: &Queen, now: i64) -> Result { + let snap = queen.snapshot(); + let cfg = queen.config(); + let rows = queen.store().drones().await?; + let active = queen.store().active_leases().await?; + + let mut states = Vec::new(); + let mut drones = Vec::new(); + for d in &rows { + let age = d.last_heartbeat.map(|t| now - t); + let fresh = age.is_some_and(|a| a <= cfg.drone_stale_secs); + if let (true, Some(cap)) = (fresh, &d.capacity) { + let mut s = DroneState::from_capacity(d.name.clone(), cap, &d.caps); + s.accepting = d.accepting; + states.push(s); + } + drones.push(DroneStatus { + name: d.name.clone(), + heartbeat_age_secs: age, + fresh, + accepting: d.accepting, + hostname: d.capacity.as_ref().map(|c| c.hostname.clone()), + kind: d.capacity.as_ref().map(|c| c.kind), + vcpu: d.capacity.as_ref().map(|c| c.vcpu), + mem_mib: d.capacity.as_ref().map(|c| c.mem_mib), + active_leases: u32::try_from(active.iter().filter(|l| l.drone == d.name).count()) + .unwrap_or(u32::MAX), + }); + } + + let slots = queen.class_slots(&states); + let mut classes = BTreeMap::new(); + for class in Class::ALL { + let entry = |scope| { + snap.queue + .iter() + .find(|q| q.class == class && q.scope == scope) + }; + let count = |state| { + u32::try_from( + active + .iter() + .filter(|l| l.class == class && l.state == state) + .count(), + ) + .unwrap_or(u32::MAX) + }; + classes.insert( + class, + ClassStatus { + queued_main: entry(Scope::Main).map_or(0, |q| q.queued), + queued_pr: entry(Scope::Pr).map_or(0, |q| q.queued), + oldest_age_secs_main: entry(Scope::Main).map_or(0, |q| q.oldest_age_secs), + oldest_age_secs_pr: entry(Scope::Pr).map_or(0, |q| q.oldest_age_secs), + slots: slots.get(&class).copied().unwrap_or(0), + leases_pending: count(LeaseStatus::Pending), + leases_running: count(LeaseStatus::Running), + }, + ); + } + Ok(Status { + now, + last_tick_at: snap.last_tick_at, + last_error: snap.last_error, + classes, + queue: snap.queue, + drones, + unmet: snap.unmet, + reserved_for: snap.reserved_for, + }) +} + +async fn status(State(queen): State>) -> Result, ApiError> { + Ok(Json( + build_status(&queen, OffsetDateTime::now_utc().unix_timestamp()).await?, + )) +} + +/// Prometheus text exposition. +pub fn render_metrics(status: &Status) -> String { + let mut out = String::new(); + let mut gauge = |name: &str, help: &str, rows: Vec<(String, f64)>| { + let _ = writeln!(out, "# HELP {name} {help}\n# TYPE {name} gauge"); + for (labels, v) in rows { + let _ = if labels.is_empty() { + writeln!(out, "{name} {v}") + } else { + writeln!(out, "{name}{{{labels}}} {v}") + }; + } + }; + let both = |f: &dyn Fn(&ClassStatus, Scope) -> f64| -> Vec<(String, f64)> { + status + .classes + .iter() + .flat_map(|(c, s)| { + Scope::ALL.map(|sc| { + ( + format!("class=\"{c}\",scope=\"{}\"", scope_name(sc)), + f(s, sc), + ) + }) + }) + .collect() + }; + gauge( + "borg_queued_jobs", + "Queued jobs per class and scope.", + both(&|s, sc| { + f64::from(if sc == Scope::Main { + s.queued_main + } else { + s.queued_pr + }) + }), + ); + gauge( + "borg_oldest_queued_age_seconds", + "Age of the oldest queued job per class and scope.", + both(&|s, sc| { + #[allow( + clippy::cast_precision_loss, + reason = "ages are far below 2^52 seconds" + )] + let v = (if sc == Scope::Main { + s.oldest_age_secs_main + } else { + s.oldest_age_secs_pr + }) as f64; + v + }), + ); + gauge( + "borg_class_slots", + "Leases of a class the fleet could hold if empty.", + status + .classes + .iter() + .map(|(c, s)| (format!("class=\"{c}\""), f64::from(s.slots))) + .collect(), + ); + let mut lease_rows = Vec::new(); + for (c, s) in &status.classes { + lease_rows.push(( + format!("class=\"{c}\",state=\"pending\""), + f64::from(s.leases_pending), + )); + lease_rows.push(( + format!("class=\"{c}\",state=\"running\""), + f64::from(s.leases_running), + )); + } + gauge( + "borg_active_leases", + "Active leases per class and state.", + lease_rows, + ); + gauge( + "borg_drones", + "Enrolled drones by freshness.", + vec![ + ( + "state=\"fresh\"".to_owned(), + f64::from( + u32::try_from(status.drones.iter().filter(|d| d.fresh).count()).unwrap_or(0), + ), + ), + ( + "state=\"stale\"".to_owned(), + f64::from( + u32::try_from(status.drones.iter().filter(|d| !d.fresh).count()).unwrap_or(0), + ), + ), + ], + ); + gauge( + "borg_reservations", + "Drones held back for starving demand.", + vec![( + String::new(), + f64::from(u32::try_from(status.reserved_for.len()).unwrap_or(0)), + )], + ); + #[allow( + clippy::cast_precision_loss, + reason = "unix timestamps fit f64 exactly" + )] + gauge( + "borg_last_tick_timestamp_seconds", + "Unix time of the last reconciliation pass.", + vec![(String::new(), status.last_tick_at.unwrap_or(0) as f64)], + ); + gauge( + "borg_last_tick_ok", + "1 if the last reconciliation pass had no error.", + vec![( + String::new(), + if status.last_error.is_none() { + 1.0 + } else { + 0.0 + }, + )], + ); + out +} + +fn scope_name(s: Scope) -> &'static str { + match s { + Scope::Main => "main", + Scope::Pr => "pr", + } +} + +async fn metrics(State(queen): State>) -> Result { + let status = build_status(&queen, OffsetDateTime::now_utc().unix_timestamp()).await?; + Ok(( + [("content-type", "text/plain; version=0.0.4; charset=utf-8")], + render_metrics(&status), + ) + .into_response()) +} + +/// Run the reconciliation loop and the HTTP server until `shutdown` resolves. +pub async fn serve( + queen: Arc, + listener: tokio::net::TcpListener, + shutdown: impl std::future::Future + Send + 'static, +) -> std::io::Result<()> { + let (stop_tx, mut stop_rx) = tokio::sync::watch::channel(false); + let ticker = { + let queen = Arc::clone(&queen); + tokio::spawn(async move { + let mut every = tokio::time::interval(std::time::Duration::from_secs( + queen.config().poll_interval_secs, + )); + every.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + // Both arms are cancel-safe: `interval.tick` and `changed` hold + // no state that a dropped future could lose. The tick itself + // runs in the arm body, outside the select, so shutdown waits + // for an in-progress pass instead of abandoning a mint. + tokio::select! { + _ = every.tick() => { + let now = OffsetDateTime::now_utc(); + if let Err(e) = queen.tick(now).await { + tracing::warn!("tick failed: {e}"); + queen.record_tick_error(now, &e); + } + } + _ = stop_rx.changed() => break, + } + } + }) + }; + let served = axum::serve(listener, router(queen)) + .with_graceful_shutdown(async move { + shutdown.await; + let _ = stop_tx.send(true); + }) + .await; + let _ = ticker.await; + served +} diff --git a/src/classes.rs b/src/classes.rs new file mode 100644 index 0000000..f06f8e0 --- /dev/null +++ b/src/classes.rs @@ -0,0 +1,297 @@ +//! Runner classes: the resources a job of each size needs. +//! +//! Jobs ask for a class through its runner label (`` for main +//! and merge-group jobs, `-pr` for pull requests). The prefix +//! and the per-class requirements are configuration with generic defaults. + +use std::fmt; + +use serde::{Deserialize, Serialize}; + +/// Default label prefix shared by every class label. +pub const DEFAULT_LABEL_PREFIX: &str = "borg-ci-"; +/// Suffix marking pull-request demand. +const PR_SUFFIX: &str = "-pr"; + +/// A runner size class. Ordered smallest to largest. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Class { + Light, + Medium, + Heavy, +} + +impl Class { + /// Every class, smallest first. + pub const ALL: [Class; 3] = [Class::Light, Class::Medium, Class::Heavy]; + + pub fn as_str(self) -> &'static str { + match self { + Class::Light => "light", + Class::Medium => "medium", + Class::Heavy => "heavy", + } + } +} + +impl fmt::Display for Class { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +/// Where a queued job comes from. Main jobs (main branch, merge groups) are +/// scheduled before pull-request jobs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Scope { + Main, + Pr, +} + +impl Scope { + pub const ALL: [Scope; 2] = [Scope::Main, Scope::Pr]; +} + +/// Resources reserved for one running job of a class. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClassSpec { + pub mem_mib: u64, + pub vcpu: u32, + /// The host must have a usable Docker daemon. + pub docker: bool, +} + +/// Requirements for all classes. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct Classes { + /// Prefix of every class label, e.g. `example-ci-` gives `example-ci-light`. + pub label_prefix: String, + pub light: ClassSpec, + pub medium: ClassSpec, + pub heavy: ClassSpec, +} + +impl Default for Classes { + fn default() -> Self { + Self { + label_prefix: DEFAULT_LABEL_PREFIX.to_owned(), + // 2.5 GiB p90, 2 vCPU. + light: ClassSpec { + mem_mib: 2560, + vcpu: 2, + docker: false, + }, + // 12 GiB, Docker. + medium: ClassSpec { + mem_mib: 12 * 1024, + vcpu: 4, + docker: true, + }, + // 24 GiB, Docker. + heavy: ClassSpec { + mem_mib: 24 * 1024, + vcpu: 8, + docker: true, + }, + } + } +} + +/// Invalid class configuration. +#[derive(Debug, thiserror::Error)] +pub enum ClassesError { + #[error("parse classes config: {0}")] + Parse(#[from] serde_json::Error), + #[error("label prefix must not be empty")] + EmptyPrefix, + #[error("class {class} must reserve at least 1 MiB and 1 vCPU")] + Empty { class: Class }, + #[error("class {larger} must not reserve less than {smaller}")] + NotMonotonic { smaller: Class, larger: Class }, +} + +impl Classes { + pub fn spec(&self, class: Class) -> &ClassSpec { + match class { + Class::Light => &self.light, + Class::Medium => &self.medium, + Class::Heavy => &self.heavy, + } + } + + /// Parse a JSON config. Omitted classes keep their defaults. + pub fn from_json(json: &str) -> Result { + let classes: Self = serde_json::from_str(json)?; + classes.validate()?; + Ok(classes) + } + + /// Reject configs the allocator cannot reason about: zero-sized classes + /// and classes that shrink as they grow. + pub fn validate(&self) -> Result<(), ClassesError> { + if self.label_prefix.is_empty() { + return Err(ClassesError::EmptyPrefix); + } + for class in Class::ALL { + let spec = self.spec(class); + if spec.mem_mib == 0 || spec.vcpu == 0 { + return Err(ClassesError::Empty { class }); + } + } + for pair in Class::ALL.windows(2) { + let (smaller, larger) = (pair[0], pair[1]); + let (s, l) = (self.spec(smaller), self.spec(larger)); + if l.mem_mib < s.mem_mib || l.vcpu < s.vcpu || (s.docker && !l.docker) { + return Err(ClassesError::NotMonotonic { smaller, larger }); + } + } + Ok(()) + } + + /// The label a job requests for `class` and `scope`. + pub fn job_label(&self, class: Class, scope: Scope) -> String { + match scope { + Scope::Main => format!("{}{class}", self.label_prefix), + Scope::Pr => format!("{}{class}{PR_SUFFIX}", self.label_prefix), + } + } + + /// Parse a job's `runs-on` label into its class and scope. + pub fn parse_label(&self, label: &str) -> Option<(Class, Scope)> { + let rest = label.strip_prefix(self.label_prefix.as_str())?; + let (name, scope) = match rest.strip_suffix(PR_SUFFIX) { + Some(name) => (name, Scope::Pr), + None => (rest, Scope::Main), + }; + let class = Class::ALL.into_iter().find(|c| c.as_str() == name)?; + Some((class, scope)) + } + + /// The exact label set a runner is minted with. + /// + /// Main grants carry only the bare label so GitHub can hand them nothing + /// but main jobs. Pull-request grants carry both labels. + pub fn runner_labels(&self, class: Class, scope: Scope) -> Vec { + match scope { + Scope::Main => vec![self.job_label(class, Scope::Main)], + Scope::Pr => vec![ + self.job_label(class, Scope::Main), + self.job_label(class, Scope::Pr), + ], + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn defaults_match_design_and_validate() { + let c = Classes::default(); + assert_eq!(c.light.mem_mib, 2560); + assert_eq!(c.light.vcpu, 2); + assert!(!c.light.docker); + assert_eq!(c.medium.mem_mib, 12 * 1024); + assert!(c.medium.docker); + assert_eq!(c.heavy.mem_mib, 24 * 1024); + assert!(c.heavy.docker); + assert!(c.validate().is_ok()); + } + + #[test] + fn labels_round_trip() { + let c = Classes::default(); + for class in Class::ALL { + for scope in Scope::ALL { + assert_eq!( + c.parse_label(&c.job_label(class, scope)), + Some((class, scope)) + ); + } + } + assert_eq!(c.job_label(Class::Heavy, Scope::Pr), "borg-ci-heavy-pr"); + assert_eq!(c.job_label(Class::Light, Scope::Main), "borg-ci-light"); + } + + #[test] + fn prefix_is_configurable() { + let c = Classes::from_json(r#"{"label_prefix":"example-ci-"}"#).unwrap(); + assert_eq!( + c.job_label(Class::Medium, Scope::Pr), + "example-ci-medium-pr" + ); + assert_eq!( + c.parse_label("example-ci-medium"), + Some((Class::Medium, Scope::Main)) + ); + assert_eq!(c.parse_label("borg-ci-medium"), None); + assert!(matches!( + Classes::from_json(r#"{"label_prefix":""}"#), + Err(ClassesError::EmptyPrefix) + )); + } + + #[test] + fn unknown_labels_do_not_parse() { + let c = Classes::default(); + for label in [ + "ubuntu-latest", + "borg-ci-", + "borg-ci-huge", + "borg-ci-light-pr-pr", + "borg-ci-pr", + ] { + assert_eq!(c.parse_label(label), None, "{label}"); + } + } + + #[test] + fn main_runner_has_only_the_bare_label() { + assert_eq!( + Classes::default().runner_labels(Class::Medium, Scope::Main), + ["borg-ci-medium"] + ); + } + + #[test] + fn pr_runner_has_both_labels() { + assert_eq!( + Classes::default().runner_labels(Class::Medium, Scope::Pr), + ["borg-ci-medium", "borg-ci-medium-pr"] + ); + } + + #[test] + fn json_overrides_one_class_and_keeps_defaults() { + let c = + Classes::from_json(r#"{"heavy":{"mem_mib":20480,"vcpu":6,"docker":true}}"#).unwrap(); + assert_eq!(c.heavy.mem_mib, 20480); + assert_eq!(c.light, Classes::default().light); + } + + #[test] + fn json_rejects_unknown_fields_and_bad_shapes() { + assert!(matches!( + Classes::from_json(r#"{"huge":{}}"#), + Err(ClassesError::Parse(_)) + )); + assert!(matches!( + Classes::from_json(r#"{"light":{"mem_mib":0,"vcpu":2,"docker":false}}"#), + Err(ClassesError::Empty { + class: Class::Light + }) + )); + assert!(matches!( + Classes::from_json(r#"{"heavy":{"mem_mib":1024,"vcpu":8,"docker":true}}"#), + Err(ClassesError::NotMonotonic { + smaller: Class::Medium, + larger: Class::Heavy + }) + )); + } +} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..fa0d270 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,196 @@ +//! Queen configuration, read from a JSON file. Nothing org-specific is baked +//! in: the organization, repositories, label prefix and class sizes all come +//! from here. + +use std::net::SocketAddr; +use std::path::{Path, PathBuf}; + +use serde::Deserialize; + +use crate::allocator::{WarmFloor, default_warm_floors}; +use crate::classes::{Classes, ClassesError}; + +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("read {path}: {source}")] + Read { + path: PathBuf, + source: std::io::Error, + }, + #[error("parse config: {0}")] + Parse(#[from] serde_json::Error), + #[error("invalid classes: {0}")] + Classes(#[from] ClassesError), + #[error("{0}")] + Invalid(String), +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct GithubSettings { + /// Organization that owns the self-hosted runners. + pub org: String, + /// GitHub App id (or client id). + pub app_id: String, + /// PEM private key file. If absent, `BORG_APP_PRIVATE_KEY` must hold the + /// PEM text. + #[serde(default)] + pub private_key_path: Option, + #[serde(default = "default_api_base")] + pub api_base: String, + #[serde(default = "default_runner_group")] + pub runner_group_id: u64, +} + +fn default_api_base() -> String { + "https://api.github.com".to_owned() +} +fn default_runner_group() -> u64 { + 1 +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct QueenConfig { + pub github: GithubSettings, + /// Repositories (`owner/name`) whose queued jobs are served. + pub repos: Vec, + #[serde(default = "default_listen")] + pub listen: SocketAddr, + pub db_path: PathBuf, + #[serde(default = "default_runner_prefix")] + pub runner_prefix: String, + #[serde(default = "default_work_folder")] + pub work_folder: String, + #[serde(default = "ten")] + pub poll_interval_secs: u64, + /// A lease whose runner has not come online within this long is expired. + #[serde(default = "default_lease_ttl")] + pub lease_ttl_secs: i64, + /// Hard ceiling on a lease's lifetime. + #[serde(default = "default_lease_max")] + pub lease_max_secs: i64, + /// A drone silent for longer than this gets no new leases. + #[serde(default = "default_drone_stale")] + pub drone_stale_secs: i64, + /// Demand older than this triggers the starvation guard; null disables. + #[serde(default = "default_starvation")] + pub starvation_after_secs: Option, + #[serde(default)] + pub classes: Classes, + #[serde(default = "default_warm_floors")] + pub warm_floors: Vec, +} + +fn default_listen() -> SocketAddr { + SocketAddr::from(([127, 0, 0, 1], 8080)) +} +fn default_runner_prefix() -> String { + "borg".to_owned() +} +fn default_work_folder() -> String { + "_work".to_owned() +} +fn ten() -> u64 { + 10 +} +fn default_lease_ttl() -> i64 { + 300 +} +fn default_lease_max() -> i64 { + 6 * 3600 +} +fn default_drone_stale() -> i64 { + 60 +} +#[allow( + clippy::unnecessary_wraps, + reason = "serde default for an Option field" +)] +fn default_starvation() -> Option { + Some(120) +} + +impl QueenConfig { + pub fn load(path: &Path) -> Result { + let text = std::fs::read_to_string(path).map_err(|source| ConfigError::Read { + path: path.to_owned(), + source, + })?; + Self::from_json(&text) + } + + pub fn from_json(json: &str) -> Result { + let cfg: Self = serde_json::from_str(json)?; + cfg.validate()?; + Ok(cfg) + } + + fn validate(&self) -> Result<(), ConfigError> { + self.classes.validate()?; + let bad = |msg: &str| Err(ConfigError::Invalid(msg.to_owned())); + if self.github.org.is_empty() || self.github.app_id.is_empty() { + return bad("github.org and github.app_id are required"); + } + if self.repos.is_empty() { + return bad("repos must name at least one owner/name repository"); + } + if self + .repos + .iter() + .any(|r| r.split('/').filter(|p| !p.is_empty()).count() != 2) + { + return bad("every repos entry must look like owner/name"); + } + let prefix_ok = !self.runner_prefix.is_empty() + && self + .runner_prefix + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'-'); + if !prefix_ok { + return bad("runner_prefix must be non-empty alphanumerics and dashes"); + } + if self.poll_interval_secs == 0 || self.lease_ttl_secs <= 0 || self.lease_max_secs <= 0 { + return bad("poll_interval_secs, lease_ttl_secs and lease_max_secs must be positive"); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const MIN: &str = r#"{"github":{"org":"example-org","app_id":"1"},"repos":["example-org/app"],"db_path":"/tmp/q.db"}"#; + + #[test] + fn minimal_config_gets_defaults() { + let c = QueenConfig::from_json(MIN).unwrap(); + assert_eq!(c.poll_interval_secs, 10); + assert_eq!(c.starvation_after_secs, Some(120)); + assert_eq!(c.github.api_base, "https://api.github.com"); + assert_eq!(c.warm_floors, default_warm_floors()); + assert_eq!(c.listen.ip().to_string(), "127.0.0.1"); + } + + #[test] + fn rejects_bad_repos_unknown_fields_and_empty_org() { + for bad in [ + MIN.replace("example-org/app", "app"), + MIN.replace("\"repos\"", "\"nope\":1,\"repos\""), + MIN.replace("\"org\":\"example-org\"", "\"org\":\"\""), + MIN.replace("\"repos\":[\"example-org/app\"]", "\"repos\":[]"), + ] { + assert!(QueenConfig::from_json(&bad).is_err(), "{bad}"); + } + } + + #[test] + fn starvation_can_be_disabled() { + let c = QueenConfig::from_json( + &MIN.replace("\"repos\"", "\"starvation_after_secs\":null,\"repos\""), + ) + .unwrap(); + assert_eq!(c.starvation_after_secs, None); + } +} diff --git a/src/drone/cleanup.rs b/src/drone/cleanup.rs new file mode 100644 index 0000000..26cac9c --- /dev/null +++ b/src/drone/cleanup.rs @@ -0,0 +1,489 @@ +//! Unconditional per-lease cleanup: the work directory (with its `_diag`) and +//! the job's docker containers, networks and volumes. Tool caches live outside +//! the work root and are never touched. +//! +//! Docker removal is conservative because other workloads share some hosts: a +//! resource is removed only after `docker inspect` confirms it belongs to the +//! lease (label or exact name prefix). There is no global prune. + +use std::collections::BTreeMap; +use std::io; +use std::path::Path; + +use crate::drone::exec::Exec; + +/// Label that marks a resource as belonging to a lease. +pub const LEASE_LABEL: &str = "borg.lease"; +const COMPOSE_PROJECT_LABEL: &str = "com.docker.compose.project"; +const PREFIX: &str = "borg-"; + +/// Lease ids come from the queen; accept only a filesystem- and docker-safe +/// alphabet before using one in a path or a docker filter. +pub fn valid_lease_id(id: &str) -> bool { + !id.is_empty() + && id.len() <= 64 + && id + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_') +} + +/// The name prefix jobs should give their docker resources (also exported as +/// `COMPOSE_PROJECT_NAME`). +pub fn docker_prefix(lease_id: &str) -> String { + format!("{PREFIX}{lease_id}") +} + +/// Which resources a cleanup may remove. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Scope<'a> { + /// Resources of one lease. + Lease(&'a str), + /// Resources of any borg lease (startup sweep after a crash). + AnyLease, +} + +fn is_hex12(s: &str) -> bool { + s.len() == 12 + && s.bytes() + .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) +} + +/// `name` is `prefix` itself or `prefix` followed by `-` or `_`. +fn name_has_prefix(name: &str, prefix: &str) -> bool { + name.strip_prefix(prefix) + .is_some_and(|rest| rest.is_empty() || rest.starts_with('-') || rest.starts_with('_')) +} + +fn belongs(scope: Scope<'_>, name: &str, labels: &BTreeMap) -> bool { + match scope { + Scope::Lease(id) => { + let prefix = docker_prefix(id); + labels.get(LEASE_LABEL).is_some_and(|v| v == id) + || labels + .get(COMPOSE_PROJECT_LABEL) + .is_some_and(|p| *p == prefix) + || name_has_prefix(name, &prefix) + } + Scope::AnyLease => { + let project_ok = labels + .get(COMPOSE_PROJECT_LABEL) + .and_then(|p| p.strip_prefix(PREFIX)) + .is_some_and(is_hex12); + let name_ok = name.strip_prefix(PREFIX).is_some_and(|rest| { + rest.get(..12).is_some_and(is_hex12) + && rest + .get(12..) + .is_some_and(|tail| tail.is_empty() || tail.starts_with(['-', '_'])) + }); + labels.contains_key(LEASE_LABEL) || project_ok || name_ok + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Kind { + Container, + Network, + Volume, +} + +impl Kind { + fn list_args(self) -> &'static [&'static str] { + match self { + Self::Container => &["ps", "-aq"], + Self::Network => &["network", "ls", "-q", "--no-trunc"], + Self::Volume => &["volume", "ls", "-q"], + } + } + + fn inspect_args(self) -> (&'static [&'static str], &'static str) { + match self { + Self::Container => (&["inspect"], "{{.Name}}\t{{json .Config.Labels}}"), + Self::Network => (&["network", "inspect"], "{{.Name}}\t{{json .Labels}}"), + Self::Volume => (&["volume", "inspect"], "{{.Name}}\t{{json .Labels}}"), + } + } + + fn remove_args(self) -> &'static [&'static str] { + match self { + Self::Container => &["rm", "-f", "-v"], + Self::Network => &["network", "rm"], + Self::Volume => &["volume", "rm", "-f"], + } + } +} + +/// What a docker cleanup did. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct DockerReport { + pub removed: Vec, + /// Listed by a filter but not provably ours, so left alone. + pub skipped: Vec, + pub errors: Vec, +} + +fn filters(scope: Scope<'_>) -> Vec { + match scope { + Scope::Lease(id) => { + let prefix = docker_prefix(id); + vec![ + format!("label={LEASE_LABEL}={id}"), + format!("label={COMPOSE_PROJECT_LABEL}={prefix}"), + format!("name={prefix}"), + ] + } + Scope::AnyLease => vec![format!("label={LEASE_LABEL}"), format!("name={PREFIX}")], + } +} + +fn parse_inspect(line: &str) -> Option<(String, BTreeMap)> { + let (name, labels) = line.split_once('\t')?; + let labels: Option> = serde_json::from_str(labels.trim()).ok()?; + Some(( + name.trim_start_matches('/').to_owned(), + labels.unwrap_or_default(), + )) +} + +/// Remove the docker containers, networks and volumes of `scope`. +pub fn cleanup_docker(exec: &dyn Exec, docker: &str, scope: Scope<'_>) -> DockerReport { + let mut report = DockerReport::default(); + if let Scope::Lease(id) = scope + && !valid_lease_id(id) + { + report + .errors + .push(format!("refusing unsafe lease id {id:?}")); + return report; + } + for kind in [Kind::Container, Kind::Network, Kind::Volume] { + let mut candidates: Vec = Vec::new(); + for filter in filters(scope) { + let mut args: Vec<&str> = kind.list_args().to_vec(); + args.extend(["--filter", &filter]); + match exec.run(docker, &args) { + Ok(out) if out.success => { + for id in out.stdout.lines().map(str::trim).filter(|l| !l.is_empty()) { + if !candidates.iter().any(|c| c == id) { + candidates.push(id.to_owned()); + } + } + } + Ok(_) => {} + Err(e) => { + report.errors.push(format!("{docker}: {e}")); + return report; + } + } + } + for id in candidates { + let (inspect, format) = kind.inspect_args(); + let mut args: Vec<&str> = inspect.to_vec(); + args.extend(["--format", format, &id]); + let owned = match exec.run(docker, &args) { + Ok(out) if out.success => { + parse_inspect(&out.stdout).is_some_and(|(n, l)| belongs(scope, &n, &l)) + } + _ => false, + }; + if !owned { + report.skipped.push(id); + continue; + } + let mut args: Vec<&str> = kind.remove_args().to_vec(); + args.push(&id); + match exec.run(docker, &args) { + Ok(out) if out.success => report.removed.push(id), + Ok(_) => report.errors.push(format!("failed to remove {id}")), + Err(e) => report.errors.push(format!("{docker}: {e}")), + } + } + } + report +} + +/// Delete `/` (a symlink at that path is unlinked, never +/// followed). Returns whether anything was removed. +pub fn remove_lease_dir(work_root: &Path, lease_id: &str) -> io::Result { + if !valid_lease_id(lease_id) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "unsafe lease id", + )); + } + let path = work_root.join(lease_id); + match std::fs::symlink_metadata(&path) { + Ok(meta) if meta.is_dir() => std::fs::remove_dir_all(&path).map(|()| true), + Ok(_) => std::fs::remove_file(&path).map(|()| true), + Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(false), + Err(e) => Err(e), + } +} + +/// Delete every lease directory under `work_root` (startup: the drone holds +/// no leases yet). Returns the removed names. +pub fn sweep_work_root(work_root: &Path) -> io::Result> { + let mut removed = Vec::new(); + let entries = match std::fs::read_dir(work_root) { + Ok(e) => e, + Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(removed), + Err(e) => return Err(e), + }; + for entry in entries { + let name = entry?.file_name().to_string_lossy().into_owned(); + if remove_lease_dir(work_root, &name)? { + removed.push(name); + } + } + Ok(removed) +} + +#[cfg(test)] +pub(crate) mod testing { + use std::collections::HashMap; + use std::io; + use std::sync::Mutex; + + use crate::drone::exec::{Exec, Output}; + + /// Scripted commands: key is the program plus space-joined args; any key + /// not present is a successful command with empty output. + #[derive(Default)] + pub struct FakeExec { + pub script: Mutex>, + pub calls: Mutex>, + } + + impl FakeExec { + pub fn on(&self, key: &str, stdout: &str) { + self.script.lock().unwrap().insert( + key.to_owned(), + Output { + success: true, + stdout: stdout.to_owned(), + }, + ); + } + pub fn fail(&self, key: &str) { + self.script.lock().unwrap().insert( + key.to_owned(), + Output { + success: false, + stdout: String::new(), + }, + ); + } + pub fn calls(&self) -> Vec { + self.calls.lock().unwrap().clone() + } + } + + impl Exec for FakeExec { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + let key = format!("{program} {}", args.join(" ")); + self.calls.lock().unwrap().push(key.clone()); + Ok(self + .script + .lock() + .unwrap() + .get(&key) + .cloned() + .unwrap_or(Output { + success: true, + stdout: String::new(), + })) + } + } +} + +#[cfg(test)] +mod tests { + use super::testing::FakeExec; + use super::*; + + const ID: &str = "0123456789ab"; + + fn inspect_key(prefix: &str, id: &str) -> String { + format!("docker {prefix}--format {{{{.Name}}}}\t{{{{json .Config.Labels}}}} {id}") + } + + #[test] + fn removes_only_containers_that_provably_belong_to_the_lease() { + let x = FakeExec::default(); + // The name filter is a substring match, so it also lists lookalikes. + x.on( + &format!("docker ps -aq --filter name=borg-{ID}"), + "c-ours\nc-lookalike\nc-longer-id\nc-other-label\nc-broken", + ); + x.on( + &inspect_key("inspect ", "c-ours"), + &format!("/borg-{ID}-db-1\tnull"), + ); + x.on( + &inspect_key("inspect ", "c-lookalike"), + &format!("/xborg-{ID}-db-1\t{{}}"), + ); + x.on( + &inspect_key("inspect ", "c-longer-id"), + &format!("/borg-{ID}9-db-1\t{{}}"), + ); + x.on( + &inspect_key("inspect ", "c-other-label"), + "/dev-db\t{\"borg.lease\":\"fedcba987654\"}", + ); + x.on(&inspect_key("inspect ", "c-broken"), "not json"); + let r = cleanup_docker(&x, "docker", Scope::Lease(ID)); + assert_eq!(r.removed, ["c-ours"]); + assert_eq!( + r.skipped, + ["c-lookalike", "c-longer-id", "c-other-label", "c-broken"] + ); + let calls = x.calls(); + assert!(calls.contains(&"docker rm -f -v c-ours".to_owned())); + assert!( + !calls + .iter() + .any(|c| c.contains("rm -f -v c-lookalike") || c.contains("prune")), + "{calls:?}" + ); + } + + #[test] + fn label_and_compose_project_matches_are_removed_even_with_other_names() { + let x = FakeExec::default(); + x.on( + &format!("docker ps -aq --filter label=borg.lease={ID}"), + "c1", + ); + x.on( + &format!("docker ps -aq --filter label=com.docker.compose.project=borg-{ID}"), + "c2", + ); + x.on( + &inspect_key("inspect ", "c1"), + &format!("/svc\t{{\"borg.lease\":\"{ID}\"}}"), + ); + x.on( + &inspect_key("inspect ", "c2"), + &format!("/svc2\t{{\"com.docker.compose.project\":\"borg-{ID}\"}}"), + ); + assert_eq!( + cleanup_docker(&x, "docker", Scope::Lease(ID)).removed, + ["c1", "c2"] + ); + } + + #[test] + fn networks_and_volumes_are_removed_with_their_own_commands() { + let x = FakeExec::default(); + x.on( + &format!("docker network ls -q --no-trunc --filter name=borg-{ID}"), + "n1", + ); + x.on( + &format!("docker volume ls -q --filter name=borg-{ID}"), + "v1", + ); + let net = "docker network inspect --format {{.Name}}\t{{json .Labels}} n1"; + let vol = "docker volume inspect --format {{.Name}}\t{{json .Labels}} v1"; + x.on(net, &format!("borg-{ID}_default\tnull")); + x.on(vol, &format!("borg-{ID}_data\t{{}}")); + let r = cleanup_docker(&x, "docker", Scope::Lease(ID)); + assert_eq!(r.removed, ["n1", "v1"]); + let calls = x.calls(); + assert!(calls.contains(&"docker network rm n1".to_owned())); + assert!(calls.contains(&"docker volume rm -f v1".to_owned())); + } + + #[test] + fn unsafe_lease_ids_make_no_docker_calls() { + let x = FakeExec::default(); + let r = cleanup_docker(&x, "docker", Scope::Lease("a b; rm")); + assert!(x.calls().is_empty()); + assert_eq!(r.errors.len(), 1); + } + + #[test] + fn startup_sweep_matches_any_borg_lease_but_not_other_workloads() { + let x = FakeExec::default(); + x.on("docker ps -aq --filter label=borg.lease", "a\nb"); + x.on("docker ps -aq --filter name=borg-", "b\nc\nd"); + x.on( + &inspect_key("inspect ", "a"), + "/job-a\t{\"borg.lease\":\"0123456789ab\"}", + ); + x.on( + &inspect_key("inspect ", "b"), + "/borg-aaaaaaaaaaaa-web-1\t{}", + ); + x.on(&inspect_key("inspect ", "c"), "/borg-queen-backup\t{}"); + x.on(&inspect_key("inspect ", "d"), "/borg-aaaaaaaaaaaaZ\tnull"); + let r = cleanup_docker(&x, "docker", Scope::AnyLease); + assert_eq!(r.removed, ["a", "b"]); + assert_eq!(r.skipped, ["c", "d"]); + } + + #[test] + fn docker_errors_are_reported_not_fatal() { + let x = FakeExec::default(); + x.on(&format!("docker ps -aq --filter name=borg-{ID}"), "c1"); + x.on(&inspect_key("inspect ", "c1"), &format!("/borg-{ID}\tnull")); + x.fail("docker rm -f -v c1"); + let r = cleanup_docker(&x, "docker", Scope::Lease(ID)); + assert!(r.removed.is_empty()); + assert_eq!(r.errors, ["failed to remove c1"]); + } + + #[test] + fn lease_dir_removal_is_contained_and_leaves_caches_alone() { + let root = tempfile::tempdir().unwrap(); + let work = root.path().join("work"); + let cache = root.path().join("cache/sccache"); + std::fs::create_dir_all(work.join(ID).join("_diag")).unwrap(); + std::fs::write(work.join(ID).join("_diag/log"), "x").unwrap(); + std::fs::create_dir_all(&cache).unwrap(); + std::fs::write(cache.join("entry"), "x").unwrap(); + std::fs::create_dir_all(work.join("keep-me")).unwrap(); + + assert!(remove_lease_dir(&work, ID).unwrap()); + assert!(!work.join(ID).exists()); + assert!( + !remove_lease_dir(&work, ID).unwrap(), + "second removal is a no-op" + ); + assert!(cache.join("entry").exists()); + assert!(work.join("keep-me").exists()); + assert!(remove_lease_dir(&work, "../cache").is_err()); + assert!(remove_lease_dir(&work, "").is_err()); + } + + #[cfg(unix)] + #[test] + fn a_symlinked_lease_dir_is_unlinked_not_followed() { + let root = tempfile::tempdir().unwrap(); + let work = root.path().join("work"); + let target = root.path().join("precious"); + std::fs::create_dir_all(&work).unwrap(); + std::fs::create_dir_all(&target).unwrap(); + std::fs::write(target.join("f"), "x").unwrap(); + std::os::unix::fs::symlink(&target, work.join(ID)).unwrap(); + assert!(remove_lease_dir(&work, ID).unwrap()); + assert!(target.join("f").exists()); + } + + #[test] + fn sweep_removes_every_lease_dir() { + let root = tempfile::tempdir().unwrap(); + for id in ["aaa", "bbb"] { + std::fs::create_dir_all(root.path().join(id).join("x")).unwrap(); + } + let mut removed = sweep_work_root(root.path()).unwrap(); + removed.sort(); + assert_eq!(removed, ["aaa", "bbb"]); + assert!( + sweep_work_root(&root.path().join("missing")) + .unwrap() + .is_empty() + ); + } +} diff --git a/src/drone/client.rs b/src/drone/client.rs new file mode 100644 index 0000000..111fa76 --- /dev/null +++ b/src/drone/client.rs @@ -0,0 +1,97 @@ +//! The drone's HTTP client for the queen API. + +use std::time::Duration; + +use crate::api::{HeartbeatRequest, HeartbeatResponse, ResultRequest}; +use crate::drone::config::DroneConfig; +use crate::secret::Secret; + +#[derive(Debug, thiserror::Error)] +pub enum ClientError { + #[error("request failed: {0}")] + Http(#[from] reqwest::Error), + #[error("queen returned {status} for {path}")] + Status { path: String, status: u16 }, + #[error("read CA certificate {path}: {source}")] + CaRead { + path: String, + source: std::io::Error, + }, + #[error("invalid CA certificate: {0}")] + CaParse(reqwest::Error), +} + +pub struct QueenClient { + http: reqwest::Client, + base: String, + token: Secret, +} + +impl QueenClient { + /// With `ca_cert_path` set, that CA is the only trust root (system roots + /// are disabled), which pins the queen's CA. + pub fn new(cfg: &DroneConfig, token: Secret) -> Result { + let mut builder = reqwest::Client::builder() + .user_agent("borg-drone") + .timeout(Duration::from_secs(30)); + if let Some(path) = &cfg.ca_cert_path { + let pem = std::fs::read(path).map_err(|source| ClientError::CaRead { + path: path.display().to_string(), + source, + })?; + let cert = reqwest::Certificate::from_pem(&pem).map_err(ClientError::CaParse)?; + builder = builder + .tls_built_in_root_certs(false) + .add_root_certificate(cert); + } + Ok(Self { + http: builder.build()?, + base: cfg.queen_url.trim_end_matches('/').to_owned(), + token, + }) + } + + pub async fn heartbeat( + &self, + req: &HeartbeatRequest, + ) -> Result { + let path = "/v1/heartbeat"; + let resp = self + .http + .post(format!("{}{path}", self.base)) + .bearer_auth(self.token.expose()) + .json(req) + .send() + .await?; + if !resp.status().is_success() { + return Err(ClientError::Status { + path: path.to_owned(), + status: resp.status().as_u16(), + }); + } + Ok(resp.json().await?) + } + + pub async fn report_result( + &self, + lease_id: &str, + result: &ResultRequest, + ) -> Result<(), ClientError> { + let path = format!("/v1/leases/{lease_id}/result"); + let resp = self + .http + .post(format!("{}{path}", self.base)) + .bearer_auth(self.token.expose()) + .json(result) + .send() + .await?; + if resp.status().is_success() { + Ok(()) + } else { + Err(ClientError::Status { + path, + status: resp.status().as_u16(), + }) + } + } +} diff --git a/src/drone/config.rs b/src/drone/config.rs new file mode 100644 index 0000000..6c42a55 --- /dev/null +++ b/src/drone/config.rs @@ -0,0 +1,201 @@ +//! Drone configuration, a JSON file written by `borg drone install`. + +use std::path::{Path, PathBuf}; + +use serde::{Deserialize, Serialize}; + +use crate::allocator::HostCaps; +use crate::secret::Secret; + +#[derive(Debug, thiserror::Error)] +pub enum DroneConfigError { + #[error("read {path}: {source}")] + Read { + path: PathBuf, + source: std::io::Error, + }, + #[error("parse config: {0}")] + Parse(#[from] serde_json::Error), + #[error("{0}")] + Invalid(String), +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct DroneConfig { + /// Base URL of the queen. Must be `https://` unless `allow_insecure`. + pub queen_url: String, + /// File holding the enrollment token. If absent, `BORG_DRONE_TOKEN`. + #[serde(default)] + pub token_file: Option, + /// PEM CA certificate that signs the queen's certificate. When set it is + /// the only trust root, which pins the queen's CA. + #[serde(default)] + pub ca_cert_path: Option, + /// Permit a plain `http://` queen URL (trusted LAN only). + #[serde(default)] + pub allow_insecure: bool, + /// Pinned actions-runner install; every lease gets a private copy of it. + pub runner_dir: PathBuf, + /// Parent of the per-lease work directories. Wiped on start and after + /// every job. + pub work_root: PathBuf, + /// Tool caches (sccache, mise, cargo). Never touched by cleanup. + pub cache_root: PathBuf, + #[serde(default = "default_docker")] + pub docker_bin: String, + #[serde(default = "default_heartbeat")] + pub heartbeat_secs: u64, + /// A lease running longer than this is killed. + #[serde(default = "default_max_job")] + pub max_job_secs: u64, + /// On SIGTERM, wait this long for in-flight jobs before killing them. + #[serde(default = "default_drain")] + pub drain_timeout_secs: u64, + /// Owner caps ("this PC gives at most 8 cores / 24 GiB"). + #[serde(default)] + pub max_vcpu: Option, + #[serde(default)] + pub max_mem_mib: Option, + #[serde(default)] + pub max_leases: Option, + #[serde(default = "default_reserve")] + pub reserve_mem_mib: u64, +} + +fn default_docker() -> String { + "docker".to_owned() +} +fn default_heartbeat() -> u64 { + 10 +} +fn default_max_job() -> u64 { + 6 * 3600 +} +fn default_drain() -> u64 { + 3600 +} +fn default_reserve() -> u64 { + HostCaps::default().reserve_mem_mib +} + +impl DroneConfig { + pub fn load(path: &Path) -> Result { + let text = std::fs::read_to_string(path).map_err(|source| DroneConfigError::Read { + path: path.to_owned(), + source, + })?; + Self::from_json(&text) + } + + pub fn from_json(json: &str) -> Result { + let cfg: Self = serde_json::from_str(json)?; + cfg.validate()?; + Ok(cfg) + } + + pub fn validate(&self) -> Result<(), DroneConfigError> { + let bad = |m: &str| Err(DroneConfigError::Invalid(m.to_owned())); + let https = self.queen_url.starts_with("https://"); + let http = self.queen_url.starts_with("http://"); + let scheme_ok = https || (http && self.allow_insecure); + if !scheme_ok { + return bad( + "queen_url must start with https:// (set allow_insecure for plain http on a trusted network)", + ); + } + if self.heartbeat_secs == 0 || self.max_job_secs == 0 { + return bad("heartbeat_secs and max_job_secs must be positive"); + } + if self.max_vcpu == Some(0) || self.max_mem_mib == Some(0) || self.max_leases == Some(0) { + return bad("owner caps must be positive when set"); + } + if self.ca_cert_path.is_some() && http { + return bad("ca_cert_path only applies to an https:// queen_url"); + } + Ok(()) + } + + /// The caps advertised to the queen with every heartbeat. + pub fn host_caps(&self) -> HostCaps { + HostCaps { + max_vcpu: self.max_vcpu, + max_mem_mib: self.max_mem_mib, + max_leases: self.max_leases, + reserve_mem_mib: self.reserve_mem_mib, + } + } + + /// The enrollment token from `token_file` or `BORG_DRONE_TOKEN`. + pub fn load_token(&self) -> Result { + let raw = match &self.token_file { + Some(path) => { + std::fs::read_to_string(path).map_err(|source| DroneConfigError::Read { + path: path.clone(), + source, + })? + } + None => std::env::var("BORG_DRONE_TOKEN").map_err(|_| { + DroneConfigError::Invalid("set token_file or BORG_DRONE_TOKEN".to_owned()) + })?, + }; + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err(DroneConfigError::Invalid( + "enrollment token is empty".to_owned(), + )); + } + Ok(Secret::new(trimmed)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const BASE: &str = r#"{"queen_url":"https://queen.example.test:8443","runner_dir":"/r","work_root":"/w","cache_root":"/c"}"#; + + #[test] + fn defaults_and_caps() { + let c = DroneConfig::from_json(BASE).unwrap(); + assert_eq!((c.heartbeat_secs, c.docker_bin.as_str()), (10, "docker")); + let with = BASE.replace('}', r#","max_vcpu":8,"max_mem_mib":24576}"#); + let caps = DroneConfig::from_json(&with).unwrap().host_caps(); + assert_eq!((caps.max_vcpu, caps.max_mem_mib), (Some(8), Some(24576))); + } + + #[test] + fn plain_http_needs_explicit_opt_in() { + let http = BASE.replace("https://", "http://"); + assert!(DroneConfig::from_json(&http).is_err()); + let allowed = http.replace('}', r#","allow_insecure":true}"#); + assert!(DroneConfig::from_json(&allowed).is_ok()); + assert!(DroneConfig::from_json(&BASE.replace("https://queen", "ftp://queen")).is_err()); + } + + #[test] + fn rejects_zero_caps_unknown_fields_and_ca_on_http() { + for extra in [r#","max_vcpu":0"#, r#","bogus":1"#] { + assert!( + DroneConfig::from_json(&BASE.replace('}', &format!("{extra}}}"))).is_err(), + "{extra}" + ); + } + let http_ca = BASE + .replace("https://", "http://") + .replace('}', r#","allow_insecure":true,"ca_cert_path":"/ca.pem"}"#); + assert!(DroneConfig::from_json(&http_ca).is_err()); + } + + #[test] + fn token_file_is_trimmed_and_must_not_be_empty() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("t"); + let mut c = DroneConfig::from_json(BASE).unwrap(); + c.token_file = Some(path.clone()); + std::fs::write(&path, "borg_abc\n").unwrap(); + assert_eq!(c.load_token().unwrap().expose(), "borg_abc"); + std::fs::write(&path, "\n").unwrap(); + assert!(c.load_token().is_err()); + } +} diff --git a/src/drone/exec.rs b/src/drone/exec.rs new file mode 100644 index 0000000..d718c34 --- /dev/null +++ b/src/drone/exec.rs @@ -0,0 +1,42 @@ +//! The drone's single process boundary for short-lived host commands +//! (docker, systemctl, tar, curl, chown ...), behind a trait so cleanup and +//! install logic is tested without touching the machine. + +use std::io; +use std::process::{Command, Stdio}; + +/// Result of a finished command. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Output { + pub success: bool, + pub stdout: String, +} + +pub trait Exec: Send + Sync { + /// Run `program` with `args` to completion. `Err` means it could not be + /// started; a non-zero exit is `Ok` with `success == false`. + fn run(&self, program: &str, args: &[&str]) -> io::Result; +} + +/// Runs commands on the local machine. Programs and arguments are chosen by +/// this crate (and the operator's config for the docker binary), never taken +/// from the queen. +pub struct SystemExec; + +impl Exec for SystemExec { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + #[allow( + clippy::disallowed_methods, + reason = "drone host commands: fixed programs and arguments built by borg, never queen-supplied" + )] + let out = Command::new(program) + .args(args) + .stdin(Stdio::null()) + .stderr(Stdio::null()) + .output()?; + Ok(Output { + success: out.status.success(), + stdout: String::from_utf8_lossy(&out.stdout).trim().to_owned(), + }) + } +} diff --git a/src/drone/install.rs b/src/drone/install.rs new file mode 100644 index 0000000..634c279 --- /dev/null +++ b/src/drone/install.rs @@ -0,0 +1,956 @@ +//! `borg drone install`: idempotently set a host up as a drone and report +//! exactly what changed. +//! +//! Linux (and WSL, which needs systemd enabled) gets a dedicated service user +//! in the docker group and a systemd unit; macOS gets a LaunchDaemon. Both get +//! a copy of this binary, the pinned actions-runner (checksum verified), a +//! config file and the enrollment token. Re-running changes only what differs +//! and restarts the service only if something changed. + +use std::fmt::Write as _; +use std::fs; +use std::io; +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; + +use sha2::{Digest, Sha256}; + +use crate::drone::config::DroneConfig; +use crate::drone::exec::Exec; +use crate::drone::pins::RunnerPin; +use crate::secret::Secret; + +const SERVICE: &str = "borg-drone.service"; +const LAUNCHD_LABEL: &str = "org.borg.drone"; +const STOP_GRACE_SECS: u64 = 120; + +#[derive(Debug, thiserror::Error)] +pub enum InstallError { + #[error("{0}")] + Preflight(String), + #[error("{path}: {source}")] + Io { path: PathBuf, source: io::Error }, + #[error("`{command}` failed")] + Command { command: String }, + #[error("checksum mismatch for {file}: expected {expected}, got {actual}")] + Checksum { + file: String, + expected: String, + actual: String, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Platform { + Linux, + MacOs, +} + +impl Platform { + pub fn current() -> Option { + match std::env::consts::OS { + "linux" => Some(Self::Linux), + "macos" => Some(Self::MacOs), + _ => None, + } + } +} + +pub struct InstallOptions { + pub queen_url: String, + pub token: Secret, + pub ca_cert: Option, + pub allow_insecure: bool, + pub max_vcpu: Option, + pub max_mem_mib: Option, + pub max_leases: Option, + /// Service account the drone and its jobs run as. + pub run_as: String, + /// Filesystem root the install is written under (`/` on a real host). + pub root: PathBuf, + pub platform: Platform, + pub pin: RunnerPin, + /// The borg binary to install. + pub binary: PathBuf, + pub dry_run: bool, + /// Skip the "must be root" check (tests, `--root` staging). + pub skip_root_check: bool, + pub drain_timeout_secs: u64, +} + +/// Where everything lives, as paths on the real host. +struct Layout { + bin: PathBuf, + config_dir: PathBuf, + data: PathBuf, + service_file: PathBuf, +} + +impl Layout { + fn new(platform: Platform) -> Self { + match platform { + Platform::Linux => Self { + bin: "/usr/local/bin/borg".into(), + config_dir: "/etc/borg".into(), + data: "/var/lib/borg".into(), + service_file: format!("/etc/systemd/system/{SERVICE}").into(), + }, + Platform::MacOs => Self { + bin: "/usr/local/bin/borg".into(), + config_dir: "/usr/local/etc/borg".into(), + data: "/usr/local/var/borg".into(), + service_file: format!("/Library/LaunchDaemons/{LAUNCHD_LABEL}.plist").into(), + }, + } + } + + fn config(&self) -> PathBuf { + self.config_dir.join("drone.json") + } + fn token_file(&self) -> PathBuf { + self.config_dir.join("drone.token") + } + fn ca_cert(&self) -> PathBuf { + self.config_dir.join("queen-ca.pem") + } + fn runner(&self) -> PathBuf { + self.data.join("runner") + } +} + +pub fn render_systemd_unit( + layout_bin: &Path, + config: &Path, + data: &Path, + user: &str, + drain_timeout_secs: u64, +) -> String { + format!( + "[Unit]\n\ + Description=borg CI drone\n\ + After=network-online.target docker.service\n\ + Wants=network-online.target\n\ + \n\ + [Service]\n\ + Type=simple\n\ + User={user}\n\ + Group={user}\n\ + WorkingDirectory={data}\n\ + ExecStart={bin} drone run --config {config}\n\ + Restart=always\n\ + RestartSec=5\n\ + KillSignal=SIGTERM\n\ + # SIGTERM goes to the drone only, so in-flight jobs can finish while it\n\ + # drains; leftovers are killed once the drone exits.\n\ + KillMode=mixed\n\ + TimeoutStopSec={stop}\n\ + \n\ + [Install]\n\ + WantedBy=multi-user.target\n", + bin = layout_bin.display(), + config = config.display(), + data = data.display(), + stop = drain_timeout_secs + STOP_GRACE_SECS, + ) +} + +pub fn render_launchd_plist( + layout_bin: &Path, + config: &Path, + data: &Path, + user: &str, + drain_timeout_secs: u64, +) -> String { + format!( + "\n\ + \n\ + \n\ + \n\ + \tLabel{LAUNCHD_LABEL}\n\ + \tProgramArguments\n\ + \t\n\ + \t\t{bin}\n\ + \t\tdrone\n\ + \t\trun\n\ + \t\t--config\n\ + \t\t{config}\n\ + \t\n\ + \tUserName{user}\n\ + \tWorkingDirectory{data}\n\ + \tRunAtLoad\n\ + \tKeepAlive\n\ + \tExitTimeOut{stop}\n\ + \tStandardOutPath{data}/drone.log\n\ + \tStandardErrorPath{data}/drone.log\n\ + \n\ + \n", + bin = layout_bin.display(), + config = config.display(), + data = data.display(), + stop = drain_timeout_secs + STOP_GRACE_SECS, + ) +} + +struct Installer<'a> { + opts: &'a InstallOptions, + layout: Layout, + exec: &'a dyn Exec, + lines: Vec, + changed: bool, +} + +fn io_err(path: &Path) -> impl FnOnce(io::Error) -> InstallError + '_ { + move |source| InstallError::Io { + path: path.to_owned(), + source, + } +} + +impl Installer<'_> { + fn on_disk(&self, real: &Path) -> PathBuf { + self.opts.root.join(real.strip_prefix("/").unwrap_or(real)) + } + + fn note(&mut self, verb: &str, what: impl std::fmt::Display) { + self.lines.push(format!("{verb} {what}")); + } + + /// Run a host command; a failure is an error unless `tolerate`. + fn cmd(&self, program: &str, args: &[&str]) -> Result { + let out = self + .exec + .run(program, args) + .map_err(|source| InstallError::Io { + path: PathBuf::from(program), + source, + })?; + Ok(out.success) + } + + fn must(&self, program: &str, args: &[&str]) -> Result<(), InstallError> { + if self.cmd(program, args)? { + Ok(()) + } else { + Err(InstallError::Command { + command: format!("{program} {}", args.join(" ")), + }) + } + } + + fn chown(&self, real: &Path, owner: &str, recursive: bool) -> Result<(), InstallError> { + let spec = format!("{owner}:{owner}"); + let path = self.on_disk(real); + let path = path.to_string_lossy(); + if recursive { + self.must("chown", &["-R", &spec, &path]) + } else { + self.must("chown", &[&spec, &path]) + } + } + + fn ensure_dir( + &mut self, + real: &Path, + mode: u32, + owner: Option<&str>, + ) -> Result<(), InstallError> { + let path = self.on_disk(real); + if path.is_dir() { + self.note("unchanged", format_args!("dir {}", real.display())); + return Ok(()); + } + self.changed = true; + if self.opts.dry_run { + self.note("would create", format_args!("dir {}", real.display())); + return Ok(()); + } + fs::create_dir_all(&path).map_err(io_err(&path))?; + fs::set_permissions(&path, fs::Permissions::from_mode(mode)).map_err(io_err(&path))?; + if let Some(owner) = owner { + self.chown(real, owner, false)?; + } + self.note("created", format_args!("dir {}", real.display())); + Ok(()) + } + + /// Write `content` unless the file already has it. Returns whether it changed. + fn write_file( + &mut self, + real: &Path, + content: &[u8], + mode: u32, + owner: Option<&str>, + ) -> Result { + let path = self.on_disk(real); + let existing = fs::read(&path).ok(); + let same = existing.as_deref() == Some(content) + && fs::metadata(&path).is_ok_and(|m| m.permissions().mode() & 0o7777 == mode); + if same { + self.note("unchanged", real.display()); + return Ok(false); + } + self.changed = true; + let verb = if existing.is_some() { + "updated" + } else { + "created" + }; + if self.opts.dry_run { + self.note( + &format!("would {}", verb.trim_end_matches('d')), + real.display(), + ); + return Ok(true); + } + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(io_err(parent))?; + } + // Replace via a unique temp file so a running binary is never truncated. + let tmp = path.with_extension("borg-new"); + fs::write(&tmp, content).map_err(io_err(&tmp))?; + fs::set_permissions(&tmp, fs::Permissions::from_mode(mode)).map_err(io_err(&tmp))?; + fs::rename(&tmp, &path).map_err(io_err(&path))?; + if let Some(owner) = owner { + self.chown(real, owner, false)?; + } + self.note(verb, real.display()); + Ok(true) + } + + fn preflight(&self) -> Result<(), InstallError> { + if !self.opts.skip_root_check && !self.opts.dry_run { + let out = self + .exec + .run("id", &["-u"]) + .map_err(|source| InstallError::Io { + path: PathBuf::from("id"), + source, + })?; + if out.stdout != "0" { + return Err(InstallError::Preflight( + "run as root (for example with sudo)".to_owned(), + )); + } + } + if self.opts.platform == Platform::Linux + && !self.on_disk(Path::new("/run/systemd/system")).is_dir() + { + let wsl = self + .on_disk(Path::new("/proc/sys/fs/binfmt_misc/WSLInterop")) + .exists() + || fs::read_to_string(self.on_disk(Path::new("/proc/version"))) + .is_ok_and(|v| v.to_ascii_lowercase().contains("microsoft")); + let hint = if wsl { + " On WSL, enable systemd: add `[boot]` and `systemd=true` to /etc/wsl.conf, then run `wsl --shutdown` from Windows and start the distribution again." + } else { + "" + }; + return Err(InstallError::Preflight(format!( + "systemd is not running on this host.{hint}" + ))); + } + Ok(()) + } + + fn ensure_user(&mut self) -> Result<(), InstallError> { + let user = self.opts.run_as.clone(); + let exists = self.cmd("id", &["-u", &user])?; + match self.opts.platform { + Platform::MacOs => { + if !exists { + return Err(InstallError::Preflight(format!( + "user {user} does not exist; create it first or pass --run-as " + ))); + } + self.note("unchanged", format_args!("user {user}")); + } + Platform::Linux => { + if exists { + self.note("unchanged", format_args!("user {user}")); + } else { + self.changed = true; + if self.opts.dry_run { + self.note("would create", format_args!("user {user}")); + } else { + let home = self.layout.data.to_string_lossy().into_owned(); + self.must( + "useradd", + &[ + "--system", + "--user-group", + "--home-dir", + &home, + "--shell", + "/usr/sbin/nologin", + &user, + ], + )?; + self.note("created", format_args!("user {user}")); + } + } + self.ensure_docker_group(&user, exists)?; + } + } + Ok(()) + } + + fn ensure_docker_group(&mut self, user: &str, user_exists: bool) -> Result<(), InstallError> { + if !self.cmd("getent", &["group", "docker"])? { + self.note("warning", "no docker group found; docker jobs fail until Docker is installed and the user joins its group"); + return Ok(()); + } + let member = user_exists + && self + .exec + .run("id", &["-nG", user]) + .is_ok_and(|o| o.success && o.stdout.split_whitespace().any(|g| g == "docker")); + if member { + self.note("unchanged", format_args!("user {user} in group docker")); + } else if self.opts.dry_run { + self.changed = true; + self.note("would add", format_args!("user {user} to group docker")); + } else { + self.changed = true; + self.must("usermod", &["-aG", "docker", user])?; + self.note("added", format_args!("user {user} to group docker")); + } + Ok(()) + } + + fn install_binary(&mut self) -> Result<(), InstallError> { + let bytes = fs::read(&self.opts.binary).map_err(io_err(&self.opts.binary))?; + let bin = self.layout.bin.clone(); + self.write_file(&bin, &bytes, 0o755, None)?; + Ok(()) + } + + fn install_runner(&mut self) -> Result<(), InstallError> { + let marker = format!("{} {}\n", self.opts.pin.version, self.opts.pin.sha256); + let runner = self.layout.runner(); + let marker_path = self.on_disk(&runner.join(".borg-runner-version")); + if fs::read_to_string(&marker_path).is_ok_and(|m| m == marker) { + self.note( + "unchanged", + format_args!( + "actions-runner {} at {}", + self.opts.pin.version, + runner.display() + ), + ); + return Ok(()); + } + self.changed = true; + if self.opts.dry_run { + self.note( + "would install", + format_args!( + "actions-runner {} at {}", + self.opts.pin.version, + runner.display() + ), + ); + return Ok(()); + } + let user = self.opts.run_as.clone(); + let data = self.layout.data.clone(); + let staging = data.join("runner.new"); + let tarball = data.join(self.opts.pin.filename()); + let (staging_fs, tarball_fs) = (self.on_disk(&staging), self.on_disk(&tarball)); + let _ = fs::remove_dir_all(&staging_fs); + fs::create_dir_all(&staging_fs).map_err(io_err(&staging_fs))?; + + let url = self.opts.pin.url(); + self.must( + "curl", + &["-fsSL", "-o", &tarball_fs.to_string_lossy(), &url], + )?; + let actual = hex::encode(Sha256::digest( + fs::read(&tarball_fs).map_err(io_err(&tarball_fs))?, + )); + if actual != self.opts.pin.sha256 { + let _ = fs::remove_file(&tarball_fs); + let _ = fs::remove_dir_all(&staging_fs); + return Err(InstallError::Checksum { + file: self.opts.pin.filename(), + expected: self.opts.pin.sha256.clone(), + actual, + }); + } + self.must( + "tar", + &[ + "-xzf", + &tarball_fs.to_string_lossy(), + "-C", + &staging_fs.to_string_lossy(), + ], + )?; + let _ = fs::remove_file(&tarball_fs); + fs::write(staging_fs.join(".borg-runner-version"), &marker).map_err(io_err(&staging_fs))?; + + let old = self.on_disk(&runner); + if old.exists() { + // The installed tree is read-only; restore write access to delete it. + self.must("chmod", &["-R", "u+w", &old.to_string_lossy()])?; + fs::remove_dir_all(&old).map_err(io_err(&old))?; + } + fs::rename(&staging_fs, &old).map_err(io_err(&old))?; + self.chown(&runner, &user, true)?; + // Leases hard-link these files; make them immutable to the runner. + self.must("chmod", &["-R", "a-w", &old.to_string_lossy()])?; + self.note( + "installed", + format_args!( + "actions-runner {} at {} (sha256 verified)", + self.opts.pin.version, + runner.display() + ), + ); + Ok(()) + } + + fn drone_config(&self) -> Result { + let o = self.opts; + let cfg = DroneConfig { + queen_url: o.queen_url.clone(), + token_file: Some(self.layout.token_file()), + ca_cert_path: o.ca_cert.as_ref().map(|_| self.layout.ca_cert()), + allow_insecure: o.allow_insecure, + runner_dir: self.layout.runner(), + work_root: self.layout.data.join("work"), + cache_root: self.layout.data.join("cache"), + docker_bin: "docker".to_owned(), + heartbeat_secs: 10, + max_job_secs: 6 * 3600, + drain_timeout_secs: o.drain_timeout_secs, + max_vcpu: o.max_vcpu, + max_mem_mib: o.max_mem_mib, + max_leases: o.max_leases, + reserve_mem_mib: crate::allocator::HostCaps::default().reserve_mem_mib, + }; + cfg.validate() + .map_err(|e| InstallError::Preflight(e.to_string()))?; + let mut json = serde_json::to_string_pretty(&cfg) + .map_err(|e| InstallError::Preflight(e.to_string()))?; + json.push('\n'); + Ok(json) + } + + fn install_service(&mut self, config_changed: bool) -> Result<(), InstallError> { + let l = &self.layout; + let (bin, cfg, data, user) = ( + l.bin.clone(), + l.config(), + l.data.clone(), + self.opts.run_as.clone(), + ); + let drain = self.opts.drain_timeout_secs; + let service_file = l.service_file.clone(); + let before = self.changed; + match self.opts.platform { + Platform::Linux => { + let unit = render_systemd_unit(&bin, &cfg, &data, &user, drain); + let unit_changed = self.write_file(&service_file, unit.as_bytes(), 0o644, None)?; + if self.opts.dry_run { + return Ok(()); + } + if unit_changed { + self.must("systemctl", &["daemon-reload"])?; + self.note("ran", "systemctl daemon-reload"); + } + if !self.cmd("systemctl", &["is-enabled", "--quiet", SERVICE])? { + self.must("systemctl", &["enable", SERVICE])?; + self.note("ran", format_args!("systemctl enable {SERVICE}")); + } + if before || self.changed || config_changed || unit_changed { + self.must("systemctl", &["restart", SERVICE])?; + self.note("ran", format_args!("systemctl restart {SERVICE}")); + } else if !self.cmd("systemctl", &["is-active", "--quiet", SERVICE])? { + self.must("systemctl", &["start", SERVICE])?; + self.note("ran", format_args!("systemctl start {SERVICE}")); + } + } + Platform::MacOs => { + let plist = render_launchd_plist(&bin, &cfg, &data, &user, drain); + let plist_changed = + self.write_file(&service_file, plist.as_bytes(), 0o644, None)?; + if self.opts.dry_run { + return Ok(()); + } + let target = format!("system/{LAUNCHD_LABEL}"); + let path = self.on_disk(&service_file); + let path = path.to_string_lossy().into_owned(); + let loaded = self.cmd("launchctl", &["print", &target])?; + if loaded && (self.changed || plist_changed) { + self.must("launchctl", &["bootout", &target])?; + self.note("ran", format_args!("launchctl bootout {target}")); + } + if !loaded || self.changed || plist_changed { + self.must("launchctl", &["bootstrap", "system", &path])?; + self.note("ran", format_args!("launchctl bootstrap system {path}")); + } + } + } + Ok(()) + } +} + +/// Install (or update) the drone. Returns one line per action taken or skipped. +pub fn install(opts: &InstallOptions, exec: &dyn Exec) -> Result, InstallError> { + let mut i = Installer { + opts, + layout: Layout::new(opts.platform), + exec, + lines: Vec::new(), + changed: false, + }; + i.preflight()?; + let run_as = opts.run_as.clone(); + let owner = Some(run_as.as_str()); + + i.ensure_user()?; + let (config_dir, data) = (i.layout.config_dir.clone(), i.layout.data.clone()); + i.ensure_dir(&config_dir, 0o755, None)?; + i.ensure_dir(&data, 0o755, owner)?; + for sub in ["work", "cache"] { + i.ensure_dir(&data.join(sub), 0o755, owner)?; + } + i.install_binary()?; + i.install_runner()?; + + let mut config_changed = false; + if let Some(ca) = &opts.ca_cert { + let pem = fs::read(ca).map_err(io_err(ca))?; + let target = i.layout.ca_cert(); + config_changed |= i.write_file(&target, &pem, 0o644, None)?; + } + let token_file = i.layout.token_file(); + let token = format!("{}\n", opts.token.expose()); + config_changed |= i.write_file(&token_file, token.as_bytes(), 0o600, owner)?; + let config = i.drone_config()?; + let config_path = i.layout.config(); + config_changed |= i.write_file(&config_path, config.as_bytes(), 0o644, None)?; + + i.install_service(config_changed)?; + + let mut lines = i.lines; + let mut summary = String::new(); + let _ = write!( + summary, + "{}", + if i.changed { + if opts.dry_run { + "dry run: changes listed above were not applied" + } else { + "done: changes applied" + } + } else { + "nothing to change" + } + ); + lines.push(summary); + Ok(lines) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::drone::cleanup::testing::FakeExec; + use crate::drone::pins::pin_for; + + fn tarball_bytes() -> Vec { + b"fake runner tarball".to_vec() + } + + struct Env { + _dir: tempfile::TempDir, + root: PathBuf, + binary: PathBuf, + pin: RunnerPin, + } + + fn env() -> Env { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("root"); + fs::create_dir_all(root.join("run/systemd/system")).unwrap(); + let binary = dir.path().join("borg-binary"); + fs::write(&binary, b"binary-v1").unwrap(); + let mut pin = pin_for("linux", "x86_64").unwrap(); + pin.sha256 = hex::encode(Sha256::digest(tarball_bytes())); + Env { + _dir: dir, + root, + binary, + pin, + } + } + + fn opts(e: &Env) -> InstallOptions { + InstallOptions { + queen_url: "https://queen.example.test:8443".into(), + token: Secret::new("borg_testtoken"), + ca_cert: None, + allow_insecure: false, + max_vcpu: Some(8), + max_mem_mib: Some(24576), + max_leases: None, + run_as: "borg".into(), + root: e.root.clone(), + platform: Platform::Linux, + pin: e.pin.clone(), + binary: e.binary.clone(), + dry_run: false, + skip_root_check: true, + drain_timeout_secs: 3600, + } + } + + /// A fake host where curl/tar produce the fixture files and the user exists. + fn exec() -> FakeExec { + let x = FakeExec::default(); + x.on("getent group docker", "docker:x:999:"); + x.on("id -nG borg", "borg docker"); + x + } + + /// Fake tar/curl side effects on the staging tree. + fn with_downloads(x: FakeExec) -> impl Exec { + struct Hooked(FakeExec); + impl Exec for Hooked { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + match program { + "curl" => fs::write(args[2], tarball_bytes())?, + "tar" => fs::write(Path::new(args[3]).join("run.sh"), "#!/bin/sh\n")?, + _ => {} + } + self.0.run(program, args) + } + } + Hooked(x) + } + + fn calls(x: &FakeExec) -> Vec { + x.calls() + } + + #[test] + fn unit_and_plist_are_rendered_with_drain_safe_settings() { + let unit = render_systemd_unit( + Path::new("/usr/local/bin/borg"), + Path::new("/etc/borg/drone.json"), + Path::new("/var/lib/borg"), + "borg", + 3600, + ); + for line in [ + "User=borg", + "KillMode=mixed", + "TimeoutStopSec=3720", + "ExecStart=/usr/local/bin/borg drone run --config /etc/borg/drone.json", + "Restart=always", + ] { + assert!(unit.contains(line), "{line}\n{unit}"); + } + let plist = render_launchd_plist( + Path::new("/usr/local/bin/borg"), + Path::new("/c.json"), + Path::new("/d"), + "builder", + 60, + ); + for s in [ + "org.borg.drone", + "UserNamebuilder", + "180", + "/c.json", + ] { + assert!(plist.contains(s), "{s}"); + } + } + + #[test] + fn first_install_creates_everything_and_second_changes_nothing() { + let e = env(); + let x = FakeExec::default(); + x.on("getent group docker", "docker:x:999:"); + x.fail("systemctl is-enabled --quiet borg-drone.service"); + x.fail("id -u borg"); + let ex = with_downloads(x); + let first = install(&opts(&e), &ex).unwrap(); + let text = first.join("\n"); + for needle in [ + "created user borg", + "added user borg to group docker", + "created /usr/local/bin/borg", + "installed actions-runner 2.337.0", + "created /etc/borg/drone.token", + "created /etc/borg/drone.json", + "created /etc/systemd/system/borg-drone.service", + "ran systemctl restart borg-drone.service", + "done: changes applied", + ] { + assert!(text.contains(needle), "missing {needle:?} in\n{text}"); + } + assert_eq!( + fs::read(e.root.join("usr/local/bin/borg")).unwrap(), + b"binary-v1" + ); + assert_eq!( + fs::metadata(e.root.join("etc/borg/drone.token")) + .unwrap() + .permissions() + .mode() + & 0o777, + 0o600 + ); + let cfg = DroneConfig::from_json( + &fs::read_to_string(e.root.join("etc/borg/drone.json")).unwrap(), + ) + .unwrap(); + assert_eq!((cfg.max_vcpu, cfg.max_mem_mib), (Some(8), Some(24576))); + assert!( + e.root + .join("var/lib/borg/runner/.borg-runner-version") + .exists() + ); + + // Second run on the now-configured host. + let y = exec(); + let ey = with_downloads(y); + let second = install(&opts(&e), &ey).unwrap(); + let text = second.join("\n"); + assert!(text.contains("nothing to change"), "{text}"); + assert!( + !text.contains("created") && !text.contains("updated") && !text.contains("ran "), + "{text}" + ); + } + + #[test] + fn rerun_restarts_only_when_something_changed() { + let e = env(); + install(&opts(&e), &with_downloads(exec())).unwrap(); + + fs::write(&e.binary, b"binary-v2").unwrap(); + let x = exec(); + let hooked = with_downloads(x); + let out = install(&opts(&e), &hooked).unwrap().join("\n"); + assert!(out.contains("updated /usr/local/bin/borg"), "{out}"); + assert!( + out.contains("ran systemctl restart borg-drone.service"), + "{out}" + ); + } + + #[test] + fn a_pin_change_reinstalls_the_runner() { + let e = env(); + install(&opts(&e), &with_downloads(exec())).unwrap(); + let mut o = opts(&e); + o.pin.version = "2.999.0".into(); + let out = install(&o, &with_downloads(exec())).unwrap().join("\n"); + assert!(out.contains("installed actions-runner 2.999.0"), "{out}"); + let marker = + fs::read_to_string(e.root.join("var/lib/borg/runner/.borg-runner-version")).unwrap(); + assert!(marker.starts_with("2.999.0 ")); + } + + #[test] + fn checksum_mismatch_aborts_before_extraction() { + let e = env(); + let mut o = opts(&e); + o.pin.sha256 = "0".repeat(64); + let x = exec(); + let err = install(&o, &with_downloads(x)).unwrap_err(); + assert!(matches!(err, InstallError::Checksum { .. }), "{err}"); + assert!(!e.root.join("var/lib/borg/runner").exists()); + assert!( + !e.root + .join("etc/systemd/system/borg-drone.service") + .exists(), + "nothing after the failed step" + ); + } + + #[test] + fn dry_run_lists_changes_and_writes_nothing() { + let e = env(); + let mut o = opts(&e); + o.dry_run = true; + let x = FakeExec::default(); + x.fail("id -u borg"); + let out = install(&o, &x).unwrap().join("\n"); + assert!( + out.contains("would create user borg") + && out.contains("would create /etc/borg/drone.json"), + "{out}" + ); + assert!(out.contains("dry run"), "{out}"); + assert!(!e.root.join("etc").exists() && !e.root.join("usr").exists()); + assert!( + calls(&x) + .iter() + .all(|c| !c.starts_with("useradd") && !c.starts_with("systemctl")), + "{:?}", + calls(&x) + ); + } + + #[test] + fn requires_systemd_and_explains_wsl() { + let e = env(); + fs::remove_dir_all(e.root.join("run/systemd")).unwrap(); + let err = install(&opts(&e), &exec()).unwrap_err().to_string(); + assert!( + err.contains("systemd is not running") && !err.contains("wsl.conf"), + "{err}" + ); + fs::create_dir_all(e.root.join("proc")).unwrap(); + fs::write( + e.root.join("proc/version"), + "Linux version 5.15-microsoft-standard-WSL2", + ) + .unwrap(); + let err = install(&opts(&e), &exec()).unwrap_err().to_string(); + assert!( + err.contains("/etc/wsl.conf") && err.contains("systemd=true"), + "{err}" + ); + } + + #[test] + fn non_root_is_refused() { + let e = env(); + let mut o = opts(&e); + o.skip_root_check = false; + let x = FakeExec::default(); + x.on("id -u", "1000"); + assert!( + install(&o, &x) + .unwrap_err() + .to_string() + .contains("run as root") + ); + } + + #[test] + fn macos_uses_launchd_and_requires_an_existing_user() { + let e = env(); + let mut o = opts(&e); + o.platform = Platform::MacOs; + o.run_as = "builder".into(); + let x = FakeExec::default(); + x.fail("launchctl print system/org.borg.drone"); + let out = install(&o, &with_downloads(x)).unwrap().join("\n"); + assert!( + out.contains("created /Library/LaunchDaemons/org.borg.drone.plist"), + "{out}" + ); + assert!(out.contains("ran launchctl bootstrap system"), "{out}"); + + let y = FakeExec::default(); + y.fail("id -u builder"); + let err = install(&o, &y).unwrap_err().to_string(); + assert!(err.contains("create it first"), "{err}"); + } +} diff --git a/src/drone/mod.rs b/src/drone/mod.rs new file mode 100644 index 0000000..289c799 --- /dev/null +++ b/src/drone/mod.rs @@ -0,0 +1,302 @@ +//! The drone: runs on every machine that joins the fleet. +//! +//! It probes the host, heartbeats to the queen with its capacity and owner +//! caps, and runs every lease the queen grants as an ephemeral runner in a +//! private directory. Cleanup after each job is unconditional. On shutdown it +//! stops accepting leases and lets in-flight jobs finish, up to a timeout. + +pub mod cleanup; +pub mod client; +pub mod config; +pub mod exec; +pub mod install; +pub mod pins; +pub mod runner; + +use std::collections::HashSet; +use std::path::Path; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use tokio::task::JoinSet; + +use crate::api::{HeartbeatRequest, LeaseGrant, ResultRequest}; +use crate::probe::{self, Host}; +use cleanup::{Scope, cleanup_docker, remove_lease_dir, sweep_work_root, valid_lease_id}; +use client::{ClientError, QueenClient}; +use config::DroneConfig; +use exec::Exec; +use runner::{RunSpec, run_lease}; + +/// Attempts to deliver a lease result before giving up (the lease then +/// expires on the queen's side). +const RESULT_ATTEMPTS: u32 = 5; + +type Held = Arc>>; + +/// Removes a lease id from the held set when its task ends, even by panic. +struct HeldGuard { + held: Held, + id: String, +} + +impl Drop for HeldGuard { + fn drop(&mut self) { + if let Ok(mut held) = self.held.lock() { + held.remove(&self.id); + } + } +} + +pub struct Drone { + cfg: Arc, + client: Arc, + host: Arc, + exec: Arc, + held: Held, + abort: tokio::sync::watch::Sender, + proc_root: Arc, +} + +impl Drone { + pub fn new( + cfg: DroneConfig, + client: QueenClient, + host: Arc, + exec: Arc, + ) -> Self { + let (abort, _) = tokio::sync::watch::channel(false); + Self { + cfg: Arc::new(cfg), + client: Arc::new(client), + host, + exec, + held: Held::default(), + abort, + proc_root: Arc::from(Path::new("/proc")), + } + } + + fn held_ids(&self) -> Vec { + let mut ids: Vec = self + .held + .lock() + .map(|h| h.iter().cloned().collect()) + .unwrap_or_default(); + ids.sort(); + ids + } + + /// Remove leftovers of a previous run: no lease is held yet, so every + /// lease directory and borg-labelled docker resource is stale. + pub async fn startup_sweep(&self) { + let (work_root, exec, docker) = ( + self.cfg.work_root.clone(), + Arc::clone(&self.exec), + self.cfg.docker_bin.clone(), + ); + let swept = tokio::task::spawn_blocking(move || { + let dirs = sweep_work_root(&work_root); + let docker = cleanup_docker(exec.as_ref(), &docker, Scope::AnyLease); + (dirs, docker) + }) + .await; + match swept { + Ok((dirs, docker)) => { + tracing::info!( + "startup sweep: {:?} lease dirs, {} docker resources", + dirs.map(|d| d.len()), + docker.removed.len() + ); + } + Err(e) => tracing::warn!("startup sweep failed: {e}"), + } + } + + /// One heartbeat: report capacity, start any newly granted leases. + /// Returns how many leases were started. + pub async fn heartbeat_once( + &self, + accepting: bool, + tasks: &mut JoinSet<()>, + ) -> Result { + let host = Arc::clone(&self.host); + let capacity = tokio::task::spawn_blocking(move || probe::probe(host.as_ref())) + .await + .map_err(|e| ClientError::Status { + path: format!("probe: {e}"), + status: 0, + })?; + let req = HeartbeatRequest { + capacity, + caps: self.cfg.host_caps(), + accepting, + active_leases: self.held_ids(), + }; + let resp = self.client.heartbeat(&req).await?; + let mut started = 0; + for grant in resp.leases { + if !accepting || !self.admit(&grant.lease_id) { + continue; + } + let ctx = LeaseContext { + cfg: Arc::clone(&self.cfg), + client: Arc::clone(&self.client), + exec: Arc::clone(&self.exec), + abort: self.abort.subscribe(), + proc_root: Arc::clone(&self.proc_root), + _held: HeldGuard { + held: Arc::clone(&self.held), + id: grant.lease_id.clone(), + }, + }; + tasks.spawn(handle_lease(ctx, grant)); + started += 1; + } + Ok(started) + } + + /// Whether a granted lease may start: well-formed, not already running, + /// and within the local lease cap. Claims the id if so. + fn admit(&self, lease_id: &str) -> bool { + if !valid_lease_id(lease_id) { + tracing::warn!("ignoring lease with unsafe id"); + return false; + } + let Ok(mut held) = self.held.lock() else { + return false; + }; + if held.contains(lease_id) { + return false; + } + if self + .cfg + .max_leases + .is_some_and(|cap| held.len() >= usize::try_from(cap).unwrap_or(usize::MAX)) + { + tracing::warn!("at local lease cap; leaving lease {lease_id} unstarted"); + return false; + } + held.insert(lease_id.to_owned()); + true + } + + /// Heartbeat until `shutdown` resolves, then drain. + pub async fn run(self, shutdown: impl std::future::Future) { + self.startup_sweep().await; + let mut tasks = JoinSet::new(); + let mut every = tokio::time::interval(Duration::from_secs(self.cfg.heartbeat_secs)); + every.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + tokio::pin!(shutdown); + // `interval.tick`, `shutdown` and `JoinSet::join_next` are cancel-safe. + loop { + tokio::select! { + _ = every.tick() => { + if let Err(e) = self.heartbeat_once(true, &mut tasks).await { + tracing::warn!("heartbeat failed: {e}"); + } + } + Some(_) = tasks.join_next(), if !tasks.is_empty() => {} + () = &mut shutdown => break, + } + } + + tracing::info!( + "draining: no new leases; waiting up to {}s for {} job(s)", + self.cfg.drain_timeout_secs, + tasks.len() + ); + let deadline = tokio::time::sleep(Duration::from_secs(self.cfg.drain_timeout_secs)); + tokio::pin!(deadline); + while !tasks.is_empty() { + tokio::select! { + _ = every.tick() => { + if let Err(e) = self.heartbeat_once(false, &mut tasks).await { + tracing::warn!("heartbeat failed: {e}"); + } + } + _ = tasks.join_next() => {} + () = &mut deadline => { + tracing::warn!("drain timeout: stopping remaining jobs"); + let _ = self.abort.send(true); + while tasks.join_next().await.is_some() {} + } + } + } + } +} + +struct LeaseContext { + cfg: Arc, + client: Arc, + exec: Arc, + abort: tokio::sync::watch::Receiver, + proc_root: Arc, + _held: HeldGuard, +} + +/// Run one lease, clean up unconditionally, report the result. +async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { + let spec = RunSpec { + lease_id: &grant.lease_id, + jit_config: &grant.jit_config, + runner_dir: &ctx.cfg.runner_dir, + work_root: &ctx.cfg.work_root, + cache_root: &ctx.cfg.cache_root, + max_job: Duration::from_secs(ctx.cfg.max_job_secs), + proc_root: &ctx.proc_root, + }; + let outcome = run_lease(&spec, ctx.exec.as_ref(), ctx.abort.clone()).await; + + // Runs whether the job passed, failed, timed out, was aborted or never started. + let (work_root, id, docker, exec) = ( + ctx.cfg.work_root.clone(), + grant.lease_id.clone(), + ctx.cfg.docker_bin.clone(), + Arc::clone(&ctx.exec), + ); + let cleaned = tokio::task::spawn_blocking(move || { + let docker = cleanup_docker(exec.as_ref(), &docker, Scope::Lease(&id)); + (remove_lease_dir(&work_root, &id), docker) + }) + .await; + match cleaned { + Ok((dir, docker)) => { + if let Err(e) = dir { + tracing::warn!("lease {} work dir cleanup failed: {e}", grant.lease_id); + } + for err in docker.errors { + tracing::warn!("lease {} docker cleanup: {err}", grant.lease_id); + } + } + Err(e) => tracing::warn!("lease {} cleanup task failed: {e}", grant.lease_id), + } + + let result = match outcome { + Ok(o) => ResultRequest { + exit_code: o.exit_code, + duration_secs: o.duration_secs, + peak_rss_mib: o.peak_rss_mib, + }, + Err(e) => { + tracing::warn!("lease {} could not run: {e}", grant.lease_id); + ResultRequest { + exit_code: -1, + duration_secs: 0, + peak_rss_mib: None, + } + } + }; + for attempt in 1..=RESULT_ATTEMPTS { + match ctx.client.report_result(&grant.lease_id, &result).await { + Ok(()) => return, + Err(e) => { + tracing::warn!( + "reporting lease {} failed (attempt {attempt}): {e}", + grant.lease_id + ); + tokio::time::sleep(Duration::from_secs(u64::from(attempt))).await; + } + } + } +} diff --git a/src/drone/pins.rs b/src/drone/pins.rs new file mode 100644 index 0000000..f46ba8d --- /dev/null +++ b/src/drone/pins.rs @@ -0,0 +1,103 @@ +//! The pinned actions-runner release and its checksums. +//! +//! The values below are the SHA-256 digests published in the +//! `actions/runner` v2.337.0 release notes and were re-computed from the +//! downloaded tarballs. To update, change the version and every digest +//! together. + +pub const RUNNER_VERSION: &str = "2.337.0"; + +/// One downloadable runner tarball. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RunnerPin { + pub version: String, + /// Platform part of the asset name, e.g. `linux-x64`. + pub platform: String, + pub sha256: String, +} + +const TABLE: [(&str, &str, &str, &str); 4] = [ + ( + "linux", + "x86_64", + "linux-x64", + "70920811a4f8ad4328818682bca5c6469c1c942fab52448868071d0063816613", + ), + ( + "linux", + "aarch64", + "linux-arm64", + "9b1dc70626422526e3c94767cf024896beb15da5342a3f4819bf2feac13e0393", + ), + ( + "macos", + "x86_64", + "osx-x64", + "d383f505d7ed041b1873ab68c35dd766fc093f2252330f95bb427be8f2c6dcfc", + ), + ( + "macos", + "aarch64", + "osx-arm64", + "5a2cd92908a93d7276a194e1de6008099f3e7946f3f8e14aa7a1a7b4a31fdec2", + ), +]; + +/// The pin for a host, by `std::env::consts` OS and architecture. +pub fn pin_for(os: &str, arch: &str) -> Option { + TABLE + .iter() + .find(|(o, a, _, _)| *o == os && *a == arch) + .map(|(_, _, platform, sha)| RunnerPin { + version: RUNNER_VERSION.to_owned(), + platform: (*platform).to_owned(), + sha256: (*sha).to_owned(), + }) +} + +impl RunnerPin { + pub fn filename(&self) -> String { + format!("actions-runner-{}-{}.tar.gz", self.platform, self.version) + } + + pub fn url(&self) -> String { + format!( + "https://github.com/actions/runner/releases/download/v{}/{}", + self.version, + self.filename() + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn every_supported_host_has_a_well_formed_pin() { + for (os, arch) in [ + ("linux", "x86_64"), + ("linux", "aarch64"), + ("macos", "x86_64"), + ("macos", "aarch64"), + ] { + let pin = pin_for(os, arch).unwrap(); + assert_eq!(pin.sha256.len(), 64); + assert!( + pin.sha256 + .bytes() + .all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase()) + ); + } + assert!(pin_for("windows", "x86_64").is_none()); + } + + #[test] + fn url_matches_the_release_asset_name() { + let pin = pin_for("linux", "x86_64").unwrap(); + assert_eq!( + pin.url(), + "https://github.com/actions/runner/releases/download/v2.337.0/actions-runner-linux-x64-2.337.0.tar.gz" + ); + } +} diff --git a/src/drone/runner.rs b/src/drone/runner.rs new file mode 100644 index 0000000..f7ef2c4 --- /dev/null +++ b/src/drone/runner.rs @@ -0,0 +1,430 @@ +//! Run one lease as an ephemeral GitHub Actions runner in a private directory. +//! +//! The pinned runner install is cloned into `/` (hard +//! links where the filesystem allows, copies otherwise), `run.sh` is started in +//! its own process group with the JIT config in the environment (never argv), +//! and the group is killed on timeout or abort. This module never cleans up; +//! the caller does, unconditionally. +//! +//! Unix only: borg drones run on Linux (including WSL) and macOS. + +use std::fs; +use std::io; +use std::os::unix::fs::symlink; +use std::path::{Path, PathBuf}; +use std::process::Stdio; +use std::time::{Duration, Instant}; + +use tokio::process::Command; + +use crate::drone::cleanup::{docker_prefix, valid_lease_id}; +use crate::drone::exec::Exec; +use crate::secret::Secret; + +const RSS_SAMPLE_EVERY: Duration = Duration::from_secs(2); + +pub struct RunSpec<'a> { + pub lease_id: &'a str, + pub jit_config: &'a Secret, + pub runner_dir: &'a Path, + pub work_root: &'a Path, + pub cache_root: &'a Path, + pub max_job: Duration, + /// `/proc` (a different root is used by tests). + pub proc_root: &'a Path, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RunOutcome { + pub exit_code: i32, + pub duration_secs: u64, + pub peak_rss_mib: Option, + pub timed_out: bool, + pub aborted: bool, +} + +/// Exit code reported when a job was killed (128 + SIGKILL). +pub const KILLED_EXIT_CODE: i32 = 137; + +/// Recreate `src` under `dst`: directories are created, files are hard-linked +/// (falling back to a copy across filesystems or when links are refused), +/// symlinks are recreated. +pub fn clone_tree(src: &Path, dst: &Path) -> io::Result<()> { + fs::create_dir_all(dst)?; + for entry in fs::read_dir(src)? { + let entry = entry?; + let (from, to) = (entry.path(), dst.join(entry.file_name())); + let kind = entry.file_type()?; + if kind.is_dir() { + clone_tree(&from, &to)?; + } else if kind.is_symlink() { + symlink(fs::read_link(&from)?, &to)?; + } else if fs::hard_link(&from, &to).is_err() { + fs::copy(&from, &to)?; + } + } + Ok(()) +} + +/// Sum of resident memory (KiB) of every process in process group `pgid`. +pub fn group_rss_kib(proc_root: &Path, pgid: u32) -> Option { + let mut total = 0; + let mut seen = false; + for entry in fs::read_dir(proc_root).ok()?.flatten() { + let name = entry.file_name(); + if !name.to_string_lossy().bytes().all(|b| b.is_ascii_digit()) { + continue; + } + let Ok(stat) = fs::read_to_string(entry.path().join("stat")) else { + continue; + }; + // Fields after the parenthesised command: state ppid pgrp ... + let Some(after) = stat.rsplit_once(')').map(|(_, rest)| rest) else { + continue; + }; + if after + .split_whitespace() + .nth(2) + .and_then(|f| f.parse::().ok()) + != Some(pgid) + { + continue; + } + let Ok(status) = fs::read_to_string(entry.path().join("status")) else { + continue; + }; + let rss = status + .lines() + .find_map(|l| l.strip_prefix("VmRSS:")) + .and_then(|v| v.split_whitespace().next()) + .and_then(|v| v.parse::().ok()); + if let Some(kib) = rss { + total += kib; + seen = true; + } + } + seen.then_some(total) +} + +/// `spawn`, retrying briefly on ETXTBSY: the script was hard-linked from an +/// install another thread may still hold open for writing. +async fn spawn_retrying(command: &mut Command) -> io::Result { + const ETXTBSY: i32 = 26; + let mut attempts = 0; + loop { + match command.spawn() { + Err(e) if e.raw_os_error() == Some(ETXTBSY) && attempts < 10 => { + attempts += 1; + tokio::time::sleep(Duration::from_millis(50)).await; + } + other => return other, + } + } +} + +fn kill_group(exec: &dyn Exec, pgid: u32) { + // Negative pid addresses the whole process group. + let target = format!("-{pgid}"); + let _ = exec.run("kill", &["-KILL", "--", &target]); +} + +/// Run the lease to completion, timeout or abort. +pub async fn run_lease( + spec: &RunSpec<'_>, + exec: &dyn Exec, + mut abort: tokio::sync::watch::Receiver, +) -> io::Result { + if !valid_lease_id(spec.lease_id) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "unsafe lease id", + )); + } + let lease_dir: PathBuf = spec.work_root.join(spec.lease_id); + let (src, dst) = (spec.runner_dir.to_owned(), lease_dir.clone()); + tokio::task::spawn_blocking(move || clone_tree(&src, &dst)) + .await + .map_err(io::Error::other)??; + + let prefix = docker_prefix(spec.lease_id); + #[allow( + clippy::disallowed_methods, + reason = "starts the pinned actions-runner run.sh from the drone-owned per-lease directory; its only argument is the queen-issued JIT config" + )] + let mut command = Command::new(lease_dir.join("run.sh")); + command + // The runner reads `ACTIONS_RUNNER_INPUT_` as a fallback for each + // command-line argument (Runner.Listener `CommandSettings`, v2.337.0), + // masks `jitconfig` as a secret and removes the variable from its own + // environment. Passing it here keeps it out of argv, which any local + // user can read from `ps` or `/proc//cmdline`. + .env("ACTIONS_RUNNER_INPUT_JITCONFIG", spec.jit_config.expose()) + .current_dir(&lease_dir) + .env("BORG_LEASE_ID", spec.lease_id) + .env("BORG_DOCKER_PREFIX", &prefix) + .env("COMPOSE_PROJECT_NAME", &prefix) + .env("BORG_CACHE_ROOT", spec.cache_root) + // The shared install is hard-linked into every lease; never let a + // runner rewrite it in place. + .env("DISABLE_RUNNER_UPDATE", "1") + .stdin(Stdio::null()) + .process_group(0) + .kill_on_drop(true); + let mut child = spawn_retrying(&mut command).await?; + let pgid = child.id(); + + let started = Instant::now(); + let deadline = tokio::time::sleep(spec.max_job); + tokio::pin!(deadline); + let mut sample = tokio::time::interval(RSS_SAMPLE_EVERY); + let (mut peak_kib, mut timed_out, mut aborted) = (0u64, false, false); + + // Every arm is cancel-safe: `Child::wait`, `Sleep`, `Interval::tick` and + // `watch::Receiver::changed` keep no partially consumed state. + let status = loop { + tokio::select! { + status = child.wait() => break status?, + () = &mut deadline => { + timed_out = true; + if let Some(pgid) = pgid { kill_group(exec, pgid); } + break child.wait().await?; + } + changed = abort.changed() => { + if changed.is_ok() && *abort.borrow() { + aborted = true; + if let Some(pgid) = pgid { kill_group(exec, pgid); } + break child.wait().await?; + } + } + _ = sample.tick() => { + if let Some(pgid) = pgid { + peak_kib = peak_kib.max(group_rss_kib(spec.proc_root, pgid).unwrap_or(0)); + } + } + } + }; + // Reap anything the runner left in its group. + if let Some(pgid) = pgid { + kill_group(exec, pgid); + } + Ok(RunOutcome { + exit_code: if timed_out || aborted { + KILLED_EXIT_CODE + } else { + status.code().unwrap_or(KILLED_EXIT_CODE) + }, + duration_secs: started.elapsed().as_secs(), + peak_rss_mib: (peak_kib > 0).then_some(peak_kib / 1024), + timed_out, + aborted, + }) +} + +#[cfg(test)] +mod tests { + use std::os::unix::fs::PermissionsExt; + + use super::*; + use crate::drone::exec::SystemExec; + + fn runner_with_script(dir: &Path, body: &str) -> PathBuf { + let runner = dir.join("runner"); + fs::create_dir_all(runner.join("bin")).unwrap(); + fs::write(runner.join("bin/marker"), "pinned").unwrap(); + let script = runner.join("run.sh"); + fs::write(&script, format!("#!/bin/sh\n{body}\n")).unwrap(); + fs::set_permissions(&script, fs::Permissions::from_mode(0o755)).unwrap(); + runner + } + + async fn run(dir: &Path, body: &str, max_job: Duration) -> (RunOutcome, PathBuf) { + let runner = runner_with_script(dir, body); + let work = dir.join("work"); + let (_tx, rx) = tokio::sync::watch::channel(false); + let secret = Secret::new("jit-secret-value"); + let spec = RunSpec { + lease_id: "lease1", + jit_config: &secret, + runner_dir: &runner, + work_root: &work, + cache_root: &dir.join("cache"), + max_job, + proc_root: Path::new("/proc"), + }; + let out = run_lease(&spec, &SystemExec, rx).await.unwrap(); + (out, work.join("lease1")) + } + + #[tokio::test] + async fn runs_run_sh_in_a_private_clone_with_the_jit_config_and_lease_env() { + let dir = tempfile::tempdir().unwrap(); + let (out, lease_dir) = run( + dir.path(), + "echo \"$@\" > args.txt; printf '%s' \"$ACTIONS_RUNNER_INPUT_JITCONFIG\" > jit.txt; echo \"$BORG_LEASE_ID $BORG_DOCKER_PREFIX $COMPOSE_PROJECT_NAME $DISABLE_RUNNER_UPDATE\" > env.txt; pwd > cwd.txt; tr '\\0' ' ' < /proc/$$/cmdline > cmdline.txt 2>/dev/null || true", + Duration::from_secs(30), + ) + .await; + assert_eq!( + (out.exit_code, out.timed_out, out.aborted), + (0, false, false) + ); + // The JIT config travels in the environment and never in argv. + assert_eq!( + fs::read_to_string(lease_dir.join("jit.txt")).unwrap(), + "jit-secret-value" + ); + assert_eq!( + fs::read_to_string(lease_dir.join("args.txt")) + .unwrap() + .trim(), + "", + "run.sh gets no arguments" + ); + let cmdline = fs::read_to_string(lease_dir.join("cmdline.txt")).unwrap_or_default(); + assert!( + !cmdline.contains("jit-secret-value"), + "secret visible in the process command line: {cmdline}" + ); + assert_eq!( + fs::read_to_string(lease_dir.join("env.txt")) + .unwrap() + .trim(), + "lease1 borg-lease1 borg-lease1 1" + ); + assert!( + fs::read_to_string(lease_dir.join("cwd.txt")) + .unwrap() + .trim() + .ends_with("work/lease1") + ); + assert_eq!( + fs::read_to_string(lease_dir.join("bin/marker")).unwrap(), + "pinned", + "install was cloned" + ); + assert!( + !dir.path().join("runner/args.txt").exists(), + "the shared install is not written to" + ); + } + + #[tokio::test] + async fn reports_the_runner_exit_code() { + let dir = tempfile::tempdir().unwrap(); + let (out, _) = run(dir.path(), "exit 7", Duration::from_secs(30)).await; + assert_eq!(out.exit_code, 7); + } + + #[tokio::test] + async fn kills_the_whole_process_group_on_timeout() { + let dir = tempfile::tempdir().unwrap(); + let (out, lease_dir) = run( + dir.path(), + "sleep 300 &\necho $! > child.pid\nwait", + Duration::from_millis(500), + ) + .await; + assert!(out.timed_out); + assert_eq!(out.exit_code, KILLED_EXIT_CODE); + let pid = fs::read_to_string(lease_dir.join("child.pid")) + .unwrap() + .trim() + .to_owned(); + tokio::time::sleep(Duration::from_millis(300)).await; + let alive = fs::read_to_string(format!("/proc/{pid}/stat")).is_ok_and(|stat| { + stat.rsplit_once(')') + .is_some_and(|(_, rest)| !rest.trim_start().starts_with('Z')) + }); + assert!( + !alive || cfg!(not(target_os = "linux")), + "background child {pid} survived" + ); + } + + #[tokio::test] + async fn abort_signal_stops_a_running_job() { + let dir = tempfile::tempdir().unwrap(); + let runner = runner_with_script(dir.path(), "sleep 300"); + let work = dir.path().join("work"); + let (tx, rx) = tokio::sync::watch::channel(false); + let secret = Secret::new("x"); + let cache = dir.path().join("cache"); + let spec = RunSpec { + lease_id: "lease2", + jit_config: &secret, + runner_dir: &runner, + work_root: &work, + cache_root: &cache, + max_job: Duration::from_secs(60), + proc_root: Path::new("/proc"), + }; + let abort = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(300)).await; + let _ = tx.send(true); + tx + }); + let out = run_lease(&spec, &SystemExec, rx).await.unwrap(); + let _tx = abort.await.unwrap(); + assert!(out.aborted && !out.timed_out); + } + + #[tokio::test] + async fn unsafe_lease_ids_never_touch_the_filesystem() { + let dir = tempfile::tempdir().unwrap(); + let runner = runner_with_script(dir.path(), "true"); + let (_tx, rx) = tokio::sync::watch::channel(false); + let secret = Secret::new("x"); + let spec = RunSpec { + lease_id: "../escape", + jit_config: &secret, + runner_dir: &runner, + work_root: &dir.path().join("work"), + cache_root: dir.path(), + max_job: Duration::from_secs(1), + proc_root: Path::new("/proc"), + }; + assert!(run_lease(&spec, &SystemExec, rx).await.is_err()); + assert!(!dir.path().join("work").exists()); + } + + #[test] + fn clone_hard_links_files_and_recreates_symlinks() { + let dir = tempfile::tempdir().unwrap(); + let src = dir.path().join("src"); + fs::create_dir_all(src.join("a/b")).unwrap(); + fs::write(src.join("a/b/f"), "data").unwrap(); + symlink("a/b/f", src.join("link")).unwrap(); + let dst = dir.path().join("dst"); + clone_tree(&src, &dst).unwrap(); + assert_eq!(fs::read_to_string(dst.join("a/b/f")).unwrap(), "data"); + assert_eq!(fs::read_link(dst.join("link")).unwrap(), Path::new("a/b/f")); + // Deleting the clone leaves the source intact. + fs::remove_dir_all(&dst).unwrap(); + assert_eq!(fs::read_to_string(src.join("a/b/f")).unwrap(), "data"); + } + + #[test] + fn group_rss_sums_members_of_the_group_only() { + let dir = tempfile::tempdir().unwrap(); + let add = |pid: &str, comm: &str, pgrp: u32, rss: &str| { + let p = dir.path().join(pid); + fs::create_dir_all(&p).unwrap(); + fs::write( + p.join("stat"), + format!("{pid} ({comm}) S 1 {pgrp} {pgrp} 0 -1\n"), + ) + .unwrap(); + fs::write( + p.join("status"), + format!("Name:\t{comm}\nVmRSS:\t{rss} kB\n"), + ) + .unwrap(); + }; + add("100", "run.sh", 100, "1000"); + add("101", "weird) name", 100, "3000"); + add("200", "other", 200, "9999"); + fs::create_dir_all(dir.path().join("self")).unwrap(); + assert_eq!(group_rss_kib(dir.path(), 100), Some(4000)); + assert_eq!(group_rss_kib(dir.path(), 999), None); + } +} diff --git a/src/enroll.rs b/src/enroll.rs new file mode 100644 index 0000000..01c5ad1 --- /dev/null +++ b/src/enroll.rs @@ -0,0 +1,54 @@ +//! Drone enrollment tokens. The queen stores only the SHA-256 of a token, so a +//! database leak does not let anyone impersonate a drone. Tokens are 256 bits +//! of OS randomness, which makes an unsalted fast hash sufficient. + +use rand::RngCore; +use rand::rngs::OsRng; +use sha2::{Digest, Sha256}; + +const PREFIX: &str = "borg_"; + +/// A fresh random enrollment token. Shown once to the operator. +pub fn new_token() -> String { + let mut bytes = [0u8; 32]; + OsRng.fill_bytes(&mut bytes); + format!("{PREFIX}{}", hex::encode(bytes)) +} + +/// The value stored and looked up for `token`. +pub fn hash_token(token: &str) -> String { + hex::encode(Sha256::digest(token.as_bytes())) +} + +/// A drone name is part of runner names and labels, so keep it to a safe set. +pub fn valid_drone_name(name: &str) -> bool { + !name.is_empty() + && name.len() <= 40 + && name + .bytes() + .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-') + && !name.starts_with('-') +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tokens_are_unique_and_hash_deterministically() { + let (a, b) = (new_token(), new_token()); + assert_ne!(a, b); + assert!(a.starts_with("borg_") && a.len() == 5 + 64); + assert_eq!(hash_token(&a), hash_token(&a)); + assert_ne!(hash_token(&a), hash_token(&b)); + assert!(!hash_token(&a).contains(&a)); + } + + #[test] + fn drone_names_are_restricted() { + assert!(valid_drone_name("host-1")); + for bad in ["", "-x", "Host", "a b", "a_b", &"x".repeat(41)] { + assert!(!valid_drone_name(bad), "{bad}"); + } + } +} diff --git a/src/github.rs b/src/github.rs new file mode 100644 index 0000000..4edd83a --- /dev/null +++ b/src/github.rs @@ -0,0 +1,367 @@ +//! GitHub App authentication and the handful of REST calls the queen needs. +//! +//! The App private key never leaves this module: it signs a short-lived JWT, +//! which is exchanged for an installation token, which authorizes everything +//! else. Drones only ever receive single-use JIT runner configs. + +use std::collections::BTreeSet; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use jsonwebtoken::{Algorithm, EncodingKey, Header}; +use reqwest::{Method, RequestBuilder}; +use serde::Deserialize; +use serde::de::DeserializeOwned; +use time::OffsetDateTime; +use time::format_description::well_known::Rfc3339; + +use crate::secret::Secret; + +/// Upper bound on pages fetched per list call (100 items per page). +const MAX_PAGES: u32 = 20; +/// Refresh the installation token this long before it expires. +const TOKEN_SLACK_SECS: i64 = 60; + +#[derive(Debug, thiserror::Error)] +pub enum GithubError { + #[error("request failed: {0}")] + Http(#[from] reqwest::Error), + #[error("{method} {path} returned {status}: {body}")] + Status { + method: &'static str, + path: String, + status: u16, + body: String, + }, + #[error("sign app jwt: {0}")] + Jwt(#[from] jsonwebtoken::errors::Error), + #[error("bad private key: {0}")] + Key(String), + #[error("unexpected response: {0}")] + Decode(String), +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QueuedJob { + pub id: u64, + pub repo: String, + pub labels: Vec, + pub created_at: OffsetDateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct Runner { + pub id: u64, + pub name: String, + pub status: String, + pub busy: bool, +} + +#[derive(Debug)] +pub struct Minted { + pub runner_id: u64, + pub jit_config: Secret, +} + +struct CachedToken { + value: Secret, + expires_at: i64, +} + +#[derive(Default)] +struct AuthState { + installation_id: Option, + token: Option, +} + +pub struct GithubClient { + http: reqwest::Client, + api_base: String, + org: String, + app_id: String, + key: EncodingKey, + auth: tokio::sync::Mutex, +} + +#[derive(serde::Serialize)] +struct Claims<'a> { + iat: i64, + exp: i64, + iss: &'a str, +} + +fn unix_now() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(0)) +} + +impl GithubClient { + pub fn new( + api_base: &str, + org: &str, + app_id: &str, + private_key_pem: &[u8], + ) -> Result { + let key = EncodingKey::from_rsa_pem(private_key_pem) + .map_err(|e| GithubError::Key(e.to_string()))?; + let http = reqwest::Client::builder() + .user_agent("borg-queen") + .timeout(Duration::from_secs(30)) + .build()?; + Ok(Self { + http, + api_base: api_base.trim_end_matches('/').to_owned(), + org: org.to_owned(), + app_id: app_id.to_owned(), + key, + auth: tokio::sync::Mutex::new(AuthState::default()), + }) + } + + fn app_jwt(&self) -> Result { + let now = unix_now(); + // GitHub allows at most 10 minutes; backdate iat for clock skew. + let claims = Claims { + iat: now - 60, + exp: now + 9 * 60, + iss: &self.app_id, + }; + Ok(jsonwebtoken::encode( + &Header::new(Algorithm::RS256), + &claims, + &self.key, + )?) + } + + fn request(&self, method: Method, path: &str) -> RequestBuilder { + self.http + .request(method, format!("{}{path}", self.api_base)) + .header("Accept", "application/vnd.github+json") + .header("X-GitHub-Api-Version", "2022-11-28") + } + + async fn send( + &self, + method: Method, + path: &str, + rb: RequestBuilder, + ) -> Result { + let resp = rb.send().await?; + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + return Err(GithubError::Status { + method: static_method(&method), + path: path.to_owned(), + status: status.as_u16(), + body: body.chars().take(200).collect(), + }); + } + let bytes = resp.bytes().await?; + serde_json::from_slice(&bytes).map_err(|e| GithubError::Decode(format!("{path}: {e}"))) + } + + /// A valid installation token, refreshed when near expiry. + async fn installation_token(&self) -> Result { + #[derive(Deserialize)] + struct Installation { + id: u64, + } + #[derive(Deserialize)] + struct Issued { + token: String, + expires_at: String, + } + + let mut auth = self.auth.lock().await; + let now = unix_now(); + if let Some(t) = &auth.token + && t.expires_at - TOKEN_SLACK_SECS > now + { + return Ok(t.value.clone()); + } + let jwt = self.app_jwt()?; + let installation_id = match auth.installation_id { + Some(id) => id, + None => { + let path = format!("/orgs/{}/installation", self.org); + let rb = self.request(Method::GET, &path).bearer_auth(&jwt); + let found: Installation = self.send(Method::GET, &path, rb).await?; + auth.installation_id = Some(found.id); + found.id + } + }; + let path = format!("/app/installations/{installation_id}/access_tokens"); + let rb = self.request(Method::POST, &path).bearer_auth(&jwt); + let issued: Issued = self.send(Method::POST, &path, rb).await?; + let expires_at = OffsetDateTime::parse(&issued.expires_at, &Rfc3339) + .map_err(|e| GithubError::Decode(format!("expires_at: {e}")))? + .unix_timestamp(); + let value = Secret::new(issued.token); + auth.token = Some(CachedToken { + value: value.clone(), + expires_at, + }); + Ok(value) + } + + async fn authed(&self, method: Method, path: &str) -> Result { + let token = self.installation_token().await?; + Ok(self.request(method, path).bearer_auth(token.expose())) + } + + async fn get(&self, path: &str) -> Result { + let rb = self.authed(Method::GET, path).await?; + self.send(Method::GET, path, rb).await + } + + /// Jobs that are `queued` in `repo`, across queued and in-progress runs. + pub async fn queued_jobs(&self, repo: &str) -> Result, GithubError> { + #[derive(Deserialize)] + struct Runs { + workflow_runs: Vec, + } + #[derive(Deserialize)] + struct RunRef { + id: u64, + } + #[derive(Deserialize)] + struct Jobs { + jobs: Vec, + } + #[derive(Deserialize)] + struct Job { + id: u64, + status: String, + labels: Vec, + created_at: String, + } + + let mut run_ids = BTreeSet::new(); + for status in ["queued", "in_progress"] { + for page in 1..=MAX_PAGES { + let runs: Runs = self + .get(&format!( + "/repos/{repo}/actions/runs?status={status}&per_page=100&page={page}" + )) + .await?; + let n = runs.workflow_runs.len(); + run_ids.extend(runs.workflow_runs.into_iter().map(|r| r.id)); + if n < 100 { + break; + } + } + } + let mut out = Vec::new(); + for run in run_ids { + for page in 1..=MAX_PAGES { + let jobs: Jobs = self + .get(&format!("/repos/{repo}/actions/runs/{run}/jobs?filter=latest&per_page=100&page={page}")) + .await?; + let n = jobs.jobs.len(); + for j in jobs.jobs.into_iter().filter(|j| j.status == "queued") { + let created_at = + OffsetDateTime::parse(&j.created_at, &Rfc3339).map_err(|e| { + GithubError::Decode(format!("job {} created_at: {e}", j.id)) + })?; + out.push(QueuedJob { + id: j.id, + repo: repo.to_owned(), + labels: j.labels, + created_at, + }); + } + if n < 100 { + break; + } + } + } + Ok(out) + } + + /// Every self-hosted runner registered to the organization. + pub async fn list_runners(&self) -> Result, GithubError> { + #[derive(Deserialize)] + struct Runners { + runners: Vec, + } + let mut out = Vec::new(); + for page in 1..=MAX_PAGES { + let r: Runners = self + .get(&format!( + "/orgs/{}/actions/runners?per_page=100&page={page}", + self.org + )) + .await?; + let n = r.runners.len(); + out.extend(r.runners); + if n < 100 { + break; + } + } + Ok(out) + } + + /// Register a one-shot (JIT) runner and return its config. + pub async fn mint_jit_runner( + &self, + name: &str, + runner_group_id: u64, + labels: &[String], + work_folder: &str, + ) -> Result { + #[derive(Deserialize)] + struct Reply { + runner: RunnerRef, + encoded_jit_config: String, + } + #[derive(Deserialize)] + struct RunnerRef { + id: u64, + } + let path = format!("/orgs/{}/actions/runners/generate-jitconfig", self.org); + let body = serde_json::json!({ + "name": name, + "runner_group_id": runner_group_id, + "labels": labels, + "work_folder": work_folder, + }); + let rb = self.authed(Method::POST, &path).await?.json(&body); + let reply: Reply = self.send(Method::POST, &path, rb).await?; + Ok(Minted { + runner_id: reply.runner.id, + jit_config: Secret::new(reply.encoded_jit_config), + }) + } + + /// Delete a runner registration. A runner that is already gone is fine. + pub async fn delete_runner(&self, runner_id: u64) -> Result<(), GithubError> { + let path = format!("/orgs/{}/actions/runners/{runner_id}", self.org); + let resp = self.authed(Method::DELETE, &path).await?.send().await?; + let status = resp.status(); + if status.is_success() || status.as_u16() == 404 { + return Ok(()); + } + Err(GithubError::Status { + method: "DELETE", + path, + status: status.as_u16(), + body: resp + .text() + .await + .unwrap_or_default() + .chars() + .take(200) + .collect(), + }) + } +} + +fn static_method(m: &Method) -> &'static str { + match *m { + Method::GET => "GET", + Method::POST => "POST", + Method::DELETE => "DELETE", + _ => "REQUEST", + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..e267324 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,21 @@ +//! borg: a self-organizing CI fleet on commodity hardware. +//! +//! * [`probe`], [`classes`] and [`allocator`] are the pure core shared by the +//! drone and the queen. +//! * The drone ([`drone`]) runs on every machine: it heartbeats to the queen, +//! runs each lease as an ephemeral runner and cleans up after it. +//! * The queen ([`queen`], [`github`], [`store`], [`api`], [`config`]) polls +//! the GitHub Actions queue, decides what to lease where, mints single-use +//! JIT runner configs and serves the drone API. + +pub mod allocator; +pub mod api; +pub mod classes; +pub mod config; +pub mod drone; +pub mod enroll; +pub mod github; +pub mod probe; +pub mod queen; +pub mod secret; +pub mod store; diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..5911fe5 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,264 @@ +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use anyhow::{Context, Result, bail}; +use borg::api; +use borg::config::QueenConfig; +use borg::drone::client::QueenClient; +use borg::drone::config::DroneConfig; +use borg::drone::exec::SystemExec; +use borg::drone::install::{self, InstallOptions, Platform}; +use borg::drone::{Drone, pins}; +use borg::enroll; +use borg::github::GithubClient; +use borg::probe::{self, system::SystemHost}; +use borg::queen::Queen; +use borg::store::Store; +use clap::{Args, Parser, Subcommand}; + +#[derive(Parser)] +#[command(name = "borg", version, about = "Self-organizing CI fleet agent")] +struct Cli { + #[command(subcommand)] + command: Command, +} + +#[derive(Subcommand)] +enum Command { + /// Probe this host and print its capacity as JSON. + Probe { + /// Emit compact single-line JSON instead of pretty-printed JSON. + #[arg(long)] + compact: bool, + }, + /// Run or install the drone (the agent on every machine). + Drone { + #[command(subcommand)] + command: DroneCommand, + }, + /// Run or administer the queen (the coordinator). + Queen { + #[command(subcommand)] + command: QueenCommand, + }, +} + +#[derive(Args)] +struct ConfigArg { + /// Path to the queen's JSON config file. + #[arg(long, env = "BORG_QUEEN_CONFIG")] + config: PathBuf, +} + +#[derive(Subcommand)] +enum QueenCommand { + /// Poll the queue, mint runners and serve the drone API. + Run(ConfigArg), + /// Enroll a drone and print its one-time token. + Enroll { + #[command(flatten)] + config: ConfigArg, + /// Drone name: lowercase letters, digits and dashes. + name: String, + }, +} + +#[derive(Args)] +struct DroneConfigArg { + /// Path to the drone's JSON config file. + #[arg(long, env = "BORG_DRONE_CONFIG")] + config: PathBuf, +} + +#[derive(Subcommand)] +enum DroneCommand { + /// Heartbeat to the queen and run the leases it grants. + Run(DroneConfigArg), + /// Install this host as a drone (systemd unit or LaunchDaemon). + Install(InstallArgs), +} + +#[derive(Args)] +struct InstallArgs { + /// Queen base URL, e.g. https://queen.example.test:8443 + #[arg(long)] + queen: String, + /// Enrollment token from `borg queen enroll`. + #[arg(long, env = "BORG_DRONE_TOKEN", hide_env_values = true)] + token: String, + /// PEM CA certificate that signs the queen's certificate (pins the CA). + #[arg(long)] + ca_cert: Option, + /// Allow a plain http:// queen URL (trusted network only). + #[arg(long)] + allow_insecure: bool, + /// Owner cap: most vCPUs this host lends. + #[arg(long)] + max_vcpu: Option, + /// Owner cap: most memory (MiB) this host lends. + #[arg(long)] + max_mem_mib: Option, + /// Owner cap: most concurrent leases. + #[arg(long)] + max_leases: Option, + /// Service account (Linux: created if missing; macOS: must exist). + #[arg(long, default_value = "borg")] + run_as: String, + /// Install under this root instead of / (staging and tests). + #[arg(long, default_value = "/")] + root: PathBuf, + /// Print what would change without changing anything. + #[arg(long)] + dry_run: bool, +} + +fn private_key(cfg: &QueenConfig) -> Result> { + match &cfg.github.private_key_path { + Some(path) => { + std::fs::read(path).with_context(|| format!("read private key {}", path.display())) + } + None => std::env::var("BORG_APP_PRIVATE_KEY") + .map(String::into_bytes) + .context("set github.private_key_path or BORG_APP_PRIVATE_KEY"), + } +} + +async fn open_store(cfg: &QueenConfig) -> Result { + Store::open(&cfg.db_path) + .await + .with_context(|| format!("open {}", cfg.db_path.display())) +} + +async fn wait_for_shutdown() { + #[cfg(unix)] + { + use tokio::signal::unix::{SignalKind, signal}; + if let Ok(mut term) = signal(SignalKind::terminate()) { + tokio::select! { + _ = term.recv() => {} + _ = tokio::signal::ctrl_c() => {} + } + return; + } + } + let _ = tokio::signal::ctrl_c().await; +} + +async fn run_queen(path: &Path) -> Result<()> { + let cfg = QueenConfig::load(path)?; + let github = GithubClient::new( + &cfg.github.api_base, + &cfg.github.org, + &cfg.github.app_id, + &private_key(&cfg)?, + )?; + let store = open_store(&cfg).await?; + let listener = tokio::net::TcpListener::bind(cfg.listen) + .await + .with_context(|| format!("bind {}", cfg.listen))?; + tracing::info!("queen listening on {}", cfg.listen); + let queen = Arc::new(Queen::new(cfg, store, github)); + api::serve(queen, listener, wait_for_shutdown()).await?; + Ok(()) +} + +async fn enroll_drone(path: &Path, name: &str) -> Result<()> { + if !enroll::valid_drone_name(name) { + bail!( + "drone name must be 1-40 lowercase letters, digits or dashes, not starting with a dash" + ); + } + let cfg = QueenConfig::load(path)?; + let store = open_store(&cfg).await?; + let one_time_code = enroll::new_token(); + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + store + .enroll(name, &enroll::hash_token(&one_time_code), now) + .await?; + eprintln!("Enrolled drone {name}. The token is shown once; the queen stores only its hash."); + println!("{one_time_code}"); + Ok(()) +} + +async fn run_drone(path: &Path) -> Result<()> { + let cfg = DroneConfig::load(path)?; + let token = cfg.load_token()?; + let client = QueenClient::new(&cfg, token)?; + let drone = Drone::new(cfg, client, Arc::new(SystemHost), Arc::new(SystemExec)); + drone.run(wait_for_shutdown()).await; + Ok(()) +} + +fn install_drone(a: InstallArgs) -> Result<()> { + let platform = + Platform::current().context("borg drone install supports Linux, WSL and macOS")?; + let pin = pins::pin_for(std::env::consts::OS, std::env::consts::ARCH).with_context(|| { + format!( + "no pinned actions-runner for {} {}", + std::env::consts::OS, + std::env::consts::ARCH + ) + })?; + let opts = InstallOptions { + queen_url: a.queen, + token: borg::secret::Secret::new(a.token), + ca_cert: a.ca_cert, + allow_insecure: a.allow_insecure, + max_vcpu: a.max_vcpu, + max_mem_mib: a.max_mem_mib, + max_leases: a.max_leases, + run_as: a.run_as, + skip_root_check: a.root != Path::new("/"), + root: a.root, + platform, + pin, + binary: std::env::current_exe().context("locate this binary")?, + dry_run: a.dry_run, + drain_timeout_secs: 3600, + }; + for line in install::install(&opts, &SystemExec)? { + println!("{line}"); + } + Ok(()) +} + +#[tokio::main] +async fn main() -> Result<()> { + match Cli::parse().command { + Command::Probe { compact } => { + let capacity = probe::probe(&SystemHost); + let json = if compact { + serde_json::to_string(&capacity) + } else { + serde_json::to_string_pretty(&capacity) + } + .context("serialize capacity")?; + println!("{json}"); + Ok(()) + } + Command::Drone { command } => { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "info".into()), + ) + .init(); + match command { + DroneCommand::Run(c) => run_drone(&c.config).await, + DroneCommand::Install(a) => install_drone(a), + } + } + Command::Queen { command } => { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "info".into()), + ) + .init(); + match command { + QueenCommand::Run(c) => run_queen(&c.config).await, + QueenCommand::Enroll { config, name } => enroll_drone(&config.config, &name).await, + } + } + } +} diff --git a/src/probe/mod.rs b/src/probe/mod.rs new file mode 100644 index 0000000..faf140c --- /dev/null +++ b/src/probe/mod.rs @@ -0,0 +1,382 @@ +//! Host probe: what this machine can lend to the fleet. +//! +//! Every host interaction goes through the [`Host`] trait so detection logic +//! is tested against fake filesystems and command outputs. + +pub mod system; + +use std::path::Path; + +use serde::{Deserialize, Serialize}; + +/// What kind of machine the drone runs on. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum HostKind { + Wsl, + Lxc, + Vm, + Bare, +} + +/// Read-only view of the machine. Reads and commands return `None` on any +/// failure; the probe degrades field by field rather than failing. +pub trait Host { + fn read_to_string(&self, path: &Path) -> Option; + fn exists(&self, path: &Path) -> bool; + /// Run `program` and return trimmed stdout if it exited 0. + fn run(&self, program: &str, args: &[&str]) -> Option; + /// Operating system, as in `std::env::consts::OS` (`linux`, `macos`). + fn os(&self) -> &'static str; + /// CPU architecture, as in `std::env::consts::ARCH`. + fn arch(&self) -> &'static str; + /// CPUs this process may use (honours cgroup quotas). + fn available_parallelism(&self) -> Option; +} + +/// The capacity record a drone reports to the queen. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Capacity { + pub hostname: String, + pub os: String, + pub arch: String, + pub kind: HostKind, + pub vcpu: u32, + pub mem_mib: u64, + /// Free space on the filesystem holding the work directory. + pub free_disk_mib: Option, + /// `docker info` reached a daemon. + pub docker_usable: bool, + pub mise_version: Option, + pub git_version: Option, +} + +/// Directory whose filesystem free space is reported. +const DISK_PROBE_PATH: &str = "/"; + +pub fn probe(host: &dyn Host) -> Capacity { + Capacity { + hostname: hostname(host), + os: host.os().to_owned(), + arch: host.arch().to_owned(), + kind: host_kind(host), + vcpu: host.available_parallelism().unwrap_or(1).max(1), + mem_mib: mem_mib(host).unwrap_or(0), + free_disk_mib: free_disk_mib(host), + docker_usable: host + .run("docker", &["info", "--format", "{{.ServerVersion}}"]) + .is_some(), + mise_version: host.run("mise", &["--version"]), + git_version: host.run("git", &["--version"]), + } +} + +fn hostname(host: &dyn Host) -> String { + host.read_to_string(Path::new("/proc/sys/kernel/hostname")) + .or_else(|| host.run("hostname", &[])) + .map(|h| h.trim().to_owned()) + .filter(|h| !h.is_empty()) + .unwrap_or_else(|| "unknown".to_owned()) +} + +/// Detect WSL, then LXC, then a hypervisor guest, else bare metal. +pub fn host_kind(host: &dyn Host) -> HostKind { + if is_wsl(host) { + HostKind::Wsl + } else if is_lxc(host) { + HostKind::Lxc + } else if is_vm(host) { + HostKind::Vm + } else { + HostKind::Bare + } +} + +fn is_wsl(host: &dyn Host) -> bool { + host.exists(Path::new("/proc/sys/fs/binfmt_misc/WSLInterop")) + || host.exists(Path::new("/run/WSL")) + || host + .read_to_string(Path::new("/proc/version")) + .is_some_and(|v| { + let v = v.to_ascii_lowercase(); + v.contains("microsoft") || v.contains("wsl") + }) +} + +fn is_lxc(host: &dyn Host) -> bool { + if host + .read_to_string(Path::new("/run/systemd/container")) + .is_some_and(|c| c.trim().starts_with("lxc")) + { + return true; + } + host.read_to_string(Path::new("/proc/1/environ")) + .is_some_and(|env| env.split('\0').any(|kv| kv.starts_with("container=lxc"))) +} + +fn is_vm(host: &dyn Host) -> bool { + if host.os() == "macos" { + return host + .run("sysctl", &["-n", "kern.hv_vmm_present"]) + .as_deref() + == Some("1"); + } + const VM_MARKERS: [&str; 7] = [ + "qemu", + "kvm", + "vmware", + "virtualbox", + "xen", + "microsoft corporation", + "bochs", + ]; + let dmi = ["product_name", "sys_vendor"].iter().any(|field| { + host.read_to_string(&Path::new("/sys/class/dmi/id").join(field)) + .is_some_and(|v| { + let v = v.to_ascii_lowercase(); + VM_MARKERS.iter().any(|m| v.contains(m)) + }) + }); + dmi || host + .read_to_string(Path::new("/proc/cpuinfo")) + .is_some_and(|c| { + c.lines() + .any(|l| l.starts_with("flags") && l.split_whitespace().any(|f| f == "hypervisor")) + }) +} + +/// Total memory in MiB, clamped by any cgroup limit (LXC and containers +/// report the host's `MemTotal` unless lxcfs is mounted). +fn mem_mib(host: &dyn Host) -> Option { + let total = if host.os() == "macos" { + host.run("sysctl", &["-n", "hw.memsize"])? + .parse::() + .ok() + .map(|bytes| bytes / (1024 * 1024))? + } else { + let meminfo = host.read_to_string(Path::new("/proc/meminfo"))?; + let kib = meminfo + .lines() + .find_map(|l| l.strip_prefix("MemTotal:"))? + .split_whitespace() + .next()? + .parse::() + .ok()?; + kib / 1024 + }; + Some(match cgroup_mem_limit_mib(host) { + Some(limit) if limit < total => limit, + _ => total, + }) +} + +fn cgroup_mem_limit_mib(host: &dyn Host) -> Option { + [ + "/sys/fs/cgroup/memory.max", + "/sys/fs/cgroup/memory/memory.limit_in_bytes", + ] + .iter() + .find_map(|path| { + host.read_to_string(Path::new(path))? + .trim() + .parse::() + .ok() + }) + .map(|bytes| bytes / (1024 * 1024)) +} + +fn free_disk_mib(host: &dyn Host) -> Option { + // POSIX `df -Pk`: header line, then `fs blocks used avail cap mount`. + let out = host.run("df", &["-Pk", DISK_PROBE_PATH])?; + let avail_kib = out + .lines() + .nth(1)? + .split_whitespace() + .nth(3)? + .parse::() + .ok()?; + Some(avail_kib / 1024) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + + #[derive(Default)] + struct FakeHost { + files: HashMap<&'static str, &'static str>, + commands: HashMap<&'static str, &'static str>, + os: &'static str, + cpus: Option, + } + + impl Host for FakeHost { + fn read_to_string(&self, path: &Path) -> Option { + self.files.get(path.to_str()?).map(|s| (*s).to_owned()) + } + fn exists(&self, path: &Path) -> bool { + path.to_str().is_some_and(|p| self.files.contains_key(p)) + } + fn run(&self, program: &str, _args: &[&str]) -> Option { + self.commands.get(program).map(|s| (*s).to_owned()) + } + fn os(&self) -> &'static str { + self.os + } + fn arch(&self) -> &'static str { + "x86_64" + } + fn available_parallelism(&self) -> Option { + self.cpus + } + } + + fn linux() -> FakeHost { + FakeHost { + os: "linux", + cpus: Some(16), + ..FakeHost::default() + } + } + + const MEMINFO_28G: &str = "MemTotal: 29360128 kB\nMemFree: 1 kB\n"; + + #[test] + fn detects_wsl_from_interop_or_kernel_string() { + let mut h = linux(); + h.files.insert("/proc/sys/fs/binfmt_misc/WSLInterop", ""); + assert_eq!(host_kind(&h), HostKind::Wsl); + + let mut h = linux(); + h.files.insert( + "/proc/version", + "Linux version 5.15.0-microsoft-standard-WSL2", + ); + assert_eq!(host_kind(&h), HostKind::Wsl); + } + + #[test] + fn detects_lxc_from_systemd_or_init_environ() { + let mut h = linux(); + h.files.insert("/run/systemd/container", "lxc\n"); + assert_eq!(host_kind(&h), HostKind::Lxc); + + let mut h = linux(); + h.files + .insert("/proc/1/environ", "PATH=/bin\0container=lxc\0TERM=linux\0"); + assert_eq!(host_kind(&h), HostKind::Lxc); + } + + #[test] + fn detects_vm_from_dmi_or_hypervisor_flag() { + let mut h = linux(); + h.files.insert("/sys/class/dmi/id/sys_vendor", "QEMU\n"); + assert_eq!(host_kind(&h), HostKind::Vm); + + let mut h = linux(); + h.files.insert( + "/proc/cpuinfo", + "processor : 0\nflags : fpu vme hypervisor lm\n", + ); + assert_eq!(host_kind(&h), HostKind::Vm); + } + + #[test] + fn detects_macos_vm_and_bare() { + let mut h = FakeHost { + os: "macos", + cpus: Some(8), + ..FakeHost::default() + }; + assert_eq!(host_kind(&h), HostKind::Bare); + h.commands.insert("sysctl", "1"); + assert_eq!(host_kind(&h), HostKind::Vm); + } + + #[test] + fn bare_metal_is_the_fallback_and_wsl_wins_over_vm() { + assert_eq!(host_kind(&linux()), HostKind::Bare); + + let mut h = linux(); + h.files.insert("/run/WSL", ""); + h.files.insert("/proc/cpuinfo", "flags : hypervisor\n"); + assert_eq!(host_kind(&h), HostKind::Wsl); + } + + #[test] + fn memory_comes_from_meminfo() { + let mut h = linux(); + h.files.insert("/proc/meminfo", MEMINFO_28G); + assert_eq!(mem_mib(&h), Some(28672)); + } + + #[test] + fn cgroup_limit_clamps_memory_but_max_does_not() { + let mut h = linux(); + h.files.insert("/proc/meminfo", MEMINFO_28G); + h.files.insert("/sys/fs/cgroup/memory.max", "8589934592\n"); + assert_eq!(mem_mib(&h), Some(8192)); + + h.files.insert("/sys/fs/cgroup/memory.max", "max\n"); + assert_eq!(mem_mib(&h), Some(28672)); + } + + #[test] + fn macos_memory_comes_from_sysctl() { + let mut h = FakeHost { + os: "macos", + cpus: Some(8), + ..FakeHost::default() + }; + h.commands.insert("sysctl", "17179869184"); + assert_eq!(mem_mib(&h), Some(16384)); + } + + #[test] + fn free_disk_parses_posix_df() { + let mut h = linux(); + h.commands.insert( + "df", + "Filesystem 1024-blocks Used Available Capacity Mounted on\n/dev/sda1 100000000 1000000 52428800 2% /", + ); + assert_eq!(free_disk_mib(&h), Some(51200)); + } + + #[test] + fn probe_assembles_capacity_and_degrades_per_field() { + let mut h = linux(); + h.files.insert("/proc/meminfo", MEMINFO_28G); + h.files.insert("/proc/sys/kernel/hostname", "host-a\n"); + h.commands.insert("docker", "27.0.1"); + h.commands.insert("git", "git version 2.45.0"); + + let c = probe(&h); + assert_eq!(c.hostname, "host-a"); + assert_eq!((c.vcpu, c.mem_mib), (16, 28672)); + assert!(c.docker_usable); + assert_eq!(c.git_version.as_deref(), Some("git version 2.45.0")); + assert_eq!(c.mise_version, None); + assert_eq!(c.free_disk_mib, None); + assert_eq!(c.kind, HostKind::Bare); + } + + #[test] + fn probe_of_an_empty_host_still_reports_something() { + let c = probe(&FakeHost { + os: "linux", + ..FakeHost::default() + }); + assert_eq!((c.hostname.as_str(), c.vcpu, c.mem_mib), ("unknown", 1, 0)); + assert!(!c.docker_usable); + } + + #[test] + fn capacity_serializes_as_json() { + let c = probe(&linux()); + let v: serde_json::Value = serde_json::to_value(&c).unwrap(); + assert_eq!(v["kind"], "bare"); + assert_eq!(v["vcpu"], 16); + } +} diff --git a/src/probe/system.rs b/src/probe/system.rs new file mode 100644 index 0000000..1d2f48f --- /dev/null +++ b/src/probe/system.rs @@ -0,0 +1,103 @@ +//! The real [`Host`]: reads the local filesystem and runs read-only commands. + +use std::fs; +use std::io::Read; +use std::path::Path; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +use super::Host; + +/// Longest any probe command may run. `docker info` can hang on a wedged +/// daemon, and a probe must never block the heartbeat. +const COMMAND_TIMEOUT: Duration = Duration::from_secs(10); +const POLL_INTERVAL: Duration = Duration::from_millis(20); + +/// The machine borg is running on. +pub struct SystemHost; + +impl Host for SystemHost { + fn read_to_string(&self, path: &Path) -> Option { + fs::read_to_string(path).ok() + } + + fn exists(&self, path: &Path) -> bool { + path.exists() + } + + fn run(&self, program: &str, args: &[&str]) -> Option { + run_bounded(program, args) + } + + fn os(&self) -> &'static str { + std::env::consts::OS + } + + fn arch(&self) -> &'static str { + std::env::consts::ARCH + } + + fn available_parallelism(&self) -> Option { + std::thread::available_parallelism() + .ok() + .and_then(|n| u32::try_from(n.get()).ok()) + } +} + +/// Run a read-only probe command and return trimmed stdout when it exits 0 +/// within [`COMMAND_TIMEOUT`]. Stdin is closed and stderr discarded. +fn run_bounded(program: &str, args: &[&str]) -> Option { + #[allow( + clippy::disallowed_methods, + reason = "borg probe runs fixed read-only host-inspection commands (git, docker, mise, df, sysctl, hostname); the program name is chosen by this crate, never by remote input" + )] + let mut child = Command::new(program) + .args(args) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .spawn() + .ok()?; + + let deadline = Instant::now() + COMMAND_TIMEOUT; + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break status, + Ok(None) if Instant::now() < deadline => std::thread::sleep(POLL_INTERVAL), + Ok(None) | Err(_) => { + // Best effort: the child is ours, so kill and reap it. + let _ = child.kill(); + let _ = child.wait(); + return None; + } + } + }; + if !status.success() { + return None; + } + let mut out = String::new(); + child.stdout.take()?.read_to_string(&mut out).ok()?; + Some(out.trim().to_owned()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn missing_program_is_none() { + assert_eq!(run_bounded("borg-no-such-program-4f1a", &[]), None); + } + + #[cfg(unix)] + #[test] + fn captures_stdout_of_successful_command() { + assert_eq!(run_bounded("echo", &["hi"]).as_deref(), Some("hi")); + } + + #[cfg(unix)] + #[test] + fn failing_command_is_none() { + assert_eq!(run_bounded("false", &[]), None); + } +} diff --git a/src/queen.rs b/src/queen.rs new file mode 100644 index 0000000..929127a --- /dev/null +++ b/src/queen.rs @@ -0,0 +1,413 @@ +//! The queen: one reconciliation pass ([`Queen::tick`]) turns the GitHub +//! queue plus drone heartbeats into minted leases. +//! +//! A tick: +//! 1. lists the organization's runners and reconciles every active lease with +//! what GitHub reports (busy, gone, never came online); +//! 2. deletes stale `-*` registrations that no lease owns; +//! 3. reads queued jobs and derives demand and queue ages per label; +//! 4. runs the pure allocator and mints a JIT runner for each grant. +//! +//! Any GitHub failure before minting aborts the tick without minting: acting +//! on a partial view of the queue would over- or under-grant. + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::RwLock; + +use rand::RngCore; +use rand::rngs::OsRng; +use serde::Serialize; +use time::OffsetDateTime; + +use crate::allocator::{ + ActiveLease, AllocatorInput, Demand, DroneState, LeaseState, QueueAge, Reservation, Unmet, + allocate, max_slots, +}; +use crate::classes::{Class, Scope}; +use crate::config::QueenConfig; +use crate::github::{GithubClient, GithubError, QueuedJob, Runner}; +use crate::probe::Capacity; +use crate::store::{LeaseRow, LeaseStatus, Store, StoreError}; + +#[derive(Debug, thiserror::Error)] +pub enum QueenError { + #[error(transparent)] + Store(#[from] StoreError), + #[error(transparent)] + Github(#[from] GithubError), +} + +/// What the last tick saw. Served by `/status` and `/metrics`. +#[derive(Debug, Clone, Default, Serialize)] +pub struct Snapshot { + pub last_tick_at: Option, + pub last_error: Option, + pub queue: Vec, + pub reserved_for: Vec, + pub unmet: Vec, +} + +#[derive(Debug, Clone, Serialize)] +pub struct QueueEntry { + pub label: String, + pub class: Class, + pub scope: Scope, + pub queued: u32, + pub oldest_age_secs: u64, +} + +pub struct Queen { + pub(crate) cfg: QueenConfig, + pub(crate) store: Store, + github: GithubClient, + snapshot: RwLock, + tick_lock: tokio::sync::Mutex<()>, +} + +/// The full label set a JIT runner is registered with: the labels GitHub +/// assigns to every self-hosted runner, then the grant's class labels. They +/// are sent explicitly because the API documents only "custom" labels. +pub fn jit_labels(os: &str, arch: &str, class_labels: &[String]) -> Vec { + let mut labels = vec!["self-hosted".to_owned()]; + match os { + "linux" => labels.push("Linux".to_owned()), + "macos" => labels.push("macOS".to_owned()), + "windows" => labels.push("Windows".to_owned()), + _ => {} + } + match arch { + "x86_64" | "x86" => labels.push("X64".to_owned()), + "aarch64" | "arm64" => labels.push("ARM64".to_owned()), + "arm" => labels.push("ARM".to_owned()), + _ => {} + } + labels.extend(class_labels.iter().cloned()); + labels +} + +fn new_lease_id() -> String { + let mut b = [0u8; 6]; + OsRng.fill_bytes(&mut b); + hex::encode(b) +} + +fn age_secs(now: OffsetDateTime, then: OffsetDateTime) -> u64 { + u64::try_from((now - then).whole_seconds()).unwrap_or(0) +} + +impl Queen { + pub fn new(cfg: QueenConfig, store: Store, github: GithubClient) -> Self { + Self { + cfg, + store, + github, + snapshot: RwLock::new(Snapshot::default()), + tick_lock: tokio::sync::Mutex::new(()), + } + } + + pub fn config(&self) -> &QueenConfig { + &self.cfg + } + + pub fn store(&self) -> &Store { + &self.store + } + + pub fn snapshot(&self) -> Snapshot { + self.snapshot.read().map(|s| s.clone()).unwrap_or_default() + } + + fn update_snapshot(&self, f: impl FnOnce(&mut Snapshot)) { + if let Ok(mut s) = self.snapshot.write() { + f(&mut s); + } + } + + pub fn record_tick_error(&self, now: OffsetDateTime, err: &QueenError) { + self.update_snapshot(|s| { + s.last_tick_at = Some(now.unix_timestamp()); + s.last_error = Some(err.to_string()); + }); + } + + fn runner_prefix(&self) -> String { + format!("{}-", self.cfg.runner_prefix) + } + + /// One reconciliation pass. + pub async fn tick(&self, now: OffsetDateTime) -> Result<(), QueenError> { + let _serial = self.tick_lock.lock().await; + let now_ts = now.unix_timestamp(); + + let runners = self.github.list_runners().await?; + let prefix = self.runner_prefix(); + let ours: HashMap<&str, &Runner> = runners + .iter() + .filter(|r| r.name.starts_with(&prefix)) + .map(|r| (r.name.as_str(), r)) + .collect(); + + let active = self.reconcile(&ours, now_ts).await?; + self.reap(&ours, &active).await; + + let mut jobs = Vec::new(); + for repo in &self.cfg.repos { + jobs.extend(self.github.queued_jobs(repo).await?); + } + let (queue, demand, queue_ages) = self.summarize(&jobs, now); + + let drone_rows = self.store.drones().await?; + let mut capacities: HashMap = HashMap::new(); + let mut drones = Vec::new(); + for d in &drone_rows { + let fresh = d + .last_heartbeat + .is_some_and(|t| now_ts - t <= self.cfg.drone_stale_secs); + if let (true, Some(cap)) = (fresh, &d.capacity) { + let mut state = DroneState::from_capacity(d.name.clone(), cap, &d.caps); + state.accepting = d.accepting; + drones.push(state); + capacities.insert(d.name.clone(), cap.clone()); + } + } + let leases = active + .iter() + .map(|l| ActiveLease { + drone_id: l.drone.clone(), + class: l.class, + scope: l.scope, + state: if l.state == LeaseStatus::Running { + LeaseState::Running + } else { + LeaseState::Pending + }, + }) + .collect(); + let allocation = allocate(&AllocatorInput { + classes: self.cfg.classes.clone(), + demand, + drones, + leases, + warm_floors: self.cfg.warm_floors.clone(), + queue_ages, + starvation_after_secs: self.cfg.starvation_after_secs, + }); + + let mut mint_error = None; + for grant in &allocation.grants { + let Some(cap) = capacities.get(&grant.drone_id) else { + continue; + }; + let lease_id = new_lease_id(); + let runner_name = format!("{prefix}{}-{lease_id}", grant.drone_id); + let labels = jit_labels(&cap.os, &cap.arch, &grant.labels); + let minted = match self + .github + .mint_jit_runner( + &runner_name, + self.cfg.github.runner_group_id, + &labels, + &self.cfg.work_folder, + ) + .await + { + Ok(m) => m, + Err(e) => { + // Likely rate limiting or permissions: stop minting this tick. + mint_error = Some(e); + break; + } + }; + let row = LeaseRow { + id: lease_id, + drone: grant.drone_id.clone(), + class: grant.class, + scope: grant.scope, + labels, + runner_name, + runner_id: minted.runner_id, + jit_config: Some(minted.jit_config), + state: LeaseStatus::Pending, + created_at: now_ts, + expires_at: now_ts + self.cfg.lease_ttl_secs, + }; + if let Err(e) = self.store.insert_lease(&row).await { + // Do not leave an unowned registration behind. + let _ = self.github.delete_runner(minted.runner_id).await; + return Err(e.into()); + } + } + + self.update_snapshot(|s| { + s.last_tick_at = Some(now_ts); + s.last_error = mint_error.as_ref().map(|e| format!("mint failed: {e}")); + s.queue = queue; + s.reserved_for = allocation.reserved_for.clone(); + s.unmet = allocation.unmet.clone(); + }); + Ok(()) + } + + /// Bring lease state in line with GitHub; returns the leases still active. + async fn reconcile( + &self, + ours: &HashMap<&str, &Runner>, + now: i64, + ) -> Result, QueenError> { + let mut still_active = Vec::new(); + for lease in self.store.active_leases().await? { + let runner = ours.get(lease.runner_name.as_str()); + let overdue = now - lease.created_at > self.cfg.lease_max_secs; + let never_came_online = + now > lease.expires_at && runner.is_none_or(|r| r.status != "online"); + match runner { + _ if overdue => { + self.end(&lease, runner.copied(), LeaseStatus::Expired, now) + .await? + } + Some(r) if r.busy => { + if lease.state == LeaseStatus::Pending { + self.store + .set_state(&lease.id, LeaseStatus::Running, now) + .await?; + } + still_active.push(lease); + } + // Ephemeral runners deregister after their one job. + None if lease.state == LeaseStatus::Running => { + self.store + .set_state(&lease.id, LeaseStatus::Done, now) + .await?; + } + _ if never_came_online => { + self.end(&lease, runner.copied(), LeaseStatus::Expired, now) + .await? + } + _ => still_active.push(lease), + } + } + Ok(still_active) + } + + async fn end( + &self, + lease: &LeaseRow, + runner: Option<&Runner>, + to: LeaseStatus, + now: i64, + ) -> Result<(), QueenError> { + self.store.set_state(&lease.id, to, now).await?; + if let Some(r) = runner { + // Best effort: the reaper retries next tick. + if let Err(e) = self.github.delete_runner(r.id).await { + tracing::warn!("delete runner {} failed: {e}", r.name); + } + } + Ok(()) + } + + /// Delete offline `-*` registrations that no active lease owns. + async fn reap(&self, ours: &HashMap<&str, &Runner>, active: &[LeaseRow]) { + let owned: HashSet<&str> = active.iter().map(|l| l.runner_name.as_str()).collect(); + for (name, r) in ours { + if !owned.contains(name) + && r.status == "offline" + && !r.busy + && let Err(e) = self.github.delete_runner(r.id).await + { + tracing::warn!("reap runner {name} failed: {e}"); + } + } + } + + fn summarize( + &self, + jobs: &[QueuedJob], + now: OffsetDateTime, + ) -> (Vec, Vec, Vec) { + let mut by: BTreeMap<(Class, Scope), (u32, u64)> = BTreeMap::new(); + for job in jobs { + let Some((class, scope)) = job + .labels + .iter() + .find_map(|l| self.cfg.classes.parse_label(l)) + else { + continue; + }; + let e = by.entry((class, scope)).or_default(); + e.0 += 1; + e.1 = e.1.max(age_secs(now, job.created_at)); + } + let queue = by + .iter() + .map(|(&(class, scope), &(queued, oldest_age_secs))| QueueEntry { + label: self.cfg.classes.job_label(class, scope), + class, + scope, + queued, + oldest_age_secs, + }) + .collect(); + let demand = by + .iter() + .map(|(&(class, scope), &(count, _))| Demand { + class, + scope, + count, + }) + .collect(); + let ages = by + .iter() + .map(|(&(class, scope), &(_, oldest_age_secs))| QueueAge { + class, + scope, + oldest_age_secs, + }) + .collect(); + (queue, demand, ages) + } + + /// Per-class totals for status and metrics: slots the fleet could hold. + pub(crate) fn class_slots(&self, drones: &[DroneState]) -> BTreeMap { + Class::ALL + .into_iter() + .map(|c| { + let spec = self.cfg.classes.spec(c); + ( + c, + drones + .iter() + .filter(|d| d.accepting) + .map(|d| max_slots(d, spec)) + .sum(), + ) + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn jit_labels_list_platform_labels_first_then_class_labels() { + let class = vec!["borg-ci-medium".to_owned(), "borg-ci-medium-pr".to_owned()]; + assert_eq!( + jit_labels("linux", "x86_64", &class), + [ + "self-hosted", + "Linux", + "X64", + "borg-ci-medium", + "borg-ci-medium-pr" + ] + ); + assert_eq!( + jit_labels("macos", "aarch64", &class[..1]), + ["self-hosted", "macOS", "ARM64", "borg-ci-medium"] + ); + assert_eq!(jit_labels("plan9", "mips", &[]), ["self-hosted"]); + } +} diff --git a/src/secret.rs b/src/secret.rs new file mode 100644 index 0000000..97e9a98 --- /dev/null +++ b/src/secret.rs @@ -0,0 +1,42 @@ +//! A string that never prints itself. + +use serde::{Deserialize, Serialize}; + +/// Secret text (installation tokens, JIT runner configs). `Debug` is +/// redacted; serialization is transparent because the value must reach the +/// one party entitled to it (GitHub or the leased drone). +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct Secret(String); + +impl Secret { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn expose(&self) -> &str { + &self.0 + } +} + +impl std::fmt::Debug for Secret { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("Secret([redacted])") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn debug_is_redacted_but_serde_is_transparent() { + let s = Secret::new("hunter2"); + assert!(!format!("{s:?}").contains("hunter2")); + assert_eq!(serde_json::to_string(&s).unwrap(), "\"hunter2\""); + assert_eq!( + serde_json::from_str::("\"x\"").unwrap().expose(), + "x" + ); + } +} diff --git a/src/store.rs b/src/store.rs new file mode 100644 index 0000000..0161d97 --- /dev/null +++ b/src/store.rs @@ -0,0 +1,522 @@ +//! SQLite bookkeeping: enrolled drones and lease history. +//! +//! The file is small and disposable in the sense that the queen re-derives +//! truth from GitHub plus drone heartbeats after a restart. Enrollment hashes +//! are the one thing worth backing up. + +use std::path::Path; +use std::str::FromStr; + +use serde::Serialize; +use serde::de::DeserializeOwned; +use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; +use sqlx::{Row, SqlitePool}; + +use crate::allocator::HostCaps; +use crate::classes::{Class, Scope}; +use crate::probe::Capacity; +use crate::secret::Secret; + +const SCHEMA: &str = " +CREATE TABLE IF NOT EXISTS drones ( + name TEXT PRIMARY KEY, + token_hash TEXT NOT NULL UNIQUE, + enrolled_at INTEGER NOT NULL, + last_heartbeat INTEGER, + capacity TEXT, + caps TEXT, + accepting INTEGER NOT NULL DEFAULT 1 +); +CREATE TABLE IF NOT EXISTS leases ( + id TEXT PRIMARY KEY, + drone TEXT NOT NULL REFERENCES drones(name), + class TEXT NOT NULL, + scope TEXT NOT NULL, + labels TEXT NOT NULL, + runner_name TEXT NOT NULL, + runner_id INTEGER NOT NULL, + jit_config TEXT, + state TEXT NOT NULL, + created_at INTEGER NOT NULL, + expires_at INTEGER NOT NULL, + finished_at INTEGER, + exit_code INTEGER, + duration_secs INTEGER, + peak_rss_mib INTEGER +); +CREATE INDEX IF NOT EXISTS leases_state ON leases(state); +"; + +#[derive(Debug, thiserror::Error)] +pub enum StoreError { + #[error("database: {0}")] + Db(#[from] sqlx::Error), + #[error("stored value is corrupt: {0}")] + Corrupt(String), + #[error("drone {0} is already enrolled")] + AlreadyEnrolled(String), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum LeaseStatus { + /// Runner minted, no job yet. + Pending, + /// Runner is executing a job. + Running, + Done, + Expired, +} + +impl LeaseStatus { + pub fn as_str(self) -> &'static str { + match self { + Self::Pending => "pending", + Self::Running => "running", + Self::Done => "done", + Self::Expired => "expired", + } + } + + fn parse(s: &str) -> Result { + Ok(match s { + "pending" => Self::Pending, + "running" => Self::Running, + "done" => Self::Done, + "expired" => Self::Expired, + other => return Err(StoreError::Corrupt(format!("lease state {other}"))), + }) + } + + pub fn is_active(self) -> bool { + matches!(self, Self::Pending | Self::Running) + } +} + +#[derive(Debug, Clone)] +pub struct DroneRow { + pub name: String, + pub last_heartbeat: Option, + pub capacity: Option, + pub caps: HostCaps, + pub accepting: bool, +} + +#[derive(Debug, Clone)] +pub struct LeaseRow { + pub id: String, + pub drone: String, + pub class: Class, + pub scope: Scope, + pub labels: Vec, + pub runner_name: String, + pub runner_id: u64, + /// Present until the drone acknowledges the lease. + pub jit_config: Option, + pub state: LeaseStatus, + pub created_at: i64, + pub expires_at: i64, +} + +#[derive(Debug, Clone)] +pub struct LeaseResult { + pub exit_code: i32, + pub duration_secs: u64, + pub peak_rss_mib: Option, +} + +#[derive(Clone)] +pub struct Store { + pool: SqlitePool, +} + +fn to_text(v: &T) -> Result { + serde_json::to_string(v).map_err(|e| StoreError::Corrupt(e.to_string())) +} + +fn from_text(s: &str) -> Result { + serde_json::from_str(s).map_err(|e| StoreError::Corrupt(e.to_string())) +} + +fn enum_text(v: &T) -> Result { + match serde_json::to_value(v) { + Ok(serde_json::Value::String(s)) => Ok(s), + _ => Err(StoreError::Corrupt("enum is not a string".to_owned())), + } +} + +fn enum_from(s: &str) -> Result { + from_text(&format!("\"{s}\"")) +} + +impl Store { + pub async fn open(path: &Path) -> Result { + let opts = SqliteConnectOptions::new() + .filename(path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal); + Self::connect(opts, 4).await + } + + pub async fn open_in_memory() -> Result { + // One connection: every pooled in-memory connection is its own database. + let opts = SqliteConnectOptions::from_str("sqlite::memory:")?; + Self::connect(opts, 1).await + } + + async fn connect(opts: SqliteConnectOptions, max: u32) -> Result { + let pool = SqlitePoolOptions::new() + .max_connections(max) + .connect_with(opts) + .await?; + sqlx::raw_sql(SCHEMA).execute(&pool).await?; + Ok(Self { pool }) + } + + pub async fn enroll(&self, name: &str, token_hash: &str, now: i64) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let r = sqlx::query("INSERT INTO drones (name, token_hash, enrolled_at) VALUES (?, ?, ?)") + .bind(name) + .bind(token_hash) + .bind(now) + .execute(&self.pool) + .await; + match r { + Ok(_) => Ok(()), + Err(sqlx::Error::Database(e)) if e.is_unique_violation() => { + Err(StoreError::AlreadyEnrolled(name.to_owned())) + } + Err(e) => Err(e.into()), + } + } + + pub async fn drone_for_token_hash( + &self, + token_hash: &str, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT name FROM drones WHERE token_hash = ?") + .bind(token_hash) + .fetch_optional(&self.pool) + .await?; + Ok(row.map(|r| r.get::("name"))) + } + + pub async fn record_heartbeat( + &self, + name: &str, + capacity: &Capacity, + caps: &HostCaps, + accepting: bool, + now: i64, + ) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE drones SET last_heartbeat = ?, capacity = ?, caps = ?, accepting = ? WHERE name = ?") + .bind(now) + .bind(to_text(capacity)?) + .bind(to_text(caps)?) + .bind(i64::from(accepting)) + .bind(name) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn drones(&self) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT name, last_heartbeat, capacity, caps, accepting FROM drones ORDER BY name", + ) + .fetch_all(&self.pool) + .await?; + rows.iter() + .map(|r| { + let capacity: Option = r.get("capacity"); + let caps: Option = r.get("caps"); + Ok(DroneRow { + name: r.get("name"), + last_heartbeat: r.get("last_heartbeat"), + capacity: capacity.as_deref().map(from_text).transpose()?, + caps: caps + .as_deref() + .map(from_text) + .transpose()? + .unwrap_or_default(), + accepting: r.get::("accepting") != 0, + }) + }) + .collect() + } + + pub async fn insert_lease(&self, l: &LeaseRow) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT INTO leases (id, drone, class, scope, labels, runner_name, runner_id, jit_config, state, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&l.id) + .bind(&l.drone) + .bind(enum_text(&l.class)?) + .bind(enum_text(&l.scope)?) + .bind(to_text(&l.labels)?) + .bind(&l.runner_name) + .bind(i64::try_from(l.runner_id).map_err(|e| StoreError::Corrupt(e.to_string()))?) + .bind(l.jit_config.as_ref().map(Secret::expose)) + .bind(l.state.as_str()) + .bind(l.created_at) + .bind(l.expires_at) + .execute(&self.pool) + .await?; + Ok(()) + } + + fn lease_from_row(r: &sqlx::sqlite::SqliteRow) -> Result { + let jit: Option = r.get("jit_config"); + Ok(LeaseRow { + id: r.get("id"), + drone: r.get("drone"), + class: enum_from(&r.get::("class"))?, + scope: enum_from(&r.get::("scope"))?, + labels: from_text(&r.get::("labels"))?, + runner_name: r.get("runner_name"), + runner_id: u64::try_from(r.get::("runner_id")) + .map_err(|e| StoreError::Corrupt(e.to_string()))?, + jit_config: jit.map(Secret::new), + state: LeaseStatus::parse(&r.get::("state"))?, + created_at: r.get("created_at"), + expires_at: r.get("expires_at"), + }) + } + + /// Leases that are pending or running. + pub async fn active_leases(&self) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT * FROM leases WHERE state IN ('pending', 'running') ORDER BY created_at, id", + ) + .fetch_all(&self.pool) + .await?; + rows.iter().map(Self::lease_from_row).collect() + } + + pub async fn lease(&self, id: &str) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT * FROM leases WHERE id = ?") + .bind(id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(Self::lease_from_row).transpose() + } + + /// Move an active lease to `to`. Finished leases never change again. + pub async fn set_state(&self, id: &str, to: LeaseStatus, now: i64) -> Result<(), StoreError> { + let finished = (!to.is_active()).then_some(now); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "UPDATE leases SET state = ?, finished_at = COALESCE(?, finished_at), jit_config = CASE WHEN ? THEN NULL ELSE jit_config END + WHERE id = ? AND state IN ('pending', 'running')", + ) + .bind(to.as_str()) + .bind(finished) + .bind(!to.is_active()) + .bind(id) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Record a drone's result for its own lease. Returns false if `drone` + /// does not own `id`. Re-reporting a finished lease is accepted. + pub async fn finish_lease( + &self, + id: &str, + drone: &str, + result: &LeaseResult, + now: i64, + ) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let owned = sqlx::query("SELECT 1 FROM leases WHERE id = ? AND drone = ?") + .bind(id) + .bind(drone) + .fetch_optional(&self.pool) + .await? + .is_some(); + if !owned { + return Ok(false); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "UPDATE leases SET state = 'done', finished_at = ?, jit_config = NULL, exit_code = ?, duration_secs = ?, peak_rss_mib = ? + WHERE id = ? AND state IN ('pending', 'running', 'expired')", + ) + .bind(now) + .bind(i64::from(result.exit_code)) + .bind(i64::try_from(result.duration_secs).unwrap_or(i64::MAX)) + .bind(result.peak_rss_mib.map(|v| i64::try_from(v).unwrap_or(i64::MAX))) + .bind(id) + .execute(&self.pool) + .await?; + Ok(true) + } + + /// Drop the JIT config of leases `drone` says it holds. + pub async fn acknowledge(&self, drone: &str, lease_ids: &[String]) -> Result<(), StoreError> { + for id in lease_ids { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE leases SET jit_config = NULL WHERE id = ? AND drone = ?") + .bind(id) + .bind(drone) + .execute(&self.pool) + .await?; + } + Ok(()) + } + + /// Active leases of `drone` whose JIT config has not been acknowledged. + pub async fn undelivered(&self, drone: &str) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT * FROM leases WHERE drone = ? AND state = 'pending' AND jit_config IS NOT NULL ORDER BY created_at, id", + ) + .bind(drone) + .fetch_all(&self.pool) + .await?; + rows.iter().map(Self::lease_from_row).collect() + } + + /// Lease counts by (class, state), for metrics. + pub async fn lease_counts(&self) -> Result, StoreError> { + let rows = + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("SELECT class, state, COUNT(*) AS n FROM leases GROUP BY class, state") + .fetch_all(&self.pool) + .await?; + rows.iter() + .map(|r| { + Ok(( + enum_from(&r.get::("class"))?, + LeaseStatus::parse(&r.get::("state"))?, + u32::try_from(r.get::("n")).unwrap_or(u32::MAX), + )) + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn lease(id: &str, drone: &str) -> LeaseRow { + LeaseRow { + id: id.into(), + drone: drone.into(), + class: Class::Medium, + scope: Scope::Pr, + labels: vec!["a".into(), "b".into()], + runner_name: format!("borg-{drone}-{id}"), + runner_id: 7, + jit_config: Some(Secret::new("jit")), + state: LeaseStatus::Pending, + created_at: 100, + expires_at: 400, + } + } + + async fn store_with(drone: &str) -> Store { + let s = Store::open_in_memory().await.unwrap(); + s.enroll(drone, &format!("hash-{drone}"), 1).await.unwrap(); + s + } + + #[tokio::test] + async fn enrollment_is_unique_and_lookup_is_by_hash() { + let s = store_with("d1").await; + assert_eq!( + s.drone_for_token_hash("hash-d1").await.unwrap().as_deref(), + Some("d1") + ); + assert_eq!(s.drone_for_token_hash("nope").await.unwrap(), None); + assert!(matches!( + s.enroll("d1", "other", 2).await, + Err(StoreError::AlreadyEnrolled(_)) + )); + assert!(matches!( + s.enroll("d2", "hash-d1", 2).await, + Err(StoreError::AlreadyEnrolled(_)) + )); + } + + #[tokio::test] + async fn jit_config_is_delivered_until_acknowledged() { + let s = store_with("d1").await; + s.insert_lease(&lease("l1", "d1")).await.unwrap(); + assert_eq!(s.undelivered("d1").await.unwrap().len(), 1); + s.acknowledge("d1", &["l1".to_owned()]).await.unwrap(); + assert!(s.undelivered("d1").await.unwrap().is_empty()); + assert_eq!( + s.active_leases().await.unwrap().len(), + 1, + "ack does not end the lease" + ); + } + + #[tokio::test] + async fn only_the_owning_drone_can_finish_and_finishing_is_idempotent() { + let s = store_with("d1").await; + s.enroll("d2", "hash-d2", 1).await.unwrap(); + s.insert_lease(&lease("l1", "d1")).await.unwrap(); + let r = LeaseResult { + exit_code: 0, + duration_secs: 5, + peak_rss_mib: Some(100), + }; + assert!(!s.finish_lease("l1", "d2", &r, 500).await.unwrap()); + assert!(s.finish_lease("l1", "d1", &r, 500).await.unwrap()); + assert!(s.finish_lease("l1", "d1", &r, 600).await.unwrap()); + let l = s.lease("l1").await.unwrap().unwrap(); + assert_eq!(l.state, LeaseStatus::Done); + assert!(l.jit_config.is_none(), "secret dropped when the lease ends"); + assert!(s.active_leases().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn set_state_never_revives_a_finished_lease() { + let s = store_with("d1").await; + s.insert_lease(&lease("l1", "d1")).await.unwrap(); + s.set_state("l1", LeaseStatus::Running, 200).await.unwrap(); + assert_eq!( + s.lease("l1").await.unwrap().unwrap().state, + LeaseStatus::Running + ); + s.set_state("l1", LeaseStatus::Expired, 300).await.unwrap(); + s.set_state("l1", LeaseStatus::Running, 400).await.unwrap(); + assert_eq!( + s.lease("l1").await.unwrap().unwrap().state, + LeaseStatus::Expired + ); + } + + #[tokio::test] + async fn heartbeat_round_trips_capacity() { + let s = store_with("d1").await; + let cap = Capacity { + hostname: "h".into(), + os: "linux".into(), + arch: "x86_64".into(), + kind: crate::probe::HostKind::Vm, + vcpu: 4, + mem_mib: 8192, + free_disk_mib: Some(10), + docker_usable: true, + mise_version: None, + git_version: None, + }; + s.record_heartbeat("d1", &cap, &HostCaps::default(), false, 123) + .await + .unwrap(); + let d = &s.drones().await.unwrap()[0]; + assert_eq!((d.last_heartbeat, d.accepting), (Some(123), false)); + assert_eq!(d.capacity.as_ref(), Some(&cap)); + } +} diff --git a/tests/drone.rs b/tests/drone.rs new file mode 100644 index 0000000..ee19b67 --- /dev/null +++ b/tests/drone.rs @@ -0,0 +1,291 @@ +mod support; + +use std::fs; +use std::io; +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use borg::drone::Drone; +use borg::drone::client::QueenClient; +use borg::drone::config::DroneConfig; +use borg::drone::exec::{Exec, Output}; +use borg::probe::Host; +use borg::secret::Secret; +use serde_json::{Value, json}; +use support::Harness; +use tokio::task::JoinSet; + +/// A 30 GiB, 16 vCPU Linux host with a working docker. +struct FixedHost; + +impl Host for FixedHost { + fn read_to_string(&self, path: &Path) -> Option { + (path == Path::new("/proc/meminfo")).then(|| "MemTotal: 31457280 kB\n".to_owned()) + } + fn exists(&self, _path: &Path) -> bool { + false + } + fn run(&self, program: &str, _args: &[&str]) -> Option { + (program == "docker").then(|| "27.0".to_owned()) + } + fn os(&self) -> &'static str { + "linux" + } + fn arch(&self) -> &'static str { + "x86_64" + } + fn available_parallelism(&self) -> Option { + Some(16) + } +} + +/// Records host commands instead of running them. +#[derive(Default)] +struct RecordingExec(Mutex>); + +impl Exec for RecordingExec { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + self.0 + .lock() + .unwrap() + .push(format!("{program} {}", args.join(" "))); + Ok(Output { + success: true, + stdout: String::new(), + }) + } +} + +struct Rig { + h: Harness, + dir: tempfile::TempDir, + token: String, + exec: Arc, +} + +impl Rig { + async fn new() -> Self { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + Self { + h, + dir: tempfile::tempdir().unwrap(), + token, + exec: Arc::default(), + } + } + + fn path(&self, p: &str) -> PathBuf { + self.dir.path().join(p) + } + + /// A fake pinned runner install whose run.sh executes `body`. + fn install_runner(&self, body: &str) { + let runner = self.path("runner"); + fs::create_dir_all(runner.join("bin")).unwrap(); + let script = runner.join("run.sh"); + fs::write( + &script, + format!("#!/bin/sh\nOUT={}\n{body}\n", self.path("out").display()), + ) + .unwrap(); + fs::set_permissions(&script, fs::Permissions::from_mode(0o755)).unwrap(); + fs::create_dir_all(self.path("out")).unwrap(); + } + + fn drone(&self, extra: Value) -> Drone { + let mut cfg = json!({ + "queen_url": self.h.url, + "allow_insecure": true, + "runner_dir": self.path("runner"), + "work_root": self.path("work"), + "cache_root": self.path("cache"), + "heartbeat_secs": 1, + "drain_timeout_secs": 30, + }); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + cfg[k] = v; + } + let cfg = DroneConfig::from_json(&cfg.to_string()).unwrap(); + let client = QueenClient::new(&cfg, Secret::new(self.token.clone())).unwrap(); + Drone::new( + cfg, + client, + Arc::new(FixedHost), + Arc::clone(&self.exec) as Arc, + ) + } + + async fn lease_state(&self) -> Vec { + let s = self.h.status().await; + let c = &s["classes"]["light"]; + vec![format!( + "pending={} running={}", + c["leases_pending"], c["leases_running"] + )] + } +} + +async fn drain(tasks: &mut JoinSet<()>) { + while tasks.join_next().await.is_some() {} +} + +#[tokio::test] +async fn runs_a_granted_lease_cleans_up_and_reports_the_result() { + let rig = Rig::new().await; + rig.install_runner( + "echo \"$@\" > $OUT/args.txt\nprintf '%s' \"$ACTIONS_RUNNER_INPUT_JITCONFIG\" > $OUT/jit.txt\ntouch junk _diag.log\nexit 0", + ); + fs::create_dir_all(rig.path("cache/sccache")).unwrap(); + fs::write(rig.path("cache/sccache/entry"), "cached").unwrap(); + let drone = rig.drone(json!({})); + let mut tasks = JoinSet::new(); + + assert_eq!( + drone.heartbeat_once(true, &mut tasks).await.unwrap(), + 0, + "registers capacity, nothing granted yet" + ); + rig.h.queue_job(1, "borg-ci-light-pr", 3); + rig.h.tick().await; + assert_eq!(drone.heartbeat_once(true, &mut tasks).await.unwrap(), 1); + drain(&mut tasks).await; + + let name = rig.h.state().mints[0]["name"].as_str().unwrap().to_owned(); + assert_eq!( + fs::read_to_string(rig.path("out/jit.txt")).unwrap(), + format!("jit-{name}") + ); + assert_eq!( + fs::read_to_string(rig.path("out/args.txt")).unwrap().trim(), + "", + "the JIT config is never an argument" + ); + let lease_id = name.rsplit('-').next().unwrap().to_owned(); + assert!( + !rig.path("work").join(&lease_id).exists(), + "work dir removed after the job" + ); + assert_eq!( + fs::read_to_string(rig.path("cache/sccache/entry")).unwrap(), + "cached", + "caches survive cleanup" + ); + assert!( + rig.h + .queen + .store() + .active_leases() + .await + .unwrap() + .is_empty(), + "result reported, lease done" + ); + let calls = rig.exec.0.lock().unwrap().clone(); + assert!( + calls + .iter() + .any(|c| c.contains(&format!("--filter name=borg-{lease_id}"))), + "{calls:?}" + ); + assert!(calls.iter().all(|c| !c.contains("prune"))); + assert_eq!( + drone.heartbeat_once(true, &mut tasks).await.unwrap(), + 0, + "nothing new to start" + ); +} + +#[tokio::test] +async fn a_failing_job_is_still_cleaned_up_and_finished() { + let rig = Rig::new().await; + rig.install_runner("exit 3"); + let drone = rig.drone(json!({})); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + rig.h.queue_job(1, "borg-ci-light", 1); + rig.h.tick().await; + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + drain(&mut tasks).await; + assert!( + rig.h + .queen + .store() + .active_leases() + .await + .unwrap() + .is_empty() + ); + assert_eq!(fs::read_dir(rig.path("work")).unwrap().count(), 0); +} + +#[tokio::test] +async fn owner_lease_cap_is_advertised_and_respected() { + let rig = Rig::new().await; + rig.install_runner("sleep 1"); + let drone = rig.drone(json!({"max_leases": 1})); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + rig.h.queue_job(1, "borg-ci-light", 1); + rig.h.queue_job(2, "borg-ci-light", 1); + rig.h.tick().await; + // The drone advertised max_leases = 1, so the queen only granted one. + assert_eq!(rig.h.state().mints.len(), 1); + assert_eq!(drone.heartbeat_once(true, &mut tasks).await.unwrap(), 1); + assert_eq!( + drone.heartbeat_once(true, &mut tasks).await.unwrap(), + 0, + "already held: not started twice" + ); + drain(&mut tasks).await; + assert_eq!(rig.lease_state().await, ["pending=0 running=0"]); +} + +#[tokio::test] +async fn shutdown_stops_taking_leases_and_lets_the_running_job_finish() { + let rig = Rig::new().await; + rig.install_runner("touch $OUT/started\nsleep 2\ntouch $OUT/finished"); + fs::create_dir_all(rig.path("work/stale-lease")).unwrap(); + let drone = rig.drone(json!({})); + let (stop_tx, stop_rx) = tokio::sync::oneshot::channel::<()>(); + let run = tokio::spawn(drone.run(async move { + let _ = stop_rx.await; + })); + + // Wait for the first heartbeat to register, then give the queen a job. + tokio::time::sleep(Duration::from_millis(300)).await; + rig.h.queue_job(1, "borg-ci-light", 1); + rig.h.tick().await; + for _ in 0..50 { + if rig.path("out/started").exists() { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!(rig.path("out/started").exists(), "lease started"); + assert!( + !rig.path("work/stale-lease").exists(), + "startup sweep removed stale lease dirs" + ); + + stop_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(30), run) + .await + .unwrap() + .unwrap(); + assert!( + rig.path("out/finished").exists(), + "in-flight job ran to completion" + ); + assert!( + rig.h + .queen + .store() + .active_leases() + .await + .unwrap() + .is_empty() + ); +} diff --git a/tests/queen.rs b/tests/queen.rs new file mode 100644 index 0000000..75ee7e4 --- /dev/null +++ b/tests/queen.rs @@ -0,0 +1,314 @@ +mod support; + +use serde_json::{Value, json}; +use support::Harness; +use time::OffsetDateTime; + +fn labels(mint: &Value) -> Vec { + mint["labels"] + .as_array() + .unwrap() + .iter() + .map(|l| l.as_str().unwrap().to_owned()) + .collect() +} + +#[tokio::test] +async fn mints_a_pr_runner_with_exact_labels_and_delivers_its_config_until_acknowledged() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.queue_job(10, "borg-ci-medium-pr", 5); + assert!(h.leases(&token, &[]).await.is_empty()); + + h.tick().await; + + let mint = h.state().mints.clone(); + assert_eq!(mint.len(), 1); + assert_eq!( + labels(&mint[0]), + [ + "self-hosted", + "Linux", + "X64", + "borg-ci-medium", + "borg-ci-medium-pr" + ], + "exact label array sent to generate-jitconfig" + ); + assert_eq!(mint[0]["runner_group_id"], 1); + assert_eq!(mint[0]["work_folder"], "_work"); + let name = mint[0]["name"].as_str().unwrap().to_owned(); + assert!(name.starts_with("borg-d1-"), "{name}"); + + let leases = h.leases(&token, &[]).await; + assert_eq!(leases.len(), 1); + assert_eq!(leases[0]["jit_config"], format!("jit-{name}")); + assert_eq!(leases[0]["runner_name"], name); + assert_eq!(leases[0]["class"], "medium"); + let id = leases[0]["lease_id"].as_str().unwrap().to_owned(); + + assert_eq!( + h.leases(&token, &[]).await.len(), + 1, + "redelivered until acknowledged" + ); + assert!(h.leases(&token, &[&id]).await.is_empty(), "acknowledged"); + assert!(h.leases(&token, &[&id]).await.is_empty()); + + // The job is still queued but the pending lease already covers it. + h.tick().await; + assert_eq!( + h.state().mints.len(), + 1, + "no second mint for the same demand" + ); +} + +#[tokio::test] +async fn main_jobs_get_a_runner_with_only_the_bare_class_label() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-light", 5); + h.tick().await; + let mints = h.state().mints.clone(); + assert_eq!( + labels(&mints[0]), + ["self-hosted", "Linux", "X64", "borg-ci-light"] + ); +} + +#[tokio::test] +async fn default_warm_floor_keeps_one_idle_main_medium_runner() { + let h = + Harness::start(json!({"warm_floors": [{"class": "medium", "scope": "main", "count": 1}]})) + .await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.tick().await; + let mints = h.state().mints.clone(); + assert_eq!(mints.len(), 1); + assert_eq!( + labels(&mints[0]).last().map(String::as_str), + Some("borg-ci-medium") + ); + h.tick().await; + assert_eq!( + h.state().mints.len(), + 1, + "floor satisfied by the pending runner" + ); +} + +#[tokio::test] +async fn drone_api_requires_a_valid_token_and_ownership() { + let h = Harness::start(json!({})).await; + let t1 = h.enroll("d1").await; + let t2 = h.enroll("d2").await; + h.heartbeat(&t1, &[]).await; + h.queue_job(1, "borg-ci-light-pr", 1); + h.tick().await; + let id = h.leases(&t1, &[]).await[0]["lease_id"] + .as_str() + .unwrap() + .to_owned(); + + assert_eq!(h.heartbeat("borg_wrong", &[]).await.status(), 401); + let no_auth = h + .http + .post(format!("{}/v1/heartbeat", h.url)) + .json(&json!({"capacity": support::capacity()})) + .send() + .await + .unwrap(); + assert_eq!(no_auth.status(), 401); + + let url = format!("{}/v1/leases/{id}/result", h.url); + let body = json!({"exit_code": 0, "duration_secs": 12, "peak_rss_mib": 900}); + let other = h + .http + .post(&url) + .bearer_auth(&t2) + .json(&body) + .send() + .await + .unwrap(); + assert_eq!( + other.status(), + 404, + "another drone cannot report this lease" + ); + let mine = h + .http + .post(&url) + .bearer_auth(&t1) + .json(&body) + .send() + .await + .unwrap(); + assert_eq!(mine.status(), 204); + let again = h + .http + .post(&url) + .bearer_auth(&t1) + .json(&body) + .send() + .await + .unwrap(); + assert_eq!(again.status(), 204, "idempotent"); + assert!(h.queen.store().active_leases().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn busy_runner_marks_lease_running_and_a_vanished_runner_finishes_it() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-medium", 1); + h.tick().await; + + h.set_runner("borg-d1-", "online", true); + h.state().jobs.clear(); + h.tick().await; + assert_eq!(h.status().await["classes"]["medium"]["leases_running"], 1); + + h.state().runners.clear(); // ephemeral runner deregistered after its job + h.tick().await; + let status = h.status().await; + assert_eq!(status["classes"]["medium"]["leases_running"], 0); + assert_eq!(status["classes"]["medium"]["leases_pending"], 0); +} + +#[tokio::test] +async fn lease_whose_runner_never_came_online_expires_and_is_deregistered() { + let h = Harness::start(json!({"lease_ttl_secs": 300})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-light", 1); + h.tick().await; + let runner_id = h.state().runners[0]["id"].as_u64().unwrap(); + h.state().jobs.clear(); + + let later = OffsetDateTime::now_utc() + time::Duration::seconds(301); + h.queen.tick(later).await.unwrap(); + assert!(h.state().deleted.contains(&runner_id)); + assert!(h.queen.store().active_leases().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn stale_borg_registrations_are_deleted_but_other_runners_are_not() { + let h = Harness::start(json!({})).await; + h.state().runners = vec![ + json!({"id": 900, "name": "borg-x-deadbeef", "status": "offline", "busy": false}), + json!({"id": 901, "name": "borg-y-busy", "status": "online", "busy": true}), + json!({"id": 902, "name": "static-runner", "status": "offline", "busy": false}), + ]; + h.tick().await; + assert_eq!(h.state().deleted, [900]); +} + +#[tokio::test] +async fn app_jwt_is_valid_and_the_installation_token_is_cached() { + let h = Harness::start(json!({})).await; + h.tick().await; + h.tick().await; + h.tick().await; + let st = h.state(); + assert_eq!(st.jwt_rejections, 0); + assert_eq!(st.token_requests, 1); +} + +#[tokio::test] +async fn starving_heavy_job_reserves_a_drone_and_shows_in_status() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + + h.queue_job(1, "borg-ci-medium", 1); + h.queue_job(2, "borg-ci-medium", 1); + h.tick().await; + assert_eq!(h.state().mints.len(), 2); + h.set_runner("borg-d1-", "online", true); + h.state().jobs.clear(); + h.tick().await; // both running; 4 GiB of the 28 GiB left + + h.queue_job(3, "borg-ci-heavy-pr", 300); + h.queue_job(4, "borg-ci-light-pr", 300); + h.tick().await; + + assert_eq!( + h.state().mints.len(), + 2, + "nothing new lands on the reserved drone" + ); + let status = h.status().await; + assert_eq!( + status["reserved_for"], + json!([{"drone_id": "d1", "class": "heavy", "scope": "pr"}]) + ); + assert_eq!( + status["classes"]["heavy"]["oldest_age_secs_pr"] + .as_u64() + .map(|a| a >= 300), + Some(true) + ); +} + +#[tokio::test] +async fn metrics_and_status_report_the_queue() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-medium-pr", 30); + h.tick().await; + + let metrics = h + .http + .get(format!("{}/metrics", h.url)) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + for line in [ + "borg_queued_jobs{class=\"medium\",scope=\"pr\"} 1", + "borg_active_leases{class=\"medium\",state=\"pending\"} 1", + "borg_drones{state=\"fresh\"} 1", + "borg_last_tick_ok 1", + ] { + assert!(metrics.contains(line), "missing {line}\n{metrics}"); + } + let status = h.status().await; + assert_eq!(status["classes"]["medium"]["queued_pr"], 1); + assert_eq!(status["drones"][0]["name"], "d1"); + assert_eq!(status["drones"][0]["fresh"], true); + assert!(status["classes"]["heavy"]["slots"].as_u64() == Some(1)); +} + +#[tokio::test] +async fn mint_failure_is_reported_and_leaves_no_lease() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-light", 1); + h.state().fail_mint = true; + h.tick().await; + assert!(h.queen.store().active_leases().await.unwrap().is_empty()); + let err = h.status().await["last_error"] + .as_str() + .unwrap_or_default() + .to_owned(); + assert!(err.contains("mint failed"), "{err}"); +} + +#[tokio::test] +async fn github_failure_before_minting_aborts_the_tick() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-light", 1); + h.state().fail_runner_list = true; + assert!(h.queen.tick(OffsetDateTime::now_utc()).await.is_err()); + assert!(h.state().mints.is_empty()); +} diff --git a/tests/support/mod.rs b/tests/support/mod.rs new file mode 100644 index 0000000..f67cb6e --- /dev/null +++ b/tests/support/mod.rs @@ -0,0 +1,345 @@ +//! A fake GitHub REST API and a running queen, for integration tests. +//! Shared by several test binaries, each of which uses only part of it. +#![allow(dead_code)] + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, MutexGuard}; + +use aws_lc_rs::encoding::AsDer; +use aws_lc_rs::rsa::{KeyPair, KeySize}; +use aws_lc_rs::signature::KeyPair as _; +use axum::extract::{Path, Query, State}; +use axum::http::{HeaderMap, StatusCode}; +use axum::routing::{delete, get, post}; +use axum::{Json, Router}; +use base64::Engine as _; +use borg::api; +use borg::config::QueenConfig; +use borg::enroll; +use borg::github::GithubClient; +use borg::probe::{Capacity, HostKind}; +use borg::queen::Queen; +use borg::store::Store; +use jsonwebtoken::{Algorithm, DecodingKey, Validation}; +use serde_json::{Value, json}; +use time::OffsetDateTime; +use time::format_description::well_known::Rfc3339; + +pub const INSTALLATION_TOKEN: &str = "ghs_installation_for_tests"; +pub const ORG: &str = "example-org"; + +#[derive(Default)] +pub struct MockState { + public_pem: String, + pub queued_runs: Vec, + pub jobs: Vec, + pub runners: Vec, + pub mints: Vec, + pub deleted: Vec, + pub token_requests: u32, + pub jwt_rejections: u32, + pub fail_mint: bool, + pub fail_runner_list: bool, + next_runner_id: u64, +} + +pub type Shared = Arc>; + +fn pem(label: &str, der: &[u8]) -> String { + let b64 = base64::engine::general_purpose::STANDARD.encode(der); + let body: Vec<&str> = b64 + .as_bytes() + .chunks(64) + .map(|c| std::str::from_utf8(c).unwrap_or_default()) + .collect(); + format!( + "-----BEGIN {label}-----\n{}\n-----END {label}-----\n", + body.join("\n") + ) +} + +/// A fresh RSA key pair as (private PEM, public PEM). +pub fn generate_key() -> (String, String) { + let kp = KeyPair::generate(KeySize::Rsa2048).unwrap(); + let private = AsDer::::as_der(&kp).unwrap(); + let public = AsDer::::as_der(kp.public_key()).unwrap(); + ( + pem("PRIVATE KEY", private.as_ref()), + pem("PUBLIC KEY", public.as_ref()), + ) +} + +fn lock(s: &Shared) -> MutexGuard<'_, MockState> { + s.lock().unwrap() +} + +#[derive(serde::Deserialize)] +struct Claims { + iss: String, + iat: i64, + exp: i64, +} + +fn check_jwt(state: &Shared, headers: &HeaderMap) -> Result<(), StatusCode> { + let presented = bearer(headers).ok_or(StatusCode::UNAUTHORIZED)?; + let pem = lock(state).public_pem.clone(); + let mut v = Validation::new(Algorithm::RS256); + v.set_required_spec_claims(&["exp", "iss"]); + let ok = DecodingKey::from_rsa_pem(pem.as_bytes()) + .ok() + .and_then(|k| jsonwebtoken::decode::(presented, &k, &v).ok()) + .is_some_and(|d| d.claims.iss == "12345" && d.claims.exp - d.claims.iat <= 600); + if ok { + Ok(()) + } else { + lock(state).jwt_rejections += 1; + Err(StatusCode::UNAUTHORIZED) + } +} + +fn check_token(headers: &HeaderMap) -> Result<(), StatusCode> { + (bearer(headers) == Some(INSTALLATION_TOKEN)) + .then_some(()) + .ok_or(StatusCode::UNAUTHORIZED) +} + +fn bearer(headers: &HeaderMap) -> Option<&str> { + headers + .get("authorization")? + .to_str() + .ok()? + .strip_prefix("Bearer ") +} + +async fn installation(State(s): State, h: HeaderMap) -> Result, StatusCode> { + check_jwt(&s, &h)?; + Ok(Json(json!({"id": 77}))) +} + +async fn access_token(State(s): State, h: HeaderMap) -> Result, StatusCode> { + check_jwt(&s, &h)?; + lock(&s).token_requests += 1; + let expires = (OffsetDateTime::now_utc() + time::Duration::hours(1)) + .format(&Rfc3339) + .unwrap(); + Ok(Json( + json!({"token": INSTALLATION_TOKEN, "expires_at": expires}), + )) +} + +async fn runs( + State(s): State, + h: HeaderMap, + Query(q): Query>, +) -> Result, StatusCode> { + check_token(&h)?; + let ids: Vec = if q.get("status").map(String::as_str) == Some("queued") { + lock(&s) + .queued_runs + .iter() + .map(|id| json!({"id": id})) + .collect() + } else { + vec![] + }; + Ok(Json(json!({"workflow_runs": ids}))) +} + +async fn jobs( + State(s): State, + h: HeaderMap, + Path(_p): Path<(String, String, u64)>, +) -> Result, StatusCode> { + check_token(&h)?; + Ok(Json(json!({"jobs": lock(&s).jobs.clone()}))) +} + +async fn list_runners(State(s): State, h: HeaderMap) -> Result, StatusCode> { + check_token(&h)?; + let st = lock(&s); + if st.fail_runner_list { + return Err(StatusCode::INTERNAL_SERVER_ERROR); + } + Ok(Json( + json!({"total_count": st.runners.len(), "runners": st.runners.clone()}), + )) +} + +async fn jitconfig( + State(s): State, + h: HeaderMap, + Json(body): Json, +) -> Result<(StatusCode, Json), StatusCode> { + check_token(&h)?; + let mut st = lock(&s); + if st.fail_mint { + return Err(StatusCode::FORBIDDEN); + } + st.next_runner_id += 1; + let id = st.next_runner_id; + let name = body["name"].as_str().unwrap_or_default().to_owned(); + st.runners + .push(json!({"id": id, "name": name, "status": "offline", "busy": false})); + st.mints.push(body); + Ok(( + StatusCode::CREATED, + Json( + json!({"runner": {"id": id, "name": name}, "encoded_jit_config": format!("jit-{name}")}), + ), + )) +} + +async fn delete_runner( + State(s): State, + h: HeaderMap, + Path((_org, id)): Path<(String, u64)>, +) -> Result { + check_token(&h)?; + let mut st = lock(&s); + st.deleted.push(id); + st.runners.retain(|r| r["id"].as_u64() != Some(id)); + Ok(StatusCode::NO_CONTENT) +} + +pub struct Harness { + pub mock: Shared, + pub queen: Arc, + pub url: String, + pub http: reqwest::Client, +} + +pub fn capacity() -> Capacity { + Capacity { + hostname: "host-a".into(), + os: "linux".into(), + arch: "x86_64".into(), + kind: HostKind::Vm, + vcpu: 16, + // 28 GiB budget after the default 2 GiB reserve. + mem_mib: 30 * 1024, + free_disk_mib: Some(100_000), + docker_usable: true, + mise_version: None, + git_version: None, + } +} + +impl Harness { + /// `extra` is merged into the queen config; warm floors default to none. + pub async fn start(extra: Value) -> Self { + let mock: Shared = Arc::default(); + let (private_pem, public_pem) = generate_key(); + lock(&mock).public_pem = public_pem; + + let app = Router::new() + .route("/orgs/{org}/installation", get(installation)) + .route("/app/installations/{id}/access_tokens", post(access_token)) + .route("/repos/{owner}/{repo}/actions/runs", get(runs)) + .route("/repos/{owner}/{repo}/actions/runs/{run}/jobs", get(jobs)) + .route("/orgs/{org}/actions/runners", get(list_runners)) + .route( + "/orgs/{org}/actions/runners/generate-jitconfig", + post(jitconfig), + ) + .route("/orgs/{org}/actions/runners/{id}", delete(delete_runner)) + .with_state(Arc::clone(&mock)); + let gh_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let gh_url = format!("http://{}", gh_listener.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(gh_listener, app).await }); + + let mut cfg = json!({ + "github": {"org": ORG, "app_id": "12345", "api_base": gh_url}, + "repos": [format!("{ORG}/app")], + "db_path": "unused.db", + "warm_floors": [] + }); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + cfg[k] = v; + } + let cfg = QueenConfig::from_json(&cfg.to_string()).unwrap(); + let github = GithubClient::new(&gh_url, ORG, "12345", private_pem.as_bytes()).unwrap(); + let store = Store::open_in_memory().await.unwrap(); + let queen = Arc::new(Queen::new(cfg, store, github)); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let router = api::router(Arc::clone(&queen)); + tokio::spawn(async move { axum::serve(listener, router).await }); + Self { + mock, + queen, + url, + http: reqwest::Client::new(), + } + } + + pub fn state(&self) -> MutexGuard<'_, MockState> { + lock(&self.mock) + } + + /// Enroll a drone directly in the store and return its token. + pub async fn enroll(&self, name: &str) -> String { + let token = enroll::new_token(); + self.queen + .store() + .enroll(name, &enroll::hash_token(&token), 0) + .await + .unwrap(); + token + } + + pub async fn heartbeat(&self, token: &str, active: &[&str]) -> reqwest::Response { + self.http + .post(format!("{}/v1/heartbeat", self.url)) + .bearer_auth(token) + .json(&json!({"capacity": capacity(), "active_leases": active})) + .send() + .await + .unwrap() + } + + pub async fn leases(&self, token: &str, active: &[&str]) -> Vec { + let body: Value = self.heartbeat(token, active).await.json().await.unwrap(); + body["leases"].as_array().cloned().unwrap_or_default() + } + + pub async fn status(&self) -> Value { + self.http + .get(format!("{}/status", self.url)) + .send() + .await + .unwrap() + .json() + .await + .unwrap() + } + + pub async fn tick(&self) { + self.queen.tick(OffsetDateTime::now_utc()).await.unwrap(); + } + + /// Queue one job of `label`, `age_secs` old. + pub fn queue_job(&self, id: u64, label: &str, age_secs: i64) { + let created = (OffsetDateTime::now_utc() - time::Duration::seconds(age_secs)) + .format(&Rfc3339) + .unwrap(); + let mut st = self.state(); + if st.queued_runs.is_empty() { + st.queued_runs.push(1); + } + st.jobs + .push(json!({"id": id, "status": "queued", "labels": [label], "created_at": created})); + } + + pub fn set_runner(&self, name_prefix: &str, status: &str, busy: bool) { + let mut st = self.state(); + for r in st.runners.iter_mut().filter(|r| { + r["name"] + .as_str() + .is_some_and(|n| n.starts_with(name_prefix)) + }) { + r["status"] = json!(status); + r["busy"] = json!(busy); + } + } +} From 96cf58f3b4b6bde3443801a24f44b72fb4408781 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 10 Oct 2026 19:10:21 +0000 Subject: [PATCH 02/27] chore: project borg from Mono 3bba595c27f8 --- .repository-projection.json | 6 +- README.md | 86 ++++++ src/api.rs | 67 ++++- src/drone/caches/bazel.rs | 103 +++++++ src/drone/caches/gc.rs | 281 +++++++++++++++++++ src/drone/caches/git.rs | 326 ++++++++++++++++++++++ src/drone/caches/kinds.rs | 241 ++++++++++++++++ src/drone/caches/mod.rs | 541 ++++++++++++++++++++++++++++++++++++ src/drone/config.rs | 205 ++++++++++++++ src/drone/install.rs | 18 ++ src/drone/mod.rs | 31 ++- src/drone/runner.rs | 6 + src/main.rs | 10 + src/probe/mod.rs | 7 +- src/store.rs | 49 +++- tests/drone.rs | 230 ++++++++++++++- 16 files changed, 2188 insertions(+), 19 deletions(-) create mode 100644 src/drone/caches/bazel.rs create mode 100644 src/drone/caches/gc.rs create mode 100644 src/drone/caches/git.rs create mode 100644 src/drone/caches/kinds.rs create mode 100644 src/drone/caches/mod.rs diff --git a/.repository-projection.json b/.repository-projection.json index fb2eacb..96b68f5 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "a5293e119fc29a43fef8e2f3fb3e06f8db5a931b", + "sourceSha": "3bba595c27f8777722f938346c87d96c2edfc42a", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "bb547084fc1e8c6827704c24b9165fa148c84945", + "priorProjectedBase": "1f067ea6a8594a1ea831b7efe50d5c857138ec6b", "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "38a3bf905bd076b969e0df500b40705791f6e6b163d9fc992462efaed7943c75", + "contentDigest": "02f5f780bf1b6ce76b937b699070ba605cf9f8854d1466938481f5327c45f1cf", "publicationEligible": true } diff --git a/README.md b/README.md index 2a558fb..35fc48a 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,92 @@ budget, so an owner can say "this PC gives at most 8 cores and 24 GiB". 5. The exit code, duration and peak RSS of the runner's process group (Linux) are reported to the queen. +### Host caches + +The drone owns one stable directory per cache kind under `cache_root`. The +path is identical for every lease on the host, is exported to every job, is +never touched by per-job cleanup, and is size-capped by least-recently-used +garbage collection that runs only while no lease is running. `cache_root` and +`work_root` must not contain one another. Everything is configured in the +`caches` section of the drone config; every kind except `git` is on by default. + +| Kind | Environment for jobs | Default cap | Collected | +| --- | --- | --- | --- | +| `sccache` | `SCCACHE_DIR`, `SCCACHE_CACHE_SIZE` | 25% of free disk at start, at most 100G | never by the drone; sccache enforces its own size | +| `cargo` | `CARGO_HOME` (registry and git), `CARGO_INSTALL_ROOT` per lease | 20G | crates, unpacked sources, git db and checkouts | +| `bazel` | `BAZELISK_HOME`; `--disk_cache` and `--repository_cache` via a managed block in `~/.bazelrc` | 50G | `disk/cas`, `disk/ac`, repository cache entries | +| `pnpm` | `PNPM_STORE_DIR`, `npm_config_store_dir` | 20G | store files (the index is kept) | +| `go` | `GOMODCACHE`, `GOCACHE` | 20G | build cache only; the module cache is read-only trees and is left alone | +| `uv` | `UV_CACHE_DIR` | 10G | entries of each versioned bucket | +| `mise` | `MISE_DATA_DIR` | 20G | installed tool versions and downloads; shims are kept | +| `git` | `BORG_GIT_MIRROR__` per mirror | none | never (objects only accumulate) | + +```json +"caches": { + "git": { "mirrors": ["example-org/app"], "remote_base": "https://github.com" }, + "sccache": { "max_size": "50G" }, + "bazel": { "max_size": "40G", "bazelrc": "/var/lib/borg/.bazelrc" }, + "uv": { "enabled": false } +} +``` + +Notes on each kind: + +* **Garbage collection** removes whole entries (a crate, a content-addressed + blob, an installed tool version), oldest use first, until the collectable data + is at most 90% of the cap. It takes a lock that every lease holds for its whole + life, so it never overlaps a job; under constant load it simply waits for an + idle moment, and caps are soft in the meantime. +* **cargo** shares only the registry and git data. `cargo install` is pointed at + a per-lease root so no installed binary ever leaks between leases. +* **bazel**: the drone writes a marked block into the service user's `~/.bazelrc` + (or `caches.bazel.bazelrc`) and leaves the rest of the file alone. The user rc + is read after the workspace rc, so these two flags win. Bazel's disk cache is + safe for concurrent builds: entries are written to a temporary file and renamed + into place, and a collision between processes is harmless. A collected entry is + a cache miss. +* **go**: the module cache is never collected; the build cache is. +* **uv** and **mise** collection is coarse: it removes entries those tools + re-create on demand. mise serializes installs of one tool version with a lock + file. +* **git mirrors**: for each `caches.git.mirrors` entry (`owner/name`) the drone + keeps a bare mirror of the branches and tags at + `/git//.git`, refreshed every + `caches.git.interval_secs` (default 600) and before a lease starts (at most + once per 15 seconds). The variable name is upper-cased with every other + character replaced by `_`, and it is set only while the mirror exists, so a job + can test it and fall back to a normal clone: + + ``` + if [ -n "${BORG_GIT_MIRROR_EXAMPLE_ORG_APP:-}" ]; then + git clone --reference "$BORG_GIT_MIRROR_EXAMPLE_ORG_APP" src + else + git clone src + fi + ``` + + Mirrors fetch from `remote_base` (default `https://github.com`) plus + `/.git`. The drone holds no GitHub credentials: for a private + repository give the service user a credential helper, or set `remote_base` to + an ssh base such as `git@host:` with a read-only deploy key. A failed refresh + keeps the previous mirror and never blocks a job. + +What the drone does not do: it does not inject `RUSTFLAGS`. A +`--remap-path-prefix` can only make paths stable if the path it names is stable, +and a lease's work directory contains the lease id; workspace crates therefore +do not share compile-cache entries across leases, while dependencies built from +the shared registry do. It does not collect sccache hit rates, because sccache +servers are job-scoped. + +Every heartbeat carries each cache's size, cap, evicted bytes and last +collection time (and, for git, how many leases were given a mirror and how many +ran without one). The queen shows them under each drone in `/status` and as +`borg_cache_*` metrics. Sizes are re-measured every +`caches.maintenance_interval_secs` (default 600) and after leases. + +`borg drone install` takes `--git-mirror owner/name` (repeatable) and +`--git-remote-base`; edit `/etc/borg/drone.json` for the rest. + On start the drone deletes leftover lease directories and borg-labelled docker resources from a previous run. On `SIGTERM` it stops accepting leases and waits up to `drain_timeout_secs` for running jobs before stopping them. diff --git a/src/api.rs b/src/api.rs index 1c831fd..b7b6ea7 100644 --- a/src/api.rs +++ b/src/api.rs @@ -38,6 +38,23 @@ pub struct HeartbeatRequest { /// Lease ids the drone currently holds; acknowledges their JIT configs. #[serde(default)] pub active_leases: Vec, + /// Size and garbage-collection state of the drone's host caches. + #[serde(default)] + pub cache_stats: Vec, +} + +/// One host cache as measured by its drone. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CacheStat { + pub name: String, + pub size_bytes: u64, + /// The cap garbage collection enforces (sccache: its own cache size). + pub max_bytes: Option, + pub evicted_bytes: u64, + pub last_gc_at: Option, + /// Git mirrors only: leases that were handed a mirror / ran without one. + pub hits: Option, + pub misses: Option, } fn yes() -> bool { @@ -132,7 +149,14 @@ async fn heartbeat( let now = OffsetDateTime::now_utc().unix_timestamp(); let store = queen.store(); store - .record_heartbeat(&drone, &req.capacity, &req.caps, req.accepting, now) + .record_heartbeat( + &drone, + &req.capacity, + &req.caps, + req.accepting, + &req.cache_stats, + now, + ) .await?; store.acknowledge(&drone, &req.active_leases).await?; let leases = store @@ -201,6 +225,7 @@ pub struct DroneStatus { pub vcpu: Option, pub mem_mib: Option, pub active_leases: u32, + pub caches: Vec, } #[derive(Debug, Serialize)] @@ -243,6 +268,7 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result mem_mib: d.capacity.as_ref().map(|c| c.mem_mib), active_leases: u32::try_from(active.iter().filter(|l| l.drone == d.name).count()) .unwrap_or(u32::MAX), + caches: d.cache_stats.clone(), }); } @@ -408,6 +434,45 @@ pub fn render_metrics(status: &Status) -> String { "Unix time of the last reconciliation pass.", vec![(String::new(), status.last_tick_at.unwrap_or(0) as f64)], ); + let cache_rows = |f: &dyn Fn(&CacheStat) -> Option| -> Vec<(String, f64)> { + status + .drones + .iter() + .flat_map(|d| { + d.caches.iter().filter_map(move |c| { + f(c).map(|v| (format!("drone=\"{}\",cache=\"{}\"", d.name, c.name), v)) + }) + }) + .collect() + }; + #[allow(clippy::cast_precision_loss, reason = "byte counts far below 2^52")] + { + gauge( + "borg_cache_size_bytes", + "Size of a drone host cache.", + cache_rows(&|c| Some(c.size_bytes as f64)), + ); + gauge( + "borg_cache_max_bytes", + "Size cap of a drone host cache.", + cache_rows(&|c| c.max_bytes.map(|v| v as f64)), + ); + gauge( + "borg_cache_evicted_bytes", + "Bytes evicted from a drone host cache by garbage collection.", + cache_rows(&|c| Some(c.evicted_bytes as f64)), + ); + gauge( + "borg_cache_hits", + "Leases that were handed a git mirror.", + cache_rows(&|c| c.hits.map(|v| v as f64)), + ); + gauge( + "borg_cache_misses", + "Leases that ran without a git mirror.", + cache_rows(&|c| c.misses.map(|v| v as f64)), + ); + } gauge( "borg_last_tick_ok", "1 if the last reconciliation pass had no error.", diff --git a/src/drone/caches/bazel.rs b/src/drone/caches/bazel.rs new file mode 100644 index 0000000..61aa582 --- /dev/null +++ b/src/drone/caches/bazel.rs @@ -0,0 +1,103 @@ +//! The drone's managed block in a Bazel user rc file. +//! +//! Bazel has no environment variable for the disk or repository cache, so the +//! drone points them at host-level directories through the service user's +//! `~/.bazelrc`. The user rc is read after the workspace rc, so it wins for +//! these two flags. Only the block between the markers is ever rewritten. + +use std::fs; +use std::io; +use std::path::Path; + +const BEGIN: &str = "# >>> borg managed cache flags (do not edit) >>>"; +const END: &str = "# <<< borg managed cache flags <<<"; + +/// The managed block for the given cache directories. +pub fn managed_block(disk_cache: &Path, repository_cache: &Path) -> String { + format!( + "{BEGIN}\ncommon --disk_cache={}\ncommon --repository_cache={}\n{END}\n", + disk_cache.display(), + repository_cache.display() + ) +} + +/// Replace (or append) the managed block in `existing`, leaving all other +/// content untouched. +pub fn with_block(existing: &str, block: &str) -> String { + match (existing.find(BEGIN), existing.find(END)) { + (Some(start), Some(end)) if end > start => { + let end = end + END.len(); + let tail = existing[end..] + .strip_prefix('\n') + .unwrap_or(&existing[end..]); + format!("{}{block}{tail}", &existing[..start]) + } + _ if existing.is_empty() => block.to_owned(), + _ if existing.ends_with('\n') => format!("{existing}\n{block}"), + _ => format!("{existing}\n\n{block}"), + } +} + +/// Write the block into `rcfile`. Returns whether the file changed. +pub fn write_rc(rcfile: &Path, disk_cache: &Path, repository_cache: &Path) -> io::Result { + let existing = match fs::read_to_string(rcfile) { + Ok(text) => text, + Err(e) if e.kind() == io::ErrorKind::NotFound => String::new(), + Err(e) => return Err(e), + }; + let updated = with_block(&existing, &managed_block(disk_cache, repository_cache)); + if updated == existing { + return Ok(false); + } + if let Some(parent) = rcfile.parent() { + fs::create_dir_all(parent)?; + } + fs::write(rcfile, updated)?; + Ok(true) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn block() -> String { + managed_block(Path::new("/c/bazel/disk"), Path::new("/c/bazel/repo")) + } + + #[test] + fn block_lists_both_caches() { + let b = block(); + assert!(b.contains("common --disk_cache=/c/bazel/disk\n")); + assert!(b.contains("common --repository_cache=/c/bazel/repo\n")); + } + + #[test] + fn appends_to_existing_content_and_is_idempotent() { + let once = with_block("build --jobs=4\n", &block()); + assert!(once.starts_with("build --jobs=4\n\n# >>> borg")); + assert_eq!(with_block(&once, &block()), once); + } + + #[test] + fn replaces_only_the_managed_block() { + let old = format!( + "before\n\n{}after\n", + managed_block(Path::new("/old/d"), Path::new("/old/r")) + ); + let updated = with_block(&old, &block()); + assert!(updated.starts_with("before\n\n")); + assert!(updated.ends_with("after\n")); + assert!(updated.contains("/c/bazel/disk") && !updated.contains("/old/")); + assert_eq!(updated.matches(BEGIN).count(), 1); + } + + #[test] + fn writes_the_file_and_reports_changes() { + let dir = tempfile::tempdir().unwrap(); + let rc = dir.path().join("home/.bazelrc"); + assert!(write_rc(&rc, Path::new("/c/bazel/disk"), Path::new("/c/bazel/repo")).unwrap()); + assert!(!write_rc(&rc, Path::new("/c/bazel/disk"), Path::new("/c/bazel/repo")).unwrap()); + assert!(write_rc(&rc, Path::new("/c/other/disk"), Path::new("/c/bazel/repo")).unwrap()); + assert!(fs::read_to_string(&rc).unwrap().contains("/c/other/disk")); + } +} diff --git a/src/drone/caches/gc.rs b/src/drone/caches/gc.rs new file mode 100644 index 0000000..3e5ca01 --- /dev/null +++ b/src/drone/caches/gc.rs @@ -0,0 +1,281 @@ +//! Size accounting and least-recently-used eviction for cache directories. +//! +//! A cache is garbage collected in whole *units*: the entries at a fixed depth +//! below a bucket directory (a downloaded crate, a content-addressed blob, an +//! installed tool version). A unit is as recent as the newest file in it, and +//! the oldest units go first until the cache is back under 90% of its cap, so +//! it is not collected again after the next job. + +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; +use std::time::SystemTime; + +/// Collection stops once the cache is at or below this share of its cap. +const TARGET_PERCENT: u64 = 90; + +/// Parse a size such as `512M` or `50G` (suffix K, M, G or T; powers of 1024). +pub fn parse_size(text: &str) -> Option { + let unit = text.chars().last()?; + let shift = match unit { + 'K' => 10, + 'M' => 20, + 'G' => 30, + 'T' => 40, + _ => return None, + }; + let digits = &text[..text.len() - 1]; + if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + digits.parse::().ok()?.checked_mul(1u64 << shift) +} + +/// Total size and the newest file modification/access time under `path`. +/// Symlinks are counted as themselves, never followed. +pub fn measure(path: &Path) -> (u64, SystemTime) { + let mut size = 0; + let mut newest = SystemTime::UNIX_EPOCH; + let mut stack = vec![path.to_owned()]; + while let Some(dir) = stack.pop() { + let Ok(meta) = fs::symlink_metadata(&dir) else { + continue; + }; + size += meta.len(); + if meta.is_dir() { + // Directory times only record when entries were added or listed; + // so recency comes from the files. + if let Ok(entries) = fs::read_dir(&dir) { + stack.extend(entries.flatten().map(|e| e.path())); + } + } else { + for t in [meta.modified(), meta.accessed()].into_iter().flatten() { + newest = newest.max(t); + } + } + } + (size, newest) +} + +/// One evictable entry. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Unit { + pub path: PathBuf, + pub size: u64, + pub recency: SystemTime, +} + +/// The entries exactly `depth` levels below each existing bucket (depth 1 is +/// the bucket's children). +pub fn collect_units(buckets: &[(PathBuf, usize)]) -> Vec { + let mut units = Vec::new(); + for (bucket, depth) in buckets { + let mut level = vec![bucket.clone()]; + for _ in 0..*depth { + level = level + .iter() + .filter(|p| fs::symlink_metadata(p).is_ok_and(|m| m.is_dir())) + .filter_map(|p| fs::read_dir(p).ok()) + .flat_map(|rd| rd.flatten().map(|e| e.path())) + .collect(); + } + for path in level { + let (size, recency) = measure(&path); + units.push(Unit { + path, + size, + recency, + }); + } + } + units +} + +/// Directory contents may be read-only (module caches); make them removable. +fn make_writable(path: &Path) { + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mut stack = vec![path.to_owned()]; + while let Some(p) = stack.pop() { + let Ok(meta) = fs::symlink_metadata(&p) else { + continue; + }; + if meta.is_dir() { + let _ = fs::set_permissions(&p, fs::Permissions::from_mode(0o755)); + if let Ok(rd) = fs::read_dir(&p) { + stack.extend(rd.flatten().map(|e| e.path())); + } + } + } + } +} + +fn remove_unit(path: &Path) -> io::Result<()> { + let is_dir = fs::symlink_metadata(path)?.is_dir(); + let attempt = |p: &Path| { + if is_dir { + fs::remove_dir_all(p) + } else { + fs::remove_file(p) + } + }; + match attempt(path) { + Err(e) if e.kind() == io::ErrorKind::PermissionDenied && is_dir => { + make_writable(path); + attempt(path) + } + other => other, + } +} + +/// Result of one collection. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct Collected { + pub removed_units: u64, + pub freed_bytes: u64, +} + +/// Evict least-recently-used `units` until `total_bytes` (the combined size of +/// all `units`; data that is never evicted does not count) is at most 90% of +/// `cap_bytes`. Does nothing while the units are within the cap. +pub fn evict_lru(mut units: Vec, total_bytes: u64, cap_bytes: u64) -> Collected { + let mut out = Collected::default(); + if total_bytes <= cap_bytes { + return out; + } + let target = cap_bytes / 100 * TARGET_PERCENT; + let mut remaining = total_bytes; + units.sort_by(|a, b| a.recency.cmp(&b.recency).then_with(|| a.path.cmp(&b.path))); + for unit in units { + if remaining <= target { + break; + } + match remove_unit(&unit.path) { + Ok(()) => { + remaining = remaining.saturating_sub(unit.size); + out.removed_units += 1; + out.freed_bytes += unit.size; + } + Err(e) => tracing::warn!("cache gc: remove {}: {e}", unit.path.display()), + } + } + out +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::*; + + fn file(path: &Path, bytes: usize, age_secs: u64) { + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(path, vec![0u8; bytes]).unwrap(); + let when = SystemTime::now() - Duration::from_secs(age_secs); + let times = std::fs::FileTimes::new() + .set_modified(when) + .set_accessed(when); + std::fs::File::options() + .write(true) + .open(path) + .unwrap() + .set_times(times) + .unwrap(); + } + + #[test] + fn sizes_parse_with_binary_suffixes() { + assert_eq!(parse_size("512M"), Some(512 << 20)); + assert_eq!(parse_size("50G"), Some(50 << 30)); + assert_eq!(parse_size("2T"), Some(2 << 40)); + for bad in ["", "G", "50", "5g", "5.5G", "-1G", "99999999999999999999G"] { + assert_eq!(parse_size(bad), None, "{bad}"); + } + } + + #[test] + fn measure_sums_a_tree_and_finds_the_newest_file() { + let dir = tempfile::tempdir().unwrap(); + file(&dir.path().join("a/old"), 100, 1000); + file(&dir.path().join("a/b/new"), 50, 10); + let (size, newest) = measure(dir.path()); + assert!(size >= 150); + assert!( + newest.elapsed().unwrap() < Duration::from_secs(500), + "newest file counts" + ); + } + + #[test] + fn units_are_the_entries_at_the_requested_depth() { + let dir = tempfile::tempdir().unwrap(); + file(&dir.path().join("cas/ab/h1"), 10, 5); + file(&dir.path().join("cas/cd/h2"), 10, 5); + file(&dir.path().join("other/x"), 10, 5); + let mut names: Vec = + collect_units(&[(dir.path().join("cas"), 2), (dir.path().join("missing"), 1)]) + .iter() + .map(|u| u.path.file_name().unwrap().to_string_lossy().into_owned()) + .collect(); + names.sort(); + assert_eq!(names, ["h1", "h2"]); + let shallow = collect_units(&[(dir.path().join("cas"), 1)]); + assert_eq!( + shallow.len(), + 2, + "depth 1 is the bucket directories themselves" + ); + } + + #[test] + fn evicts_oldest_first_down_to_ninety_percent_of_the_cap() { + let dir = tempfile::tempdir().unwrap(); + for (name, age) in [("oldest", 400), ("old", 300), ("mid", 200), ("new", 100)] { + file(&dir.path().join(format!("b/{name}/data")), 1000, age); + } + let units = collect_units(&[(dir.path().join("b"), 1)]); + let total: u64 = units.iter().map(|u| u.size).sum(); + // A cap of half the total must drop the two oldest units. + let got = evict_lru(units, total, total / 2); + assert!(got.removed_units >= 2, "{got:?}"); + assert!(!dir.path().join("b/oldest").exists() && !dir.path().join("b/old").exists()); + assert!( + dir.path().join("b/new").exists(), + "the newest unit survives" + ); + } + + #[test] + fn nothing_is_evicted_within_the_cap() { + let dir = tempfile::tempdir().unwrap(); + file(&dir.path().join("b/u/data"), 1000, 100); + let units = collect_units(&[(dir.path().join("b"), 1)]); + let total: u64 = units.iter().map(|u| u.size).sum(); + assert_eq!(evict_lru(units, total, total), Collected::default()); + assert!(dir.path().join("b/u").exists()); + } + + #[cfg(unix)] + #[test] + fn read_only_trees_are_still_removed() { + use std::os::unix::fs::PermissionsExt; + let dir = tempfile::tempdir().unwrap(); + file(&dir.path().join("b/mod@v1/src/lib.go"), 1000, 100); + fs::set_permissions( + dir.path().join("b/mod@v1/src"), + fs::Permissions::from_mode(0o555), + ) + .unwrap(); + fs::set_permissions( + dir.path().join("b/mod@v1"), + fs::Permissions::from_mode(0o555), + ) + .unwrap(); + let units = collect_units(&[(dir.path().join("b"), 1)]); + let total: u64 = units.iter().map(|u| u.size).sum(); + let got = evict_lru(units, total, 1); + assert_eq!(got.removed_units, 1); + assert!(!dir.path().join("b/mod@v1").exists()); + } +} diff --git a/src/drone/caches/git.rs b/src/drone/caches/git.rs new file mode 100644 index 0000000..81cde39 --- /dev/null +++ b/src/drone/caches/git.rs @@ -0,0 +1,326 @@ +//! Bare git mirrors: one per configured repository under +//! `/git//.git`, refreshed on a timer and before a +//! lease starts. Jobs learn the path from `BORG_GIT_MIRROR__`, +//! set only while the mirror exists, so they can fall back to a plain clone. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use tokio::sync::Mutex; + +use crate::drone::config::GitCaches; +use crate::drone::exec::Exec; + +/// A mirror refreshed this recently is not refreshed again for a new lease. +const MIN_REFRESH_GAP: Duration = Duration::from_secs(15); +/// Upper bound for one mirror refresh (the first clone of a large repository included). +const REFRESH_TIMEOUT: Duration = Duration::from_secs(1800); +const PREFIX: &str = "BORG_GIT_MIRROR_"; + +/// Environment variable naming the mirror of `owner/repo`: upper-cased with +/// every non-alphanumeric character replaced by `_` +/// (`example-org/app` becomes `BORG_GIT_MIRROR_EXAMPLE_ORG_APP`). +pub fn mirror_var(repo: &str) -> String { + let body: String = repo + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() { + c.to_ascii_uppercase() + } else { + '_' + } + }) + .collect(); + format!("{PREFIX}{body}") +} + +/// `/git//.git` +pub fn mirror_path(cache_root: &Path, repo: &str) -> PathBuf { + cache_root.join("git").join(format!("{repo}.git")) +} + +/// The URL a mirror fetches from: `base` and `owner/repo.git`, joined by `/` +/// unless `base` already ends in `/` or `:` (scp-style `git@host:`). +pub fn remote_url(base: &str, repo: &str) -> String { + if base.ends_with('/') || base.ends_with(':') { + format!("{base}{repo}.git") + } else { + format!("{base}/{repo}.git") + } +} + +fn run_git(exec: &dyn Exec, git: &str, args: &[&str]) -> Result<(), String> { + match exec.run(git, args) { + Ok(out) if out.success => Ok(()), + Ok(_) => Err(format!("`git {}` failed", args.join(" "))), + Err(e) => Err(format!("{git}: {e}")), + } +} + +/// Create the mirror if needed, point it at `url`, and fetch every branch and +/// tag into it. Idempotent, so an interrupted first clone heals on the next run. +fn refresh_mirror(exec: &dyn Exec, git: &str, path: &Path, url: &str) -> Result<(), String> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(|e| format!("{}: {e}", parent.display()))?; + } + let p = path.to_string_lossy(); + if !path.join("HEAD").exists() { + run_git(exec, git, &["init", "--bare", "--quiet", &p])?; + } + run_git(exec, git, &["-C", &p, "config", "remote.origin.url", url])?; + run_git( + exec, + git, + &[ + "-C", + &p, + "config", + "remote.origin.fetch", + "+refs/heads/*:refs/heads/*", + ], + )?; + // Objects only accumulate: a job borrowing them via alternates must never + // see them pruned underneath it. + run_git(exec, git, &["-C", &p, "config", "gc.auto", "0"])?; + run_git( + exec, + git, + &[ + "-C", + &p, + "-c", + "http.lowSpeedLimit=1000", + "-c", + "http.lowSpeedTime=60", + "fetch", + "--prune", + "--tags", + "--force", + "--quiet", + "origin", + ], + ) +} + +struct Mirror { + repo: String, + path: PathBuf, + url: String, + refreshed: Mutex>, +} + +pub struct GitMirrors { + git: String, + exec: Arc, + mirrors: Vec, + /// Leases that were handed a mirror / ran without one. + served: AtomicU64, + missed: AtomicU64, +} + +impl GitMirrors { + pub fn new(cfg: &GitCaches, cache_root: &Path, exec: Arc) -> Self { + let mirrors = cfg + .mirrors + .iter() + .map(|repo| Mirror { + repo: repo.clone(), + path: mirror_path(cache_root, repo), + url: remote_url(&cfg.remote_base, repo), + refreshed: Mutex::new(None), + }) + .collect(); + Self { + git: cfg.bin.clone(), + exec, + mirrors, + served: AtomicU64::new(0), + missed: AtomicU64::new(0), + } + } + + pub fn is_empty(&self) -> bool { + self.mirrors.is_empty() + } + + async fn refresh(&self, mirror: &Mirror, force: bool) { + let mut last = mirror.refreshed.lock().await; + if !force && last.is_some_and(|t| t.elapsed() < MIN_REFRESH_GAP) { + return; + } + let (exec, git, path, url) = ( + Arc::clone(&self.exec), + self.git.clone(), + mirror.path.clone(), + mirror.url.clone(), + ); + let job = + tokio::task::spawn_blocking(move || refresh_mirror(exec.as_ref(), &git, &path, &url)); + match tokio::time::timeout(REFRESH_TIMEOUT, job).await { + Ok(Ok(Ok(()))) => *last = Some(Instant::now()), + Ok(Ok(Err(e))) => tracing::warn!("git mirror {} refresh failed: {e}", mirror.repo), + Ok(Err(e)) => tracing::warn!("git mirror {} refresh task failed: {e}", mirror.repo), + Err(_) => tracing::warn!("git mirror {} refresh timed out", mirror.repo), + } + } + + pub async fn refresh_stale(&self) { + for m in &self.mirrors { + self.refresh(m, false).await; + } + } + + pub async fn refresh_all(&self) { + for m in &self.mirrors { + self.refresh(m, true).await; + } + } + + /// The `BORG_GIT_MIRROR_*` variables for mirrors that exist, counting the + /// lease as served or not. + pub fn env(&self) -> Vec<(String, String)> { + let mut env = Vec::new(); + for m in &self.mirrors { + if m.path.join("HEAD").exists() { + self.served.fetch_add(1, Ordering::Relaxed); + env.push((mirror_var(&m.repo), m.path.to_string_lossy().into_owned())); + } else { + self.missed.fetch_add(1, Ordering::Relaxed); + } + } + env + } + + pub fn hit_counts(&self) -> (u64, u64) { + ( + self.served.load(Ordering::Relaxed), + self.missed.load(Ordering::Relaxed), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::drone::cleanup::testing::FakeExec; + + fn cfg(mirrors: &[&str]) -> GitCaches { + GitCaches { + mirrors: mirrors.iter().map(|m| (*m).to_owned()).collect(), + remote_base: "https://git.example.test".to_owned(), + ..GitCaches::default() + } + } + + #[test] + fn variable_names_paths_and_urls() { + assert_eq!( + mirror_var("example-org/app"), + "BORG_GIT_MIRROR_EXAMPLE_ORG_APP" + ); + assert_eq!(mirror_var("a.b/c_d"), "BORG_GIT_MIRROR_A_B_C_D"); + assert_eq!( + mirror_path(Path::new("/cache"), "example-org/app"), + Path::new("/cache/git/example-org/app.git") + ); + assert_eq!( + remote_url("https://git.example.test", "o/r"), + "https://git.example.test/o/r.git" + ); + assert_eq!( + remote_url("https://git.example.test/", "o/r"), + "https://git.example.test/o/r.git" + ); + assert_eq!( + remote_url("git@git.example.test:", "o/r"), + "git@git.example.test:o/r.git" + ); + } + + #[tokio::test] + async fn refresh_runs_the_expected_steps_and_is_rate_limited() { + let root = tempfile::tempdir().unwrap(); + let x = Arc::new(FakeExec::default()); + let m = GitMirrors::new(&cfg(&["example-org/app"]), root.path(), x.clone()); + let path = mirror_path(root.path(), "example-org/app") + .display() + .to_string(); + + m.refresh_stale().await; + let calls = x.calls(); + assert_eq!(calls[0], format!("git init --bare --quiet {path}")); + assert_eq!( + calls[1], + format!( + "git -C {path} config remote.origin.url https://git.example.test/example-org/app.git" + ) + ); + assert_eq!( + calls[2], + format!("git -C {path} config remote.origin.fetch +refs/heads/*:refs/heads/*") + ); + assert_eq!(calls[3], format!("git -C {path} config gc.auto 0")); + assert!( + calls[4].contains("fetch --prune --tags --force --quiet origin"), + "{}", + calls[4] + ); + assert_eq!(calls.len(), 5); + + m.refresh_stale().await; + assert_eq!( + x.calls().len(), + 5, + "a lease right after a refresh does not refresh again" + ); + m.refresh_all().await; + assert!( + x.calls().len() > 5, + "the timer refresh ignores the rate limit" + ); + } + + #[tokio::test] + async fn a_failed_refresh_is_retried_by_the_next_lease() { + let root = tempfile::tempdir().unwrap(); + let x = Arc::new(FakeExec::default()); + let path = mirror_path(root.path(), "example-org/app") + .display() + .to_string(); + x.fail(&format!( + "git -C {path} -c http.lowSpeedLimit=1000 -c http.lowSpeedTime=60 fetch --prune --tags --force --quiet origin" + )); + let m = GitMirrors::new(&cfg(&["example-org/app"]), root.path(), x.clone()); + m.refresh_stale().await; + let after_first = x.calls().len(); + m.refresh_stale().await; + assert!( + x.calls().len() > after_first, + "failure does not start the rate-limit window" + ); + } + + #[test] + fn env_lists_only_existing_mirrors_and_counts_hits_and_misses() { + let root = tempfile::tempdir().unwrap(); + let m = GitMirrors::new( + &cfg(&["example-org/app", "example-org/other"]), + root.path(), + Arc::new(FakeExec::default()), + ); + let present = mirror_path(root.path(), "example-org/app"); + std::fs::create_dir_all(&present).unwrap(); + std::fs::write(present.join("HEAD"), "ref: refs/heads/main\n").unwrap(); + let env = m.env(); + assert_eq!( + env, + [( + "BORG_GIT_MIRROR_EXAMPLE_ORG_APP".to_owned(), + present.display().to_string() + )] + ); + assert_eq!(m.hit_counts(), (1, 1)); + } +} diff --git a/src/drone/caches/kinds.rs b/src/drone/caches/kinds.rs new file mode 100644 index 0000000..3013c39 --- /dev/null +++ b/src/drone/caches/kinds.rs @@ -0,0 +1,241 @@ +//! The cache kinds the drone manages: where each lives under `cache_root`, the +//! environment that points jobs at it, and which parts may be garbage +//! collected (and in what units). +//! +//! Concurrency, by kind. Leases on one host run at the same time and share +//! these directories: +//! * sccache, cargo, go, uv, pnpm: documented as safe for concurrent use. +//! * bazel: the disk cache is safe for concurrent Bazel processes; its source +//! (`DiskCacheClient`) writes entries to a temporary file and renames them +//! into place, and states that concurrent processes "can safely retrieve and +//! store entries in a shared disk cache, even when they collide". The +//! repository cache is content addressed and shared across workspaces. +//! * mise: installs take a per-tool-version lock file (`mise` source, +//! `tool_version_lock`); read from the source, not stress tested here. + +use std::fs; +use std::path::{Path, PathBuf}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Kind { + Sccache, + Cargo, + Bazel, + Pnpm, + Go, + Uv, + Mise, +} + +const GIB: u64 = 1 << 30; + +impl Kind { + pub const ALL: [Kind; 7] = [ + Kind::Sccache, + Kind::Cargo, + Kind::Bazel, + Kind::Pnpm, + Kind::Go, + Kind::Uv, + Kind::Mise, + ]; + + /// Config key, stats name and directory name under `cache_root`. + pub fn name(self) -> &'static str { + match self { + Kind::Sccache => "sccache", + Kind::Cargo => "cargo", + Kind::Bazel => "bazel", + Kind::Pnpm => "pnpm", + Kind::Go => "go", + Kind::Uv => "uv", + Kind::Mise => "mise", + } + } + + pub fn dir(self, cache_root: &Path) -> PathBuf { + cache_root.join(self.name()) + } + + /// Default size cap. sccache is absent: it enforces `SCCACHE_CACHE_SIZE` + /// itself, so the drone never deletes from its directory. + pub fn default_max_bytes(self) -> Option { + match self { + Kind::Sccache => None, + Kind::Bazel => Some(50 * GIB), + Kind::Uv => Some(10 * GIB), + Kind::Cargo | Kind::Pnpm | Kind::Go | Kind::Mise => Some(20 * GIB), + } + } + + /// Directories to create up front. + pub fn subdirs(self) -> &'static [&'static str] { + match self { + Kind::Bazel => &["disk", "repo", "bazelisk"], + Kind::Go => &["mod", "build"], + Kind::Cargo => &["registry", "git"], + _ => &[""], + } + } + + /// The environment that points a lease at this cache. `lease_dir` is the + /// lease's private directory, for anything that must not be shared. + pub fn env( + self, + cache_root: &Path, + lease_dir: &Path, + sccache_size: &str, + ) -> Vec<(String, String)> { + let dir = self.dir(cache_root); + let path = |p: PathBuf| p.to_string_lossy().into_owned(); + let pair = |k: &str, v: String| (k.to_owned(), v); + match self { + Kind::Sccache => vec![ + pair("SCCACHE_DIR", path(dir)), + pair("SCCACHE_CACHE_SIZE", sccache_size.to_owned()), + ], + // Registry and git data are shared. `cargo install` writes to + // `$CARGO_HOME/bin` unless told otherwise, so its root is per + // lease and no installed binary leaks between leases. + Kind::Cargo => vec![ + pair("CARGO_HOME", path(dir)), + pair("CARGO_INSTALL_ROOT", path(lease_dir.join("cargo-install"))), + ], + Kind::Bazel => vec![ + pair("BAZELISK_HOME", path(dir.join("bazelisk"))), + pair("BORG_BAZEL_DISK_CACHE", path(dir.join("disk"))), + pair("BORG_BAZEL_REPOSITORY_CACHE", path(dir.join("repo"))), + ], + Kind::Pnpm => vec![ + pair("PNPM_STORE_DIR", path(dir.clone())), + pair("npm_config_store_dir", path(dir)), + ], + Kind::Go => vec![ + pair("GOMODCACHE", path(dir.join("mod"))), + pair("GOCACHE", path(dir.join("build"))), + ], + Kind::Uv => vec![pair("UV_CACHE_DIR", path(dir))], + Kind::Mise => vec![pair("MISE_DATA_DIR", path(dir))], + } + } + + /// Where garbage collection may delete: `(bucket, depth)` pairs whose + /// entries `depth` levels down are the eviction units. Everything not + /// listed (indexes, shims, module caches, sccache) is never deleted. + pub fn gc_buckets(self, cache_root: &Path) -> Vec<(PathBuf, usize)> { + let dir = self.dir(cache_root); + let children = |p: &Path| -> Vec { + fs::read_dir(p) + .map(|rd| rd.flatten().map(|e| e.path()).collect()) + .unwrap_or_default() + }; + match self { + Kind::Sccache => Vec::new(), + // Downloaded crates, unpacked sources, git databases and checkouts. + Kind::Cargo => { + let mut b: Vec<(PathBuf, usize)> = Vec::new(); + for sub in ["registry/cache", "registry/src"] { + b.extend(children(&dir.join(sub)).into_iter().map(|p| (p, 1))); + } + b.push((dir.join("git/db"), 1)); + b.push((dir.join("git/checkouts"), 1)); + b + } + // Content-addressed blobs: action results, outputs, downloads. + Kind::Bazel => vec![ + (dir.join("disk/cas"), 2), + (dir.join("disk/ac"), 2), + (dir.join("repo/content_addressable"), 2), + ], + // `//files//`; the index is kept. + Kind::Pnpm => children(&dir) + .into_iter() + .map(|v| (v.join("files"), 2)) + .collect(), + // The build cache only: module directories are read-only trees + // that other builds may be reading, so they are left alone. + Kind::Go => vec![(dir.join("build"), 2)], + // Entries of each versioned bucket (`wheels-v5`, `archive-v0`...). + Kind::Uv => children(&dir) + .into_iter() + .filter(|p| { + p.file_name() + .is_some_and(|n| n.to_string_lossy().contains("-v")) + }) + .map(|p| (p, 1)) + .collect(), + // Installed tool versions and downloads, never shims or config. + Kind::Mise => { + let mut b: Vec<(PathBuf, usize)> = children(&dir.join("installs")) + .into_iter() + .map(|t| (t, 1)) + .collect(); + b.push((dir.join("downloads"), 2)); + b + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn env_points_every_kind_at_a_stable_path_under_the_cache_root() { + let root = Path::new("/cache"); + let lease = Path::new("/work/abc"); + let all: Vec<(String, String)> = Kind::ALL + .iter() + .flat_map(|k| k.env(root, lease, "9G")) + .collect(); + let get = |k: &str| all.iter().find(|(n, _)| n == k).map(|(_, v)| v.as_str()); + assert_eq!(get("SCCACHE_DIR"), Some("/cache/sccache")); + assert_eq!(get("SCCACHE_CACHE_SIZE"), Some("9G")); + assert_eq!(get("CARGO_HOME"), Some("/cache/cargo")); + assert_eq!(get("BAZELISK_HOME"), Some("/cache/bazel/bazelisk")); + assert_eq!(get("PNPM_STORE_DIR"), Some("/cache/pnpm")); + assert_eq!(get("npm_config_store_dir"), Some("/cache/pnpm")); + assert_eq!(get("GOMODCACHE"), Some("/cache/go/mod")); + assert_eq!(get("GOCACHE"), Some("/cache/go/build")); + assert_eq!(get("UV_CACHE_DIR"), Some("/cache/uv")); + assert_eq!(get("MISE_DATA_DIR"), Some("/cache/mise")); + } + + #[test] + fn cargo_install_root_is_per_lease_so_bin_is_never_shared() { + let a = Kind::Cargo.env(Path::new("/c"), Path::new("/work/a"), "1G"); + let b = Kind::Cargo.env(Path::new("/c"), Path::new("/work/b"), "1G"); + let root = |env: &[(String, String)]| { + env.iter() + .find(|(k, _)| k == "CARGO_INSTALL_ROOT") + .unwrap() + .1 + .clone() + }; + assert_ne!(root(&a), root(&b)); + assert!(root(&a).starts_with("/work/a/")); + assert_eq!(a[0], b[0], "CARGO_HOME itself is the same for every lease"); + } + + #[test] + fn only_sccache_has_no_default_cap_and_no_gc_buckets() { + assert_eq!(Kind::Sccache.default_max_bytes(), None); + assert!(Kind::Sccache.gc_buckets(Path::new("/c")).is_empty()); + for k in Kind::ALL.into_iter().filter(|k| *k != Kind::Sccache) { + assert!(k.default_max_bytes().is_some(), "{k:?}"); + } + } + + #[test] + fn go_module_cache_and_bazelisk_are_never_gc_buckets() { + let go = Kind::Go.gc_buckets(Path::new("/c")); + assert_eq!(go, [(PathBuf::from("/c/go/build"), 2)]); + let bazel = Kind::Bazel.gc_buckets(Path::new("/c")); + assert!( + bazel + .iter() + .all(|(p, _)| !p.starts_with("/c/bazel/bazelisk")) + ); + } +} diff --git a/src/drone/caches/mod.rs b/src/drone/caches/mod.rs new file mode 100644 index 0000000..6e211c0 --- /dev/null +++ b/src/drone/caches/mod.rs @@ -0,0 +1,541 @@ +//! Host-level caches the drone owns so they outlive every lease: one stable +//! path per cache kind under `cache_root`, injected into every job's +//! environment, excluded from per-lease cleanup, size-capped by LRU garbage +//! collection that only runs while no lease is running, and reported in +//! heartbeats. +//! +//! Which kinds exist and how each is wired is in [`kinds`]; git mirrors are in +//! [`git`]; eviction is in [`gc`]; Bazel's rc file is in [`bazel`]. + +pub mod bazel; +pub mod gc; +pub mod git; +pub mod kinds; + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use tokio::sync::{OnceCell, OwnedRwLockReadGuard, RwLock}; + +use crate::api::CacheStat; +use crate::drone::config::DroneConfig; +use crate::drone::exec::Exec; +use crate::probe::{self, Host}; +use git::GitMirrors; +use kinds::Kind; + +/// Cache size when free space cannot be measured. +const FALLBACK_CACHE_SIZE: &str = "10G"; +const MAX_SCCACHE_GIB: u64 = 100; + +/// `SCCACHE_CACHE_SIZE` value: the configured size, else 25% of the free disk +/// capped at 100G (in M below 1G, so tiny disks still get a bounded cache). +pub fn cache_size_for(free_mib: Option, configured: Option<&str>) -> String { + if let Some(size) = configured { + return size.to_owned(); + } + let Some(free) = free_mib else { + return FALLBACK_CACHE_SIZE.to_owned(); + }; + let quarter = free / 4; + if quarter >= 1024 { + format!("{}G", (quarter / 1024).min(MAX_SCCACHE_GIB)) + } else { + format!("{}M", quarter.max(64)) + } +} + +#[derive(Default, Clone)] +struct Track { + size_bytes: u64, + evicted_bytes: u64, + last_gc_at: Option, +} + +pub struct Caches { + cfg: Arc, + host: Arc, + sccache_size: OnceCell, + git: GitMirrors, + /// Leases hold the read side for their whole life; garbage collection only + /// runs when it can take the write side, so it never overlaps a lease. + gate: Arc>, + tracks: Mutex>, +} + +fn unix_now() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(0)) +} + +impl Caches { + pub fn new( + cfg: Arc, + exec: Arc, + host: Arc, + ) -> Self { + let git = GitMirrors::new(&cfg.caches.git, &cfg.cache_root, exec); + Self { + cfg, + host, + sccache_size: OnceCell::new(), + git, + gate: Arc::default(), + tracks: Mutex::new(BTreeMap::new()), + } + } + + fn enabled(&self, kind: Kind) -> bool { + let c = &self.cfg.caches; + match kind { + Kind::Sccache => c.sccache.enabled, + Kind::Cargo => c.cargo.enabled, + Kind::Bazel => c.bazel.enabled, + Kind::Pnpm => c.pnpm.enabled, + Kind::Go => c.go.enabled, + Kind::Uv => c.uv.enabled, + Kind::Mise => c.mise.enabled, + } + } + + fn configured_max(&self, kind: Kind) -> Option<&str> { + let c = &self.cfg.caches; + match kind { + Kind::Sccache => c.sccache.max_size.as_deref(), + Kind::Cargo => c.cargo.max_size.as_deref(), + Kind::Bazel => c.bazel.max_size.as_deref(), + Kind::Pnpm => c.pnpm.max_size.as_deref(), + Kind::Go => c.go.max_size.as_deref(), + Kind::Uv => c.uv.max_size.as_deref(), + Kind::Mise => c.mise.max_size.as_deref(), + } + } + + /// Create the directories, size the compile cache and write Bazel's rc + /// file, once. Returns the `SCCACHE_CACHE_SIZE` value. + async fn init(&self) -> &str { + self.sccache_size + .get_or_init(|| async { + let (cfg, host) = (Arc::clone(&self.cfg), Arc::clone(&self.host)); + let enabled: Vec = + Kind::ALL.into_iter().filter(|k| self.enabled(*k)).collect(); + let made = tokio::task::spawn_blocking(move || { + for kind in &enabled { + for sub in kind.subdirs() { + let path = kind.dir(&cfg.cache_root).join(sub); + if let Err(e) = std::fs::create_dir_all(&path) { + tracing::warn!("create {}: {e}", path.display()); + } + } + } + if let Err(e) = std::fs::create_dir_all(cfg.cache_root.join("git")) { + tracing::warn!("create git cache dir: {e}"); + } + if cfg.caches.bazel.enabled { + write_bazelrc(&cfg); + } + let free = probe::free_disk_mib_at(host.as_ref(), &cfg.cache_root); + cache_size_for(free, cfg.caches.sccache.max_size.as_deref()) + }) + .await; + made.unwrap_or_else(|_| FALLBACK_CACHE_SIZE.to_owned()) + }) + .await + } + + /// Held for the whole life of a lease (refresh, run, cleanup). + pub async fn begin_lease(&self) -> OwnedRwLockReadGuard<()> { + Arc::clone(&self.gate).read_owned().await + } + + /// Refresh stale git mirrors, then return the cache environment for a new + /// lease running in `lease_dir`. + pub async fn prepare_lease(&self, lease_dir: &Path) -> Vec<(String, String)> { + let size = self.init().await.to_owned(); + self.git.refresh_stale().await; + let mut env = Vec::new(); + for kind in Kind::ALL.into_iter().filter(|k| self.enabled(*k)) { + env.extend(kind.env(&self.cfg.cache_root, lease_dir, &size)); + } + env.extend(self.git.env()); + env + } + + /// Re-measure every cache and, if no lease is running, evict least recently + /// used units from caches over their cap. Skips silently while a lease is + /// active. + pub async fn maintain(&self) { + let _ = self.init().await; + let Ok(_idle) = self.gate.try_write() else { + return; + }; + let plan: Vec<(Kind, Option)> = Kind::ALL + .into_iter() + .filter(|k| self.enabled(*k)) + .map(|k| { + let cap = self + .configured_max(k) + .and_then(gc::parse_size) + .or_else(|| k.default_max_bytes()); + (k, cap) + }) + .collect(); + let root = self.cfg.cache_root.clone(); + let work = tokio::task::spawn_blocking(move || { + let mut out: Vec<(&'static str, u64, u64)> = Vec::new(); + for (kind, cap) in plan { + let (mut size, _) = gc::measure(&kind.dir(&root)); + let mut freed = 0; + if let Some(cap) = cap + && size > cap + { + let units = gc::collect_units(&kind.gc_buckets(&root)); + let pool: u64 = units.iter().map(|u| u.size).sum(); + freed = gc::evict_lru(units, pool, cap).freed_bytes; + size = size.saturating_sub(freed); + } + out.push((kind.name(), size, freed)); + } + let (git_size, _) = gc::measure(&root.join("git")); + out.push(("git", git_size, 0)); + out + }) + .await; + let Ok(results) = work else { return }; + let now = unix_now(); + if let Ok(mut tracks) = self.tracks.lock() { + for (name, size, freed) in results { + let t = tracks.entry(name).or_default(); + t.size_bytes = size; + if freed > 0 { + t.evicted_bytes += freed; + t.last_gc_at = Some(now); + } + } + } + } + + /// Per-cache size, cap, eviction and (git) hit statistics for heartbeats. + pub fn stats(&self) -> Vec { + let tracks = self.tracks.lock().map(|t| t.clone()).unwrap_or_default(); + let sccache_cap = self.sccache_size.get().and_then(|s| gc::parse_size(s)); + let (served, missed) = self.git.hit_counts(); + tracks + .into_iter() + .map(|(name, t)| { + let kind = Kind::ALL.into_iter().find(|k| k.name() == name); + let max_bytes = match kind { + Some(Kind::Sccache) => sccache_cap, + Some(k) => self + .configured_max(k) + .and_then(gc::parse_size) + .or_else(|| k.default_max_bytes()), + None => None, + }; + let git = name == "git"; + CacheStat { + name: name.to_owned(), + size_bytes: t.size_bytes, + max_bytes, + evicted_bytes: t.evicted_bytes, + last_gc_at: t.last_gc_at, + hits: git.then_some(served), + misses: git.then_some(missed), + } + }) + .collect() + } + + /// Refresh git mirrors and run maintenance on a timer until aborted. + pub fn spawn_background(self: &Arc) -> tokio::task::JoinHandle<()> { + let caches = Arc::clone(self); + tokio::spawn(async move { + let _ = caches.init().await; + let every = Duration::from_secs( + caches + .cfg + .caches + .maintenance_interval_secs + .min(caches.cfg.caches.git.interval_secs), + ); + loop { + caches.git.refresh_all().await; + caches.maintain().await; + tokio::time::sleep(every).await; + } + }) + } +} + +/// Point Bazel at the host-level disk and repository caches. +fn write_bazelrc(cfg: &DroneConfig) { + let rc: Option = cfg + .caches + .bazel + .bazelrc + .clone() + .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".bazelrc"))); + let Some(rc) = rc else { + tracing::warn!( + "HOME is unset and caches.bazel.bazelrc is not configured; skipping Bazel cache flags" + ); + return; + }; + let dir = Kind::Bazel.dir(&cfg.cache_root); + if let Err(e) = bazel::write_rc(&rc, &dir.join("disk"), &dir.join("repo")) { + tracing::warn!("write {}: {e}", rc.display()); + } +} + +#[cfg(test)] +mod tests { + use std::time::SystemTime; + + use super::*; + use crate::drone::cleanup::testing::FakeExec; + + struct NoHost; + impl Host for NoHost { + fn read_to_string(&self, _: &Path) -> Option { + None + } + fn exists(&self, _: &Path) -> bool { + false + } + fn run(&self, program: &str, _: &[&str]) -> Option { + (program == "df").then(|| { + "Filesystem 1024-blocks Used Available Capacity Mounted on\n/dev/x 1 1 419430400 1% /".to_owned() + }) + } + fn os(&self) -> &'static str { + "linux" + } + fn arch(&self) -> &'static str { + "x86_64" + } + fn available_parallelism(&self) -> Option { + Some(1) + } + } + + fn caches(root: &Path, caches_json: &str) -> Arc { + let json = format!( + r#"{{"queen_url":"https://q.example.test","runner_dir":"/r","work_root":"/w","cache_root":"{}", + "caches":{caches_json}}}"#, + root.display() + ); + let cfg = Arc::new(DroneConfig::from_json(&json).unwrap()); + Arc::new(Caches::new( + cfg, + Arc::new(FakeExec::default()), + Arc::new(NoHost), + )) + } + + fn rc(root: &Path) -> String { + format!( + r#""bazel":{{"bazelrc":"{}/home/.bazelrc"}}"#, + root.display() + ) + } + + fn blob(path: &Path, bytes: usize, age_secs: u64) { + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(path, vec![1u8; bytes]).unwrap(); + let when = SystemTime::now() - Duration::from_secs(age_secs); + let times = std::fs::FileTimes::new() + .set_modified(when) + .set_accessed(when); + std::fs::File::options() + .write(true) + .open(path) + .unwrap() + .set_times(times) + .unwrap(); + } + + #[test] + fn cache_size_is_a_quarter_of_free_disk_capped_at_100g() { + let gib = |n: u64| Some(n * 1024); + assert_eq!(cache_size_for(gib(400), None), "100G"); + assert_eq!(cache_size_for(gib(2000), None), "100G"); + assert_eq!(cache_size_for(gib(200), None), "50G"); + assert_eq!(cache_size_for(gib(100), None), "25G"); + assert_eq!(cache_size_for(Some(2048), None), "512M"); + assert_eq!(cache_size_for(Some(100), None), "64M"); + assert_eq!(cache_size_for(None, None), "10G"); + assert_eq!( + cache_size_for(gib(400), Some("7G")), + "7G", + "configured size wins" + ); + } + + #[tokio::test] + async fn every_lease_gets_the_same_cache_paths_and_bazel_flags_are_written() { + let root = tempfile::tempdir().unwrap(); + let c = caches(root.path(), &format!("{{{}}}", rc(root.path()))); + let a = c.prepare_lease(Path::new("/work/a")).await; + let b = c.prepare_lease(Path::new("/work/b")).await; + let strip = |env: &[(String, String)]| -> Vec<(String, String)> { + env.iter() + .filter(|(k, _)| k != "CARGO_INSTALL_ROOT") + .cloned() + .collect() + }; + assert_eq!( + strip(&a), + strip(&b), + "shared paths are identical across leases" + ); + let get = |k: &str| a.iter().find(|(n, _)| n == k).map(|(_, v)| v.clone()); + assert_eq!( + get("SCCACHE_DIR"), + Some(format!("{}/sccache", root.path().display())) + ); + assert_eq!(get("SCCACHE_CACHE_SIZE").as_deref(), Some("100G")); + assert_eq!( + get("CARGO_INSTALL_ROOT").as_deref(), + Some("/work/a/cargo-install") + ); + for dir in [ + "sccache", + "cargo/registry", + "bazel/disk", + "bazel/repo", + "go/mod", + "go/build", + "uv", + "mise", + "pnpm", + "git", + ] { + assert!(root.path().join(dir).is_dir(), "{dir} created"); + } + let bazelrc = std::fs::read_to_string(root.path().join("home/.bazelrc")).unwrap(); + assert!(bazelrc.contains(&format!( + "--disk_cache={}/bazel/disk", + root.path().display() + ))); + assert!(bazelrc.contains(&format!( + "--repository_cache={}/bazel/repo", + root.path().display() + ))); + } + + #[tokio::test] + async fn disabled_kinds_get_no_env_no_dir_and_no_bazel_flags() { + let root = tempfile::tempdir().unwrap(); + let c = caches( + root.path(), + &format!( + r#"{{"sccache":{{"enabled":false}},"uv":{{"enabled":false}},"bazel":{{"enabled":false,"bazelrc":"{}/rc"}}}}"#, + root.path().display() + ), + ); + let env = c.prepare_lease(Path::new("/work/a")).await; + for var in ["SCCACHE_DIR", "UV_CACHE_DIR", "BAZELISK_HOME"] { + assert!(!env.iter().any(|(k, _)| k == var), "{var}"); + } + assert!(env.iter().any(|(k, _)| k == "CARGO_HOME")); + assert!(!root.path().join("sccache").exists() && !root.path().join("rc").exists()); + } + + #[tokio::test] + async fn git_mirror_variables_appear_only_for_mirrors_that_exist() { + let root = tempfile::tempdir().unwrap(); + let c = caches( + root.path(), + r#"{"git":{"mirrors":["example-org/app"],"remote_base":"https://git.example.test"}}"#, + ); + let present = git::mirror_path(root.path(), "example-org/app"); + std::fs::create_dir_all(&present).unwrap(); + std::fs::write(present.join("HEAD"), "ref: refs/heads/main\n").unwrap(); + let env = c.prepare_lease(Path::new("/work/a")).await; + assert!(env.contains(&( + "BORG_GIT_MIRROR_EXAMPLE_ORG_APP".to_owned(), + present.display().to_string() + ))); + c.maintain().await; + let git = c.stats().into_iter().find(|s| s.name == "git").unwrap(); + assert_eq!((git.hits, git.misses), (Some(1), Some(0))); + } + + #[tokio::test] + async fn gc_evicts_least_recently_used_between_leases_and_never_during_one() { + let root = tempfile::tempdir().unwrap(); + let c = caches(root.path(), r#"{"cargo":{"max_size":"1M"}}"#); + let reg = root.path().join("cargo/registry/cache/index-abc"); + for (name, age) in [("old.crate", 900), ("mid.crate", 500), ("new.crate", 100)] { + blob(®.join(name), 600 * 1024, age); + } + // Not collectable: the sccache dir and Go's module cache, however big. + blob(&root.path().join("sccache/big"), 3 * 1024 * 1024, 1000); + blob( + &root.path().join("go/mod/example.test/m@v1/f.go"), + 3 * 1024 * 1024, + 1000, + ); + + let lease = c.begin_lease().await; + c.maintain().await; + assert!( + reg.join("old.crate").exists(), + "no collection while a lease is active" + ); + drop(lease); + + c.maintain().await; + assert!(!reg.join("old.crate").exists()); + assert!(reg.join("new.crate").exists(), "recently used entries stay"); + assert!( + root.path().join("sccache/big").exists() + && root.path().join("go/mod/example.test/m@v1/f.go").exists() + ); + + let stats = c.stats(); + let cargo = stats.iter().find(|s| s.name == "cargo").unwrap(); + assert!(cargo.evicted_bytes >= 600 * 1024 && cargo.last_gc_at.is_some()); + assert_eq!(cargo.max_bytes, Some(1 << 20)); + assert!(cargo.size_bytes <= (1 << 20), "{}", cargo.size_bytes); + let sccache = stats.iter().find(|s| s.name == "sccache").unwrap(); + assert!(sccache.size_bytes >= 3 * 1024 * 1024 && sccache.evicted_bytes == 0); + assert_eq!( + sccache.max_bytes, + Some(100 << 30), + "sccache reports its own cap" + ); + } + + #[tokio::test] + async fn bazel_collection_leaves_the_bazelisk_download_cache_alone() { + let root = tempfile::tempdir().unwrap(); + let c = caches( + root.path(), + &format!( + r#"{{"bazel":{{"max_size":"1M","bazelrc":"{}/rc"}}}}"#, + root.path().display() + ), + ); + let cas = root.path().join("bazel/disk/cas"); + for (name, age) in [("aa/h1", 900), ("bb/h2", 100), ("cc/h3", 50)] { + blob(&cas.join(name), 700 * 1024, age); + } + blob( + &root.path().join("bazel/bazelisk/downloads/bazel-8"), + 3 * 1024 * 1024, + 1000, + ); + c.maintain().await; + assert!(!cas.join("aa/h1").exists()); + assert!(cas.join("cc/h3").exists()); + assert!( + root.path() + .join("bazel/bazelisk/downloads/bazel-8") + .exists() + ); + } +} diff --git a/src/drone/config.rs b/src/drone/config.rs index 6c42a55..576e275 100644 --- a/src/drone/config.rs +++ b/src/drone/config.rs @@ -20,6 +20,109 @@ pub enum DroneConfigError { Invalid(String), } +/// One cache directory the drone manages. Present-by-default; set +/// `enabled: false` to skip it. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct CacheEntry { + pub enabled: bool, + /// Size cap such as `512M` or `50G`. Enforced by LRU garbage collection + /// between leases (sccache enforces its own via `SCCACHE_CACHE_SIZE`). + /// Each kind has a default; sccache's is 25% of the free disk, at most 100G. + pub max_size: Option, +} + +impl Default for CacheEntry { + fn default() -> Self { + Self { + enabled: true, + max_size: None, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct BazelCache { + pub enabled: bool, + /// Cap for the disk cache and repository cache together. + pub max_size: Option, + /// The rc file the drone writes `--disk_cache` and `--repository_cache` + /// into (a managed block). Default: `$HOME/.bazelrc`. + pub bazelrc: Option, +} + +impl Default for BazelCache { + fn default() -> Self { + Self { + enabled: true, + max_size: None, + bazelrc: None, + } + } +} + +/// Bare git mirrors, one per repository. Off unless `mirrors` is non-empty. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct GitCaches { + /// Repositories (`owner/name`). Jobs get each mirror's path in + /// `BORG_GIT_MIRROR__`. + pub mirrors: Vec, + /// Where mirrors fetch from; `//.git`. Credentials are + /// the host's business (a credential helper, or an ssh base such as + /// `git@host:` with a read-only deploy key for the service user). + pub remote_base: String, + pub bin: String, + pub interval_secs: u64, +} + +impl Default for GitCaches { + fn default() -> Self { + Self { + mirrors: Vec::new(), + remote_base: "https://github.com".to_owned(), + bin: "git".to_owned(), + interval_secs: 600, + } + } +} + +/// Every cache kind the drone can manage. All kinds except git are on by +/// default; a cache nobody uses costs only the disk it fills, up to its cap. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct CachesConfig { + pub git: GitCaches, + pub sccache: CacheEntry, + /// `CARGO_HOME` registry and git data. `bin` is never shared between leases. + pub cargo: CacheEntry, + pub bazel: BazelCache, + pub pnpm: CacheEntry, + pub go: CacheEntry, + pub uv: CacheEntry, + pub mise: CacheEntry, + /// How often sizes are re-measured and garbage collection is attempted + /// while the drone is idle. + pub maintenance_interval_secs: u64, +} + +impl Default for CachesConfig { + fn default() -> Self { + Self { + git: GitCaches::default(), + sccache: CacheEntry::default(), + cargo: CacheEntry::default(), + bazel: BazelCache::default(), + pnpm: CacheEntry::default(), + go: CacheEntry::default(), + uv: CacheEntry::default(), + mise: CacheEntry::default(), + maintenance_interval_secs: 600, + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct DroneConfig { @@ -61,6 +164,9 @@ pub struct DroneConfig { pub max_leases: Option, #[serde(default = "default_reserve")] pub reserve_mem_mib: u64, + /// Host-level caches the drone manages under `cache_root`. + #[serde(default)] + pub caches: CachesConfig, } fn default_docker() -> String { @@ -113,6 +219,63 @@ impl DroneConfig { if self.ca_cert_path.is_some() && http { return bad("ca_cert_path only applies to an https:// queen_url"); } + self.validate_caches() + } + + fn validate_caches(&self) -> Result<(), DroneConfigError> { + let bad = |m: String| Err(DroneConfigError::Invalid(m)); + let c = &self.caches; + if c.git.interval_secs == 0 || c.maintenance_interval_secs == 0 { + return bad("cache intervals must be positive".to_owned()); + } + // The work root is wiped on start and after every job; caches must + // never sit inside it, nor it inside them. + if self.cache_root.starts_with(&self.work_root) + || self.work_root.starts_with(&self.cache_root) + { + return bad("cache_root and work_root must not contain one another".to_owned()); + } + let sizes = [ + ("sccache", &c.sccache.max_size), + ("cargo", &c.cargo.max_size), + ("bazel", &c.bazel.max_size), + ("pnpm", &c.pnpm.max_size), + ("go", &c.go.max_size), + ("uv", &c.uv.max_size), + ("mise", &c.mise.max_size), + ]; + for (name, size) in sizes { + if let Some(size) = size + && crate::drone::caches::gc::parse_size(size).is_none() + { + return bad(format!( + "caches.{name}.max_size {size:?} must look like 512M or 50G" + )); + } + } + let mut vars = Vec::new(); + for repo in &c.git.mirrors { + let parts: Vec<&str> = repo.split('/').collect(); + let part_ok = |p: &&str| { + !p.is_empty() + && *p != "." + && *p != ".." + && p.bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.')) + }; + if parts.len() != 2 || !parts.iter().all(part_ok) { + return bad(format!( + "caches.git.mirrors entry {repo:?} must look like owner/name" + )); + } + let var = crate::drone::caches::git::mirror_var(repo); + if vars.contains(&var) { + return bad(format!( + "caches.git.mirrors entry {repo:?} maps to the same variable {var} as another entry" + )); + } + vars.push(var); + } Ok(()) } @@ -187,6 +350,48 @@ mod tests { assert!(DroneConfig::from_json(&http_ca).is_err()); } + #[test] + fn cache_settings_have_defaults_and_are_validated() { + let with = |extra: &str| BASE.replace('}', &format!(",{extra}}}")); + let defaults = DroneConfig::from_json(BASE).unwrap().caches; + assert!(defaults.sccache.enabled && defaults.bazel.enabled && defaults.mise.enabled); + assert!(defaults.git.mirrors.is_empty()); + assert_eq!( + ( + defaults.git.remote_base.as_str(), + defaults.git.interval_secs + ), + ("https://github.com", 600) + ); + + let ok = with( + r#""caches":{"git":{"mirrors":["example-org/app"]},"sccache":{"max_size":"50G"},"uv":{"enabled":false}}"#, + ); + let cfg = DroneConfig::from_json(&ok).unwrap().caches; + assert_eq!(cfg.git.mirrors, ["example-org/app"]); + assert!(!cfg.uv.enabled && cfg.go.enabled); + for bad in [ + r#""caches":{"git":{"mirrors":["noslash"]}}"#, + r#""caches":{"git":{"mirrors":["a/b/c"]}}"#, + r#""caches":{"git":{"mirrors":["../x"]}}"#, + r#""caches":{"git":{"mirrors":["a-b/c","a/b-c"]}}"#, + r#""caches":{"sccache":{"max_size":"50"}}"#, + r#""caches":{"bazel":{"max_size":"G"}}"#, + r#""caches":{"git":{"interval_secs":0}}"#, + r#""caches":{"npm":{}}"#, + ] { + assert!(DroneConfig::from_json(&with(bad)).is_err(), "{bad}"); + } + } + + #[test] + fn caches_may_not_overlap_the_wiped_work_root() { + let nested = BASE.replace(r#""cache_root":"/c""#, r#""cache_root":"/w/cache""#); + assert!(DroneConfig::from_json(&nested).is_err()); + let parent = BASE.replace(r#""work_root":"/w""#, r#""work_root":"/c/work""#); + assert!(DroneConfig::from_json(&parent).is_err()); + } + #[test] fn token_file_is_trimmed_and_must_not_be_empty() { let dir = tempfile::tempdir().unwrap(); diff --git a/src/drone/install.rs b/src/drone/install.rs index 634c279..6a470bd 100644 --- a/src/drone/install.rs +++ b/src/drone/install.rs @@ -76,6 +76,10 @@ pub struct InstallOptions { /// Skip the "must be root" check (tests, `--root` staging). pub skip_root_check: bool, pub drain_timeout_secs: u64, + /// Repositories (`owner/name`) to mirror on the host. + pub git_mirrors: Vec, + /// Remote base for mirrors; the drone default (`https://github.com`) if none. + pub git_remote_base: Option, } /// Where everything lives, as paths on the real host. @@ -527,6 +531,17 @@ impl Installer<'_> { max_mem_mib: o.max_mem_mib, max_leases: o.max_leases, reserve_mem_mib: crate::allocator::HostCaps::default().reserve_mem_mib, + caches: crate::drone::config::CachesConfig { + git: crate::drone::config::GitCaches { + mirrors: o.git_mirrors.clone(), + remote_base: o + .git_remote_base + .clone() + .unwrap_or_else(|| "https://github.com".to_owned()), + ..Default::default() + }, + ..Default::default() + }, }; cfg.validate() .map_err(|e| InstallError::Preflight(e.to_string()))?; @@ -702,6 +717,8 @@ mod tests { dry_run: false, skip_root_check: true, drain_timeout_secs: 3600, + git_mirrors: vec!["example-org/app".into()], + git_remote_base: None, } } @@ -808,6 +825,7 @@ mod tests { ) .unwrap(); assert_eq!((cfg.max_vcpu, cfg.max_mem_mib), (Some(8), Some(24576))); + assert_eq!(cfg.caches.git.mirrors, ["example-org/app"]); assert!( e.root .join("var/lib/borg/runner/.borg-runner-version") diff --git a/src/drone/mod.rs b/src/drone/mod.rs index 289c799..d0de4c2 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -5,6 +5,7 @@ //! private directory. Cleanup after each job is unconditional. On shutdown it //! stops accepting leases and lets in-flight jobs finish, up to a timeout. +pub mod caches; pub mod cleanup; pub mod client; pub mod config; @@ -22,6 +23,7 @@ use tokio::task::JoinSet; use crate::api::{HeartbeatRequest, LeaseGrant, ResultRequest}; use crate::probe::{self, Host}; +use caches::Caches; use cleanup::{Scope, cleanup_docker, remove_lease_dir, sweep_work_root, valid_lease_id}; use client::{ClientError, QueenClient}; use config::DroneConfig; @@ -56,6 +58,7 @@ pub struct Drone { held: Held, abort: tokio::sync::watch::Sender, proc_root: Arc, + caches: Arc, } impl Drone { @@ -66,11 +69,18 @@ impl Drone { exec: Arc, ) -> Self { let (abort, _) = tokio::sync::watch::channel(false); + let cfg = Arc::new(cfg); + let caches = Arc::new(Caches::new( + Arc::clone(&cfg), + Arc::clone(&exec), + Arc::clone(&host), + )); Self { - cfg: Arc::new(cfg), + cfg, client: Arc::new(client), host, exec, + caches, held: Held::default(), abort, proc_root: Arc::from(Path::new("/proc")), @@ -132,6 +142,7 @@ impl Drone { caps: self.cfg.host_caps(), accepting, active_leases: self.held_ids(), + cache_stats: self.caches.stats(), }; let resp = self.client.heartbeat(&req).await?; let mut started = 0; @@ -145,6 +156,7 @@ impl Drone { exec: Arc::clone(&self.exec), abort: self.abort.subscribe(), proc_root: Arc::clone(&self.proc_root), + caches: Arc::clone(&self.caches), _held: HeldGuard { held: Arc::clone(&self.held), id: grant.lease_id.clone(), @@ -185,6 +197,7 @@ impl Drone { pub async fn run(self, shutdown: impl std::future::Future) { self.startup_sweep().await; let mut tasks = JoinSet::new(); + let cache_background = self.caches.spawn_background(); let mut every = tokio::time::interval(Duration::from_secs(self.cfg.heartbeat_secs)); every.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); tokio::pin!(shutdown); @@ -223,6 +236,7 @@ impl Drone { } } } + cache_background.abort(); } } @@ -232,17 +246,25 @@ struct LeaseContext { exec: Arc, abort: tokio::sync::watch::Receiver, proc_root: Arc, + caches: Arc, _held: HeldGuard, } /// Run one lease, clean up unconditionally, report the result. async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { + // Held for the whole lease: cache garbage collection waits for it. + let lease_guard = ctx.caches.begin_lease().await; + // Refresh stale git mirrors and learn the shared cache environment. Cache + // trouble never blocks a job; it just runs without the speedup. + let lease_dir = ctx.cfg.work_root.join(&grant.lease_id); + let extra_env = ctx.caches.prepare_lease(&lease_dir).await; let spec = RunSpec { lease_id: &grant.lease_id, jit_config: &grant.jit_config, runner_dir: &ctx.cfg.runner_dir, work_root: &ctx.cfg.work_root, cache_root: &ctx.cfg.cache_root, + extra_env: &extra_env, max_job: Duration::from_secs(ctx.cfg.max_job_secs), proc_root: &ctx.proc_root, }; @@ -272,6 +294,13 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { Err(e) => tracing::warn!("lease {} cleanup task failed: {e}", grant.lease_id), } + // Between leases: let the caches re-measure and, if nothing else is + // running, evict. Skipped by `maintain` itself while another lease holds + // the gate. + drop(lease_guard); + let caches = Arc::clone(&ctx.caches); + tokio::spawn(async move { caches.maintain().await }); + let result = match outcome { Ok(o) => ResultRequest { exit_code: o.exit_code, diff --git a/src/drone/runner.rs b/src/drone/runner.rs index f7ef2c4..edced8f 100644 --- a/src/drone/runner.rs +++ b/src/drone/runner.rs @@ -29,6 +29,8 @@ pub struct RunSpec<'a> { pub runner_dir: &'a Path, pub work_root: &'a Path, pub cache_root: &'a Path, + /// Extra job environment (cache directories, git mirrors). + pub extra_env: &'a [(String, String)], pub max_job: Duration, /// `/proc` (a different root is used by tests). pub proc_root: &'a Path, @@ -160,6 +162,7 @@ pub async fn run_lease( // user can read from `ps` or `/proc//cmdline`. .env("ACTIONS_RUNNER_INPUT_JITCONFIG", spec.jit_config.expose()) .current_dir(&lease_dir) + .envs(spec.extra_env.iter().map(|(k, v)| (k, v))) .env("BORG_LEASE_ID", spec.lease_id) .env("BORG_DOCKER_PREFIX", &prefix) .env("COMPOSE_PROJECT_NAME", &prefix) @@ -248,6 +251,7 @@ mod tests { runner_dir: &runner, work_root: &work, cache_root: &dir.join("cache"), + extra_env: &[], max_job, proc_root: Path::new("/proc"), }; @@ -355,6 +359,7 @@ mod tests { runner_dir: &runner, work_root: &work, cache_root: &cache, + extra_env: &[], max_job: Duration::from_secs(60), proc_root: Path::new("/proc"), }; @@ -380,6 +385,7 @@ mod tests { runner_dir: &runner, work_root: &dir.path().join("work"), cache_root: dir.path(), + extra_env: &[], max_job: Duration::from_secs(1), proc_root: Path::new("/proc"), }; diff --git a/src/main.rs b/src/main.rs index 5911fe5..6c7a03d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -107,6 +107,14 @@ struct InstallArgs { /// Install under this root instead of / (staging and tests). #[arg(long, default_value = "/")] root: PathBuf, + /// Keep a host-local bare git mirror of this repository (owner/name); + /// repeatable. Jobs get its path in BORG_GIT_MIRROR__. + #[arg(long = "git-mirror")] + git_mirrors: Vec, + /// Where mirrors fetch from (default https://github.com). Use an ssh base + /// such as git@host: with a read-only deploy key for private repositories. + #[arg(long)] + git_remote_base: Option, /// Print what would change without changing anything. #[arg(long)] dry_run: bool, @@ -215,6 +223,8 @@ fn install_drone(a: InstallArgs) -> Result<()> { binary: std::env::current_exe().context("locate this binary")?, dry_run: a.dry_run, drain_timeout_secs: 3600, + git_mirrors: a.git_mirrors, + git_remote_base: a.git_remote_base, }; for line in install::install(&opts, &SystemExec)? { println!("{line}"); diff --git a/src/probe/mod.rs b/src/probe/mod.rs index faf140c..b617cc2 100644 --- a/src/probe/mod.rs +++ b/src/probe/mod.rs @@ -186,8 +186,13 @@ fn cgroup_mem_limit_mib(host: &dyn Host) -> Option { } fn free_disk_mib(host: &dyn Host) -> Option { + free_disk_mib_at(host, Path::new(DISK_PROBE_PATH)) +} + +/// Free space, in MiB, on the filesystem holding `path`. +pub fn free_disk_mib_at(host: &dyn Host, path: &Path) -> Option { // POSIX `df -Pk`: header line, then `fs blocks used avail cap mount`. - let out = host.run("df", &["-Pk", DISK_PROBE_PATH])?; + let out = host.run("df", &["-Pk", &path.to_string_lossy()])?; let avail_kib = out .lines() .nth(1)? diff --git a/src/store.rs b/src/store.rs index 0161d97..7d6124a 100644 --- a/src/store.rs +++ b/src/store.rs @@ -13,6 +13,7 @@ use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; use sqlx::{Row, SqlitePool}; use crate::allocator::HostCaps; +use crate::api::CacheStat; use crate::classes::{Class, Scope}; use crate::probe::Capacity; use crate::secret::Secret; @@ -25,7 +26,8 @@ CREATE TABLE IF NOT EXISTS drones ( last_heartbeat INTEGER, capacity TEXT, caps TEXT, - accepting INTEGER NOT NULL DEFAULT 1 + accepting INTEGER NOT NULL DEFAULT 1, + cache_stats TEXT ); CREATE TABLE IF NOT EXISTS leases ( id TEXT PRIMARY KEY, @@ -100,6 +102,7 @@ pub struct DroneRow { pub capacity: Option, pub caps: HostCaps, pub accepting: bool, + pub cache_stats: Vec, } #[derive(Debug, Clone)] @@ -170,6 +173,15 @@ impl Store { .connect_with(opts) .await?; sqlx::raw_sql(SCHEMA).execute(&pool).await?; + // Databases created before cache stats existed lack the column. + // sqlx-runtime-query: one-off additive migration, no data involved + if let Err(e) = sqlx::query("ALTER TABLE drones ADD COLUMN cache_stats TEXT") + .execute(&pool) + .await + && !e.to_string().contains("duplicate column") + { + return Err(e.into()); + } Ok(Self { pool }) } @@ -208,14 +220,16 @@ impl Store { capacity: &Capacity, caps: &HostCaps, accepting: bool, + cache_stats: &[CacheStat], now: i64, ) -> Result<(), StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database - sqlx::query("UPDATE drones SET last_heartbeat = ?, capacity = ?, caps = ?, accepting = ? WHERE name = ?") + sqlx::query("UPDATE drones SET last_heartbeat = ?, capacity = ?, caps = ?, accepting = ?, cache_stats = ? WHERE name = ?") .bind(now) .bind(to_text(capacity)?) .bind(to_text(caps)?) .bind(i64::from(accepting)) + .bind(to_text(&cache_stats)?) .bind(name) .execute(&self.pool) .await?; @@ -225,7 +239,7 @@ impl Store { pub async fn drones(&self) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let rows = sqlx::query( - "SELECT name, last_heartbeat, capacity, caps, accepting FROM drones ORDER BY name", + "SELECT name, last_heartbeat, capacity, caps, accepting, cache_stats FROM drones ORDER BY name", ) .fetch_all(&self.pool) .await?; @@ -233,6 +247,7 @@ impl Store { .map(|r| { let capacity: Option = r.get("capacity"); let caps: Option = r.get("caps"); + let stats: Option = r.get("cache_stats"); Ok(DroneRow { name: r.get("name"), last_heartbeat: r.get("last_heartbeat"), @@ -243,6 +258,11 @@ impl Store { .transpose()? .unwrap_or_default(), accepting: r.get::("accepting") != 0, + cache_stats: stats + .as_deref() + .map(from_text) + .transpose()? + .unwrap_or_default(), }) }) .collect() @@ -512,11 +532,28 @@ mod tests { mise_version: None, git_version: None, }; - s.record_heartbeat("d1", &cap, &HostCaps::default(), false, 123) - .await - .unwrap(); + let stat = CacheStat { + name: "cargo".into(), + size_bytes: 5, + max_bytes: Some(10), + evicted_bytes: 1, + last_gc_at: Some(100), + hits: None, + misses: None, + }; + s.record_heartbeat( + "d1", + &cap, + &HostCaps::default(), + false, + std::slice::from_ref(&stat), + 123, + ) + .await + .unwrap(); let d = &s.drones().await.unwrap()[0]; assert_eq!((d.last_heartbeat, d.accepting), (Some(123), false)); + assert_eq!(d.cache_stats, [stat]); assert_eq!(d.capacity.as_ref(), Some(&cap)); } } diff --git a/tests/drone.rs b/tests/drone.rs index ee19b67..cd1c55f 100644 --- a/tests/drone.rs +++ b/tests/drone.rs @@ -10,7 +10,7 @@ use std::time::Duration; use borg::drone::Drone; use borg::drone::client::QueenClient; use borg::drone::config::DroneConfig; -use borg::drone::exec::{Exec, Output}; +use borg::drone::exec::{Exec, Output, SystemExec}; use borg::probe::Host; use borg::secret::Secret; use serde_json::{Value, json}; @@ -96,6 +96,10 @@ impl Rig { } fn drone(&self, extra: Value) -> Drone { + self.drone_with(extra, Arc::clone(&self.exec) as Arc) + } + + fn drone_with(&self, extra: Value, exec: Arc) -> Drone { let mut cfg = json!({ "queen_url": self.h.url, "allow_insecure": true, @@ -108,14 +112,17 @@ impl Rig { for (k, v) in extra.as_object().cloned().unwrap_or_default() { cfg[k] = v; } + // Never let a test write the real user's ~/.bazelrc. + if cfg["caches"].is_null() { + cfg["caches"] = json!({}); + } + if cfg["caches"]["bazel"].is_null() { + cfg["caches"]["bazel"] = json!({}); + } + cfg["caches"]["bazel"]["bazelrc"] = json!(self.path("bazelrc")); let cfg = DroneConfig::from_json(&cfg.to_string()).unwrap(); let client = QueenClient::new(&cfg, Secret::new(self.token.clone())).unwrap(); - Drone::new( - cfg, - client, - Arc::new(FixedHost), - Arc::clone(&self.exec) as Arc, - ) + Drone::new(cfg, client, Arc::new(FixedHost), exec) } async fn lease_state(&self) -> Vec { @@ -289,3 +296,212 @@ async fn shutdown_stops_taking_leases_and_lets_the_running_job_finish() { .is_empty() ); } + +/// Real `git` for mirror work; everything else is recorded, not executed. +struct GitOnlyExec(RecordingExec); + +impl Exec for GitOnlyExec { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + if program == "git" { + SystemExec.run(program, args) + } else { + self.0.run(program, args) + } + } +} + +fn git(args: &[&str]) -> String { + let out = SystemExec.run("git", args).unwrap(); + assert!(out.success, "git {args:?} failed"); + out.stdout +} + +/// A bare origin at `/example-org/app.git` with one commit on `main`; +/// returns (base, work clone dir). +fn origin(root: &Path) -> (PathBuf, PathBuf) { + let base = root.join("origin"); + let bare = base.join("example-org/app.git"); + let work = root.join("work-clone"); + fs::create_dir_all(bare.parent().unwrap()).unwrap(); + git(&[ + "init", + "--bare", + "--quiet", + "-b", + "main", + bare.to_str().unwrap(), + ]); + git(&["init", "--quiet", "-b", "main", work.to_str().unwrap()]); + git(&[ + "-C", + work.to_str().unwrap(), + "remote", + "add", + "origin", + bare.to_str().unwrap(), + ]); + commit(&work, "one"); + (base, work) +} + +fn commit(work: &Path, msg: &str) -> String { + let w = work.to_str().unwrap(); + git(&[ + "-C", + w, + "-c", + "user.name=t", + "-c", + "user.email=t@example.test", + "commit", + "--quiet", + "--allow-empty", + "-m", + msg, + ]); + git(&["-C", w, "push", "--quiet", "origin", "main"]); + git(&["-C", w, "rev-parse", "HEAD"]) +} + +#[tokio::test] +async fn jobs_share_cache_dirs_and_a_git_mirror_that_survive_cleanup() { + let rig = Rig::new().await; + let (base, _work) = origin(rig.dir.path()); + rig.install_runner( + "printf '%s' \"$BORG_GIT_MIRROR_EXAMPLE_ORG_APP\" > $OUT/mirror-$BORG_LEASE_ID.txt\n\ + git --git-dir \"$BORG_GIT_MIRROR_EXAMPLE_ORG_APP\" rev-parse main > $OUT/head-$BORG_LEASE_ID.txt\n\ + echo \"$SCCACHE_DIR|$CARGO_HOME|$SCCACHE_CACHE_SIZE\" > $OUT/caches-$BORG_LEASE_ID.txt\n\ + echo \"$UV_CACHE_DIR|$GOCACHE|$GOMODCACHE|$MISE_DATA_DIR|$PNPM_STORE_DIR|$BAZELISK_HOME\" > $OUT/more-$BORG_LEASE_ID.txt\n\ + echo \"$CARGO_INSTALL_ROOT\" > $OUT/install-$BORG_LEASE_ID.txt\n\ + touch \"$SCCACHE_DIR/entry-$BORG_LEASE_ID\"", + ); + let exec = Arc::new(GitOnlyExec(RecordingExec::default())); + let drone = rig.drone_with( + json!({"caches": { + "git": {"mirrors": ["example-org/app"], "remote_base": base.to_str().unwrap()}, + "sccache": {"max_size": "7G"}, + }}), + exec, + ); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + + for job in [1, 2] { + rig.h.state().jobs.clear(); + rig.h.queue_job(job, "borg-ci-light", 1); + rig.h.tick().await; + assert_eq!(drone.heartbeat_once(true, &mut tasks).await.unwrap(), 1); + drain(&mut tasks).await; + } + + let read = |prefix: &str| -> Vec { + let mut found: Vec = fs::read_dir(rig.path("out")) + .unwrap() + .filter_map(|e| e.ok()) + .filter(|e| e.file_name().to_string_lossy().starts_with(prefix)) + .map(|e| fs::read_to_string(e.path()).unwrap().trim().to_owned()) + .collect(); + found.sort(); + found + }; + let cache = rig.path("cache"); + let mirror = cache.join("git/example-org/app.git").display().to_string(); + assert_eq!( + read("mirror-"), + [mirror.clone(), mirror], + "both leases see the mirror" + ); + let caches = read("caches-"); + assert_eq!(caches.len(), 2); + assert_eq!( + caches[0], caches[1], + "same cache paths and size for every lease" + ); + assert_eq!( + caches[0], + format!("{0}/sccache|{0}/cargo|7G", cache.display()) + ); + let more = read("more-"); + assert_eq!(more[0], more[1]); + assert_eq!( + more[0], + format!( + "{0}/uv|{0}/go/build|{0}/go/mod|{0}/mise|{0}/pnpm|{0}/bazel/bazelisk", + cache.display() + ) + ); + let installs = read("install-"); + assert_ne!( + installs[0], installs[1], + "cargo install roots are per lease" + ); + let bazelrc = fs::read_to_string(rig.path("bazelrc")).unwrap(); + assert!(bazelrc.contains(&format!("--disk_cache={}/bazel/disk", cache.display()))); + let heads = read("head-"); + assert_eq!(heads[0], heads[1]); + assert_eq!(heads[0].len(), 40, "mirror serves the origin commit"); + + assert!( + cache.join("git/example-org/app.git/HEAD").exists(), + "mirror survives cleanup" + ); + let entries = fs::read_dir(cache.join("sccache")).unwrap().count(); + assert_eq!( + entries, 2, + "compile cache entries from both leases survive cleanup" + ); + assert_eq!(fs::read_dir(rig.path("work")).unwrap().count(), 0); +} + +#[tokio::test] +async fn the_mirror_timer_picks_up_new_commits() { + let rig = Rig::new().await; + let (base, work) = origin(rig.dir.path()); + let json_cfg = json!({ + "queen_url": "http://127.0.0.1:1", "allow_insecure": true, + "runner_dir": rig.path("runner"), "work_root": rig.path("work"), "cache_root": rig.path("cache"), + "caches": { + "git": {"mirrors": ["example-org/app"], "remote_base": base.to_str().unwrap(), "interval_secs": 1}, + "bazel": {"enabled": false}, + }, + }); + let cfg = Arc::new(DroneConfig::from_json(&json_cfg.to_string()).unwrap()); + let caches = Arc::new(borg::drone::caches::Caches::new( + cfg, + Arc::new(SystemExec), + Arc::new(FixedHost), + )); + let timer = caches.spawn_background(); + let mirror = rig.path("cache/git/example-org/app.git"); + let head = |expected: &str| { + let m = mirror.to_str().unwrap().to_owned(); + let expected = expected.to_owned(); + async move { + for _ in 0..100 { + let out = SystemExec.run("git", &["--git-dir", &m, "rev-parse", "main"]); + if out.is_ok_and(|o| o.success && o.stdout == expected) { + return true; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + false + } + }; + let first = git(&["-C", work.to_str().unwrap(), "rev-parse", "HEAD"]); + assert!(head(&first).await, "initial clone"); + let second = commit(&work, "two"); + assert!(head(&second).await, "timer refresh fetched the new commit"); + timer.abort(); + + // A job can borrow the mirror's objects. + let clone = rig.path("job-clone"); + git(&[ + "clone", + "--quiet", + "--reference", + mirror.to_str().unwrap(), + base.join("example-org/app.git").to_str().unwrap(), + clone.to_str().unwrap(), + ]); + assert!(clone.join(".git/objects/info/alternates").exists()); +} From b273b085e7cba4d7f015f984a55fb10783dd7ca4 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 10 Oct 2026 19:58:51 +0000 Subject: [PATCH 03/27] chore: project borg from Mono 95eee04b8127 --- .repository-projection.json | 6 +- README.md | 50 ++++++++++-- src/drone/caches/mod.rs | 112 +++++++++++++++++++++++-- src/drone/caches/rustflags.rs | 148 ++++++++++++++++++++++++++++++++++ src/drone/caches/sccache.rs | 110 +++++++++++++++++++++++++ src/drone/config.rs | 133 +++++++++++++++++++++++++++++- src/drone/install.rs | 1 + src/drone/mod.rs | 49 ++++++++--- src/drone/runner.rs | 30 ++++--- src/drone/slots.rs | 112 +++++++++++++++++++++++++ tests/drone.rs | 82 +++++++++++++++++-- tests/sccache_slots.rs | 123 ++++++++++++++++++++++++++++ 12 files changed, 910 insertions(+), 46 deletions(-) create mode 100644 src/drone/caches/rustflags.rs create mode 100644 src/drone/caches/sccache.rs create mode 100644 src/drone/slots.rs create mode 100644 tests/sccache_slots.rs diff --git a/.repository-projection.json b/.repository-projection.json index 96b68f5..4df2a52 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "3bba595c27f8777722f938346c87d96c2edfc42a", + "sourceSha": "95eee04b812703f5b31b593f8d10e7d55d4a3a0f", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "1f067ea6a8594a1ea831b7efe50d5c857138ec6b", + "priorProjectedBase": "101caa77b052562b6862c9689347e746a301735b", "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "02f5f780bf1b6ce76b937b699070ba605cf9f8854d1466938481f5327c45f1cf", + "contentDigest": "955724be218525b4b86854d0976c551258e5f616da68af35eeea95dbe293bfe5", "publicationEligible": true } diff --git a/README.md b/README.md index 35fc48a..f4c56ed 100644 --- a/README.md +++ b/README.md @@ -159,8 +159,8 @@ garbage collection that runs only while no lease is running. `cache_root` and | Kind | Environment for jobs | Default cap | Collected | | --- | --- | --- | --- | -| `sccache` | `SCCACHE_DIR`, `SCCACHE_CACHE_SIZE` | 25% of free disk at start, at most 100G | never by the drone; sccache enforces its own size | -| `cargo` | `CARGO_HOME` (registry and git), `CARGO_INSTALL_ROOT` per lease | 20G | crates, unpacked sources, git db and checkouts | +| `sccache` | `SCCACHE_DIR`, `SCCACHE_CACHE_SIZE` (and `SCCACHE_SERVER_PORT` with the optional host server) | 25% of free disk at start, at most 100G | never by the drone; sccache enforces its own size | +| `cargo` | `CARGO_HOME` (registry and git), `CARGO_INSTALL_ROOT` per lease, `BORG_RUSTFLAGS` and `RUSTFLAGS` | 20G | crates, unpacked sources, git db and checkouts | | `bazel` | `BAZELISK_HOME`; `--disk_cache` and `--repository_cache` via a managed block in `~/.bazelrc` | 50G | `disk/cas`, `disk/ac`, repository cache entries | | `pnpm` | `PNPM_STORE_DIR`, `npm_config_store_dir` | 20G | store files (the index is kept) | | `go` | `GOMODCACHE`, `GOCACHE` | 20G | build cache only; the module cache is read-only trees and is left alone | @@ -218,12 +218,46 @@ Notes on each kind: an ssh base such as `git@host:` with a read-only deploy key. A failed refresh keeps the previous mirror and never blocks a job. -What the drone does not do: it does not inject `RUSTFLAGS`. A -`--remap-path-prefix` can only make paths stable if the path it names is stable, -and a lease's work directory contains the lease id; workspace crates therefore -do not share compile-cache entries across leases, while dependencies built from -the shared registry do. It does not collect sccache hit rates, because sccache -servers are job-scoped. +#### Slots, path remapping and compile-cache keys + +A lease runs in `/slot-` and takes the lowest free slot +(`slots` in the drone config, default `max_leases`, else 16; it is also the most +leases that run at once). The path never contains the lease id, concurrent +leases always get distinct slots, and a slot is fully cleaned before it is +reused, so the next lease in a slot sees exactly the same absolute paths. The +lease-to-slot mapping is logged (`lease X -> slot N`) and kept in drone state. +Jobs see `BORG_SLOT`. + +Jobs also get `BORG_RUSTFLAGS` and `RUSTFLAGS`: `--remap-path-prefix` flags that +map **every** slot directory onto `/borg/work`, and `CARGO_HOME`, its registry +and its git data onto fixed virtual prefixes. The string lists all slots, so it +is byte-identical for every lease (a per-slot flag would itself differ between +slots and change the cache key). The same flags are written to +`$CARGO_HOME/config.toml` as `[build] rustflags`, as a fallback. Two cautions: + +* Cargo ignores `build.rustflags` and `target.*.rustflags` from config files + whenever `RUSTFLAGS` is set. A workflow that sets its own `RUSTFLAGS` replaces + ours; it should append the remap instead, e.g. + `RUSTFLAGS: "-D warnings ${BORG_RUSTFLAGS}"`. Repositories that depend on + config-file rustflags can turn the export off with + `caches.cargo.export_rustflags: false`. +* The remap does **not** make workspace crates hit the compile cache across + different slots. sccache hashes the compile's working directory and every + `CARGO_*` variable, including `CARGO_MANIFEST_DIR` (read from sccache's + `generate_hash_key`, v0.16), and no rustc flag changes those. What slots buy is + that a *reused* slot recompiles at identical paths and hits, and that + registry dependencies, whose paths are in the shared `CARGO_HOME`, hit from any + slot. `tests/sccache_slots.rs` compiles one crate in slot 0, again in a cleaned + slot 0, and in slot 1, and prints the hit counts for each. + +sccache keeps hit and miss counters only in the server's memory, so after a +lease's own server exits there is nothing on disk to read. Set +`caches.sccache.server: {"port": 4226}` to have the drone run one long-lived +sccache server per host (its configuration then applies to all jobs, and a job +whose sccache version differs from the host binary's restarts it); the drone +reads its totals with `sccache --show-stats --stats-format json` and reports +hits and misses for `sccache` in heartbeats, `/status` and `borg_cache_hits`. +Without it those fields stay empty. Every heartbeat carries each cache's size, cap, evicted bytes and last collection time (and, for git, how many leases were given a mirror and how many diff --git a/src/drone/caches/mod.rs b/src/drone/caches/mod.rs index 6e211c0..3527437 100644 --- a/src/drone/caches/mod.rs +++ b/src/drone/caches/mod.rs @@ -11,6 +11,8 @@ pub mod bazel; pub mod gc; pub mod git; pub mod kinds; +pub mod rustflags; +pub mod sccache; use std::collections::BTreeMap; use std::path::{Path, PathBuf}; @@ -52,10 +54,14 @@ struct Track { size_bytes: u64, evicted_bytes: u64, last_gc_at: Option, + /// sccache only, from the host server. + hits: Option, + misses: Option, } pub struct Caches { cfg: Arc, + exec: Arc, host: Arc, sccache_size: OnceCell, git: GitMirrors, @@ -77,9 +83,10 @@ impl Caches { exec: Arc, host: Arc, ) -> Self { - let git = GitMirrors::new(&cfg.caches.git, &cfg.cache_root, exec); + let git = GitMirrors::new(&cfg.caches.git, &cfg.cache_root, Arc::clone(&exec)); Self { cfg, + exec, host, sccache_size: OnceCell::new(), git, @@ -119,7 +126,12 @@ impl Caches { async fn init(&self) -> &str { self.sccache_size .get_or_init(|| async { - let (cfg, host) = (Arc::clone(&self.cfg), Arc::clone(&self.host)); + let (cfg, host, exec) = ( + Arc::clone(&self.cfg), + Arc::clone(&self.host), + Arc::clone(&self.exec), + ); + let flags = self.remap_flags(); let enabled: Vec = Kind::ALL.into_iter().filter(|k| self.enabled(*k)).collect(); let made = tokio::task::spawn_blocking(move || { @@ -137,8 +149,32 @@ impl Caches { if cfg.caches.bazel.enabled { write_bazelrc(&cfg); } + if cfg.caches.cargo.enabled + && cfg.caches.cargo.export_rustflags + && rustflags::flags_are_passable(&flags) + && let Err(e) = + rustflags::write_config(&Kind::Cargo.dir(&cfg.cache_root), &flags) + { + tracing::warn!("cargo config.toml: {e}"); + } let free = probe::free_disk_mib_at(host.as_ref(), &cfg.cache_root); - cache_size_for(free, cfg.caches.sccache.max_size.as_deref()) + let size = cache_size_for(free, cfg.caches.sccache.max_size.as_deref()); + if let (true, Some(server)) = + (cfg.caches.sccache.enabled, &cfg.caches.sccache.server) + { + let dir = Kind::Sccache + .dir(&cfg.cache_root) + .to_string_lossy() + .into_owned(); + let env = sccache::server_env(&dir, &size, server.port); + if !sccache::start_server(exec.as_ref(), &server.bin, &env) { + tracing::warn!( + "could not start the sccache server on port {}", + server.port + ); + } + } + size }) .await; made.unwrap_or_else(|_| FALLBACK_CACHE_SIZE.to_owned()) @@ -146,6 +182,15 @@ impl Caches { .await } + /// The remap flags for every slot and the cargo home (identical for all leases). + fn remap_flags(&self) -> Vec { + rustflags::remap_flags( + &self.cfg.work_root, + self.cfg.slot_count(), + &Kind::Cargo.dir(&self.cfg.cache_root), + ) + } + /// Held for the whole life of a lease (refresh, run, cleanup). pub async fn begin_lease(&self) -> OwnedRwLockReadGuard<()> { Arc::clone(&self.gate).read_owned().await @@ -161,14 +206,65 @@ impl Caches { env.extend(kind.env(&self.cfg.cache_root, lease_dir, &size)); } env.extend(self.git.env()); + if self.cfg.caches.cargo.enabled && self.cfg.caches.cargo.export_rustflags { + let flags = self.remap_flags(); + if rustflags::flags_are_passable(&flags) { + let joined = flags.join(" "); + env.push(("BORG_RUSTFLAGS".to_owned(), joined.clone())); + env.push(("RUSTFLAGS".to_owned(), joined)); + } else { + tracing::warn!( + "work_root or cache_root contains whitespace; not exporting RUSTFLAGS" + ); + } + } + if let (true, Some(server)) = ( + self.cfg.caches.sccache.enabled, + &self.cfg.caches.sccache.server, + ) { + env.push(("SCCACHE_SERVER_PORT".to_owned(), server.port.to_string())); + env.push(("SCCACHE_IDLE_TIMEOUT".to_owned(), "0".to_owned())); + } env } + /// Cumulative hits and misses of the host sccache server, if configured. + async fn sccache_counts(&self) -> Option<(u64, u64)> { + let (true, Some(server)) = ( + self.cfg.caches.sccache.enabled, + self.cfg.caches.sccache.server.clone(), + ) else { + return None; + }; + let size = self.init().await.to_owned(); + let exec = Arc::clone(&self.exec); + let dir = Kind::Sccache + .dir(&self.cfg.cache_root) + .to_string_lossy() + .into_owned(); + tokio::task::spawn_blocking(move || { + let env = sccache::server_env(&dir, &size, server.port); + sccache::show_stats(exec.as_ref(), &server.bin, &env) + }) + .await + .ok() + .flatten() + } + /// Re-measure every cache and, if no lease is running, evict least recently /// used units from caches over their cap. Skips silently while a lease is /// active. pub async fn maintain(&self) { let _ = self.init().await; + // Hit counters are cheap to read and live in the server, so they are + // refreshed even while leases run. + if let Some((hits, misses)) = self.sccache_counts().await + && let Ok(mut tracks) = self.tracks.lock() + { + let t = tracks.entry("sccache").or_default(); + t.hits = Some(hits); + t.misses = Some(misses); + } let Ok(_idle) = self.gate.try_write() else { return; }; @@ -235,15 +331,19 @@ impl Caches { .or_else(|| k.default_max_bytes()), None => None, }; - let git = name == "git"; + let (hits, misses) = if name == "git" { + (Some(served), Some(missed)) + } else { + (t.hits, t.misses) + }; CacheStat { name: name.to_owned(), size_bytes: t.size_bytes, max_bytes, evicted_bytes: t.evicted_bytes, last_gc_at: t.last_gc_at, - hits: git.then_some(served), - misses: git.then_some(missed), + hits, + misses, } }) .collect() diff --git a/src/drone/caches/rustflags.rs b/src/drone/caches/rustflags.rs new file mode 100644 index 0000000..891b4a1 --- /dev/null +++ b/src/drone/caches/rustflags.rs @@ -0,0 +1,148 @@ +//! Stable path remapping for Rust builds. +//! +//! Every lease runs in `/slot-`, so a lease's absolute paths are +//! the same each time its slot is reused. The remap flags below make the paths +//! rustc embeds (debuginfo, panic locations, rlib contents) identical across +//! *all* slots and leases by mapping every slot directory, and the shared cargo +//! directories, onto fixed virtual prefixes. The flag string lists every slot, +//! so it is byte-identical for every lease; a per-slot flag would itself differ +//! between slots and change the compile-cache key. +//! +//! What this cannot do: sccache also hashes the compile's working directory and +//! every `CARGO_*` variable, including `CARGO_MANIFEST_DIR`, and rustc flags +//! cannot change those. Workspace crates therefore hit the compile cache when a +//! slot is reused, not across different slots. + +use std::fs; +use std::io; +use std::path::Path; + +/// Virtual prefix a slot directory is mapped to. +pub const VIRTUAL_WORK: &str = "/borg/work"; +const VIRTUAL_CARGO: &str = "/borg/cargo"; +const VIRTUAL_REGISTRY: &str = "/borg/registry"; +const VIRTUAL_GIT: &str = "/borg/git"; +const MARKER: &str = "# managed by the borg drone; edits are overwritten"; + +/// The `--remap-path-prefix` flags for `slots` slot directories and the cargo +/// home. Later flags take precedence in rustc, so the more specific registry +/// and git prefixes follow the cargo home. +pub fn remap_flags(work_root: &Path, slots: usize, cargo_home: &Path) -> Vec { + let mut flags: Vec = (0..slots) + .map(|n| { + format!( + "--remap-path-prefix={}/slot-{n}={VIRTUAL_WORK}", + work_root.display() + ) + }) + .collect(); + let home = cargo_home.display(); + flags.push(format!("--remap-path-prefix={home}={VIRTUAL_CARGO}")); + flags.push(format!( + "--remap-path-prefix={home}/registry={VIRTUAL_REGISTRY}" + )); + flags.push(format!("--remap-path-prefix={home}/git={VIRTUAL_GIT}")); + flags +} + +/// `RUSTFLAGS` splits on whitespace, so paths containing it cannot be passed. +pub fn flags_are_passable(flags: &[String]) -> bool { + flags.iter().all(|f| !f.contains(char::is_whitespace)) +} + +/// Content of `$CARGO_HOME/config.toml`: the same flags as `[build] rustflags`. +/// Cargo ignores this when `RUSTFLAGS` is set, so it covers builds that run +/// without the lease environment. +pub fn config_toml(flags: &[String]) -> String { + let items: Vec = flags + .iter() + .map(|f| format!(" \"{}\",", f.replace('\\', "\\\\").replace('"', "\\\""))) + .collect(); + format!( + "{MARKER}\n[build]\nrustflags = [\n{}\n]\n", + items.join("\n") + ) +} + +/// Write `config.toml` unless the file exists and is not ours. Returns whether +/// the file was written. +pub fn write_config(cargo_home: &Path, flags: &[String]) -> io::Result { + let path = cargo_home.join("config.toml"); + let wanted = config_toml(flags); + match fs::read_to_string(&path) { + Ok(existing) if existing == wanted => return Ok(false), + Ok(existing) if !existing.starts_with(MARKER) => { + return Err(io::Error::other(format!( + "{} exists and is not managed by borg; left alone", + path.display() + ))); + } + Ok(_) => {} + Err(e) if e.kind() == io::ErrorKind::NotFound => {} + Err(e) => return Err(e), + } + fs::create_dir_all(cargo_home)?; + fs::write(&path, wanted)?; + Ok(true) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn flags_cover_every_slot_and_the_cargo_dirs_identically_for_all_leases() { + let flags = remap_flags( + Path::new("/var/lib/borg/work"), + 3, + Path::new("/cache/cargo"), + ); + assert_eq!( + flags, + [ + "--remap-path-prefix=/var/lib/borg/work/slot-0=/borg/work", + "--remap-path-prefix=/var/lib/borg/work/slot-1=/borg/work", + "--remap-path-prefix=/var/lib/borg/work/slot-2=/borg/work", + "--remap-path-prefix=/cache/cargo=/borg/cargo", + "--remap-path-prefix=/cache/cargo/registry=/borg/registry", + "--remap-path-prefix=/cache/cargo/git=/borg/git", + ] + ); + assert!(flags_are_passable(&flags)); + assert!(!flags_are_passable(&remap_flags( + Path::new("/with space"), + 1, + Path::new("/c") + ))); + } + + #[test] + fn config_toml_is_valid_toml_with_the_flags() { + let flags = remap_flags(Path::new("/w"), 2, Path::new("/c")); + let text = config_toml(&flags); + assert!(text.starts_with(MARKER) && text.contains("[build]\nrustflags = [")); + for f in &flags { + assert!(text.contains(&format!("\"{f}\"")), "{f}"); + } + } + + #[test] + fn write_is_idempotent_and_never_clobbers_a_foreign_config() { + let dir = tempfile::tempdir().unwrap(); + let flags = remap_flags(Path::new("/w"), 1, dir.path()); + assert!(write_config(dir.path(), &flags).unwrap()); + assert!(!write_config(dir.path(), &flags).unwrap()); + let more = remap_flags(Path::new("/w"), 2, dir.path()); + assert!( + write_config(dir.path(), &more).unwrap(), + "our own file is updated" + ); + + fs::write(dir.path().join("config.toml"), "[net]\nretry = 5\n").unwrap(); + assert!(write_config(dir.path(), &flags).is_err()); + assert_eq!( + fs::read_to_string(dir.path().join("config.toml")).unwrap(), + "[net]\nretry = 5\n" + ); + } +} diff --git a/src/drone/caches/sccache.rs b/src/drone/caches/sccache.rs new file mode 100644 index 0000000..1fbe2c3 --- /dev/null +++ b/src/drone/caches/sccache.rs @@ -0,0 +1,110 @@ +//! Optional host-level sccache server, for hit statistics. +//! +//! sccache keeps its hit and miss counters in the server's memory only; a +//! stopped server leaves nothing on disk to read (a new server starts at zero). +//! Because a job's own server dies with its lease, the drone can only report +//! hit rates by running one long-lived server per host that all leases share. +//! That is opt-in (`caches.sccache.server`): the server's configuration then +//! applies to every job, and a job whose sccache version differs from this +//! binary's will restart it. + +use crate::drone::exec::Exec; + +/// Environment that configures (and addresses) the host server. +pub fn server_env(dir: &str, size: &str, port: u16) -> Vec<(String, String)> { + vec![ + ("SCCACHE_DIR".to_owned(), dir.to_owned()), + ("SCCACHE_CACHE_SIZE".to_owned(), size.to_owned()), + ("SCCACHE_SERVER_PORT".to_owned(), port.to_string()), + // The server must outlive any single job. + ("SCCACHE_IDLE_TIMEOUT".to_owned(), "0".to_owned()), + ] +} + +/// Total cache hits and misses from `sccache --show-stats --stats-format json`. +pub fn parse_stats_json(json: &str) -> Option<(u64, u64)> { + let value: serde_json::Value = serde_json::from_str(json).ok()?; + let sum = |key: &str| -> Option { + let counts = value.get("stats")?.get(key)?.get("counts")?.as_object()?; + Some(counts.values().filter_map(serde_json::Value::as_u64).sum()) + }; + Some((sum("cache_hits")?, sum("cache_misses")?)) +} + +fn with_env(env: &[(String, String)], bin: &str, args: &[&str]) -> (String, Vec) { + let mut argv: Vec = env.iter().map(|(k, v)| format!("{k}={v}")).collect(); + argv.push(bin.to_owned()); + argv.extend(args.iter().map(|a| (*a).to_owned())); + ("env".to_owned(), argv) +} + +fn run(exec: &dyn Exec, env: &[(String, String)], bin: &str, args: &[&str]) -> Option { + let (program, argv) = with_env(env, bin, args); + let argv: Vec<&str> = argv.iter().map(String::as_str).collect(); + exec.run(&program, &argv) + .ok() + .filter(|o| o.success) + .map(|o| o.stdout) +} + +/// Start the host server. Succeeds when one is already running. +pub fn start_server(exec: &dyn Exec, bin: &str, env: &[(String, String)]) -> bool { + run(exec, env, bin, &["--start-server"]).is_some() +} + +/// Cumulative `(hits, misses)` of the host server, if sccache answers. +pub fn show_stats(exec: &dyn Exec, bin: &str, env: &[(String, String)]) -> Option<(u64, u64)> { + parse_stats_json(&run( + exec, + env, + bin, + &["--show-stats", "--stats-format", "json"], + )?) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::drone::cleanup::testing::FakeExec; + + // Shape captured from sccache 0.16.0. + const SAMPLE: &str = r#"{"stats":{"compile_requests":8,"requests_executed":3, + "cache_hits":{"counts":{"Rust":2,"C/C++":1},"adv_counts":{"rust":2}}, + "cache_misses":{"counts":{"Rust":4},"adv_counts":{"rust":4}}},"cache_size":9}"#; + + #[test] + fn hits_and_misses_are_summed_over_languages() { + assert_eq!(parse_stats_json(SAMPLE), Some((3, 4))); + let empty = r#"{"stats":{"cache_hits":{"counts":{}},"cache_misses":{"counts":{}}}}"#; + assert_eq!(parse_stats_json(empty), Some((0, 0))); + for bad in ["", "{}", "[]", r#"{"stats":{}}"#, "not json"] { + assert_eq!(parse_stats_json(bad), None, "{bad}"); + } + } + + #[test] + fn commands_run_through_env_with_the_server_configuration() { + let x = FakeExec::default(); + let env = server_env("/c/sccache", "50G", 4999); + x.on( + "env SCCACHE_DIR=/c/sccache SCCACHE_CACHE_SIZE=50G SCCACHE_SERVER_PORT=4999 SCCACHE_IDLE_TIMEOUT=0 sccache --show-stats --stats-format json", + SAMPLE, + ); + assert!(start_server(&x, "sccache", &env)); + assert_eq!(show_stats(&x, "sccache", &env), Some((3, 4))); + assert!(x.calls()[0].ends_with("sccache --start-server")); + } + + #[test] + fn a_missing_or_failing_sccache_is_just_no_stats() { + let x = FakeExec::default(); + let env = server_env("/c", "1G", 1); + x.fail("env SCCACHE_DIR=/c SCCACHE_CACHE_SIZE=1G SCCACHE_SERVER_PORT=1 SCCACHE_IDLE_TIMEOUT=0 sccache --start-server"); + assert!(!start_server(&x, "sccache", &env)); + assert_eq!( + show_stats(&x, "sccache", &env), + None, + "empty stdout is not a stats document" + ); + } +} diff --git a/src/drone/config.rs b/src/drone/config.rs index 576e275..64b1a47 100644 --- a/src/drone/config.rs +++ b/src/drone/config.rs @@ -41,6 +41,65 @@ impl Default for CacheEntry { } } +/// A host-level sccache server, so hit rates can be reported. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SccacheServer { + /// Port the server listens on; exported to jobs as `SCCACHE_SERVER_PORT`. + pub port: u16, + /// The sccache binary the drone runs (default `sccache`, from its PATH). + #[serde(default = "default_sccache_bin")] + pub bin: String, +} + +fn default_sccache_bin() -> String { + "sccache".to_owned() +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct SccacheCache { + pub enabled: bool, + /// `SCCACHE_CACHE_SIZE`; default 25% of the free disk, at most 100G. + pub max_size: Option, + /// Run one long-lived sccache server per host and report its hit/miss + /// counts. Off by default: without it each job's own server dies with the + /// lease and its counters are lost. + pub server: Option, +} + +impl Default for SccacheCache { + fn default() -> Self { + Self { + enabled: true, + max_size: None, + server: None, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct CargoCache { + pub enabled: bool, + pub max_size: Option, + /// Export `BORG_RUSTFLAGS` and `RUSTFLAGS` (path remaps that make compile + /// output independent of the slot) and write them to `$CARGO_HOME/config.toml`. + /// Cargo ignores `build.rustflags` and `target.*.rustflags` from config files + /// whenever `RUSTFLAGS` is set, so turn this off for repositories that rely on them. + pub export_rustflags: bool, +} + +impl Default for CargoCache { + fn default() -> Self { + Self { + enabled: true, + max_size: None, + export_rustflags: true, + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct BazelCache { @@ -94,9 +153,9 @@ impl Default for GitCaches { #[serde(default, deny_unknown_fields)] pub struct CachesConfig { pub git: GitCaches, - pub sccache: CacheEntry, + pub sccache: SccacheCache, /// `CARGO_HOME` registry and git data. `bin` is never shared between leases. - pub cargo: CacheEntry, + pub cargo: CargoCache, pub bazel: BazelCache, pub pnpm: CacheEntry, pub go: CacheEntry, @@ -111,8 +170,8 @@ impl Default for CachesConfig { fn default() -> Self { Self { git: GitCaches::default(), - sccache: CacheEntry::default(), - cargo: CacheEntry::default(), + sccache: SccacheCache::default(), + cargo: CargoCache::default(), bazel: BazelCache::default(), pnpm: CacheEntry::default(), go: CacheEntry::default(), @@ -164,11 +223,21 @@ pub struct DroneConfig { pub max_leases: Option, #[serde(default = "default_reserve")] pub reserve_mem_mib: u64, + /// Number of slot directories (`/slot-`) leases run in. A + /// lease takes the lowest free slot, so its path never contains the lease + /// id and a reused slot has an identical path. Also the most leases that + /// run at once. Default: `max_leases`, else 16. + #[serde(default)] + pub slots: Option, /// Host-level caches the drone manages under `cache_root`. #[serde(default)] pub caches: CachesConfig, } +/// Slot directories when neither `slots` nor `max_leases` is set. +const DEFAULT_SLOTS: u32 = 16; +const MAX_SLOTS: u32 = 256; + fn default_docker() -> String { "docker".to_owned() } @@ -225,6 +294,14 @@ impl DroneConfig { fn validate_caches(&self) -> Result<(), DroneConfigError> { let bad = |m: String| Err(DroneConfigError::Invalid(m)); let c = &self.caches; + if self.slots == Some(0) || self.slots.is_some_and(|s| s > MAX_SLOTS) { + return bad(format!("slots must be between 1 and {MAX_SLOTS}")); + } + if let (Some(leases), Some(slots)) = (self.max_leases, self.slots) + && leases > slots + { + return bad("max_leases must not exceed slots".to_owned()); + } if c.git.interval_secs == 0 || c.maintenance_interval_secs == 0 { return bad("cache intervals must be positive".to_owned()); } @@ -279,6 +356,12 @@ impl DroneConfig { Ok(()) } + /// How many slot directories exist (and so how many leases can run at once). + pub fn slot_count(&self) -> usize { + usize::try_from(self.slots.or(self.max_leases).unwrap_or(DEFAULT_SLOTS)) + .unwrap_or(usize::MAX) + } + /// The caps advertised to the queen with every heartbeat. pub fn host_caps(&self) -> HostCaps { HostCaps { @@ -384,6 +467,48 @@ mod tests { } } + #[test] + fn slots_default_to_max_leases_then_sixteen_and_are_validated() { + let with = |extra: &str| BASE.replace('}', &format!(",{extra}}}")); + assert_eq!(DroneConfig::from_json(BASE).unwrap().slot_count(), 16); + assert_eq!( + DroneConfig::from_json(&with(r#""max_leases":4"#)) + .unwrap() + .slot_count(), + 4 + ); + assert_eq!( + DroneConfig::from_json(&with(r#""max_leases":4,"slots":8"#)) + .unwrap() + .slot_count(), + 8 + ); + for bad in [ + r#""slots":0"#, + r#""slots":257"#, + r#""max_leases":9,"slots":8"#, + ] { + assert!(DroneConfig::from_json(&with(bad)).is_err(), "{bad}"); + } + } + + #[test] + fn sccache_server_and_cargo_rustflags_are_configurable() { + let with = |extra: &str| BASE.replace('}', &format!(",{extra}}}")); + let c = DroneConfig::from_json(&with( + r#""caches":{"sccache":{"server":{"port":4999}},"cargo":{"export_rustflags":false}}"#, + )) + .unwrap() + .caches; + assert_eq!( + c.sccache.server.as_ref().map(|s| (s.port, s.bin.as_str())), + Some((4999, "sccache")) + ); + assert!(!c.cargo.export_rustflags); + let d = DroneConfig::from_json(BASE).unwrap().caches; + assert!(d.sccache.server.is_none() && d.cargo.export_rustflags); + } + #[test] fn caches_may_not_overlap_the_wiped_work_root() { let nested = BASE.replace(r#""cache_root":"/c""#, r#""cache_root":"/w/cache""#); diff --git a/src/drone/install.rs b/src/drone/install.rs index 6a470bd..fa1564d 100644 --- a/src/drone/install.rs +++ b/src/drone/install.rs @@ -531,6 +531,7 @@ impl Installer<'_> { max_mem_mib: o.max_mem_mib, max_leases: o.max_leases, reserve_mem_mib: crate::allocator::HostCaps::default().reserve_mem_mib, + slots: None, caches: crate::drone::config::CachesConfig { git: crate::drone::config::GitCaches { mirrors: o.git_mirrors.clone(), diff --git a/src/drone/mod.rs b/src/drone/mod.rs index d0de4c2..d74bc2f 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -13,6 +13,7 @@ pub mod exec; pub mod install; pub mod pins; pub mod runner; +pub mod slots; use std::collections::HashSet; use std::path::Path; @@ -29,6 +30,7 @@ use client::{ClientError, QueenClient}; use config::DroneConfig; use exec::Exec; use runner::{RunSpec, run_lease}; +use slots::{SlotGuard, Slots, slot_dir, slot_name}; /// Attempts to deliver a lease result before giving up (the lease then /// expires on the queen's side). @@ -59,6 +61,7 @@ pub struct Drone { abort: tokio::sync::watch::Sender, proc_root: Arc, caches: Arc, + slots: Arc, } impl Drone { @@ -69,6 +72,7 @@ impl Drone { exec: Arc, ) -> Self { let (abort, _) = tokio::sync::watch::channel(false); + let slot_count = cfg.slot_count(); let cfg = Arc::new(cfg); let caches = Arc::new(Caches::new( Arc::clone(&cfg), @@ -81,12 +85,18 @@ impl Drone { host, exec, caches, + slots: Slots::new(slot_count), held: Held::default(), abort, proc_root: Arc::from(Path::new("/proc")), } } + /// `(slot, lease id)` for every lease currently holding a slot. + pub fn slot_assignments(&self) -> Vec<(usize, String)> { + self.slots.assignments() + } + fn held_ids(&self) -> Vec { let mut ids: Vec = self .held @@ -150,6 +160,17 @@ impl Drone { if !accepting || !self.admit(&grant.lease_id) { continue; } + let held = HeldGuard { + held: Arc::clone(&self.held), + id: grant.lease_id.clone(), + }; + let Some(slot) = self.slots.acquire(&grant.lease_id) else { + tracing::warn!( + "no free slot for lease {}; leaving it unstarted", + grant.lease_id + ); + continue; + }; let ctx = LeaseContext { cfg: Arc::clone(&self.cfg), client: Arc::clone(&self.client), @@ -157,10 +178,8 @@ impl Drone { abort: self.abort.subscribe(), proc_root: Arc::clone(&self.proc_root), caches: Arc::clone(&self.caches), - _held: HeldGuard { - held: Arc::clone(&self.held), - id: grant.lease_id.clone(), - }, + slot, + _held: held, }; tasks.spawn(handle_lease(ctx, grant)); started += 1; @@ -181,6 +200,13 @@ impl Drone { if held.contains(lease_id) { return false; } + if held.len() >= self.cfg.slot_count() { + tracing::warn!( + "all {} slots are busy; leaving lease {lease_id} unstarted", + self.cfg.slot_count() + ); + return false; + } if self .cfg .max_leases @@ -247,6 +273,8 @@ struct LeaseContext { abort: tokio::sync::watch::Receiver, proc_root: Arc, caches: Arc, + /// Dropped before `_held`, so the slot is free whenever the id is. + slot: SlotGuard, _held: HeldGuard, } @@ -256,13 +284,15 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { let lease_guard = ctx.caches.begin_lease().await; // Refresh stale git mirrors and learn the shared cache environment. Cache // trouble never blocks a job; it just runs without the speedup. - let lease_dir = ctx.cfg.work_root.join(&grant.lease_id); - let extra_env = ctx.caches.prepare_lease(&lease_dir).await; + let slot = ctx.slot.index(); + let lease_dir = slot_dir(&ctx.cfg.work_root, slot); + let mut extra_env = ctx.caches.prepare_lease(&lease_dir).await; + extra_env.push(("BORG_SLOT".to_owned(), slot.to_string())); let spec = RunSpec { lease_id: &grant.lease_id, jit_config: &grant.jit_config, runner_dir: &ctx.cfg.runner_dir, - work_root: &ctx.cfg.work_root, + lease_dir: &lease_dir, cache_root: &ctx.cfg.cache_root, extra_env: &extra_env, max_job: Duration::from_secs(ctx.cfg.max_job_secs), @@ -271,15 +301,16 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { let outcome = run_lease(&spec, ctx.exec.as_ref(), ctx.abort.clone()).await; // Runs whether the job passed, failed, timed out, was aborted or never started. - let (work_root, id, docker, exec) = ( + let (work_root, id, slot_name, docker, exec) = ( ctx.cfg.work_root.clone(), grant.lease_id.clone(), + slot_name(slot), ctx.cfg.docker_bin.clone(), Arc::clone(&ctx.exec), ); let cleaned = tokio::task::spawn_blocking(move || { let docker = cleanup_docker(exec.as_ref(), &docker, Scope::Lease(&id)); - (remove_lease_dir(&work_root, &id), docker) + (remove_lease_dir(&work_root, &slot_name), docker) }) .await; match cleaned { diff --git a/src/drone/runner.rs b/src/drone/runner.rs index edced8f..cac7a5e 100644 --- a/src/drone/runner.rs +++ b/src/drone/runner.rs @@ -1,6 +1,6 @@ //! Run one lease as an ephemeral GitHub Actions runner in a private directory. //! -//! The pinned runner install is cloned into `/` (hard +//! The pinned runner install is cloned into the lease's slot directory (hard //! links where the filesystem allows, copies otherwise), `run.sh` is started in //! its own process group with the JIT config in the environment (never argv), //! and the group is killed on timeout or abort. This module never cleans up; @@ -27,7 +27,9 @@ pub struct RunSpec<'a> { pub lease_id: &'a str, pub jit_config: &'a Secret, pub runner_dir: &'a Path, - pub work_root: &'a Path, + /// The lease's private directory, `/slot-`; it never + /// contains the lease id, so a reused slot has an identical path. + pub lease_dir: &'a Path, pub cache_root: &'a Path, /// Extra job environment (cache directories, git mirrors). pub extra_env: &'a [(String, String)], @@ -142,11 +144,17 @@ pub async fn run_lease( "unsafe lease id", )); } - let lease_dir: PathBuf = spec.work_root.join(spec.lease_id); + let lease_dir: PathBuf = spec.lease_dir.to_owned(); let (src, dst) = (spec.runner_dir.to_owned(), lease_dir.clone()); - tokio::task::spawn_blocking(move || clone_tree(&src, &dst)) - .await - .map_err(io::Error::other)??; + tokio::task::spawn_blocking(move || { + // A slot is reused, so make sure nothing of an earlier lease remains. + if dst.symlink_metadata().is_ok() { + fs::remove_dir_all(&dst)?; + } + clone_tree(&src, &dst) + }) + .await + .map_err(io::Error::other)??; let prefix = docker_prefix(spec.lease_id); #[allow( @@ -249,14 +257,14 @@ mod tests { lease_id: "lease1", jit_config: &secret, runner_dir: &runner, - work_root: &work, + lease_dir: &work.join("slot-0"), cache_root: &dir.join("cache"), extra_env: &[], max_job, proc_root: Path::new("/proc"), }; let out = run_lease(&spec, &SystemExec, rx).await.unwrap(); - (out, work.join("lease1")) + (out, work.join("slot-0")) } #[tokio::test] @@ -299,7 +307,7 @@ mod tests { fs::read_to_string(lease_dir.join("cwd.txt")) .unwrap() .trim() - .ends_with("work/lease1") + .ends_with("work/slot-0") ); assert_eq!( fs::read_to_string(lease_dir.join("bin/marker")).unwrap(), @@ -357,7 +365,7 @@ mod tests { lease_id: "lease2", jit_config: &secret, runner_dir: &runner, - work_root: &work, + lease_dir: &work.join("slot-0"), cache_root: &cache, extra_env: &[], max_job: Duration::from_secs(60), @@ -383,7 +391,7 @@ mod tests { lease_id: "../escape", jit_config: &secret, runner_dir: &runner, - work_root: &dir.path().join("work"), + lease_dir: &dir.path().join("work/slot-0"), cache_root: dir.path(), extra_env: &[], max_job: Duration::from_secs(1), diff --git a/src/drone/slots.rs b/src/drone/slots.rs new file mode 100644 index 0000000..1d0e8c4 --- /dev/null +++ b/src/drone/slots.rs @@ -0,0 +1,112 @@ +//! Slot directories. A lease runs in `/slot-` and takes the +//! lowest free slot, so concurrent leases always get distinct slots, a slot is +//! reused once its lease has been cleaned up, and the path never contains the +//! lease id. Identical paths are what let compile caches hit when a slot is +//! reused. + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; + +pub struct Slots { + capacity: usize, + /// slot index -> lease id + held: Mutex>, +} + +/// A slot held by a lease; released on drop. +pub struct SlotGuard { + slots: Arc, + index: usize, +} + +impl SlotGuard { + pub fn index(&self) -> usize { + self.index + } +} + +impl Drop for SlotGuard { + fn drop(&mut self) { + if let Ok(mut held) = self.slots.held.lock() { + held.remove(&self.index); + } + } +} + +impl Slots { + pub fn new(capacity: usize) -> Arc { + Arc::new(Self { + capacity, + held: Mutex::new(BTreeMap::new()), + }) + } + + /// The lowest free slot, or `None` when all are taken. + pub fn acquire(self: &Arc, lease_id: &str) -> Option { + let mut held = self.held.lock().ok()?; + let index = (0..self.capacity).find(|i| !held.contains_key(i))?; + held.insert(index, lease_id.to_owned()); + tracing::info!("lease {lease_id} -> slot {index}"); + Some(SlotGuard { + slots: Arc::clone(self), + index, + }) + } + + /// `(slot, lease id)` for every held slot. + pub fn assignments(&self) -> Vec<(usize, String)> { + self.held + .lock() + .map(|h| h.iter().map(|(i, l)| (*i, l.clone())).collect()) + .unwrap_or_default() + } +} + +/// `/slot-` +pub fn slot_dir(work_root: &Path, index: usize) -> PathBuf { + work_root.join(slot_name(index)) +} + +pub fn slot_name(index: usize) -> String { + format!("slot-{index}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn concurrent_leases_get_distinct_lowest_free_slots() { + let slots = Slots::new(3); + let a = slots.acquire("a").unwrap(); + let b = slots.acquire("b").unwrap(); + let c = slots.acquire("c").unwrap(); + assert_eq!((a.index(), b.index(), c.index()), (0, 1, 2)); + assert!(slots.acquire("d").is_none(), "no slot beyond capacity"); + assert_eq!( + slots.assignments(), + [(0, "a".into()), (1, "b".into()), (2, "c".into())] + ); + } + + #[test] + fn a_released_slot_is_reused_lowest_first() { + let slots = Slots::new(3); + let a = slots.acquire("a").unwrap(); + let _b = slots.acquire("b").unwrap(); + drop(a); + let c = slots.acquire("c").unwrap(); + assert_eq!(c.index(), 0, "the freed low slot, not slot 2"); + assert_eq!(slots.assignments().len(), 2); + } + + #[test] + fn paths_never_contain_the_lease_id() { + let slots = Slots::new(2); + let g = slots.acquire("0123456789ab").unwrap(); + let dir = slot_dir(Path::new("/var/lib/borg/work"), g.index()); + assert_eq!(dir, Path::new("/var/lib/borg/work/slot-0")); + assert!(!dir.to_string_lossy().contains("0123456789ab")); + } +} diff --git a/tests/drone.rs b/tests/drone.rs index cd1c55f..c64e01c 100644 --- a/tests/drone.rs +++ b/tests/drone.rs @@ -172,8 +172,8 @@ async fn runs_a_granted_lease_cleans_up_and_reports_the_result() { ); let lease_id = name.rsplit('-').next().unwrap().to_owned(); assert!( - !rig.path("work").join(&lease_id).exists(), - "work dir removed after the job" + !rig.path("work/slot-0").exists(), + "slot dir removed after the job" ); assert_eq!( fs::read_to_string(rig.path("cache/sccache/entry")).unwrap(), @@ -431,9 +431,11 @@ async fn jobs_share_cache_dirs_and_a_git_mirror_that_survive_cleanup() { ) ); let installs = read("install-"); - assert_ne!( - installs[0], installs[1], - "cargo install roots are per lease" + let slot_install = format!("{}/slot-0/cargo-install", rig.path("work").display()); + assert_eq!( + installs, + [slot_install.clone(), slot_install], + "cargo install roots live in the lease's slot, never in the shared CARGO_HOME" ); let bazelrc = fs::read_to_string(rig.path("bazelrc")).unwrap(); assert!(bazelrc.contains(&format!("--disk_cache={}/bazel/disk", cache.display()))); @@ -505,3 +507,73 @@ async fn the_mirror_timer_picks_up_new_commits() { ]); assert!(clone.join(".git/objects/info/alternates").exists()); } + +#[tokio::test] +async fn leases_get_distinct_slots_while_concurrent_and_the_same_path_when_a_slot_is_reused() { + let rig = Rig::new().await; + rig.install_runner( + "pwd > $OUT/pwd-$BORG_LEASE_ID.txt\necho $BORG_SLOT > $OUT/slot-$BORG_LEASE_ID.txt\nsleep 1", + ); + let drone = rig.drone(json!({})); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + + rig.h.queue_job(1, "borg-ci-light", 1); + rig.h.queue_job(2, "borg-ci-light", 1); + rig.h.tick().await; + assert_eq!(drone.heartbeat_once(true, &mut tasks).await.unwrap(), 2); + let during = drone.slot_assignments(); + assert_eq!( + during.iter().map(|(s, _)| *s).collect::>(), + [0, 1], + "distinct slots while concurrent" + ); + drain(&mut tasks).await; + assert!( + drone.slot_assignments().is_empty(), + "slots are released after cleanup" + ); + + let read_all = |prefix: &str| -> Vec<(String, String)> { + let mut v: Vec<(String, String)> = fs::read_dir(rig.path("out")) + .unwrap() + .filter_map(|e| e.ok()) + .filter(|e| e.file_name().to_string_lossy().starts_with(prefix)) + .map(|e| { + let name = e.file_name().to_string_lossy().into_owned(); + ( + name, + fs::read_to_string(e.path()).unwrap().trim().to_owned(), + ) + }) + .collect(); + v.sort(); + v + }; + let first_pwds: Vec = read_all("pwd-").into_iter().map(|(_, p)| p).collect(); + let work = rig.path("work"); + let mut want = vec![ + work.join("slot-0").display().to_string(), + work.join("slot-1").display().to_string(), + ]; + let mut got = first_pwds.clone(); + got.sort(); + want.sort(); + assert_eq!(got, want); + + // A later lease alone takes slot 0 again, at exactly the same path. + rig.h.state().jobs.clear(); + rig.h.queue_job(3, "borg-ci-light", 1); + rig.h.tick().await; + assert_eq!(drone.heartbeat_once(true, &mut tasks).await.unwrap(), 1); + drain(&mut tasks).await; + let all = read_all("pwd-"); + assert_eq!(all.len(), 3); + let slot_files = read_all("slot-"); + assert_eq!( + slot_files.iter().filter(|(_, s)| s == "0").count(), + 2, + "slot 0 was used twice" + ); + assert_eq!(fs::read_dir(rig.path("work")).unwrap().count(), 0); +} diff --git a/tests/sccache_slots.rs b/tests/sccache_slots.rs new file mode 100644 index 0000000..b62180a --- /dev/null +++ b/tests/sccache_slots.rs @@ -0,0 +1,123 @@ +//! Compiles one small crate through sccache in slot directories, to show what +//! slot-stable work directories do for the compile cache. Skipped when `cargo` +//! or `sccache` is not installed. + +use std::fs; +use std::path::Path; + +use borg::drone::caches::{rustflags, sccache}; +use borg::drone::exec::{Exec, SystemExec}; + +fn available(program: &str) -> bool { + SystemExec + .run(program, &["--version"]) + .is_ok_and(|o| o.success) +} + +fn free_port() -> u16 { + std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port() +} + +/// Run `program args` under `env`, through `env(1)` since the executor takes no environment. +fn run_with(env: &[(String, String)], program: &str, args: &[&str]) -> bool { + let mut argv: Vec = env.iter().map(|(k, v)| format!("{k}={v}")).collect(); + argv.push(program.to_owned()); + argv.extend(args.iter().map(|a| (*a).to_owned())); + let argv: Vec<&str> = argv.iter().map(String::as_str).collect(); + SystemExec.run("env", &argv).is_ok_and(|o| o.success) +} + +fn write_crate(dir: &Path) { + fs::create_dir_all(dir.join("src")).unwrap(); + fs::write( + dir.join("Cargo.toml"), + "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n", + ) + .unwrap(); + fs::write( + dir.join("src/lib.rs"), + "pub fn add(a: u64, b: u64) -> u64 { a + b }\npub fn fail() { panic!(\"boom\") }\n", + ) + .unwrap(); +} + +#[test] +fn workspace_crates_hit_the_compile_cache_when_a_slot_is_reused() { + if !(available("cargo") && available("sccache")) { + println!("skipping: cargo and sccache are required"); + return; + } + let tmp = tempfile::tempdir().unwrap(); + let work = tmp.path().join("work"); + let cargo_home = tmp.path().join("cargo"); + let port = free_port(); + let dir = tmp.path().join("sccache").to_string_lossy().into_owned(); + let flags = rustflags::remap_flags(&work, 2, &cargo_home).join(" "); + + let mut env = sccache::server_env(&dir, "1G", port); + env.extend([ + ("RUSTC_WRAPPER".to_owned(), "sccache".to_owned()), + ("CARGO_INCREMENTAL".to_owned(), "0".to_owned()), + ( + "CARGO_HOME".to_owned(), + cargo_home.to_string_lossy().into_owned(), + ), + ("RUSTFLAGS".to_owned(), flags), + ]); + assert!( + sccache::start_server(&SystemExec, "sccache", &env), + "sccache server starts" + ); + + let build = |slot: usize| { + let crate_dir = work.join(format!("slot-{slot}")).join("demo"); + write_crate(&crate_dir); + let manifest = crate_dir.join("Cargo.toml"); + assert!(run_with( + &env, + "cargo", + &[ + "build", + "--lib", + "--quiet", + "--manifest-path", + manifest.to_str().unwrap() + ] + )); + sccache::show_stats(&SystemExec, "sccache", &env).expect("sccache answers") + }; + + let first = build(0); + fs::remove_dir_all(work.join("slot-0")).unwrap(); // what drone cleanup does between leases + let reused = build(0); + let other = build(1); + println!( + "slot-0 first compile: hits={} misses={}", + first.0, first.1 + ); + println!( + "slot-0 reused (cleaned): hits={} misses={}", + reused.0, reused.1 + ); + println!( + "slot-1 (other slot): hits={} misses={}", + other.0, other.1 + ); + + assert_eq!(first.0, 0, "a cold cache cannot hit"); + assert!(first.1 >= 1); + assert!( + reused.0 > first.0, + "recompiling in a reused slot hits: {first:?} -> {reused:?}" + ); + // The other slot is deliberately not asserted: sccache hashes the compile's + // working directory and every CARGO_* variable (CARGO_MANIFEST_DIR), which + // differ between slots whatever rustc flags are passed. The numbers above + // record what happens. + + let _ = run_with(&env, "sccache", &["--stop-server"]); +} From cf8a4c5c243ca668e397b29f635c38369fe17737 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 10 Oct 2026 20:40:07 +0000 Subject: [PATCH 04/27] chore: project borg from Mono d85864165c24 --- .repository-projection.json | 6 +- README.md | 7 ++ src/api.rs | 142 ++++++++++++++++++++++++++++++++---- src/config.rs | 4 + src/drone/mod.rs | 1 + src/main.rs | 38 ++++++++++ src/queen.rs | 2 +- src/store.rs | 88 +++++++++++++++++++--- tests/queen.rs | 78 ++++++++++++++++++++ tests/support/mod.rs | 22 ++++++ 10 files changed, 361 insertions(+), 27 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 4df2a52..28685de 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "95eee04b812703f5b31b593f8d10e7d55d4a3a0f", + "sourceSha": "d85864165c244cbf6502331236b6a5210b69e5d8", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "101caa77b052562b6862c9689347e746a301735b", + "priorProjectedBase": "f988d4adbc65ec8f5552020fce0a802ad3658466", "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "955724be218525b4b86854d0976c551258e5f616da68af35eeea95dbe293bfe5", + "contentDigest": "a221daff41b64481a81d0fb48bb5146630416ec0630abd6b94ef54db81a87420", "publicationEligible": true } diff --git a/README.md b/README.md index f4c56ed..afe0823 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,13 @@ optional): } ``` +Operators can stop new leases to one drone with `borg queen quarantine +--config queen.json --reason "" ` and lift it with `borg queen +release --config queen.json `. Set `rollout_target_version` in the queen +config to have `/metrics` and `/status` report how many drones run that borg +version (drones report theirs in every heartbeat; the queen never forces an +update). + The GitHub App needs organization self-hosted runners: read and write, plus read access to Actions on the listed repositories. Set `BORG_APP_PRIVATE_KEY` instead of `private_key_path` to pass the PEM in the diff --git a/src/api.rs b/src/api.rs index b7b6ea7..b9f4cd9 100644 --- a/src/api.rs +++ b/src/api.rs @@ -41,6 +41,15 @@ pub struct HeartbeatRequest { /// Size and garbage-collection state of the drone's host caches. #[serde(default)] pub cache_stats: Vec, + /// The borg version the drone runs (reported for rollout tracking). + #[serde(default)] + pub version: Option, +} + +/// A drone-reported version is shown in metric labels, so accept only short +/// printable ASCII. +fn sanitize_version(v: Option<&str>) -> Option<&str> { + v.filter(|v| !v.is_empty() && v.len() <= 64 && v.bytes().all(|b| b.is_ascii_graphic())) } /// One host cache as measured by its drone. @@ -158,6 +167,9 @@ async fn heartbeat( now, ) .await?; + store + .set_version(&drone, sanitize_version(req.version.as_deref())) + .await?; store.acknowledge(&drone, &req.active_leases).await?; let leases = store .undelivered(&drone) @@ -226,6 +238,21 @@ pub struct DroneStatus { pub mem_mib: Option, pub active_leases: u32, pub caches: Vec, + pub quarantined: bool, + pub quarantine_reason: Option, + pub version: Option, +} + +/// How many enrolled drones run the configured target version. +#[derive(Debug, Default, Serialize)] +pub struct RolloutStatus { + pub target_version: Option, + /// Drones reporting exactly the target. + pub current: u32, + /// Drones reporting another version. + pub outdated: u32, + /// Drones that report no version. + pub unknown: u32, } #[derive(Debug, Serialize)] @@ -239,6 +266,7 @@ pub struct Status { pub unmet: Vec, /// Drones held back for starving demand. pub reserved_for: Vec, + pub rollout: RolloutStatus, } pub async fn build_status(queen: &Queen, now: i64) -> Result { @@ -254,7 +282,7 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result let fresh = age.is_some_and(|a| a <= cfg.drone_stale_secs); if let (true, Some(cap)) = (fresh, &d.capacity) { let mut s = DroneState::from_capacity(d.name.clone(), cap, &d.caps); - s.accepting = d.accepting; + s.accepting = d.schedulable(); states.push(s); } drones.push(DroneStatus { @@ -269,6 +297,9 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result active_leases: u32::try_from(active.iter().filter(|l| l.drone == d.name).count()) .unwrap_or(u32::MAX), caches: d.cache_stats.clone(), + quarantined: d.quarantined, + quarantine_reason: d.quarantine_reason.clone(), + version: d.version.clone(), }); } @@ -302,6 +333,17 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result }, ); } + let mut rollout = RolloutStatus { + target_version: cfg.rollout_target_version.clone(), + ..RolloutStatus::default() + }; + for d in &drones { + match (&d.version, &cfg.rollout_target_version) { + (None, _) => rollout.unknown += 1, + (Some(v), Some(t)) if v == t => rollout.current += 1, + (Some(_), _) => rollout.outdated += 1, + } + } Ok(Status { now, last_tick_at: snap.last_tick_at, @@ -311,6 +353,7 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result drones, unmet: snap.unmet, reserved_for: snap.reserved_for, + rollout, }) } @@ -399,24 +442,81 @@ pub fn render_metrics(status: &Status) -> String { "Active leases per class and state.", lease_rows, ); + let count = |f: &dyn Fn(&DroneStatus) -> bool| { + f64::from(u32::try_from(status.drones.iter().filter(|d| f(d)).count()).unwrap_or(0)) + }; gauge( "borg_drones", - "Enrolled drones by freshness.", + "Enrolled drones by state. fresh and stale split by heartbeat age; quarantined \ + counts operator-quarantined drones regardless of heartbeat.", vec![ + ("state=\"fresh\"".to_owned(), count(&|d| d.fresh)), + ("state=\"stale\"".to_owned(), count(&|d| !d.fresh)), ( - "state=\"fresh\"".to_owned(), - f64::from( - u32::try_from(status.drones.iter().filter(|d| d.fresh).count()).unwrap_or(0), - ), - ), - ( - "state=\"stale\"".to_owned(), - f64::from( - u32::try_from(status.drones.iter().filter(|d| !d.fresh).count()).unwrap_or(0), - ), + "state=\"quarantined\"".to_owned(), + count(&|d| d.quarantined), ), ], ); + gauge( + "borg_drone_quarantined", + "1 while the drone is quarantined and gets no new leases.", + status + .drones + .iter() + .map(|d| { + ( + format!("drone=\"{}\"", label_value(&d.name)), + f64::from(u8::from(d.quarantined)), + ) + }) + .collect(), + ); + gauge( + "borg_drone_info", + "The borg version each drone reported (always 1).", + status + .drones + .iter() + .filter_map(|d| { + d.version.as_ref().map(|v| { + ( + format!( + "drone=\"{}\",version=\"{}\"", + label_value(&d.name), + label_value(v) + ), + 1.0, + ) + }) + }) + .collect(), + ); + gauge( + "borg_rollout_target_info", + "The configured target borg version (always 1; absent when unset).", + status + .rollout + .target_version + .iter() + .map(|v| (format!("version=\"{}\"", label_value(v)), 1.0)) + .collect(), + ); + gauge( + "borg_rollout_drones", + "Drones by rollout state against the target version (absent when no target is set).", + if status.rollout.target_version.is_some() { + [ + ("current", status.rollout.current), + ("outdated", status.rollout.outdated), + ("unknown", status.rollout.unknown), + ] + .map(|(state, n)| (format!("state=\"{state}\""), f64::from(n))) + .to_vec() + } else { + Vec::new() + }, + ); gauge( "borg_reservations", "Drones held back for starving demand.", @@ -440,7 +540,16 @@ pub fn render_metrics(status: &Status) -> String { .iter() .flat_map(|d| { d.caches.iter().filter_map(move |c| { - f(c).map(|v| (format!("drone=\"{}\",cache=\"{}\"", d.name, c.name), v)) + f(c).map(|v| { + ( + format!( + "drone=\"{}\",cache=\"{}\"", + label_value(&d.name), + label_value(&c.name) + ), + v, + ) + }) }) }) .collect() @@ -488,6 +597,13 @@ pub fn render_metrics(status: &Status) -> String { out } +/// Escape a Prometheus label value. +fn label_value(v: &str) -> String { + v.replace('\\', "\\\\") + .replace('"', "\\\"") + .replace('\n', "\\n") +} + fn scope_name(s: Scope) -> &'static str { match s { Scope::Main => "main", diff --git a/src/config.rs b/src/config.rs index fa0d270..d0a7542 100644 --- a/src/config.rs +++ b/src/config.rs @@ -80,6 +80,10 @@ pub struct QueenConfig { pub classes: Classes, #[serde(default = "default_warm_floors")] pub warm_floors: Vec, + /// The borg version drones should run. Only reported (`/metrics`, + /// `/status`): drones update themselves, the queen never forces it. + #[serde(default)] + pub rollout_target_version: Option, } fn default_listen() -> SocketAddr { diff --git a/src/drone/mod.rs b/src/drone/mod.rs index d74bc2f..d04df49 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -153,6 +153,7 @@ impl Drone { accepting, active_leases: self.held_ids(), cache_stats: self.caches.stats(), + version: Some(env!("CARGO_PKG_VERSION").to_owned()), }; let resp = self.client.heartbeat(&req).await?; let mut started = 0; diff --git a/src/main.rs b/src/main.rs index 6c7a03d..e2a7f6c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -61,6 +61,23 @@ enum QueenCommand { /// Drone name: lowercase letters, digits and dashes. name: String, }, + /// Stop granting a drone new leases (running jobs finish). + Quarantine { + #[command(flatten)] + config: ConfigArg, + /// Why the drone is quarantined; shown in `/status`. + #[arg(long, default_value = "operator")] + reason: String, + /// Drone name. + name: String, + }, + /// Lift a quarantine. + Release { + #[command(flatten)] + config: ConfigArg, + /// Drone name. + name: String, + }, } #[derive(Args)] @@ -188,6 +205,19 @@ async fn enroll_drone(path: &Path, name: &str) -> Result<()> { Ok(()) } +async fn set_quarantine(path: &Path, name: &str, reason: Option<&str>) -> Result<()> { + let cfg = QueenConfig::load(path)?; + let store = open_store(&cfg).await?; + if !store.set_quarantine(name, reason).await? { + bail!("no enrolled drone named {name}"); + } + match reason { + Some(r) => eprintln!("Quarantined drone {name} ({r}); it gets no new leases."), + None => eprintln!("Released drone {name}."), + } + Ok(()) +} + async fn run_drone(path: &Path) -> Result<()> { let cfg = DroneConfig::load(path)?; let token = cfg.load_token()?; @@ -268,6 +298,14 @@ async fn main() -> Result<()> { match command { QueenCommand::Run(c) => run_queen(&c.config).await, QueenCommand::Enroll { config, name } => enroll_drone(&config.config, &name).await, + QueenCommand::Quarantine { + config, + reason, + name, + } => set_quarantine(&config.config, &name, Some(&reason)).await, + QueenCommand::Release { config, name } => { + set_quarantine(&config.config, &name, None).await + } } } } diff --git a/src/queen.rs b/src/queen.rs index 929127a..1e36edf 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -166,7 +166,7 @@ impl Queen { .is_some_and(|t| now_ts - t <= self.cfg.drone_stale_secs); if let (true, Some(cap)) = (fresh, &d.capacity) { let mut state = DroneState::from_capacity(d.name.clone(), cap, &d.caps); - state.accepting = d.accepting; + state.accepting = d.schedulable(); drones.push(state); capacities.insert(d.name.clone(), cap.clone()); } diff --git a/src/store.rs b/src/store.rs index 7d6124a..bca4707 100644 --- a/src/store.rs +++ b/src/store.rs @@ -27,7 +27,10 @@ CREATE TABLE IF NOT EXISTS drones ( capacity TEXT, caps TEXT, accepting INTEGER NOT NULL DEFAULT 1, - cache_stats TEXT + cache_stats TEXT, + quarantined INTEGER NOT NULL DEFAULT 0, + quarantine_reason TEXT, + version TEXT ); CREATE TABLE IF NOT EXISTS leases ( id TEXT PRIMARY KEY, @@ -103,6 +106,18 @@ pub struct DroneRow { pub caps: HostCaps, pub accepting: bool, pub cache_stats: Vec, + /// Set by an operator; a quarantined drone gets no new leases. + pub quarantined: bool, + pub quarantine_reason: Option, + /// The borg version the drone last reported, if it reports one. + pub version: Option, +} + +impl DroneRow { + /// Whether the allocator may hand this drone new leases. + pub fn schedulable(&self) -> bool { + self.accepting && !self.quarantined + } } #[derive(Debug, Clone)] @@ -173,14 +188,19 @@ impl Store { .connect_with(opts) .await?; sqlx::raw_sql(SCHEMA).execute(&pool).await?; - // Databases created before cache stats existed lack the column. - // sqlx-runtime-query: one-off additive migration, no data involved - if let Err(e) = sqlx::query("ALTER TABLE drones ADD COLUMN cache_stats TEXT") - .execute(&pool) - .await - && !e.to_string().contains("duplicate column") - { - return Err(e.into()); + // Databases created before these columns existed lack them. + for ddl in [ + "ALTER TABLE drones ADD COLUMN cache_stats TEXT", + "ALTER TABLE drones ADD COLUMN quarantined INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE drones ADD COLUMN quarantine_reason TEXT", + "ALTER TABLE drones ADD COLUMN version TEXT", + ] { + // sqlx-runtime-query: one-off additive migration, no data involved + if let Err(e) = sqlx::query(ddl).execute(&pool).await + && !e.to_string().contains("duplicate column") + { + return Err(e.into()); + } } Ok(Self { pool }) } @@ -239,7 +259,7 @@ impl Store { pub async fn drones(&self) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let rows = sqlx::query( - "SELECT name, last_heartbeat, capacity, caps, accepting, cache_stats FROM drones ORDER BY name", + "SELECT name, last_heartbeat, capacity, caps, accepting, cache_stats, quarantined, quarantine_reason, version FROM drones ORDER BY name", ) .fetch_all(&self.pool) .await?; @@ -263,11 +283,43 @@ impl Store { .map(from_text) .transpose()? .unwrap_or_default(), + quarantined: r.get::("quarantined") != 0, + quarantine_reason: r.get("quarantine_reason"), + version: r.get("version"), }) }) .collect() } + /// Quarantine (`Some(reason)`) or release (`None`) a drone. False if no + /// such drone is enrolled. + pub async fn set_quarantine( + &self, + name: &str, + reason: Option<&str>, + ) -> Result { + const SQL: &str = "UPDATE drones SET quarantined = ?, quarantine_reason = ? WHERE name = ?"; + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let r = sqlx::query(SQL) + .bind(i64::from(reason.is_some())) + .bind(reason) + .bind(name) + .execute(&self.pool) + .await?; + Ok(r.rows_affected() > 0) + } + + /// Record the borg version a drone reported in its heartbeat. + pub async fn set_version(&self, name: &str, version: Option<&str>) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE drones SET version = ? WHERE name = ?") + .bind(version) + .bind(name) + .execute(&self.pool) + .await?; + Ok(()) + } + pub async fn insert_lease(&self, l: &LeaseRow) -> Result<(), StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database sqlx::query( @@ -517,6 +569,22 @@ mod tests { ); } + #[tokio::test] + async fn quarantine_and_version_round_trip_and_survive_heartbeats() { + let s = store_with("d1").await; + assert!(!s.set_quarantine("nope", Some("x")).await.unwrap()); + assert!(s.set_quarantine("d1", Some("bad disk")).await.unwrap()); + s.set_version("d1", Some("1.2.3")).await.unwrap(); + let d = &s.drones().await.unwrap()[0]; + assert!(d.accepting && d.quarantined && !d.schedulable()); + assert_eq!(d.quarantine_reason.as_deref(), Some("bad disk")); + assert_eq!(d.version.as_deref(), Some("1.2.3")); + assert!(s.set_quarantine("d1", None).await.unwrap()); + let d = &s.drones().await.unwrap()[0]; + assert!(!d.quarantined && d.schedulable()); + assert_eq!(d.quarantine_reason, None); + } + #[tokio::test] async fn heartbeat_round_trips_capacity() { let s = store_with("d1").await; diff --git a/tests/queen.rs b/tests/queen.rs index 75ee7e4..a34d365 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -286,6 +286,84 @@ async fn metrics_and_status_report_the_queue() { assert!(status["classes"]["heavy"]["slots"].as_u64() == Some(1)); } +#[tokio::test] +async fn quarantined_drones_get_no_leases_and_are_counted() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + assert!( + h.queen + .store() + .set_quarantine("d1", Some("bad disk")) + .await + .unwrap() + ); + h.queue_job(1, "borg-ci-medium-pr", 30); + h.tick().await; + + assert!(h.leases(&token, &[]).await.is_empty()); + let metrics = h.metrics().await; + for line in [ + "borg_drones{state=\"fresh\"} 1", + "borg_drones{state=\"quarantined\"} 1", + "borg_drone_quarantined{drone=\"d1\"} 1", + "borg_class_slots{class=\"medium\"} 0", + "borg_oldest_queued_age_seconds{class=\"medium\",scope=\"pr\"} 30", + ] { + assert!(metrics.contains(line), "missing {line}\n{metrics}"); + } + let status = h.status().await; + assert_eq!(status["drones"][0]["quarantined"], true); + assert_eq!(status["drones"][0]["quarantine_reason"], "bad disk"); + + h.queen.store().set_quarantine("d1", None).await.unwrap(); + h.tick().await; + assert_eq!(h.leases(&token, &[]).await.len(), 1); + assert!( + h.metrics() + .await + .contains("borg_drones{state=\"quarantined\"} 0") + ); +} + +#[tokio::test] +async fn rollout_metrics_compare_reported_versions_with_the_target() { + let h = Harness::start(json!({"rollout_target_version": "2.0.0"})).await; + let (a, b, c) = ( + h.enroll("d1").await, + h.enroll("d2").await, + h.enroll("d3").await, + ); + h.heartbeat_version(&a, "2.0.0").await; + h.heartbeat_version(&b, "1.9.0\"x").await; + h.heartbeat(&c, &[]).await; + + let metrics = h.metrics().await; + for line in [ + "borg_rollout_target_info{version=\"2.0.0\"} 1", + "borg_rollout_drones{state=\"current\"} 1", + "borg_rollout_drones{state=\"outdated\"} 1", + "borg_rollout_drones{state=\"unknown\"} 1", + "borg_drone_info{drone=\"d1\",version=\"2.0.0\"} 1", + ] { + assert!(metrics.contains(line), "missing {line}\n{metrics}"); + } + // A quote in a drone-reported version must be escaped, not break the format. + assert!(metrics.contains("version=\"1.9.0\\\"x\""), "{metrics}"); + assert_eq!(h.status().await["rollout"]["current"], 1); +} + +#[tokio::test] +async fn no_rollout_series_without_a_target() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat_version(&token, "1.0.0").await; + let metrics = h.metrics().await; + assert!(!metrics.contains("borg_rollout_drones{"), "{metrics}"); + assert!(!metrics.contains("borg_rollout_target_info{"), "{metrics}"); + assert!(metrics.contains("borg_drone_info{drone=\"d1\",version=\"1.0.0\"} 1")); +} + #[tokio::test] async fn mint_failure_is_reported_and_leaves_no_lease() { let h = Harness::start(json!({})).await; diff --git a/tests/support/mod.rs b/tests/support/mod.rs index f67cb6e..be7188c 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -298,6 +298,28 @@ impl Harness { .unwrap() } + /// Heartbeat reporting `version`. + pub async fn heartbeat_version(&self, token: &str, version: &str) -> reqwest::Response { + self.http + .post(format!("{}/v1/heartbeat", self.url)) + .bearer_auth(token) + .json(&json!({"capacity": capacity(), "version": version})) + .send() + .await + .unwrap() + } + + pub async fn metrics(&self) -> String { + self.http + .get(format!("{}/metrics", self.url)) + .send() + .await + .unwrap() + .text() + .await + .unwrap() + } + pub async fn leases(&self, token: &str, active: &[&str]) -> Vec { let body: Value = self.heartbeat(token, active).await.json().await.unwrap(); body["leases"].as_array().cloned().unwrap_or_default() From df40113e356b4dd48d692f29e9bee66d098a9d4a Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 10 Oct 2026 21:31:43 +0000 Subject: [PATCH 05/27] chore: project borg from Mono ef5e513984ed --- .repository-projection.json | 6 +- Cargo.toml | 2 +- README.md | 36 ++ src/api.rs | 91 ++++- src/config.rs | 80 ++++ src/drone/mod.rs | 2 + src/janitor.rs | 300 ++++++++++++++ src/lib.rs | 3 + src/main.rs | 97 ++++- src/queen.rs | 29 +- src/schedule.rs | 294 ++++++++++++++ src/store.rs | 18 +- src/telemetry.rs | 769 ++++++++++++++++++++++++++++++++++++ tests/janitor.rs | 116 ++++++ tests/queen.rs | 158 ++++++++ tests/support/mod.rs | 103 ++++- 16 files changed, 2093 insertions(+), 11 deletions(-) create mode 100644 src/janitor.rs create mode 100644 src/schedule.rs create mode 100644 src/telemetry.rs create mode 100644 tests/janitor.rs diff --git a/.repository-projection.json b/.repository-projection.json index 28685de..ebc7a0d 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "d85864165c244cbf6502331236b6a5210b69e5d8", + "sourceSha": "ef5e513984edf1a1fbf8d2caeac0d809e6818ab0", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "f988d4adbc65ec8f5552020fce0a802ad3658466", + "priorProjectedBase": "5088bfcbe596a86566c86db9a446883b66215ad5", "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "a221daff41b64481a81d0fb48bb5146630416ec0630abd6b94ef54db81a87420", + "contentDigest": "fbb22932e292fa1fc25b21a331a41d441dafacc2cb8ecc66b11c1d9b238f0250", "publicationEligible": true } diff --git a/Cargo.toml b/Cargo.toml index d99b00f..8a868e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,7 +30,7 @@ serde_json = "1.0" sha2 = "0.10" sqlx = { version = "0.8", default-features = false, features = ["runtime-tokio-rustls", "sqlite"] } thiserror = "2.0" -time = { version = "0.3", features = ["formatting", "parsing"] } +time = { version = "0.3", features = ["formatting", "macros", "parsing"] } tokio = { version = "1.48", features = ["macros", "net", "rt-multi-thread", "process", "signal", "sync", "time"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } diff --git a/README.md b/README.md index afe0823..dff4939 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,42 @@ config to have `/metrics` and `/status` report how many drones run that borg version (drones report theirs in every heartbeat; the queen never forces an update). +`warm_floors` keep idle runners alive per class and scope; the optional +`warm_schedule` raises them in time-of-day windows (`{"class": "light", "scope": +"pr", "count": 4, "start": "08:00", "end": "20:00", "days": ["mon", "tue"]}`) +at a fixed `utc_offset_minutes`. A window never lowers a floor. + +### Job telemetry + +A drone may add a `telemetry` object to `POST /v1/leases/{id}/result`: + +```json +{ + "exit_code": 1, "duration_secs": 412, + "telemetry": { + "job_name": "rust-test (shard 2)", "repo": "example-org/app", "run_id": 123456789, + "conclusion": "failure", "failure_class": "host", "queue_wait_secs": 18, + "steps": [{ "name": "cargo test", "duration_secs": 401, "conclusion": "failure" }] + } +} +``` + +`conclusion` is `success`, `failure`, `cancelled`, `skipped`, `timed_out` or +`unknown`; `failure_class` is `host` (machine or runner failure), `code` (the +job's own failure) or `unknown`. Unrecognized values read as `unknown`; at most +64 steps and 120 characters per name are kept; `job_name` and a `owner/name` +`repo` are required, otherwise the telemetry is dropped and the result still +recorded. The queen keeps it for `telemetry_retention_days` (default 30). +`borg report --queen ` (or `--config queen.json` on the queen host) and +`GET /report` summarize time by job and step, queue wait per class, the +cancellation share and host-caused failures per drone. + +The optional supersede janitor cancels queued pull-request workflow runs whose +head commit is no longer the pull request's head. Add a `supersede_janitor` +block (`token_path` or `BORG_JANITOR_TOKEN` for a token with `actions: write`, +`dry_run` defaulting to true, `interval_secs`, `max_cancels_per_pass`) to turn +it on. It never cancels a run that has started. + The GitHub App needs organization self-hosted runners: read and write, plus read access to Actions on the listed repositories. Set `BORG_APP_PRIVATE_KEY` instead of `private_key_path` to pass the PEM in the diff --git a/src/api.rs b/src/api.rs index b9f4cd9..3b48fe7 100644 --- a/src/api.rs +++ b/src/api.rs @@ -11,7 +11,7 @@ use std::collections::BTreeMap; use std::fmt::Write as _; use std::sync::Arc; -use axum::extract::{Path, State}; +use axum::extract::{Path, Query, State}; use axum::http::{HeaderMap, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; @@ -26,6 +26,7 @@ use crate::probe::Capacity; use crate::queen::{Queen, QueenError, QueueEntry}; use crate::secret::Secret; use crate::store::{LeaseResult, LeaseStatus, StoreError}; +use crate::telemetry::{JobTelemetry, MAX_REPORT_ROWS, Report, build_report}; #[derive(Debug, Serialize, Deserialize)] pub struct HeartbeatRequest { @@ -92,6 +93,9 @@ pub struct ResultRequest { pub duration_secs: u64, #[serde(default)] pub peak_rss_mib: Option, + /// Compact per-job summary; see [`crate::telemetry`] for the wire format. + #[serde(default)] + pub telemetry: Option, } #[derive(Debug)] @@ -130,6 +134,7 @@ pub fn router(queen: Arc) -> Router { Router::new() .route("/status", get(status)) .route("/metrics", get(metrics)) + .route("/report", get(report)) .route("/v1/heartbeat", post(heartbeat)) .route("/v1/leases/{id}/result", post(lease_result)) .with_state(queen) @@ -208,6 +213,19 @@ async fn lease_result( .finish_lease(&id, &drone, &result, now) .await? { + if let Some(t) = req.telemetry { + match t.sanitized() { + Some(t) => { + if let Some(lease) = queen.store().lease(&id).await? { + queen + .store() + .record_telemetry(&lease, &t, req.duration_secs, now) + .await?; + } + } + None => tracing::warn!("lease {id}: dropped unusable telemetry"), + } + } Ok(StatusCode::NO_CONTENT) } else { Err(ApiError::NotFound) @@ -267,6 +285,7 @@ pub struct Status { /// Drones held back for starving demand. pub reserved_for: Vec, pub rollout: RolloutStatus, + pub warm_floors: Vec, } pub async fn build_status(queen: &Queen, now: i64) -> Result { @@ -354,9 +373,38 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result unmet: snap.unmet, reserved_for: snap.reserved_for, rollout, + warm_floors: snap.warm_floors, }) } +#[derive(Debug, Deserialize)] +struct ReportQuery { + since_secs: Option, + limit: Option, +} + +/// `GET /report?since_secs=86400&limit=20`: job time, queue wait, cancellation +/// share and host failures from stored telemetry. Unauthenticated like +/// `/status`. +async fn report( + State(queen): State>, + Query(q): Query, +) -> Result, ApiError> { + let now = OffsetDateTime::now_utc().unix_timestamp(); + let max = i64::from(queen.config().telemetry_retention_days) * 86_400; + let window = q.since_secs.unwrap_or(86_400).clamp(1, max); + let rows = queen + .store() + .telemetry_since(now - window, MAX_REPORT_ROWS) + .await?; + Ok(Json(build_report( + &rows, + now, + window, + q.limit.unwrap_or(20).min(500), + ))) +} + async fn status(State(queen): State>) -> Result, ApiError> { Ok(Json( build_status(&queen, OffsetDateTime::now_utc().unix_timestamp()).await?, @@ -517,6 +565,20 @@ pub fn render_metrics(status: &Status) -> String { Vec::new() }, ); + gauge( + "borg_warm_floor", + "Idle runners the queen keeps warm per class and scope right now (base floors plus the schedule).", + status + .warm_floors + .iter() + .map(|f| { + ( + format!("class=\"{}\",scope=\"{}\"", f.class, scope_name(f.scope)), + f64::from(f.count), + ) + }) + .collect(), + ); gauge( "borg_reservations", "Drones held back for starving demand.", @@ -627,6 +689,7 @@ pub async fn serve( shutdown: impl std::future::Future + Send + 'static, ) -> std::io::Result<()> { let (stop_tx, mut stop_rx) = tokio::sync::watch::channel(false); + let stop_rx_janitor = stop_rx.clone(); let ticker = { let queen = Arc::clone(&queen); tokio::spawn(async move { @@ -652,6 +715,29 @@ pub async fn serve( } }) }; + let janitor_task = queen.janitor().map(|_| { + let queen = Arc::clone(&queen); + let mut stop_rx = stop_rx_janitor; + tokio::spawn(async move { + let Some(janitor) = queen.janitor() else { + return; + }; + let mut every = tokio::time::interval(std::time::Duration::from_secs( + janitor.settings().interval_secs, + )); + every.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + // Cancel-safe for the same reason as the ticker above. + tokio::select! { + _ = every.tick() => { + let report = janitor.pass(&queen.config().repos).await; + tracing::debug!("janitor pass: {report:?}"); + } + _ = stop_rx.changed() => break, + } + } + }) + }); let served = axum::serve(listener, router(queen)) .with_graceful_shutdown(async move { shutdown.await; @@ -659,5 +745,8 @@ pub async fn serve( }) .await; let _ = ticker.await; + if let Some(task) = janitor_task { + let _ = task.await; + } served } diff --git a/src/config.rs b/src/config.rs index d0a7542..2e7b0f5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -9,6 +9,7 @@ use serde::Deserialize; use crate::allocator::{WarmFloor, default_warm_floors}; use crate::classes::{Classes, ClassesError}; +use crate::schedule::WarmSchedule; #[derive(Debug, thiserror::Error)] pub enum ConfigError { @@ -49,6 +50,36 @@ fn default_runner_group() -> u64 { 1 } +/// The opt-in supersede janitor (see [`crate::janitor`]). Present means +/// enabled; it starts in dry-run so enabling it cannot cancel anything until +/// the operator has read its log and set `dry_run` to false. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct JanitorSettings { + /// File holding a token with `actions: write` on the listed repositories. + /// If absent, `BORG_JANITOR_TOKEN` must hold the token. + #[serde(default)] + pub token_path: Option, + /// Log what would be cancelled without cancelling. + #[serde(default = "yes")] + pub dry_run: bool, + #[serde(default = "default_janitor_interval")] + pub interval_secs: u64, + /// Hard cap on cancellations (or dry-run matches) per pass. + #[serde(default = "default_janitor_max")] + pub max_cancels_per_pass: u32, +} + +fn yes() -> bool { + true +} +fn default_janitor_interval() -> u64 { + 60 +} +fn default_janitor_max() -> u32 { + 20 +} + #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] pub struct QueenConfig { @@ -80,10 +111,19 @@ pub struct QueenConfig { pub classes: Classes, #[serde(default = "default_warm_floors")] pub warm_floors: Vec, + /// Time-of-day raises to the warm floors; see [`crate::schedule`]. + #[serde(default)] + pub warm_schedule: WarmSchedule, + /// Job telemetry older than this is deleted. + #[serde(default = "default_telemetry_days")] + pub telemetry_retention_days: u32, /// The borg version drones should run. Only reported (`/metrics`, /// `/status`): drones update themselves, the queen never forces it. #[serde(default)] pub rollout_target_version: Option, + /// Absent (the default) disables the supersede janitor. + #[serde(default)] + pub supersede_janitor: Option, } fn default_listen() -> SocketAddr { @@ -111,6 +151,9 @@ fn default_drone_stale() -> i64 { clippy::unnecessary_wraps, reason = "serde default for an Option field" )] +fn default_telemetry_days() -> u32 { + 30 +} fn default_starvation() -> Option { Some(120) } @@ -157,6 +200,16 @@ impl QueenConfig { if self.poll_interval_secs == 0 || self.lease_ttl_secs <= 0 || self.lease_max_secs <= 0 { return bad("poll_interval_secs, lease_ttl_secs and lease_max_secs must be positive"); } + self.warm_schedule + .validate() + .map_err(ConfigError::Invalid)?; + if let Some(j) = &self.supersede_janitor + && (j.interval_secs == 0 || j.max_cancels_per_pass == 0) + { + return bad( + "supersede_janitor interval_secs and max_cancels_per_pass must be positive", + ); + } Ok(()) } } @@ -189,6 +242,33 @@ mod tests { } } + #[test] + fn janitor_is_off_by_default_and_dry_run_when_enabled() { + assert!( + QueenConfig::from_json(MIN) + .unwrap() + .supersede_janitor + .is_none() + ); + let c = + QueenConfig::from_json(&MIN.replace("\"repos\"", "\"supersede_janitor\":{},\"repos\"")) + .unwrap(); + let j = c.supersede_janitor.unwrap(); + assert!(j.dry_run); + assert_eq!((j.interval_secs, j.max_cancels_per_pass), (60, 20)); + for bad in [ + "{\"interval_secs\":0}", + "{\"max_cancels_per_pass\":0}", + "{\"nope\":1}", + ] { + let json = MIN.replace( + "\"repos\"", + &format!("\"supersede_janitor\":{bad},\"repos\""), + ); + assert!(QueenConfig::from_json(&json).is_err(), "{bad}"); + } + } + #[test] fn starvation_can_be_disabled() { let c = QueenConfig::from_json( diff --git a/src/drone/mod.rs b/src/drone/mod.rs index d04df49..b3c89e3 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -338,6 +338,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { exit_code: o.exit_code, duration_secs: o.duration_secs, peak_rss_mib: o.peak_rss_mib, + telemetry: None, }, Err(e) => { tracing::warn!("lease {} could not run: {e}", grant.lease_id); @@ -345,6 +346,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { exit_code: -1, duration_secs: 0, peak_rss_mib: None, + telemetry: None, } } }; diff --git a/src/janitor.rs b/src/janitor.rs new file mode 100644 index 0000000..46df700 --- /dev/null +++ b/src/janitor.rs @@ -0,0 +1,300 @@ +//! The supersede janitor: cancels queued pull-request workflow runs whose head +//! commit is no longer the pull request's head. +//! +//! A push to a pull request branch makes every run of the older commit +//! pointless, but GitHub keeps them queued and they occupy runner demand. The +//! janitor is opt-in (`supersede_janitor` in the queen config) and uses its own +//! token, because cancelling runs needs `actions: write` and the queen's App +//! deliberately does not have it. +//! +//! Safety rules, all enforced here: +//! * only runs that are `queued` are candidates, and the run is fetched again +//! immediately before cancelling, so a run that started in between is left +//! alone; +//! * only `pull_request` runs that name exactly one pull request are touched; +//! runs from forks (no pull request attached) and merge groups never are; +//! * the pull request's current head comes from the pulls API, not from the +//! run's own snapshot; +//! * at most `max_cancels_per_pass` cancellations per pass, one pass per +//! `interval_secs`; +//! * a cancel that finds the run already finished (409) is a no-op; +//! * every cancellation (or, in dry-run, every would-be cancellation) is logged. + +use std::collections::BTreeMap; +use std::time::Duration; + +use reqwest::{Method, RequestBuilder}; +use serde::Deserialize; +use serde::de::DeserializeOwned; + +use crate::config::JanitorSettings; +use crate::github::GithubError; +use crate::secret::Secret; + +/// Upper bound on pages of queued runs read per repository (100 per page). +const MAX_PAGES: u32 = 5; + +/// What one pass did. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct PassReport { + /// Queued pull-request runs examined. + pub considered: u32, + /// Runs whose head is no longer the pull request's head. + pub superseded: u32, + /// Runs cancelled. + pub cancelled: u32, + /// Runs that dry-run would have cancelled. + pub would_cancel: u32, + /// Superseded runs left for a later pass by the per-pass cap. + pub deferred: u32, + /// Runs that changed state before the cancel and were left alone. + pub raced: u32, + /// Per-run or per-pull-request failures (logged, not fatal). + pub errors: u32, +} + +#[derive(Debug, Deserialize)] +struct Run { + id: u64, + status: String, + event: String, + head_sha: String, + #[serde(default)] + pull_requests: Vec, +} + +#[derive(Debug, Deserialize)] +struct PullRef { + number: u64, +} + +#[derive(Debug, Deserialize)] +struct Pull { + head: PullHead, +} + +#[derive(Debug, Deserialize)] +struct PullHead { + sha: String, +} + +pub struct Janitor { + http: reqwest::Client, + api_base: String, + token: Secret, + settings: JanitorSettings, +} + +impl Janitor { + pub fn new( + api_base: &str, + token: Secret, + settings: JanitorSettings, + ) -> Result { + let http = reqwest::Client::builder() + .user_agent("borg-queen-janitor") + .timeout(Duration::from_secs(30)) + .build()?; + Ok(Self { + http, + api_base: api_base.trim_end_matches('/').to_owned(), + token, + settings, + }) + } + + pub fn settings(&self) -> &JanitorSettings { + &self.settings + } + + fn request(&self, method: Method, path: &str) -> RequestBuilder { + self.http + .request(method, format!("{}{path}", self.api_base)) + .header("Accept", "application/vnd.github+json") + .header("X-GitHub-Api-Version", "2022-11-28") + .bearer_auth(self.token.expose()) + } + + async fn get(&self, path: &str) -> Result { + let resp = self.request(Method::GET, path).send().await?; + let status = resp.status(); + if !status.is_success() { + return Err(status_error("GET", path, status.as_u16(), resp).await); + } + let bytes = resp.bytes().await?; + serde_json::from_slice(&bytes).map_err(|e| GithubError::Decode(format!("{path}: {e}"))) + } + + /// Cancel a run. `Ok(false)` means GitHub says it can no longer be + /// cancelled (already finished), which is fine. + async fn cancel(&self, repo: &str, run: u64) -> Result { + let path = format!("/repos/{repo}/actions/runs/{run}/cancel"); + let resp = self.request(Method::POST, &path).send().await?; + let status = resp.status(); + match status.as_u16() { + 200..=299 => Ok(true), + 409 => Ok(false), + code => Err(status_error("POST", &path, code, resp).await), + } + } + + async fn queued_pr_runs(&self, repo: &str) -> Result, GithubError> { + #[derive(Deserialize)] + struct Runs { + workflow_runs: Vec, + } + let mut out = Vec::new(); + for page in 1..=MAX_PAGES { + let runs: Runs = self + .get(&format!( + "/repos/{repo}/actions/runs?status=queued&event=pull_request&per_page=100&page={page}" + )) + .await?; + let n = runs.workflow_runs.len(); + out.extend(runs.workflow_runs); + if n < 100 { + break; + } + } + Ok(out) + } + + /// One pass over `repos`. A failure to list a repository's runs aborts + /// that repository only; the error is counted and logged. + pub async fn pass(&self, repos: &[String]) -> PassReport { + let mut report = PassReport::default(); + let mut budget = self.settings.max_cancels_per_pass; + for repo in repos { + if let Err(e) = self.pass_repo(repo, &mut budget, &mut report).await { + tracing::warn!("janitor: {repo}: {e}"); + report.errors += 1; + } + } + report + } + + async fn pass_repo( + &self, + repo: &str, + budget: &mut u32, + report: &mut PassReport, + ) -> Result<(), GithubError> { + let runs = self.queued_pr_runs(repo).await?; + let mut heads: BTreeMap> = BTreeMap::new(); + for run in &runs { + if run.status != "queued" || run.event != "pull_request" { + continue; + } + // A fork's run carries no pull request; with more than one the + // answer is ambiguous. Leave both alone. + let [pr] = run.pull_requests.as_slice() else { + continue; + }; + report.considered += 1; + let head = if let Some(cached) = heads.get(&pr.number) { + cached.clone() + } else { + let fetched = match self + .get::(&format!("/repos/{repo}/pulls/{}", pr.number)) + .await + { + Ok(p) => Some(p.head.sha), + Err(e) => { + tracing::warn!("janitor: {repo}#{}: {e}", pr.number); + report.errors += 1; + None + } + }; + heads.insert(pr.number, fetched.clone()); + fetched + }; + let Some(current) = head else { + continue; + }; + if run.head_sha == current { + continue; + } + report.superseded += 1; + if *budget == 0 { + report.deferred += 1; + continue; + } + match self.supersede(repo, run, pr.number, ¤t).await { + Ok(Outcome::Cancelled) => { + *budget -= 1; + report.cancelled += 1; + } + Ok(Outcome::DryRun) => { + *budget -= 1; + report.would_cancel += 1; + } + Ok(Outcome::Raced) => report.raced += 1, + Err(e) => { + tracing::warn!("janitor: {repo} run {}: {e}", run.id); + report.errors += 1; + } + } + } + Ok(()) + } + + async fn supersede( + &self, + repo: &str, + run: &Run, + pr: u64, + current: &str, + ) -> Result { + // Re-read: the run may have started since the list was fetched. + let fresh: Run = self + .get(&format!("/repos/{repo}/actions/runs/{}", run.id)) + .await?; + if fresh.status != "queued" || fresh.head_sha != run.head_sha { + return Ok(Outcome::Raced); + } + if self.settings.dry_run { + tracing::info!( + "janitor (dry-run): would cancel {repo} run {} for #{pr}: head {} superseded by {current}", + run.id, + run.head_sha + ); + return Ok(Outcome::DryRun); + } + if self.cancel(repo, run.id).await? { + tracing::info!( + "janitor: cancelled {repo} run {} for #{pr}: head {} superseded by {current}", + run.id, + run.head_sha + ); + Ok(Outcome::Cancelled) + } else { + Ok(Outcome::Raced) + } + } +} + +enum Outcome { + Cancelled, + DryRun, + Raced, +} + +async fn status_error( + method: &'static str, + path: &str, + status: u16, + resp: reqwest::Response, +) -> GithubError { + GithubError::Status { + method, + path: path.to_owned(), + status, + body: resp + .text() + .await + .unwrap_or_default() + .chars() + .take(200) + .collect(), + } +} diff --git a/src/lib.rs b/src/lib.rs index e267324..5be323d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,7 +15,10 @@ pub mod config; pub mod drone; pub mod enroll; pub mod github; +pub mod janitor; pub mod probe; pub mod queen; +pub mod schedule; pub mod secret; pub mod store; +pub mod telemetry; diff --git a/src/main.rs b/src/main.rs index e2a7f6c..2eaa67c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use anyhow::{Context, Result, bail}; use borg::api; -use borg::config::QueenConfig; +use borg::config::{JanitorSettings, QueenConfig}; use borg::drone::client::QueenClient; use borg::drone::config::DroneConfig; use borg::drone::exec::SystemExec; @@ -11,6 +11,7 @@ use borg::drone::install::{self, InstallOptions, Platform}; use borg::drone::{Drone, pins}; use borg::enroll; use borg::github::GithubClient; +use borg::janitor::Janitor; use borg::probe::{self, system::SystemHost}; use borg::queen::Queen; use borg::store::Store; @@ -36,6 +37,9 @@ enum Command { #[command(subcommand)] command: DroneCommand, }, + /// Summarize stored job telemetry: time by job and step, queue wait per + /// class, cancellation share and host-caused failures per drone. + Report(ReportArgs), /// Run or administer the queen (the coordinator). Queen { #[command(subcommand)] @@ -43,6 +47,25 @@ enum Command { }, } +#[derive(Args)] +struct ReportArgs { + /// Fetch from a running queen, e.g. http://queen.internal:8080 + #[arg(long, conflicts_with = "config", required_unless_present = "config")] + queen: Option, + /// Read the queen's SQLite file directly (run on the queen host). + #[arg(long)] + config: Option, + /// Look back this many hours (capped by the queen's retention). + #[arg(long, default_value_t = 24)] + since_hours: u32, + /// Rows in the by-job and by-step lists. + #[arg(long, default_value_t = 20)] + limit: usize, + /// Print the JSON report instead of text. + #[arg(long)] + json: bool, +} + #[derive(Args)] struct ConfigArg { /// Path to the queen's JSON config file. @@ -148,6 +171,20 @@ fn private_key(cfg: &QueenConfig) -> Result> { } } +fn janitor_token(settings: &JanitorSettings) -> Result { + let raw = match &settings.token_path { + Some(path) => std::fs::read_to_string(path) + .with_context(|| format!("read janitor token {}", path.display()))?, + None => std::env::var("BORG_JANITOR_TOKEN") + .context("set supersede_janitor.token_path or BORG_JANITOR_TOKEN")?, + }; + let token = raw.trim(); + if token.is_empty() { + bail!("the janitor token is empty"); + } + Ok(borg::secret::Secret::new(token)) +} + async fn open_store(cfg: &QueenConfig) -> Result { Store::open(&cfg.db_path) .await @@ -182,7 +219,26 @@ async fn run_queen(path: &Path) -> Result<()> { .await .with_context(|| format!("bind {}", cfg.listen))?; tracing::info!("queen listening on {}", cfg.listen); - let queen = Arc::new(Queen::new(cfg, store, github)); + let janitor = match &cfg.supersede_janitor { + Some(settings) => { + let token = janitor_token(settings)?; + tracing::info!( + "supersede janitor enabled ({})", + if settings.dry_run { + "dry-run" + } else { + "cancelling" + } + ); + Some(Janitor::new(&cfg.github.api_base, token, settings.clone())?) + } + None => None, + }; + let mut queen = Queen::new(cfg, store, github); + if let Some(j) = janitor { + queen = queen.with_janitor(j); + } + let queen = Arc::new(queen); api::serve(queen, listener, wait_for_shutdown()).await?; Ok(()) } @@ -205,6 +261,42 @@ async fn enroll_drone(path: &Path, name: &str) -> Result<()> { Ok(()) } +async fn report(a: ReportArgs) -> Result<()> { + let window = i64::from(a.since_hours.max(1)) * 3600; + let report = match (&a.queen, &a.config) { + (Some(url), _) => { + let url = format!( + "{}/report?since_secs={window}&limit={}", + url.trim_end_matches('/'), + a.limit + ); + reqwest::get(&url) + .await + .and_then(reqwest::Response::error_for_status) + .with_context(|| format!("GET {url}"))? + .json::() + .await + .context("decode report")? + } + (None, Some(path)) => { + let cfg = QueenConfig::load(path)?; + let store = open_store(&cfg).await?; + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + let rows = store + .telemetry_since(now - window, borg::telemetry::MAX_REPORT_ROWS) + .await?; + borg::telemetry::build_report(&rows, now, window, a.limit) + } + (None, None) => bail!("give --queen or --config"), + }; + if a.json { + println!("{}", serde_json::to_string_pretty(&report)?); + } else { + print!("{}", borg::telemetry::render_text(&report)); + } + Ok(()) +} + async fn set_quarantine(path: &Path, name: &str, reason: Option<&str>) -> Result<()> { let cfg = QueenConfig::load(path)?; let store = open_store(&cfg).await?; @@ -276,6 +368,7 @@ async fn main() -> Result<()> { println!("{json}"); Ok(()) } + Command::Report(a) => report(a).await, Command::Drone { command } => { tracing_subscriber::fmt() .with_env_filter( diff --git a/src/queen.rs b/src/queen.rs index 1e36edf..44be470 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -21,11 +21,12 @@ use time::OffsetDateTime; use crate::allocator::{ ActiveLease, AllocatorInput, Demand, DroneState, LeaseState, QueueAge, Reservation, Unmet, - allocate, max_slots, + WarmFloor, allocate, max_slots, }; use crate::classes::{Class, Scope}; use crate::config::QueenConfig; use crate::github::{GithubClient, GithubError, QueuedJob, Runner}; +use crate::janitor::Janitor; use crate::probe::Capacity; use crate::store::{LeaseRow, LeaseStatus, Store, StoreError}; @@ -45,6 +46,8 @@ pub struct Snapshot { pub queue: Vec, pub reserved_for: Vec, pub unmet: Vec, + /// The warm floors in force at the last pass (base plus schedule). + pub warm_floors: Vec, } #[derive(Debug, Clone, Serialize)] @@ -61,6 +64,7 @@ pub struct Queen { pub(crate) store: Store, github: GithubClient, snapshot: RwLock, + janitor: Option, tick_lock: tokio::sync::Mutex<()>, } @@ -102,10 +106,22 @@ impl Queen { store, github, snapshot: RwLock::new(Snapshot::default()), + janitor: None, tick_lock: tokio::sync::Mutex::new(()), } } + /// Attach the opt-in supersede janitor; `api::serve` runs it. + #[must_use] + pub fn with_janitor(mut self, janitor: Janitor) -> Self { + self.janitor = Some(janitor); + self + } + + pub fn janitor(&self) -> Option<&Janitor> { + self.janitor.as_ref() + } + pub fn config(&self) -> &QueenConfig { &self.cfg } @@ -184,12 +200,16 @@ impl Queen { }, }) .collect(); + let warm_floors = self + .cfg + .warm_schedule + .effective_floors(&self.cfg.warm_floors, now); let allocation = allocate(&AllocatorInput { classes: self.cfg.classes.clone(), demand, drones, leases, - warm_floors: self.cfg.warm_floors.clone(), + warm_floors: warm_floors.clone(), queue_ages, starvation_after_secs: self.cfg.starvation_after_secs, }); @@ -245,7 +265,12 @@ impl Queen { s.queue = queue; s.reserved_for = allocation.reserved_for.clone(); s.unmet = allocation.unmet.clone(); + s.warm_floors = warm_floors; }); + let keep = i64::from(self.cfg.telemetry_retention_days) * 86_400; + if let Err(e) = self.store.prune_telemetry(now_ts - keep).await { + tracing::warn!("prune telemetry failed: {e}"); + } Ok(()) } diff --git a/src/schedule.rs b/src/schedule.rs new file mode 100644 index 0000000..9e036f7 --- /dev/null +++ b/src/schedule.rs @@ -0,0 +1,294 @@ +//! Time-of-day warm runner pools. +//! +//! The base warm floors (`warm_floors`) keep a number of idle runners alive per +//! class and scope all day. A schedule raises them during configured windows: +//! for example more warm light pull-request runners from 08:00 to 20:00 on +//! weekdays. The result is still just a list of [`WarmFloor`]s for the +//! allocator, so every allocator guarantee (never over-grant, main first) +//! holds unchanged. +//! +//! Times are wall-clock times at a fixed UTC offset +//! (`utc_offset_minutes`). The `time` crate has no zone database; the offset +//! does not follow daylight saving, so change it twice a year or write the +//! windows in UTC. + +use std::collections::BTreeMap; + +use serde::Deserialize; +use time::{OffsetDateTime, UtcOffset, Weekday}; + +use crate::allocator::WarmFloor; +use crate::classes::{Class, Scope}; + +const MINUTES_PER_DAY: u16 = 24 * 60; + +/// Minutes since midnight, parsed from `HH:MM` (`24:00` is allowed as an end). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(try_from = "String")] +pub struct TimeOfDay(u16); + +impl TryFrom for TimeOfDay { + type Error = String; + + fn try_from(s: String) -> Result { + let bad = || format!("{s:?} is not a time of day like 08:00"); + let (h, m) = s.split_once(':').ok_or_else(bad)?; + let digits = |p: &str| p.len() == 2 && p.bytes().all(|b| b.is_ascii_digit()); + if !digits(h) || !digits(m) { + return Err(bad()); + } + let (h, m): (u16, u16) = (h.parse().map_err(|_| bad())?, m.parse().map_err(|_| bad())?); + let total = h * 60 + m; + if m > 59 || total > MINUTES_PER_DAY { + return Err(bad()); + } + Ok(Self(total)) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Day { + Mon, + Tue, + Wed, + Thu, + Fri, + Sat, + Sun, +} + +impl Day { + fn of(w: Weekday) -> Self { + match w { + Weekday::Monday => Self::Mon, + Weekday::Tuesday => Self::Tue, + Weekday::Wednesday => Self::Wed, + Weekday::Thursday => Self::Thu, + Weekday::Friday => Self::Fri, + Weekday::Saturday => Self::Sat, + Weekday::Sunday => Self::Sun, + } + } +} + +/// Keep `count` warm runners of `(class, scope)` between `start` and `end`. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WarmWindow { + pub class: Class, + pub scope: Scope, + pub count: u32, + pub start: TimeOfDay, + /// Exclusive. An `end` before `start` wraps past midnight. + pub end: TimeOfDay, + /// Days the window starts on; every day when absent. + #[serde(default)] + pub days: Option>, +} + +impl WarmWindow { + fn active(&self, day: Day, prev_day: Day, minute: u16) -> bool { + let on = |d: Day| self.days.as_ref().is_none_or(|ds| ds.contains(&d)); + let (start, end) = (self.start.0, self.end.0); + if start < end { + on(day) && (start..end).contains(&minute) + } else { + // Wraps midnight: the evening part belongs to the starting day, + // the morning part to the day before. + (on(day) && minute >= start) || (on(prev_day) && minute < end) + } + } +} + +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WarmSchedule { + /// Offset of the wall clock the windows are written in, in minutes east of + /// UTC (US Pacific standard time is -480). + #[serde(default)] + pub utc_offset_minutes: i32, + #[serde(default)] + pub windows: Vec, +} + +impl WarmSchedule { + pub fn validate(&self) -> Result<(), String> { + if UtcOffset::from_whole_seconds(self.utc_offset_minutes.saturating_mul(60)).is_err() { + return Err("warm_schedule.utc_offset_minutes is out of range".to_owned()); + } + for w in &self.windows { + if w.start == w.end { + return Err("a warm_schedule window needs start different from end".to_owned()); + } + if w.days.as_ref().is_some_and(Vec::is_empty) { + return Err("a warm_schedule window's days may not be empty".to_owned()); + } + } + Ok(()) + } + + /// The floors in force at `now`: per `(class, scope)` the larger of the + /// base floor and every active window. Sorted by class and scope. + pub fn effective_floors(&self, base: &[WarmFloor], now: OffsetDateTime) -> Vec { + let mut floors: BTreeMap<(Class, Scope), u32> = BTreeMap::new(); + for f in base { + let e = floors.entry((f.class, f.scope)).or_default(); + *e = (*e).max(f.count); + } + let offset = UtcOffset::from_whole_seconds(self.utc_offset_minutes.saturating_mul(60)) + .unwrap_or(UtcOffset::UTC); + let local = now.to_offset(offset); + let (day, prev_day) = ( + Day::of(local.weekday()), + Day::of(local.weekday().previous()), + ); + let minute = u16::from(local.hour()) * 60 + u16::from(local.minute()); + for w in self + .windows + .iter() + .filter(|w| w.active(day, prev_day, minute)) + { + let e = floors.entry((w.class, w.scope)).or_default(); + *e = (*e).max(w.count); + } + floors + .into_iter() + .map(|((class, scope), count)| WarmFloor { + class, + scope, + count, + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use time::macros::datetime; + + use super::*; + use crate::allocator::default_warm_floors; + + fn schedule(json: &str) -> WarmSchedule { + let s: WarmSchedule = serde_json::from_str(json).unwrap(); + s.validate().unwrap(); + s + } + + fn light_pr(floors: &[WarmFloor]) -> u32 { + floors + .iter() + .find(|f| f.class == Class::Light && f.scope == Scope::Pr) + .map_or(0, |f| f.count) + } + + const DAYTIME: &str = r#"{"utc_offset_minutes": -480, "windows": [ + {"class":"light","scope":"pr","count":3,"start":"08:00","end":"20:00", + "days":["mon","tue","wed","thu","fri"]}]}"#; + + #[test] + fn window_applies_at_local_time_and_only_on_listed_days() { + let s = schedule(DAYTIME); + let base = default_warm_floors(); + // 2026-10-12 is a Monday. 15:59 UTC is 07:59 at UTC-8: not yet. + assert_eq!( + light_pr(&s.effective_floors(&base, datetime!(2026-10-12 15:59 UTC))), + 0 + ); + assert_eq!( + light_pr(&s.effective_floors(&base, datetime!(2026-10-12 16:00 UTC))), + 3 + ); + // End is exclusive: 20:00 local is 04:00 UTC the next day. + assert_eq!( + light_pr(&s.effective_floors(&base, datetime!(2026-10-13 03:59 UTC))), + 3 + ); + assert_eq!( + light_pr(&s.effective_floors(&base, datetime!(2026-10-13 04:00 UTC))), + 0 + ); + // Saturday noon local. + assert_eq!( + light_pr(&s.effective_floors(&base, datetime!(2026-10-17 20:00 UTC))), + 0 + ); + } + + #[test] + fn windows_only_raise_the_base_floor_never_lower_it() { + let s = schedule( + r#"{"windows":[{"class":"medium","scope":"main","count":0,"start":"00:00","end":"24:00"}]}"#, + ); + let floors = s.effective_floors(&default_warm_floors(), datetime!(2026-10-12 12:00 UTC)); + assert_eq!(floors, default_warm_floors()); + let s = schedule( + r#"{"windows":[{"class":"medium","scope":"main","count":4,"start":"00:00","end":"24:00"}]}"#, + ); + let floors = s.effective_floors(&default_warm_floors(), datetime!(2026-10-12 12:00 UTC)); + assert_eq!(floors[0].count, 4); + assert_eq!(floors.len(), 1); + } + + #[test] + fn overlapping_windows_take_the_largest_count() { + let s = schedule( + r#"{"windows":[ + {"class":"light","scope":"pr","count":2,"start":"08:00","end":"20:00"}, + {"class":"light","scope":"pr","count":5,"start":"12:00","end":"13:00"}]}"#, + ); + assert_eq!( + light_pr(&s.effective_floors(&[], datetime!(2026-10-12 11:00 UTC))), + 2 + ); + assert_eq!( + light_pr(&s.effective_floors(&[], datetime!(2026-10-12 12:30 UTC))), + 5 + ); + } + + #[test] + fn a_window_may_wrap_midnight_and_belongs_to_its_starting_day() { + let s = schedule( + r#"{"windows":[{"class":"light","scope":"pr","count":2, + "start":"22:00","end":"06:00","days":["fri"]}]}"#, + ); + // 2026-10-16 is a Friday. + let at = |dt| light_pr(&s.effective_floors(&[], dt)); + assert_eq!(at(datetime!(2026-10-16 21:59 UTC)), 0); + assert_eq!(at(datetime!(2026-10-16 22:00 UTC)), 2); + assert_eq!(at(datetime!(2026-10-17 05:59 UTC)), 2); // Saturday morning + assert_eq!(at(datetime!(2026-10-17 06:00 UTC)), 0); + assert_eq!(at(datetime!(2026-10-16 05:00 UTC)), 0); // Friday morning + assert_eq!(at(datetime!(2026-10-17 22:30 UTC)), 0); // Saturday night + } + + #[test] + fn no_schedule_means_the_base_floors_unchanged() { + let s = WarmSchedule::default(); + assert_eq!( + s.effective_floors(&default_warm_floors(), datetime!(2026-10-12 12:00 UTC)), + default_warm_floors() + ); + } + + #[test] + fn bad_times_and_windows_are_rejected() { + for t in ["8:00", "08:60", "25:00", "24:01", "0800", "08:0a", ""] { + let json = format!( + r#"{{"windows":[{{"class":"light","scope":"pr","count":1,"start":"{t}","end":"09:00"}}]}}"# + ); + assert!(serde_json::from_str::(&json).is_err(), "{t}"); + } + for json in [ + r#"{"windows":[{"class":"light","scope":"pr","count":1,"start":"08:00","end":"08:00"}]}"#, + r#"{"windows":[{"class":"light","scope":"pr","count":1,"start":"08:00","end":"09:00","days":[]}]}"#, + r#"{"utc_offset_minutes": 100000}"#, + ] { + let s: WarmSchedule = serde_json::from_str(json).unwrap(); + assert!(s.validate().is_err(), "{json}"); + } + assert!(serde_json::from_str::(r#"{"windows":[],"x":1}"#).is_err()); + } +} diff --git a/src/store.rs b/src/store.rs index bca4707..751eab3 100644 --- a/src/store.rs +++ b/src/store.rs @@ -50,6 +50,22 @@ CREATE TABLE IF NOT EXISTS leases ( peak_rss_mib INTEGER ); CREATE INDEX IF NOT EXISTS leases_state ON leases(state); +CREATE TABLE IF NOT EXISTS job_telemetry ( + lease_id TEXT PRIMARY KEY, + drone TEXT NOT NULL, + class TEXT NOT NULL, + scope TEXT NOT NULL, + repo TEXT NOT NULL, + run_id INTEGER NOT NULL, + job_name TEXT NOT NULL, + conclusion TEXT NOT NULL, + failure_class TEXT NOT NULL, + duration_secs INTEGER NOT NULL, + queue_wait_secs INTEGER, + steps TEXT NOT NULL, + finished_at INTEGER NOT NULL +); +CREATE INDEX IF NOT EXISTS job_telemetry_finished ON job_telemetry(finished_at); "; #[derive(Debug, thiserror::Error)] @@ -145,7 +161,7 @@ pub struct LeaseResult { #[derive(Clone)] pub struct Store { - pool: SqlitePool, + pub(crate) pool: SqlitePool, } fn to_text(v: &T) -> Result { diff --git a/src/telemetry.rs b/src/telemetry.rs new file mode 100644 index 0000000..ae0580e --- /dev/null +++ b/src/telemetry.rs @@ -0,0 +1,769 @@ +//! Per-lease job telemetry: the wire format drones report in a lease result, +//! its validation, and its storage. +//! +//! # Wire format +//! +//! `POST /v1/leases/{id}/result` accepts an optional `telemetry` object next to +//! `exit_code`, `duration_secs` and `peak_rss_mib`: +//! +//! ```json +//! { +//! "exit_code": 0, +//! "duration_secs": 412, +//! "telemetry": { +//! "job_name": "rust-test (shard 2)", +//! "repo": "example-org/app", +//! "run_id": 123456789, +//! "conclusion": "failure", +//! "failure_class": "host", +//! "queue_wait_secs": 18, +//! "steps": [ +//! { "name": "Set up job", "duration_secs": 3, "conclusion": "success" }, +//! { "name": "cargo test", "duration_secs": 401, "conclusion": "failure" } +//! ] +//! } +//! } +//! ``` +//! +//! * `job_name` (required, at most 120 characters) and `repo` (required, +//! `owner/name`) identify the job; `run_id` is the workflow run id. +//! * `conclusion`: `success`, `failure`, `cancelled`, `skipped`, `timed_out` or +//! `unknown` (the default; any unrecognized value also reads as `unknown`). +//! * `failure_class`: `host` (the machine or runner failed: out of disk, +//! runner lost, docker daemon down), `code` (the job's own steps failed) or +//! `unknown` (the default). Only meaningful for a job that did not succeed; +//! the queen stores `unknown` for `success` and `skipped`. +//! * `queue_wait_secs` (optional): seconds the job waited between being queued +//! and starting on this runner. +//! * `steps` (optional, at most 64, extra entries are dropped): name (at most +//! 120 characters), `duration_secs`, `conclusion`. +//! +//! Everything is optional from the queen's point of view: a result without +//! `telemetry` is accepted exactly as before, and a `telemetry` object the +//! queen cannot use (empty `job_name`, malformed `repo`) is dropped with a +//! warning while the result itself is still recorded. Control characters are +//! stripped from every string. Reporting the same lease twice replaces the +//! first report. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use sqlx::Row; + +use crate::classes::{Class, Scope}; +use crate::store::{LeaseRow, Store, StoreError}; + +pub const MAX_STEPS: usize = 64; +pub const MAX_NAME_CHARS: usize = 120; +/// Rows the report reads at most (newest first). +pub const MAX_REPORT_ROWS: i64 = 100_000; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Conclusion { + Success, + Failure, + Cancelled, + Skipped, + TimedOut, + #[default] + #[serde(other)] + Unknown, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum FailureClass { + Host, + Code, + #[default] + #[serde(other)] + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct StepTiming { + pub name: String, + pub duration_secs: u64, + #[serde(default)] + pub conclusion: Conclusion, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct JobTelemetry { + pub job_name: String, + pub repo: String, + #[serde(default)] + pub run_id: u64, + #[serde(default)] + pub conclusion: Conclusion, + #[serde(default)] + pub failure_class: FailureClass, + #[serde(default)] + pub queue_wait_secs: Option, + #[serde(default)] + pub steps: Vec, +} + +fn clean(s: &str) -> String { + s.chars() + .filter(|c| !c.is_control()) + .take(MAX_NAME_CHARS) + .collect::() + .trim() + .to_owned() +} + +impl JobTelemetry { + /// The report as the queen stores it, or `None` if it is unusable. + pub fn sanitized(mut self) -> Option { + self.job_name = clean(&self.job_name); + self.repo = clean(&self.repo); + let repo_ok = self.repo.split('/').filter(|p| !p.is_empty()).count() == 2 + && !self.repo.contains(char::is_whitespace); + if self.job_name.is_empty() || !repo_ok { + return None; + } + self.steps.truncate(MAX_STEPS); + for s in &mut self.steps { + s.name = clean(&s.name); + } + if matches!(self.conclusion, Conclusion::Success | Conclusion::Skipped) { + self.failure_class = FailureClass::Unknown; + } + Some(self) + } +} + +/// One stored report with the lease facts the queen added. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TelemetryRow { + pub lease_id: String, + pub drone: String, + pub class: Class, + pub scope: Scope, + pub duration_secs: u64, + pub finished_at: i64, + pub telemetry: JobTelemetry, +} + +fn text(v: &T) -> Result { + match serde_json::to_value(v) { + Ok(serde_json::Value::String(s)) => Ok(s), + _ => Err(StoreError::Corrupt("enum is not a string".to_owned())), + } +} + +fn parse Deserialize<'de>>(s: &str) -> Result { + serde_json::from_str(&format!("\"{s}\"")).map_err(|e| StoreError::Corrupt(e.to_string())) +} + +fn to_i64(v: u64) -> i64 { + i64::try_from(v).unwrap_or(i64::MAX) +} + +fn to_u64(v: i64) -> u64 { + u64::try_from(v).unwrap_or(0) +} + +impl Store { + /// Store (or replace) the telemetry of a finished lease. + pub async fn record_telemetry( + &self, + lease: &LeaseRow, + t: &JobTelemetry, + duration_secs: u64, + now: i64, + ) -> Result<(), StoreError> { + let steps = + serde_json::to_string(&t.steps).map_err(|e| StoreError::Corrupt(e.to_string()))?; + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT OR REPLACE INTO job_telemetry + (lease_id, drone, class, scope, repo, run_id, job_name, conclusion, failure_class, + duration_secs, queue_wait_secs, steps, finished_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&lease.id) + .bind(&lease.drone) + .bind(text(&lease.class)?) + .bind(text(&lease.scope)?) + .bind(&t.repo) + .bind(to_i64(t.run_id)) + .bind(&t.job_name) + .bind(text(&t.conclusion)?) + .bind(text(&t.failure_class)?) + .bind(to_i64(duration_secs)) + .bind(t.queue_wait_secs.map(to_i64)) + .bind(steps) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Reports finished at or after `since`, newest first, at most `limit`. + pub async fn telemetry_since( + &self, + since: i64, + limit: i64, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT * FROM job_telemetry WHERE finished_at >= ? ORDER BY finished_at DESC LIMIT ?", + ) + .bind(since) + .bind(limit) + .fetch_all(&self.pool) + .await?; + rows.iter() + .map(|r| { + let steps: String = r.get("steps"); + Ok(TelemetryRow { + lease_id: r.get("lease_id"), + drone: r.get("drone"), + class: parse(&r.get::("class"))?, + scope: parse(&r.get::("scope"))?, + duration_secs: to_u64(r.get("duration_secs")), + finished_at: r.get("finished_at"), + telemetry: JobTelemetry { + job_name: r.get("job_name"), + repo: r.get("repo"), + run_id: to_u64(r.get("run_id")), + conclusion: parse(&r.get::("conclusion"))?, + failure_class: parse(&r.get::("failure_class"))?, + queue_wait_secs: r.get::, _>("queue_wait_secs").map(to_u64), + steps: serde_json::from_str(&steps) + .map_err(|e| StoreError::Corrupt(e.to_string()))?, + }, + }) + }) + .collect() + } + + /// Delete reports finished before `before`; returns how many. + pub async fn prune_telemetry(&self, before: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let r = sqlx::query("DELETE FROM job_telemetry WHERE finished_at < ?") + .bind(before) + .execute(&self.pool) + .await?; + Ok(r.rows_affected()) + } +} + +// ---------------------------------------------------------------- report + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct JobTime { + pub repo: String, + pub job_name: String, + pub runs: u32, + pub total_secs: u64, + pub avg_secs: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct StepTime { + pub job_name: String, + pub step: String, + pub runs: u32, + pub total_secs: u64, + pub avg_secs: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct QueueWait { + pub class: Class, + /// Jobs that reported a queue wait. + pub samples: u32, + pub avg_secs: u64, + pub p95_secs: u64, + pub max_secs: u64, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Cancellation { + pub jobs: u32, + pub cancelled: u32, + /// `cancelled / jobs`, 0 when there are no jobs. + pub share: f64, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct DroneHealth { + pub drone: String, + pub jobs: u32, + /// Jobs that failed or timed out (any class). + pub failed: u32, + /// Jobs whose failure was caused by the host. + pub host_failures: u32, + pub code_failures: u32, + pub cancelled: u32, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct Report { + pub now: i64, + pub window_secs: i64, + pub jobs: u32, + pub by_job: Vec, + pub by_step: Vec, + pub queue_wait: Vec, + pub cancellation: Cancellation, + pub drones: Vec, +} + +#[derive(Default)] +struct Sum { + runs: u32, + total: u64, +} + +impl Sum { + fn add(&mut self, secs: u64) { + self.runs += 1; + self.total = self.total.saturating_add(secs); + } + + fn avg(&self) -> u64 { + self.total / u64::from(self.runs.max(1)) + } +} + +fn count(n: usize) -> u32 { + u32::try_from(n).unwrap_or(u32::MAX) +} + +/// Aggregate `rows` (all finished within the last `window_secs`). The by-job +/// and by-step lists hold the `limit` entries with the most total time. +pub fn build_report(rows: &[TelemetryRow], now: i64, window_secs: i64, limit: usize) -> Report { + let mut jobs: BTreeMap<(&str, &str), Sum> = BTreeMap::new(); + let mut steps: BTreeMap<(&str, &str), Sum> = BTreeMap::new(); + let mut waits: BTreeMap> = BTreeMap::new(); + let mut drones: BTreeMap<&str, DroneHealth> = BTreeMap::new(); + let mut cancelled = 0u32; + + for r in rows { + let t = &r.telemetry; + jobs.entry((t.repo.as_str(), t.job_name.as_str())) + .or_default() + .add(r.duration_secs); + for s in &t.steps { + steps + .entry((t.job_name.as_str(), s.name.as_str())) + .or_default() + .add(s.duration_secs); + } + if let Some(w) = t.queue_wait_secs { + waits.entry(r.class).or_default().push(w); + } + let d = drones + .entry(r.drone.as_str()) + .or_insert_with(|| DroneHealth { + drone: r.drone.clone(), + jobs: 0, + failed: 0, + host_failures: 0, + code_failures: 0, + cancelled: 0, + }); + d.jobs += 1; + let failed = matches!(t.conclusion, Conclusion::Failure | Conclusion::TimedOut); + if failed { + d.failed += 1; + } + if t.conclusion == Conclusion::Cancelled { + d.cancelled += 1; + cancelled += 1; + } + match t.failure_class { + FailureClass::Host => d.host_failures += 1, + FailureClass::Code => d.code_failures += 1, + FailureClass::Unknown => {} + } + } + + let mut by_job: Vec = jobs + .into_iter() + .map(|((repo, job), s)| JobTime { + repo: repo.to_owned(), + job_name: job.to_owned(), + runs: s.runs, + total_secs: s.total, + avg_secs: s.avg(), + }) + .collect(); + by_job.sort_by_key(|j| std::cmp::Reverse(j.total_secs)); + by_job.truncate(limit); + + let mut by_step: Vec = steps + .into_iter() + .map(|((job, step), s)| StepTime { + job_name: job.to_owned(), + step: step.to_owned(), + runs: s.runs, + total_secs: s.total, + avg_secs: s.avg(), + }) + .collect(); + by_step.sort_by_key(|s| std::cmp::Reverse(s.total_secs)); + by_step.truncate(limit); + + let queue_wait = waits + .into_iter() + .map(|(class, mut v)| { + v.sort_unstable(); + let n = v.len(); + // Nearest-rank p95. + let rank = (n * 95).div_ceil(100).max(1); + QueueWait { + class, + samples: count(n), + avg_secs: v.iter().sum::() / u64::try_from(n).unwrap_or(1).max(1), + p95_secs: v.get(rank - 1).copied().unwrap_or(0), + max_secs: v.last().copied().unwrap_or(0), + } + }) + .collect(); + + let total = count(rows.len()); + let share = if total == 0 { + 0.0 + } else { + f64::from(cancelled) / f64::from(total) + }; + let mut drones: Vec = drones.into_values().collect(); + drones.sort_by(|a, b| { + b.host_failures + .cmp(&a.host_failures) + .then_with(|| a.drone.cmp(&b.drone)) + }); + Report { + now, + window_secs, + jobs: total, + by_job, + by_step, + queue_wait, + cancellation: Cancellation { + jobs: total, + cancelled, + share, + }, + drones, + } +} + +fn dur(secs: u64) -> String { + match secs { + 0..60 => format!("{secs}s"), + 60..3600 => format!("{}m{:02}s", secs / 60, secs % 60), + _ => format!("{}h{:02}m", secs / 3600, secs % 3600 / 60), + } +} + +/// Plain-text rendering for `borg report`. +pub fn render_text(r: &Report) -> String { + use std::fmt::Write as _; + let mut o = String::new(); + let _ = writeln!( + o, + "borg report: last {}, {} job(s) with telemetry\n", + dur(u64::try_from(r.window_secs).unwrap_or(0)), + r.jobs + ); + let _ = writeln!( + o, + "Time by job\n {:>9} {:>5} {:>8} job", + "total", "runs", "avg" + ); + for j in &r.by_job { + let _ = writeln!( + o, + " {:>9} {:>5} {:>8} {} / {}", + dur(j.total_secs), + j.runs, + dur(j.avg_secs), + j.repo, + j.job_name + ); + } + let _ = writeln!( + o, + "\nTime by step\n {:>9} {:>5} {:>8} step", + "total", "runs", "avg" + ); + for s in &r.by_step { + let _ = writeln!( + o, + " {:>9} {:>5} {:>8} {} / {}", + dur(s.total_secs), + s.runs, + dur(s.avg_secs), + s.job_name, + s.step + ); + } + let _ = writeln!( + o, + "\nQueue wait by class\n {:<8} {:>7} {:>8} {:>8} {:>8}", + "class", "samples", "avg", "p95", "max" + ); + for q in &r.queue_wait { + let _ = writeln!( + o, + " {:<8} {:>7} {:>8} {:>8} {:>8}", + q.class, + q.samples, + dur(q.avg_secs), + dur(q.p95_secs), + dur(q.max_secs) + ); + } + let c = &r.cancellation; + let _ = writeln!( + o, + "\nCancellations: {} of {} job(s) ({:.1}%)", + c.cancelled, + c.jobs, + c.share * 100.0 + ); + let _ = writeln!( + o, + "\nDrones\n {:<20} {:>5} {:>7} {:>5} {:>5} {:>9}", + "drone", "jobs", "failed", "host", "code", "cancelled" + ); + for d in &r.drones { + let _ = writeln!( + o, + " {:<20} {:>5} {:>7} {:>5} {:>5} {:>9}", + d.drone, d.jobs, d.failed, d.host_failures, d.code_failures, d.cancelled + ); + } + o +} + +#[cfg(test)] +mod tests { + use super::*; + + fn row( + drone: &str, + class: Class, + job: &str, + conclusion: Conclusion, + fc: FailureClass, + ) -> TelemetryRow { + TelemetryRow { + lease_id: format!("{drone}-{job}"), + drone: drone.to_owned(), + class, + scope: Scope::Pr, + duration_secs: 100, + finished_at: 1000, + telemetry: JobTelemetry { + job_name: job.to_owned(), + repo: "o/r".to_owned(), + run_id: 1, + conclusion, + failure_class: fc, + queue_wait_secs: Some(10), + steps: vec![ + StepTiming { + name: "build".into(), + duration_secs: 60, + conclusion, + }, + StepTiming { + name: "test".into(), + duration_secs: 30, + conclusion, + }, + ], + }, + } + } + + #[test] + fn wire_format_reads_defaults_and_unknown_values() { + let t: JobTelemetry = serde_json::from_str( + r#"{"job_name":"j","repo":"o/r","conclusion":"weird","failure_class":"cosmic", + "steps":[{"name":"s","duration_secs":3}]}"#, + ) + .unwrap(); + assert_eq!(t.conclusion, Conclusion::Unknown); + assert_eq!(t.failure_class, FailureClass::Unknown); + assert_eq!(t.steps[0].conclusion, Conclusion::Unknown); + assert_eq!((t.run_id, t.queue_wait_secs), (0, None)); + let timed: Conclusion = serde_json::from_str("\"timed_out\"").unwrap(); + assert_eq!(timed, Conclusion::TimedOut); + } + + #[test] + fn sanitizing_bounds_and_cleans() { + let mut t = row( + "d", + Class::Light, + "j", + Conclusion::Success, + FailureClass::Host, + ) + .telemetry; + t.job_name = format!("a\u{7}b{}", "x".repeat(500)); + t.steps = vec![ + StepTiming { + name: "s".into(), + duration_secs: 1, + conclusion: Conclusion::Success + }; + 200 + ]; + let s = t.clone().sanitized().unwrap(); + assert_eq!(s.job_name.chars().count(), MAX_NAME_CHARS); + assert!(!s.job_name.contains('\u{7}')); + assert_eq!(s.steps.len(), MAX_STEPS); + // A successful job has no failure class. + assert_eq!(s.failure_class, FailureClass::Unknown); + for bad in [("", "o/r"), ("j", "nope"), ("j", "a/b/c"), ("j", "a b/c")] { + let mut t = t.clone(); + (t.job_name, t.repo) = (bad.0.into(), bad.1.into()); + assert!(t.sanitized().is_none(), "{bad:?}"); + } + } + + #[test] + fn report_aggregates_time_waits_cancellations_and_host_failures() { + let rows = vec![ + row( + "d1", + Class::Light, + "lint", + Conclusion::Success, + FailureClass::Unknown, + ), + row( + "d1", + Class::Light, + "lint", + Conclusion::Failure, + FailureClass::Host, + ), + row( + "d2", + Class::Heavy, + "test", + Conclusion::Failure, + FailureClass::Code, + ), + row( + "d2", + Class::Heavy, + "test", + Conclusion::Cancelled, + FailureClass::Unknown, + ), + ]; + let r = build_report(&rows, 2000, 86_400, 10); + assert_eq!(r.jobs, 4); + assert_eq!(r.by_job.len(), 2); + assert_eq!( + ( + r.by_job[0].runs, + r.by_job[0].total_secs, + r.by_job[0].avg_secs + ), + (2, 200, 100) + ); + let build = r + .by_step + .iter() + .find(|s| s.job_name == "lint" && s.step == "build") + .unwrap(); + assert_eq!((build.runs, build.total_secs, build.avg_secs), (2, 120, 60)); + assert_eq!(r.queue_wait.len(), 2); + assert_eq!( + ( + r.queue_wait[0].class, + r.queue_wait[0].samples, + r.queue_wait[0].p95_secs + ), + (Class::Light, 2, 10) + ); + assert_eq!(r.cancellation.cancelled, 1); + assert!((r.cancellation.share - 0.25).abs() < f64::EPSILON); + // Worst host first. + assert_eq!(r.drones[0].drone, "d1"); + assert_eq!((r.drones[0].host_failures, r.drones[0].failed), (1, 1)); + assert_eq!((r.drones[1].code_failures, r.drones[1].cancelled), (1, 1)); + let text = render_text(&r); + assert!( + text.contains("Cancellations: 1 of 4 job(s) (25.0%)"), + "{text}" + ); + assert!(text.contains("o/r / lint"), "{text}"); + } + + #[test] + fn limit_keeps_the_biggest_and_an_empty_report_is_zeroed() { + let mut rows = vec![row( + "d", + Class::Light, + "small", + Conclusion::Success, + FailureClass::Unknown, + )]; + let mut big = row( + "d", + Class::Light, + "big", + Conclusion::Success, + FailureClass::Unknown, + ); + big.duration_secs = 900; + rows.push(big); + let r = build_report(&rows, 0, 60, 1); + assert_eq!(r.by_job.len(), 1); + assert_eq!(r.by_job[0].job_name, "big"); + let empty = build_report(&[], 0, 60, 5); + assert_eq!((empty.jobs, empty.cancellation.share), (0, 0.0)); + assert!(render_text(&empty).contains("0 job(s)")); + } + + #[tokio::test] + async fn store_round_trips_replaces_and_prunes() { + let s = Store::open_in_memory().await.unwrap(); + s.enroll("d1", "h", 0).await.unwrap(); + let lease = LeaseRow { + id: "l1".into(), + drone: "d1".into(), + class: Class::Medium, + scope: Scope::Main, + labels: vec![], + runner_name: "r".into(), + runner_id: 1, + jit_config: None, + state: crate::store::LeaseStatus::Done, + created_at: 0, + expires_at: 0, + }; + let t = row( + "d1", + Class::Medium, + "j", + Conclusion::Failure, + FailureClass::Host, + ) + .telemetry; + s.record_telemetry(&lease, &t, 77, 500).await.unwrap(); + s.record_telemetry(&lease, &t, 88, 600).await.unwrap(); + let rows = s.telemetry_since(0, 10).await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!((rows[0].duration_secs, rows[0].finished_at), (88, 600)); + assert_eq!(rows[0].telemetry, t); + assert_eq!((rows[0].class, rows[0].scope), (Class::Medium, Scope::Main)); + assert!(s.telemetry_since(601, 10).await.unwrap().is_empty()); + assert_eq!(s.prune_telemetry(601).await.unwrap(), 1); + assert!(s.telemetry_since(0, 10).await.unwrap().is_empty()); + } +} diff --git a/tests/janitor.rs b/tests/janitor.rs new file mode 100644 index 0000000..d5dc208 --- /dev/null +++ b/tests/janitor.rs @@ -0,0 +1,116 @@ +mod support; + +use serde_json::json; +use support::Harness; + +fn enabled(dry_run: bool, max: u32) -> serde_json::Value { + json!({"supersede_janitor": {"dry_run": dry_run, "max_cancels_per_pass": max}}) +} + +#[tokio::test] +async fn cancels_only_queued_runs_of_superseded_heads_and_is_idempotent() { + let h = Harness::start(enabled(false, 20)).await; + h.queue_pr_run(1, 10, "old", "new"); + h.queue_pr_run(2, 10, "new", "new"); + h.queue_pr_run(3, 11, "old11", "new11"); + + let r = h.janitor_pass().await; + assert_eq!( + (r.considered, r.superseded, r.cancelled), + (3, 2, 2), + "{r:?}" + ); + let mut cancelled = h.state().cancelled.clone(); + cancelled.sort_unstable(); + assert_eq!(cancelled, [1, 3]); + + // A second pass finds nothing left to do. + let r = h.janitor_pass().await; + assert_eq!( + (r.considered, r.superseded, r.cancelled), + (1, 0, 0), + "{r:?}" + ); + assert_eq!(h.state().cancelled.len(), 2); +} + +#[tokio::test] +async fn dry_run_cancels_nothing() { + let h = Harness::start(json!({"supersede_janitor": {}})).await; + h.queue_pr_run(1, 10, "old", "new"); + let r = h.janitor_pass().await; + assert_eq!((r.would_cancel, r.cancelled), (1, 0), "{r:?}"); + assert!(h.state().cancelled.is_empty()); +} + +#[tokio::test] +async fn a_run_that_started_after_listing_is_left_alone() { + let h = Harness::start(enabled(false, 20)).await; + h.queue_pr_run(1, 10, "old", "new"); + h.state().fresh_status.insert(1, "in_progress".into()); + let r = h.janitor_pass().await; + assert_eq!((r.raced, r.cancelled), (1, 0), "{r:?}"); + assert!(h.state().cancelled.is_empty()); +} + +#[tokio::test] +async fn an_already_finished_run_answering_409_is_not_an_error() { + let h = Harness::start(enabled(false, 20)).await; + h.queue_pr_run(1, 10, "old", "new"); + h.state().uncancellable.push(1); + let r = h.janitor_pass().await; + assert_eq!((r.raced, r.cancelled, r.errors), (1, 0, 0), "{r:?}"); +} + +#[tokio::test] +async fn the_per_pass_cap_defers_the_rest_to_later_passes() { + let h = Harness::start(enabled(false, 2)).await; + for id in 1..=5 { + h.queue_pr_run(id, 10 + id, "old", "new"); + } + let r = h.janitor_pass().await; + assert_eq!((r.cancelled, r.deferred), (2, 3), "{r:?}"); + let r = h.janitor_pass().await; + assert_eq!((r.cancelled, r.deferred), (2, 1), "{r:?}"); + let r = h.janitor_pass().await; + assert_eq!((r.cancelled, r.deferred), (1, 0), "{r:?}"); + assert_eq!(h.state().cancelled.len(), 5); +} + +#[tokio::test] +async fn runs_without_exactly_one_pull_request_are_never_touched() { + let h = Harness::start(enabled(false, 20)).await; + { + let mut st = h.state(); + st.pulls.insert(10, "new".into()); + st.pr_runs + .push(json!({"id": 1, "status": "queued", "event": "pull_request", + "head_sha": "old", "pull_requests": []})); + st.pr_runs + .push(json!({"id": 2, "status": "queued", "event": "pull_request", + "head_sha": "old", "pull_requests": [{"number": 10}, {"number": 11}]})); + st.pr_runs + .push(json!({"id": 3, "status": "queued", "event": "merge_group", + "head_sha": "old", "pull_requests": [{"number": 10}]})); + } + let r = h.janitor_pass().await; + assert_eq!((r.considered, r.cancelled), (0, 0), "{r:?}"); +} + +#[tokio::test] +async fn an_unknown_pull_request_is_an_error_not_a_cancellation() { + let h = Harness::start(enabled(false, 20)).await; + h.queue_pr_run(1, 10, "old", "new"); + h.state().pulls.clear(); + let r = h.janitor_pass().await; + assert_eq!((r.errors, r.cancelled), (1, 0), "{r:?}"); +} + +#[tokio::test] +async fn the_janitor_uses_its_own_token() { + // The fake accepts only the janitor token on these endpoints; a pass that + // used the queen's installation token would see 401s and cancel nothing. + let h = Harness::start(enabled(false, 20)).await; + h.queue_pr_run(1, 10, "old", "new"); + assert_eq!(h.janitor_pass().await.cancelled, 1); +} diff --git a/tests/queen.rs b/tests/queen.rs index a34d365..201db54 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -390,3 +390,161 @@ async fn github_failure_before_minting_aborts_the_tick() { assert!(h.queen.tick(OffsetDateTime::now_utc()).await.is_err()); assert!(h.state().mints.is_empty()); } + +#[tokio::test] +async fn warm_schedule_raises_the_floor_only_inside_its_window() { + use time::macros::datetime; + + // The ticks use a fake clock days from the real heartbeat time. + let h = Harness::start( + json!({"drone_stale_secs": 100_000_000, "warm_schedule": {"windows": [ + {"class": "light", "scope": "pr", "count": 2, "start": "08:00", "end": "20:00"} + ]}}), + ) + .await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + + h.queen.tick(datetime!(2026-10-12 07:59 UTC)).await.unwrap(); + assert!(h.state().mints.is_empty()); + assert_eq!(h.status().await["warm_floors"], json!([])); + + h.queen.tick(datetime!(2026-10-12 08:00 UTC)).await.unwrap(); + assert_eq!(h.state().mints.len(), 2); + let metrics = h.metrics().await; + assert!( + metrics.contains("borg_warm_floor{class=\"light\",scope=\"pr\"} 2"), + "{metrics}" + ); + + // The window closing does not revoke leases already minted. + h.queen.tick(datetime!(2026-10-12 20:00 UTC)).await.unwrap(); + assert_eq!(h.state().mints.len(), 2); +} + +async fn lease_for(h: &Harness, token: &str, label: &str, id: u64) -> String { + h.queue_job(id, label, 5); + h.tick().await; + h.leases(token, &[]).await[0]["lease_id"] + .as_str() + .unwrap() + .to_owned() +} + +async fn post_result(h: &Harness, token: &str, lease: &str, body: Value) -> reqwest::StatusCode { + h.http + .post(format!("{}/v1/leases/{lease}/result", h.url)) + .bearer_auth(token) + .json(&body) + .send() + .await + .unwrap() + .status() +} + +#[tokio::test] +async fn telemetry_is_stored_and_reported_per_job_class_and_drone() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let lease = lease_for(&h, &token, "borg-ci-light-pr", 1).await; + + let status = post_result( + &h, + &token, + &lease, + json!({"exit_code": 1, "duration_secs": 120, "telemetry": { + "job_name": "lint", "repo": "example-org/app", "run_id": 9, + "conclusion": "failure", "failure_class": "host", "queue_wait_secs": 25, + "steps": [{"name": "setup", "duration_secs": 20, "conclusion": "success"}, + {"name": "cargo clippy", "duration_secs": 90, "conclusion": "failure"}]}}), + ) + .await; + assert_eq!(status, 204); + + let r: Value = h + .http + .get(format!("{}/report?since_secs=3600", h.url)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(r["jobs"], 1); + assert_eq!(r["by_job"][0]["job_name"], "lint"); + assert_eq!(r["by_job"][0]["total_secs"], 120); + assert_eq!(r["by_step"][0]["step"], "cargo clippy"); + assert_eq!(r["queue_wait"][0]["class"], "light"); + assert_eq!(r["queue_wait"][0]["max_secs"], 25); + assert_eq!(r["drones"][0]["drone"], "d1"); + assert_eq!(r["drones"][0]["host_failures"], 1); + assert_eq!(r["cancellation"]["cancelled"], 0); +} + +#[tokio::test] +async fn results_without_or_with_unusable_telemetry_are_still_accepted() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let lease = lease_for(&h, &token, "borg-ci-light-pr", 1).await; + let status = post_result( + &h, + &token, + &lease, + json!({"exit_code": 0, "duration_secs": 5, "telemetry": {"job_name": "", "repo": "x"}}), + ) + .await; + assert_eq!(status, 204); + assert_eq!( + post_result( + &h, + &token, + &lease, + json!({"exit_code": 0, "duration_secs": 5}) + ) + .await, + 204 + ); + let r: Value = h + .http + .get(format!("{}/report", h.url)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(r["jobs"], 0); +} + +#[tokio::test] +async fn old_telemetry_is_pruned_by_the_tick() { + let h = Harness::start(json!({"telemetry_retention_days": 1})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let lease = lease_for(&h, &token, "borg-ci-light-pr", 1).await; + post_result( + &h, + &token, + &lease, + json!({"exit_code": 0, "duration_secs": 5, + "telemetry": {"job_name": "j", "repo": "o/r", "conclusion": "success"}}), + ) + .await; + assert_eq!( + h.queen.store().telemetry_since(0, 10).await.unwrap().len(), + 1 + ); + // Two days later the one-day retention drops it. + let later = OffsetDateTime::now_utc() + time::Duration::days(2); + h.queen.tick(later).await.unwrap(); + assert!( + h.queen + .store() + .telemetry_since(0, 10) + .await + .unwrap() + .is_empty() + ); +} diff --git a/tests/support/mod.rs b/tests/support/mod.rs index be7188c..89952f0 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -26,6 +26,7 @@ use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; pub const INSTALLATION_TOKEN: &str = "ghs_installation_for_tests"; +pub const JANITOR_TOKEN: &str = "ghp_janitor_for_tests"; pub const ORG: &str = "example-org"; #[derive(Default)] @@ -40,6 +41,16 @@ pub struct MockState { pub jwt_rejections: u32, pub fail_mint: bool, pub fail_runner_list: bool, + /// Queued `pull_request` runs as the janitor lists them. + pub pr_runs: Vec, + /// Current head sha per pull request number. + pub pulls: HashMap, + /// Status the single-run endpoint reports instead of the listed one. + pub fresh_status: HashMap, + /// Runs the janitor cancelled. + pub cancelled: Vec, + /// Runs whose cancel answers 409. + pub uncancellable: Vec, next_runner_id: u64, } @@ -127,11 +138,21 @@ async fn access_token(State(s): State, h: HeaderMap) -> Result Result<(), StatusCode> { + (bearer(headers) == Some(JANITOR_TOKEN)) + .then_some(()) + .ok_or(StatusCode::UNAUTHORIZED) +} + async fn runs( State(s): State, h: HeaderMap, Query(q): Query>, ) -> Result, StatusCode> { + if q.get("event").map(String::as_str) == Some("pull_request") { + check_janitor(&h)?; + return Ok(Json(json!({"workflow_runs": lock(&s).pr_runs.clone()}))); + } check_token(&h)?; let ids: Vec = if q.get("status").map(String::as_str) == Some("queued") { lock(&s) @@ -154,6 +175,56 @@ async fn jobs( Ok(Json(json!({"jobs": lock(&s).jobs.clone()}))) } +async fn one_run( + State(s): State, + h: HeaderMap, + Path((_o, _r, id)): Path<(String, String, u64)>, +) -> Result, StatusCode> { + check_janitor(&h)?; + let st = lock(&s); + let mut run = st + .pr_runs + .iter() + .find(|r| r["id"].as_u64() == Some(id)) + .cloned() + .ok_or(StatusCode::NOT_FOUND)?; + if let Some(status) = st.fresh_status.get(&id) { + run["status"] = json!(status); + } + Ok(Json(run)) +} + +async fn cancel_run( + State(s): State, + h: HeaderMap, + Path((_o, _r, id)): Path<(String, String, u64)>, +) -> StatusCode { + if check_janitor(&h).is_err() { + return StatusCode::UNAUTHORIZED; + } + let mut st = lock(&s); + if st.uncancellable.contains(&id) { + return StatusCode::CONFLICT; + } + st.cancelled.push(id); + st.pr_runs.retain(|r| r["id"].as_u64() != Some(id)); + StatusCode::ACCEPTED +} + +async fn pull( + State(s): State, + h: HeaderMap, + Path((_o, _r, n)): Path<(String, String, u64)>, +) -> Result, StatusCode> { + check_janitor(&h)?; + let sha = lock(&s) + .pulls + .get(&n) + .cloned() + .ok_or(StatusCode::NOT_FOUND)?; + Ok(Json(json!({"state": "open", "head": {"sha": sha}}))) +} + async fn list_runners(State(s): State, h: HeaderMap) -> Result, StatusCode> { check_token(&h)?; let st = lock(&s); @@ -236,6 +307,12 @@ impl Harness { .route("/app/installations/{id}/access_tokens", post(access_token)) .route("/repos/{owner}/{repo}/actions/runs", get(runs)) .route("/repos/{owner}/{repo}/actions/runs/{run}/jobs", get(jobs)) + .route("/repos/{owner}/{repo}/actions/runs/{run}", get(one_run)) + .route( + "/repos/{owner}/{repo}/actions/runs/{run}/cancel", + post(cancel_run), + ) + .route("/repos/{owner}/{repo}/pulls/{n}", get(pull)) .route("/orgs/{org}/actions/runners", get(list_runners)) .route( "/orgs/{org}/actions/runners/generate-jitconfig", @@ -259,7 +336,15 @@ impl Harness { let cfg = QueenConfig::from_json(&cfg.to_string()).unwrap(); let github = GithubClient::new(&gh_url, ORG, "12345", private_pem.as_bytes()).unwrap(); let store = Store::open_in_memory().await.unwrap(); - let queen = Arc::new(Queen::new(cfg, store, github)); + let janitor = cfg.supersede_janitor.clone().map(|settings| { + borg::janitor::Janitor::new(&gh_url, borg::secret::Secret::new(JANITOR_TOKEN), settings) + .unwrap() + }); + let mut queen = Queen::new(cfg, store, github); + if let Some(j) = janitor { + queen = queen.with_janitor(j); + } + let queen = Arc::new(queen); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let url = format!("http://{}", listener.local_addr().unwrap()); @@ -336,6 +421,22 @@ impl Harness { .unwrap() } + /// Add a queued pull-request run on `sha` for pull request `pr`; `head` + /// is the pull request's current head. + pub fn queue_pr_run(&self, id: u64, pr: u64, sha: &str, head: &str) { + let mut st = self.state(); + st.pulls.insert(pr, head.to_owned()); + st.pr_runs.push(json!({ + "id": id, "status": "queued", "event": "pull_request", + "head_sha": sha, "pull_requests": [{"number": pr}] + })); + } + + pub async fn janitor_pass(&self) -> borg::janitor::PassReport { + let repos = self.queen.config().repos.clone(); + self.queen.janitor().unwrap().pass(&repos).await + } + pub async fn tick(&self) { self.queen.tick(OffsetDateTime::now_utc()).await.unwrap(); } From 7e22b640b6731e647abf95a5e415a151abbd077f Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sun, 11 Oct 2026 02:28:40 +0000 Subject: [PATCH 06/27] chore: project borg from Mono 7c0e3ee59b04 --- .repository-projection.json | 6 +- Cargo.toml | 2 +- README.md | 363 ++++++++++ src/admin.rs | 200 ++++++ src/allocator.rs | 480 ++++++++++++- src/api.rs | 901 +++++++++++++++++++++++- src/bid.rs | 339 +++++++++ src/classes.rs | 106 +++ src/config.rs | 95 +++ src/drone/caches/bazel.rs | 45 +- src/drone/caches/gc.rs | 59 +- src/drone/caches/git.rs | 46 +- src/drone/caches/mod.rs | 42 +- src/drone/caches/rustflags.rs | 16 +- src/drone/cleanup.rs | 35 +- src/drone/config.rs | 136 ++++ src/drone/fsops.rs | 196 ++++++ src/drone/install.rs | 577 +++++++++++++-- src/drone/jobuser.rs | 162 +++++ src/drone/mod.rs | 159 ++++- src/drone/runner.rs | 69 +- src/drone/selfupdate.rs | 235 +++++++ src/envcontract.rs | 1024 +++++++++++++++++++++++++++ src/github.rs | 8 + src/health.rs | 833 ++++++++++++++++++++++ src/homeostasis.rs | 511 ++++++++++++++ src/lib.rs | 14 + src/main.rs | 406 ++++++++++- src/profile.rs | 423 +++++++++++ src/queen.rs | 714 ++++++++++++++++++- src/reclaim.rs | 934 +++++++++++++++++++++++++ src/rollout.rs | 851 ++++++++++++++++++++++ src/run.rs | 500 +++++++++++++ src/store.rs | 643 ++++++++++++++++- src/tailnet.rs | 376 ++++++++++ src/units.rs | 1238 +++++++++++++++++++++++++++++++++ src/units_api.rs | 335 +++++++++ src/update/cli.rs | 168 +++++ src/update/config.rs | 171 +++++ src/update/health.rs | 65 ++ src/update/manifest.rs | 260 +++++++ src/update/mod.rs | 266 +++++++ src/update/runner.rs | 185 +++++ src/update/service.rs | 107 +++ src/update/swap.rs | 213 ++++++ src/wake.rs | 533 ++++++++++++++ tests/admin_cli.rs | 113 +++ tests/capacity.rs | 129 ++++ tests/contract.rs | 287 ++++++++ tests/drone.rs | 219 +++++- tests/health.rs | 318 +++++++++ tests/job_user.rs | 94 +++ tests/queen.rs | 141 ++++ tests/rollout.rs | 229 ++++++ tests/run.rs | 449 ++++++++++++ tests/support/mod.rs | 59 +- tests/tailnet.rs | 116 +++ tests/units.rs | 599 ++++++++++++++++ tests/update.rs | 640 +++++++++++++++++ tests/wake.rs | 89 +++ 60 files changed, 18282 insertions(+), 247 deletions(-) create mode 100644 src/admin.rs create mode 100644 src/bid.rs create mode 100644 src/drone/fsops.rs create mode 100644 src/drone/jobuser.rs create mode 100644 src/drone/selfupdate.rs create mode 100644 src/envcontract.rs create mode 100644 src/health.rs create mode 100644 src/homeostasis.rs create mode 100644 src/profile.rs create mode 100644 src/reclaim.rs create mode 100644 src/rollout.rs create mode 100644 src/run.rs create mode 100644 src/tailnet.rs create mode 100644 src/units.rs create mode 100644 src/units_api.rs create mode 100644 src/update/cli.rs create mode 100644 src/update/config.rs create mode 100644 src/update/health.rs create mode 100644 src/update/manifest.rs create mode 100644 src/update/mod.rs create mode 100644 src/update/runner.rs create mode 100644 src/update/service.rs create mode 100644 src/update/swap.rs create mode 100644 src/wake.rs create mode 100644 tests/admin_cli.rs create mode 100644 tests/capacity.rs create mode 100644 tests/contract.rs create mode 100644 tests/health.rs create mode 100644 tests/job_user.rs create mode 100644 tests/rollout.rs create mode 100644 tests/run.rs create mode 100644 tests/tailnet.rs create mode 100644 tests/units.rs create mode 100644 tests/update.rs create mode 100644 tests/wake.rs diff --git a/.repository-projection.json b/.repository-projection.json index ebc7a0d..aa3ad6d 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "ef5e513984edf1a1fbf8d2caeac0d809e6818ab0", + "sourceSha": "7c0e3ee59b046f3f95c09c34cb8e769a031a7fc7", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "5088bfcbe596a86566c86db9a446883b66215ad5", + "priorProjectedBase": "68b3777037221a85f3b0683621dc7a07f1a8e3a1", "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "fbb22932e292fa1fc25b21a331a41d441dafacc2cb8ecc66b11c1d9b238f0250", + "contentDigest": "aa7636d26420d00c729a3d64cc047b6d9ef57bf433e1c8756954b9b2b830da52", "publicationEligible": true } diff --git a/Cargo.toml b/Cargo.toml index 8a868e4..2b40cc4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ sha2 = "0.10" sqlx = { version = "0.8", default-features = false, features = ["runtime-tokio-rustls", "sqlite"] } thiserror = "2.0" time = { version = "0.3", features = ["formatting", "macros", "parsing"] } -tokio = { version = "1.48", features = ["macros", "net", "rt-multi-thread", "process", "signal", "sync", "time"] } +tokio = { version = "1.48", features = ["io-util", "macros", "net", "rt-multi-thread", "process", "signal", "sync", "time"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } diff --git a/README.md b/README.md index dff4939..9aface8 100644 --- a/README.md +++ b/README.md @@ -109,6 +109,294 @@ recorded. The queen keeps it for `telemetry_retention_days` (default 30). `GET /report` summarize time by job and step, queue wait per class, the cancellation share and host-caused failures per drone. +#### What a drone-side parser should produce + +The queen validates but does not interpret these fields, so the drone owns +their meaning. Keep it consistent across drones: + +- **Identity.** `repo`, `run_id` and `job_name` come from the job the runner + executed (its job context: repository, run id, job display name), not from + log text a workflow can print. `job_name` is the name as GitHub shows it, + including matrix suffixes. +- **Steps.** One entry per step in execution order, `duration_secs` rounded + down to whole seconds, `conclusion` mapped from the step result (succeeded to + `success`, failed to `failure`, canceled to `cancelled`, skipped to + `skipped`, anything else to `unknown`). Report names only; never log lines, + environment values or step output. +- **Conclusion.** The job's result, mapped the same way; a job the drone killed + at `max_job_secs` is `timed_out`. +- **`queue_wait_secs`.** Seconds from the job being queued on GitHub to the + runner starting it. Omit it when it cannot be determined; do not send 0 as a + placeholder, because the report averages the values it receives. +- **`failure_class`.** `host` when the machine or the runner failed rather than + the job: the runner process was killed by a signal or the kernel (OOM), the + disk filled, Docker was unreachable, the runner lost its connection, or the + lease was torn down by the drone. `code` when the job's own step exited + non-zero on a healthy host. `unknown` when the drone cannot tell. Leave it + `unknown` for jobs that did not fail. +- **Failure behaviour.** Parsing is best effort and must never delay or fail + the result: if anything cannot be read, send the result without `telemetry` + (or with fewer fields), and send it in the same request as `exit_code` and + `duration_secs`. A retried result may be sent again; the queen keeps the + latest. + +### Drone bids and demand pressure + +A drone may add `bids` to its heartbeat: at most one per `(class, scope)` it is +willing to serve. The queen answers every heartbeat with `pressure`, the +demand pressure (0 to 1000) per class, so the drone can set its bids against it. + +```json +{ "class": "medium", "scope": "pr", "available_slots": 2, + "estimated_start_ms": 1500, "cache_affinity": 800, "resource_pressure": 150, + "estimated_cost": 1000, "confidence": 900, "min_pressure": 0, "expires_at": 1760000000 } +``` + +`cache_affinity`, `resource_pressure`, `confidence` and `min_pressure` are +permille (0 to 1000; larger values are clamped); `estimated_cost` is relative +(1000 is fleet typical); `expires_at` is Unix seconds, after which the queen +ignores the bid. `min_pressure` is the drone's response threshold: it does not +want this work while the class's pressure is below it, so a warm idle host +reports 0, a cold host something higher, and a workstation whose owner is +active something high. Starving demand overrides every threshold. + +The allocator drops a bid with no `available_slots` or above the current +pressure, and scores the rest with integers: +`3*cache_affinity - 3*resource_pressure - 2*start - 1*(cost-1000)`, where +`start` is `estimated_start_ms` as permille of 60 s, all scaled by +`confidence/1000`. The highest score wins; ties go to the tightest memory fit, +then the drone name. A drone that sends no bids, or only expired ones, is +treated as before (always eligible, neutral score); a drone that bids but not +for a key declines that key. GitHub still chooses which queued job a runner +takes: bids steer only which host receives a freshly minted runner. + +### Wake-on-LAN + +Optional (`wake` block, off by default). When demand for a class stays unmet for +`after_secs`, the queen sends a magic packet to a listed host that can serve the +class and whose drone is not fresh, rate limited per pass, per hour and per host. +Targets are IPv4 addresses or `addr/len` networks. See `/status` `wake`. + +### Build environment contract + +Optional (`env_contract.enabled`, off by default). Before a CI runner lease or a +work unit is placed on a drone, the queen checks what the drone last reported +about its build environment. Four checks, each `required`, `advisory` or `off` +globally and per class (`env_contract.checks`, `env_contract.classes.`): + +* `compiler_cache`: `RUSTC_WRAPPER` is an allowed wrapper (`allowed_wrappers`, + by file name: `sccache` by default, plus any single-flight wrapper in front of + it), `SCCACHE_DIR` is set, the sccache server is reachable. +* `toolchain`: when the work declares a Rust toolchain, the drone has it. A unit + declares it (`rust_toolchain`, which `borg run` reads from `rust-toolchain.toml` + or `rust/rust-toolchain.toml` at the unit's tree); a class can declare it in + config (`classes..rust_toolchain`). A mismatch is not admitted. +* `cache_store`: stats work, errors per thousand requests are under + `max_cache_error_pm`, the cache directory has at least + `min_cache_dir_free_bytes` free, and a write works. A hit rate at or under + `hit_rate_warn_pm` after `min_requests_for_hit_rate` requests is a *warning* in + `/status`, never a block. +* `resources`: memory available covers the class's memory (`mem_available_pct`) + and CPU pressure is at most `max_cpu_psi_pm`. (Free capacity after active + leases is the allocator's own rule.) + +A required check that fails **or is unknown** makes the drone ineligible for the +class; nothing else changes about placement. Facts older than +`facts_max_age_secs` (180) count as unknown. A drone that has never reported is +skipped unless `require_report`, so enabling the contract does not cut off +drones that predate it. Defaults when enabled: `toolchain` and `resources` +required, `compiler_cache` and `cache_store` advisory. + +`/status` shows each drone's `contract` (the checks with state, reason, observed +values and warnings, and the admission per class with the reason) and +`units.not_admitted` (why a queued unit is not admitted on each matching drone). +Metrics: `borg_env_contract_check{drone,check,state}` and +`borg_env_contract_admitted{drone,class}`. + +**Drone probe spec** (the heartbeat field `env_facts`). The drone only measures; +verdicts and thresholds are the queen's. Every field is optional, and `null` / +absent means "not measured" (never "fine"). Probe as the **job user**, with the +environment a job will get, and cache results about 60 seconds (resend the cached +object with its original `observed_at`; the queen treats facts over 180 seconds +old as unknown, so probe at least that often): + +| Field | How | +| --- | --- | +| `observed_at` | Unix seconds the probes ran. | +| `rustc_wrapper` | `RUSTC_WRAPPER` as a job sees it, resolved to a path through `PATH` if it is a bare name; absent if unset or empty. | +| `sccache_dir` | `SCCACHE_DIR` as a job sees it; absent if unset. | +| `sccache_server_reachable` | Whether the sccache server answers (connect to `SCCACHE_SERVER_PORT` or its socket, 5 s timeout). Never start a server to find out. | +| `sccache_stats_ok` | `sccache --show-stats --stats-format json` exits 0 within 5 s. | +| `cache_requests`, `cache_hits`, `cache_errors` | From that JSON: `compile_requests`, the sum of the `cache_hits.counts` values and the sum of the `cache_errors.counts` values (check against the installed sccache's schema). | +| `cache_dir_free_bytes` | Available bytes (`f_bavail * f_frsize`) on the filesystem holding `SCCACHE_DIR`. | +| `cache_write_ok` | Create, write and remove a small file named `.borg-probe-` in `SCCACHE_DIR`. | +| `toolchains` | For each toolchain `rustup toolchain list` shows, `{channel, rustc_version}`: the channel as `rust-toolchain.toml` would name it (`1.95.0`, `stable`, `nightly-2026-01-01`; strip the host triple) and the `release:` line of `rustc + -vV`. `null` if rustup could not be queried. | +| `mem_available_mib` | Linux `MemAvailable` from `/proc/meminfo`; macOS (free + inactive + speculative pages) times the page size, in MiB. | +| `cpu_psi_some_avg10_pm` | Linux `/proc/pressure/cpu` line `some`, `avg10` times ten, rounded; absent where PSI does not exist. | + +A unit's toolchain (`rust_toolchain`) matches an installed `channel` or +`rustc_version` exactly. + +### `borg run` + +``` +borg run [--index | --tree ] [--publish] [--class medium] [--wait-secs 60] -- +``` + +Runs a command on the fleet when a drone has capacity, and otherwise right here; +the exit code is the command's either way. With no `BORG_QUEEN` set it just runs +the command locally, so wrappers can call it unconditionally. + +* Environment: `BORG_QUEEN` (queen URL), `BORG_CLIENT_TOKEN` or + `BORG_CLIENT_TOKEN_FILE` (default `~/.config/borg/client-token`; the token comes + from `borg queen enroll-client`), `BORG_RUN_WAIT_SECS`, `BORG_PUBLISH`, + `BORG_TOOLCHAIN`. +* The tree is `HEAD^{tree}` by default, `git write-tree` of the index with + `--index` (pre-commit hooks), or `--tree `. The unit's identity is the repo, + tree, argv, platform and toolchain, so an identical command that already passed + for this tree is not run again (the receipt is reported, exit 0) and one that is + running is joined. +* The drone has to be able to get the tree. A tree from a pushed commit is in + its git mirror. Otherwise `--publish` uploads a `git bundle` to the queen + (`PUT /v1/units/{id}/bundle`): a synthetic commit of the tree whose parent is the + merge base with `origin/main`, bundled without everything that merge base + already reaches, so it is small. Nothing is pushed to any remote. Without + `--publish`, an unpublished tree runs locally. +* Falls back to running here when the queen is unset or unreachable, refuses the + unit, no drone takes it within `--wait-secs` (the unit is cancelled), the queen + is lost for 60 seconds mid-run, or the unit ends for infrastructure reasons + (cancelled, timed out, out of attempts). A command that ran on a drone and + exited nonzero is a real failure and its code is returned. Ctrl-C cancels a + unit this invocation created. +### Work units (gate and hook runs) + +Optional (`units.enabled`, off by default). Clients such as `borg run`, the +pre-commit hook and `mq` submit whole commands; drones run them. Enroll a client +with `borg queen enroll-client ` (prints a token once). + +Client API (bearer: client token): + +* `POST /v1/units` with `{repo, tree_sha, argv, platform, toolchain, class, timeout_secs?, force?}` + returns `201` `{outcome: "new", unit_id, ...}`, or `200` with `outcome: + "joined"` (an identical unit is queued or running) or `"receipt"` (an + identical unit succeeded recently; `receipt` has the proof). The key is + (repo, tree SHA, argv, platform, toolchain); only successes are reused, never + for another tree, and `force` skips reuse. `platform` is `-` + (`linux-x86_64`, `macos-aarch64`). +* `GET /v1/units/{id}` returns the unit (state `queued|leased|running|done|failed`, + `exit_code`, `receipt`). +* `GET /v1/units/{id}/log?from=N&wait=S` returns log bytes from offset `N`, + waiting up to `S` seconds (at most 30) for new bytes while the unit is active. + Headers: `x-borg-next` (the offset to ask for next), `x-borg-state`, and + `x-borg-exit` once finished and fully read. +* `POST /v1/units/{id}/cancel`. +* `PUT /v1/units/{id}/bundle` (the unit's own client; body is a git bundle, at most + `units.max_bundle_bytes`, 64 MiB by default). A unit submitted with + `"bundle": true` is not placed until the bundle arrives, and fails after + `units.bundle_wait_secs` (60) without it. Bundles are deleted with their unit + (7 days, the life of its receipt) or at once if the unit fails or is cancelled. + +Drone protocol (bearer: drone token): + +1. The heartbeat response carries `units: [{unit_id, repo, tree_sha, argv, class, + toolchain, timeout_secs, attempt}]` for units placed on this drone, and + `cancel_units: [id]`. Placed units are resent until acknowledged. +2. Acknowledge by listing the unit id in `active_units` in the heartbeat request; + that starts the unit. Keep listing it while it runs. A running unit the drone + stops listing for 60 seconds is treated as lost and re-queued (up to + `units.max_attempts`). +3. Run it: get `tree_sha` into the local git mirror (if `git cat-file -e + ` fails and the grant has `bundle_bytes > 0`, fetch + `GET /v1/units/{id}/bundle` and `git bundle unbundle` it into the mirror; its + prerequisite is a commit the mirror already has, and if that fails the unit + fails and the client runs locally), check the tree out into a slot directory + (`git clean` it first), run `argv` as the unprivileged job user with the same + isolation rules as CI jobs, and enforce `timeout_secs`. +4. Stream output with `POST /v1/units/{id}/log?offset=N` (raw bytes; `N` is the + byte offset of the first byte of the body). Repeating a chunk is harmless; + `409` means the queen holds fewer than `N` bytes, so resend from its length + (the `200` body of any append is the queen's length). +5. Finish with `POST /v1/units/{id}/result` `{exit_code, duration_secs}`. + Idempotent. Stop and report on a `cancel_units` entry. + +Units are placed by the runner allocator (class capacity, owner caps, bids, +profiles, quarantine and probation) on drones whose `-` matches, using +capacity left after CI runner demand. + +### Build-box reclaim + +A drone on a shared build host can reclaim abandoned git worktrees and cargo +`target/` and `incremental/` directories by disk pressure (`reclaim` block of the +drone config; off by default, dry-run by default). `reclaim::run_cycle` is the +policy; the drone implements `ReclaimHost` and reports `ReclaimStats` in its +heartbeat (`reclaim`), which the queen audits (`/status` `reclaim`). + +What a `ReclaimHost` implementation must guarantee, because the policy relies on +it: `probe` reports **busy** if any process has its working directory or an open +file inside the path (check `/proc/*/cwd` and `fd`, or `lsof +D` on macOS); +**dirty** if `git status --porcelain` is non-empty, untracked files included; +**unpushed** if `git rev-list HEAD --not --remotes` is non-empty (or HEAD is +detached on such commits); **protected** for the main checkout (first entry of +`git worktree list`), a locked worktree, or a repository with stashes; **not_git** +when the path is not a work tree; and **unknown** whenever any step failed or +timed out. `delete` removes the directory (for a worktree, `git worktree remove` +or removal followed by `git worktree prune`) and returns the bytes freed. `scan` +lists each child of a root, and for `cargo_target` roots also each +`incremental/` directory inside (with `within` set to its target directory). + +### Quorum rollout + +`rollout_target_version` plus `rollout.previous_version` turns on a gated rollout +of a new drone version: the queen tells one Linux and one macOS canary the +target through `desired_version` in the heartbeat response, promotes through +waves only while success rate, host-failure rate against the control group and +heartbeat health agree, and halts and tells every drone the previous version +otherwise. State is persisted; see `rollout` in `/status` and +`borg queen rollout halt|reset`. Drones update between jobs when +`desired_version` differs from what they run, and may downgrade. + +### Learned profiles + +The queen builds per-drone, per-class profiles from telemetry: speed against the +fleet median for the same job family (the job name without matrix suffixes) and +reliability from host-caused failures only. They add to the bid score in +proportion to the evidence, so proven faster and more reliable hosts are +preferred when placing a runner. See `profiles` in the queen config and +`profiles[]` in `/status`. + +### Warm-floor controller + +Optional (`homeostasis.enabled`). A damped feedback loop raises the warm floor of +a class and scope when queue latency stays above a target and idle room exists, +and lowers it gradually when demand subsides, with hysteresis, per-key cooldowns, +a step size, an hourly change limit and per-class ceilings. The clock schedule +(`warm_schedule`) remains a minimum. State appears in `/status` and the +`borg_warm_floor*` metrics. + +### Host health + +The queen scores each drone's host-caused failures (from telemetry with +`failure_class: host`) with a half-life decay. Code failures never count. It +recommends quarantine on moderate evidence, and quarantines automatically only +when failures from several distinct runs corroborate, the failures do not look +fleet-wide, and enough other hosts remain in service. A quarantined host returns +after a cooldown, on probation (one lease at a time), once a few jobs succeed. +Tune it with a `health` block in the queen config (`auto_quarantine`, +`half_life_secs`, `window_secs`, `corroboration_runs`, `cooldown_secs`, +`proof_jobs`, and the bounds); `/status` shows each drone's `health`. + +### Tailnet enrollment + +Optional and off by default. With a `tailnet_enroll` block +(`allowed_tags`, `allowed_users`, optional `socket_path` for the local +tailscaled), `POST /v1/enroll` identifies the caller through the local +tailscaled (`whois` by remote address) and, if the node carries an allowed tag +(or, for an untagged node, belongs to an allowed user), enrolls a drone named +after the node and returns `{"name", "token"}` once. The queen must see the +caller's tailnet address directly (bind it to the tailnet address, no proxy in +front), and its service user needs access to the tailscaled socket +(`tailscale set --operator=`). A node that is already enrolled gets 409. + The optional supersede janitor cancels queued pull-request workflow runs whose head commit is no longer the pull request's head. Add a `supersede_janitor` block (`token_path` or `BORG_JANITOR_TOKEN` for a token with `actions: write`, @@ -128,6 +416,9 @@ put a TLS terminator in front. | `POST /v1/heartbeat` | Drone API. Reports capacity; returns new leases until their JIT config is acknowledged. | | `POST /v1/leases/{id}/result` | Drone API. Reports a lease's exit code, duration and peak RSS. | +Manage drones with `borg queen drone list` and `borg queen drone remove ` +(which revokes the token). `borg queen enroll` prints only the token to stdout. + Drone requests carry `Authorization: Bearer `. The queen stores only the SHA-256 of each token. @@ -165,6 +456,34 @@ It sets up: The queen URL must be `https://`; `ca_cert_path` makes that CA the only trust root. A plain `http://` URL is refused unless `allow_insecure` is set. +### Separate job user (Linux) + +Leases run code from pull requests, so they must not run as the account that +holds the enrollment token. `borg drone install` creates two accounts: the +drone daemon (owns `/etc/borg`, mode 0700, with the config and token at 0600) +and an unprivileged job user (`job_user` in the config; default `borg-job`) +that runs every lease and owns the work and cache directories. The daemon +becomes the job user with `setpriv --reuid ... --init-groups`, which drops every +capability before the job starts; the systemd unit grants the daemon only +`CAP_SETUID`, `CAP_SETGID` and `CAP_SETPCAP` for that. File operations on +job-owned directories (cloning the runner, cleanup, cache collection) also run +as the job user. `tests/job_user.rs` checks, as root, that a lease's uid differs +from the drone's, that the token and config directory are unreadable to it, and +that it holds no capability. Without `job_user` (macOS, or `--no-job-user`) +jobs run as the drone's own account and the drone logs a warning. + +### Labels, environment, hooks + +* `extra_labels` are added to every runner the drone registers (for lanes that + still route by label). The queen drops any that equal a job class label, so + they can never change whether a runner serves main or pull-request jobs. +* `runner_env` is a map injected into every lease's environment (keys starting + `BORG_` or `ACTIONS_RUNNER_`, and a few others, are reserved). +* `job_started_hook` and `job_completed_hook` are absolute paths ending in `.sh`, + exported as `ACTIONS_RUNNER_HOOK_JOB_STARTED` and `ACTIONS_RUNNER_HOOK_JOB_COMPLETED`. +* `borg drone run --token-file F` and `borg drone install --token-file F` read + the token from a file instead of the command line. + ### Owner caps `--max-vcpu`, `--max-mem-mib` and `--max-leases` (config keys of the same @@ -315,6 +634,50 @@ On start the drone deletes leftover lease directories and borg-labelled docker resources from a previous run. On `SIGTERM` it stops accepting leases and waits up to `drain_timeout_secs` for running jobs before stopping them. +## Updating a drone + +A drone follows the queen's `desired_version` (see Quorum rollout) between +jobs, in either direction. When the version differs it stops accepting leases, +finishes the jobs it holds, and then asks the updater: it writes `update-request` +into its state directory (`state_dir` in `drone.json`). The drone itself cannot +replace its binary, because it runs unprivileged and stopping it is what drains +it. A root-owned path unit (`borg-update.path` on Linux, the `org.borg.update` +launchd daemon on macOS) starts `borg update --from-request`, which: + +1. fetches the release manifest and downloads every artifact it needs, checking + each SHA-256 against the manifest **before** anything is stopped; +2. stops the drone service (which drains running jobs), swaps the binary + atomically and keeps the previous one as `borg.prev`; +3. starts the service and waits for a fresh `health.json`, which the new drone + writes after each heartbeat the queen answered; +4. rolls back to `borg.prev` if the service does not start or does not report + healthy within `health_timeout_secs`. + +The outcome is left in `update-result.json` and sent to the queen in the next +heartbeat as `update_report`. A target that failed or rolled back on a host is +not retried there until the queen asks for a different version. A queen that is +down is retried with growing delays, and job results are retried for several +minutes, so a restart of the queen loses nothing. + +The updater reads `/etc/borg-update/update.json`, which only root can write, so +a compromised drone cannot point it at another release. The manifest is a JSON +document (`schema`, `version`, per-platform `url` and `sha256` for +`linux-x64-musl`, `linux-arm64-musl` and `darwin-arm64`, and an optional pinned +actions-runner). `https://` is required unless the config sets `allow_insecure`. +There are no signatures: the trust anchor is HTTPS to the configured manifest +URL plus the pinned digests. + +```text +borg update --check # what would change +borg update --version 0.5.0 # install exactly this version +borg update --runner-only # replace just the pinned actions-runner, no restart +``` + +`borg drone install` writes the updater configuration and units; +`--update-manifest-url` points them at another manifest and `--no-self-update` +skips them. Drones installed before this feature have to be updated once by +rerunning `borg drone install` with the new binary. + ## Contributing Issues are welcome here. Code changes land in the source-of-truth repository diff --git a/src/admin.rs b/src/admin.rs new file mode 100644 index 0000000..f6cc959 --- /dev/null +++ b/src/admin.rs @@ -0,0 +1,200 @@ +//! Queen administration: enrolling, listing and removing drones. The CLI is a +//! thin wrapper so the behaviour is testable without a process. + +use crate::enroll; +use crate::store::{DroneRow, Store, StoreError}; + +#[derive(Debug, thiserror::Error)] +pub enum AdminError { + #[error( + "drone name must be 1-40 lowercase letters, digits or dashes, not starting with a dash" + )] + BadName, + #[error("no drone named {0}")] + NoSuchDrone(String), + #[error(transparent)] + Store(#[from] StoreError), +} + +/// Enroll `name` and return its one-time token (the queen keeps only the hash). +pub async fn enroll_drone(store: &Store, name: &str, now: i64) -> Result { + if !enroll::valid_drone_name(name) { + return Err(AdminError::BadName); + } + let token = enroll::new_token(); + store.enroll(name, &enroll::hash_token(&token), now).await?; + Ok(token) +} + +/// Remove a drone. Its token stops working immediately (the hash is deleted) +/// and its lease records are dropped, which lets the queen deregister any +/// runners still offline. Returns how many unfinished leases it had. +pub async fn remove_drone(store: &Store, name: &str) -> Result { + store + .remove_drone(name) + .await? + .ok_or_else(|| AdminError::NoSuchDrone(name.to_owned())) +} + +/// One line per drone: name, state, heartbeat age, host, size and labels. +pub fn list_lines(rows: &[DroneRow], now: i64, stale_secs: i64) -> Vec { + let mut lines = vec![format!( + "{:<24} {:<8} {:>10} {:<24} {:>5} {:>9} {}", + "NAME", "STATE", "HEARTBEAT", "HOST", "VCPU", "MEM(MiB)", "LABELS" + )]; + for d in rows { + let (state, age) = match d.last_heartbeat { + None => ("never", "-".to_owned()), + Some(t) if now - t > stale_secs => ("stale", format!("{}s", now - t)), + Some(t) => ( + if d.accepting { "active" } else { "draining" }, + format!("{}s", now - t), + ), + }; + let cap = d.capacity.as_ref(); + lines.push(format!( + "{:<24} {:<8} {:>10} {:<24} {:>5} {:>9} {}", + d.name, + state, + age, + cap.map_or("-", |c| c.hostname.as_str()), + cap.map_or_else(|| "-".to_owned(), |c| c.vcpu.to_string()), + cap.map_or_else(|| "-".to_owned(), |c| c.mem_mib.to_string()), + d.extra_labels.join(","), + )); + } + lines +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::allocator::HostCaps; + use crate::probe::{Capacity, HostKind}; + + fn capacity() -> Capacity { + Capacity { + hostname: "host-a".into(), + os: "linux".into(), + arch: "x86_64".into(), + kind: HostKind::Vm, + vcpu: 4, + mem_mib: 8192, + free_disk_mib: None, + docker_usable: true, + mise_version: None, + git_version: None, + } + } + + async fn heartbeat(store: &Store, name: &str, accepting: bool, now: i64) { + store + .record_heartbeat(name, &capacity(), &HostCaps::default(), accepting, &[], now) + .await + .unwrap(); + store + .set_extra_labels(name, &["on-prem".to_owned()]) + .await + .unwrap(); + } + + #[tokio::test] + async fn enroll_returns_only_the_token_and_stores_only_its_hash() { + let store = Store::open_in_memory().await.unwrap(); + let token = enroll_drone(&store, "host-a", 1).await.unwrap(); + assert!( + token.starts_with("borg_") + && token.len() == 5 + 64 + && !token.contains(char::is_whitespace) + ); + assert_eq!( + store + .drone_for_token_hash(&enroll::hash_token(&token)) + .await + .unwrap() + .as_deref(), + Some("host-a") + ); + assert_eq!( + store.drone_for_token_hash(&token).await.unwrap(), + None, + "the plaintext is not stored" + ); + assert!(matches!( + enroll_drone(&store, "Bad Name", 1).await, + Err(AdminError::BadName) + )); + assert!( + enroll_drone(&store, "host-a", 2).await.is_err(), + "names are unique" + ); + } + + #[tokio::test] + async fn remove_revokes_the_token_and_drops_unfinished_leases() { + use crate::classes::{Class, Scope}; + use crate::secret::Secret; + use crate::store::{LeaseRow, LeaseStatus}; + + let store = Store::open_in_memory().await.unwrap(); + let token = enroll_drone(&store, "host-a", 1).await.unwrap(); + store + .insert_lease(&LeaseRow { + id: "l1".into(), + drone: "host-a".into(), + class: Class::Light, + scope: Scope::Pr, + labels: vec![], + runner_name: "borg-host-a-l1".into(), + runner_id: 1, + jit_config: Some(Secret::new("jit")), + state: LeaseStatus::Pending, + created_at: 1, + expires_at: 100, + }) + .await + .unwrap(); + + assert_eq!(remove_drone(&store, "host-a").await.unwrap(), 1); + assert_eq!( + store + .drone_for_token_hash(&enroll::hash_token(&token)) + .await + .unwrap(), + None, + "token revoked" + ); + assert!( + store.lease("l1").await.unwrap().is_none(), + "the lease record, and its JIT config, are gone" + ); + assert!(store.active_leases().await.unwrap().is_empty()); + assert!(store.drones().await.unwrap().is_empty()); + assert!(matches!( + remove_drone(&store, "host-a").await, + Err(AdminError::NoSuchDrone(_)) + )); + } + + #[tokio::test] + async fn list_shows_state_from_heartbeat_age_and_acceptance() { + let store = Store::open_in_memory().await.unwrap(); + for n in ["a-never", "b-active", "c-draining", "d-stale"] { + enroll_drone(&store, n, 1).await.unwrap(); + } + heartbeat(&store, "b-active", true, 990).await; + heartbeat(&store, "c-draining", false, 995).await; + heartbeat(&store, "d-stale", true, 100).await; + let lines = list_lines(&store.drones().await.unwrap(), 1000, 60); + assert!(lines[0].starts_with("NAME")); + let find = |n: &str| lines.iter().find(|l| l.starts_with(n)).unwrap().clone(); + assert!(find("a-never").contains("never")); + assert!( + find("b-active").contains("active") + && find("b-active").contains("10s") + && find("b-active").contains("host-a") + ); + assert!(find("c-draining").contains("draining")); + assert!(find("d-stale").contains("stale") && find("d-stale").contains("on-prem")); + } +} diff --git a/src/allocator.rs b/src/allocator.rs index 790f0c8..88f8dff 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -26,6 +26,9 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; +use crate::bid::{ + ClassPressure, DroneBid, ProfileScore, ScoreWeights, demand_pressure, profile_term, score, +}; use crate::classes::{Class, ClassSpec, Classes, Scope}; use crate::probe::Capacity; @@ -62,6 +65,17 @@ pub struct DroneState { /// False while the drone drains (SIGTERM) or yields to its owner. pub accepting: bool, pub max_leases: Option, + /// Live (unexpired) bids. Empty means a legacy drone: always eligible, + /// neutral score. See [`crate::bid`]. + #[serde(default)] + pub bids: Vec, + /// Learned per-class profiles; add to the score, see [`crate::profile`]. + #[serde(default)] + pub profiles: Vec, + /// Classes the build environment contract does not admit on this drone + /// right now; see [`crate::envcontract`]. The allocator never places them. + #[serde(default)] + pub blocked_classes: Vec, } impl DroneState { @@ -77,6 +91,9 @@ impl DroneState { docker: capacity.docker_usable, accepting: true, max_leases: caps.max_leases, + bids: Vec::new(), + profiles: Vec::new(), + blocked_classes: Vec::new(), } } } @@ -146,6 +163,7 @@ pub struct AllocatorInput { /// Demand that has waited at least this long is starving; `None` turns /// the guard off. pub starvation_after_secs: Option, + pub weights: ScoreWeights, } /// An instruction to mint one runner and hand it to one drone. @@ -182,6 +200,17 @@ pub struct Allocation { pub grants: Vec, pub unmet: Vec, pub reserved_for: Vec, + /// Demand pressure per class this pass, to publish to drones. + pub pressure: Vec, + /// Slots each class could still be granted before this pass's grants. + pub idle: Vec, +} + +/// Idle room for one class. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct ClassIdle { + pub class: Class, + pub idle_slots: u32, } /// How many leases of `spec` this drone could hold if it were empty. @@ -195,6 +224,42 @@ pub fn max_slots(drone: &DroneState, spec: &ClassSpec) -> u32 { u32::try_from(by_mem.min(by_cpu).min(by_cap)).unwrap_or(u32::MAX) } +/// What is left of a drone's budget after the leases it holds. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct Room { + pub mem_mib: u64, + pub vcpu: u32, + pub leases_left: Option, +} + +/// The room left on `drone` after `held` (leases on other drones are ignored). +pub fn free_room(drone: &DroneState, held: &[ActiveLease], classes: &Classes) -> Room { + let mut room = Room { + mem_mib: drone.mem_mib, + vcpu: drone.vcpu, + leases_left: drone.max_leases, + }; + for l in held.iter().filter(|l| l.drone_id == drone.id) { + let spec = classes.spec(l.class); + room.mem_mib = room.mem_mib.saturating_sub(spec.mem_mib); + room.vcpu = room.vcpu.saturating_sub(spec.vcpu); + room.leases_left = room.leases_left.map(|n| n.saturating_sub(1)); + } + room +} + +/// Whole leases of `spec` that fit in `room` right now; 0 if the drone takes +/// no new leases or lacks Docker for the class. +pub fn free_slots(drone: &DroneState, room: &Room, spec: &ClassSpec) -> u32 { + if !drone.accepting || (spec.docker && !drone.docker) { + return 0; + } + let by_mem = room.mem_mib / spec.mem_mib.max(1); + let by_cpu = u64::from(room.vcpu / spec.vcpu.max(1)); + let by_cap = room.leases_left.map_or(u64::MAX, u64::from); + u32::try_from(by_mem.min(by_cpu).min(by_cap)).unwrap_or(u32::MAX) +} + /// Remaining room on a drone while allocating. struct Free<'a> { drone: &'a DroneState, @@ -202,9 +267,39 @@ struct Free<'a> { vcpu: u32, leases_left: Option, in_flight: u32, + /// Grants made to this drone this pass, per key (bounds `available_slots`). + granted: BTreeMap<(Class, Scope), u32>, } impl Free<'_> { + /// The score this drone offers for `(class, scope)` at `pressure`, or + /// `None` if it declines. + fn offer( + &self, + class: Class, + scope: Scope, + pressure: u16, + weights: &ScoreWeights, + ) -> Option { + let learned = self + .drone + .profiles + .iter() + .find(|p| p.class == class) + .map_or(0, |p| profile_term(p, weights)); + if self.drone.bids.is_empty() { + return Some(learned); + } + let bid = self + .drone + .bids + .iter() + .find(|b| b.class == class && b.scope == scope)?; + let used = self.granted.get(&(class, scope)).copied().unwrap_or(0); + (used < bid.available_slots && pressure >= bid.min_pressure.min(1000)) + .then(|| score(bid, weights).saturating_add(learned)) + } + fn fits(&self, spec: &ClassSpec) -> bool { self.drone.accepting && (!spec.docker || self.drone.docker) @@ -214,6 +309,14 @@ impl Free<'_> { } /// Could this drone ever hold `spec` if it were empty? + /// Whole leases of `spec` the remaining room holds. + fn slots(&self, spec: &ClassSpec) -> u32 { + let by_mem = self.mem_mib / spec.mem_mib.max(1); + let by_cpu = u64::from(self.vcpu / spec.vcpu.max(1)); + let by_cap = self.leases_left.map_or(u64::MAX, u64::from); + u32::try_from(by_mem.min(by_cpu).min(by_cap)).unwrap_or(u32::MAX) + } + fn could_ever_fit(&self, spec: &ClassSpec) -> bool { self.drone.accepting && (!spec.docker || self.drone.docker) @@ -245,27 +348,36 @@ fn sum_counts>(items: I) -> Counts { fn place( input: &AllocatorInput, free: &mut BTreeMap<&str, Free<'_>>, + pressure: &BTreeMap, reserved: &BTreeMap, (scope, class): (Scope, Class), wanted: u32, out: &mut Vec, ) -> u32 { + let level = pressure.get(&class).copied().unwrap_or(0); let spec = input.classes.spec(class); let mut placed = 0; for _ in 0..wanted { - // Tightest fit, then id: keeps hosts packed and output stable. + // Highest bid score; then tightest fit, which keeps hosts packed; then + // id, which keeps the output stable. let best = free .iter_mut() .filter(|(id, f)| { f.fits(spec) + && !f.drone.blocked_classes.contains(&class) && reserved .get(**id) .is_none_or(|held| *held == (class, scope)) }) - .min_by_key(|(id, f)| (f.mem_mib - spec.mem_mib, **id)) - .map(|(_, f)| f); + .filter_map(|(id, f)| { + let s = f.offer(class, scope, level, &input.weights)?; + Some((s, id, f)) + }) + .min_by_key(|(s, id, f)| (std::cmp::Reverse(*s), f.mem_mib - spec.mem_mib, **id)) + .map(|(_, _, f)| f); let Some(f) = best else { break }; f.take(spec); + *f.granted.entry((class, scope)).or_default() += 1; placed += 1; out.push(Grant { drone_id: f.drone.id.clone(), @@ -291,6 +403,7 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { vcpu: drone.vcpu, leases_left: drone.max_leases, in_flight: 0, + granted: BTreeMap::new(), }); } for lease in &input.leases { @@ -338,20 +451,51 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { } let initial_need = need.clone(); - let mut out = Allocation::default(); + let aged = |class, scope| { + input.starvation_after_secs.is_some_and(|after| { + input + .queue_ages + .iter() + .any(|a| a.class == class && a.scope == scope && a.oldest_age_secs >= after) + }) + }; + let mut pressure: BTreeMap = BTreeMap::new(); + let mut idle_by_class: Vec = Vec::new(); + for class in Class::ALL { + let spec = input.classes.spec(class); + let mut shortfall = 0u32; + let mut starving = false; + for scope in Scope::ALL { + let queued_here = get(&queued, class, scope); + shortfall += queued_here.min(initial_need.get(&(scope, class)).copied().unwrap_or(0)); + starving |= queued_here > 0 && aged(class, scope); + } + let idle = free + .values() + .filter(|f| f.fits(spec)) + .map(|f| f.slots(spec)) + .fold(0u32, u32::saturating_add); + pressure.insert(class, demand_pressure(shortfall, idle, starving)); + idle_by_class.push(ClassIdle { + class, + idle_slots: idle, + }); + } + let mut out = Allocation { + pressure: pressure + .iter() + .map(|(&class, &pressure)| ClassPressure { class, pressure }) + .collect(), + idle: idle_by_class, + ..Allocation::default() + }; let mut placed: BTreeMap<(Scope, Class), u32> = BTreeMap::new(); let mut reserved: BTreeMap = BTreeMap::new(); // Starvation guard. Demand that has waited too long is placed first; if it // still cannot fit, hold back the drone that will free up soonest so the // stream of smaller jobs cannot keep it full. - if let Some(after) = input.starvation_after_secs { - let aged = |class, scope| { - input - .queue_ages - .iter() - .any(|a| a.class == class && a.scope == scope && a.oldest_age_secs >= after) - }; + if input.starvation_after_secs.is_some() { for scope in Scope::ALL { for class in Class::ALL.into_iter().rev() { let wanted = need.get(&(scope, class)).copied().unwrap_or(0); @@ -361,6 +505,7 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { let n = place( input, &mut free, + &pressure, &reserved, (scope, class), wanted, @@ -372,7 +517,11 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { let spec = input.classes.spec(class); let target = free .iter() - .filter(|(id, f)| !reserved.contains_key(**id) && f.could_ever_fit(spec)) + .filter(|(id, f)| { + !reserved.contains_key(**id) + && f.could_ever_fit(spec) + && !f.drone.blocked_classes.contains(&class) + }) .min_by_key(|(id, f)| (f.in_flight, **id)) .map(|(id, _)| (*id).to_owned()); if let Some(id) = target { @@ -394,6 +543,7 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { let n = place( input, &mut free, + &pressure, &reserved, (scope, class), wanted, @@ -434,6 +584,9 @@ mod tests { docker: true, accepting: true, max_leases: None, + bids: vec![], + profiles: vec![], + blocked_classes: vec![], } } @@ -453,6 +606,7 @@ mod tests { warm_floors: vec![], queue_ages: vec![], starvation_after_secs: None, + weights: ScoreWeights::default(), } } @@ -513,10 +667,59 @@ mod tests { } } + #[test] + fn free_room_and_slots_subtract_what_the_drone_already_holds() { + let classes = Classes::default(); + let d = host("a"); // 28 GiB, 16 vCPU, docker + let held = vec![ + lease("a", Class::Heavy, Scope::Pr, LeaseState::Running), + lease("b", Class::Heavy, Scope::Pr, LeaseState::Running), + ]; + let room = free_room(&d, &held, &classes); + let heavy = classes.spec(Class::Heavy); + assert_eq!(room.mem_mib, d.mem_mib - heavy.mem_mib); + assert_eq!(room.vcpu, d.vcpu - heavy.vcpu); + // Theoretical: one heavy fits an empty host; with one held, none do. + let empty = free_room(&d, &[], &classes); + assert_eq!(free_slots(&d, &empty, heavy), max_slots(&d, heavy)); + assert_eq!(free_slots(&d, &room, heavy), 0); + assert!(free_slots(&d, &room, classes.spec(Class::Light)) > 0); + // A draining drone and a docker-less one have no free slots. + let draining = DroneState { + accepting: false, + ..host("c") + }; + assert_eq!(free_slots(&draining, &empty, heavy), 0); + let nodocker = DroneState { + docker: false, + ..host("e") + }; + assert_eq!( + free_slots(&nodocker, &empty, classes.spec(Class::Medium)), + 0 + ); + assert!(free_slots(&nodocker, &empty, classes.spec(Class::Light)) > 0); + // The lease cap counts too. + let capped = DroneState { + max_leases: Some(1), + ..host("f") + }; + let one = vec![lease("f", Class::Light, Scope::Pr, LeaseState::Running)]; + assert_eq!( + free_slots( + &capped, + &free_room(&capped, &one, &classes), + classes.spec(Class::Light) + ), + 0 + ); + } + #[test] fn no_demand_no_floor_no_grants() { let a = allocate(&input(&[], vec![host("a")])); - assert_eq!(a, Allocation::default()); + assert!(a.grants.is_empty() && a.unmet.is_empty() && a.reserved_for.is_empty()); + assert!(a.pressure.iter().all(|p| p.pressure == 0)); } #[test] @@ -1000,6 +1203,9 @@ mod tests { docker: false, accepting: true, max_leases: Some(2), + bids: vec![], + profiles: vec![], + blocked_classes: vec![], }, DroneState { id: "c".into(), @@ -1008,6 +1214,9 @@ mod tests { docker: true, accepting: true, max_leases: None, + bids: vec![], + profiles: vec![], + blocked_classes: vec![], }, ]; for lm in 0..3 { @@ -1035,4 +1244,249 @@ mod tests { } } } + + // ---- bids ----------------------------------------------------------- + + use crate::bid::DroneBid; + + fn bid(class: Class, scope: Scope) -> DroneBid { + DroneBid { + class, + scope, + available_slots: 8, + estimated_start_ms: 0, + cache_affinity: 0, + resource_pressure: 0, + estimated_cost: 1000, + confidence: 1000, + expires_at: i64::MAX, + min_pressure: 0, + } + } + + fn bidder(id: &str, b: DroneBid) -> DroneState { + DroneState { + bids: vec![b], + ..host(id) + } + } + + fn only_drone(a: &Allocation) -> &str { + assert_eq!(a.grants.len(), 1, "{a:?}"); + &a.grants[0].drone_id + } + + #[test] + fn the_warmest_bid_wins_over_the_tightest_fit() { + let mut cold = bid(Class::Light, Scope::Pr); + cold.cache_affinity = 100; + let mut warm = bid(Class::Light, Scope::Pr); + warm.cache_affinity = 900; + // "a" is the emptier host and would lose a pure tightest-fit contest + // only if it were fuller; give "b" a lease so tightest-fit prefers b. + let mut i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![bidder("a", warm), bidder("b", cold)], + ); + i.leases = vec![lease("b", Class::Medium, Scope::Pr, LeaseState::Running)]; + let a = allocate(&i); + assert_eq!(only_drone(&a), "a"); + } + + #[test] + fn loaded_slow_expensive_hosts_lose_and_ties_fall_back_to_fit_then_id() { + let mut loaded = bid(Class::Light, Scope::Pr); + loaded.resource_pressure = 800; + let i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![ + bidder("a", loaded), + bidder("b", bid(Class::Light, Scope::Pr)), + ], + ); + assert_eq!(only_drone(&allocate(&i)), "b"); + + // Equal bids: tightest fit (b has a lease), then id. + let mut i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![ + bidder("a", bid(Class::Light, Scope::Pr)), + bidder("b", bid(Class::Light, Scope::Pr)), + ], + ); + i.leases = vec![lease("b", Class::Light, Scope::Pr, LeaseState::Running)]; + assert_eq!(only_drone(&allocate(&i)), "b"); + i.leases.clear(); + assert_eq!(only_drone(&allocate(&i)), "a"); + } + + #[test] + fn a_bidding_drone_declines_keys_it_did_not_bid_on_and_exhausted_bids() { + let i = input( + &[(Class::Medium, Scope::Pr, 1)], + vec![bidder("a", bid(Class::Light, Scope::Pr))], + ); + assert!(allocate(&i).grants.is_empty()); + + let mut one = bid(Class::Light, Scope::Pr); + one.available_slots = 2; + let i = input(&[(Class::Light, Scope::Pr, 5)], vec![bidder("a", one)]); + let a = allocate(&i); + assert_eq!(count(&a, Class::Light, Scope::Pr), 2); + assert_eq!(a.unmet[0].count, 3); + } + + #[test] + fn legacy_drones_without_bids_stay_eligible_with_a_neutral_score() { + let mut bad = bid(Class::Light, Scope::Pr); + bad.resource_pressure = 900; + // The legacy drone ("a", no bids) beats a heavily loaded bidder. + let i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![host("a"), bidder("b", bad)], + ); + assert_eq!(only_drone(&allocate(&i)), "a"); + } + + #[test] + fn response_thresholds_gate_participation_on_demand_pressure() { + let mut shy = bid(Class::Light, Scope::Pr); + shy.min_pressure = 600; + // One queued job against eight idle slots: low pressure, shy host out. + let i = input(&[(Class::Light, Scope::Pr, 1)], vec![bidder("a", shy)]); + let a = allocate(&i); + assert!(a.grants.is_empty()); + assert!( + a.pressure + .iter() + .any(|p| p.class == Class::Light && p.pressure < 600) + ); + + // Many queued jobs against the same room: pressure crosses the + // threshold and the shy host joins. + let i = input(&[(Class::Light, Scope::Pr, 30)], vec![bidder("a", shy)]); + let a = allocate(&i); + assert!(!a.grants.is_empty()); + assert_eq!( + a.pressure + .iter() + .find(|p| p.class == Class::Light) + .map(|p| p.pressure), + Some(1000) + ); + } + + #[test] + fn starving_demand_overrides_every_threshold() { + let mut shy = bid(Class::Light, Scope::Pr); + shy.min_pressure = 1000; + let mut i = input(&[(Class::Light, Scope::Pr, 1)], vec![bidder("a", shy)]); + assert!(allocate(&i).grants.is_empty()); + aged(&mut i, Class::Light, Scope::Pr, 500); + assert_eq!(only_drone(&allocate(&i)), "a"); + } + + #[test] + fn a_warm_floor_is_placed_only_on_always_willing_hosts() { + let mut shy = bid(Class::Medium, Scope::Main); + shy.min_pressure = 200; + let mut i = input(&[], vec![bidder("a", shy)]); + i.warm_floors = default_warm_floors(); + assert!(allocate(&i).grants.is_empty()); + let eager = bid(Class::Medium, Scope::Main); + i.drones = vec![bidder("b", eager)]; + assert_eq!(only_drone(&allocate(&i)), "b"); + } + + #[test] + fn bid_scoring_is_deterministic_and_never_over_grants() { + let mut drones = Vec::new(); + for (n, id) in ["a", "b", "c", "d"].into_iter().enumerate() { + let mut b = bid(Class::Light, Scope::Pr); + b.cache_affinity = u16::try_from(n * 250).unwrap(); + b.available_slots = 3; + drones.push(bidder(id, b)); + } + let i = input(&[(Class::Light, Scope::Pr, 9)], drones); + let first = allocate(&i); + assert_eq!(first, allocate(&i)); + assert_within_budget(&i, &first); + assert_eq!(count(&first, Class::Light, Scope::Pr), 9); + // The warmest host fills its three slots first. + assert_eq!(first.grants.iter().filter(|g| g.drone_id == "d").count(), 3); + } + + // ---- learned profiles ------------------------------------------------- + + fn profiled(id: &str, class: Class, speed: u32, conf: u32, rel: u32) -> DroneState { + DroneState { + profiles: vec![ProfileScore { + class, + speed_permille: speed, + speed_confidence_permille: conf, + reliability_permille: rel, + }], + ..host(id) + } + } + + #[test] + fn a_proven_faster_host_wins_ties_that_would_go_to_the_lower_id() { + let i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![ + profiled("a", Class::Light, 1000, 0, 1000), + profiled("b", Class::Light, 1600, 1000, 1000), + ], + ); + assert_eq!(only_drone(&allocate(&i)), "b"); + } + + #[test] + fn unreliable_hosts_lose_and_unproven_profiles_change_nothing() { + let i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![ + profiled("a", Class::Light, 1000, 1000, 700), + profiled("b", Class::Light, 1000, 1000, 1000), + ], + ); + assert_eq!(only_drone(&allocate(&i)), "b"); + // Speed with no confidence is ignored: the id tie-break decides. + let i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![ + profiled("a", Class::Light, 1000, 0, 1000), + profiled("b", Class::Light, 2000, 0, 1000), + ], + ); + assert_eq!(only_drone(&allocate(&i)), "a"); + } + + #[test] + fn profiles_are_per_class_and_stack_with_bids() { + // The profile for medium does not help a light placement. + let i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![ + profiled("a", Class::Light, 1000, 0, 1000), + profiled("b", Class::Medium, 2000, 1000, 1000), + ], + ); + assert_eq!(only_drone(&allocate(&i)), "a"); + // A bidding drone's score includes its profile. + let mut warm = bid(Class::Light, Scope::Pr); + warm.cache_affinity = 100; + let fast = DroneState { + bids: vec![bid(Class::Light, Scope::Pr)], + ..profiled("b", Class::Light, 2000, 1000, 1000) + }; + let i = input( + &[(Class::Light, Scope::Pr, 1)], + vec![bidder("a", warm), fast], + ); + // a: 3*100 = 300; b: 1000 from the profile. + assert_eq!(only_drone(&allocate(&i)), "b"); + assert_eq!(allocate(&i), allocate(&i)); + } } diff --git a/src/api.rs b/src/api.rs index 3b48fe7..953d101 100644 --- a/src/api.rs +++ b/src/api.rs @@ -11,7 +11,7 @@ use std::collections::BTreeMap; use std::fmt::Write as _; use std::sync::Arc; -use axum::extract::{Path, Query, State}; +use axum::extract::{ConnectInfo, Path, Query, State}; use axum::http::{HeaderMap, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; @@ -19,14 +19,25 @@ use axum::{Json, Router}; use serde::{Deserialize, Serialize}; use time::OffsetDateTime; -use crate::allocator::{DroneState, HostCaps, Reservation, Unmet}; +use crate::allocator::{ + ActiveLease, HostCaps, LeaseState, Reservation, Unmet, free_room, free_slots, max_slots, +}; +use crate::bid::{ClassPressure, DroneBid}; use crate::classes::{Class, Scope}; use crate::enroll::hash_token; +use crate::envcontract::{CheckState, EnvFacts, Need, Verdict, blocked_classes, evaluate}; +use crate::health::{DroneHealth, HealthState}; +use crate::homeostasis::KeyReport; use crate::probe::Capacity; +use crate::queen::RolloutView; use crate::queen::{Queen, QueenError, QueueEntry}; +use crate::reclaim::ReclaimStats; +use crate::rollout::Phase as RolloutPhase; use crate::secret::Secret; use crate::store::{LeaseResult, LeaseStatus, StoreError}; +use crate::tailnet::drone_name_for; use crate::telemetry::{JobTelemetry, MAX_REPORT_ROWS, Report, build_report}; +use crate::units::UnitGrant; #[derive(Debug, Serialize, Deserialize)] pub struct HeartbeatRequest { @@ -42,9 +53,63 @@ pub struct HeartbeatRequest { /// Size and garbage-collection state of the drone's host caches. #[serde(default)] pub cache_stats: Vec, + /// Labels the drone asks to have added to its runners. + #[serde(default)] + pub extra_labels: Vec, + /// Outcome of an update this drone just ran. The queen may use it as + /// rollout evidence or ignore it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub update_report: Option, /// The borg version the drone runs (reported for rollout tracking). #[serde(default)] pub version: Option, + /// One bid per `(class, scope)` the drone is willing to serve; see + /// [`crate::bid`]. Empty for drones that do not bid. + #[serde(default)] + pub bids: Vec, + /// Build-box reclaim stats and new audit entries; see [`crate::reclaim`]. + #[serde(default)] + pub reclaim: Option, + /// Work units the drone currently holds; acknowledges new ones. + #[serde(default)] + pub active_units: Vec, + /// Facts about the drone's build environment; see [`crate::envcontract`]. + #[serde(default)] + pub env_facts: Option, +} + +/// Bound the free text a drone reports; the observation time cannot be in the +/// future. +fn sanitize_facts(f: &EnvFacts, now: i64) -> EnvFacts { + let clip = |s: &String| { + s.chars() + .filter(|c| !c.is_control()) + .take(300) + .collect::() + }; + let mut out = f.clone(); + out.observed_at = out.observed_at.min(now); + out.rustc_wrapper = out.rustc_wrapper.as_ref().map(clip); + out.sccache_dir = out.sccache_dir.as_ref().map(clip); + if let Some(list) = &mut out.toolchains { + list.truncate(32); + for t in list { + t.channel = clip(&t.channel); + t.rustc_version = clip(&t.rustc_version); + } + } + out +} + +/// At most one bid per key (the first), clamped, at most one per key overall. +fn sanitize_bids(bids: &[DroneBid]) -> Vec { + let mut out: Vec = Vec::new(); + for b in bids { + if !out.iter().any(|o| o.class == b.class && o.scope == b.scope) { + out.push(b.clamped()); + } + } + out } /// A drone-reported version is shown in metric labels, so accept only short @@ -67,6 +132,9 @@ pub struct CacheStat { pub misses: Option, } +/// A running unit the drone stops listing is considered lost after this. +const LOST_UNIT_GRACE_SECS: i64 = 60; + fn yes() -> bool { true } @@ -82,9 +150,33 @@ pub struct LeaseGrant { pub expires_at: i64, } +/// A drone's report of an update it ran. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct UpdateReport { + /// The version the update aimed for. + pub target: String, + /// `installed`, `runner_updated`, `up_to_date`, `rolled_back` or `failed`. + pub state: String, + pub detail: String, +} + #[derive(Debug, Serialize, Deserialize)] pub struct HeartbeatResponse { pub leases: Vec, + /// Work units to run (resent until acknowledged in `active_units`). + #[serde(default)] + pub units: Vec, + /// Units the client cancelled; stop them. + #[serde(default)] + pub cancel_units: Vec, + /// The borg version this drone should run, from the rollout gate (or the + /// configured target when the rollout is not gated). Absent: no + /// instruction, keep the current version. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub desired_version: Option, + /// Demand pressure per class, for the drone to set its bids against. + #[serde(default)] + pub pressure: Vec, } #[derive(Debug, Serialize, Deserialize)] @@ -100,8 +192,14 @@ pub struct ResultRequest { #[derive(Debug)] pub enum ApiError { + BadRequest(String), + TooMany, + TooLarge, Unauthorized, NotFound, + Forbidden, + Conflict, + Unprocessable, Internal(String), } @@ -120,8 +218,14 @@ impl From for ApiError { impl IntoResponse for ApiError { fn into_response(self) -> Response { match self { + Self::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg).into_response(), + Self::TooMany => StatusCode::TOO_MANY_REQUESTS.into_response(), + Self::TooLarge => StatusCode::PAYLOAD_TOO_LARGE.into_response(), Self::Unauthorized => StatusCode::UNAUTHORIZED.into_response(), Self::NotFound => StatusCode::NOT_FOUND.into_response(), + Self::Forbidden => StatusCode::FORBIDDEN.into_response(), + Self::Conflict => StatusCode::CONFLICT.into_response(), + Self::Unprocessable => StatusCode::UNPROCESSABLE_ENTITY.into_response(), Self::Internal(msg) => { tracing::error!("api error: {msg}"); StatusCode::INTERNAL_SERVER_ERROR.into_response() @@ -135,13 +239,15 @@ pub fn router(queen: Arc) -> Router { .route("/status", get(status)) .route("/metrics", get(metrics)) .route("/report", get(report)) + .route("/v1/enroll", post(tailnet_enroll)) .route("/v1/heartbeat", post(heartbeat)) .route("/v1/leases/{id}/result", post(lease_result)) + .merge(crate::units_api::routes()) .with_state(queen) } /// The drone name behind the request's bearer token. -async fn authenticate(queen: &Queen, headers: &HeaderMap) -> Result { +pub(crate) async fn authenticate(queen: &Queen, headers: &HeaderMap) -> Result { let presented = headers .get("authorization") .and_then(|v| v.to_str().ok()) @@ -154,6 +260,67 @@ async fn authenticate(queen: &Queen, headers: &HeaderMap) -> Result>, + ConnectInfo(peer): ConnectInfo, +) -> Result, ApiError> { + let Some(enroll) = queen.tailnet_enroll() else { + return Err(ApiError::NotFound); + }; + let identity = match enroll.verifier.whois(peer).await { + Ok(Some(id)) if enroll.settings.allows(&id) => id, + Ok(found) => { + tracing::warn!( + "tailnet enroll refused for {peer}: {}", + found.map_or_else( + || "not a tailnet node".to_owned(), + |id| format!("{} (tags {:?}) is not allowlisted", id.node_name, id.tags) + ) + ); + return Err(ApiError::Forbidden); + } + Err(e) => { + tracing::warn!("tailnet enroll: whois for {peer} failed: {e}"); + return Err(ApiError::Internal(format!("whois: {e}"))); + } + }; + let Some(name) = drone_name_for(&identity.node_name) else { + return Err(ApiError::Unprocessable); + }; + let token = crate::enroll::new_token(); + let now = OffsetDateTime::now_utc().unix_timestamp(); + match queen + .store() + .enroll(&name, &crate::enroll::hash_token(&token), now) + .await + { + Ok(()) => { + tracing::info!( + "tailnet enroll: enrolled drone {name} for {} (tags {:?}, user {:?}) from {peer}", + identity.node_name, + identity.tags, + identity.user + ); + Ok(Json(EnrollResponse { + name, + token: Secret::new(token), + })) + } + Err(StoreError::AlreadyEnrolled(_)) => Err(ApiError::Conflict), + Err(e) => Err(e.into()), + } +} + async fn heartbeat( State(queen): State>, headers: HeaderMap, @@ -173,9 +340,50 @@ async fn heartbeat( ) .await?; store - .set_version(&drone, sanitize_version(req.version.as_deref())) + .set_version(&drone, sanitize_version(req.version.as_deref()), now) .await?; + store.set_extra_labels(&drone, &req.extra_labels).await?; + store.set_bids(&drone, &sanitize_bids(&req.bids)).await?; + if let Some(facts) = &req.env_facts { + store + .set_env_facts(&drone, &sanitize_facts(facts, now)) + .await?; + } + if let Some(stats) = &req.reclaim { + store.record_reclaim(&drone, stats).await?; + } store.acknowledge(&drone, &req.active_leases).await?; + let (units, cancel_units) = if queen.config().units.enabled { + store.ack_units(&drone, &req.active_units, now).await?; + store + .requeue_lost_units( + &drone, + &req.active_units, + now - LOST_UNIT_GRACE_SECS, + queen.config().units.max_attempts, + now, + ) + .await?; + let grants = store + .unacked_units(&drone) + .await? + .into_iter() + .map(|u| UnitGrant { + unit_id: u.id, + repo: u.repo, + tree_sha: u.tree_sha, + argv: u.argv, + class: u.class, + toolchain: u.toolchain, + timeout_secs: u.timeout_secs, + attempt: u.attempt, + bundle_bytes: u.bundle_bytes, + }) + .collect(); + (grants, store.cancel_requests(&drone).await?) + } else { + (Vec::new(), Vec::new()) + }; let leases = store .undelivered(&drone) .await? @@ -192,7 +400,13 @@ async fn heartbeat( }) }) .collect(); - Ok(Json(HeartbeatResponse { leases })) + Ok(Json(HeartbeatResponse { + leases, + units, + cancel_units, + desired_version: queen.desired_version(&drone), + pressure: queen.snapshot().pressure, + })) } async fn lease_result( @@ -238,12 +452,35 @@ pub struct ClassStatus { pub queued_pr: u32, pub oldest_age_secs_main: u64, pub oldest_age_secs_pr: u64, - /// Leases of this class the fleet could hold if empty. + /// Leases of this class the fleet could hold if it were empty (theoretical + /// capacity, not what is free). pub slots: u32, + /// Leases of this class that fit right now: capacity minus every active + /// lease, on drones that are fresh and take new leases. + pub free_slots: u32, + pub demand_main: ScopeDemand, + pub demand_pr: ScopeDemand, pub leases_pending: u32, pub leases_running: u32, } +/// Queued demand for one label, split by who can serve it. The split is an +/// estimate: GitHub, not borg, picks which runner takes a job, and a queued job +/// carries no runner yet, so borg counts the idle online runners it did not mint +/// that carry the job's label. +#[derive(Debug, Default, Serialize)] +pub struct ScopeDemand { + pub queued: u32, + /// Online runners not minted by borg with this label. + pub static_online: u32, + /// Of those, the ones not running a job. + pub static_idle: u32, + /// `min(queued, static_idle)`: likely taken by those runners. + pub claimable_by_static_estimate: u32, + /// `queued - claimable_by_static_estimate`: what borg has to place. + pub borg_placeable: u32, +} + #[derive(Debug, Serialize)] pub struct DroneStatus { pub name: String, @@ -253,12 +490,48 @@ pub struct DroneStatus { pub hostname: Option, pub kind: Option, pub vcpu: Option, + /// Raw host memory and vCPU as probed. pub mem_mib: Option, + /// Memory and vCPU borg may hand out (host less the OS reserve, under the + /// owner's caps): the theoretical budget. + pub budget_mem_mib: Option, + pub budget_vcpu: Option, + /// Budget minus the leases the drone holds right now. + pub free_mem_mib: Option, + pub free_vcpu: Option, + /// Leases per class that fit right now / that would fit if it were empty. + pub free_slots: BTreeMap, + pub theoretical_slots: BTreeMap, pub active_leases: u32, pub caches: Vec, pub quarantined: bool, pub quarantine_reason: Option, pub version: Option, + /// Evidence-based host health; absent until the first assessment. + pub health: Option, + /// Build-box reclaim stats, when the drone reports them. + pub reclaim: Option, + /// The build environment contract as it stands for this drone; absent when + /// the contract is off. + pub contract: Option, +} + +/// One drone's build environment contract. +#[derive(Debug, Serialize)] +pub struct ContractStatus { + pub reported: bool, + pub age_secs: Option, + /// The checks as evaluated for the lightest class (the least demanding + /// resources need); the admission per class is below. + pub checks: Vec, + /// Whether the drone is admitted for each class, and why not. + pub admission: BTreeMap, +} + +#[derive(Debug, Serialize)] +pub struct ClassAdmission { + pub admit: bool, + pub why: Option, } /// How many enrolled drones run the configured target version. @@ -271,6 +544,10 @@ pub struct RolloutStatus { pub outdated: u32, /// Drones that report no version. pub unknown: u32, + /// The quorum rollout gate, when `rollout.previous_version` is set. + pub gate: Option, + /// Recent rollout transitions, newest first. + pub log: Vec<(i64, String)>, } #[derive(Debug, Serialize)] @@ -281,11 +558,124 @@ pub struct Status { pub classes: BTreeMap, pub queue: Vec, pub drones: Vec, + /// Demand the allocator could not place. Counts jobs that non-borg runners + /// may serve; see `unmet_placeable`. pub unmet: Vec, + /// `unmet` less the idle non-borg runners carrying each label (an estimate). + pub unmet_placeable: Vec, /// Drones held back for starving demand. pub reserved_for: Vec, pub rollout: RolloutStatus, pub warm_floors: Vec, + pub pressure: Vec, + /// Host failures look fleet-wide; automatic quarantine is paused. + pub health_systemic: bool, + pub homeostasis: Vec, + pub profiles: Vec, + pub wake: crate::wake::WakeView, + /// Build-box reclaim: bytes ever really deleted and the newest audit + /// entries. + pub reclaim: ReclaimSummary, + pub units: UnitsStatus, + /// The newest quarantine and health audit entries. + pub health_events: Vec, +} + +#[derive(Debug, Serialize)] +pub struct ReclaimSummary { + pub reclaimed_bytes_logged: u64, + pub events: Vec, +} + +#[derive(Debug, Serialize)] +pub struct UnitsStatus { + pub enabled: bool, + pub counts: std::collections::BTreeMap<&'static str, u32>, + /// Submissions that joined a running unit since the queen started. + pub joins: u64, + /// Submissions answered with a stored receipt since the queen started. + pub receipt_hits: u64, + pub oldest_queued_age_secs: u64, + pub recent: Vec, + /// Why queued units are not admitted on drones that otherwise match them + /// (platform, fresh): the build environment contract's verdict. + pub not_admitted: Vec, +} + +#[derive(Debug, Serialize)] +pub struct UnitBlock { + pub unit_id: String, + pub drone: String, + pub why: String, +} + +fn scope_demand( + snap: &crate::queen::Snapshot, + cfg: &crate::config::QueenConfig, + class: Class, + scope: Scope, +) -> ScopeDemand { + let queued = snap + .queue + .iter() + .find(|q| q.class == class && q.scope == scope) + .map_or(0, |q| q.queued); + let statics = snap + .static_runners + .get(&cfg.classes.job_label(class, scope)) + .copied() + .unwrap_or_default(); + let claimable = queued.min(statics.idle); + ScopeDemand { + queued, + static_online: statics.online, + static_idle: statics.idle, + claimable_by_static_estimate: claimable, + borg_placeable: queued - claimable, + } +} + +/// For each queued unit, the drones of its platform that the contract does not +/// admit for it, and why (at most 50 entries). +fn unit_blocks( + cfg: &crate::config::QueenConfig, + rows: &[crate::store::DroneRow], + queued: &[crate::units::Unit], + now: i64, +) -> Vec { + if !cfg.env_contract.enabled { + return Vec::new(); + } + let mut out = Vec::new(); + for u in queued.iter().filter(|u| !u.awaiting_bundle) { + for d in rows { + let Some(cap) = &d.capacity else { + continue; + }; + let fresh = d + .last_heartbeat + .is_some_and(|t| now - t <= cfg.drone_stale_secs); + if !fresh || crate::queen::platform_of(cap) != u.platform { + continue; + } + let need = Need::for_class( + &cfg.env_contract, + u.class, + cfg.classes.spec(u.class), + u.rust_toolchain.as_deref(), + ); + let v = evaluate(&cfg.env_contract, d.env_facts.as_ref(), &need, now); + if let (false, Some(why)) = (v.admit, v.why) { + out.push(UnitBlock { + unit_id: u.id.clone(), + drone: d.name.clone(), + why, + }); + } + } + } + out.truncate(50); + out } pub async fn build_status(queen: &Queen, now: i64) -> Result { @@ -294,15 +684,72 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result let rows = queen.store().drones().await?; let active = queen.store().active_leases().await?; + let mut held: Vec = active + .iter() + .map(|l| ActiveLease { + drone_id: l.drone.clone(), + class: l.class, + scope: l.scope, + state: LeaseState::Running, + }) + .collect(); + // Placed work units hold capacity like running leases. + if cfg.units.enabled { + use crate::units::UnitState; + for u in queen + .store() + .units_in(&[UnitState::Leased, UnitState::Running]) + .await? + { + if let Some(drone_id) = u.drone { + held.push(ActiveLease { + drone_id, + class: u.class, + scope: Scope::Pr, + state: LeaseState::Running, + }); + } + } + } let mut states = Vec::new(); let mut drones = Vec::new(); + let mut free_by_class: BTreeMap = BTreeMap::new(); for d in &rows { let age = d.last_heartbeat.map(|t| now - t); let fresh = age.is_some_and(|a| a <= cfg.drone_stale_secs); + let mut budget = None; + let mut free_slots_here = BTreeMap::new(); + let mut theoretical_here = BTreeMap::new(); if let (true, Some(cap)) = (fresh, &d.capacity) { - let mut s = DroneState::from_capacity(d.name.clone(), cap, &d.caps); - s.accepting = d.schedulable(); - states.push(s); + let mut state = d.allocator_state(cap, now); + state.blocked_classes = blocked_classes( + &cfg.env_contract, + d.env_facts.as_ref(), + &cfg.classes, + None, + now, + ); + let room = free_room(&state, &held, &cfg.classes); + for class in Class::ALL { + let spec = cfg.classes.spec(class); + let free = if state.blocked_classes.contains(&class) { + 0 + } else { + free_slots(&state, &room, spec) + }; + free_slots_here.insert(class, free); + theoretical_here.insert( + class, + if state.accepting { + max_slots(&state, spec) + } else { + 0 + }, + ); + *free_by_class.entry(class).or_default() += free; + } + budget = Some((state.mem_mib, state.vcpu, room)); + states.push(state); } drones.push(DroneStatus { name: d.name.clone(), @@ -313,12 +760,46 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result kind: d.capacity.as_ref().map(|c| c.kind), vcpu: d.capacity.as_ref().map(|c| c.vcpu), mem_mib: d.capacity.as_ref().map(|c| c.mem_mib), + budget_mem_mib: budget.map(|b| b.0), + budget_vcpu: budget.map(|b| b.1), + free_mem_mib: budget.map(|b| b.2.mem_mib), + free_vcpu: budget.map(|b| b.2.vcpu), + free_slots: free_slots_here, + theoretical_slots: theoretical_here, active_leases: u32::try_from(active.iter().filter(|l| l.drone == d.name).count()) .unwrap_or(u32::MAX), caches: d.cache_stats.clone(), quarantined: d.quarantined, quarantine_reason: d.quarantine_reason.clone(), version: d.version.clone(), + health: snap.health.get(&d.name).cloned(), + reclaim: d.reclaim.clone(), + contract: cfg.env_contract.enabled.then(|| { + let verdict = |class| { + let need = + Need::for_class(&cfg.env_contract, class, cfg.classes.spec(class), None); + evaluate(&cfg.env_contract, d.env_facts.as_ref(), &need, now) + }; + let light: Verdict = verdict(Class::Light); + ContractStatus { + reported: light.reported, + age_secs: light.age_secs, + checks: light.checks.clone(), + admission: Class::ALL + .into_iter() + .map(|c| { + let v = verdict(c); + ( + c, + ClassAdmission { + admit: v.admit, + why: v.why, + }, + ) + }) + .collect(), + } + }), }); } @@ -347,6 +828,9 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result oldest_age_secs_main: entry(Scope::Main).map_or(0, |q| q.oldest_age_secs), oldest_age_secs_pr: entry(Scope::Pr).map_or(0, |q| q.oldest_age_secs), slots: slots.get(&class).copied().unwrap_or(0), + free_slots: free_by_class.get(&class).copied().unwrap_or(0), + demand_main: scope_demand(&snap, cfg, class, Scope::Main), + demand_pr: scope_demand(&snap, cfg, class, Scope::Pr), leases_pending: count(LeaseStatus::Pending), leases_running: count(LeaseStatus::Running), }, @@ -354,6 +838,8 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result } let mut rollout = RolloutStatus { target_version: cfg.rollout_target_version.clone(), + gate: snap.rollout.clone(), + log: queen.store().rollout_log(20).await?, ..RolloutStatus::default() }; for d in &drones { @@ -371,9 +857,36 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result queue: snap.queue, drones, unmet: snap.unmet, + unmet_placeable: snap.unmet_placeable, reserved_for: snap.reserved_for, rollout, warm_floors: snap.warm_floors, + pressure: snap.pressure, + health_systemic: snap.health_systemic, + homeostasis: snap.homeostasis, + profiles: snap.profiles, + wake: snap.wake, + units: { + let (joins, receipt_hits) = queen.unit_reuse_counts(); + let store = queen.store(); + let queued = store.units_in(&[crate::units::UnitState::Queued]).await?; + UnitsStatus { + enabled: cfg.units.enabled, + counts: store.unit_counts().await?, + joins, + receipt_hits, + oldest_queued_age_secs: queued + .first() + .map_or(0, |u| u64::try_from(now - u.created_at).unwrap_or(0)), + recent: store.recent_units(20).await?, + not_admitted: unit_blocks(cfg, &rows, &queued, now), + } + }, + reclaim: ReclaimSummary { + reclaimed_bytes_logged: queen.store().reclaim_bytes_logged().await?, + events: queen.store().reclaim_log(20).await?, + }, + health_events: queen.store().health_log(20).await?, }) } @@ -485,6 +998,109 @@ pub fn render_metrics(status: &Status) -> String { f64::from(s.leases_running), )); } + gauge( + "borg_env_contract_check", + "Build environment contract checks per drone (evaluated for the lightest class): 1 for the current state.", + status + .drones + .iter() + .filter_map(|d| d.contract.as_ref().map(|c| (d, c))) + .flat_map(|(d, c)| { + c.checks.iter().flat_map(move |o| { + CheckState::ALL.into_iter().map(move |st| { + ( + format!( + "drone=\"{}\",check=\"{}\",state=\"{}\"", + label_value(&d.name), + o.kind.as_str(), + st.as_str() + ), + f64::from(u8::from(o.state == st)), + ) + }) + }) + }) + .collect(), + ); + gauge( + "borg_env_contract_admitted", + "1 if the build environment contract admits the class on the drone.", + status + .drones + .iter() + .filter_map(|d| d.contract.as_ref().map(|c| (d, c))) + .flat_map(|(d, c)| { + c.admission.iter().map(move |(class, a)| { + ( + format!("drone=\"{}\",class=\"{class}\"", label_value(&d.name)), + f64::from(u8::from(a.admit)), + ) + }) + }) + .collect(), + ); + gauge( + "borg_class_free_slots", + "Leases of a class that fit right now (capacity minus active leases).", + status + .classes + .iter() + .map(|(c, s)| (format!("class=\"{c}\""), f64::from(s.free_slots))) + .collect(), + ); + gauge( + "borg_drone_free_mem_mib", + "Memory budget left on a drone after its active leases.", + status + .drones + .iter() + .filter_map(|d| { + d.free_mem_mib.map(|m| { + #[allow(clippy::cast_precision_loss, reason = "MiB far below 2^52")] + let v = m as f64; + (format!("drone=\"{}\"", label_value(&d.name)), v) + }) + }) + .collect(), + ); + gauge( + "borg_drone_free_slots", + "Leases of a class that fit on a drone right now.", + status + .drones + .iter() + .flat_map(|d| { + d.free_slots.iter().map(move |(c, n)| { + ( + format!("drone=\"{}\",class=\"{c}\"", label_value(&d.name)), + f64::from(*n), + ) + }) + }) + .collect(), + ); + gauge( + "borg_queued_borg_placeable", + "Queued jobs borg has to place: queued minus idle non-borg runners with the label (estimate).", + both(&|s, sc| { + f64::from(if sc == Scope::Main { + s.demand_main.borg_placeable + } else { + s.demand_pr.borg_placeable + }) + }), + ); + gauge( + "borg_queued_claimable_by_static", + "Queued jobs idle non-borg runners with the label can take (estimate).", + both(&|s, sc| { + f64::from(if sc == Scope::Main { + s.demand_main.claimable_by_static_estimate + } else { + s.demand_pr.claimable_by_static_estimate + }) + }), + ); gauge( "borg_active_leases", "Active leases per class and state.", @@ -520,6 +1136,48 @@ pub fn render_metrics(status: &Status) -> String { }) .collect(), ); + gauge( + "borg_drone_health_state", + "1 for the drone's current evidence-based health state (healthy, recommended, quarantined, probation).", + status + .drones + .iter() + .filter_map(|d| d.health.as_ref().map(|h| (d, h))) + .flat_map(|(d, h)| { + HealthState::ALL.into_iter().map(move |st| { + ( + format!( + "drone=\"{}\",state=\"{}\"", + label_value(&d.name), + st.as_str() + ), + f64::from(u8::from(h.state == st)), + ) + }) + }) + .collect(), + ); + gauge( + "borg_drone_host_failure_score", + "Decayed host-caused failure evidence per drone (1 = one fresh failure).", + status + .drones + .iter() + .filter_map(|d| { + d.health.as_ref().map(|h| { + ( + format!("drone=\"{}\"", label_value(&d.name)), + f64::from(h.score_permille) / 1000.0, + ) + }) + }) + .collect(), + ); + gauge( + "borg_health_systemic", + "1 while host failures look fleet-wide and automatic quarantine is paused.", + vec![(String::new(), f64::from(u8::from(status.health_systemic)))], + ); gauge( "borg_drone_info", "The borg version each drone reported (always 1).", @@ -550,6 +1208,51 @@ pub fn render_metrics(status: &Status) -> String { .map(|v| (format!("version=\"{}\"", label_value(v)), 1.0)) .collect(), ); + gauge( + "borg_rollout_phase", + "1 for the quorum rollout's current phase (canary, wave, complete, halted); absent without a gated rollout.", + status + .rollout + .gate + .iter() + .flat_map(|g| { + RolloutPhase::ALL.into_iter().map(move |p| { + ( + format!("phase=\"{}\"", p.as_str()), + f64::from(u8::from(g.state.phase == p)), + ) + }) + }) + .collect(), + ); + gauge( + "borg_rollout_members", + "Drones currently told to run the target version (gated rollout).", + status + .rollout + .gate + .iter() + .map(|g| { + ( + String::new(), + f64::from(u32::try_from(g.state.members.len()).unwrap_or(0)), + ) + }) + .collect(), + ); + gauge( + "borg_rollout_jobs_observed", + "Jobs observed since the drones got the target version, by group (gated rollout).", + status + .rollout + .gate + .iter() + .flat_map(|g| { + [("new", g.new_group.jobs), ("control", g.control_group.jobs)] + .map(|(group, n)| (format!("group=\"{group}\""), f64::from(n))) + }) + .collect(), + ); gauge( "borg_rollout_drones", "Drones by rollout state against the target version (absent when no target is set).", @@ -579,6 +1282,171 @@ pub fn render_metrics(status: &Status) -> String { }) .collect(), ); + gauge( + "borg_demand_pressure", + "Demand pressure per class (0-1000) published to drones for their bids.", + status + .pressure + .iter() + .map(|p| (format!("class=\"{}\"", p.class), f64::from(p.pressure))) + .collect(), + ); + gauge( + "borg_warm_floor_dynamic", + "The warm-floor controller's own level per class and scope (0 when idle or off).", + status + .homeostasis + .iter() + .map(|r| { + ( + format!("class=\"{}\",scope=\"{}\"", r.class, scope_name(r.scope)), + f64::from(r.level), + ) + }) + .collect(), + ); + gauge( + "borg_warm_floor_changes_last_hour", + "Warm-floor controller changes in the last hour per class and scope.", + status + .homeostasis + .iter() + .map(|r| { + ( + format!("class=\"{}\",scope=\"{}\"", r.class, scope_name(r.scope)), + f64::from(r.changes_last_hour), + ) + }) + .collect(), + ); + gauge( + "borg_drone_profile_speed", + "Learned speed of a drone for a class against the fleet median (1 = fleet speed).", + status + .profiles + .iter() + .map(|p| { + ( + format!("drone=\"{}\",class=\"{}\"", label_value(&p.drone), p.class), + f64::from(p.speed_permille) / 1000.0, + ) + }) + .collect(), + ); + gauge( + "borg_drone_profile_reliability", + "Learned reliability of a drone for a class (1 = no host-caused failures).", + status + .profiles + .iter() + .map(|p| { + ( + format!("drone=\"{}\",class=\"{}\"", label_value(&p.drone), p.class), + f64::from(p.reliability_permille) / 1000.0, + ) + }) + .collect(), + ); + let reclaim_rows = |f: &dyn Fn(&ReclaimStats) -> Option| -> Vec<(String, f64)> { + status + .drones + .iter() + .filter_map(|d| { + d.reclaim + .as_ref() + .and_then(f) + .map(|v| (format!("drone=\"{}\"", label_value(&d.name)), v)) + }) + .collect() + }; + #[allow(clippy::cast_precision_loss, reason = "byte counts far below 2^52")] + { + gauge( + "borg_reclaim_bytes_total", + "Bytes the drone's build-box reclaim really deleted since it started.", + reclaim_rows(&|r| Some(r.reclaimed_bytes_total as f64)), + ); + gauge( + "borg_reclaim_would_reclaim_bytes", + "Bytes the last dry-run reclaim cycle would have deleted.", + reclaim_rows(&|r| Some(r.last_run_would_reclaim_bytes as f64)), + ); + } + gauge( + "borg_reclaim_disk_usage_percent", + "Disk usage the drone's last reclaim cycle saw.", + reclaim_rows(&|r| r.usage_pct.map(f64::from)), + ); + gauge( + "borg_reclaim_errors", + "Deletions that failed in the drone's last reclaim cycle.", + reclaim_rows(&|r| Some(f64::from(r.last_run_errors))), + ); + #[allow( + clippy::cast_precision_loss, + reason = "unix timestamps fit f64 exactly" + )] + gauge( + "borg_reclaim_last_run_timestamp_seconds", + "Unix time of the drone's last reclaim cycle.", + reclaim_rows(&|r| r.last_run_at.map(|t| t as f64)), + ); + gauge( + "borg_wake_attempts_total", + "Wake-on-LAN packets sent per host since the queen started.", + status + .wake + .attempts + .iter() + .map(|(h, n)| { + #[allow(clippy::cast_precision_loss, reason = "counts far below 2^52")] + let v = *n as f64; + (format!("host=\"{}\"", label_value(h)), v) + }) + .collect(), + ); + gauge( + "borg_wake_unmet_seconds", + "How long demand for a class has stayed unmet (wake-on-LAN trigger).", + status + .wake + .unmet_secs + .iter() + .map(|(c, s)| { + #[allow(clippy::cast_precision_loss, reason = "durations far below 2^52")] + let v = *s as f64; + (format!("class=\"{c}\""), v) + }) + .collect(), + ); + gauge( + "borg_units", + "Work units by state.", + status + .units + .counts + .iter() + .map(|(s, n)| (format!("state=\"{s}\""), f64::from(*n))) + .collect(), + ); + #[allow(clippy::cast_precision_loss, reason = "counts and ages far below 2^52")] + { + gauge( + "borg_unit_joins_total", + "Unit submissions that joined an identical running unit since the queen started.", + vec![(String::new(), status.units.joins as f64)], + ); + gauge( + "borg_unit_receipt_hits_total", + "Unit submissions answered with a stored receipt since the queen started.", + vec![(String::new(), status.units.receipt_hits as f64)], + ); + gauge( + "borg_unit_oldest_queued_age_seconds", + "Age of the oldest queued work unit.", + vec![(String::new(), status.units.oldest_queued_age_secs as f64)], + ); + } gauge( "borg_reservations", "Drones held back for starving demand.", @@ -738,12 +1606,15 @@ pub async fn serve( } }) }); - let served = axum::serve(listener, router(queen)) - .with_graceful_shutdown(async move { - shutdown.await; - let _ = stop_tx.send(true); - }) - .await; + let served = axum::serve( + listener, + router(queen).into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(async move { + shutdown.await; + let _ = stop_tx.send(true); + }) + .await; let _ = ticker.await; if let Some(task) = janitor_task { let _ = task.await; diff --git a/src/bid.rs b/src/bid.rs new file mode 100644 index 0000000..3b807ab --- /dev/null +++ b/src/bid.rs @@ -0,0 +1,339 @@ +//! Drone bids and demand pressure: the "response threshold" half of the +//! allocator. +//! +//! A drone that knows its own state (warm caches, CPU pressure, owner activity, +//! owner caps) reports a [`DroneBid`] per `(class, scope)` in its heartbeat. The +//! queen computes a demand pressure per class ([`demand_pressure`]) and +//! publishes it to drones in the heartbeat response and in `/status`. The +//! allocator ([`crate::allocator::allocate`]) then: +//! +//! 1. drops a bid when the class's pressure is below the bid's `min_pressure` +//! (the drone's response threshold: low for a warm idle host, higher for a +//! cold one, high for a workstation whose owner is active); +//! 2. drops a bid with no `available_slots` left; +//! 3. scores the rest with [`score`], an integer function of the bid, and picks +//! the highest score, breaking ties by tightest memory fit and then by id. +//! +//! A drone that reports no bids at all (an older drone, or one whose bids have +//! expired) is a *legacy* drone: always eligible, neutral score. A drone that +//! reports bids but none for a given `(class, scope)` declines that key. +//! +//! Everything is integers (permille, milliseconds) so scoring is deterministic +//! and `Eq`, and the allocator stays a pure function. This module does not read +//! a clock: the queen drops expired bids (`expires_at`) before calling the +//! allocator. + +use serde::{Deserialize, Serialize}; + +use crate::classes::{Class, Scope}; + +/// 1.0 in permille. +pub const PERMILLE: u16 = 1000; + +/// One drone's offer for one `(class, scope)`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct DroneBid { + pub class: Class, + pub scope: Scope, + /// How many more leases of this key the drone would take right now. + pub available_slots: u32, + /// Estimated time from lease grant to the runner taking a job. + #[serde(default)] + pub estimated_start_ms: u64, + /// Share (permille) of the work this class typically needs that is warm in + /// the host caches. 1000 = fully warm. + #[serde(default)] + pub cache_affinity: u16, + /// Host load (permille): CPU/IO pressure, owner activity folded in. + #[serde(default)] + pub resource_pressure: u16, + /// Relative cost of a slot-hour on this host; 1000 = fleet typical. + #[serde(default = "typical_cost")] + pub estimated_cost: u32, + /// How much the drone trusts its own numbers (permille); the score is + /// shrunk toward neutral as confidence falls. + #[serde(default = "full")] + pub confidence: u16, + /// Unix time after which the queen ignores this bid. + pub expires_at: i64, + /// Response threshold: the demand pressure (permille) below which the drone + /// does not want this work. 0 = always willing. + #[serde(default)] + pub min_pressure: u16, +} + +fn typical_cost() -> u32 { + 1000 +} +fn full() -> u16 { + PERMILLE +} + +impl DroneBid { + /// Clamp the permille fields to range. Applied to everything a drone + /// sends. + #[must_use] + pub fn clamped(mut self) -> Self { + self.cache_affinity = self.cache_affinity.min(PERMILLE); + self.resource_pressure = self.resource_pressure.min(PERMILLE); + self.confidence = self.confidence.min(PERMILLE); + self.min_pressure = self.min_pressure.min(PERMILLE); + self.estimated_cost = self.estimated_cost.min(100_000); + self + } +} + +/// Weights of the score terms. Integers, so the result is exact. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ScoreWeights { + pub cache_affinity: i64, + pub resource_pressure: i64, + /// Applied to the start estimate scaled to permille of [`START_CAP_MS`]. + pub start: i64, + /// Applied to the cost's excess over typical (permille); cheaper hosts + /// score higher. + pub cost: i64, + /// Applied to a learned profile's speed excess over the fleet (permille), + /// scaled by its confidence. + pub profile_speed: i64, + /// Applied to a learned profile's reliability shortfall (permille). + pub profile_reliability: i64, +} + +impl Default for ScoreWeights { + fn default() -> Self { + Self { + cache_affinity: 3, + resource_pressure: 3, + start: 2, + cost: 1, + profile_speed: 1, + profile_reliability: 3, + } + } +} + +/// Start estimates are capped here before weighing. +pub const START_CAP_MS: u64 = 60_000; + +/// Deterministic score of a bid; higher wins. +/// +/// `weights.cache * affinity - weights.pressure * pressure - weights.start * +/// start - weights.cost * (cost - typical)`, then scaled by confidence. +pub fn score(bid: &DroneBid, w: &ScoreWeights) -> i64 { + let bid = bid.clamped(); + let start = i64::try_from(bid.estimated_start_ms.min(START_CAP_MS) * 1000 / START_CAP_MS) + .unwrap_or(1000); + let cost = i64::from(bid.estimated_cost) - 1000; + let raw = w.cache_affinity * i64::from(bid.cache_affinity) + - w.resource_pressure * i64::from(bid.resource_pressure) + - w.start * start + - w.cost * cost; + raw * i64::from(bid.confidence) / i64::from(PERMILLE) +} + +/// What the queen has learned about a drone for one class, reduced to what the +/// score needs. See [`crate::profile`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProfileScore { + pub class: Class, + /// 1000 = fleet speed. + pub speed_permille: u32, + pub speed_confidence_permille: u32, + /// 1000 = no host-caused failures. + pub reliability_permille: u32, +} + +/// The score contribution of a learned profile: faster than the fleet adds (in +/// proportion to the evidence), host-caused failures subtract. +pub fn profile_term(p: &ProfileScore, w: &ScoreWeights) -> i64 { + let speed = i64::from(p.speed_permille.clamp(500, 2000)) - 1000; + let conf = i64::from(p.speed_confidence_permille.min(1000)); + let rel = i64::from(p.reliability_permille.min(1000)) - 1000; + w.profile_speed * speed * conf / 1000 + w.profile_reliability * rel +} + +/// Demand pressure for one class, as published to drones. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct ClassPressure { + pub class: Class, + /// 0 (nothing waiting for a runner) to 1000 (saturated or starving). + pub pressure: u16, +} + +/// Pressure (permille) from queued jobs that still need a runner +/// (`shortfall`), the idle slots the fleet has for the class, and whether any +/// queued job of the class has starved. 0 with no shortfall; 1000 when +/// starving or when the shortfall is at least twice the idle slots. +pub fn demand_pressure(shortfall: u32, idle_slots: u32, starving: bool) -> u16 { + if shortfall == 0 { + return 0; + } + if starving { + return PERMILLE; + } + let p = u64::from(shortfall) * 500 / u64::from(idle_slots.max(1)); + u16::try_from(p.min(u64::from(PERMILLE))).unwrap_or(PERMILLE) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn bid() -> DroneBid { + DroneBid { + class: Class::Light, + scope: Scope::Pr, + available_slots: 4, + estimated_start_ms: 0, + cache_affinity: 0, + resource_pressure: 0, + estimated_cost: 1000, + confidence: PERMILLE, + expires_at: 0, + min_pressure: 0, + } + } + + #[test] + fn warmer_idler_cheaper_faster_hosts_score_higher() { + let w = ScoreWeights::default(); + let base = score(&bid(), &w); + assert_eq!(base, 0); + assert!( + score( + &DroneBid { + cache_affinity: 800, + ..bid() + }, + &w + ) > base + ); + assert!( + score( + &DroneBid { + resource_pressure: 500, + ..bid() + }, + &w + ) < base + ); + assert!( + score( + &DroneBid { + estimated_start_ms: 30_000, + ..bid() + }, + &w + ) < base + ); + assert!( + score( + &DroneBid { + estimated_cost: 500, + ..bid() + }, + &w + ) > base + ); + assert!( + score( + &DroneBid { + estimated_cost: 2000, + ..bid() + }, + &w + ) < base + ); + } + + #[test] + fn exact_values_and_confidence_shrinkage() { + let w = ScoreWeights::default(); + let b = DroneBid { + cache_affinity: 1000, + resource_pressure: 200, + estimated_start_ms: 30_000, + estimated_cost: 1100, + ..bid() + }; + // 3*1000 - 3*200 - 2*500 - 1*100 = 1300 + assert_eq!(score(&b, &w), 1300); + assert_eq!( + score( + &DroneBid { + confidence: 500, + ..b + }, + &w + ), + 650 + ); + assert_eq!(score(&DroneBid { confidence: 0, ..b }, &w), 0); + } + + #[test] + fn hostile_numbers_are_clamped_not_overflowed() { + let w = ScoreWeights::default(); + let b = DroneBid { + cache_affinity: u16::MAX, + resource_pressure: u16::MAX, + estimated_start_ms: u64::MAX, + estimated_cost: u32::MAX, + confidence: u16::MAX, + min_pressure: u16::MAX, + ..bid() + }; + let c = b.clamped(); + assert_eq!( + ( + c.cache_affinity, + c.resource_pressure, + c.confidence, + c.min_pressure + ), + (1000, 1000, 1000, 1000) + ); + let _ = score(&b, &w); + } + + #[test] + fn pressure_rises_with_shortfall_and_falls_with_idle_capacity() { + assert_eq!(demand_pressure(0, 0, true), 0); + assert_eq!(demand_pressure(1, 8, false), 62); + assert_eq!(demand_pressure(4, 8, false), 250); + assert_eq!(demand_pressure(1, 0, false), 500); + assert_eq!(demand_pressure(2, 0, false), 1000); + assert_eq!(demand_pressure(100, 1, false), 1000); + assert_eq!(demand_pressure(1, 100, true), 1000); + } + + #[test] + fn profiles_add_for_speed_in_proportion_to_evidence_and_subtract_for_unreliability() { + let w = ScoreWeights::default(); + let p = |speed, conf, rel| ProfileScore { + class: Class::Light, + speed_permille: speed, + speed_confidence_permille: conf, + reliability_permille: rel, + }; + assert_eq!(profile_term(&p(1000, 1000, 1000), &w), 0); + assert_eq!(profile_term(&p(1500, 1000, 1000), &w), 500); + assert_eq!(profile_term(&p(1500, 500, 1000), &w), 250); + assert_eq!(profile_term(&p(1500, 0, 1000), &w), 0); + assert_eq!(profile_term(&p(1000, 1000, 800), &w), -600); + assert_eq!(profile_term(&p(u32::MAX, 5000, u32::MAX), &w), 1000); + } + + #[test] + fn wire_format_defaults() { + let b: DroneBid = serde_json::from_str( + r#"{"class":"light","scope":"pr","available_slots":2,"expires_at":99}"#, + ) + .unwrap(); + assert_eq!( + (b.confidence, b.estimated_cost, b.min_pressure), + (1000, 1000, 0) + ); + } +} diff --git a/src/classes.rs b/src/classes.rs index f06f8e0..a65824f 100644 --- a/src/classes.rs +++ b/src/classes.rs @@ -10,6 +10,8 @@ use serde::{Deserialize, Serialize}; /// Default label prefix shared by every class label. pub const DEFAULT_LABEL_PREFIX: &str = "borg-ci-"; +/// Most extra labels one drone may add to its runners. +pub const MAX_EXTRA_LABELS: usize = 16; /// Suffix marking pull-request demand. const PR_SUFFIX: &str = "-pr"; @@ -152,6 +154,35 @@ impl Classes { Ok(()) } + /// Keep the extra runner labels a drone asked for that are safe to add: + /// well-formed, not `self-hosted`, not a label of any class (which could + /// change which jobs the runner is offered), no duplicates, at most 16. + pub fn sanitize_extra_labels(&self, labels: &[String]) -> Vec { + let mut out: Vec = Vec::new(); + for label in labels { + if valid_extra_label(label) && self.parse_label(label).is_none() && !out.contains(label) + { + out.push(label.clone()); + } + } + out.truncate(MAX_EXTRA_LABELS); + out + } + + /// The class and scope a runner minted with `labels` will serve: pull + /// request runners carry the `-pr` label, main-only runners only the bare + /// one. Labels that are not class labels never affect the answer. + pub fn labels_scope(&self, labels: &[String]) -> Option<(Class, Scope)> { + let mut found: Option<(Class, Scope)> = None; + for (class, scope) in labels.iter().filter_map(|l| self.parse_label(l)) { + found = match found { + Some((c, Scope::Pr)) if c == class => Some((c, Scope::Pr)), + _ => Some((class, scope)), + }; + } + found + } + /// The label a job requests for `class` and `scope`. pub fn job_label(&self, class: Class, scope: Scope) -> String { match scope { @@ -186,6 +217,17 @@ impl Classes { } } +/// A label the drone may add to its runners: 1 to 64 characters of letters, +/// digits, `.`, `_` and `-`, and not the reserved `self-hosted`. +pub fn valid_extra_label(label: &str) -> bool { + !label.is_empty() + && label.len() <= 64 + && label + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-')) + && !label.eq_ignore_ascii_case("self-hosted") +} + #[cfg(test)] mod tests { use super::*; @@ -294,4 +336,68 @@ mod tests { }) )); } + + #[test] + fn extra_labels_are_filtered_and_never_change_the_scope() { + let c = Classes::default(); + let extras: Vec = [ + "legacy-lane", + "on-prem", + "borg-ci-heavy-pr", + "borg-ci-light", + "self-hosted", + "bad label", + "on-prem", + ] + .iter() + .map(|s| (*s).to_owned()) + .collect(); + assert_eq!( + c.sanitize_extra_labels(&extras), + ["legacy-lane", "on-prem"], + "class labels, duplicates and malformed labels are dropped" + ); + + for (class, scope) in [ + (Class::Medium, Scope::Pr), + (Class::Medium, Scope::Main), + (Class::Light, Scope::Pr), + ] { + let mut labels = c.runner_labels(class, scope); + labels.extend(c.sanitize_extra_labels(&extras)); + assert_eq!( + c.labels_scope(&labels), + Some((class, scope)), + "{class} {scope:?}" + ); + } + // Without the sanitizer a smuggled class label would flip a PR runner to the main lane. + let mut smuggled = c.runner_labels(Class::Medium, Scope::Main); + smuggled.push("borg-ci-medium-pr".to_owned()); + assert_eq!(c.labels_scope(&smuggled), Some((Class::Medium, Scope::Pr))); + assert_eq!( + c.sanitize_extra_labels(&["borg-ci-medium-pr".to_owned()]), + Vec::::new() + ); + } + + #[test] + fn extra_label_syntax() { + assert!(valid_extra_label("on-prem") && valid_extra_label("a.b_c-1")); + for bad in [ + "", + "self-hosted", + "Self-Hosted", + "a b", + "a/b", + &"x".repeat(65), + ] { + assert!(!valid_extra_label(bad), "{bad}"); + } + let many: Vec = (0..40).map(|i| format!("l{i}")).collect(); + assert_eq!( + Classes::default().sanitize_extra_labels(&many).len(), + MAX_EXTRA_LABELS + ); + } } diff --git a/src/config.rs b/src/config.rs index 2e7b0f5..b8a73a1 100644 --- a/src/config.rs +++ b/src/config.rs @@ -9,7 +9,15 @@ use serde::Deserialize; use crate::allocator::{WarmFloor, default_warm_floors}; use crate::classes::{Classes, ClassesError}; +use crate::envcontract::ContractSettings; +use crate::health::HealthSettings; +use crate::homeostasis::HomeostasisSettings; +use crate::profile::ProfileSettings; +use crate::rollout::RolloutSettings; use crate::schedule::WarmSchedule; +use crate::tailnet::TailnetEnrollSettings; +use crate::units::UnitSettings; +use crate::wake::WakeSettings; #[derive(Debug, thiserror::Error)] pub enum ConfigError { @@ -114,6 +122,35 @@ pub struct QueenConfig { /// Time-of-day raises to the warm floors; see [`crate::schedule`]. #[serde(default)] pub warm_schedule: WarmSchedule, + /// Absent (the default) disables `POST /v1/enroll`; see [`crate::tailnet`]. + #[serde(default)] + pub tailnet_enroll: Option, + /// Demand-responsive warm floors; off by default. See + /// [`crate::homeostasis`]. + #[serde(default)] + pub homeostasis: HomeostasisSettings, + /// The build environment contract checked before work is placed on a drone; + /// off by default. See [`crate::envcontract`]. + #[serde(default)] + pub env_contract: ContractSettings, + /// Work units (gate and hook runs) and their API; off by default. See + /// [`crate::units`]. + #[serde(default)] + pub units: UnitSettings, + /// Wake-on-LAN for sleeping hosts when demand goes unmet; off by default. + /// See [`crate::wake`]. + #[serde(default)] + pub wake: WakeSettings, + /// Gating of the rollout of `rollout_target_version`; see + /// [`crate::rollout`]. Gated only when `previous_version` is set. + #[serde(default)] + pub rollout: RolloutSettings, + /// Learned per-drone profiles that feed placement; see [`crate::profile`]. + #[serde(default)] + pub profiles: ProfileSettings, + /// Host-health thresholds, decay and cooldown; see [`crate::health`]. + #[serde(default)] + pub health: HealthSettings, /// Job telemetry older than this is deleted. #[serde(default = "default_telemetry_days")] pub telemetry_retention_days: u32, @@ -200,6 +237,19 @@ impl QueenConfig { if self.poll_interval_secs == 0 || self.lease_ttl_secs <= 0 || self.lease_max_secs <= 0 { return bad("poll_interval_secs, lease_ttl_secs and lease_max_secs must be positive"); } + if self.telemetry_retention_days == 0 { + return bad("telemetry_retention_days must be at least 1"); + } + self.env_contract.validate().map_err(ConfigError::Invalid)?; + self.units.validate().map_err(ConfigError::Invalid)?; + self.wake.validate().map_err(ConfigError::Invalid)?; + self.rollout.validate().map_err(ConfigError::Invalid)?; + self.profiles.validate().map_err(ConfigError::Invalid)?; + self.homeostasis.validate().map_err(ConfigError::Invalid)?; + self.health.validate().map_err(ConfigError::Invalid)?; + if let Some(t) = &self.tailnet_enroll { + t.validate().map_err(ConfigError::Invalid)?; + } self.warm_schedule .validate() .map_err(ConfigError::Invalid)?; @@ -269,6 +319,51 @@ mod tests { } } + #[test] + fn health_settings_default_and_validate() { + let c = QueenConfig::from_json(MIN).unwrap(); + assert!(c.health.auto_quarantine); + let with = |h: &str| { + QueenConfig::from_json(&MIN.replace("\"repos\"", &format!("\"health\":{h},\"repos\""))) + }; + assert!(with(r#"{"half_life_secs":0}"#).is_err()); + assert!(with(r#"{"nope":1}"#).is_err()); + assert!( + !with(r#"{"auto_quarantine":false}"#) + .unwrap() + .health + .auto_quarantine + ); + } + + #[test] + fn zero_telemetry_retention_is_rejected() { + let json = MIN.replace("\"repos\"", "\"telemetry_retention_days\":0,\"repos\""); + assert!(QueenConfig::from_json(&json).is_err()); + } + + #[test] + fn tailnet_enroll_is_off_by_default_and_needs_an_allowlist() { + assert!( + QueenConfig::from_json(MIN) + .unwrap() + .tailnet_enroll + .is_none() + ); + let with = |t: &str| { + QueenConfig::from_json( + &MIN.replace("\"repos\"", &format!("\"tailnet_enroll\":{t},\"repos\"")), + ) + }; + assert!(with("{}").is_err()); + assert!(with(r#"{"allowed_tags":["ci"]}"#).is_err()); + let ok = with(r#"{"allowed_tags":["tag:ci"]}"#).unwrap(); + assert_eq!( + ok.tailnet_enroll.unwrap().socket_path.to_str(), + Some("/var/run/tailscale/tailscaled.sock") + ); + } + #[test] fn starvation_can_be_disabled() { let c = QueenConfig::from_json( diff --git a/src/drone/caches/bazel.rs b/src/drone/caches/bazel.rs index 61aa582..9719282 100644 --- a/src/drone/caches/bazel.rs +++ b/src/drone/caches/bazel.rs @@ -9,6 +9,8 @@ use std::fs; use std::io; use std::path::Path; +use crate::drone::fsops::FsOps; + const BEGIN: &str = "# >>> borg managed cache flags (do not edit) >>>"; const END: &str = "# <<< borg managed cache flags <<<"; @@ -39,7 +41,12 @@ pub fn with_block(existing: &str, block: &str) -> String { } /// Write the block into `rcfile`. Returns whether the file changed. -pub fn write_rc(rcfile: &Path, disk_cache: &Path, repository_cache: &Path) -> io::Result { +pub fn write_rc( + fs_ops: &dyn FsOps, + rcfile: &Path, + disk_cache: &Path, + repository_cache: &Path, +) -> io::Result { let existing = match fs::read_to_string(rcfile) { Ok(text) => text, Err(e) if e.kind() == io::ErrorKind::NotFound => String::new(), @@ -49,16 +56,14 @@ pub fn write_rc(rcfile: &Path, disk_cache: &Path, repository_cache: &Path) -> io if updated == existing { return Ok(false); } - if let Some(parent) = rcfile.parent() { - fs::create_dir_all(parent)?; - } - fs::write(rcfile, updated)?; + fs_ops.write_file(rcfile, &updated)?; Ok(true) } #[cfg(test)] mod tests { use super::*; + use crate::drone::fsops::LocalFs; fn block() -> String { managed_block(Path::new("/c/bazel/disk"), Path::new("/c/bazel/repo")) @@ -95,9 +100,33 @@ mod tests { fn writes_the_file_and_reports_changes() { let dir = tempfile::tempdir().unwrap(); let rc = dir.path().join("home/.bazelrc"); - assert!(write_rc(&rc, Path::new("/c/bazel/disk"), Path::new("/c/bazel/repo")).unwrap()); - assert!(!write_rc(&rc, Path::new("/c/bazel/disk"), Path::new("/c/bazel/repo")).unwrap()); - assert!(write_rc(&rc, Path::new("/c/other/disk"), Path::new("/c/bazel/repo")).unwrap()); + assert!( + write_rc( + &LocalFs, + &rc, + Path::new("/c/bazel/disk"), + Path::new("/c/bazel/repo") + ) + .unwrap() + ); + assert!( + !write_rc( + &LocalFs, + &rc, + Path::new("/c/bazel/disk"), + Path::new("/c/bazel/repo") + ) + .unwrap() + ); + assert!( + write_rc( + &LocalFs, + &rc, + Path::new("/c/other/disk"), + Path::new("/c/bazel/repo") + ) + .unwrap() + ); assert!(fs::read_to_string(&rc).unwrap().contains("/c/other/disk")); } } diff --git a/src/drone/caches/gc.rs b/src/drone/caches/gc.rs index 3e5ca01..70f9aff 100644 --- a/src/drone/caches/gc.rs +++ b/src/drone/caches/gc.rs @@ -7,10 +7,11 @@ //! it is not collected again after the next job. use std::fs; -use std::io; use std::path::{Path, PathBuf}; use std::time::SystemTime; +use crate::drone::fsops::FsOps; + /// Collection stops once the cache is at or below this share of its cap. const TARGET_PERCENT: u64 = 90; @@ -91,44 +92,6 @@ pub fn collect_units(buckets: &[(PathBuf, usize)]) -> Vec { units } -/// Directory contents may be read-only (module caches); make them removable. -fn make_writable(path: &Path) { - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let mut stack = vec![path.to_owned()]; - while let Some(p) = stack.pop() { - let Ok(meta) = fs::symlink_metadata(&p) else { - continue; - }; - if meta.is_dir() { - let _ = fs::set_permissions(&p, fs::Permissions::from_mode(0o755)); - if let Ok(rd) = fs::read_dir(&p) { - stack.extend(rd.flatten().map(|e| e.path())); - } - } - } - } -} - -fn remove_unit(path: &Path) -> io::Result<()> { - let is_dir = fs::symlink_metadata(path)?.is_dir(); - let attempt = |p: &Path| { - if is_dir { - fs::remove_dir_all(p) - } else { - fs::remove_file(p) - } - }; - match attempt(path) { - Err(e) if e.kind() == io::ErrorKind::PermissionDenied && is_dir => { - make_writable(path); - attempt(path) - } - other => other, - } -} - /// Result of one collection. #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] pub struct Collected { @@ -139,7 +102,12 @@ pub struct Collected { /// Evict least-recently-used `units` until `total_bytes` (the combined size of /// all `units`; data that is never evicted does not count) is at most 90% of /// `cap_bytes`. Does nothing while the units are within the cap. -pub fn evict_lru(mut units: Vec, total_bytes: u64, cap_bytes: u64) -> Collected { +pub fn evict_lru( + fs: &dyn FsOps, + mut units: Vec, + total_bytes: u64, + cap_bytes: u64, +) -> Collected { let mut out = Collected::default(); if total_bytes <= cap_bytes { return out; @@ -151,7 +119,7 @@ pub fn evict_lru(mut units: Vec, total_bytes: u64, cap_bytes: u64) -> Coll if remaining <= target { break; } - match remove_unit(&unit.path) { + match fs.remove_all(&unit.path) { Ok(()) => { remaining = remaining.saturating_sub(unit.size); out.removed_units += 1; @@ -237,7 +205,7 @@ mod tests { let units = collect_units(&[(dir.path().join("b"), 1)]); let total: u64 = units.iter().map(|u| u.size).sum(); // A cap of half the total must drop the two oldest units. - let got = evict_lru(units, total, total / 2); + let got = evict_lru(&crate::drone::fsops::LocalFs, units, total, total / 2); assert!(got.removed_units >= 2, "{got:?}"); assert!(!dir.path().join("b/oldest").exists() && !dir.path().join("b/old").exists()); assert!( @@ -252,7 +220,10 @@ mod tests { file(&dir.path().join("b/u/data"), 1000, 100); let units = collect_units(&[(dir.path().join("b"), 1)]); let total: u64 = units.iter().map(|u| u.size).sum(); - assert_eq!(evict_lru(units, total, total), Collected::default()); + assert_eq!( + evict_lru(&crate::drone::fsops::LocalFs, units, total, total), + Collected::default() + ); assert!(dir.path().join("b/u").exists()); } @@ -274,7 +245,7 @@ mod tests { .unwrap(); let units = collect_units(&[(dir.path().join("b"), 1)]); let total: u64 = units.iter().map(|u| u.size).sum(); - let got = evict_lru(units, total, 1); + let got = evict_lru(&crate::drone::fsops::LocalFs, units, total, 1); assert_eq!(got.removed_units, 1); assert!(!dir.path().join("b/mod@v1").exists()); } diff --git a/src/drone/caches/git.rs b/src/drone/caches/git.rs index 81cde39..e668329 100644 --- a/src/drone/caches/git.rs +++ b/src/drone/caches/git.rs @@ -12,6 +12,7 @@ use tokio::sync::Mutex; use crate::drone::config::GitCaches; use crate::drone::exec::Exec; +use crate::drone::fsops::FsOps; /// A mirror refreshed this recently is not refreshed again for a new lease. const MIN_REFRESH_GAP: Duration = Duration::from_secs(15); @@ -61,9 +62,17 @@ fn run_git(exec: &dyn Exec, git: &str, args: &[&str]) -> Result<(), String> { /// Create the mirror if needed, point it at `url`, and fetch every branch and /// tag into it. Idempotent, so an interrupted first clone heals on the next run. -fn refresh_mirror(exec: &dyn Exec, git: &str, path: &Path, url: &str) -> Result<(), String> { +fn refresh_mirror( + exec: &dyn Exec, + fs_ops: &dyn FsOps, + git: &str, + path: &Path, + url: &str, +) -> Result<(), String> { if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent).map_err(|e| format!("{}: {e}", parent.display()))?; + fs_ops + .mkdir_p(parent) + .map_err(|e| format!("{}: {e}", parent.display()))?; } let p = path.to_string_lossy(); if !path.join("HEAD").exists() { @@ -114,6 +123,7 @@ struct Mirror { pub struct GitMirrors { git: String, exec: Arc, + fs_ops: Arc, mirrors: Vec, /// Leases that were handed a mirror / ran without one. served: AtomicU64, @@ -121,7 +131,12 @@ pub struct GitMirrors { } impl GitMirrors { - pub fn new(cfg: &GitCaches, cache_root: &Path, exec: Arc) -> Self { + pub fn new( + cfg: &GitCaches, + cache_root: &Path, + exec: Arc, + fs_ops: Arc, + ) -> Self { let mirrors = cfg .mirrors .iter() @@ -135,6 +150,7 @@ impl GitMirrors { Self { git: cfg.bin.clone(), exec, + fs_ops, mirrors, served: AtomicU64::new(0), missed: AtomicU64::new(0), @@ -150,14 +166,16 @@ impl GitMirrors { if !force && last.is_some_and(|t| t.elapsed() < MIN_REFRESH_GAP) { return; } - let (exec, git, path, url) = ( + let (exec, fs_ops, git, path, url) = ( Arc::clone(&self.exec), + Arc::clone(&self.fs_ops), self.git.clone(), mirror.path.clone(), mirror.url.clone(), ); - let job = - tokio::task::spawn_blocking(move || refresh_mirror(exec.as_ref(), &git, &path, &url)); + let job = tokio::task::spawn_blocking(move || { + refresh_mirror(exec.as_ref(), fs_ops.as_ref(), &git, &path, &url) + }); match tokio::time::timeout(REFRESH_TIMEOUT, job).await { Ok(Ok(Ok(()))) => *last = Some(Instant::now()), Ok(Ok(Err(e))) => tracing::warn!("git mirror {} refresh failed: {e}", mirror.repo), @@ -205,6 +223,7 @@ impl GitMirrors { mod tests { use super::*; use crate::drone::cleanup::testing::FakeExec; + use crate::drone::fsops::LocalFs; fn cfg(mirrors: &[&str]) -> GitCaches { GitCaches { @@ -243,7 +262,12 @@ mod tests { async fn refresh_runs_the_expected_steps_and_is_rate_limited() { let root = tempfile::tempdir().unwrap(); let x = Arc::new(FakeExec::default()); - let m = GitMirrors::new(&cfg(&["example-org/app"]), root.path(), x.clone()); + let m = GitMirrors::new( + &cfg(&["example-org/app"]), + root.path(), + x.clone(), + Arc::new(LocalFs), + ); let path = mirror_path(root.path(), "example-org/app") .display() .to_string(); @@ -292,7 +316,12 @@ mod tests { x.fail(&format!( "git -C {path} -c http.lowSpeedLimit=1000 -c http.lowSpeedTime=60 fetch --prune --tags --force --quiet origin" )); - let m = GitMirrors::new(&cfg(&["example-org/app"]), root.path(), x.clone()); + let m = GitMirrors::new( + &cfg(&["example-org/app"]), + root.path(), + x.clone(), + Arc::new(LocalFs), + ); m.refresh_stale().await; let after_first = x.calls().len(); m.refresh_stale().await; @@ -309,6 +338,7 @@ mod tests { &cfg(&["example-org/app", "example-org/other"]), root.path(), Arc::new(FakeExec::default()), + Arc::new(LocalFs), ); let present = mirror_path(root.path(), "example-org/app"); std::fs::create_dir_all(&present).unwrap(); diff --git a/src/drone/caches/mod.rs b/src/drone/caches/mod.rs index 3527437..f025b52 100644 --- a/src/drone/caches/mod.rs +++ b/src/drone/caches/mod.rs @@ -24,6 +24,7 @@ use tokio::sync::{OnceCell, OwnedRwLockReadGuard, RwLock}; use crate::api::CacheStat; use crate::drone::config::DroneConfig; use crate::drone::exec::Exec; +use crate::drone::fsops::FsOps; use crate::probe::{self, Host}; use git::GitMirrors; use kinds::Kind; @@ -62,6 +63,7 @@ struct Track { pub struct Caches { cfg: Arc, exec: Arc, + fs_ops: Arc, host: Arc, sccache_size: OnceCell, git: GitMirrors, @@ -81,12 +83,19 @@ impl Caches { pub fn new( cfg: Arc, exec: Arc, + fs_ops: Arc, host: Arc, ) -> Self { - let git = GitMirrors::new(&cfg.caches.git, &cfg.cache_root, Arc::clone(&exec)); + let git = GitMirrors::new( + &cfg.caches.git, + &cfg.cache_root, + Arc::clone(&exec), + Arc::clone(&fs_ops), + ); Self { cfg, exec, + fs_ops, host, sccache_size: OnceCell::new(), git, @@ -126,10 +135,11 @@ impl Caches { async fn init(&self) -> &str { self.sccache_size .get_or_init(|| async { - let (cfg, host, exec) = ( + let (cfg, host, exec, fs_ops) = ( Arc::clone(&self.cfg), Arc::clone(&self.host), Arc::clone(&self.exec), + Arc::clone(&self.fs_ops), ); let flags = self.remap_flags(); let enabled: Vec = @@ -138,22 +148,25 @@ impl Caches { for kind in &enabled { for sub in kind.subdirs() { let path = kind.dir(&cfg.cache_root).join(sub); - if let Err(e) = std::fs::create_dir_all(&path) { + if let Err(e) = fs_ops.mkdir_p(&path) { tracing::warn!("create {}: {e}", path.display()); } } } - if let Err(e) = std::fs::create_dir_all(cfg.cache_root.join("git")) { + if let Err(e) = fs_ops.mkdir_p(&cfg.cache_root.join("git")) { tracing::warn!("create git cache dir: {e}"); } if cfg.caches.bazel.enabled { - write_bazelrc(&cfg); + write_bazelrc(fs_ops.as_ref(), &cfg); } if cfg.caches.cargo.enabled && cfg.caches.cargo.export_rustflags && rustflags::flags_are_passable(&flags) - && let Err(e) = - rustflags::write_config(&Kind::Cargo.dir(&cfg.cache_root), &flags) + && let Err(e) = rustflags::write_config( + fs_ops.as_ref(), + &Kind::Cargo.dir(&cfg.cache_root), + &flags, + ) { tracing::warn!("cargo config.toml: {e}"); } @@ -280,6 +293,7 @@ impl Caches { }) .collect(); let root = self.cfg.cache_root.clone(); + let fs_ops = Arc::clone(&self.fs_ops); let work = tokio::task::spawn_blocking(move || { let mut out: Vec<(&'static str, u64, u64)> = Vec::new(); for (kind, cap) in plan { @@ -290,7 +304,7 @@ impl Caches { { let units = gc::collect_units(&kind.gc_buckets(&root)); let pool: u64 = units.iter().map(|u| u.size).sum(); - freed = gc::evict_lru(units, pool, cap).freed_bytes; + freed = gc::evict_lru(fs_ops.as_ref(), units, pool, cap).freed_bytes; size = size.saturating_sub(freed); } out.push((kind.name(), size, freed)); @@ -371,13 +385,18 @@ impl Caches { } /// Point Bazel at the host-level disk and repository caches. -fn write_bazelrc(cfg: &DroneConfig) { +fn write_bazelrc(fs_ops: &dyn FsOps, cfg: &DroneConfig) { + // Bazel reads the rc of whoever runs it: the job user when there is one. + let home = match &cfg.job_user { + Some(user) => Some(user.home.clone()), + None => std::env::var_os("HOME").map(PathBuf::from), + }; let rc: Option = cfg .caches .bazel .bazelrc .clone() - .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".bazelrc"))); + .or_else(|| home.map(|h| h.join(".bazelrc"))); let Some(rc) = rc else { tracing::warn!( "HOME is unset and caches.bazel.bazelrc is not configured; skipping Bazel cache flags" @@ -385,7 +404,7 @@ fn write_bazelrc(cfg: &DroneConfig) { return; }; let dir = Kind::Bazel.dir(&cfg.cache_root); - if let Err(e) = bazel::write_rc(&rc, &dir.join("disk"), &dir.join("repo")) { + if let Err(e) = bazel::write_rc(fs_ops, &rc, &dir.join("disk"), &dir.join("repo")) { tracing::warn!("write {}: {e}", rc.display()); } } @@ -431,6 +450,7 @@ mod tests { Arc::new(Caches::new( cfg, Arc::new(FakeExec::default()), + Arc::new(crate::drone::fsops::LocalFs), Arc::new(NoHost), )) } diff --git a/src/drone/caches/rustflags.rs b/src/drone/caches/rustflags.rs index 891b4a1..6e24235 100644 --- a/src/drone/caches/rustflags.rs +++ b/src/drone/caches/rustflags.rs @@ -17,6 +17,8 @@ use std::fs; use std::io; use std::path::Path; +use crate::drone::fsops::FsOps; + /// Virtual prefix a slot directory is mapped to. pub const VIRTUAL_WORK: &str = "/borg/work"; const VIRTUAL_CARGO: &str = "/borg/cargo"; @@ -66,7 +68,7 @@ pub fn config_toml(flags: &[String]) -> String { /// Write `config.toml` unless the file exists and is not ours. Returns whether /// the file was written. -pub fn write_config(cargo_home: &Path, flags: &[String]) -> io::Result { +pub fn write_config(fs_ops: &dyn FsOps, cargo_home: &Path, flags: &[String]) -> io::Result { let path = cargo_home.join("config.toml"); let wanted = config_toml(flags); match fs::read_to_string(&path) { @@ -81,14 +83,14 @@ pub fn write_config(cargo_home: &Path, flags: &[String]) -> io::Result { Err(e) if e.kind() == io::ErrorKind::NotFound => {} Err(e) => return Err(e), } - fs::create_dir_all(cargo_home)?; - fs::write(&path, wanted)?; + fs_ops.write_file(&path, &wanted)?; Ok(true) } #[cfg(test)] mod tests { use super::*; + use crate::drone::fsops::LocalFs; #[test] fn flags_cover_every_slot_and_the_cargo_dirs_identically_for_all_leases() { @@ -130,16 +132,16 @@ mod tests { fn write_is_idempotent_and_never_clobbers_a_foreign_config() { let dir = tempfile::tempdir().unwrap(); let flags = remap_flags(Path::new("/w"), 1, dir.path()); - assert!(write_config(dir.path(), &flags).unwrap()); - assert!(!write_config(dir.path(), &flags).unwrap()); + assert!(write_config(&LocalFs, dir.path(), &flags).unwrap()); + assert!(!write_config(&LocalFs, dir.path(), &flags).unwrap()); let more = remap_flags(Path::new("/w"), 2, dir.path()); assert!( - write_config(dir.path(), &more).unwrap(), + write_config(&LocalFs, dir.path(), &more).unwrap(), "our own file is updated" ); fs::write(dir.path().join("config.toml"), "[net]\nretry = 5\n").unwrap(); - assert!(write_config(dir.path(), &flags).is_err()); + assert!(write_config(&LocalFs, dir.path(), &flags).is_err()); assert_eq!( fs::read_to_string(dir.path().join("config.toml")).unwrap(), "[net]\nretry = 5\n" diff --git a/src/drone/cleanup.rs b/src/drone/cleanup.rs index 26cac9c..fa23dbf 100644 --- a/src/drone/cleanup.rs +++ b/src/drone/cleanup.rs @@ -11,6 +11,7 @@ use std::io; use std::path::Path; use crate::drone::exec::Exec; +use crate::drone::fsops::FsOps; /// Label that marks a resource as belonging to a lease. pub const LEASE_LABEL: &str = "borg.lease"; @@ -202,19 +203,18 @@ pub fn cleanup_docker(exec: &dyn Exec, docker: &str, scope: Scope<'_>) -> Docker report } -/// Delete `/` (a symlink at that path is unlinked, never -/// followed). Returns whether anything was removed. -pub fn remove_lease_dir(work_root: &Path, lease_id: &str) -> io::Result { - if !valid_lease_id(lease_id) { +/// Delete `/` (a symlink at that path is unlinked, never +/// followed). Returns whether anything was there to remove. +pub fn remove_lease_dir(fs: &dyn FsOps, work_root: &Path, name: &str) -> io::Result { + if !valid_lease_id(name) { return Err(io::Error::new( io::ErrorKind::InvalidInput, - "unsafe lease id", + "unsafe directory name", )); } - let path = work_root.join(lease_id); + let path = work_root.join(name); match std::fs::symlink_metadata(&path) { - Ok(meta) if meta.is_dir() => std::fs::remove_dir_all(&path).map(|()| true), - Ok(_) => std::fs::remove_file(&path).map(|()| true), + Ok(_) => fs.remove_all(&path).map(|()| true), Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(false), Err(e) => Err(e), } @@ -222,7 +222,7 @@ pub fn remove_lease_dir(work_root: &Path, lease_id: &str) -> io::Result { /// Delete every lease directory under `work_root` (startup: the drone holds /// no leases yet). Returns the removed names. -pub fn sweep_work_root(work_root: &Path) -> io::Result> { +pub fn sweep_work_root(fs: &dyn FsOps, work_root: &Path) -> io::Result> { let mut removed = Vec::new(); let entries = match std::fs::read_dir(work_root) { Ok(e) => e, @@ -231,7 +231,7 @@ pub fn sweep_work_root(work_root: &Path) -> io::Result> { }; for entry in entries { let name = entry?.file_name().to_string_lossy().into_owned(); - if remove_lease_dir(work_root, &name)? { + if remove_lease_dir(fs, work_root, &name)? { removed.push(name); } } @@ -300,6 +300,7 @@ pub(crate) mod testing { mod tests { use super::testing::FakeExec; use super::*; + use crate::drone::fsops::LocalFs; const ID: &str = "0123456789ab"; @@ -445,16 +446,16 @@ mod tests { std::fs::write(cache.join("entry"), "x").unwrap(); std::fs::create_dir_all(work.join("keep-me")).unwrap(); - assert!(remove_lease_dir(&work, ID).unwrap()); + assert!(remove_lease_dir(&LocalFs, &work, ID).unwrap()); assert!(!work.join(ID).exists()); assert!( - !remove_lease_dir(&work, ID).unwrap(), + !remove_lease_dir(&LocalFs, &work, ID).unwrap(), "second removal is a no-op" ); assert!(cache.join("entry").exists()); assert!(work.join("keep-me").exists()); - assert!(remove_lease_dir(&work, "../cache").is_err()); - assert!(remove_lease_dir(&work, "").is_err()); + assert!(remove_lease_dir(&LocalFs, &work, "../cache").is_err()); + assert!(remove_lease_dir(&LocalFs, &work, "").is_err()); } #[cfg(unix)] @@ -467,7 +468,7 @@ mod tests { std::fs::create_dir_all(&target).unwrap(); std::fs::write(target.join("f"), "x").unwrap(); std::os::unix::fs::symlink(&target, work.join(ID)).unwrap(); - assert!(remove_lease_dir(&work, ID).unwrap()); + assert!(remove_lease_dir(&LocalFs, &work, ID).unwrap()); assert!(target.join("f").exists()); } @@ -477,11 +478,11 @@ mod tests { for id in ["aaa", "bbb"] { std::fs::create_dir_all(root.path().join(id).join("x")).unwrap(); } - let mut removed = sweep_work_root(root.path()).unwrap(); + let mut removed = sweep_work_root(&LocalFs, root.path()).unwrap(); removed.sort(); assert_eq!(removed, ["aaa", "bbb"]); assert!( - sweep_work_root(&root.path().join("missing")) + sweep_work_root(&LocalFs, &root.path().join("missing")) .unwrap() .is_empty() ); diff --git a/src/drone/config.rs b/src/drone/config.rs index 64b1a47..e61d0b9 100644 --- a/src/drone/config.rs +++ b/src/drone/config.rs @@ -1,10 +1,12 @@ //! Drone configuration, a JSON file written by `borg drone install`. +use std::collections::BTreeMap; use std::path::{Path, PathBuf}; use serde::{Deserialize, Serialize}; use crate::allocator::HostCaps; +use crate::drone::jobuser::JobUser; use crate::secret::Secret; #[derive(Debug, thiserror::Error)] @@ -229,9 +231,36 @@ pub struct DroneConfig { /// run at once. Default: `max_leases`, else 16. #[serde(default)] pub slots: Option, + /// Directory for self-update hand-off and health evidence (`update-request`, + /// `update-result.json`, `health.json`). Absent: no self-update. + #[serde(default)] + pub state_dir: Option, + /// Run leases as this unprivileged user instead of the drone's own, so job + /// code cannot read the drone's token or config. Linux only. + #[serde(default)] + pub job_user: Option, + /// Labels added to every runner this drone registers, for lanes that still + /// route by label (for example `on-prem`). Labels of the job classes are + /// ignored by the queen so they cannot change a runner's scope. + #[serde(default)] + pub extra_labels: Vec, + /// Environment injected into every lease's runner. + #[serde(default)] + pub runner_env: BTreeMap, + /// Absolute path of a script (`.sh`) the runner runs when a job starts; + /// exported as `ACTIONS_RUNNER_HOOK_JOB_STARTED`. + #[serde(default)] + pub job_started_hook: Option, + /// Same for job completion (`ACTIONS_RUNNER_HOOK_JOB_COMPLETED`). + #[serde(default)] + pub job_completed_hook: Option, /// Host-level caches the drone manages under `cache_root`. #[serde(default)] pub caches: CachesConfig, + /// Build-box reclaim of abandoned worktrees and cargo directories; off + /// unless a host opts in, and dry-run by default. See [`crate::reclaim`]. + #[serde(default)] + pub reclaim: crate::reclaim::ReclaimSettings, } /// Slot directories when neither `slots` nor `max_leases` is set. @@ -288,9 +317,77 @@ impl DroneConfig { if self.ca_cert_path.is_some() && http { return bad("ca_cert_path only applies to an https:// queen_url"); } + if let Some(user) = &self.job_user { + user.validate().map_err(DroneConfigError::Invalid)?; + } + self.validate_job_settings()?; + self.reclaim.validate().map_err(DroneConfigError::Invalid)?; self.validate_caches() } + fn validate_job_settings(&self) -> Result<(), DroneConfigError> { + let bad = |m: String| Err(DroneConfigError::Invalid(m)); + for label in &self.extra_labels { + if !crate::classes::valid_extra_label(label) { + return bad(format!( + "extra_labels entry {label:?} is not a valid runner label" + )); + } + } + if self.extra_labels.len() > crate::classes::MAX_EXTRA_LABELS { + return bad(format!( + "at most {} extra_labels", + crate::classes::MAX_EXTRA_LABELS + )); + } + for (key, value) in &self.runner_env { + let name_ok = key + .chars() + .next() + .is_some_and(|c| c.is_ascii_alphabetic() || c == '_') + && key.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_'); + let reserved = key.starts_with("BORG_") + || key.starts_with("ACTIONS_RUNNER_") + || key == "COMPOSE_PROJECT_NAME" + || key == "DISABLE_RUNNER_UPDATE"; + if !name_ok || reserved || value.contains('\0') { + return bad(format!("runner_env key {key:?} is invalid or reserved")); + } + } + for (name, hook) in [ + ("job_started_hook", &self.job_started_hook), + ("job_completed_hook", &self.job_completed_hook), + ] { + if let Some(path) = hook + && !(path.is_absolute() && path.extension().is_some_and(|e| e == "sh")) + { + return bad(format!("{name} must be an absolute path ending in .sh")); + } + } + Ok(()) + } + + /// The runner hook and `runner_env` variables for a lease. + pub fn runner_environment(&self) -> Vec<(String, String)> { + let mut env: Vec<(String, String)> = self + .runner_env + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + for (var, hook) in [ + ("ACTIONS_RUNNER_HOOK_JOB_STARTED", &self.job_started_hook), + ( + "ACTIONS_RUNNER_HOOK_JOB_COMPLETED", + &self.job_completed_hook, + ), + ] { + if let Some(path) = hook { + env.push((var.to_owned(), path.to_string_lossy().into_owned())); + } + } + env + } + fn validate_caches(&self) -> Result<(), DroneConfigError> { let bad = |m: String| Err(DroneConfigError::Invalid(m)); let c = &self.caches; @@ -509,6 +606,45 @@ mod tests { assert!(d.sccache.server.is_none() && d.cargo.export_rustflags); } + #[test] + fn job_settings_are_validated_and_exported() { + let with = |extra: &str| BASE.replace('}', &format!(",{extra}}}")); + let ok = with( + r#""extra_labels":["on-prem","legacy-lane"],"runner_env":{"HTTP_PROXY":"http://proxy.example.test:3128"},"job_started_hook":"/etc/borg/hooks/start.sh","job_completed_hook":"/etc/borg/hooks/done.sh","job_user":{"name":"borg-job","home":"/var/lib/borg-job"}"#, + ); + let cfg = DroneConfig::from_json(&ok).unwrap(); + assert_eq!( + cfg.job_user.as_ref().map(|u| u.priv_bin.as_str()), + Some("setpriv") + ); + let env = cfg.runner_environment(); + assert!(env.contains(&( + "HTTP_PROXY".to_owned(), + "http://proxy.example.test:3128".to_owned() + ))); + assert!(env.contains(&( + "ACTIONS_RUNNER_HOOK_JOB_STARTED".to_owned(), + "/etc/borg/hooks/start.sh".to_owned() + ))); + assert!(env.contains(&( + "ACTIONS_RUNNER_HOOK_JOB_COMPLETED".to_owned(), + "/etc/borg/hooks/done.sh".to_owned() + ))); + for bad in [ + r#""job_started_hook":"/etc/hook.py""#, + r#""job_started_hook":"hook.sh""#, + r#""job_completed_hook":"/etc/hook""#, + r#""runner_env":{"BORG_LEASE_ID":"x"}"#, + r#""runner_env":{"ACTIONS_RUNNER_INPUT_JITCONFIG":"x"}"#, + r#""runner_env":{"A B":"x"}"#, + r#""extra_labels":["has space"]"#, + r#""extra_labels":["self-hosted"]"#, + r#""job_user":{"name":"x","home":"relative"}"#, + ] { + assert!(DroneConfig::from_json(&with(bad)).is_err(), "{bad}"); + } + } + #[test] fn caches_may_not_overlap_the_wiped_work_root() { let nested = BASE.replace(r#""cache_root":"/c""#, r#""cache_root":"/w/cache""#); diff --git a/src/drone/fsops.rs b/src/drone/fsops.rs new file mode 100644 index 0000000..ee9fc55 --- /dev/null +++ b/src/drone/fsops.rs @@ -0,0 +1,196 @@ +//! File operations on directories the job user owns. +//! +//! Work and cache directories belong to the job user so jobs can write them, +//! which means the drone daemon (a different user) cannot create, clone or +//! delete inside them directly. With a job user configured these operations run +//! as that user through [`JobExec`](crate::drone::jobuser::JobExec); without +//! one they are plain in-process file operations. + +use std::fs; +use std::io; +use std::path::Path; +use std::sync::Arc; + +use crate::drone::exec::Exec; + +pub trait FsOps: Send + Sync { + fn mkdir_p(&self, path: &Path) -> io::Result<()>; + /// Replace `path` with `content` (via a temporary file and rename). + fn write_file(&self, path: &Path, content: &str) -> io::Result<()>; + /// Delete a file or directory tree; a missing path is fine. + fn remove_all(&self, path: &Path) -> io::Result<()>; + /// Recreate `src` under `dst`, hard-linking files where possible. + fn clone_tree(&self, src: &Path, dst: &Path) -> io::Result<()>; +} + +/// Plain in-process file operations. +pub struct LocalFs; + +/// Directory contents may be read-only (module caches); make them removable. +fn make_writable(path: &Path) { + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let mut stack = vec![path.to_owned()]; + while let Some(p) = stack.pop() { + let Ok(meta) = fs::symlink_metadata(&p) else { + continue; + }; + if meta.is_dir() { + let _ = fs::set_permissions(&p, fs::Permissions::from_mode(0o755)); + if let Ok(rd) = fs::read_dir(&p) { + stack.extend(rd.flatten().map(|e| e.path())); + } + } + } + } +} + +impl FsOps for LocalFs { + fn mkdir_p(&self, path: &Path) -> io::Result<()> { + fs::create_dir_all(path) + } + + fn write_file(&self, path: &Path, content: &str) -> io::Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + let tmp = path.with_extension("borg-tmp"); + fs::write(&tmp, content)?; + fs::rename(&tmp, path) + } + + fn remove_all(&self, path: &Path) -> io::Result<()> { + let is_dir = match fs::symlink_metadata(path) { + Ok(m) => m.is_dir(), + Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(()), + Err(e) => return Err(e), + }; + let attempt = |p: &Path| { + if is_dir { + fs::remove_dir_all(p) + } else { + fs::remove_file(p) + } + }; + match attempt(path) { + Err(e) if e.kind() == io::ErrorKind::PermissionDenied && is_dir => { + make_writable(path); + attempt(path) + } + other => other, + } + } + + fn clone_tree(&self, src: &Path, dst: &Path) -> io::Result<()> { + crate::drone::runner::clone_tree(src, dst) + } +} + +/// File operations carried out by commands run through `exec` (the job user). +pub struct ExecFs { + exec: Arc, +} + +impl ExecFs { + pub fn new(exec: Arc) -> Self { + Self { exec } + } + + fn run(&self, program: &str, args: &[&str]) -> io::Result<()> { + let out = self.exec.run(program, args)?; + if out.success { + Ok(()) + } else { + Err(io::Error::other(format!( + "`{program} {}` failed", + args.join(" ") + ))) + } + } +} + +impl FsOps for ExecFs { + fn mkdir_p(&self, path: &Path) -> io::Result<()> { + self.run("mkdir", &["-p", "--", &path.to_string_lossy()]) + } + + fn write_file(&self, path: &Path, content: &str) -> io::Result<()> { + // $1 = content, $2 = path: nothing is interpolated into the script. + self.run( + "sh", + &[ + "-c", + "mkdir -p -- \"$(dirname -- \"$2\")\" && printf '%s' \"$1\" > \"$2.borg-tmp\" && mv -f -- \"$2.borg-tmp\" \"$2\"", + "sh", + content, + &path.to_string_lossy(), + ], + ) + } + + fn remove_all(&self, path: &Path) -> io::Result<()> { + self.run("rm", &["-rf", "--", &path.to_string_lossy()]) + } + + fn clone_tree(&self, src: &Path, dst: &Path) -> io::Result<()> { + if let Some(parent) = dst.parent() { + self.mkdir_p(parent)?; + } + // GNU cp: archive mode, hard links where the filesystem allows. + self.run( + "cp", + &["-al", "--", &src.to_string_lossy(), &dst.to_string_lossy()], + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::drone::cleanup::testing::FakeExec; + + #[test] + fn local_fs_round_trip() { + let dir = tempfile::tempdir().unwrap(); + let fs = LocalFs; + let file = dir.path().join("a/b/c.txt"); + fs.write_file(&file, "hello").unwrap(); + assert_eq!(std::fs::read_to_string(&file).unwrap(), "hello"); + fs.write_file(&file, "again").unwrap(); + assert_eq!(std::fs::read_to_string(&file).unwrap(), "again"); + fs.remove_all(&dir.path().join("a")).unwrap(); + assert!(!dir.path().join("a").exists()); + fs.remove_all(&dir.path().join("a")).unwrap(); + } + + #[test] + fn exec_fs_issues_argument_only_commands() { + let fake = Arc::new(FakeExec::default()); + let fs = ExecFs::new(fake.clone()); + fs.mkdir_p(Path::new("/c/git")).unwrap(); + fs.remove_all(Path::new("/w/slot-0")).unwrap(); + fs.clone_tree(Path::new("/r"), Path::new("/w/slot-0")) + .unwrap(); + fs.write_file(Path::new("/h/.bazelrc"), "common --x\n; rm -rf /") + .unwrap(); + let calls = fake.calls(); + assert_eq!(calls[0], "mkdir -p -- /c/git"); + assert_eq!(calls[1], "rm -rf -- /w/slot-0"); + assert_eq!(calls[2], "mkdir -p -- /w"); + assert_eq!(calls[3], "cp -al -- /r /w/slot-0"); + assert!(calls[4].starts_with("sh -c "), "{}", calls[4]); + assert!( + calls[4].ends_with("sh common --x\n; rm -rf / /h/.bazelrc"), + "content is an argument, not script text: {}", + calls[4] + ); + } + + #[test] + fn a_failing_command_is_an_error() { + let fake = Arc::new(FakeExec::default()); + fake.fail("rm -rf -- /w"); + assert!(ExecFs::new(fake).remove_all(Path::new("/w")).is_err()); + } +} diff --git a/src/drone/install.rs b/src/drone/install.rs index fa1564d..0a124d7 100644 --- a/src/drone/install.rs +++ b/src/drone/install.rs @@ -1,12 +1,16 @@ //! `borg drone install`: idempotently set a host up as a drone and report //! exactly what changed. //! -//! Linux (and WSL, which needs systemd enabled) gets a dedicated service user -//! in the docker group and a systemd unit; macOS gets a LaunchDaemon. Both get -//! a copy of this binary, the pinned actions-runner (checksum verified), a -//! config file and the enrollment token. Re-running changes only what differs -//! and restarts the service only if something changed. +//! Linux (and WSL, which needs systemd enabled) gets two service users and a +//! systemd unit: the drone daemon user (owns `/etc/borg`, mode 0700, holding +//! the enrollment token) and a separate unprivileged job user that every lease +//! runs as, so job code cannot read the token. Both are in the docker group +//! when it exists. macOS gets a LaunchDaemon and a single user. Both get a copy +//! of this binary, the pinned actions-runner (checksum verified), a config file +//! and the enrollment token. Re-running changes only what differs and restarts +//! the service only if something changed. +use std::collections::BTreeMap; use std::fmt::Write as _; use std::fs; use std::io; @@ -17,8 +21,12 @@ use sha2::{Digest, Sha256}; use crate::drone::config::DroneConfig; use crate::drone::exec::Exec; +use crate::drone::jobuser::JobUser; use crate::drone::pins::RunnerPin; use crate::secret::Secret; +use crate::update::config::{ + DEFAULT_DRONE_CONFIG, DEFAULT_MANIFEST_URL, DEFAULT_VERSIONED_MANIFEST_URL, UpdateConfig, +}; const SERVICE: &str = "borg-drone.service"; const LAUNCHD_LABEL: &str = "org.borg.drone"; @@ -64,8 +72,21 @@ pub struct InstallOptions { pub max_vcpu: Option, pub max_mem_mib: Option, pub max_leases: Option, - /// Service account the drone and its jobs run as. + /// Service account the drone daemon runs as (and, without a job user, the + /// jobs too). pub run_as: String, + /// Linux: the separate unprivileged account leases run as; `None` runs jobs + /// as `run_as` (they can then read the drone's token). Ignored on macOS. + pub job_user: Option, + /// Where host caches live; default `/cache`. + pub cache_root: Option, + /// Take ownership of an existing `cache_root` (recursively) instead of + /// leaving its files as they are. + pub adopt_cache: bool, + pub extra_labels: Vec, + pub runner_env: BTreeMap, + pub job_started_hook: Option, + pub job_completed_hook: Option, /// Filesystem root the install is written under (`/` on a real host). pub root: PathBuf, pub platform: Platform, @@ -73,6 +94,10 @@ pub struct InstallOptions { /// The borg binary to install. pub binary: PathBuf, pub dry_run: bool, + /// Release manifest the updater reads; the public release by default. + pub update_manifest_url: Option, + /// Do not install the self-update path unit and updater configuration. + pub no_self_update: bool, /// Skip the "must be root" check (tests, `--root` staging). pub skip_root_check: bool, pub drain_timeout_secs: u64, @@ -86,25 +111,43 @@ pub struct InstallOptions { struct Layout { bin: PathBuf, config_dir: PathBuf, + /// The drone's own state and working directory. data: PathBuf, + /// Home of whoever runs leases: holds the runner install, work directories + /// and (by default) the caches. + job_home: PathBuf, + cache_root: PathBuf, service_file: PathBuf, } impl Layout { - fn new(platform: Platform) -> Self { - match platform { - Platform::Linux => Self { - bin: "/usr/local/bin/borg".into(), - config_dir: "/etc/borg".into(), - data: "/var/lib/borg".into(), - service_file: format!("/etc/systemd/system/{SERVICE}").into(), - }, - Platform::MacOs => Self { - bin: "/usr/local/bin/borg".into(), - config_dir: "/usr/local/etc/borg".into(), - data: "/usr/local/var/borg".into(), - service_file: format!("/Library/LaunchDaemons/{LAUNCHD_LABEL}.plist").into(), - }, + fn new(platform: Platform, separate_job_user: bool, cache_root: Option<&Path>) -> Self { + let (bin, config_dir, data, service_file): (&str, &str, &str, String) = match platform { + Platform::Linux => ( + "/usr/local/bin/borg", + "/etc/borg", + "/var/lib/borg", + format!("/etc/systemd/system/{SERVICE}"), + ), + Platform::MacOs => ( + "/usr/local/bin/borg", + "/usr/local/etc/borg", + "/usr/local/var/borg", + format!("/Library/LaunchDaemons/{LAUNCHD_LABEL}.plist"), + ), + }; + let job_home: PathBuf = if separate_job_user { + "/var/lib/borg-job".into() + } else { + data.into() + }; + Self { + bin: bin.into(), + config_dir: config_dir.into(), + data: data.into(), + cache_root: cache_root.map_or_else(|| job_home.join("cache"), Path::to_path_buf), + job_home, + service_file: service_file.into(), } } @@ -118,7 +161,10 @@ impl Layout { self.config_dir.join("queen-ca.pem") } fn runner(&self) -> PathBuf { - self.data.join("runner") + self.job_home.join("runner") + } + fn work(&self) -> PathBuf { + self.job_home.join("work") } } @@ -128,7 +174,17 @@ pub fn render_systemd_unit( data: &Path, user: &str, drain_timeout_secs: u64, + separate_job_user: bool, ) -> String { + // The daemon starts job processes as the job user with `setpriv`, which + // needs these three capabilities; `setpriv` drops them again before the + // job runs. `NoNewPrivileges` must stay off for the same reason. + let caps = if separate_job_user { + "AmbientCapabilities=CAP_SETUID CAP_SETGID CAP_SETPCAP\n\ + CapabilityBoundingSet=CAP_SETUID CAP_SETGID CAP_SETPCAP\n" + } else { + "" + }; format!( "[Unit]\n\ Description=borg CI drone\n\ @@ -141,6 +197,7 @@ pub fn render_systemd_unit( Group={user}\n\ WorkingDirectory={data}\n\ ExecStart={bin} drone run --config {config}\n\ + {caps}\ Restart=always\n\ RestartSec=5\n\ KillSignal=SIGTERM\n\ @@ -195,6 +252,45 @@ pub fn render_launchd_plist( ) } +const UPDATE_PATH_UNIT: &str = "borg-update.path"; +const UPDATE_SERVICE_UNIT: &str = "borg-update.service"; +const UPDATE_LAUNCHD_LABEL: &str = "org.borg.update"; + +/// systemd path unit: starts the updater when the drone writes a request. +pub fn render_update_path_unit(request: &Path) -> String { + format!( + "[Unit]\nDescription=borg drone update requests\n\n[Path]\nPathChanged={}\nUnit={UPDATE_SERVICE_UNIT}\n\n[Install]\nWantedBy=multi-user.target\n", + request.display() + ) +} + +/// The privileged one-shot updater, started by the path unit. +pub fn render_update_service_unit(bin: &Path, config: &Path, timeout_secs: u64) -> String { + format!( + "[Unit]\nDescription=borg drone self-update\n\n[Service]\nType=oneshot\nUser=root\nExecStart={} update --config {} --from-request\nTimeoutStartSec={timeout_secs}\n", + bin.display(), + config.display() + ) +} + +/// launchd equivalent: a root daemon that runs when the request file changes. +pub fn render_update_plist(bin: &Path, config: &Path, request: &Path) -> String { + format!( + "\n\ + \n\ + \n\n\ + \tLabel{UPDATE_LAUNCHD_LABEL}\n\ + \tProgramArguments\n\t\n\ + \t\t{}\n\t\tupdate\n\t\t--config\n\t\t{}\n\t\t--from-request\n\ + \t\n\ + \tWatchPaths\n\t\n\t\t{}\n\t\n\ + \n\n", + bin.display(), + config.display(), + request.display() + ) +} + struct Installer<'a> { opts: &'a InstallOptions, layout: Layout, @@ -211,6 +307,118 @@ fn io_err(path: &Path) -> impl FnOnce(io::Error) -> InstallError + '_ { } impl Installer<'_> { + /// The updater: root-owned configuration (outside the drone's own + /// directory, so a compromised drone cannot redirect it) and the unit that + /// runs it when the drone asks for an update. + fn install_updater(&mut self) -> Result { + let config_path = PathBuf::from(DEFAULT_DRONE_CONFIG); + let config_dir = config_path.parent().unwrap_or(Path::new("/")).to_owned(); + self.ensure_dir(&config_dir, 0o755, Some("root"))?; + let update_cfg = UpdateConfig { + manifest_url: self + .opts + .update_manifest_url + .clone() + .unwrap_or_else(|| DEFAULT_MANIFEST_URL.to_owned()), + versioned_manifest_url: self + .opts + .update_manifest_url + .is_none() + .then(|| DEFAULT_VERSIONED_MANIFEST_URL.to_owned()), + allow_insecure: false, + service: match self.opts.platform { + Platform::Linux => SERVICE.to_owned(), + Platform::MacOs => LAUNCHD_LABEL.to_owned(), + }, + launchd_plist: (self.opts.platform == Platform::MacOs) + .then(|| self.layout.service_file.to_string_lossy().into_owned()), + bin: self.layout.bin.clone(), + health_file: self.layout.data.join("health.json"), + state_dir: self.layout.data.clone(), + runner_dir: Some(self.layout.runner()), + runner_owner: Some(self.job_owner()), + health_timeout_secs: 120, + }; + let mut json = serde_json::to_string_pretty(&update_cfg) + .map_err(|e| InstallError::Preflight(e.to_string()))?; + json.push('\n'); + let mut changed = self.write_file(&config_path, json.as_bytes(), 0o644, Some("root"))?; + let request = crate::update::request_path(&self.layout.data); + let bin = self.layout.bin.clone(); + let timeout = self.opts.drain_timeout_secs + 600; + match self.opts.platform { + Platform::Linux => { + let dir = self + .layout + .service_file + .parent() + .unwrap_or(Path::new("/")) + .to_owned(); + let path_unit = dir.join(UPDATE_PATH_UNIT); + let service_unit = dir.join(UPDATE_SERVICE_UNIT); + changed |= self.write_file( + &path_unit, + render_update_path_unit(&request).as_bytes(), + 0o644, + None, + )?; + changed |= self.write_file( + &service_unit, + render_update_service_unit(&bin, &config_path, timeout).as_bytes(), + 0o644, + None, + )?; + if !self.opts.dry_run { + if changed { + self.must("systemctl", &["daemon-reload"])?; + self.note("ran", "systemctl daemon-reload"); + } + if !self.cmd("systemctl", &["is-enabled", "--quiet", UPDATE_PATH_UNIT])? { + self.must("systemctl", &["enable", "--now", UPDATE_PATH_UNIT])?; + self.note( + "ran", + format_args!("systemctl enable --now {UPDATE_PATH_UNIT}"), + ); + } + } + } + Platform::MacOs => { + let plist = PathBuf::from(format!( + "/Library/LaunchDaemons/{UPDATE_LAUNCHD_LABEL}.plist" + )); + let content = render_update_plist(&bin, &config_path, &request); + let plist_changed = self.write_file(&plist, content.as_bytes(), 0o644, None)?; + changed |= plist_changed; + if !self.opts.dry_run { + let target = format!("system/{UPDATE_LAUNCHD_LABEL}"); + let loaded = self.cmd("launchctl", &["print", &target])?; + if loaded && plist_changed { + self.must("launchctl", &["bootout", &target])?; + } + if !loaded || plist_changed { + let path = self.on_disk(&plist).to_string_lossy().into_owned(); + self.must("launchctl", &["bootstrap", "system", &path])?; + self.note("ran", format_args!("launchctl bootstrap system {path}")); + } + } + } + } + Ok(changed) + } + + /// The separate job user, on platforms that support one. + fn job_user(&self) -> Option<&str> { + match self.opts.platform { + Platform::Linux => self.opts.job_user.as_deref(), + Platform::MacOs => None, + } + } + + /// Who owns job-facing files: the job user, else the drone's user. + fn job_owner(&self) -> String { + self.job_user().unwrap_or(&self.opts.run_as).to_owned() + } + fn on_disk(&self, real: &Path) -> PathBuf { self.opts.root.join(real.strip_prefix("/").unwrap_or(real)) } @@ -353,12 +561,27 @@ impl Installer<'_> { "systemd is not running on this host.{hint}" ))); } + if self.job_user().is_some() && !self.cmd("setpriv", &["--version"])? { + return Err(InstallError::Preflight( + "setpriv (util-linux) is required to run jobs as a separate user; install it or pass --no-job-user".to_owned(), + )); + } + Ok(()) + } + + fn ensure_users(&mut self) -> Result<(), InstallError> { + let daemon = self.opts.run_as.clone(); + self.ensure_user(&daemon, &self.layout.data.clone())?; + if let Some(job) = self.job_user().map(str::to_owned) + && job != daemon + { + self.ensure_user(&job, &self.layout.job_home.clone())?; + } Ok(()) } - fn ensure_user(&mut self) -> Result<(), InstallError> { - let user = self.opts.run_as.clone(); - let exists = self.cmd("id", &["-u", &user])?; + fn ensure_user(&mut self, user: &str, home: &Path) -> Result<(), InstallError> { + let exists = self.cmd("id", &["-u", user])?; match self.opts.platform { Platform::MacOs => { if !exists { @@ -376,7 +599,7 @@ impl Installer<'_> { if self.opts.dry_run { self.note("would create", format_args!("user {user}")); } else { - let home = self.layout.data.to_string_lossy().into_owned(); + let home = home.to_string_lossy().into_owned(); self.must( "useradd", &[ @@ -386,13 +609,15 @@ impl Installer<'_> { &home, "--shell", "/usr/sbin/nologin", - &user, + user, ], )?; self.note("created", format_args!("user {user}")); } } - self.ensure_docker_group(&user, exists)?; + // The job user runs docker; the daemon also needs it to probe + // the engine and remove a lease's containers. + self.ensure_docker_group(user, exists)?; } } Ok(()) @@ -421,6 +646,39 @@ impl Installer<'_> { Ok(()) } + /// Hand an existing cache directory (for example a host's warm compile + /// cache) to the job user, recursively, unless it already belongs to them. + fn adopt_cache(&mut self, real: &Path, owner: &str) -> Result<(), InstallError> { + let path = self.on_disk(real); + let current = self + .exec + .run("stat", &["-c", "%U", &path.to_string_lossy()]) + .ok() + .filter(|o| o.success) + .map(|o| o.stdout); + if current.as_deref() == Some(owner) { + self.note( + "unchanged", + format_args!("cache {} already owned by {owner}", real.display()), + ); + return Ok(()); + } + self.changed = true; + if self.opts.dry_run { + self.note( + "would adopt", + format_args!("cache {} (chown -R {owner}:{owner})", real.display()), + ); + return Ok(()); + } + self.chown(real, owner, true)?; + self.note( + "adopted", + format_args!("cache {} (chown -R {owner}:{owner})", real.display()), + ); + Ok(()) + } + fn install_binary(&mut self) -> Result<(), InstallError> { let bytes = fs::read(&self.opts.binary).map_err(io_err(&self.opts.binary))?; let bin = self.layout.bin.clone(); @@ -455,8 +713,8 @@ impl Installer<'_> { ); return Ok(()); } - let user = self.opts.run_as.clone(); - let data = self.layout.data.clone(); + let user = self.job_owner(); + let data = self.layout.job_home.clone(); let staging = data.join("runner.new"); let tarball = data.join(self.opts.pin.filename()); let (staging_fs, tarball_fs) = (self.on_disk(&staging), self.on_disk(&tarball)); @@ -515,14 +773,22 @@ impl Installer<'_> { fn drone_config(&self) -> Result { let o = self.opts; + // A host's reclaim policy is edited in place; reinstalling must not + // reset it. + let reclaim = std::fs::read_to_string(self.on_disk(&self.layout.config())) + .ok() + .and_then(|t| serde_json::from_str::(&t).ok()) + .and_then(|v| v.get("reclaim").cloned()) + .and_then(|r| serde_json::from_value(r).ok()) + .unwrap_or_default(); let cfg = DroneConfig { queen_url: o.queen_url.clone(), token_file: Some(self.layout.token_file()), ca_cert_path: o.ca_cert.as_ref().map(|_| self.layout.ca_cert()), allow_insecure: o.allow_insecure, runner_dir: self.layout.runner(), - work_root: self.layout.data.join("work"), - cache_root: self.layout.data.join("cache"), + work_root: self.layout.work(), + cache_root: self.layout.cache_root.clone(), docker_bin: "docker".to_owned(), heartbeat_secs: 10, max_job_secs: 6 * 3600, @@ -532,6 +798,17 @@ impl Installer<'_> { max_leases: o.max_leases, reserve_mem_mib: crate::allocator::HostCaps::default().reserve_mem_mib, slots: None, + state_dir: Some(self.layout.data.clone()), + job_user: self.job_user().map(|name| JobUser { + name: name.to_owned(), + home: self.layout.job_home.clone(), + priv_bin: "setpriv".to_owned(), + }), + extra_labels: o.extra_labels.clone(), + runner_env: o.runner_env.clone(), + job_started_hook: o.job_started_hook.clone(), + job_completed_hook: o.job_completed_hook.clone(), + reclaim, caches: crate::drone::config::CachesConfig { git: crate::drone::config::GitCaches { mirrors: o.git_mirrors.clone(), @@ -565,7 +842,8 @@ impl Installer<'_> { let before = self.changed; match self.opts.platform { Platform::Linux => { - let unit = render_systemd_unit(&bin, &cfg, &data, &user, drain); + let unit = + render_systemd_unit(&bin, &cfg, &data, &user, drain, self.job_user().is_some()); let unit_changed = self.write_file(&service_file, unit.as_bytes(), 0o644, None)?; if self.opts.dry_run { return Ok(()); @@ -615,21 +893,36 @@ impl Installer<'_> { pub fn install(opts: &InstallOptions, exec: &dyn Exec) -> Result, InstallError> { let mut i = Installer { opts, - layout: Layout::new(opts.platform), + layout: Layout::new( + opts.platform, + opts.platform == Platform::Linux && opts.job_user.is_some(), + opts.cache_root.as_deref(), + ), exec, lines: Vec::new(), changed: false, }; i.preflight()?; let run_as = opts.run_as.clone(); - let owner = Some(run_as.as_str()); + let job_owner = i.job_owner(); - i.ensure_user()?; - let (config_dir, data) = (i.layout.config_dir.clone(), i.layout.data.clone()); - i.ensure_dir(&config_dir, 0o755, None)?; - i.ensure_dir(&data, 0o755, owner)?; - for sub in ["work", "cache"] { - i.ensure_dir(&data.join(sub), 0o755, owner)?; + i.ensure_users()?; + let (config_dir, data, job_home) = ( + i.layout.config_dir.clone(), + i.layout.data.clone(), + i.layout.job_home.clone(), + ); + // The config directory holds the enrollment token: the drone's user only. + i.ensure_dir(&config_dir, 0o700, Some(&run_as))?; + i.ensure_dir(&data, 0o755, Some(&run_as))?; + if job_home != data { + i.ensure_dir(&job_home, 0o755, Some(&job_owner))?; + } + let (work, cache_root) = (i.layout.work(), i.layout.cache_root.clone()); + i.ensure_dir(&work, 0o755, Some(&job_owner))?; + i.ensure_dir(&cache_root, 0o755, Some(&job_owner))?; + if opts.adopt_cache { + i.adopt_cache(&cache_root, &job_owner)?; } i.install_binary()?; i.install_runner()?; @@ -638,15 +931,18 @@ pub fn install(opts: &InstallOptions, exec: &dyn Exec) -> Result, In if let Some(ca) = &opts.ca_cert { let pem = fs::read(ca).map_err(io_err(ca))?; let target = i.layout.ca_cert(); - config_changed |= i.write_file(&target, &pem, 0o644, None)?; + config_changed |= i.write_file(&target, &pem, 0o600, Some(&run_as))?; } let token_file = i.layout.token_file(); let token = format!("{}\n", opts.token.expose()); - config_changed |= i.write_file(&token_file, token.as_bytes(), 0o600, owner)?; + config_changed |= i.write_file(&token_file, token.as_bytes(), 0o600, Some(&run_as))?; let config = i.drone_config()?; let config_path = i.layout.config(); - config_changed |= i.write_file(&config_path, config.as_bytes(), 0o644, None)?; + config_changed |= i.write_file(&config_path, config.as_bytes(), 0o600, Some(&run_as))?; + if !opts.no_self_update { + i.install_updater()?; + } i.install_service(config_changed)?; let mut lines = i.lines; @@ -673,6 +969,7 @@ mod tests { use super::*; use crate::drone::cleanup::testing::FakeExec; use crate::drone::pins::pin_for; + use std::sync::Arc; fn tarball_bytes() -> Vec { b"fake runner tarball".to_vec() @@ -711,11 +1008,20 @@ mod tests { max_mem_mib: Some(24576), max_leases: None, run_as: "borg".into(), + job_user: Some("borg-job".into()), + cache_root: None, + adopt_cache: false, + extra_labels: vec!["on-prem".into()], + runner_env: BTreeMap::new(), + job_started_hook: None, + job_completed_hook: None, root: e.root.clone(), platform: Platform::Linux, pin: e.pin.clone(), binary: e.binary.clone(), dry_run: false, + update_manifest_url: None, + no_self_update: false, skip_root_check: true, drain_timeout_secs: 3600, git_mirrors: vec!["example-org/app".into()], @@ -728,12 +1034,13 @@ mod tests { let x = FakeExec::default(); x.on("getent group docker", "docker:x:999:"); x.on("id -nG borg", "borg docker"); + x.on("id -nG borg-job", "borg-job docker"); x } /// Fake tar/curl side effects on the staging tree. - fn with_downloads(x: FakeExec) -> impl Exec { - struct Hooked(FakeExec); + fn with_downloads(x: impl Into>) -> impl Exec { + struct Hooked(Arc); impl Exec for Hooked { fn run(&self, program: &str, args: &[&str]) -> io::Result { match program { @@ -744,7 +1051,7 @@ mod tests { self.0.run(program, args) } } - Hooked(x) + Hooked(x.into()) } fn calls(x: &FakeExec) -> Vec { @@ -759,9 +1066,12 @@ mod tests { Path::new("/var/lib/borg"), "borg", 3600, + true, ); for line in [ "User=borg", + "AmbientCapabilities=CAP_SETUID CAP_SETGID CAP_SETPCAP", + "CapabilityBoundingSet=CAP_SETUID CAP_SETGID CAP_SETPCAP", "KillMode=mixed", "TimeoutStopSec=3720", "ExecStart=/usr/local/bin/borg drone run --config /etc/borg/drone.json", @@ -769,6 +1079,18 @@ mod tests { ] { assert!(unit.contains(line), "{line}\n{unit}"); } + let single = render_systemd_unit( + Path::new("/b"), + Path::new("/c"), + Path::new("/d"), + "borg", + 60, + false, + ); + assert!( + !single.contains("Capabilit"), + "no capabilities without a job user" + ); let plist = render_launchd_plist( Path::new("/usr/local/bin/borg"), Path::new("/c.json"), @@ -793,12 +1115,15 @@ mod tests { x.on("getent group docker", "docker:x:999:"); x.fail("systemctl is-enabled --quiet borg-drone.service"); x.fail("id -u borg"); + x.fail("id -u borg-job"); let ex = with_downloads(x); let first = install(&opts(&e), &ex).unwrap(); let text = first.join("\n"); for needle in [ - "created user borg", + "created user borg\n", + "created user borg-job", "added user borg to group docker", + "added user borg-job to group docker", "created /usr/local/bin/borg", "installed actions-runner 2.337.0", "created /etc/borg/drone.token", @@ -829,10 +1154,24 @@ mod tests { assert_eq!(cfg.caches.git.mirrors, ["example-org/app"]); assert!( e.root - .join("var/lib/borg/runner/.borg-runner-version") + .join("var/lib/borg-job/runner/.borg-runner-version") .exists() ); + // A reinstall keeps a reclaim policy the operator added to the config. + let path = e.root.join("etc/borg/drone.json"); + let mut v: serde_json::Value = + serde_json::from_str(&fs::read_to_string(&path).unwrap()).unwrap(); + v["reclaim"] = serde_json::json!({ + "enabled": true, "dry_run": false, + "roots": [{"path": "/srv/builds", "kind": "worktree"}] + }); + fs::write(&path, serde_json::to_string(&v).unwrap()).unwrap(); + let again = install(&opts(&e), &with_downloads(exec())).unwrap(); + let kept = DroneConfig::from_json(&fs::read_to_string(&path).unwrap()).unwrap(); + assert!(kept.reclaim.enabled && !kept.reclaim.dry_run, "{again:?}"); + assert_eq!(kept.reclaim.roots.len(), 1); + // Second run on the now-configured host. let y = exec(); let ey = with_downloads(y); @@ -870,7 +1209,8 @@ mod tests { let out = install(&o, &with_downloads(exec())).unwrap().join("\n"); assert!(out.contains("installed actions-runner 2.999.0"), "{out}"); let marker = - fs::read_to_string(e.root.join("var/lib/borg/runner/.borg-runner-version")).unwrap(); + fs::read_to_string(e.root.join("var/lib/borg-job/runner/.borg-runner-version")) + .unwrap(); assert!(marker.starts_with("2.999.0 ")); } @@ -882,7 +1222,7 @@ mod tests { let x = exec(); let err = install(&o, &with_downloads(x)).unwrap_err(); assert!(matches!(err, InstallError::Checksum { .. }), "{err}"); - assert!(!e.root.join("var/lib/borg/runner").exists()); + assert!(!e.root.join("var/lib/borg-job/runner").exists()); assert!( !e.root .join("etc/systemd/system/borg-drone.service") @@ -972,4 +1312,141 @@ mod tests { let err = install(&o, &y).unwrap_err().to_string(); assert!(err.contains("create it first"), "{err}"); } + + fn mode(root: &Path, rel: &str) -> u32 { + fs::metadata(root.join(rel)).unwrap().permissions().mode() & 0o777 + } + + #[test] + fn the_job_user_is_separate_and_cannot_reach_the_drones_files() { + let e = env(); + let x = Arc::new(FakeExec::default()); + x.fail("id -u borg"); + x.fail("id -u borg-job"); + install(&opts(&e), &with_downloads(Arc::clone(&x))).unwrap(); + + assert_eq!( + mode(&e.root, "etc/borg"), + 0o700, + "config dir is the drone's alone" + ); + assert_eq!(mode(&e.root, "etc/borg/drone.token"), 0o600); + assert_eq!(mode(&e.root, "etc/borg/drone.json"), 0o600); + let cfg = DroneConfig::from_json( + &fs::read_to_string(e.root.join("etc/borg/drone.json")).unwrap(), + ) + .unwrap(); + let job = cfg.job_user.expect("job user configured"); + assert_eq!( + (job.name.as_str(), job.home.as_path()), + ("borg-job", Path::new("/var/lib/borg-job")) + ); + assert_eq!(cfg.work_root, Path::new("/var/lib/borg-job/work")); + assert_eq!(cfg.cache_root, Path::new("/var/lib/borg-job/cache")); + assert_eq!(cfg.extra_labels, ["on-prem"]); + + let calls = x.calls(); + let root = e.root.display(); + for owned in [ + format!("chown borg:borg {root}/etc/borg"), + format!("chown borg:borg {root}/etc/borg/drone.token"), + format!("chown borg:borg {root}/etc/borg/drone.json"), + format!("chown borg-job:borg-job {root}/var/lib/borg-job/work"), + format!("chown borg-job:borg-job {root}/var/lib/borg-job/cache"), + format!("chown -R borg-job:borg-job {root}/var/lib/borg-job/runner"), + ] { + assert!(calls.contains(&owned), "missing `{owned}` in {calls:?}"); + } + assert!( + !calls.iter().any(|c| c.starts_with("chown") + && c.contains("borg-job") + && c.contains("/etc/borg")), + "the job user is never given the drone's files: {calls:?}" + ); + assert!(calls.contains(&"usermod -aG docker borg-job".to_owned())); + let unit = + fs::read_to_string(e.root.join("etc/systemd/system/borg-drone.service")).unwrap(); + assert!( + unit.contains("User=borg\n") && unit.contains("AmbientCapabilities="), + "{unit}" + ); + } + + #[test] + fn without_a_job_user_everything_is_one_account_and_unit_has_no_capabilities() { + let e = env(); + let mut o = opts(&e); + o.job_user = None; + let x = Arc::new(exec()); + let out = install(&o, &with_downloads(Arc::clone(&x))) + .unwrap() + .join("\n"); + assert!(!out.contains("borg-job"), "{out}"); + let cfg = DroneConfig::from_json( + &fs::read_to_string(e.root.join("etc/borg/drone.json")).unwrap(), + ) + .unwrap(); + assert!(cfg.job_user.is_none()); + assert_eq!(cfg.work_root, Path::new("/var/lib/borg/work")); + let unit = + fs::read_to_string(e.root.join("etc/systemd/system/borg-drone.service")).unwrap(); + assert!(!unit.contains("Capabilit")); + assert!(!x.calls().iter().any(|c| c.starts_with("setpriv"))); + } + + #[test] + fn a_missing_setpriv_is_a_clear_preflight_error() { + let e = env(); + let x = exec(); + x.fail("setpriv --version"); + let err = install(&opts(&e), &x).unwrap_err().to_string(); + assert!( + err.contains("setpriv") && err.contains("--no-job-user"), + "{err}" + ); + } + + #[test] + fn an_existing_cache_can_be_adopted_once() { + let e = env(); + let mut o = opts(&e); + o.cache_root = Some("/srv/example-cache".into()); + o.adopt_cache = true; + let cache = e.root.join("srv/example-cache"); + fs::create_dir_all(&cache).unwrap(); + fs::write(cache.join("warm-entry"), "21 GB of sccache").unwrap(); + + let x = Arc::new(exec()); + let first = install(&o, &with_downloads(Arc::clone(&x))) + .unwrap() + .join("\n"); + assert!( + first.contains("adopted cache /srv/example-cache (chown -R borg-job:borg-job)"), + "{first}" + ); + assert!( + x.calls() + .contains(&format!("chown -R borg-job:borg-job {}", cache.display())) + ); + assert_eq!( + fs::read_to_string(cache.join("warm-entry")).unwrap(), + "21 GB of sccache", + "contents are kept" + ); + let cfg = DroneConfig::from_json( + &fs::read_to_string(e.root.join("etc/borg/drone.json")).unwrap(), + ) + .unwrap(); + assert_eq!(cfg.cache_root, Path::new("/srv/example-cache")); + + // Already owned by the job user: nothing to do on the next run. + let y = exec(); + y.on(&format!("stat -c %U {}", cache.display()), "borg-job"); + let second = install(&o, &with_downloads(y)).unwrap().join("\n"); + assert!( + second.contains("cache /srv/example-cache already owned by borg-job"), + "{second}" + ); + assert!(second.contains("nothing to change"), "{second}"); + } } diff --git a/src/drone/jobuser.rs b/src/drone/jobuser.rs new file mode 100644 index 0000000..b848276 --- /dev/null +++ b/src/drone/jobuser.rs @@ -0,0 +1,162 @@ +//! Running work as a separate, unprivileged job user. +//! +//! The drone daemon holds the enrollment token and talks to the queen; a lease +//! runs code from pull requests. If both were one user, a job could read the +//! token and impersonate the drone. With a job user configured, every lease +//! (and every file operation on job-owned directories) runs as that user, which +//! has no access to the drone's config or token. +//! +//! The daemon changes user with `setpriv`, which needs `CAP_SETUID`, +//! `CAP_SETGID` and `CAP_SETPCAP`; the systemd unit grants exactly those as +//! ambient capabilities and `setpriv` drops them again before the job starts, +//! so the job holds no capability. + +use std::io; +use std::path::PathBuf; +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; + +use crate::drone::exec::{Exec, Output}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct JobUser { + /// Account (and primary group) leases run as. Numeric ids are accepted. + pub name: String, + /// The job user's home directory, exported to jobs as `HOME`. + pub home: PathBuf, + /// The `setpriv` binary (util-linux). + #[serde(default = "default_priv")] + pub priv_bin: String, +} + +fn default_priv() -> String { + "setpriv".to_owned() +} + +impl JobUser { + /// `setpriv` arguments that become the job user, load its supplementary + /// groups (for example `docker`) and drop every capability, ending in `--`. + pub fn setpriv_args(&self) -> Vec { + vec![ + "--reuid".to_owned(), + self.name.clone(), + "--regid".to_owned(), + self.name.clone(), + "--init-groups".to_owned(), + "--inh-caps=-all".to_owned(), + "--ambient-caps=-all".to_owned(), + "--bounding-set=-all".to_owned(), + "--".to_owned(), + ] + } + + pub fn validate(&self) -> Result<(), String> { + let name_ok = !self.name.is_empty() + && self.name.len() <= 32 + && self + .name + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.')); + if !name_ok { + return Err(format!( + "job_user.name {:?} must be a plain account name or numeric id", + self.name + )); + } + if !self.home.is_absolute() { + return Err("job_user.home must be an absolute path".to_owned()); + } + if self.priv_bin.is_empty() { + return Err("job_user.priv_bin must not be empty".to_owned()); + } + Ok(()) + } +} + +/// An [`Exec`] that runs every command as the job user. +pub struct JobExec { + inner: Arc, + user: JobUser, +} + +impl JobExec { + pub fn new(inner: Arc, user: JobUser) -> Self { + Self { inner, user } + } +} + +impl Exec for JobExec { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + let mut argv = self.user.setpriv_args(); + argv.push(program.to_owned()); + argv.extend(args.iter().map(|a| (*a).to_owned())); + let argv: Vec<&str> = argv.iter().map(String::as_str).collect(); + self.inner.run(&self.user.priv_bin, &argv) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::drone::cleanup::testing::FakeExec; + + fn user() -> JobUser { + JobUser { + name: "borg-job".into(), + home: "/var/lib/borg-job".into(), + priv_bin: "setpriv".into(), + } + } + + #[test] + fn commands_become_setpriv_invocations_that_drop_capabilities() { + let fake = Arc::new(FakeExec::default()); + let exec = JobExec::new(fake.clone(), user()); + exec.run("rm", &["-rf", "--", "/w/slot-0"]).unwrap(); + assert_eq!( + fake.calls(), + [ + "setpriv --reuid borg-job --regid borg-job --init-groups --inh-caps=-all --ambient-caps=-all --bounding-set=-all -- rm -rf -- /w/slot-0" + ] + ); + } + + #[test] + fn validation_rejects_bad_names_and_relative_homes() { + assert!(user().validate().is_ok()); + for bad in [ + JobUser { + name: String::new(), + ..user() + }, + JobUser { + name: "a b".into(), + ..user() + }, + JobUser { + name: "x;rm".into(), + ..user() + }, + JobUser { + home: "relative".into(), + ..user() + }, + JobUser { + priv_bin: String::new(), + ..user() + }, + ] { + assert!(bad.validate().is_err(), "{bad:?}"); + } + assert!( + JobUser { + name: "65534".into(), + ..user() + } + .validate() + .is_ok() + ); + } +} diff --git a/src/drone/mod.rs b/src/drone/mod.rs index b3c89e3..0b9958a 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -10,13 +10,17 @@ pub mod cleanup; pub mod client; pub mod config; pub mod exec; +pub mod fsops; pub mod install; +pub mod jobuser; pub mod pins; pub mod runner; +pub mod selfupdate; pub mod slots; use std::collections::HashSet; use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -29,12 +33,16 @@ use cleanup::{Scope, cleanup_docker, remove_lease_dir, sweep_work_root, valid_le use client::{ClientError, QueenClient}; use config::DroneConfig; use exec::Exec; +use fsops::{ExecFs, FsOps, LocalFs}; +use jobuser::JobExec; use runner::{RunSpec, run_lease}; +use selfupdate::{Handled, Plan, SelfUpdate}; use slots::{SlotGuard, Slots, slot_dir, slot_name}; /// Attempts to deliver a lease result before giving up (the lease then -/// expires on the queen's side). -const RESULT_ATTEMPTS: u32 = 5; +/// expires on the queen's side). The waits grow to a minute, so a queen that +/// is down for a few minutes does not lose results. +const RESULT_ATTEMPTS: u32 = 20; type Held = Arc>>; @@ -57,11 +65,19 @@ pub struct Drone { client: Arc, host: Arc, exec: Arc, + /// Runs commands as the job user (or as the drone when none is configured). + job_exec: Arc, + /// File operations on job-owned directories. + fs_ops: Arc, held: Held, abort: tokio::sync::watch::Sender, proc_root: Arc, caches: Arc, slots: Arc, + /// Self-update hand-off; present when `state_dir` is configured. + selfupdate: Option, + /// The queen wants another version: stop taking leases until idle. + update_wanted: AtomicBool, } impl Drone { @@ -74,16 +90,36 @@ impl Drone { let (abort, _) = tokio::sync::watch::channel(false); let slot_count = cfg.slot_count(); let cfg = Arc::new(cfg); + let (job_exec, fs_ops): (Arc, Arc) = match &cfg.job_user { + Some(user) => { + let job_exec: Arc = + Arc::new(JobExec::new(Arc::clone(&exec), user.clone())); + let fs_ops: Arc = Arc::new(ExecFs::new(Arc::clone(&job_exec))); + (job_exec, fs_ops) + } + None => { + tracing::warn!( + "no job_user configured: jobs run as the drone's own user and can read its token" + ); + (Arc::clone(&exec), Arc::new(LocalFs)) + } + }; let caches = Arc::new(Caches::new( Arc::clone(&cfg), - Arc::clone(&exec), + Arc::clone(&job_exec), + Arc::clone(&fs_ops), Arc::clone(&host), )); + let selfupdate = cfg.state_dir.clone().map(SelfUpdate::new); Self { cfg, + selfupdate, + update_wanted: AtomicBool::new(false), client: Arc::new(client), host, exec, + job_exec, + fs_ops, caches, slots: Slots::new(slot_count), held: Held::default(), @@ -110,13 +146,14 @@ impl Drone { /// Remove leftovers of a previous run: no lease is held yet, so every /// lease directory and borg-labelled docker resource is stale. pub async fn startup_sweep(&self) { - let (work_root, exec, docker) = ( + let (work_root, exec, docker, fs_ops) = ( self.cfg.work_root.clone(), Arc::clone(&self.exec), self.cfg.docker_bin.clone(), + Arc::clone(&self.fs_ops), ); let swept = tokio::task::spawn_blocking(move || { - let dirs = sweep_work_root(&work_root); + let dirs = sweep_work_root(fs_ops.as_ref(), &work_root); let docker = cleanup_docker(exec.as_ref(), &docker, Scope::AnyLease); (dirs, docker) }) @@ -140,6 +177,9 @@ impl Drone { accepting: bool, tasks: &mut JoinSet<()>, ) -> Result { + // Between jobs the drone obeys `desired_version`: once the queen wants + // another version it stops accepting and updates when idle. + let accepting = accepting && !self.update_wanted.load(Ordering::Relaxed); let host = Arc::clone(&self.host); let capacity = tokio::task::spawn_blocking(move || probe::probe(host.as_ref())) .await @@ -154,11 +194,21 @@ impl Drone { active_leases: self.held_ids(), cache_stats: self.caches.stats(), version: Some(env!("CARGO_PKG_VERSION").to_owned()), + extra_labels: self.cfg.extra_labels.clone(), + bids: Vec::new(), + reclaim: None, + update_report: self + .selfupdate + .as_ref() + .and_then(SelfUpdate::pending_report), + active_units: Vec::new(), + env_facts: None, }; let resp = self.client.heartbeat(&req).await?; + let defer = self.after_heartbeat(&req, resp.desired_version.as_deref()); let mut started = 0; for grant in resp.leases { - if !accepting || !self.admit(&grant.lease_id) { + if !accepting || defer || !self.admit(&grant.lease_id) { continue; } let held = HeldGuard { @@ -176,6 +226,8 @@ impl Drone { cfg: Arc::clone(&self.cfg), client: Arc::clone(&self.client), exec: Arc::clone(&self.exec), + job_exec: Arc::clone(&self.job_exec), + fs_ops: Arc::clone(&self.fs_ops), abort: self.abort.subscribe(), proc_root: Arc::clone(&self.proc_root), caches: Arc::clone(&self.caches), @@ -188,6 +240,60 @@ impl Drone { Ok(started) } + /// What a successful heartbeat means for self-update: the queen has the + /// updater's report, the running binary is demonstrably working, and the + /// queen may want another version (older ones included). Returns true when + /// leases granted in this response must be left unstarted because an update + /// is being handed to the updater. + fn after_heartbeat(&self, req: &HeartbeatRequest, desired: Option<&str>) -> bool { + let Some(su) = &self.selfupdate else { + return false; + }; + let own = env!("CARGO_PKG_VERSION"); + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + if let Err(e) = su.record_health(own, now) { + tracing::warn!("write health record: {e}"); + } + if req.update_report.is_some() { + su.acknowledge_report(); + } + let held = self.held_ids().len(); + let desired = match selfupdate::plan(desired, own, held) { + Plan::Stay => { + self.update_wanted.store(false, Ordering::Relaxed); + return false; + } + Plan::FinishJobsFirst => { + // No new jobs meanwhile. + self.update_wanted.store(true, Ordering::Relaxed); + return false; + } + Plan::UpdateNow => desired.unwrap_or(own), + }; + match su.handle(desired, own) { + Ok(Handled::Requested) => { + tracing::info!("update to {desired} requested"); + self.update_wanted.store(true, Ordering::Relaxed); + true + } + Ok(Handled::InFlight) => { + self.update_wanted.store(true, Ordering::Relaxed); + true + } + Ok(Handled::Blocked) => { + tracing::warn!("not retrying the failed update to {desired}"); + self.update_wanted.store(false, Ordering::Relaxed); + false + } + Ok(Handled::UpToDate) => false, + Err(e) => { + tracing::warn!("request update to {desired}: {e}"); + self.update_wanted.store(false, Ordering::Relaxed); + false + } + } + } + /// Whether a granted lease may start: well-formed, not already running, /// and within the local lease cap. Claims the id if so. fn admit(&self, lease_id: &str) -> bool { @@ -229,11 +335,24 @@ impl Drone { every.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); tokio::pin!(shutdown); // `interval.tick`, `shutdown` and `JoinSet::join_next` are cancel-safe. + // While the queen is unreachable, running leases carry on, no new lease + // arrives, and heartbeats are retried with backoff. + let mut failures = 0u32; + let mut retry_at = tokio::time::Instant::now(); loop { tokio::select! { _ = every.tick() => { - if let Err(e) = self.heartbeat_once(true, &mut tasks).await { - tracing::warn!("heartbeat failed: {e}"); + if tokio::time::Instant::now() < retry_at { + continue; + } + match self.heartbeat_once(true, &mut tasks).await { + Ok(_) => failures = 0, + Err(e) => { + failures = failures.saturating_add(1); + let wait = selfupdate::backoff(self.cfg.heartbeat_secs, failures); + retry_at = tokio::time::Instant::now() + Duration::from_secs(wait); + tracing::warn!("heartbeat failed ({failures} in a row), retrying in {wait}s: {e}"); + } } } Some(_) = tasks.join_next(), if !tasks.is_empty() => {} @@ -271,6 +390,8 @@ struct LeaseContext { cfg: Arc, client: Arc, exec: Arc, + job_exec: Arc, + fs_ops: Arc, abort: tokio::sync::watch::Receiver, proc_root: Arc, caches: Arc, @@ -287,7 +408,8 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { // trouble never blocks a job; it just runs without the speedup. let slot = ctx.slot.index(); let lease_dir = slot_dir(&ctx.cfg.work_root, slot); - let mut extra_env = ctx.caches.prepare_lease(&lease_dir).await; + let mut extra_env = ctx.cfg.runner_environment(); + extra_env.extend(ctx.caches.prepare_lease(&lease_dir).await); extra_env.push(("BORG_SLOT".to_owned(), slot.to_string())); let spec = RunSpec { lease_id: &grant.lease_id, @@ -297,21 +419,32 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { cache_root: &ctx.cfg.cache_root, extra_env: &extra_env, max_job: Duration::from_secs(ctx.cfg.max_job_secs), + job_user: ctx.cfg.job_user.as_ref(), proc_root: &ctx.proc_root, }; - let outcome = run_lease(&spec, ctx.exec.as_ref(), ctx.abort.clone()).await; + let outcome = run_lease( + &spec, + ctx.job_exec.as_ref(), + Arc::clone(&ctx.fs_ops), + ctx.abort.clone(), + ) + .await; // Runs whether the job passed, failed, timed out, was aborted or never started. - let (work_root, id, slot_name, docker, exec) = ( + let (work_root, id, slot_name, docker, exec, fs_ops) = ( ctx.cfg.work_root.clone(), grant.lease_id.clone(), slot_name(slot), ctx.cfg.docker_bin.clone(), Arc::clone(&ctx.exec), + Arc::clone(&ctx.fs_ops), ); let cleaned = tokio::task::spawn_blocking(move || { let docker = cleanup_docker(exec.as_ref(), &docker, Scope::Lease(&id)); - (remove_lease_dir(&work_root, &slot_name), docker) + ( + remove_lease_dir(fs_ops.as_ref(), &work_root, &slot_name), + docker, + ) }) .await; match cleaned { @@ -358,7 +491,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { "reporting lease {} failed (attempt {attempt}): {e}", grant.lease_id ); - tokio::time::sleep(Duration::from_secs(u64::from(attempt))).await; + tokio::time::sleep(Duration::from_secs((u64::from(attempt) * 5).min(60))).await; } } } diff --git a/src/drone/runner.rs b/src/drone/runner.rs index cac7a5e..5355d4f 100644 --- a/src/drone/runner.rs +++ b/src/drone/runner.rs @@ -13,12 +13,15 @@ use std::io; use std::os::unix::fs::symlink; use std::path::{Path, PathBuf}; use std::process::Stdio; +use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::process::Command; use crate::drone::cleanup::{docker_prefix, valid_lease_id}; use crate::drone::exec::Exec; +use crate::drone::fsops::FsOps; +use crate::drone::jobuser::JobUser; use crate::secret::Secret; const RSS_SAMPLE_EVERY: Duration = Duration::from_secs(2); @@ -34,6 +37,9 @@ pub struct RunSpec<'a> { /// Extra job environment (cache directories, git mirrors). pub extra_env: &'a [(String, String)], pub max_job: Duration, + /// Run the job as this user (see [`crate::drone::jobuser`]); `None` runs it + /// as the drone's own user. + pub job_user: Option<&'a JobUser>, /// `/proc` (a different root is used by tests). pub proc_root: &'a Path, } @@ -136,6 +142,7 @@ fn kill_group(exec: &dyn Exec, pgid: u32) { pub async fn run_lease( spec: &RunSpec<'_>, exec: &dyn Exec, + fs_ops: Arc, mut abort: tokio::sync::watch::Receiver, ) -> io::Result { if !valid_lease_id(spec.lease_id) { @@ -149,19 +156,37 @@ pub async fn run_lease( tokio::task::spawn_blocking(move || { // A slot is reused, so make sure nothing of an earlier lease remains. if dst.symlink_metadata().is_ok() { - fs::remove_dir_all(&dst)?; + fs_ops.remove_all(&dst)?; } - clone_tree(&src, &dst) + fs_ops.clone_tree(&src, &dst) }) .await .map_err(io::Error::other)??; let prefix = docker_prefix(spec.lease_id); + let run_sh = lease_dir.join("run.sh"); + // With a job user the drone starts `setpriv ... run.sh`, which becomes the + // runner under the job user's uid with no capabilities. + let (program, leading): (PathBuf, Vec) = match spec.job_user { + Some(user) => { + let mut args = user.setpriv_args(); + args.push(run_sh.to_string_lossy().into_owned()); + (PathBuf::from(&user.priv_bin), args) + } + None => (run_sh, Vec::new()), + }; #[allow( clippy::disallowed_methods, - reason = "starts the pinned actions-runner run.sh from the drone-owned per-lease directory; its only argument is the queen-issued JIT config" + reason = "starts the pinned actions-runner run.sh (directly, or through setpriv as the job user) from the drone-owned per-lease directory; its only secret input is the queen-issued JIT config, passed in the environment" )] - let mut command = Command::new(lease_dir.join("run.sh")); + let mut command = Command::new(&program); + command.args(&leading); + if let Some(user) = spec.job_user { + command + .env("HOME", &user.home) + .env("USER", &user.name) + .env("LOGNAME", &user.name); + } command // The runner reads `ACTIONS_RUNNER_INPUT_` as a fallback for each // command-line argument (Runner.Listener `CommandSettings`, v2.337.0), @@ -180,7 +205,9 @@ pub async fn run_lease( .env("DISABLE_RUNNER_UPDATE", "1") .stdin(Stdio::null()) .process_group(0) - .kill_on_drop(true); + // With a job user the drone cannot signal the job directly; its group + // is killed through `exec` instead. + .kill_on_drop(spec.job_user.is_none()); let mut child = spawn_retrying(&mut command).await?; let pgid = child.id(); @@ -260,10 +287,18 @@ mod tests { lease_dir: &work.join("slot-0"), cache_root: &dir.join("cache"), extra_env: &[], + job_user: None, max_job, proc_root: Path::new("/proc"), }; - let out = run_lease(&spec, &SystemExec, rx).await.unwrap(); + let out = run_lease( + &spec, + &SystemExec, + Arc::new(crate::drone::fsops::LocalFs), + rx, + ) + .await + .unwrap(); (out, work.join("slot-0")) } @@ -368,6 +403,7 @@ mod tests { lease_dir: &work.join("slot-0"), cache_root: &cache, extra_env: &[], + job_user: None, max_job: Duration::from_secs(60), proc_root: Path::new("/proc"), }; @@ -376,7 +412,14 @@ mod tests { let _ = tx.send(true); tx }); - let out = run_lease(&spec, &SystemExec, rx).await.unwrap(); + let out = run_lease( + &spec, + &SystemExec, + Arc::new(crate::drone::fsops::LocalFs), + rx, + ) + .await + .unwrap(); let _tx = abort.await.unwrap(); assert!(out.aborted && !out.timed_out); } @@ -394,10 +437,20 @@ mod tests { lease_dir: &dir.path().join("work/slot-0"), cache_root: dir.path(), extra_env: &[], + job_user: None, max_job: Duration::from_secs(1), proc_root: Path::new("/proc"), }; - assert!(run_lease(&spec, &SystemExec, rx).await.is_err()); + assert!( + run_lease( + &spec, + &SystemExec, + Arc::new(crate::drone::fsops::LocalFs), + rx + ) + .await + .is_err() + ); assert!(!dir.path().join("work").exists()); } diff --git a/src/drone/selfupdate.rs b/src/drone/selfupdate.rs new file mode 100644 index 0000000..671a445 --- /dev/null +++ b/src/drone/selfupdate.rs @@ -0,0 +1,235 @@ +//! The drone's half of self-update. +//! +//! The drone cannot replace its own binary or restart itself: it runs as an +//! unprivileged user, and stopping it is what drains it. So it only *asks*: it +//! writes `update-request` into its state directory, and a root-owned path unit +//! (systemd) or watch (launchd) starts `borg update --from-request`, which +//! downloads, verifies, drains, swaps, restarts and rolls back if needed. The +//! updater leaves `update-result.json`; the restarted drone reports it to the +//! queen in its next heartbeat. Health evidence for the updater is the +//! `health.json` this module writes after every successful heartbeat. + +use std::fs; +use std::io; +use std::path::PathBuf; + +use crate::api::UpdateReport; +use crate::update::health::{self, HealthRecord}; +use crate::update::{Request, UpdateResult, request_path, result_path}; + +/// What to do about the queen's `desired_version` right now. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Plan { + /// Nothing to do: no instruction, or already on that version. + Stay, + /// Leases are held: take no new ones, finish these, update after. + FinishJobsFirst, + /// Idle: hand the update to the updater now. + UpdateNow, +} + +/// Between jobs the drone obeys `desired_version`, in either direction. +pub fn plan(desired: Option<&str>, own_version: &str, leases_held: usize) -> Plan { + match desired { + None => Plan::Stay, + Some(d) if d == own_version => Plan::Stay, + Some(_) if leases_held > 0 => Plan::FinishJobsFirst, + Some(_) => Plan::UpdateNow, + } +} + +/// A request the updater has not consumed after this long is rewritten. +const STALE_REQUEST: std::time::Duration = std::time::Duration::from_secs(30 * 60); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Handled { + /// Already running the requested version. + UpToDate, + /// A request was written; the updater takes it from here. + Requested, + /// An update is already in flight; nothing more is asked. + InFlight, + /// The same target already failed or rolled back here; not retried. + Blocked, +} + +pub struct SelfUpdate { + state_dir: PathBuf, +} + +impl SelfUpdate { + pub fn new(state_dir: PathBuf) -> Self { + Self { state_dir } + } + + fn result(&self) -> Option { + serde_json::from_slice(&fs::read(result_path(&self.state_dir)).ok()?).ok() + } + + /// The updater's outcome, until the queen has acknowledged it. + pub fn pending_report(&self) -> Option { + self.result().map(|r| UpdateReport { + target: r.target, + state: r.state, + detail: r.detail, + }) + } + + /// The queen has the report; forget it. + pub fn acknowledge_report(&self) { + let _ = fs::remove_file(result_path(&self.state_dir)); + } + + /// Act on the queen's `desired_version`. Any other version is a target, + /// older ones included. At most one update is in flight: while a request + /// file exists nothing further is written. + pub fn handle(&self, desired: &str, own_version: &str) -> io::Result { + if desired == own_version { + return Ok(Handled::UpToDate); + } + let request = request_path(&self.state_dir); + if let Ok(meta) = fs::metadata(&request) { + let age = meta.modified().ok().and_then(|m| m.elapsed().ok()); + if age.is_none_or(|a| a < STALE_REQUEST) { + return Ok(Handled::InFlight); + } + // The updater never picked it up (or died); ask again. + fs::remove_file(&request)?; + } + if self.result().is_some_and(|r| { + r.target == desired && matches!(r.state.as_str(), "rolled_back" | "failed") + }) { + return Ok(Handled::Blocked); + } + let tmp = request.with_extension("tmp"); + fs::write( + &tmp, + serde_json::to_vec(&Request { + version: desired.to_owned(), + }) + .map_err(io::Error::other)?, + )?; + fs::rename(&tmp, &request)?; + Ok(Handled::Requested) + } + + /// Record a successful heartbeat for the updater's health check. + pub fn record_health(&self, version: &str, now: i64) -> io::Result<()> { + health::write( + &self.health_path(), + &HealthRecord { + version: version.to_owned(), + at: now, + pid: std::process::id(), + }, + ) + } + + pub fn health_path(&self) -> PathBuf { + self.state_dir.join("health.json") + } +} + +/// Delay before the next heartbeat after `failures` consecutive failures: the +/// normal interval, doubling per failure, capped at a minute, so a queen that +/// is down is retried gently and one that comes back is found quickly. +pub fn backoff(interval_secs: u64, failures: u32) -> u64 { + interval_secs + .saturating_mul(1u64 << failures.min(6)) + .min(60) + .max(interval_secs) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_directive_writes_one_request_and_then_waits() { + let dir = tempfile::tempdir().unwrap(); + let su = SelfUpdate::new(dir.path().to_owned()); + assert_eq!(su.handle("1.0.0", "1.0.0").unwrap(), Handled::UpToDate); + assert!(!request_path(dir.path()).exists()); + + assert_eq!(su.handle("2.0.0", "1.0.0").unwrap(), Handled::Requested); + let request: Request = + serde_json::from_slice(&fs::read(request_path(dir.path())).unwrap()).unwrap(); + assert_eq!(request.version, "2.0.0"); + assert_eq!( + su.handle("2.0.0", "1.0.0").unwrap(), + Handled::InFlight, + "at most one in flight" + ); + assert_eq!(su.handle("3.0.0", "1.0.0").unwrap(), Handled::InFlight); + } + + #[test] + fn a_failed_target_is_not_retried_but_a_newer_one_is() { + let dir = tempfile::tempdir().unwrap(); + let su = SelfUpdate::new(dir.path().to_owned()); + let result = UpdateResult { + target: "2.0.0".into(), + state: "rolled_back".into(), + detail: "no health".into(), + at: 1, + }; + fs::write( + result_path(dir.path()), + serde_json::to_vec(&result).unwrap(), + ) + .unwrap(); + assert_eq!(su.pending_report().unwrap().state, "rolled_back"); + assert_eq!(su.handle("2.0.0", "1.0.0").unwrap(), Handled::Blocked); + assert_eq!(su.handle("2.1.0", "1.0.0").unwrap(), Handled::Requested); + su.acknowledge_report(); + assert_eq!(su.pending_report(), None); + } + + #[test] + fn health_evidence_is_written_for_the_updater() { + let dir = tempfile::tempdir().unwrap(); + let su = SelfUpdate::new(dir.path().to_owned()); + su.record_health("1.2.3", 77).unwrap(); + let rec = health::read(&su.health_path()).unwrap(); + assert_eq!((rec.version.as_str(), rec.at), ("1.2.3", 77)); + } + + #[test] + fn the_drone_updates_only_between_jobs_and_in_either_direction() { + assert_eq!(plan(None, "1.0.0", 3), Plan::Stay); + assert_eq!(plan(Some("1.0.0"), "1.0.0", 0), Plan::Stay); + assert_eq!(plan(Some("2.0.0"), "1.0.0", 2), Plan::FinishJobsFirst); + assert_eq!(plan(Some("2.0.0"), "1.0.0", 0), Plan::UpdateNow); + assert_eq!( + plan(Some("0.9.0"), "1.0.0", 0), + Plan::UpdateNow, + "downgrades too" + ); + } + + #[test] + fn a_stale_request_is_asked_again() { + let dir = tempfile::tempdir().unwrap(); + let su = SelfUpdate::new(dir.path().to_owned()); + assert_eq!(su.handle("2.0.0", "1.0.0").unwrap(), Handled::Requested); + let old = std::time::SystemTime::now() - std::time::Duration::from_secs(3600); + fs::File::options() + .write(true) + .open(request_path(dir.path())) + .unwrap() + .set_modified(old) + .unwrap(); + assert_eq!(su.handle("2.0.0", "1.0.0").unwrap(), Handled::Requested); + assert_eq!(su.handle("2.0.0", "1.0.0").unwrap(), Handled::InFlight); + } + + #[test] + fn heartbeat_backoff_grows_and_is_capped() { + assert_eq!(backoff(10, 0), 10); + assert_eq!(backoff(10, 1), 20); + assert_eq!(backoff(10, 2), 40); + assert_eq!(backoff(10, 3), 60); + assert_eq!(backoff(10, 30), 60); + assert_eq!(backoff(90, 4), 90, "never below the configured interval"); + } +} diff --git a/src/envcontract.rs b/src/envcontract.rs new file mode 100644 index 0000000..62746c8 --- /dev/null +++ b/src/envcontract.rs @@ -0,0 +1,1024 @@ +//! The build environment contract: facts a drone reports about its build +//! environment, and the queen's pure admission check over them. +//! +//! Before work (a CI runner lease or a work unit) is placed on a drone, four +//! checks are evaluated against what the drone last reported: +//! +//! 1. **compiler_cache**: `RUSTC_WRAPPER` resolves to an allowed wrapper +//! (sccache, or the single-flight wrapper in front of it), `SCCACHE_DIR` is +//! set and the sccache server is reachable; +//! 2. **toolchain**: when the work declares a Rust toolchain (the class config, +//! or `rust-toolchain.toml` at the unit's tree), the drone has it installed; +//! 3. **cache_store**: the sccache stats call works, the error rate is under a +//! threshold, the cache directory has free space above a floor, and a write +//! works. A pathological hit rate is *reported* as a warning, never blocks; +//! 4. **resources**: memory available now covers the class's memory, and live +//! CPU pressure (PSI) is under a ceiling. (Free capacity after active leases +//! is enforced by the allocator itself.) +//! +//! The drone is a dumb prober: it reports *facts* ([`EnvFacts`]), every one of +//! which may be unknown (`None`), cached about a minute on its side. Thresholds +//! live in the queen's config, so they change without touching drones. Each +//! check is `required`, `advisory` or `off`, globally and per class. A required +//! check that fails or is unknown makes the drone ineligible for that class; +//! [`evaluate`] is a pure function of the settings, the facts and the clock, so +//! the allocator stays pure (it only sees `DroneState::blocked_classes`). +//! +//! Nothing is enforced unless `env_contract.enabled`; a drone that has never +//! reported is skipped unless `require_report` is set, so enabling the contract +//! does not cut off drones that predate it. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::classes::{Class, ClassSpec, Classes}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CheckKind { + CompilerCache, + Toolchain, + CacheStore, + Resources, +} + +impl CheckKind { + pub const ALL: [Self; 4] = [ + Self::CompilerCache, + Self::Toolchain, + Self::CacheStore, + Self::Resources, + ]; + + pub fn as_str(self) -> &'static str { + match self { + Self::CompilerCache => "compiler_cache", + Self::Toolchain => "toolchain", + Self::CacheStore => "cache_store", + Self::Resources => "resources", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Mode { + /// A failing or unknown check makes the drone ineligible. + Required, + /// Evaluated and reported, never blocks. + Advisory, + Off, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CheckState { + Pass, + Unknown, + Fail, +} + +impl CheckState { + pub fn as_str(self) -> &'static str { + match self { + Self::Pass => "pass", + Self::Unknown => "unknown", + Self::Fail => "fail", + } + } + + pub const ALL: [Self; 3] = [Self::Pass, Self::Fail, Self::Unknown]; +} + +/// An installed Rust toolchain, as `rustc -vV` reports it for the channel. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Toolchain { + /// The channel as `rust-toolchain.toml` names it (`1.95.0`, `stable`, + /// `nightly-2026-01-01`). + pub channel: String, + /// The `release:` line of `rustc -vV` for that channel. + pub rustc_version: String, +} + +/// What a drone reports about its build environment. Every fact may be +/// unknown; `None` means "not measured", never "fine". +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct EnvFacts { + /// Unix time the probes ran (the drone may serve cached probes for ~60 s). + pub observed_at: i64, + /// `RUSTC_WRAPPER` as a job on this drone would see it; `None` if unset. + #[serde(default)] + pub rustc_wrapper: Option, + /// `SCCACHE_DIR` as a job would see it; `None` if unset. + #[serde(default)] + pub sccache_dir: Option, + /// Whether the sccache server answered. + #[serde(default)] + pub sccache_server_reachable: Option, + /// Whether `sccache --show-stats` succeeded. + #[serde(default)] + pub sccache_stats_ok: Option, + /// Compile requests, hits and errors from the stats. + #[serde(default)] + pub cache_requests: Option, + #[serde(default)] + pub cache_hits: Option, + #[serde(default)] + pub cache_errors: Option, + /// Free bytes on the filesystem holding the cache directory. + #[serde(default)] + pub cache_dir_free_bytes: Option, + /// Whether a test write into the cache directory succeeded. + #[serde(default)] + pub cache_write_ok: Option, + /// Installed toolchains; `None` when the probe could not run. + #[serde(default)] + pub toolchains: Option>, + /// `MemAvailable` (Linux) or the platform equivalent, in MiB. + #[serde(default)] + pub mem_available_mib: Option, + /// Linux PSI `cpu some avg10`, in permille (percent times ten); `None` + /// where PSI does not exist. + #[serde(default)] + pub cpu_psi_some_avg10_pm: Option, +} + +/// Settings of one class (or the defaults for all). +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ClassContract { + /// Check modes overriding the global ones. + pub checks: BTreeMap, + /// The Rust toolchain every build of this class needs. + pub rust_toolchain: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ContractSettings { + /// Off unless the queen opts in. + pub enabled: bool, + /// Treat a drone that has never reported facts as failing the required + /// checks. Off by default so enabling the contract does not cut off drones + /// that predate it. + pub require_report: bool, + /// Facts older than this count as unknown. + pub facts_max_age_secs: i64, + pub allowed_wrappers: Vec, + pub min_cache_dir_free_bytes: u64, + /// Cache errors per thousand requests above which the store fails. + pub max_cache_error_pm: u32, + /// A hit rate at or below this per thousand, after `min_requests_for_hit_rate` + /// requests, is reported as a warning (never blocks). + pub hit_rate_warn_pm: u32, + pub min_requests_for_hit_rate: u64, + /// CPU pressure (permille) above which resources fail. + pub max_cpu_psi_pm: u16, + /// Memory available must be at least this percent of the class's memory. + pub mem_available_pct: u64, + pub checks: BTreeMap, + pub classes: BTreeMap, +} + +impl Default for ContractSettings { + fn default() -> Self { + Self { + enabled: false, + require_report: false, + facts_max_age_secs: 180, + allowed_wrappers: vec!["sccache".into()], + min_cache_dir_free_bytes: 10 << 30, + max_cache_error_pm: 50, + hit_rate_warn_pm: 10, + min_requests_for_hit_rate: 200, + max_cpu_psi_pm: 700, + mem_available_pct: 100, + // Conservative: only the checks that protect a build from failing + // outright block by default; the cache checks report. + checks: BTreeMap::from([ + (CheckKind::CompilerCache, Mode::Advisory), + (CheckKind::Toolchain, Mode::Required), + (CheckKind::CacheStore, Mode::Advisory), + (CheckKind::Resources, Mode::Required), + ]), + classes: BTreeMap::new(), + } + } +} + +impl ContractSettings { + pub fn validate(&self) -> Result<(), String> { + if self.facts_max_age_secs <= 0 + || self.allowed_wrappers.is_empty() + || self.mem_available_pct == 0 + { + return Err("env_contract needs positive facts_max_age_secs and mem_available_pct and at least one allowed wrapper".into()); + } + if self.max_cache_error_pm > 1000 + || self.max_cpu_psi_pm > 1000 + || self.hit_rate_warn_pm > 1000 + { + return Err("env_contract permille values are at most 1000".into()); + } + let bad_tc = |t: &str| { + t.is_empty() + || t.len() > 64 + || !t + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_')) + }; + if self + .classes + .values() + .filter_map(|c| c.rust_toolchain.as_deref()) + .any(bad_tc) + { + return Err( + "env_contract rust_toolchain must look like 1.95.0 or nightly-2026-01-01".into(), + ); + } + Ok(()) + } + + fn mode(&self, class: Class, kind: CheckKind) -> Mode { + self.classes + .get(&class) + .and_then(|c| c.checks.get(&kind)) + .or_else(|| self.checks.get(&kind)) + .copied() + .unwrap_or(Mode::Off) + } + + /// The toolchain a class needs by configuration. + pub fn class_toolchain(&self, class: Class) -> Option<&str> { + self.classes + .get(&class) + .and_then(|c| c.rust_toolchain.as_deref()) + } +} + +/// What the work being admitted needs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Need<'a> { + pub class: Class, + pub mem_mib: u64, + pub rust_toolchain: Option<&'a str>, +} + +impl<'a> Need<'a> { + /// The need of a class; `unit_toolchain` (from the unit's tree) wins over + /// the class configuration. + pub fn for_class( + s: &'a ContractSettings, + class: Class, + spec: &ClassSpec, + unit_toolchain: Option<&'a str>, + ) -> Self { + Self { + class, + mem_mib: spec.mem_mib, + rust_toolchain: unit_toolchain.or_else(|| s.class_toolchain(class)), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct CheckOutcome { + pub kind: CheckKind, + pub mode: Mode, + pub state: CheckState, + pub reason: String, + /// Values the verdict was based on. + pub observed: BTreeMap, + /// Things worth knowing that never block (e.g. a pathological hit rate). + pub warnings: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct Verdict { + /// Whether any facts have been reported. + pub reported: bool, + pub age_secs: Option, + pub checks: Vec, + pub admit: bool, + /// Why the work is not admitted. + pub why: Option, +} + +fn worst(a: CheckState, b: CheckState) -> CheckState { + a.max(b) +} + +fn basename(p: &str) -> &str { + p.rsplit(['/', '\\']).next().unwrap_or(p) +} + +struct Acc { + state: CheckState, + reasons: Vec, + observed: BTreeMap, + warnings: Vec, +} + +impl Acc { + fn new() -> Self { + Self { + state: CheckState::Pass, + reasons: Vec::new(), + observed: BTreeMap::new(), + warnings: Vec::new(), + } + } + + fn note(&mut self, state: CheckState, why: impl Into) { + if state != CheckState::Pass { + self.state = worst(self.state, state); + self.reasons.push(why.into()); + } + } + + fn see(&mut self, k: &str, v: impl ToString) { + self.observed.insert(k.to_owned(), v.to_string()); + } + + fn finish(self, kind: CheckKind, mode: Mode) -> CheckOutcome { + let reason = if self.reasons.is_empty() { + "ok".to_owned() + } else { + self.reasons.join("; ") + }; + CheckOutcome { + kind, + mode, + state: self.state, + reason, + observed: self.observed, + warnings: self.warnings, + } + } +} + +fn compiler_cache(s: &ContractSettings, f: &EnvFacts) -> Acc { + let mut a = Acc::new(); + match &f.rustc_wrapper { + None => a.note(CheckState::Fail, "RUSTC_WRAPPER is not set"), + Some(w) => { + a.see("rustc_wrapper", w); + if !s.allowed_wrappers.iter().any(|ok| ok == basename(w)) { + a.note( + CheckState::Fail, + format!("RUSTC_WRAPPER={w} is not an allowed wrapper"), + ); + } + } + } + match &f.sccache_dir { + None => a.note(CheckState::Fail, "SCCACHE_DIR is not set"), + Some(d) => a.see("sccache_dir", d), + } + match f.sccache_server_reachable { + Some(true) => a.see("sccache_server_reachable", true), + Some(false) => a.note(CheckState::Fail, "the sccache server is not reachable"), + None => a.note( + CheckState::Unknown, + "sccache server reachability is unknown", + ), + } + a +} + +fn toolchain(f: &EnvFacts, need: &Need<'_>) -> Acc { + let mut a = Acc::new(); + let Some(want) = need.rust_toolchain else { + a.see("required", "none"); + return a; + }; + a.see("required", want); + match &f.toolchains { + None => a.note(CheckState::Unknown, "installed toolchains are unknown"), + Some(have) => { + let names: Vec = have.iter().map(|t| t.channel.clone()).collect(); + a.see("installed", names.join(",")); + if !have + .iter() + .any(|t| t.channel == want || t.rustc_version == want) + { + a.note( + CheckState::Fail, + format!( + "toolchain {want} is not installed (have {})", + names.join(", ") + ), + ); + } + } + } + a +} + +fn cache_store(s: &ContractSettings, f: &EnvFacts) -> Acc { + let mut a = Acc::new(); + match f.sccache_stats_ok { + Some(true) => {} + Some(false) => a.note(CheckState::Fail, "sccache --show-stats failed"), + None => a.note(CheckState::Unknown, "cache stats were not measured"), + } + if let (Some(req), Some(err)) = (f.cache_requests, f.cache_errors) { + a.see("cache_requests", req); + a.see("cache_errors", err); + if req > 0 && err.saturating_mul(1000) / req > u64::from(s.max_cache_error_pm) { + a.note( + CheckState::Fail, + format!( + "cache error rate {}‰ is over {}‰", + err.saturating_mul(1000) / req, + s.max_cache_error_pm + ), + ); + } + } + match f.cache_dir_free_bytes { + Some(free) => { + a.see("cache_dir_free_bytes", free); + if free < s.min_cache_dir_free_bytes { + a.note( + CheckState::Fail, + format!( + "cache directory has {free} bytes free, under the floor of {}", + s.min_cache_dir_free_bytes + ), + ); + } + } + None => a.note(CheckState::Unknown, "cache directory free space is unknown"), + } + match f.cache_write_ok { + Some(true) => {} + Some(false) => a.note(CheckState::Fail, "writing to the cache directory failed"), + None => a.note(CheckState::Unknown, "cache writability was not measured"), + } + if let (Some(req), Some(hits)) = (f.cache_requests, f.cache_hits) + && req >= s.min_requests_for_hit_rate + { + let pm = hits.saturating_mul(1000) / req.max(1); + a.see("hit_rate_pm", pm); + if pm <= u64::from(s.hit_rate_warn_pm) { + a.warnings + .push(format!("hit rate is {pm}‰ after {req} requests")); + } + } + a +} + +fn resources(s: &ContractSettings, f: &EnvFacts, need: &Need<'_>) -> Acc { + let mut a = Acc::new(); + let mut known = 0; + if let Some(avail) = f.mem_available_mib { + known += 1; + a.see("mem_available_mib", avail); + a.see("mem_needed_mib", need.mem_mib); + if avail.saturating_mul(100) < need.mem_mib.saturating_mul(s.mem_available_pct) { + a.note( + CheckState::Fail, + format!( + "{avail} MiB available, {} class needs {} MiB", + need.class, need.mem_mib + ), + ); + } + } + if let Some(psi) = f.cpu_psi_some_avg10_pm { + known += 1; + a.see("cpu_psi_some_avg10_pm", psi); + if psi > s.max_cpu_psi_pm { + a.note( + CheckState::Fail, + format!("cpu pressure {psi}‰ is over {}‰", s.max_cpu_psi_pm), + ); + } + } + if known == 0 { + a.note( + CheckState::Unknown, + "neither available memory nor cpu pressure was measured", + ); + } + a +} + +/// Evaluate the contract for `need` against `facts` at `now`. Pure. +pub fn evaluate( + s: &ContractSettings, + facts: Option<&EnvFacts>, + need: &Need<'_>, + now: i64, +) -> Verdict { + if !s.enabled { + return Verdict { + reported: facts.is_some(), + age_secs: facts.map(|f| now - f.observed_at), + checks: Vec::new(), + admit: true, + why: None, + }; + } + let modes: Vec<(CheckKind, Mode)> = CheckKind::ALL + .into_iter() + .map(|k| (k, s.mode(need.class, k))) + .filter(|(_, m)| *m != Mode::Off) + .collect(); + let Some(f) = facts else { + let blocked = s.require_report && modes.iter().any(|(_, m)| *m == Mode::Required); + return Verdict { + reported: false, + age_secs: None, + checks: modes + .iter() + .map(|&(kind, mode)| CheckOutcome { + kind, + mode, + state: CheckState::Unknown, + reason: "the drone has not reported its environment".into(), + observed: BTreeMap::new(), + warnings: Vec::new(), + }) + .collect(), + admit: !blocked, + why: blocked.then(|| "the drone has not reported its environment".to_owned()), + }; + }; + let age = now - f.observed_at; + let stale = age > s.facts_max_age_secs; + let mut checks = Vec::new(); + for (kind, mode) in modes { + let acc = if stale { + let mut a = Acc::new(); + a.note( + CheckState::Unknown, + format!("facts are {age}s old (limit {}s)", s.facts_max_age_secs), + ); + a + } else { + match kind { + CheckKind::CompilerCache => compiler_cache(s, f), + CheckKind::Toolchain => toolchain(f, need), + CheckKind::CacheStore => cache_store(s, f), + CheckKind::Resources => resources(s, f, need), + } + }; + checks.push(acc.finish(kind, mode)); + } + let blockers: Vec = checks + .iter() + .filter(|c| c.mode == Mode::Required && c.state != CheckState::Pass) + .map(|c| format!("{} {}: {}", c.kind.as_str(), c.state.as_str(), c.reason)) + .collect(); + Verdict { + reported: true, + age_secs: Some(age), + checks, + admit: blockers.is_empty(), + why: (!blockers.is_empty()).then(|| blockers.join("; ")), + } +} + +/// The classes `facts` do not admit, for `DroneState::blocked_classes`. +pub fn blocked_classes( + s: &ContractSettings, + facts: Option<&EnvFacts>, + classes: &Classes, + unit_toolchain: Option<&str>, + now: i64, +) -> Vec { + if !s.enabled { + return Vec::new(); + } + Class::ALL + .into_iter() + .filter(|&c| { + let need = Need::for_class(s, c, classes.spec(c), unit_toolchain); + !evaluate(s, facts, &need, now).admit + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + const NOW: i64 = 10_000; + + fn on() -> ContractSettings { + ContractSettings { + enabled: true, + ..ContractSettings::default() + } + } + + fn good() -> EnvFacts { + EnvFacts { + observed_at: NOW - 10, + rustc_wrapper: Some("/usr/local/bin/sccache".into()), + sccache_dir: Some("/var/lib/borg/cache/sccache".into()), + sccache_server_reachable: Some(true), + sccache_stats_ok: Some(true), + cache_requests: Some(1000), + cache_hits: Some(600), + cache_errors: Some(2), + cache_dir_free_bytes: Some(50 << 30), + cache_write_ok: Some(true), + toolchains: Some(vec![Toolchain { + channel: "1.95.0".into(), + rustc_version: "1.95.0".into(), + }]), + mem_available_mib: Some(20_000), + cpu_psi_some_avg10_pm: Some(100), + } + } + + fn need(class: Class, tc: Option<&str>) -> Need<'_> { + Need { + class, + mem_mib: Classes::default().spec(class).mem_mib, + rust_toolchain: tc, + } + } + + fn check(v: &Verdict, kind: CheckKind) -> &CheckOutcome { + v.checks.iter().find(|c| c.kind == kind).unwrap() + } + + fn all_required() -> ContractSettings { + ContractSettings { + checks: CheckKind::ALL + .into_iter() + .map(|k| (k, Mode::Required)) + .collect(), + ..on() + } + } + + #[test] + fn a_healthy_drone_passes_every_check_and_is_admitted() { + let v = evaluate( + &all_required(), + Some(&good()), + &need(Class::Medium, Some("1.95.0")), + NOW, + ); + assert!(v.admit, "{v:?}"); + assert!( + v.checks.iter().all(|c| c.state == CheckState::Pass), + "{v:?}" + ); + assert_eq!(v.age_secs, Some(10)); + assert_eq!( + check(&v, CheckKind::Toolchain).observed["installed"], + "1.95.0" + ); + } + + #[test] + fn compiler_cache_requires_an_allowed_wrapper_a_dir_and_a_reachable_server() { + let s = all_required(); + let n = need(Class::Light, None); + let run = |f: EnvFacts| evaluate(&s, Some(&f), &n, NOW); + let v = run(EnvFacts { + rustc_wrapper: None, + ..good() + }); + assert!(!v.admit); + assert!( + check(&v, CheckKind::CompilerCache) + .reason + .contains("RUSTC_WRAPPER is not set") + ); + let v = run(EnvFacts { + rustc_wrapper: Some("/usr/bin/ccache".into()), + ..good() + }); + assert!( + check(&v, CheckKind::CompilerCache) + .reason + .contains("not an allowed wrapper") + ); + // The single-flight wrapper is allowed when configured. + let s2 = ContractSettings { + allowed_wrappers: vec!["sccache".into(), "borg-sccache".into()], + ..all_required() + }; + let f = EnvFacts { + rustc_wrapper: Some("/opt/borg/bin/borg-sccache".into()), + ..good() + }; + assert!(evaluate(&s2, Some(&f), &n, NOW).admit); + let v = run(EnvFacts { + sccache_dir: None, + ..good() + }); + assert!( + check(&v, CheckKind::CompilerCache) + .reason + .contains("SCCACHE_DIR") + ); + let v = run(EnvFacts { + sccache_server_reachable: Some(false), + ..good() + }); + assert_eq!(check(&v, CheckKind::CompilerCache).state, CheckState::Fail); + let v = run(EnvFacts { + sccache_server_reachable: None, + ..good() + }); + assert_eq!( + check(&v, CheckKind::CompilerCache).state, + CheckState::Unknown + ); + assert!(!v.admit, "unknown on a required check blocks"); + } + + #[test] + fn toolchain_matches_by_channel_or_version_and_only_when_declared() { + let s = all_required(); + let f = good(); + assert!(evaluate(&s, Some(&f), &need(Class::Light, None), NOW).admit); + let v = evaluate(&s, Some(&f), &need(Class::Light, Some("1.96.0")), NOW); + assert!(!v.admit); + assert!( + check(&v, CheckKind::Toolchain) + .reason + .contains("1.96.0 is not installed (have 1.95.0)") + ); + // A named channel resolved to a version matches by version too. + let g = EnvFacts { + toolchains: Some(vec![Toolchain { + channel: "stable".into(), + rustc_version: "1.95.0".into(), + }]), + ..good() + }; + assert!(evaluate(&s, Some(&g), &need(Class::Light, Some("1.95.0")), NOW).admit); + let u = EnvFacts { + toolchains: None, + ..good() + }; + let v = evaluate(&s, Some(&u), &need(Class::Light, Some("1.95.0")), NOW); + assert_eq!(check(&v, CheckKind::Toolchain).state, CheckState::Unknown); + // The class config supplies a requirement; a unit's own wins. + let mut cs = all_required(); + cs.classes.insert( + Class::Heavy, + ClassContract { + rust_toolchain: Some("1.90.0".into()), + ..ClassContract::default() + }, + ); + let n = Need::for_class( + &cs, + Class::Heavy, + Classes::default().spec(Class::Heavy), + None, + ); + assert_eq!(n.rust_toolchain, Some("1.90.0")); + let n = Need::for_class( + &cs, + Class::Heavy, + Classes::default().spec(Class::Heavy), + Some("1.95.0"), + ); + assert_eq!(n.rust_toolchain, Some("1.95.0")); + } + + #[test] + fn cache_store_checks_errors_free_space_writes_and_only_warns_about_hit_rate() { + let s = all_required(); + let n = need(Class::Light, None); + let run = |f: EnvFacts| evaluate(&s, Some(&f), &n, NOW); + let v = run(EnvFacts { + cache_errors: Some(100), + ..good() + }); + assert!( + check(&v, CheckKind::CacheStore) + .reason + .contains("error rate 100‰") + ); + let v = run(EnvFacts { + cache_dir_free_bytes: Some(1 << 30), + ..good() + }); + assert!( + check(&v, CheckKind::CacheStore) + .reason + .contains("under the floor") + ); + let v = run(EnvFacts { + cache_write_ok: Some(false), + ..good() + }); + assert_eq!(check(&v, CheckKind::CacheStore).state, CheckState::Fail); + let v = run(EnvFacts { + sccache_stats_ok: Some(false), + ..good() + }); + assert_eq!(check(&v, CheckKind::CacheStore).state, CheckState::Fail); + // A 0% hit rate after enough requests warns and does not block. + let v = run(EnvFacts { + cache_hits: Some(0), + ..good() + }); + assert!(v.admit, "{v:?}"); + assert!(check(&v, CheckKind::CacheStore).warnings[0].contains("hit rate is 0")); + // Too few requests to judge: no warning. + let v = run(EnvFacts { + cache_requests: Some(10), + cache_hits: Some(0), + cache_errors: Some(0), + ..good() + }); + assert!(check(&v, CheckKind::CacheStore).warnings.is_empty()); + } + + #[test] + fn resources_compare_available_memory_with_the_class_and_cpu_pressure_with_the_ceiling() { + let s = all_required(); + // Heavy needs 24 GiB. + let heavy = need(Class::Heavy, None); + let v = evaluate( + &s, + Some(&EnvFacts { + mem_available_mib: Some(10_000), + ..good() + }), + &heavy, + NOW, + ); + assert!(!v.admit); + assert!( + check(&v, CheckKind::Resources) + .reason + .contains("10000 MiB available") + ); + let v = evaluate( + &s, + Some(&EnvFacts { + mem_available_mib: Some(30_000), + ..good() + }), + &heavy, + NOW, + ); + assert!(v.admit); + let v = evaluate( + &s, + Some(&EnvFacts { + cpu_psi_some_avg10_pm: Some(900), + ..good() + }), + &need(Class::Light, None), + NOW, + ); + assert!( + check(&v, CheckKind::Resources) + .reason + .contains("cpu pressure 900") + ); + // No PSI (macOS) but memory known: still decidable. + let v = evaluate( + &s, + Some(&EnvFacts { + cpu_psi_some_avg10_pm: None, + ..good() + }), + &need(Class::Light, None), + NOW, + ); + assert!(v.admit); + // Nothing measured: unknown, which a required check treats as blocking. + let v = evaluate( + &s, + Some(&EnvFacts { + cpu_psi_some_avg10_pm: None, + mem_available_mib: None, + ..good() + }), + &need(Class::Light, None), + NOW, + ); + assert_eq!(check(&v, CheckKind::Resources).state, CheckState::Unknown); + assert!(!v.admit); + } + + #[test] + fn advisory_and_off_checks_never_block_and_class_overrides_apply() { + let bad = EnvFacts { + rustc_wrapper: None, + ..good() + }; + let n = need(Class::Light, None); + // Default: compiler_cache is advisory, so the failure is reported only. + let v = evaluate(&on(), Some(&bad), &n, NOW); + assert!(v.admit); + assert_eq!(check(&v, CheckKind::CompilerCache).state, CheckState::Fail); + // Required for one class only. + let mut s = on(); + s.classes.insert( + Class::Heavy, + ClassContract { + checks: BTreeMap::from([(CheckKind::CompilerCache, Mode::Required)]), + ..ClassContract::default() + }, + ); + assert!(evaluate(&s, Some(&bad), &need(Class::Light, None), NOW).admit); + assert!(!evaluate(&s, Some(&bad), &need(Class::Heavy, None), NOW).admit); + // Off removes the check from the report. + s.checks.insert(CheckKind::CompilerCache, Mode::Off); + let v = evaluate(&s, Some(&bad), &need(Class::Light, None), NOW); + assert!(v.checks.iter().all(|c| c.kind != CheckKind::CompilerCache)); + } + + #[test] + fn stale_or_missing_facts_are_unknown_and_the_report_requirement_is_opt_in() { + let s = all_required(); + let old = EnvFacts { + observed_at: NOW - 1000, + ..good() + }; + let v = evaluate(&s, Some(&old), &need(Class::Light, None), NOW); + assert!(!v.admit); + assert!(v.checks.iter().all(|c| c.state == CheckState::Unknown)); + assert!(v.checks[0].reason.contains("1000s old")); + // A drone that never reported: admitted by default, blocked on request. + let v = evaluate(&s, None, &need(Class::Light, None), NOW); + assert!(v.admit && !v.reported); + let strict = ContractSettings { + require_report: true, + ..all_required() + }; + let v = evaluate(&strict, None, &need(Class::Light, None), NOW); + assert!(!v.admit); + assert!(v.why.unwrap().contains("not reported")); + } + + #[test] + fn a_disabled_contract_admits_everything_and_blocked_classes_are_per_class() { + let off = ContractSettings::default(); + assert!( + evaluate( + &off, + Some(&EnvFacts::default()), + &need(Class::Heavy, None), + NOW + ) + .admit + ); + assert!(blocked_classes(&off, None, &Classes::default(), None, NOW).is_empty()); + // 10 GiB available: light and medium fit (2.5 and 12 GiB: medium does not). + let s = all_required(); + let f = EnvFacts { + mem_available_mib: Some(10_000), + ..good() + }; + let blocked = blocked_classes(&s, Some(&f), &Classes::default(), None, NOW); + assert_eq!(blocked, [Class::Medium, Class::Heavy]); + // A unit's toolchain requirement blocks every class on a drone without it. + let blocked = blocked_classes(&s, Some(&good()), &Classes::default(), Some("9.9.9"), NOW); + assert_eq!(blocked, Class::ALL); + } + + #[test] + fn settings_validate_and_evaluation_is_deterministic() { + assert!(ContractSettings::default().validate().is_ok()); + for bad in [ + ContractSettings { + facts_max_age_secs: 0, + ..ContractSettings::default() + }, + ContractSettings { + allowed_wrappers: vec![], + ..ContractSettings::default() + }, + ContractSettings { + max_cpu_psi_pm: 2000, + ..ContractSettings::default() + }, + ] { + assert!(bad.validate().is_err()); + } + let mut s = on(); + s.classes.insert( + Class::Light, + ClassContract { + rust_toolchain: Some("bad toolchain!".into()), + ..ClassContract::default() + }, + ); + assert!(s.validate().is_err()); + let n = need(Class::Medium, Some("1.95.0")); + assert_eq!( + evaluate(&all_required(), Some(&good()), &n, NOW), + evaluate(&all_required(), Some(&good()), &n, NOW) + ); + } +} diff --git a/src/github.rs b/src/github.rs index 4edd83a..4aa6db7 100644 --- a/src/github.rs +++ b/src/github.rs @@ -54,6 +54,14 @@ pub struct Runner { pub name: String, pub status: String, pub busy: bool, + /// Labels GitHub reports for the runner (`self-hosted`, `Linux`, ...). + #[serde(default)] + pub labels: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct RunnerLabel { + pub name: String, } #[derive(Debug)] diff --git a/src/health.rs b/src/health.rs new file mode 100644 index 0000000..c39854f --- /dev/null +++ b/src/health.rs @@ -0,0 +1,833 @@ +//! Host health from job telemetry: decayed evidence, quarantine +//! recommendations, bounded automatic quarantine, and a cooldown with a health +//! proof before a host returns. +//! +//! [`evaluate`] is a pure function of the clock value it is given, the drones' +//! current holds and the recent job outcomes (from the telemetry drones report +//! with each lease result). The queen applies the [`Action`]s it returns. +//! +//! Rules: +//! * Only host-caused failures (`failure_class: host`) count against a host. +//! Code-caused failures, successes and unclassified failures never do. +//! * Evidence decays with a half-life ([`decay_permille`]); a failure from a +//! long time ago weighs next to nothing, and a single fresh failure is below +//! every threshold. +//! * Corroboration: automatic quarantine needs `corroboration_runs` host +//! failures from distinct workflow runs inside the window, plus a decayed +//! score at or above `quarantine_score_permille`. Retrying one run does not +//! count twice. +//! * Bounds: no automatic quarantine while the evidence looks fleet-wide +//! (many hosts failing at once points at GitHub, the network or a bad +//! workflow, not at one host), beyond `max_auto_quarantined` held hosts, or +//! if it would leave fewer than `min_schedulable_remaining` fresh hosts in +//! service. In those cases the host is only *recommended* for quarantine. +//! * Return: after `cooldown_secs` (doubling per earlier strike, capped) an +//! automatically quarantined host enters probation: at most one lease at a +//! time. After `proof_jobs` successful jobs it is released; a host-caused +//! failure during probation sends it back with one more strike. Operator +//! quarantines are never lifted automatically. + +use std::collections::{BTreeMap, BTreeSet}; + +use serde::{Deserialize, Serialize}; + +/// Tunables; every field has a default. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct HealthSettings { + /// Quarantine automatically when evidence is strong; otherwise only + /// recommend. + pub auto_quarantine: bool, + /// Half-life of a failure's weight. + pub half_life_secs: u64, + /// Only failures newer than this are considered at all. + pub window_secs: u64, + /// Distinct runs with a host failure needed to quarantine automatically. + pub corroboration_runs: u32, + /// Decayed score (1000 = one fresh failure) needed to quarantine. + pub quarantine_score_permille: u32, + /// Decayed score at which a quarantine is recommended. + pub recommend_score_permille: u32, + /// Base cooldown before an auto-quarantined host enters probation. + pub cooldown_secs: u64, + /// Successful jobs on probation needed to return to service. + pub proof_jobs: u32, + /// At most this many hosts held (quarantined or on probation) by the + /// automatic path at once. + pub max_auto_quarantined: u32, + /// Never auto-quarantine if fewer than this many fresh, unheld hosts + /// would remain. + pub min_schedulable_remaining: u32, + /// Fleet-wide guard: with at least 3 hosts that ran jobs, if this share + /// (percent) had host failures, nothing is auto-quarantined. + pub systemic_share_pct: u32, + /// Cooldown multiplier cap (cooldown doubles per strike up to this). + pub max_cooldown_multiplier: u32, +} + +impl Default for HealthSettings { + fn default() -> Self { + Self { + auto_quarantine: true, + half_life_secs: 6 * 3600, + window_secs: 12 * 3600, + corroboration_runs: 3, + quarantine_score_permille: 2500, + recommend_score_permille: 1500, + cooldown_secs: 1800, + proof_jobs: 2, + max_auto_quarantined: 2, + min_schedulable_remaining: 1, + systemic_share_pct: 50, + max_cooldown_multiplier: 8, + } + } +} + +impl HealthSettings { + pub fn validate(&self) -> Result<(), String> { + if self.half_life_secs == 0 || self.window_secs == 0 || self.cooldown_secs == 0 { + return Err( + "health half_life_secs, window_secs and cooldown_secs must be positive".into(), + ); + } + if self.corroboration_runs == 0 || self.proof_jobs == 0 || self.max_cooldown_multiplier == 0 + { + return Err( + "health corroboration_runs, proof_jobs and max_cooldown_multiplier must be positive" + .into(), + ); + } + if self.recommend_score_permille > self.quarantine_score_permille { + return Err( + "health recommend_score_permille may not exceed quarantine_score_permille".into(), + ); + } + if self.systemic_share_pct == 0 || self.systemic_share_pct > 100 { + return Err("health systemic_share_pct must be 1 to 100".into()); + } + Ok(()) + } + + /// Cooldown after `strikes` earlier failed returns. + pub fn cooldown_for(&self, strikes: u32) -> i64 { + let mult = 1u64 + .checked_shl(strikes.min(16)) + .unwrap_or(u64::MAX) + .min(u64::from(self.max_cooldown_multiplier)); + i64::try_from(self.cooldown_secs.saturating_mul(mult)).unwrap_or(i64::MAX) + } +} + +/// Weight (permille) of evidence `age_secs` old: 1000 now, half after one +/// half-life, a quarter after two. Piecewise linear between half-lives so it +/// stays integer and monotone. +pub fn decay_permille(age_secs: i64, half_life_secs: u64) -> u32 { + let Ok(age) = u64::try_from(age_secs) else { + return 1000; + }; + let hl = half_life_secs.max(1); + let whole = age / hl; + if whole >= 10 { + return 0; + } + let frac = age % hl; + let base = 1000u64 >> whole; + let w = base - base * frac / (2 * hl); + u32::try_from(w).unwrap_or(0) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutcomeKind { + /// The job succeeded. + Ok, + /// The job failed because the host or runner failed. + HostFailure, + /// Anything else: a code failure, a cancellation, an unclassified failure. + Other, +} + +/// One finished job, from telemetry. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Outcome { + pub drone: String, + pub lease_id: String, + pub repo: String, + pub run_id: u64, + pub job_name: String, + pub finished_at: i64, + pub kind: OutcomeKind, +} + +/// Why a drone is held back from new leases. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Hold { + None, + /// Quarantined by an operator: never lifted automatically. + Operator, + Auto { + since: i64, + strikes: u32, + }, + Probation { + since: i64, + strikes: u32, + }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DroneInfo { + pub name: String, + pub fresh: bool, + pub hold: Hold, + /// Evidence at or before this time was cleared by a return to service. + pub cleared_at: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum HealthState { + Healthy, + /// Evidence suggests a quarantine, but it was not applied. + Recommended, + Quarantined, + Probation, +} + +impl HealthState { + pub const ALL: [Self; 4] = [ + Self::Healthy, + Self::Recommended, + Self::Quarantined, + Self::Probation, + ]; + + pub fn as_str(self) -> &'static str { + match self { + Self::Healthy => "healthy", + Self::Recommended => "recommended", + Self::Quarantined => "quarantined", + Self::Probation => "probation", + } + } +} + +/// One drone's health as shown in `/status` and metrics. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct DroneHealth { + pub state: HealthState, + /// `operator` or `auto` while held. + pub source: Option, + /// Decayed host-failure evidence, 1000 = one fresh failure. + pub score_permille: u32, + pub host_failures: u32, + pub corroborating_runs: u32, + pub strikes: u32, + /// Unix time the hold (or probation) started. + pub since: Option, + /// When an auto-quarantined host becomes eligible for probation. + pub probation_at: Option, + /// Successful jobs counted toward the health proof. + pub proof_jobs_done: u32, + /// Why a recommendation was not applied automatically. + pub note: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Evidence { + pub lease_id: String, + pub repo: String, + pub run_id: u64, + pub job_name: String, + pub finished_at: i64, + pub weight_permille: u32, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Action { + /// Quarantine on strong, corroborated evidence. + Quarantine { + drone: String, + reason: String, + evidence: Vec, + strikes: u32, + }, + /// Cooldown elapsed: allow one lease at a time. + BeginProbation { drone: String }, + /// A host failure during probation: back to quarantine. + Requarantine { + drone: String, + reason: String, + evidence: Vec, + strikes: u32, + }, + /// Health proven: return to service. + Release { drone: String }, +} + +impl Action { + pub fn drone(&self) -> &str { + match self { + Self::Quarantine { drone, .. } + | Self::BeginProbation { drone } + | Self::Requarantine { drone, .. } + | Self::Release { drone } => drone, + } + } + + pub fn event(&self) -> &'static str { + match self { + Self::Quarantine { .. } => "auto_quarantine", + Self::BeginProbation { .. } => "probation_start", + Self::Requarantine { .. } => "requarantine", + Self::Release { .. } => "release", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct Evaluation { + pub drones: BTreeMap, + pub actions: Vec, + /// Host failures look fleet-wide; automatic quarantine is paused. + pub systemic: bool, +} + +fn run_key(o: &Outcome) -> String { + if o.run_id == 0 { + format!("lease:{}", o.lease_id) + } else { + format!("{}#{}", o.repo, o.run_id) + } +} + +fn score_text(permille: u32) -> String { + format!("{}.{}", permille / 1000, permille % 1000 / 100) +} + +struct Failures { + score: u32, + count: u32, + runs: u32, + evidence: Vec, +} + +fn failures_of( + now: i64, + s: &HealthSettings, + drone: &str, + since: Option, + outcomes: &[Outcome], +) -> Failures { + let mut f = Failures { + score: 0, + count: 0, + runs: 0, + evidence: Vec::new(), + }; + let mut runs = BTreeSet::new(); + let window = i64::try_from(s.window_secs).unwrap_or(i64::MAX); + for o in outcomes.iter().filter(|o| { + o.drone == drone + && o.kind == OutcomeKind::HostFailure + && now.saturating_sub(o.finished_at) <= window + && since.is_none_or(|t| o.finished_at > t) + }) { + let weight = decay_permille(now - o.finished_at, s.half_life_secs); + f.score = f.score.saturating_add(weight); + f.count += 1; + runs.insert(run_key(o)); + f.evidence.push(Evidence { + lease_id: o.lease_id.clone(), + repo: o.repo.clone(), + run_id: o.run_id, + job_name: o.job_name.clone(), + finished_at: o.finished_at, + weight_permille: weight, + }); + } + f.runs = u32::try_from(runs.len()).unwrap_or(u32::MAX); + f.evidence.sort_by_key(|e| std::cmp::Reverse(e.finished_at)); + f.evidence.truncate(20); + f +} + +/// Evaluate every drone. Pure: the same inputs give the same actions. +pub fn evaluate( + now: i64, + s: &HealthSettings, + drones: &[DroneInfo], + outcomes: &[Outcome], +) -> Evaluation { + let window = i64::try_from(s.window_secs).unwrap_or(i64::MAX); + let recent = |o: &&Outcome| now.saturating_sub(o.finished_at) <= window; + + // Fleet-wide guard. + let ran: BTreeSet<&str> = outcomes + .iter() + .filter(recent) + .map(|o| o.drone.as_str()) + .collect(); + let hit: BTreeSet<&str> = outcomes + .iter() + .filter(recent) + .filter(|o| o.kind == OutcomeKind::HostFailure) + .map(|o| o.drone.as_str()) + .collect(); + let systemic = ran.len() >= 3 + && u64::try_from(hit.len()).unwrap_or(0) * 100 + >= u64::try_from(ran.len()).unwrap_or(0) * u64::from(s.systemic_share_pct); + + let mut out = Evaluation { + systemic, + ..Evaluation::default() + }; + let mut auto_held = u32::try_from( + drones + .iter() + .filter(|d| matches!(d.hold, Hold::Auto { .. } | Hold::Probation { .. })) + .count(), + ) + .unwrap_or(u32::MAX); + let mut in_service = u32::try_from( + drones + .iter() + .filter(|d| d.fresh && d.hold == Hold::None) + .count(), + ) + .unwrap_or(u32::MAX); + + let health = |state, source: Option<&str>, f: &Failures, strikes, since| DroneHealth { + state, + source: source.map(str::to_owned), + score_permille: f.score, + host_failures: f.count, + corroborating_runs: f.runs, + strikes, + since, + probation_at: None, + proof_jobs_done: 0, + note: None, + }; + + let mut candidates: Vec<(u32, &DroneInfo, Failures)> = Vec::new(); + for d in drones { + let f = failures_of(now, s, &d.name, d.cleared_at, outcomes); + match d.hold { + Hold::None => { + if f.runs >= s.corroboration_runs && f.score >= s.quarantine_score_permille { + candidates.push((f.score, d, f)); + } else { + let state = if f.score >= s.recommend_score_permille { + HealthState::Recommended + } else { + HealthState::Healthy + }; + out.drones + .insert(d.name.clone(), health(state, None, &f, 0, None)); + } + } + Hold::Operator => { + out.drones.insert( + d.name.clone(), + health(HealthState::Quarantined, Some("operator"), &f, 0, None), + ); + } + Hold::Auto { since, strikes } => { + let ready = since + s.cooldown_for(strikes); + let mut h = health( + HealthState::Quarantined, + Some("auto"), + &f, + strikes, + Some(since), + ); + h.probation_at = Some(ready); + if now >= ready { + out.actions.push(Action::BeginProbation { + drone: d.name.clone(), + }); + } + out.drones.insert(d.name.clone(), h); + } + Hold::Probation { since, strikes } => { + let mine = |o: &&Outcome| o.drone == d.name && o.finished_at >= since; + let failed: Vec<&Outcome> = outcomes + .iter() + .filter(mine) + .filter(|o| o.kind == OutcomeKind::HostFailure) + .collect(); + let ok = u32::try_from( + outcomes + .iter() + .filter(mine) + .filter(|o| o.kind == OutcomeKind::Ok) + .count(), + ) + .unwrap_or(u32::MAX); + let mut h = health( + HealthState::Probation, + Some("auto"), + &f, + strikes, + Some(since), + ); + h.proof_jobs_done = ok; + if !failed.is_empty() { + let evidence = failures_of(now, s, &d.name, Some(since - 1), outcomes).evidence; + out.actions.push(Action::Requarantine { + drone: d.name.clone(), + reason: format!( + "auto: host-caused failure during probation ({} job(s))", + failed.len() + ), + evidence, + strikes: strikes + 1, + }); + } else if ok >= s.proof_jobs { + out.actions.push(Action::Release { + drone: d.name.clone(), + }); + } + out.drones.insert(d.name.clone(), h); + } + } + } + + candidates.sort_by(|a, b| b.0.cmp(&a.0).then_with(|| a.1.name.cmp(&b.1.name))); + for (_, d, f) in candidates { + let after = in_service.saturating_sub(u32::from(d.fresh)); + let blocked = if !s.auto_quarantine { + Some("automatic quarantine is off") + } else if systemic { + Some("failures look fleet-wide") + } else if auto_held >= s.max_auto_quarantined { + Some("automatic quarantine limit reached") + } else if after < s.min_schedulable_remaining { + Some("would leave too few hosts in service") + } else { + None + }; + match blocked { + None => { + auto_held += 1; + in_service = after; + let reason = format!( + "auto: {} host-caused failures in {} distinct runs, decayed score {}", + f.count, + f.runs, + score_text(f.score) + ); + let mut h = health(HealthState::Quarantined, Some("auto"), &f, 0, Some(now)); + h.probation_at = Some(now + s.cooldown_for(0)); + out.drones.insert(d.name.clone(), h); + out.actions.push(Action::Quarantine { + drone: d.name.clone(), + reason, + evidence: f.evidence, + strikes: 0, + }); + } + Some(why) => { + let mut h = health(HealthState::Recommended, None, &f, 0, None); + h.note = Some(why.to_owned()); + out.drones.insert(d.name.clone(), h); + } + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + const HL: u64 = 3600; + + fn settings() -> HealthSettings { + HealthSettings { + half_life_secs: HL, + window_secs: 12 * 3600, + ..HealthSettings::default() + } + } + + fn drone(name: &str) -> DroneInfo { + DroneInfo { + name: name.into(), + fresh: true, + hold: Hold::None, + cleared_at: None, + } + } + + fn out(drone: &str, run: u64, at: i64, kind: OutcomeKind) -> Outcome { + Outcome { + drone: drone.into(), + lease_id: format!("{drone}-{run}-{at}"), + repo: "o/r".into(), + run_id: run, + job_name: "job".into(), + finished_at: at, + kind, + } + } + + fn fails(drone: &str, runs: &[u64], at: i64) -> Vec { + runs.iter() + .map(|&r| out(drone, r, at, OutcomeKind::HostFailure)) + .collect() + } + + #[test] + fn decay_halves_per_half_life_and_is_monotone() { + assert_eq!(decay_permille(0, HL), 1000); + assert_eq!(decay_permille(-50, HL), 1000); + assert_eq!(decay_permille(3600, HL), 500); + assert_eq!(decay_permille(7200, HL), 250); + assert_eq!(decay_permille(1800, HL), 750); + assert_eq!(decay_permille(36_000, HL), 0); + let mut last = 1000; + for t in (0..40_000).step_by(300) { + let w = decay_permille(t, HL); + assert!(w <= last, "t={t}"); + last = w; + } + } + + #[test] + fn one_failure_never_triggers_anything_and_two_only_recommend() { + let s = settings(); + let d = [drone("a"), drone("b")]; + let e = evaluate(1000, &s, &d, &fails("a", &[1], 1000)); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].state, HealthState::Healthy); + let e = evaluate(1000, &s, &d, &fails("a", &[1, 2], 1000)); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].state, HealthState::Recommended); + } + + #[test] + fn three_corroborated_fresh_failures_quarantine_with_evidence() { + let s = settings(); + let d = [drone("a"), drone("b")]; + let e = evaluate(1000, &s, &d, &fails("a", &[1, 2, 3], 990)); + assert_eq!(e.drones["a"].state, HealthState::Quarantined); + let Action::Quarantine { + drone, + reason, + evidence, + strikes, + } = &e.actions[0] + else { + panic!("{:?}", e.actions); + }; + assert_eq!((drone.as_str(), *strikes, evidence.len()), ("a", 0, 3)); + assert!( + reason.contains("3 host-caused failures in 3 distinct runs"), + "{reason}" + ); + assert_eq!(e.drones["b"].state, HealthState::Healthy); + } + + #[test] + fn retries_of_one_run_do_not_corroborate() { + let s = settings(); + let d = [drone("a"), drone("b")]; + let e = evaluate(1000, &s, &d, &fails("a", &[7, 7, 7, 7], 990)); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].state, HealthState::Recommended); + } + + #[test] + fn code_failures_and_cancellations_never_count_against_a_host() { + let s = settings(); + let d = [drone("a"), drone("b")]; + let outcomes: Vec = (1..=10) + .map(|r| out("a", r, 990, OutcomeKind::Other)) + .collect(); + let e = evaluate(1000, &s, &d, &outcomes); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].score_permille, 0); + } + + #[test] + fn old_evidence_decays_below_the_thresholds() { + let s = settings(); + let d = [drone("a"), drone("b")]; + // Three failures, five half-lives ago: weight 31 each. + let e = evaluate(5 * 3600 + 1000, &s, &d, &fails("a", &[1, 2, 3], 1000)); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].state, HealthState::Healthy); + // Outside the window they are ignored entirely. + let e = evaluate(13 * 3600, &s, &d, &fails("a", &[1, 2, 3], 0)); + assert_eq!(e.drones["a"].host_failures, 0); + } + + #[test] + fn fleet_wide_failures_only_recommend() { + let s = settings(); + let d = [drone("a"), drone("b"), drone("c")]; + let mut o = fails("a", &[1, 2, 3], 990); + o.extend(fails("b", &[4], 990)); + o.push(out("c", 5, 990, OutcomeKind::Ok)); + let e = evaluate(1000, &s, &d, &o); + assert!(e.systemic); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].state, HealthState::Recommended); + assert!( + e.drones["a"] + .note + .as_deref() + .unwrap() + .contains("fleet-wide") + ); + } + + #[test] + fn bounds_keep_the_last_hosts_and_the_auto_limit() { + let s = settings(); + // The only fresh host is never auto-quarantined. + let e = evaluate(1000, &s, &[drone("a")], &fails("a", &[1, 2, 3], 990)); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].state, HealthState::Recommended); + // Five hosts, two bad (40%, below the fleet-wide share): both can go + // (limit 2); the worse goes first. + let d = [drone("a"), drone("b"), drone("c"), drone("d"), drone("e")]; + let mut o = fails("a", &[1, 2, 3], 990); + o.extend(fails("b", &[4, 5, 6, 7], 990)); + o.extend((10..13).map(|r| out("c", r, 990, OutcomeKind::Ok))); + o.extend((20..23).map(|r| out("d", r, 990, OutcomeKind::Ok))); + o.extend((30..33).map(|r| out("e", r, 990, OutcomeKind::Ok))); + let e = evaluate(1000, &s, &d, &o); + let names: Vec<&str> = e.actions.iter().map(Action::drone).collect(); + assert_eq!(names, ["b", "a"]); + // Limit of one: only the worst. + let one = HealthSettings { + max_auto_quarantined: 1, + ..s.clone() + }; + let e = evaluate(1000, &one, &d, &o); + assert_eq!(e.actions.len(), 1); + assert_eq!(e.drones["a"].state, HealthState::Recommended); + // Auto off: recommendations only. + let off = HealthSettings { + auto_quarantine: false, + ..s + }; + let e = evaluate(1000, &off, &d, &o); + assert!(e.actions.is_empty()); + } + + #[test] + fn cooldown_then_probation_then_proof_releases_and_failure_requarantines() { + let s = settings(); + let mut a = drone("a"); + let d_b = drone("b"); + a.hold = Hold::Auto { + since: 1000, + strikes: 0, + }; + // Before the cooldown ends: nothing. + let e = evaluate(1000 + 1799, &s, &[a.clone(), d_b.clone()], &[]); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].probation_at, Some(2800)); + // After: probation. + let e = evaluate(2800, &s, &[a.clone(), d_b.clone()], &[]); + assert_eq!(e.actions, [Action::BeginProbation { drone: "a".into() }]); + + a.hold = Hold::Probation { + since: 2800, + strikes: 0, + }; + // One success is not enough; old successes before probation do not count. + let o = vec![ + out("a", 1, 2000, OutcomeKind::Ok), + out("a", 2, 2900, OutcomeKind::Ok), + ]; + let e = evaluate(3000, &s, &[a.clone(), d_b.clone()], &o); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].proof_jobs_done, 1); + // Two successes: released. + let o = vec![ + out("a", 2, 2900, OutcomeKind::Ok), + out("a", 3, 2950, OutcomeKind::Ok), + ]; + let e = evaluate(3000, &s, &[a.clone(), d_b.clone()], &o); + assert_eq!(e.actions, [Action::Release { drone: "a".into() }]); + // A host failure on probation wins over successes. + let mut o = o; + o.push(out("a", 4, 2960, OutcomeKind::HostFailure)); + let e = evaluate(3000, &s, &[a, d_b], &o); + let Action::Requarantine { strikes, .. } = &e.actions[0] else { + panic!("{:?}", e.actions); + }; + assert_eq!(*strikes, 1); + } + + #[test] + fn cooldown_doubles_per_strike_up_to_the_cap() { + let s = settings(); + assert_eq!(s.cooldown_for(0), 1800); + assert_eq!(s.cooldown_for(1), 3600); + assert_eq!(s.cooldown_for(2), 7200); + assert_eq!(s.cooldown_for(3), 14_400); + assert_eq!(s.cooldown_for(10), 14_400); + assert_eq!(s.cooldown_for(u32::MAX), 14_400); + } + + #[test] + fn evidence_before_a_return_to_service_is_forgotten() { + let s = settings(); + let mut a = drone("a"); + a.cleared_at = Some(995); + let d = [a, drone("b")]; + let e = evaluate(1000, &s, &d, &fails("a", &[1, 2, 3], 990)); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].host_failures, 0); + } + + #[test] + fn operator_quarantines_are_never_lifted_automatically() { + let s = settings(); + let mut a = drone("a"); + a.hold = Hold::Operator; + let e = evaluate(10_000_000, &s, &[a], &[]); + assert!(e.actions.is_empty()); + assert_eq!(e.drones["a"].source.as_deref(), Some("operator")); + } + + #[test] + fn evaluation_is_deterministic() { + let s = settings(); + let d = [drone("a"), drone("b"), drone("c"), drone("d")]; + let mut o = fails("a", &[1, 2, 3], 990); + o.extend(fails("b", &[4, 5, 6], 990)); + assert_eq!(evaluate(1000, &s, &d, &o), evaluate(1000, &s, &d, &o)); + } + + #[test] + fn settings_validate() { + assert!(HealthSettings::default().validate().is_ok()); + for bad in [ + HealthSettings { + half_life_secs: 0, + ..HealthSettings::default() + }, + HealthSettings { + corroboration_runs: 0, + ..HealthSettings::default() + }, + HealthSettings { + recommend_score_permille: 9000, + ..HealthSettings::default() + }, + HealthSettings { + systemic_share_pct: 0, + ..HealthSettings::default() + }, + ] { + assert!(bad.validate().is_err()); + } + } +} diff --git a/src/homeostasis.rs b/src/homeostasis.rs new file mode 100644 index 0000000..1c40e13 --- /dev/null +++ b/src/homeostasis.rs @@ -0,0 +1,511 @@ +//! Homeostasis: a damped controller for the warm runner floors. +//! +//! The base `warm_floors` and the clock-based `warm_schedule` say how many idle +//! runners to keep warm. This controller adds a negative-feedback loop on +//! queue latency: when the oldest queued job of a `(class, scope)` has waited +//! longer than the target for a sustained period and the fleet has idle room, it +//! raises that key's dynamic floor one step; when demand has subsided it lowers +//! it one at a time. The effective floor is `max(schedule floor, dynamic +//! level)`, so the schedule stays a minimum. +//! +//! Damping, all of it configurable: +//! * **Sustained, not instantaneous.** The latency must stay above the target +//! for `raise_after_secs` without a break (and below the low-water mark for +//! `lower_after_secs` to lower). A reading in between resets both timers: that +//! dead band is the hysteresis. +//! * **Step size and cooldown.** One `step` up or one runner down per change, +//! and `cooldown_secs` between changes of a key. +//! * **Rate limit.** At most `max_changes_per_hour` changes per key. +//! * **Bounds.** A per-class `max_floor`, and no raise without idle room for it. +//! +//! The controller reads no clock (the caller passes `now`) and keeps its state +//! in memory: after a restart it starts from the schedule and re-learns. + +use std::collections::{BTreeMap, VecDeque}; + +use serde::{Deserialize, Serialize}; + +use crate::classes::{Class, Scope}; + +/// Tunables; opt-in via `enabled`. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct HomeostasisSettings { + pub enabled: bool, + /// Queue latency (oldest queued job age) the controller steers below. + pub target_secs: u64, + /// Latency must exceed the target this long without a break to raise. + pub raise_after_secs: u64, + /// Latency must stay below `low_pct` of the target (or nothing queued) + /// this long without a break to lower. + pub lower_after_secs: u64, + pub low_pct: u32, + /// Runners added per raise. + pub step: u32, + /// Minimum time between two changes of the same key. + pub cooldown_secs: u64, + pub max_changes_per_hour: u32, + /// Ceiling on the dynamic floor per class. + pub max_floor: BTreeMap, +} + +impl Default for HomeostasisSettings { + fn default() -> Self { + Self { + enabled: false, + target_secs: 60, + raise_after_secs: 120, + lower_after_secs: 900, + low_pct: 25, + step: 1, + cooldown_secs: 120, + max_changes_per_hour: 6, + max_floor: BTreeMap::from([(Class::Light, 6), (Class::Medium, 3), (Class::Heavy, 1)]), + } + } +} + +impl HomeostasisSettings { + pub fn validate(&self) -> Result<(), String> { + if self.target_secs == 0 || self.raise_after_secs == 0 || self.lower_after_secs == 0 { + return Err( + "homeostasis target_secs, raise_after_secs and lower_after_secs must be positive" + .into(), + ); + } + if self.step == 0 || self.max_changes_per_hour == 0 { + return Err("homeostasis step and max_changes_per_hour must be positive".into()); + } + if self.low_pct >= 100 { + return Err("homeostasis low_pct must be below 100".into()); + } + Ok(()) + } + + fn cap(&self, class: Class) -> u32 { + self.max_floor.get(&class).copied().unwrap_or(0) + } + + fn low_water_secs(&self) -> u64 { + self.target_secs * u64::from(self.low_pct) / 100 + } +} + +/// One key's reading for a step. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct KeyInput { + pub class: Class, + pub scope: Scope, + /// Age of the oldest queued job; 0 with nothing queued. + pub latency_secs: u64, + pub queued: u32, + /// The floor from the base config and the clock schedule. + pub schedule_floor: u32, + /// Slots the fleet could still hand out for the class. + pub idle_slots: u32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Phase { + /// Latency within the dead band. + Steady, + /// Latency above the target, timing toward a raise. + Over, + /// Latency at or below the low-water mark, timing toward a lower. + Under, +} + +/// What `/status` shows per key. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct KeyReport { + pub class: Class, + pub scope: Scope, + pub latency_secs: u64, + pub schedule_floor: u32, + /// The controller's own level (0 when idle). + pub level: u32, + /// What the allocator uses: `max(schedule_floor, level)`. + pub effective_floor: u32, + pub phase: Phase, + /// How long the current phase has lasted without a break. + pub phase_secs: u64, + pub last_change_at: Option, + pub changes_last_hour: u32, + /// Why a due change was held back. + pub blocked: Option, +} + +#[derive(Debug, Clone, Default)] +struct KeyState { + level: u32, + over_since: Option, + under_since: Option, + last_change: Option, + changes: VecDeque, +} + +#[derive(Debug, Default)] +pub struct Controller { + state: BTreeMap<(Class, Scope), KeyState>, +} + +impl Controller { + pub fn new() -> Self { + Self::default() + } + + /// One control step for every input key. Returns the report per key. + pub fn step( + &mut self, + now: i64, + s: &HomeostasisSettings, + inputs: &[KeyInput], + ) -> Vec { + inputs.iter().map(|i| self.step_key(now, s, i)).collect() + } + + fn step_key(&mut self, now: i64, s: &HomeostasisSettings, i: &KeyInput) -> KeyReport { + let st = self.state.entry((i.class, i.scope)).or_default(); + while st.changes.front().is_some_and(|&t| now - t >= 3600) { + st.changes.pop_front(); + } + let effective = i.schedule_floor.max(st.level); + let over = i.latency_secs > s.target_secs; + let under = i.queued == 0 || i.latency_secs <= s.low_water_secs(); + + // Phase timers: a reading outside a phase breaks it. + if over { + st.over_since.get_or_insert(now); + st.under_since = None; + } else { + st.over_since = None; + if under { + st.under_since.get_or_insert(now); + } else { + st.under_since = None; + } + } + let (phase, since) = match (st.over_since, st.under_since) { + (Some(t), _) => (Phase::Over, Some(t)), + (_, Some(t)) => (Phase::Under, Some(t)), + _ => (Phase::Steady, None), + }; + + let cooled = st + .last_change + .is_none_or(|t| now - t >= i64::try_from(s.cooldown_secs).unwrap_or(i64::MAX)); + let limited = u32::try_from(st.changes.len()).unwrap_or(u32::MAX) >= s.max_changes_per_hour; + let held = |why: &str| Some(why.to_owned()); + let mut blocked = None; + + match phase { + Phase::Over + if since.is_some_and(|t| { + now - t >= i64::try_from(s.raise_after_secs).unwrap_or(i64::MAX) + }) => + { + let want = (effective + s.step).min(s.cap(i.class)); + if want <= effective { + blocked = held("at the floor ceiling for the class"); + } else if i.idle_slots < want - effective { + blocked = held("no idle capacity for another warm runner"); + } else if !cooled { + blocked = held("cooling down"); + } else if limited { + blocked = held("hourly change limit reached"); + } else { + st.level = want; + st.last_change = Some(now); + st.changes.push_back(now); + // Another step needs another sustained period. + st.over_since = Some(now); + } + } + Phase::Under + if since.is_some_and(|t| { + now - t >= i64::try_from(s.lower_after_secs).unwrap_or(i64::MAX) + }) && st.level > i.schedule_floor => + { + if !cooled { + blocked = held("cooling down"); + } else if limited { + blocked = held("hourly change limit reached"); + } else { + st.level -= 1; + st.last_change = Some(now); + st.changes.push_back(now); + st.under_since = Some(now); + } + } + _ => {} + } + // A level at or below the schedule is inert; drop it so it cannot + // resurface when the schedule window closes. + if st.level <= i.schedule_floor { + st.level = 0; + } + + let level = st.level; + KeyReport { + class: i.class, + scope: i.scope, + latency_secs: i.latency_secs, + schedule_floor: i.schedule_floor, + level, + effective_floor: i.schedule_floor.max(level), + phase, + phase_secs: since.map_or(0, |t| u64::try_from(now - t).unwrap_or(0)), + last_change_at: st.last_change, + changes_last_hour: u32::try_from(st.changes.len()).unwrap_or(u32::MAX), + blocked, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn settings() -> HomeostasisSettings { + HomeostasisSettings { + enabled: true, + ..HomeostasisSettings::default() + } + } + + fn key(latency: u64, queued: u32, schedule: u32, idle: u32) -> KeyInput { + KeyInput { + class: Class::Light, + scope: Scope::Pr, + latency_secs: latency, + queued, + schedule_floor: schedule, + idle_slots: idle, + } + } + + fn one(c: &mut Controller, now: i64, s: &HomeostasisSettings, k: KeyInput) -> KeyReport { + c.step(now, s, &[k]).remove(0) + } + + #[test] + fn a_sustained_overshoot_raises_one_step_then_waits_for_another_sustained_period() { + let s = settings(); + let mut c = Controller::new(); + let hot = key(200, 3, 0, 8); + assert_eq!(one(&mut c, 0, &s, hot).effective_floor, 0); + assert_eq!(one(&mut c, 60, &s, hot).effective_floor, 0); + // 120s of continuous overshoot: raise by one step. + let r = one(&mut c, 120, &s, hot); + assert_eq!((r.level, r.effective_floor, r.phase), (1, 1, Phase::Over)); + // Still hot, but a further step needs another full sustained period. + assert_eq!(one(&mut c, 180, &s, hot).level, 1); + assert_eq!(one(&mut c, 240, &s, hot).level, 2); + assert_eq!(one(&mut c, 360, &s, hot).level, 3); + } + + #[test] + fn a_blip_never_raises_and_flapping_readings_reset_the_timer() { + let s = settings(); + let mut c = Controller::new(); + for t in (0..3600).step_by(10) { + // Alternates above and below the target every tick. + let latency = if (t / 10) % 2 == 0 { 300 } else { 40 }; + let r = one(&mut c, t, &s, key(latency, 1, 0, 8)); + assert_eq!(r.level, 0, "t={t}"); + } + } + + #[test] + fn the_ceiling_and_idle_capacity_bound_the_raise() { + let s = settings(); + let mut c = Controller::new(); + // No idle room: held back, with the reason. + let r = (0..=12) + .map(|n| one(&mut c, n * 60, &s, key(500, 5, 0, 0))) + .last() + .unwrap(); + assert_eq!(r.level, 0); + assert!(r.blocked.as_deref().unwrap().contains("idle capacity")); + // Heavy has a ceiling of 1. + let mut c = Controller::new(); + let heavy = KeyInput { + class: Class::Heavy, + ..key(500, 5, 0, 8) + }; + let r = (0..=20) + .map(|n| one(&mut c, n * 130, &s, heavy)) + .last() + .unwrap(); + assert_eq!(r.level, 1); + assert!(r.blocked.as_deref().unwrap().contains("ceiling")); + } + + #[test] + fn demand_subsiding_lowers_one_runner_at_a_time_after_a_quiet_period() { + // Room in the hourly limit, so the lowering is not rate limited. + let s = HomeostasisSettings { + max_changes_per_hour: 20, + ..settings() + }; + let mut c = Controller::new(); + for n in 0..=6 { + one(&mut c, n * 130, &s, key(500, 5, 0, 8)); + } + let peak = one(&mut c, 800, &s, key(500, 5, 0, 8)).level; + assert!(peak >= 3, "{peak}"); + // Quiet: nothing for a while, then one step down per quiet period. + let t0 = 1000; + assert_eq!(one(&mut c, t0, &s, key(0, 0, 0, 8)).level, peak); + assert_eq!(one(&mut c, t0 + 899, &s, key(0, 0, 0, 8)).level, peak); + assert_eq!(one(&mut c, t0 + 900, &s, key(0, 0, 0, 8)).level, peak - 1); + assert_eq!(one(&mut c, t0 + 1799, &s, key(0, 0, 0, 8)).level, peak - 1); + assert_eq!(one(&mut c, t0 + 1800, &s, key(0, 0, 0, 8)).level, peak - 2); + } + + #[test] + fn the_dead_band_between_low_water_and_target_changes_nothing() { + let s = settings(); + let mut c = Controller::new(); + for n in 0..=3 { + one(&mut c, n * 130, &s, key(500, 5, 0, 8)); + } + let level = one(&mut c, 400, &s, key(500, 5, 0, 8)).level; + // 40s is below the 60s target and above the 15s low-water mark. + for t in (500..10_000).step_by(30) { + let r = one(&mut c, t, &s, key(40, 2, 0, 8)); + assert_eq!((r.level, r.phase), (level, Phase::Steady), "t={t}"); + } + } + + #[test] + fn the_schedule_is_a_minimum_and_an_inert_level_does_not_resurface() { + let s = settings(); + let mut c = Controller::new(); + // Schedule floor 3 covers the controller's level of 1: effective 3. + for n in 0..=2 { + one(&mut c, n * 130, &s, key(500, 5, 3, 8)); + } + let r = one(&mut c, 300, &s, key(500, 5, 3, 8)); + assert_eq!(r.effective_floor, 4.max(r.level)); + // The window closes: the dynamic level, not the schedule, now decides. + let mut c = Controller::new(); + let r = one(&mut c, 0, &s, key(0, 0, 3, 8)); + assert_eq!((r.level, r.effective_floor), (0, 3)); + let r = one(&mut c, 10, &s, key(0, 0, 0, 8)); + assert_eq!(r.effective_floor, 0); + } + + #[test] + fn cooldown_and_the_hourly_limit_cap_the_change_rate() { + let s = HomeostasisSettings { + raise_after_secs: 10, + cooldown_secs: 300, + max_changes_per_hour: 3, + max_floor: BTreeMap::from([(Class::Light, 50)]), + ..settings() + }; + let mut c = Controller::new(); + let mut changes = Vec::new(); + let mut last = 0; + for t in (0..3600).step_by(10) { + let r = one(&mut c, t, &s, key(500, 5, 0, 100)); + if r.level != last { + changes.push(t); + last = r.level; + } + } + assert_eq!(changes.len(), 3, "{changes:?}"); + assert!( + changes.windows(2).all(|w| w[1] - w[0] >= 300), + "{changes:?}" + ); + // An hour after the first change the limit window slides. + let r = one(&mut c, 3700, &s, key(500, 5, 0, 100)); + assert!(r.changes_last_hour <= 3); + } + + #[test] + fn a_square_wave_load_does_not_thrash() { + let s = settings(); + let mut c = Controller::new(); + let mut events: Vec<(i64, i64)> = Vec::new(); // (time, +1/-1) + let mut last = 0u32; + // Eight hours: 10 minutes of heavy queueing, 10 minutes quiet, repeat. + for t in (0..8 * 3600).step_by(10) { + let busy = (t / 600) % 2 == 0; + let k = if busy { + key(300, 6, 0, 20) + } else { + key(0, 0, 0, 20) + }; + let r = one(&mut c, t, &s, k); + if r.level != last { + events.push((t, if r.level > last { 1 } else { -1 })); + last = r.level; + } + } + // Bounded by the rate limit, and no two changes inside the cooldown. + assert!(events.len() <= 8 * 6, "{events:?}"); + assert!( + events.windows(2).all(|w| w[1].0 - w[0].0 >= 120), + "{events:?}" + ); + // Opposite-direction changes never happen back to back within the + // lower-after period: a raise is never undone by a quick lower. + assert!( + events + .windows(2) + .all(|w| w[0].1 == w[1].1 || w[1].0 - w[0].0 >= 900), + "{events:?}" + ); + } + + #[test] + fn keys_are_independent_and_the_controller_is_deterministic() { + let s = settings(); + let run = || { + let mut c = Controller::new(); + let mut out = Vec::new(); + for t in (0..1200).step_by(30) { + let inputs = [ + key(300, 4, 0, 8), + KeyInput { + scope: Scope::Main, + ..key(0, 0, 1, 8) + }, + ]; + out.push(c.step(t, &s, &inputs)); + } + out + }; + let a = run(); + assert_eq!(a, run()); + let last = a.last().unwrap(); + assert!(last[0].level >= 1); + assert_eq!((last[1].level, last[1].effective_floor), (0, 1)); + } + + #[test] + fn settings_validate() { + assert!(HomeostasisSettings::default().validate().is_ok()); + for bad in [ + HomeostasisSettings { + target_secs: 0, + ..HomeostasisSettings::default() + }, + HomeostasisSettings { + step: 0, + ..HomeostasisSettings::default() + }, + HomeostasisSettings { + low_pct: 100, + ..HomeostasisSettings::default() + }, + ] { + assert!(bad.validate().is_err()); + } + assert!(!HomeostasisSettings::default().enabled); + } +} diff --git a/src/lib.rs b/src/lib.rs index 5be323d..6c778b0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,17 +8,31 @@ //! the GitHub Actions queue, decides what to lease where, mints single-use //! JIT runner configs and serves the drone API. +pub mod admin; pub mod allocator; pub mod api; +pub mod bid; pub mod classes; pub mod config; pub mod drone; pub mod enroll; +pub mod envcontract; pub mod github; +pub mod health; +pub mod homeostasis; pub mod janitor; pub mod probe; +pub mod profile; pub mod queen; +pub mod reclaim; +pub mod rollout; +pub mod run; pub mod schedule; pub mod secret; pub mod store; +pub mod tailnet; pub mod telemetry; +pub mod units; +pub mod units_api; +pub mod update; +pub mod wake; diff --git a/src/main.rs b/src/main.rs index 2eaa67c..a752cde 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use anyhow::{Context, Result, bail}; +use borg::admin; use borg::api; use borg::config::{JanitorSettings, QueenConfig}; use borg::drone::client::QueenClient; @@ -15,6 +16,8 @@ use borg::janitor::Janitor; use borg::probe::{self, system::SystemHost}; use borg::queen::Queen; use borg::store::Store; +use borg::tailnet::{TailnetEnroll, TailscaledWhois}; +use borg::update; use clap::{Args, Parser, Subcommand}; #[derive(Parser)] @@ -25,6 +28,10 @@ struct Cli { } #[derive(Subcommand)] +#[allow( + clippy::large_enum_variant, + reason = "parsed once at startup; boxing a clap Args payload only adds noise" +)] enum Command { /// Probe this host and print its capacity as JSON. Probe { @@ -40,6 +47,13 @@ enum Command { /// Summarize stored job telemetry: time by job and step, queue wait per /// class, cancellation share and host-caused failures per drone. Report(ReportArgs), + /// Run a command on the fleet when capacity is free, otherwise here. + /// Everything after `--` is the command. + Run(RunArgs), + /// Update this host's borg (and, with --runner-only, just its actions-runner). + /// Downloads and verifies the release, drains the service, swaps the binary + /// (keeping borg.prev), restarts it, and rolls back if it does not come up. + Update(UpdateArgs), /// Run or administer the queen (the coordinator). Queen { #[command(subcommand)] @@ -47,6 +61,76 @@ enum Command { }, } +#[derive(Args)] +struct UpdateArgs { + /// Updater configuration (root-owned). + #[arg(long, default_value = borg::update::config::DEFAULT_DRONE_CONFIG)] + config: PathBuf, + /// Install exactly this version (default: the newest release). + #[arg(long)] + version: Option, + /// Only report what would change. + #[arg(long)] + check: bool, + /// Update only the pinned actions-runner; no restart. + #[arg(long, conflicts_with = "version")] + runner_only: bool, + /// Read the target version from the drone's update request and leave a + /// result for it (how the update path unit runs this). + #[arg(long)] + from_request: bool, + /// Print the outcome as JSON. + #[arg(long)] + json: bool, +} + +#[derive(Args)] +struct RunArgs { + /// Queen base URL. Without it the command simply runs here. + #[arg(long, env = "BORG_QUEEN")] + queen: Option, + /// File holding the client token (from `borg queen enroll-client`). + /// Default: BORG_CLIENT_TOKEN, else ~/.config/borg/client-token. + #[arg(long, env = "BORG_CLIENT_TOKEN_FILE")] + token_file: Option, + /// Run against the staged content (`git write-tree`), for pre-commit hooks. + /// Default is HEAD's tree. + #[arg(long, conflicts_with = "tree")] + index: bool, + /// Run against this tree or commit instead. + #[arg(long)] + tree: Option, + /// Repository as owner/name (default: from the origin remote). + #[arg(long)] + repo: Option, + /// Platform the command needs (default: this machine's, e.g. linux-x86_64). + #[arg(long)] + platform: Option, + /// Runner class that sizes the unit. + #[arg(long, default_value = "medium")] + class: String, + /// Digest of the toolchain the result depends on (part of the cache key). + #[arg(long, env = "BORG_TOOLCHAIN", default_value = "")] + toolchain: String, + /// Time limit for the command on the fleet, in seconds. + #[arg(long)] + timeout: Option, + /// Run even if an identical command already passed for this tree. + #[arg(long)] + force: bool, + /// Seconds to wait for a drone before running here. + #[arg(long, env = "BORG_RUN_WAIT_SECS", default_value_t = 60)] + wait_secs: u64, + /// Upload a git bundle of the tree (since the merge base with origin/main) + /// to the queen when no remote has it, so a drone can get it. Nothing is + /// pushed anywhere. + #[arg(long, env = "BORG_PUBLISH")] + publish: bool, + /// The command and its arguments. + #[arg(last = true, required = true)] + argv: Vec, +} + #[derive(Args)] struct ReportArgs { /// Fetch from a running queen, e.g. http://queen.internal:8080 @@ -77,13 +161,22 @@ struct ConfigArg { enum QueenCommand { /// Poll the queue, mint runners and serve the drone API. Run(ConfigArg), - /// Enroll a drone and print its one-time token. + /// Enroll a drone. Prints only the one-time token on stdout (messages go + /// to stderr), so `TOKEN=$(borg queen enroll ...)` is safe. Enroll { #[command(flatten)] config: ConfigArg, /// Drone name: lowercase letters, digits and dashes. name: String, }, + /// Enroll a client (a developer machine or agent that submits work units) + /// and print its one-time token. + EnrollClient { + #[command(flatten)] + config: ConfigArg, + /// Client name: lowercase letters, digits and dashes. + name: String, + }, /// Stop granting a drone new leases (running jobs finish). Quarantine { #[command(flatten)] @@ -94,6 +187,11 @@ enum QueenCommand { /// Drone name. name: String, }, + /// Administer the quorum rollout gate. + Rollout { + #[command(subcommand)] + command: RolloutCommand, + }, /// Lift a quarantine. Release { #[command(flatten)] @@ -101,6 +199,37 @@ enum QueenCommand { /// Drone name. name: String, }, + /// List or remove enrolled drones. + Drone { + #[command(subcommand)] + command: QueenDroneCommand, + }, +} + +#[derive(Subcommand)] +enum QueenDroneCommand { + /// Show every enrolled drone: state, last heartbeat, host and size. + List(ConfigArg), + /// Remove a drone: revokes its token and drops its unfinished leases. + Remove { + #[command(flatten)] + config: ConfigArg, + name: String, + }, +} + +#[derive(Subcommand)] +enum RolloutCommand { + /// Halt the current rollout now and tell every drone the previous version. + Halt { + #[command(flatten)] + config: ConfigArg, + #[arg(long, default_value = "operator")] + reason: String, + }, + /// Forget the current rollout so the next pass starts it over (use after + /// fixing the cause of a halt, to retry the same target). + Reset(ConfigArg), } #[derive(Args)] @@ -108,9 +237,16 @@ struct DroneConfigArg { /// Path to the drone's JSON config file. #[arg(long, env = "BORG_DRONE_CONFIG")] config: PathBuf, + /// File holding the enrollment token (overrides `token_file` in the config). + #[arg(long)] + token_file: Option, } #[derive(Subcommand)] +#[allow( + clippy::large_enum_variant, + reason = "parsed once at startup; boxing a clap Args payload only adds noise" +)] enum DroneCommand { /// Heartbeat to the queen and run the leases it grants. Run(DroneConfigArg), @@ -123,9 +259,13 @@ struct InstallArgs { /// Queen base URL, e.g. https://queen.example.test:8443 #[arg(long)] queen: String, - /// Enrollment token from `borg queen enroll`. + /// Enrollment token from `borg queen enroll`. Prefer `--token-file`: a + /// token on the command line is visible in `ps`. #[arg(long, env = "BORG_DRONE_TOKEN", hide_env_values = true)] - token: String, + token: Option, + /// File holding the enrollment token. + #[arg(long, conflicts_with = "token")] + token_file: Option, /// PEM CA certificate that signs the queen's certificate (pins the CA). #[arg(long)] ca_cert: Option, @@ -141,9 +281,36 @@ struct InstallArgs { /// Owner cap: most concurrent leases. #[arg(long)] max_leases: Option, - /// Service account (Linux: created if missing; macOS: must exist). + /// Account the drone daemon runs as (Linux: created if missing; macOS: must + /// exist). It owns the config and the enrollment token. #[arg(long, default_value = "borg")] run_as: String, + /// Linux: the separate unprivileged account leases run as, so job code + /// cannot read the drone's token (created if missing). + #[arg(long, default_value = "borg-job")] + job_user: String, + /// Run leases as the daemon's own account (they can then read its token). + #[arg(long)] + no_job_user: bool, + /// Where host caches live (default: under the job user's home). + #[arg(long)] + cache_root: Option, + /// Take ownership of an existing --cache-root, keeping its contents (for + /// example a host's warm compile cache). + #[arg(long, requires = "cache_root")] + adopt_cache: bool, + /// Add this label to every runner the drone registers; repeatable. + #[arg(long = "extra-label")] + extra_labels: Vec, + /// KEY=VALUE injected into every lease's environment; repeatable. + #[arg(long = "runner-env", value_parser = parse_env_pair)] + runner_env: Vec<(String, String)>, + /// Absolute path of a .sh script run when a job starts. + #[arg(long)] + job_started_hook: Option, + /// Absolute path of a .sh script run when a job completes. + #[arg(long)] + job_completed_hook: Option, /// Install under this root instead of / (staging and tests). #[arg(long, default_value = "/")] root: PathBuf, @@ -155,11 +322,24 @@ struct InstallArgs { /// such as git@host: with a read-only deploy key for private repositories. #[arg(long)] git_remote_base: Option, + /// Release manifest `borg update` reads (default: the public release). + #[arg(long)] + update_manifest_url: Option, + /// Do not install the self-update units. + #[arg(long)] + no_self_update: bool, /// Print what would change without changing anything. #[arg(long)] dry_run: bool, } +fn parse_env_pair(s: &str) -> Result<(String, String), String> { + s.split_once('=') + .filter(|(k, _)| !k.is_empty()) + .map(|(k, v)| (k.to_owned(), v.to_owned())) + .ok_or_else(|| "expected KEY=VALUE".to_owned()) +} + fn private_key(cfg: &QueenConfig) -> Result> { match &cfg.github.private_key_path { Some(path) => { @@ -234,7 +414,20 @@ async fn run_queen(path: &Path) -> Result<()> { } None => None, }; + let tailnet = cfg.tailnet_enroll.clone().map(|settings| { + tracing::info!( + "tailnet enrollment enabled via {}", + settings.socket_path.display() + ); + TailnetEnroll { + verifier: Arc::new(TailscaledWhois::new(settings.socket_path.clone())), + settings, + } + }); let mut queen = Queen::new(cfg, store, github); + if let Some(t) = tailnet { + queen = queen.with_tailnet_enroll(t); + } if let Some(j) = janitor { queen = queen.with_janitor(j); } @@ -244,19 +437,12 @@ async fn run_queen(path: &Path) -> Result<()> { } async fn enroll_drone(path: &Path, name: &str) -> Result<()> { - if !enroll::valid_drone_name(name) { - bail!( - "drone name must be 1-40 lowercase letters, digits or dashes, not starting with a dash" - ); - } let cfg = QueenConfig::load(path)?; let store = open_store(&cfg).await?; - let one_time_code = enroll::new_token(); let now = time::OffsetDateTime::now_utc().unix_timestamp(); - store - .enroll(name, &enroll::hash_token(&one_time_code), now) - .await?; + let one_time_code = admin::enroll_drone(&store, name, now).await?; eprintln!("Enrolled drone {name}. The token is shown once; the queen stores only its hash."); + // stdout carries the token and nothing else. println!("{one_time_code}"); Ok(()) } @@ -297,6 +483,24 @@ async fn report(a: ReportArgs) -> Result<()> { Ok(()) } +async fn enroll_client(path: &Path, name: &str) -> Result<()> { + if !enroll::valid_drone_name(name) { + bail!( + "client name must be 1-40 lowercase letters, digits or dashes, not starting with a dash" + ); + } + let cfg = QueenConfig::load(path)?; + let store = open_store(&cfg).await?; + let token = enroll::new_token(); + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + store + .enroll_client(name, &enroll::hash_token(&token), now) + .await?; + eprintln!("Enrolled client {name}. The token is shown once; the queen stores only its hash."); + println!("{token}"); + Ok(()) +} + async fn set_quarantine(path: &Path, name: &str, reason: Option<&str>) -> Result<()> { let cfg = QueenConfig::load(path)?; let store = open_store(&cfg).await?; @@ -310,8 +514,52 @@ async fn set_quarantine(path: &Path, name: &str, reason: Option<&str>) -> Result Ok(()) } -async fn run_drone(path: &Path) -> Result<()> { - let cfg = DroneConfig::load(path)?; +async fn rollout_admin(path: &Path, halt: Option<&str>) -> Result<()> { + let cfg = QueenConfig::load(path)?; + let store = open_store(&cfg).await?; + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + match halt { + Some(reason) => { + let Some(mut st) = store.rollout_state().await? else { + bail!("no rollout is in progress"); + }; + st.phase = borg::rollout::Phase::Halted; + st.halted_reason = Some(format!("operator: {reason}")); + store.save_rollout_state(&st, now).await?; + store + .log_rollout( + now, + &format!("rollout of {} HALTED by an operator: {reason}", st.target), + ) + .await?; + eprintln!( + "Halted the rollout of {}; drones are told {}.", + st.target, st.previous + ); + } + None => { + store.clear_rollout_state().await?; + store + .log_rollout(now, "rollout state reset by an operator") + .await?; + eprintln!("Rollout state cleared; the next pass starts over."); + } + } + Ok(()) +} + +async fn run_update(args: update::cli::Args) -> Result<()> { + match update::cli::run(args).await? { + 0 => Ok(()), + _ => bail!("the update was rolled back; see the line above"), + } +} + +async fn run_drone(path: &Path, token_file: Option) -> Result<()> { + let mut cfg = DroneConfig::load(path)?; + if token_file.is_some() { + cfg.token_file = token_file; + } let token = cfg.load_token()?; let client = QueenClient::new(&cfg, token)?; let drone = Drone::new(cfg, client, Arc::new(SystemHost), Arc::new(SystemExec)); @@ -329,15 +577,34 @@ fn install_drone(a: InstallArgs) -> Result<()> { std::env::consts::ARCH ) })?; + let token = match (&a.token, &a.token_file) { + (Some(t), _) => t.clone(), + (None, Some(path)) => std::fs::read_to_string(path) + .with_context(|| format!("read {}", path.display()))? + .trim() + .to_owned(), + (None, None) => bail!("pass --token-file (preferred) or --token"), + }; + let job_user = (platform == Platform::Linux && !a.no_job_user).then(|| a.job_user.clone()); + if platform == Platform::MacOs { + eprintln!("note: macOS runs jobs as the same account as the drone (no separate job user)"); + } let opts = InstallOptions { queen_url: a.queen, - token: borg::secret::Secret::new(a.token), + token: borg::secret::Secret::new(token), ca_cert: a.ca_cert, allow_insecure: a.allow_insecure, max_vcpu: a.max_vcpu, max_mem_mib: a.max_mem_mib, max_leases: a.max_leases, run_as: a.run_as, + job_user, + cache_root: a.cache_root, + adopt_cache: a.adopt_cache, + extra_labels: a.extra_labels, + runner_env: a.runner_env.into_iter().collect(), + job_started_hook: a.job_started_hook, + job_completed_hook: a.job_completed_hook, skip_root_check: a.root != Path::new("/"), root: a.root, platform, @@ -347,6 +614,8 @@ fn install_drone(a: InstallArgs) -> Result<()> { drain_timeout_secs: 3600, git_mirrors: a.git_mirrors, git_remote_base: a.git_remote_base, + update_manifest_url: a.update_manifest_url, + no_self_update: a.no_self_update, }; for line in install::install(&opts, &SystemExec)? { println!("{line}"); @@ -354,9 +623,98 @@ fn install_drone(a: InstallArgs) -> Result<()> { Ok(()) } +async fn queen_drone_admin(command: QueenDroneCommand) -> Result<()> { + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + match command { + QueenDroneCommand::List(c) => { + let cfg = QueenConfig::load(&c.config)?; + let store = open_store(&cfg).await?; + for line in admin::list_lines(&store.drones().await?, now, cfg.drone_stale_secs) { + println!("{line}"); + } + } + QueenDroneCommand::Remove { config, name } => { + let cfg = QueenConfig::load(&config.config)?; + let store = open_store(&cfg).await?; + let expired = admin::remove_drone(&store, &name).await?; + println!( + "removed drone {name}; its token no longer works; {expired} unfinished lease(s) dropped" + ); + } + } + Ok(()) +} + +fn client_token(file: Option<&Path>) -> Option { + if let Ok(t) = std::env::var("BORG_CLIENT_TOKEN") { + return Some(t.trim().to_owned()).filter(|t| !t.is_empty()); + } + let path = file.map(Path::to_path_buf).or_else(|| { + std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".config/borg/client-token")) + })?; + let text = std::fs::read_to_string(path).ok()?; + Some(text.trim().to_owned()).filter(|t| !t.is_empty()) +} + +async fn run_command(a: RunArgs) -> u8 { + let dir = match std::env::current_dir() { + Ok(d) => d, + Err(e) => { + eprintln!("borg run: {e}"); + return 2; + } + }; + let opts = borg::run::RunOptions { + token: a + .queen + .as_ref() + .and_then(|_| client_token(a.token_file.as_deref())) + .map(borg::secret::Secret::new), + queen: a.queen, + dir, + tree: match (a.index, a.tree) { + (_, Some(t)) => borg::run::TreeSource::Explicit(t), + (true, None) => borg::run::TreeSource::Index, + (false, None) => borg::run::TreeSource::Head, + }, + repo: a.repo, + platform: a.platform.unwrap_or_else(borg::run::local_platform), + class: a.class, + toolchain: a.toolchain, + timeout_secs: a.timeout, + force: a.force, + wait_secs: a.wait_secs, + publish: a.publish, + argv: a.argv, + }; + let code = borg::run::run(&opts, &mut std::io::stdout()).await; + u8::try_from(code).unwrap_or(if code < 0 { 1 } else { u8::MAX }) +} + #[tokio::main] -async fn main() -> Result<()> { - match Cli::parse().command { +async fn main() -> Result { + let cli = Cli::parse(); + if let Command::Run(a) = cli.command { + return Ok(std::process::ExitCode::from(run_command(a).await)); + } + dispatch(cli.command).await?; + Ok(std::process::ExitCode::SUCCESS) +} + +async fn dispatch(command: Command) -> Result<()> { + match command { + Command::Run(_) => Ok(()), + Command::Update(a) => { + run_update(update::cli::Args { + config: a.config, + version: a.version, + check: a.check, + runner_only: a.runner_only, + from_request: a.from_request, + json: a.json, + }) + .await + } Command::Probe { compact } => { let capacity = probe::probe(&SystemHost); let json = if compact { @@ -377,7 +735,7 @@ async fn main() -> Result<()> { ) .init(); match command { - DroneCommand::Run(c) => run_drone(&c.config).await, + DroneCommand::Run(c) => run_drone(&c.config, c.token_file).await, DroneCommand::Install(a) => install_drone(a), } } @@ -396,9 +754,19 @@ async fn main() -> Result<()> { reason, name, } => set_quarantine(&config.config, &name, Some(&reason)).await, + QueenCommand::EnrollClient { config, name } => { + enroll_client(&config.config, &name).await + } + QueenCommand::Rollout { command } => match command { + RolloutCommand::Halt { config, reason } => { + rollout_admin(&config.config, Some(&reason)).await + } + RolloutCommand::Reset(config) => rollout_admin(&config.config, None).await, + }, QueenCommand::Release { config, name } => { set_quarantine(&config.config, &name, None).await } + QueenCommand::Drone { command } => queen_drone_admin(command).await, } } } diff --git a/src/profile.rs b/src/profile.rs new file mode 100644 index 0000000..0da64f0 --- /dev/null +++ b/src/profile.rs @@ -0,0 +1,423 @@ +//! Learned per-drone, per-class performance profiles. +//! +//! The queen learns which hosts are good at which work from the telemetry +//! drones report, with no host names in configuration. For each drone and +//! class it measures: +//! +//! * **speed**: for each *job family* the host has run enough times, its median +//! successful duration against the fleet median for the same class and +//! family (1000 = as fast as the fleet, 1500 = one and a half times faster), +//! averaged over families by sample count; +//! * **reliability**: one minus the host-caused failure rate (1000 = none), +//! shrunk toward 1000 while there are few samples. Code failures do not +//! count against a host. +//! +//! Profiles feed bid scoring ([`crate::bid::profile_term`]): a faster, more +//! reliable host scores higher, in proportion to how much evidence there is. +//! With too little evidence a profile is neutral and has no effect. +//! +//! A *job family* is the job name with matrix suffixes removed (everything from +//! the first `(` or `[`, and trailing digits and separators), lowercased: +//! `rust-test (shard 2)` and `rust-test (shard 7)` are one family, `rust-test`. +//! +//! GitHub, not the queen, chooses which queued job a runner takes, so profiles +//! steer which *host* gets a freshly minted runner; they cannot assign a job. + +use std::collections::{BTreeMap, BTreeSet}; + +use serde::{Deserialize, Serialize}; + +use crate::bid::ProfileScore; +use crate::classes::Class; +use crate::telemetry::{Conclusion, FailureClass, TelemetryRow}; + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ProfileSettings { + /// Learn profiles and let them influence placement. + pub enabled: bool, + /// Telemetry newer than this many days is used. + pub window_days: u32, + /// Samples a host needs per family (and in total for reliability) before + /// the evidence counts. + pub min_samples: u32, + /// Recompute at most this often. + pub refresh_secs: u64, +} + +impl Default for ProfileSettings { + fn default() -> Self { + Self { + enabled: true, + window_days: 7, + min_samples: 5, + refresh_secs: 300, + } + } +} + +impl ProfileSettings { + pub fn validate(&self) -> Result<(), String> { + if self.window_days == 0 || self.min_samples == 0 || self.refresh_secs == 0 { + return Err( + "profiles window_days, min_samples and refresh_secs must be positive".into(), + ); + } + Ok(()) + } +} + +/// Evidence for one family on one host. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FamilyStat { + pub family: String, + pub jobs: u32, + pub median_secs: u64, + pub fleet_median_secs: u64, +} + +/// One drone's profile for one class. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ClassProfile { + pub drone: String, + pub class: Class, + /// Jobs in the window (successes, host failures and other failures). + pub jobs: u32, + /// 1000 = fleet speed; clamped to 500..=2000. + pub speed_permille: u32, + /// How much evidence backs `speed_permille` (0..=1000). + pub speed_confidence_permille: u32, + /// 1000 = no host-caused failures. + pub reliability_permille: u32, + pub families: Vec, +} + +impl ClassProfile { + pub fn score(&self) -> ProfileScore { + ProfileScore { + class: self.class, + speed_permille: self.speed_permille, + speed_confidence_permille: self.speed_confidence_permille, + reliability_permille: self.reliability_permille, + } + } +} + +/// The job family of a job name; see the module docs. +pub fn family_of(job_name: &str) -> String { + let lower = job_name.to_lowercase(); + let cut = lower + .find(['(', '[']) + .map_or(lower.as_str(), |i| &lower[..i]); + let mut s = cut.trim().to_owned(); + loop { + let t = + s.trim_end_matches(|c: char| c.is_ascii_digit() || matches!(c, '-' | '_' | ' ' | '/')); + if t.len() == s.len() { + break; + } + s = t.to_owned(); + } + s +} + +fn median(v: &mut [u64]) -> u64 { + v.sort_unstable(); + let n = v.len(); + if n == 0 { + return 0; + } + if n % 2 == 1 { + v[n / 2] + } else { + (v[n / 2 - 1] + v[n / 2]) / 2 + } +} + +const SPEED_MIN: u64 = 500; +const SPEED_MAX: u64 = 2000; +/// Samples at which confidence is full. +const FULL_CONFIDENCE_AT: u32 = 20; + +/// Build every profile from `rows`. Deterministic; sorted by drone and class. +pub fn build_profiles(rows: &[TelemetryRow], s: &ProfileSettings) -> Vec { + // Successful durations per (class, family, drone). + let mut durations: BTreeMap<(Class, String, &str), Vec> = BTreeMap::new(); + // Job counts per (drone, class): total and host failures. + let mut counts: BTreeMap<(&str, Class), (u32, u32)> = BTreeMap::new(); + for r in rows { + let t = &r.telemetry; + if t.conclusion == Conclusion::Skipped { + continue; + } + let c = counts.entry((r.drone.as_str(), r.class)).or_default(); + c.0 += 1; + if t.conclusion != Conclusion::Success && t.failure_class == FailureClass::Host { + c.1 += 1; + } + if t.conclusion == Conclusion::Success { + durations + .entry((r.class, family_of(&t.job_name), r.drone.as_str())) + .or_default() + .push(r.duration_secs); + } + } + + // Fleet medians per (class, family), only where at least two hosts ran it. + let mut fleet: BTreeMap<(Class, &str), (u64, usize)> = BTreeMap::new(); + let mut families: BTreeSet<(Class, &str)> = BTreeSet::new(); + for (class, family, _) in durations.keys() { + families.insert((*class, family.as_str())); + } + for (class, family) in families { + let mut all: Vec = Vec::new(); + let mut drones = 0usize; + for ((c, f, _), v) in &durations { + if *c == class && f == family { + all.extend(v); + drones += 1; + } + } + if drones >= 2 { + fleet.insert((class, family), (median(&mut all), drones)); + } + } + + let min = usize::try_from(s.min_samples).unwrap_or(usize::MAX); + let mut out = Vec::new(); + for (&(drone, class), &(total, host)) in &counts { + let mut fams = Vec::new(); + let mut weighted = 0u64; + let mut samples = 0u32; + for ((c, family, d), v) in &durations { + if *c != class || *d != drone || v.len() < min { + continue; + } + let Some(&(fleet_med, _)) = fleet.get(&(class, family.as_str())) else { + continue; + }; + let mut mine = v.clone(); + let med = median(&mut mine); + let n = u32::try_from(v.len()).unwrap_or(u32::MAX); + let ratio = (fleet_med.saturating_mul(1000) / med.max(1)).clamp(SPEED_MIN, SPEED_MAX); + weighted += ratio * u64::from(n); + samples = samples.saturating_add(n); + fams.push(FamilyStat { + family: family.clone(), + jobs: n, + median_secs: med, + fleet_median_secs: fleet_med, + }); + } + let speed = if samples == 0 { + 1000 + } else { + u32::try_from(weighted / u64::from(samples)).unwrap_or(1000) + }; + let reliability = if total >= s.min_samples { + let failure_permille = u64::from(host) * 1000 / u64::from(total); + let weight = u64::from(total.min(FULL_CONFIDENCE_AT)); + let shrunk = failure_permille * weight / u64::from(FULL_CONFIDENCE_AT); + u32::try_from(1000u64.saturating_sub(shrunk)).unwrap_or(1000) + } else { + 1000 + }; + out.push(ClassProfile { + drone: drone.to_owned(), + class, + jobs: total, + speed_permille: speed, + speed_confidence_permille: samples.min(FULL_CONFIDENCE_AT) * 1000 / FULL_CONFIDENCE_AT, + reliability_permille: reliability, + families: fams, + }); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::classes::Scope; + use crate::telemetry::JobTelemetry; + + fn row( + drone: &str, + class: Class, + job: &str, + secs: u64, + c: Conclusion, + f: FailureClass, + ) -> TelemetryRow { + TelemetryRow { + lease_id: format!("{drone}-{job}-{secs}"), + drone: drone.into(), + class, + scope: Scope::Pr, + duration_secs: secs, + finished_at: 0, + telemetry: JobTelemetry { + job_name: job.into(), + repo: "o/r".into(), + run_id: 1, + conclusion: c, + failure_class: f, + queue_wait_secs: None, + steps: vec![], + }, + } + } + + fn ok(drone: &str, class: Class, job: &str, secs: u64) -> TelemetryRow { + row( + drone, + class, + job, + secs, + Conclusion::Success, + FailureClass::Unknown, + ) + } + + fn profile<'a>(p: &'a [ClassProfile], drone: &str, class: Class) -> &'a ClassProfile { + p.iter() + .find(|x| x.drone == drone && x.class == class) + .unwrap() + } + + #[test] + fn families_drop_matrix_suffixes() { + assert_eq!(family_of("rust-test (shard 2)"), "rust-test"); + assert_eq!(family_of("Rust-Test (shard 7)"), "rust-test"); + assert_eq!(family_of("build [linux, x64]"), "build"); + assert_eq!(family_of("test-shard-12"), "test-shard"); + assert_eq!( + family_of("rust-lint-and-release-security"), + "rust-lint-and-release-security" + ); + assert_eq!(family_of(" lint "), "lint"); + assert_eq!(family_of("42"), ""); + } + + #[test] + fn a_faster_host_scores_above_one_and_a_slower_below() { + let s = ProfileSettings::default(); + let mut rows = Vec::new(); + for i in 0..6 { + rows.push(ok("fast", Class::Medium, "rust-test (shard 1)", 100 + i)); + rows.push(ok("slow", Class::Medium, "rust-test (shard 2)", 200 + i)); + } + let p = build_profiles(&rows, &s); + let fast = profile(&p, "fast", Class::Medium); + let slow = profile(&p, "slow", Class::Medium); + assert!(fast.speed_permille > 1200, "{fast:?}"); + assert!(slow.speed_permille < 900, "{slow:?}"); + assert_eq!(fast.families[0].family, "rust-test"); + assert_eq!(fast.speed_confidence_permille, 300); // 6 of 20 samples + assert_eq!(fast.reliability_permille, 1000); + } + + #[test] + fn too_little_evidence_or_a_lone_host_is_neutral() { + let s = ProfileSettings::default(); + // Four samples each: below min_samples. + let mut rows = Vec::new(); + for _ in 0..4 { + rows.push(ok("a", Class::Light, "lint", 10)); + rows.push(ok("b", Class::Light, "lint", 100)); + } + for p in build_profiles(&rows, &s) { + assert_eq!( + (p.speed_permille, p.speed_confidence_permille), + (1000, 0), + "{p:?}" + ); + } + // One host only: nothing to compare with. + let rows: Vec<_> = (0..10).map(|_| ok("a", Class::Light, "lint", 10)).collect(); + let p = build_profiles(&rows, &s); + assert_eq!(profile(&p, "a", Class::Light).speed_confidence_permille, 0); + } + + #[test] + fn host_failures_lower_reliability_and_code_failures_do_not() { + let s = ProfileSettings::default(); + let mut rows = Vec::new(); + for _ in 0..10 { + rows.push(row( + "code", + Class::Light, + "lint", + 5, + Conclusion::Failure, + FailureClass::Code, + )); + } + for _ in 0..5 { + rows.push(ok("flaky", Class::Light, "lint", 5)); + } + for _ in 0..5 { + rows.push(row( + "flaky", + Class::Light, + "lint", + 5, + Conclusion::Failure, + FailureClass::Host, + )); + } + let p = build_profiles(&rows, &s); + assert_eq!(profile(&p, "code", Class::Light).reliability_permille, 1000); + // 50% host failures over 10 jobs: shrunk by 10/20 to 25% -> 750. + assert_eq!(profile(&p, "flaky", Class::Light).reliability_permille, 750); + } + + #[test] + fn classes_are_profiled_separately_and_skipped_jobs_are_ignored() { + let s = ProfileSettings::default(); + let mut rows = Vec::new(); + for _ in 0..6 { + rows.push(ok("a", Class::Light, "lint", 10)); + rows.push(ok("b", Class::Light, "lint", 20)); + rows.push(ok("a", Class::Heavy, "build", 500)); + rows.push(row( + "a", + Class::Light, + "lint", + 1, + Conclusion::Skipped, + FailureClass::Unknown, + )); + } + let p = build_profiles(&rows, &s); + assert_eq!(profile(&p, "a", Class::Light).jobs, 6); + assert_eq!(profile(&p, "a", Class::Heavy).speed_confidence_permille, 0); + } + + #[test] + fn extreme_durations_are_clamped_and_the_result_is_deterministic() { + let s = ProfileSettings::default(); + let mut rows = Vec::new(); + for _ in 0..6 { + rows.push(ok("a", Class::Light, "lint", 1)); + rows.push(ok("b", Class::Light, "lint", u64::MAX / 4000)); + } + let p = build_profiles(&rows, &s); + assert_eq!(profile(&p, "a", Class::Light).speed_permille, 2000); + assert!(profile(&p, "b", Class::Light).speed_permille >= 500); + assert_eq!(p, build_profiles(&rows, &s)); + } + + #[test] + fn settings_validate() { + assert!(ProfileSettings::default().validate().is_ok()); + assert!( + ProfileSettings { + min_samples: 0, + ..ProfileSettings::default() + } + .validate() + .is_err() + ); + } +} diff --git a/src/queen.rs b/src/queen.rs index 44be470..0207e55 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -11,8 +11,10 @@ //! Any GitHub failure before minting aborts the tick without minting: acting //! on a partial view of the queue would over- or under-grant. -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::Arc; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; use rand::RngCore; use rand::rngs::OsRng; @@ -20,15 +22,28 @@ use serde::Serialize; use time::OffsetDateTime; use crate::allocator::{ - ActiveLease, AllocatorInput, Demand, DroneState, LeaseState, QueueAge, Reservation, Unmet, - WarmFloor, allocate, max_slots, + ActiveLease, AllocatorInput, ClassIdle, Demand, DroneState, Grant, LeaseState, QueueAge, + Reservation, Unmet, WarmFloor, allocate, max_slots, }; +use crate::bid::{ClassPressure, ScoreWeights}; use crate::classes::{Class, Scope}; use crate::config::QueenConfig; +use crate::envcontract::blocked_classes; use crate::github::{GithubClient, GithubError, QueuedJob, Runner}; +use crate::health::{DroneHealth, DroneInfo, Outcome, OutcomeKind, evaluate}; +use crate::homeostasis::{Controller, KeyInput, KeyReport}; use crate::janitor::Janitor; use crate::probe::Capacity; +use crate::profile::{ClassProfile, build_profiles}; +use crate::rollout::{ + Check as RolloutCheck, DroneView, Observed as RolloutObserved, RolloutState, advance, + desired_for, +}; use crate::store::{LeaseRow, LeaseStatus, Store, StoreError}; +use crate::tailnet::TailnetEnroll; +use crate::telemetry::{Conclusion, FailureClass, MAX_REPORT_ROWS, TelemetryRow}; +use crate::units::{Unit, UnitState}; +use crate::wake::{UdpWakeSender, WakePlanner, WakeSender, WakeView, parse_mac}; #[derive(Debug, thiserror::Error)] pub enum QueenError { @@ -48,6 +63,82 @@ pub struct Snapshot { pub unmet: Vec, /// The warm floors in force at the last pass (base plus schedule). pub warm_floors: Vec, + /// Demand pressure per class from the last pass; published to drones. + pub pressure: Vec, + /// Per-drone host health from the last assessment. + pub health: BTreeMap, + /// Host failures look fleet-wide, so automatic quarantine is paused. + pub health_systemic: bool, + /// Slots each class could still be granted, from the last pass. + pub idle: Vec, + /// Warm-floor controller state per key (empty when off). + pub homeostasis: Vec, + /// Learned per-drone, per-class profiles (empty when off). + pub profiles: Vec, + /// The gated rollout after the last pass; `None` when not gated. + pub rollout: Option, + pub wake: WakeView, + /// Non-borg online runners by label, from the last pass. + pub static_runners: BTreeMap, + /// `unmet` minus what idle non-borg runners with the label can take. + pub unmet_placeable: Vec, +} + +/// Runners in the organization that borg did not mint, for one label. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize)] +pub struct StaticRunners { + pub online: u32, + /// Online and not running a job: able to take a queued job right now. + pub idle: u32, +} + +/// Group the organization's non-borg online runners by label. +fn static_runners(runners: &[Runner], borg_prefix: &str) -> BTreeMap { + let mut out: BTreeMap = BTreeMap::new(); + for r in runners + .iter() + .filter(|r| !r.name.starts_with(borg_prefix) && r.status == "online") + { + for l in &r.labels { + let e = out.entry(l.name.clone()).or_default(); + e.online += 1; + if !r.busy { + e.idle += 1; + } + } + } + out +} + +/// Demand borg itself has to place: each unmet count less the idle non-borg +/// runners that carry the job's label (an estimate: GitHub picks the runner). +fn placeable_unmet( + unmet: &[Unmet], + statics: &BTreeMap, + classes: &crate::classes::Classes, +) -> Vec { + unmet + .iter() + .filter_map(|u| { + let idle = statics + .get(&classes.job_label(u.class, u.scope)) + .map_or(0, |s| s.idle); + let count = u.count.saturating_sub(idle); + (count > 0).then_some(Unmet { count, ..*u }) + }) + .collect() +} + +/// What `/status` shows of a gated rollout. +#[derive(Debug, Clone, Serialize)] +pub struct RolloutView { + pub state: RolloutState, + pub checks: Vec, + pub waiting: Option, + /// Jobs on drones told the target, since they got it. + pub new_group: RolloutObserved, + /// Jobs on the other drones in the same period. + pub control_group: RolloutObserved, } #[derive(Debug, Clone, Serialize)] @@ -65,13 +156,20 @@ pub struct Queen { github: GithubClient, snapshot: RwLock, janitor: Option, + tailnet: Option, + homeostasis: std::sync::Mutex, + profile_cache: std::sync::Mutex)>>, + wake: std::sync::Mutex, + unit_joins: AtomicU64, + unit_receipts: AtomicU64, + wake_sender: Arc, tick_lock: tokio::sync::Mutex<()>, } /// The full label set a JIT runner is registered with: the labels GitHub /// assigns to every self-hosted runner, then the grant's class labels. They /// are sent explicitly because the API documents only "custom" labels. -pub fn jit_labels(os: &str, arch: &str, class_labels: &[String]) -> Vec { +pub fn jit_labels(os: &str, arch: &str, extra: &[String], class_labels: &[String]) -> Vec { let mut labels = vec!["self-hosted".to_owned()]; match os { "linux" => labels.push("Linux".to_owned()), @@ -85,10 +183,51 @@ pub fn jit_labels(os: &str, arch: &str, class_labels: &[String]) -> Vec "arm" => labels.push("ARM".to_owned()), _ => {} } + labels.extend(extra.iter().cloned()); labels.extend(class_labels.iter().cloned()); labels } +/// Reclaim audit entries are kept this long (90 days). +const RECLAIM_LOG_KEEP_SECS: i64 = 90 * 86_400; + +/// How a finished job counts for host health. Skipped jobs say nothing. +fn outcome_of(row: &TelemetryRow) -> Option { + let t = &row.telemetry; + let kind = match (t.conclusion, t.failure_class) { + (Conclusion::Skipped, _) => return None, + (Conclusion::Success, _) => OutcomeKind::Ok, + (_, FailureClass::Host) => OutcomeKind::HostFailure, + _ => OutcomeKind::Other, + }; + Some(Outcome { + drone: row.drone.clone(), + lease_id: row.lease_id.clone(), + repo: t.repo.clone(), + run_id: t.run_id, + job_name: t.job_name.clone(), + finished_at: row.finished_at, + kind, + }) +} + +/// A fresh work unit id. +pub fn new_unit_id() -> String { + let mut b = [0u8; 8]; + OsRng.fill_bytes(&mut b); + format!("u-{}", hex::encode(b)) +} + +/// `-` as units name platforms (`linux-x86_64`, `macos-aarch64`). +pub fn platform_of(cap: &Capacity) -> String { + let arch = match cap.arch.as_str() { + "arm64" => "aarch64", + "amd64" => "x86_64", + other => other, + }; + format!("{}-{arch}", cap.os) +} + fn new_lease_id() -> String { let mut b = [0u8; 6]; OsRng.fill_bytes(&mut b); @@ -107,6 +246,13 @@ impl Queen { github, snapshot: RwLock::new(Snapshot::default()), janitor: None, + tailnet: None, + homeostasis: std::sync::Mutex::new(Controller::new()), + profile_cache: std::sync::Mutex::new(None), + wake: std::sync::Mutex::new(WakePlanner::new()), + unit_joins: AtomicU64::new(0), + unit_receipts: AtomicU64::new(0), + wake_sender: Arc::new(UdpWakeSender), tick_lock: tokio::sync::Mutex::new(()), } } @@ -118,6 +264,40 @@ impl Queen { self } + /// Enable `POST /v1/enroll` with `verifier` deciding who the caller is. + #[must_use] + pub fn with_tailnet_enroll(mut self, enroll: TailnetEnroll) -> Self { + self.tailnet = Some(enroll); + self + } + + /// Replace how magic packets are sent (tests use a recording fake). + #[must_use] + pub fn with_wake_sender(mut self, sender: Arc) -> Self { + self.wake_sender = sender; + self + } + + pub fn count_unit_join(&self) { + self.unit_joins.fetch_add(1, Ordering::Relaxed); + } + + pub fn count_unit_receipt(&self) { + self.unit_receipts.fetch_add(1, Ordering::Relaxed); + } + + /// Submissions that joined a running unit / reused a receipt since start. + pub fn unit_reuse_counts(&self) -> (u64, u64) { + ( + self.unit_joins.load(Ordering::Relaxed), + self.unit_receipts.load(Ordering::Relaxed), + ) + } + + pub fn tailnet_enroll(&self) -> Option<&TailnetEnroll> { + self.tailnet.as_ref() + } + pub fn janitor(&self) -> Option<&Janitor> { self.janitor.as_ref() } @@ -164,6 +344,7 @@ impl Queen { .map(|r| (r.name.as_str(), r)) .collect(); + let statics = static_runners(&runners, &prefix); let active = self.reconcile(&ours, now_ts).await?; self.reap(&ours, &active).await; @@ -173,20 +354,51 @@ impl Queen { } let (queue, demand, queue_ages) = self.summarize(&jobs, now); + if let Err(e) = self.assess_health(now_ts).await { + tracing::warn!("health assessment failed: {e}"); + } + if let Err(e) = self.advance_rollout(now_ts).await { + tracing::warn!("rollout pass failed: {e}"); + } + let profiles = self.profiles(now_ts).await; let drone_rows = self.store.drones().await?; + let mut fresh_names: BTreeSet = BTreeSet::new(); let mut capacities: HashMap = HashMap::new(); + let mut extra_labels: HashMap> = HashMap::new(); let mut drones = Vec::new(); for d in &drone_rows { let fresh = d .last_heartbeat .is_some_and(|t| now_ts - t <= self.cfg.drone_stale_secs); + if fresh { + fresh_names.insert(d.name.clone()); + } if let (true, Some(cap)) = (fresh, &d.capacity) { - let mut state = DroneState::from_capacity(d.name.clone(), cap, &d.caps); - state.accepting = d.schedulable(); + let mut state = d.allocator_state(cap, now_ts); + state.blocked_classes = blocked_classes( + &self.cfg.env_contract, + d.env_facts.as_ref(), + &self.cfg.classes, + None, + now_ts, + ); + state.profiles = profiles + .iter() + .filter(|p| p.drone == d.name) + .map(ClassProfile::score) + .collect(); drones.push(state); capacities.insert(d.name.clone(), cap.clone()); + extra_labels.insert( + d.name.clone(), + self.cfg.classes.sanitize_extra_labels(&d.extra_labels), + ); } } + if let Err(e) = self.advance_units(now_ts, &fresh_names).await { + tracing::warn!("unit pass failed: {e}"); + } + let unit_leases = self.unit_leases().await; let leases = active .iter() .map(|l| ActiveLease { @@ -199,11 +411,15 @@ impl Queen { LeaseState::Pending }, }) - .collect(); - let warm_floors = self + .chain(unit_leases) + .collect::>(); + let schedule_floors = self .cfg .warm_schedule .effective_floors(&self.cfg.warm_floors, now); + let (warm_floors, homeostasis) = self.regulate(now_ts, &schedule_floors, &queue); + let placement_drones = drones.clone(); + let placement_leases = leases.clone(); let allocation = allocate(&AllocatorInput { classes: self.cfg.classes.clone(), demand, @@ -212,8 +428,25 @@ impl Queen { warm_floors: warm_floors.clone(), queue_ages, starvation_after_secs: self.cfg.starvation_after_secs, + weights: ScoreWeights::default(), }); + if let Err(e) = self + .place_units( + now_ts, + &placement_drones, + &capacities, + placement_leases, + &allocation.grants, + ) + .await + { + tracing::warn!("unit placement failed: {e}"); + } + // Demand that idle non-borg runners will take is not borg's to wake + // hosts for. + let unmet_placeable = placeable_unmet(&allocation.unmet, &statics, &self.cfg.classes); + let wake = self.wake_pass(now_ts, &unmet_placeable, &fresh_names).await; let mut mint_error = None; for grant in &allocation.grants { let Some(cap) = capacities.get(&grant.drone_id) else { @@ -221,7 +454,10 @@ impl Queen { }; let lease_id = new_lease_id(); let runner_name = format!("{prefix}{}-{lease_id}", grant.drone_id); - let labels = jit_labels(&cap.os, &cap.arch, &grant.labels); + let extra = extra_labels + .get(&grant.drone_id) + .map_or(&[][..], Vec::as_slice); + let labels = jit_labels(&cap.os, &cap.arch, extra, &grant.labels); let minted = match self .github .mint_jit_runner( @@ -266,7 +502,21 @@ impl Queen { s.reserved_for = allocation.reserved_for.clone(); s.unmet = allocation.unmet.clone(); s.warm_floors = warm_floors; + s.pressure = allocation.pressure.clone(); + s.idle = allocation.idle.clone(); + s.homeostasis = homeostasis; + s.profiles = profiles; + s.wake = wake; + s.static_runners = statics; + s.unmet_placeable = unmet_placeable; }); + if let Err(e) = self + .store + .prune_reclaim_log(now_ts - RECLAIM_LOG_KEEP_SECS) + .await + { + tracing::warn!("prune reclaim log failed: {e}"); + } let keep = i64::from(self.cfg.telemetry_retention_days) * 86_400; if let Err(e) = self.store.prune_telemetry(now_ts - keep).await { tracing::warn!("prune telemetry failed: {e}"); @@ -274,6 +524,441 @@ impl Queen { Ok(()) } + /// The version `drone` should run: from the gated rollout when it is on, + /// otherwise the configured target (report-only mode), otherwise no + /// instruction. + pub fn desired_version(&self, drone: &str) -> Option { + let target = self.cfg.rollout_target_version.as_ref()?; + if self.cfg.rollout.previous_version.is_none() { + return Some(target.clone()); + } + self.snapshot() + .rollout + .map(|r| desired_for(&r.state, drone).to_owned()) + } + + /// One pass of the quorum rollout gate (only when both the target and + /// `rollout.previous_version` are configured). + async fn advance_rollout(&self, now: i64) -> Result<(), QueenError> { + let (Some(target), Some(previous)) = ( + self.cfg.rollout_target_version.as_deref(), + self.cfg.rollout.previous_version.as_deref(), + ) else { + return Ok(()); + }; + let stored = self.store.rollout_state().await?; + let rows = self.store.drones().await?; + let since = stored + .as_ref() + .filter(|s| s.target == target) + .map_or(now.saturating_sub(86_400), |s| s.started_at); + let outcomes: Vec = self + .store + .telemetry_since(since, MAX_REPORT_ROWS) + .await? + .iter() + .filter_map(outcome_of) + .collect(); + let views: Vec = rows + .iter() + .map(|d| DroneView { + name: d.name.clone(), + os: d + .capacity + .as_ref() + .map(|c| c.os.clone()) + .unwrap_or_default(), + fresh: d + .last_heartbeat + .is_some_and(|t| now - t <= self.cfg.drone_stale_secs), + held: d.held(), + version: d.version.clone(), + version_since: d.version_since, + }) + .collect(); + let adv = advance( + now, + &self.cfg.rollout, + previous, + target, + stored.clone(), + &views, + &outcomes, + ); + if stored.as_ref() != Some(&adv.state) { + self.store.save_rollout_state(&adv.state, now).await?; + } + for e in &adv.events { + tracing::warn!("rollout: {e}"); + self.store.log_rollout(now, e).await?; + } + self.update_snapshot(|s| { + s.rollout = Some(RolloutView { + state: adv.state, + checks: adv.checks, + waiting: adv.waiting, + new_group: adv.new_group, + control_group: adv.control_group, + }); + }); + Ok(()) + } + + /// Units already placed hold capacity like running leases. + async fn unit_leases(&self) -> Vec { + if !self.cfg.units.enabled { + return Vec::new(); + } + match self + .store + .units_in(&[UnitState::Leased, UnitState::Running]) + .await + { + Ok(units) => units + .into_iter() + .filter_map(|u| { + Some(ActiveLease { + drone_id: u.drone?, + class: u.class, + scope: Scope::Pr, + state: LeaseState::Running, + }) + }) + .collect(), + Err(e) => { + tracing::warn!("reading units failed: {e}"); + Vec::new() + } + } + } + + /// Time out, re-queue or fail units whose drone went quiet, and prune old + /// finished ones. + async fn advance_units(&self, now: i64, fresh: &BTreeSet) -> Result<(), QueenError> { + let cfg = &self.cfg.units; + if !cfg.enabled { + return Ok(()); + } + for u in self + .store + .units_in(&[UnitState::Leased, UnitState::Running]) + .await? + { + let silent = u.drone.as_ref().is_none_or(|d| !fresh.contains(d)); + let unacked = u.state == UnitState::Leased + && u.leased_at.is_some_and(|t| now - t > cfg.lease_ttl_secs); + let overdue = u.state == UnitState::Running + && u.started_at + .is_some_and(|t| now - t > i64::try_from(u.timeout_secs).unwrap_or(i64::MAX)); + if overdue { + self.store + .fail_unit(&u.id, &format!("timed out after {}s", u.timeout_secs), now) + .await?; + } else if silent || unacked { + let why = if silent { + "drone went silent" + } else { + "drone never started it" + }; + self.store + .retry_or_fail(&u.id, u.attempt, cfg.max_attempts, why, now) + .await?; + } + } + for u in self.store.units_in(&[UnitState::Queued]).await? { + if u.awaiting_bundle && now - u.created_at > cfg.bundle_wait_secs { + self.store + .fail_unit(&u.id, "the client's bundle never arrived", now) + .await?; + } + } + self.store.prune_units(now - cfg.keep_secs).await?; + Ok(()) + } + + /// Place queued units with the runner allocator, on drones of each unit's + /// platform, in the capacity left after CI runner grants. + async fn place_units( + &self, + now: i64, + drones: &[DroneState], + capacities: &HashMap, + leases: Vec, + runner_grants: &[Grant], + ) -> Result<(), QueenError> { + if !self.cfg.units.enabled { + return Ok(()); + } + let queued: Vec = self + .store + .units_in(&[UnitState::Queued]) + .await? + .into_iter() + .filter(|u| !u.awaiting_bundle) + .collect(); + if queued.is_empty() { + return Ok(()); + } + let mut base = leases; + base.extend(runner_grants.iter().map(|g| ActiveLease { + drone_id: g.drone_id.clone(), + class: g.class, + scope: g.scope, + state: LeaseState::Pending, + })); + // Units of one platform and Rust toolchain share a candidate set: the + // drones of that platform whose build environment contract admits the + // class for that toolchain. + let mut by_platform: BTreeMap<(&str, Option<&str>), Vec<&Unit>> = BTreeMap::new(); + for u in &queued { + by_platform + .entry((u.platform.as_str(), u.rust_toolchain.as_deref())) + .or_default() + .push(u); + } + let rows = self.store.drones().await?; + for ((platform, toolchain), units) in by_platform { + let group: Vec = drones + .iter() + .filter(|d| { + capacities + .get(&d.id) + .is_some_and(|c| platform_of(c) == platform) + }) + .map(|d| { + let mut d = d.clone(); + d.blocked_classes = blocked_classes( + &self.cfg.env_contract, + rows.iter() + .find(|r| r.name == d.id) + .and_then(|r| r.env_facts.as_ref()), + &self.cfg.classes, + toolchain, + now, + ); + d + }) + .collect(); + if group.is_empty() { + continue; + } + let mut demand: BTreeMap = BTreeMap::new(); + for u in &units { + *demand.entry(u.class).or_default() += 1; + } + let allocation = allocate(&AllocatorInput { + classes: self.cfg.classes.clone(), + demand: demand + .into_iter() + .map(|(class, count)| Demand { + class, + scope: Scope::Pr, + count, + }) + .collect(), + drones: group, + leases: base.clone(), + warm_floors: Vec::new(), + queue_ages: Vec::new(), + starvation_after_secs: None, + weights: ScoreWeights::default(), + }); + // Grants of a class are interchangeable: hand them to the units + // of that class in submission order. + let mut grants: Vec<&Grant> = allocation.grants.iter().collect(); + for u in units { + if let Some(i) = grants.iter().position(|g| g.class == u.class) { + let g = grants.remove(i); + if self.store.lease_unit(&u.id, &g.drone_id, now).await? { + base.push(ActiveLease { + drone_id: g.drone_id.clone(), + class: g.class, + scope: Scope::Pr, + state: LeaseState::Running, + }); + tracing::info!("unit {} placed on {}", u.id, g.drone_id); + } + } + } + } + Ok(()) + } + + /// Send wake-on-LAN packets for demand that has stayed unmet. + async fn wake_pass(&self, now: i64, unmet: &[Unmet], fresh: &BTreeSet) -> WakeView { + let settings = &self.cfg.wake; + let classes: BTreeSet = unmet + .iter() + .filter(|u| u.count > 0) + .map(|u| u.class) + .collect(); + let awake: BTreeSet<&str> = fresh.iter().map(String::as_str).collect(); + let decisions = self + .wake + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .step(now, settings, &classes, &awake); + for (i, event) in decisions { + let host = &settings.hosts[i]; + let Some(mac) = parse_mac(&host.mac) else { + continue; + }; + match self.wake_sender.send(mac, &settings.destinations()).await { + Ok(sent) => tracing::warn!( + "wake: sent a magic packet for {} ({}) to {} destination(s)", + host.name, + event.reason, + sent.len() + ), + Err(e) => tracing::warn!("wake: sending to {} failed: {e}", host.name), + } + } + self.wake + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .view(now, settings.enabled) + } + + /// Learned drone profiles, recomputed from telemetry at most every + /// `profiles.refresh_secs`. Empty when profiles are off. + async fn profiles(&self, now: i64) -> Vec { + let settings = &self.cfg.profiles; + if !settings.enabled { + return Vec::new(); + } + let refresh = i64::try_from(settings.refresh_secs).unwrap_or(i64::MAX); + let cached = self + .profile_cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + if let Some((at, profiles)) = &cached + && (0..refresh).contains(&(now - at)) + { + return profiles.clone(); + } + let since = now.saturating_sub(i64::from(settings.window_days) * 86_400); + match self.store.telemetry_since(since, MAX_REPORT_ROWS).await { + Ok(rows) => { + let profiles = build_profiles(&rows, settings); + *self + .profile_cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some((now, profiles.clone())); + profiles + } + Err(e) => { + tracing::warn!("profile refresh failed: {e}"); + cached.map(|(_, p)| p).unwrap_or_default() + } + } + } + + /// Run the warm-floor controller: the schedule's floors plus the dynamic + /// level per key. With homeostasis off this is the schedule unchanged. + fn regulate( + &self, + now: i64, + schedule: &[WarmFloor], + queue: &[QueueEntry], + ) -> (Vec, Vec) { + let settings = &self.cfg.homeostasis; + if !settings.enabled { + return (schedule.to_vec(), Vec::new()); + } + let idle = self.snapshot().idle; + let mut inputs = Vec::new(); + for class in Class::ALL { + for scope in Scope::ALL { + let entry = queue.iter().find(|q| q.class == class && q.scope == scope); + inputs.push(KeyInput { + class, + scope, + latency_secs: entry.map_or(0, |q| q.oldest_age_secs), + queued: entry.map_or(0, |q| q.queued), + schedule_floor: schedule + .iter() + .filter(|f| f.class == class && f.scope == scope) + .map(|f| f.count) + .max() + .unwrap_or(0), + idle_slots: idle + .iter() + .find(|i| i.class == class) + .map_or(0, |i| i.idle_slots), + }); + } + } + let reports = self + .homeostasis + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .step(now, settings, &inputs); + for r in reports.iter().filter(|r| r.last_change_at == Some(now)) { + tracing::info!( + "homeostasis: {} {} warm floor now {} (latency {}s, target {}s)", + r.class, + format!("{:?}", r.scope).to_lowercase(), + r.effective_floor, + r.latency_secs, + settings.target_secs + ); + } + let floors = reports + .iter() + .filter(|r| r.effective_floor > 0) + .map(|r| WarmFloor { + class: r.class, + scope: r.scope, + count: r.effective_floor, + }) + .collect(); + (floors, reports) + } + + /// Evaluate host health from recent telemetry and apply the automatic + /// quarantine, probation and release actions it decides. + async fn assess_health(&self, now: i64) -> Result<(), QueenError> { + let settings = &self.cfg.health; + let rows = self.store.drones().await?; + let window = i64::try_from(settings.window_secs).unwrap_or(i64::MAX); + // Probation proof needs outcomes since the probation began, which can + // predate the window if no work came. + let since = rows + .iter() + .filter_map(|d| d.probation_since) + .fold(now.saturating_sub(window), i64::min); + let outcomes: Vec = self + .store + .telemetry_since(since, MAX_REPORT_ROWS) + .await? + .iter() + .filter_map(outcome_of) + .collect(); + let infos: Vec = rows + .iter() + .map(|d| DroneInfo { + name: d.name.clone(), + fresh: d + .last_heartbeat + .is_some_and(|t| now - t <= self.cfg.drone_stale_secs), + hold: d.hold(), + cleared_at: d.health_cleared_at, + }) + .collect(); + let eval = evaluate(now, settings, &infos, &outcomes); + for action in &eval.actions { + self.store.apply_health_action(action, now).await?; + tracing::warn!("health: {} {}", action.event(), action.drone()); + } + self.update_snapshot(|s| { + s.health = eval.drones; + s.health_systemic = eval.systemic; + }); + Ok(()) + } + /// Bring lease state in line with GitHub; returns the leases still active. async fn reconcile( &self, @@ -420,7 +1105,7 @@ mod tests { fn jit_labels_list_platform_labels_first_then_class_labels() { let class = vec!["borg-ci-medium".to_owned(), "borg-ci-medium-pr".to_owned()]; assert_eq!( - jit_labels("linux", "x86_64", &class), + jit_labels("linux", "x86_64", &[], &class), [ "self-hosted", "Linux", @@ -430,9 +1115,14 @@ mod tests { ] ); assert_eq!( - jit_labels("macos", "aarch64", &class[..1]), + jit_labels("macos", "aarch64", &[], &class[..1]), ["self-hosted", "macOS", "ARM64", "borg-ci-medium"] ); - assert_eq!(jit_labels("plan9", "mips", &[]), ["self-hosted"]); + assert_eq!(jit_labels("plan9", "mips", &[], &[]), ["self-hosted"]); + assert_eq!( + jit_labels("linux", "x86_64", &["on-prem".to_owned()], &class[..1]), + ["self-hosted", "Linux", "X64", "on-prem", "borg-ci-medium"], + "extra labels sit between the platform and class labels" + ); } } diff --git a/src/reclaim.rs b/src/reclaim.rs new file mode 100644 index 0000000..e29e50d --- /dev/null +++ b/src/reclaim.rs @@ -0,0 +1,934 @@ +//! Build-box reclaim: disk-pressure driven cleanup of abandoned git +//! worktrees and cargo build directories on hosts where people and agents +//! build by hand. +//! +//! This module is the *policy* and the *report*. It is pure apart from the +//! [`ReclaimHost`] trait, which the drone implements against the real +//! filesystem, git and process table; tests use a fake host. A cycle is +//! [`run_cycle`]. +//! +//! What it removes, cheapest first (tiers): +//! 1. `incremental/` directories under cargo target directories (regenerable); +//! 2. whole cargo target directories; +//! 3. git worktrees. +//! +//! When it runs: +//! * **Pressure.** Disk usage at or above `high_water_pct` starts reclaiming; +//! it frees the oldest eligible items first until usage is projected at or +//! below `low_water_pct`. An item is eligible once it has been untouched for +//! `min_age_secs` (6 hours). At or above `critical_pct` the age drops to +//! `critical_min_age_secs`. +//! * **Hygiene.** Below the high-water mark only items untouched for +//! `idle_max_age_secs` (7 days; 0 turns it off) are removed. +//! +//! What it never removes (checked immediately before each deletion, so state +//! that changed since the scan is respected; an unreadable or unknown state +//! counts as unsafe): +//! * anything a process is using: any process with its working directory or an +//! open file inside it; +//! * a worktree with uncommitted changes, **including untracked files**; +//! * a worktree with commits that are on no remote-tracking branch; +//! * a worktree that is the main checkout, locked, or whose repository has +//! stashes; +//! * a worktree that is not a git checkout. +//! +//! It is **dry-run by default** and enabled per host. Every deletion, real or +//! would-be, becomes a [`DeletionRecord`] (path, size, reason, outcome) that +//! the drone logs and reports to the queen in its heartbeat. + +use std::collections::BTreeMap; +use std::path::PathBuf; + +use serde::{Deserialize, Serialize}; + +/// What a candidate directory is. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Kind { + /// `incremental/` inside a cargo target directory. + CargoIncremental, + /// A whole cargo target directory. + CargoTarget, + /// A git worktree. + Worktree, +} + +impl Kind { + pub fn as_str(self) -> &'static str { + match self { + Self::CargoIncremental => "cargo_incremental", + Self::CargoTarget => "cargo_target", + Self::Worktree => "worktree", + } + } +} + +/// A directory the drone scans for candidates: each child is one candidate +/// (cargo target roots also offer the `incremental/` directories inside). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Root { + pub path: PathBuf, + /// `worktree` or `cargo_target` (an `incremental` root is not a thing). + pub kind: Kind, +} + +/// Per-host settings (the `reclaim` block of the drone config). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ReclaimSettings { + /// Off unless a host opts in. + pub enabled: bool, + /// Log and report what would be deleted without deleting anything. + pub dry_run: bool, + pub roots: Vec, + /// Start reclaiming at this disk usage (percent). + pub high_water_pct: u8, + /// Reclaim until usage is projected at or below this (percent). + pub low_water_pct: u8, + /// At or above this usage the minimum age drops to + /// `critical_min_age_secs`. + pub critical_pct: u8, + pub min_age_secs: u64, + pub critical_min_age_secs: u64, + /// Below high water, remove items untouched this long; 0 = never. + pub idle_max_age_secs: u64, + /// At most this many deletions per cycle. + pub max_deletions_per_cycle: u32, + /// Seconds between cycles. + pub interval_secs: u64, +} + +impl Default for ReclaimSettings { + fn default() -> Self { + Self { + enabled: false, + dry_run: true, + roots: Vec::new(), + high_water_pct: 85, + low_water_pct: 70, + critical_pct: 95, + min_age_secs: 6 * 3600, + critical_min_age_secs: 3600, + idle_max_age_secs: 7 * 24 * 3600, + max_deletions_per_cycle: 50, + interval_secs: 600, + } + } +} + +impl ReclaimSettings { + pub fn validate(&self) -> Result<(), String> { + if !(self.low_water_pct < self.high_water_pct + && self.high_water_pct <= self.critical_pct + && self.critical_pct <= 100 + && self.low_water_pct > 0) + { + return Err( + "reclaim needs 0 < low_water_pct < high_water_pct <= critical_pct <= 100".into(), + ); + } + if self.critical_min_age_secs > self.min_age_secs { + return Err("reclaim critical_min_age_secs may not exceed min_age_secs".into()); + } + if self.max_deletions_per_cycle == 0 || self.interval_secs == 0 { + return Err( + "reclaim max_deletions_per_cycle and interval_secs must be positive".into(), + ); + } + if self.enabled && self.roots.is_empty() { + return Err("reclaim is enabled but has no roots".into()); + } + for r in &self.roots { + if !r.path.is_absolute() { + return Err("reclaim root paths must be absolute".into()); + } + if r.kind == Kind::CargoIncremental { + return Err("reclaim roots are worktree or cargo_target".into()); + } + } + Ok(()) + } +} + +/// A directory that could be reclaimed, as the scan found it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Candidate { + pub path: PathBuf, + pub kind: Kind, + pub size_bytes: u64, + /// Seconds since anything in it was modified or used. + pub idle_secs: u64, + /// For an `incremental/`, the target directory that contains it. + pub within: Option, +} + +/// What the safety probe found. Anything the probe could not determine must be +/// reported as unknown, which blocks deletion. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct Safety { + /// A process has its cwd or an open file inside. + pub busy: bool, + /// Worktree: uncommitted changes including untracked files. + pub dirty: bool, + /// Worktree: commits reachable from no remote-tracking branch. + pub unpushed: bool, + /// Worktree: not a git checkout. + pub not_git: bool, + /// Worktree: main checkout, locked, or stashes exist. + pub protected: Option, + /// The probe failed or could not tell. + pub unknown: Option, +} + +impl Safety { + /// Why this candidate must be kept, if it must. + pub fn blocker(&self, kind: Kind) -> Option { + if let Some(why) = &self.unknown { + return Some(format!("unknown: {why}")); + } + if self.busy { + return Some("busy".into()); + } + if kind == Kind::Worktree { + if self.not_git { + return Some("not a git checkout".into()); + } + if self.dirty { + return Some("uncommitted changes".into()); + } + if self.unpushed { + return Some("unpushed commits".into()); + } + if let Some(why) = &self.protected { + return Some(why.clone()); + } + } + None + } +} + +/// Disk state of the filesystem holding the roots. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Disk { + pub total_bytes: u64, + pub free_bytes: u64, +} + +impl Disk { + fn used(&self) -> u64 { + self.total_bytes.saturating_sub(self.free_bytes) + } + + /// Usage in whole percent, rounded down. + pub fn usage_pct(&self) -> u8 { + if self.total_bytes == 0 { + return 0; + } + u8::try_from(u128::from(self.used()) * 100 / u128::from(self.total_bytes)).unwrap_or(100) + } +} + +/// The filesystem, git and process table, as the drone provides them. +pub trait ReclaimHost { + fn disk(&self) -> Result; + /// List candidates under the configured roots. + fn scan(&self, roots: &[Root]) -> Result, String>; + /// Probe a candidate immediately before it is deleted. + fn probe(&self, c: &Candidate) -> Safety; + /// Delete a candidate; returns the bytes freed. + fn delete(&self, c: &Candidate) -> Result; +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Outcome { + Deleted, + WouldDelete, + Failed, +} + +/// One audit entry: a deletion (or would-be deletion) and why. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeletionRecord { + /// Per-drone sequence number; the queen ignores repeats. + pub seq: u64, + pub at: i64, + pub path: String, + pub kind: Kind, + pub size_bytes: u64, + pub reason: String, + pub outcome: Outcome, + /// Set for `failed`. + #[serde(default)] + pub error: Option, +} + +/// Why a cycle ran and what it concluded. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Mode { + /// Below high water and no idle hygiene candidates due. + Idle, + Hygiene, + Pressure, + Critical, +} + +/// The result of one cycle. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Cycle { + pub mode: Mode, + pub usage_pct_before: u8, + /// Projected after the deletions (exact when not a dry run). + pub usage_pct_after: u8, + pub records: Vec, + /// Candidates kept, by reason. + pub kept: BTreeMap, + pub reclaimed_bytes: u64, + pub would_reclaim_bytes: u64, + pub errors: u32, +} + +/// Reported to the queen in every heartbeat. +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] +pub struct ReclaimStats { + pub enabled: bool, + pub dry_run: bool, + /// Bytes really deleted since the drone started. + pub reclaimed_bytes_total: u64, + pub last_run_at: Option, + pub last_run_reclaimed_bytes: u64, + /// What the last run would have deleted (dry run) . + pub last_run_would_reclaim_bytes: u64, + pub last_run_errors: u32, + /// Disk usage the last run saw. + pub usage_pct: Option, + pub last_run_mode: Option, + /// Candidates the last run kept, by reason. + pub kept: BTreeMap, + /// Audit entries created since the previous heartbeat; the queen stores + /// them and ignores any it already has (by `seq`). + #[serde(default)] + pub recent: Vec, +} + +impl ReclaimStats { + /// Fold a finished cycle into the running stats. + pub fn absorb(&mut self, settings: &ReclaimSettings, now: i64, cycle: &Cycle) { + self.enabled = settings.enabled; + self.dry_run = settings.dry_run; + self.reclaimed_bytes_total = self + .reclaimed_bytes_total + .saturating_add(cycle.reclaimed_bytes); + self.last_run_at = Some(now); + self.last_run_reclaimed_bytes = cycle.reclaimed_bytes; + self.last_run_would_reclaim_bytes = cycle.would_reclaim_bytes; + self.last_run_errors = cycle.errors; + self.usage_pct = Some(cycle.usage_pct_before); + self.last_run_mode = Some(cycle.mode); + self.kept = cycle.kept.clone(); + self.recent.extend(cycle.records.iter().cloned()); + } + + /// Take the audit entries for the next heartbeat (at most `max`; the rest + /// stay for the following one). + pub fn drain_recent(&mut self, max: usize) -> Vec { + let n = self.recent.len().min(max); + self.recent.drain(..n).collect() + } +} + +fn tier(k: Kind) -> u8 { + match k { + Kind::CargoIncremental => 0, + Kind::CargoTarget => 1, + Kind::Worktree => 2, + } +} + +fn gib(bytes: u64) -> String { + format!( + "{}.{} GiB", + bytes >> 30, + ((bytes & ((1 << 30) - 1)) * 10) >> 30 + ) +} + +fn hours(secs: u64) -> String { + if secs >= 86_400 { + format!("{}d", secs / 86_400) + } else { + format!("{}h", secs / 3600) + } +} + +/// The mode for this disk state. +pub fn mode_for(s: &ReclaimSettings, disk: &Disk) -> Mode { + let u = disk.usage_pct(); + if u >= s.critical_pct { + Mode::Critical + } else if u >= s.high_water_pct { + Mode::Pressure + } else if s.idle_max_age_secs > 0 { + Mode::Hygiene + } else { + Mode::Idle + } +} + +/// Candidates in the order they should be tried, each with its eligibility +/// reason; ineligible candidates are left out. Deterministic. +pub fn order( + s: &ReclaimSettings, + mode: Mode, + usage_pct: u8, + cands: &[Candidate], +) -> Vec<(Candidate, String)> { + let min_age = match mode { + Mode::Critical => s.critical_min_age_secs, + Mode::Pressure => s.min_age_secs, + Mode::Hygiene => s.idle_max_age_secs, + Mode::Idle => return Vec::new(), + }; + let mut out: Vec<(Candidate, String)> = cands + .iter() + .filter(|c| c.idle_secs >= min_age) + .map(|c| { + let why = match mode { + Mode::Hygiene => format!( + "idle {} (>= {} hygiene limit)", + hours(c.idle_secs), + hours(min_age) + ), + _ => format!( + "disk {usage_pct}% >= {}% high water ({mode:?}); idle {} (>= {})", + s.high_water_pct, + hours(c.idle_secs), + hours(min_age) + ), + }; + (c.clone(), why) + }) + .collect(); + out.sort_by(|(a, _), (b, _)| { + tier(a.kind) + .cmp(&tier(b.kind)) + .then(b.idle_secs.cmp(&a.idle_secs)) + .then(b.size_bytes.cmp(&a.size_bytes)) + .then(a.path.cmp(&b.path)) + }); + out +} + +/// Run one cycle. `seq` numbers the audit records (it is advanced). Returns +/// the cycle; nothing is deleted when `settings.dry_run` is true. +pub fn run_cycle( + host: &dyn ReclaimHost, + settings: &ReclaimSettings, + now: i64, + seq: &mut u64, +) -> Result { + let disk = host.disk()?; + let before = disk.usage_pct(); + let mode = mode_for(settings, &disk); + let mut cycle = Cycle { + mode, + usage_pct_before: before, + usage_pct_after: before, + records: Vec::new(), + kept: BTreeMap::new(), + reclaimed_bytes: 0, + would_reclaim_bytes: 0, + errors: 0, + }; + if mode == Mode::Idle { + return Ok(cycle); + } + let cands = host.scan(&settings.roots)?; + // Under pressure, free until usage is projected at or below low water. + // Hygiene has no byte target: it removes everything that is old enough. + let target_free = match mode { + Mode::Hygiene => u64::MAX, + _ => disk + .used() + .saturating_sub(disk.total_bytes / 100 * u64::from(settings.low_water_pct)), + }; + let mut freed: u64 = 0; + // Bytes already freed inside a target directory (its incremental/), so a + // later deletion of the whole directory is not counted twice. + let mut freed_within: BTreeMap = BTreeMap::new(); + let mut deletions = 0u32; + + for (c, eligible) in order(settings, mode, before, &cands) { + if freed >= target_free || deletions >= settings.max_deletions_per_cycle { + break; + } + // Re-probe now, not at scan time. + let safety = host.probe(&c); + if let Some(why) = safety.blocker(c.kind) { + *cycle.kept.entry(why).or_default() += 1; + continue; + } + let already = freed_within.get(&c.path).copied().unwrap_or(0); + let size = c.size_bytes.saturating_sub(already); + let reason = format!( + "{eligible}; {} {}; tier {}", + c.kind.as_str(), + gib(size), + tier(c.kind) + 1 + ); + *seq += 1; + let mut rec = DeletionRecord { + seq: *seq, + at: now, + path: c.path.display().to_string(), + kind: c.kind, + size_bytes: size, + reason, + outcome: Outcome::WouldDelete, + error: None, + }; + if settings.dry_run { + cycle.would_reclaim_bytes = cycle.would_reclaim_bytes.saturating_add(size); + } else { + match host.delete(&c) { + Ok(bytes) => { + rec.outcome = Outcome::Deleted; + // What the host says it freed, but never more than the + // estimate net of what an earlier deletion already took. + rec.size_bytes = bytes.min(size); + cycle.reclaimed_bytes = cycle.reclaimed_bytes.saturating_add(rec.size_bytes); + } + Err(e) => { + rec.outcome = Outcome::Failed; + rec.error = Some(e); + cycle.errors += 1; + cycle.records.push(rec); + continue; + } + } + } + let credited = if settings.dry_run { + size + } else { + rec.size_bytes + }; + freed = freed.saturating_add(credited); + deletions += 1; + if let Some(parent) = &c.within { + *freed_within.entry(parent.clone()).or_default() += credited; + } + cycle.records.push(rec); + } + let after_used = disk.used().saturating_sub(freed); + cycle.usage_pct_after = if disk.total_bytes == 0 { + 0 + } else { + u8::try_from(u128::from(after_used) * 100 / u128::from(disk.total_bytes)).unwrap_or(100) + }; + Ok(cycle) +} + +#[cfg(test)] +mod tests { + use std::cell::RefCell; + + use super::*; + + const GIB: u64 = 1 << 30; + const H: u64 = 3600; + + struct Fake { + disk: RefCell, + cands: Vec, + safety: BTreeMap, + fail: Vec, + deleted: RefCell>, + probes: RefCell, + } + + impl Fake { + fn new(total_gib: u64, used_gib: u64, cands: Vec) -> Self { + Self { + disk: RefCell::new(Disk { + total_bytes: total_gib * GIB, + free_bytes: (total_gib - used_gib) * GIB, + }), + cands, + safety: BTreeMap::new(), + fail: Vec::new(), + deleted: RefCell::new(Vec::new()), + probes: RefCell::new(0), + } + } + } + + impl ReclaimHost for Fake { + fn disk(&self) -> Result { + Ok(*self.disk.borrow()) + } + fn scan(&self, _: &[Root]) -> Result, String> { + Ok(self.cands.clone()) + } + fn probe(&self, c: &Candidate) -> Safety { + *self.probes.borrow_mut() += 1; + self.safety.get(&c.path).cloned().unwrap_or_default() + } + fn delete(&self, c: &Candidate) -> Result { + if self.fail.contains(&c.path) { + return Err("permission denied".into()); + } + self.deleted.borrow_mut().push(c.path.clone()); + Ok(c.size_bytes) + } + } + + fn cand(path: &str, kind: Kind, gib: u64, idle_h: u64) -> Candidate { + Candidate { + path: PathBuf::from(path), + kind, + size_bytes: gib * GIB, + idle_secs: idle_h * H, + within: None, + } + } + + fn live(settings: ReclaimSettings) -> ReclaimSettings { + ReclaimSettings { + enabled: true, + dry_run: false, + roots: vec![Root { + path: "/b".into(), + kind: Kind::Worktree, + }], + ..settings + } + } + + fn run(f: &Fake, s: &ReclaimSettings) -> Cycle { + run_cycle(f, s, 1000, &mut 0).unwrap() + } + + #[test] + fn below_high_water_only_week_old_items_go_and_nothing_else() { + let f = Fake::new( + 100, + 50, + vec![ + cand("/b/new", Kind::Worktree, 5, 30), + cand("/b/old", Kind::Worktree, 5, 24 * 8), + ], + ); + let c = run(&f, &live(ReclaimSettings::default())); + assert_eq!(c.mode, Mode::Hygiene); + assert_eq!(f.deleted.borrow().as_slice(), [PathBuf::from("/b/old")]); + // Hygiene off: nothing at all below high water. + let off = ReclaimSettings { + idle_max_age_secs: 0, + ..ReclaimSettings::default() + }; + let f2 = Fake::new(100, 50, vec![cand("/b/old", Kind::Worktree, 5, 24 * 30)]); + let c = run(&f2, &live(off)); + assert_eq!((c.mode, c.records.len()), (Mode::Idle, 0)); + } + + #[test] + fn pressure_frees_oldest_first_cheapest_tier_first_until_low_water() { + // 100 GiB disk at 90%: need to get to 70%, i.e. free 20 GiB. + let f = Fake::new( + 100, + 90, + vec![ + cand("/b/wt-old", Kind::Worktree, 8, 200), + cand("/t/a", Kind::CargoTarget, 12, 10), + cand("/t/b", Kind::CargoTarget, 12, 50), + cand("/t/b/incremental", Kind::CargoIncremental, 4, 50), + cand("/t/young", Kind::CargoTarget, 30, 2), + ], + ); + let c = run(&f, &live(ReclaimSettings::default())); + assert_eq!(c.mode, Mode::Pressure); + let order: Vec = f + .deleted + .borrow() + .iter() + .map(|p| p.display().to_string()) + .collect(); + // Incremental first, then target dirs oldest first; stops at 20 GiB; + // the worktree is never reached; the 2h-old dir is below min age. + assert_eq!(order, ["/t/b/incremental", "/t/b", "/t/a"]); + assert!(c.reclaimed_bytes >= 20 * GIB, "{}", c.reclaimed_bytes); + assert!(c.usage_pct_after <= 70, "{}", c.usage_pct_after); + assert_eq!(c.usage_pct_before, 90); + } + + #[test] + fn a_target_dir_is_not_double_counted_after_its_incremental_went() { + let mut inc = cand("/t/b/incremental", Kind::CargoIncremental, 4, 50); + inc.within = Some("/t/b".into()); + let f = Fake::new(100, 90, vec![inc, cand("/t/b", Kind::CargoTarget, 12, 50)]); + let c = run(&f, &live(ReclaimSettings::default())); + let sizes: Vec = c.records.iter().map(|r| r.size_bytes / GIB).collect(); + // 4 GiB for incremental, then 12 - 4 = 8 GiB credited for the rest. + assert_eq!(sizes, [4, 8]); + assert_eq!(c.reclaimed_bytes, 12 * GIB); + } + + #[test] + fn critical_usage_relaxes_the_minimum_age() { + let items = vec![cand("/t/recent", Kind::CargoTarget, 10, 2)]; + let f = Fake::new(100, 96, items.clone()); + let c = run(&f, &live(ReclaimSettings::default())); + assert_eq!(c.mode, Mode::Critical); + assert_eq!(c.records.len(), 1); + // At pressure (not critical) the same item is too young. + let f = Fake::new(100, 90, items); + assert!( + run(&f, &live(ReclaimSettings::default())) + .records + .is_empty() + ); + } + + #[test] + fn unsafe_worktrees_are_never_deleted_and_the_reason_is_counted() { + let paths = [ + "busy", "dirty", "unpushed", "nongit", "main", "unknown", "ok", + ]; + let cands: Vec = paths + .iter() + .map(|p| cand(&format!("/b/{p}"), Kind::Worktree, 5, 100)) + .collect(); + let mut f = Fake::new(100, 99, cands); + let set = |f: &mut Fake, p: &str, s: Safety| { + f.safety.insert(PathBuf::from(format!("/b/{p}")), s); + }; + set( + &mut f, + "busy", + Safety { + busy: true, + ..Safety::default() + }, + ); + set( + &mut f, + "dirty", + Safety { + dirty: true, + ..Safety::default() + }, + ); + set( + &mut f, + "unpushed", + Safety { + unpushed: true, + ..Safety::default() + }, + ); + set( + &mut f, + "nongit", + Safety { + not_git: true, + ..Safety::default() + }, + ); + set( + &mut f, + "main", + Safety { + protected: Some("main checkout".into()), + ..Safety::default() + }, + ); + set( + &mut f, + "unknown", + Safety { + unknown: Some("git failed".into()), + ..Safety::default() + }, + ); + let c = run(&f, &live(ReclaimSettings::default())); + assert_eq!(f.deleted.borrow().as_slice(), [PathBuf::from("/b/ok")]); + assert_eq!(c.kept["busy"], 1); + assert_eq!(c.kept["uncommitted changes"], 1); + assert_eq!(c.kept["unpushed commits"], 1); + assert_eq!(c.kept["not a git checkout"], 1); + assert_eq!(c.kept["main checkout"], 1); + assert_eq!(c.kept["unknown: git failed"], 1); + } + + #[test] + fn git_state_does_not_block_cargo_directories_but_busy_does() { + let mut f = Fake::new( + 100, + 99, + vec![ + cand("/t/a", Kind::CargoTarget, 10, 100), + cand("/t/b", Kind::CargoTarget, 10, 100), + ], + ); + f.safety.insert( + "/t/a".into(), + Safety { + dirty: true, + not_git: true, + ..Safety::default() + }, + ); + f.safety.insert( + "/t/b".into(), + Safety { + busy: true, + ..Safety::default() + }, + ); + let c = run(&f, &live(ReclaimSettings::default())); + assert_eq!(f.deleted.borrow().as_slice(), [PathBuf::from("/t/a")]); + assert_eq!(c.kept["busy"], 1); + } + + #[test] + fn dry_run_deletes_nothing_but_reports_every_would_be_deletion() { + let f = Fake::new(100, 90, vec![cand("/t/a", Kind::CargoTarget, 25, 100)]); + let s = ReclaimSettings { + dry_run: true, + ..live(ReclaimSettings::default()) + }; + let c = run(&f, &s); + assert!(f.deleted.borrow().is_empty()); + assert_eq!((c.reclaimed_bytes, c.would_reclaim_bytes), (0, 25 * GIB)); + assert_eq!(c.records[0].outcome, Outcome::WouldDelete); + assert!( + c.records[0].reason.contains("90% >= 85% high water"), + "{}", + c.records[0].reason + ); + // Defaults are dry-run and disabled. + assert!(ReclaimSettings::default().dry_run && !ReclaimSettings::default().enabled); + } + + #[test] + fn failures_are_recorded_and_the_cycle_continues_and_the_cap_applies() { + let mut f = Fake::new( + 100, + 99, + vec![ + cand("/t/a", Kind::CargoTarget, 5, 100), + cand("/t/b", Kind::CargoTarget, 5, 90), + cand("/t/c", Kind::CargoTarget, 5, 80), + ], + ); + f.fail.push("/t/a".into()); + let s = ReclaimSettings { + max_deletions_per_cycle: 2, + ..live(ReclaimSettings::default()) + }; + let c = run(&f, &s); + assert_eq!(c.errors, 1); + assert_eq!(c.records[0].outcome, Outcome::Failed); + assert_eq!(c.records[0].error.as_deref(), Some("permission denied")); + // The failure does not count toward the cap: b and c are deleted. + assert_eq!(f.deleted.borrow().len(), 2); + } + + #[test] + fn safety_is_probed_immediately_before_each_deletion_only_for_tried_items() { + let f = Fake::new( + 100, + 90, + vec![ + cand("/t/a", Kind::CargoTarget, 25, 100), + cand("/t/b", Kind::CargoTarget, 25, 90), + ], + ); + run(&f, &live(ReclaimSettings::default())); + // One deletion reached the target; the second item was never probed. + assert_eq!(*f.probes.borrow(), 1); + } + + #[test] + fn records_are_numbered_and_stats_hand_them_over_once() { + let f = Fake::new(100, 90, vec![cand("/t/a", Kind::CargoTarget, 25, 100)]); + let s = live(ReclaimSettings::default()); + let mut seq = 41; + let c = run_cycle(&f, &s, 1000, &mut seq).unwrap(); + assert_eq!(c.records[0].seq, 42); + let mut stats = ReclaimStats::default(); + stats.absorb(&s, 1000, &c); + assert_eq!( + (stats.reclaimed_bytes_total, stats.last_run_at), + (25 * GIB, Some(1000)) + ); + assert_eq!(stats.drain_recent(100).len(), 1); + assert!(stats.drain_recent(100).is_empty()); + } + + #[test] + fn ordering_is_deterministic_and_settings_validate() { + let cands = vec![ + cand("/b/y", Kind::Worktree, 1, 50), + cand("/b/x", Kind::Worktree, 1, 50), + cand("/t/z", Kind::CargoTarget, 1, 50), + ]; + let s = ReclaimSettings::default(); + let a = order(&s, Mode::Pressure, 90, &cands); + assert_eq!(a, order(&s, Mode::Pressure, 90, &cands)); + let paths: Vec<&str> = a.iter().map(|(c, _)| c.path.to_str().unwrap()).collect(); + assert_eq!(paths, ["/t/z", "/b/x", "/b/y"]); + + assert!(live(ReclaimSettings::default()).validate().is_ok()); + for bad in [ + ReclaimSettings { + low_water_pct: 90, + ..ReclaimSettings::default() + }, + ReclaimSettings { + critical_pct: 80, + ..ReclaimSettings::default() + }, + ReclaimSettings { + critical_min_age_secs: u64::MAX, + ..ReclaimSettings::default() + }, + ReclaimSettings { + enabled: true, + ..ReclaimSettings::default() + }, + ReclaimSettings { + max_deletions_per_cycle: 0, + ..ReclaimSettings::default() + }, + ReclaimSettings { + roots: vec![Root { + path: "relative".into(), + kind: Kind::Worktree, + }], + ..ReclaimSettings::default() + }, + ] { + assert!(bad.validate().is_err(), "{bad:?}"); + } + } + + #[test] + fn disk_usage_rounds_down_and_handles_empty_disks() { + let d = Disk { + total_bytes: 1000, + free_bytes: 151, + }; + assert_eq!(d.usage_pct(), 84); + assert_eq!( + Disk { + total_bytes: 0, + free_bytes: 0 + } + .usage_pct(), + 0 + ); + } +} diff --git a/src/rollout.rs b/src/rollout.rs new file mode 100644 index 0000000..fbd97f8 --- /dev/null +++ b/src/rollout.rs @@ -0,0 +1,851 @@ +//! Quorum rollout of a new drone version, queen side. +//! +//! The queen decides which drones are *told* to run the target version +//! (`desired_version` in the heartbeat response); drones update themselves. A +//! version number alone never promotes anything: each step is gated on what the +//! fleet observed. +//! +//! Stages (configured by [`RolloutSettings`]): +//! +//! 1. **Canary.** One fresh, schedulable drone per canary OS (default one Linux +//! and one macOS host, chosen by name and then kept) is told the target. The +//! stage passes when every canary reports the target and has completed +//! `canary_jobs` successful jobs on it, with all checks agreeing. +//! 2. **Waves.** Then `waves_pct` percent of the eligible fleet at a time +//! (default 25, 50, 100). A wave passes after `soak_secs` and `wave_jobs` +//! jobs on the new version, with all checks agreeing. After the last wave the +//! rollout is complete and every drone, including new ones, is told the +//! target. +//! +//! Checks, evaluated every pass against a control group (the drones still on the +//! previous version, same period): +//! * **update**: a told drone must report the target within `update_timeout_secs`; +//! * **heartbeat health**: drones on the target stay fresh and out of +//! quarantine; +//! * **host failures**: the host-failure rate on the target may not exceed the +//! control's by more than `max_host_failure_increase_pm`; with too little +//! control data an absolute cap applies; +//! * **success rate**: the job success rate may not fall below the control's by +//! more than `max_success_drop_pm`; +//! * **independence**: in the canary stage every canary host is also checked on +//! its own, so one clean host cannot vouch for a failing one. +//! +//! Any failing check **halts** the rollout and **rolls back**: every drone is +//! told the previous version. A halted rollout stays halted until an operator +//! resets it or the target changes. The module is pure: the caller supplies the +//! clock, the fleet view and the job outcomes, and persists the returned state. + +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +use crate::health::{Outcome, OutcomeKind}; + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct RolloutSettings { + /// The version drones run today. Setting it (together with the top-level + /// `rollout_target_version`) turns on the gated rollout; without it the + /// target is only reported and every drone is told the target. + pub previous_version: Option, + /// OSes that each need a canary host, as reported by the drone probe. + pub canary_oses: Vec, + /// Successful jobs each canary must complete on the target. + pub canary_jobs: u32, + /// Cumulative share of eligible drones per wave; ascending, ends at 100. + pub waves_pct: Vec, + /// Jobs the new-version drones must complete in a wave. + pub wave_jobs: u32, + /// Minimum duration of a wave. + pub soak_secs: u64, + /// How long a told drone has to report the target (and heartbeat). + pub update_timeout_secs: u64, + pub max_host_failure_increase_pm: u32, + pub max_success_drop_pm: u32, + /// Host-failure cap (permille) when the control group is too small. + pub absolute_host_failure_pm: u32, + /// Control-group jobs needed to compare against. + pub min_baseline_jobs: u32, +} + +impl Default for RolloutSettings { + fn default() -> Self { + Self { + previous_version: None, + canary_oses: vec!["linux".into(), "macos".into()], + canary_jobs: 5, + waves_pct: vec![25, 50, 100], + wave_jobs: 10, + soak_secs: 1800, + update_timeout_secs: 900, + max_host_failure_increase_pm: 20, + max_success_drop_pm: 50, + absolute_host_failure_pm: 100, + min_baseline_jobs: 5, + } + } +} + +impl RolloutSettings { + pub fn validate(&self) -> Result<(), String> { + if self.canary_oses.is_empty() || self.canary_jobs == 0 { + return Err( + "rollout canary_oses and canary_jobs must be non-empty and positive".into(), + ); + } + let ascending = self.waves_pct.windows(2).all(|w| w[0] < w[1]); + if self.waves_pct.is_empty() + || !ascending + || self.waves_pct.last() != Some(&100) + || self.waves_pct.contains(&0) + { + return Err("rollout waves_pct must be ascending percentages ending at 100".into()); + } + if self.update_timeout_secs == 0 || self.wave_jobs == 0 { + return Err("rollout update_timeout_secs and wave_jobs must be positive".into()); + } + Ok(()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Phase { + Canary, + Wave, + Complete, + Halted, +} + +impl Phase { + pub const ALL: [Self; 4] = [Self::Canary, Self::Wave, Self::Complete, Self::Halted]; + + pub fn as_str(self) -> &'static str { + match self { + Self::Canary => "canary", + Self::Wave => "wave", + Self::Complete => "complete", + Self::Halted => "halted", + } + } +} + +/// The persisted state of one rollout. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RolloutState { + pub target: String, + pub previous: String, + pub phase: Phase, + /// Canary hosts chosen so far, in `canary_oses` order. + pub canaries: Vec, + /// Drones told to run the target (grows monotonically). + pub members: Vec, + /// Index into `waves_pct` while in `Wave`. + pub wave: usize, + pub started_at: i64, + pub stage_started_at: i64, + pub halted_reason: Option, +} + +/// The fleet as the rollout sees it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DroneView { + pub name: String, + pub os: String, + pub fresh: bool, + /// Quarantined or otherwise not schedulable. + pub held: bool, + pub version: Option, + /// When the queen first saw the drone report `version`. + pub version_since: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct Check { + pub name: &'static str, + pub ok: bool, + pub detail: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize)] +pub struct Observed { + pub jobs: u32, + pub ok: u32, + pub host_failures: u32, +} + +impl Observed { + fn rate(part: u32, whole: u32) -> u32 { + if whole == 0 { + 0 + } else { + u32::try_from(u64::from(part) * 1000 / u64::from(whole)).unwrap_or(1000) + } + } + + fn host_pm(&self) -> u32 { + Self::rate(self.host_failures, self.jobs) + } + + fn success_pm(&self) -> u32 { + Self::rate(self.ok, self.jobs) + } +} + +/// The result of one pass. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Advance { + pub state: RolloutState, + /// Audit-trail lines for transitions made in this pass. + pub events: Vec, + pub checks: Vec, + /// Why the rollout is not moving yet. + pub waiting: Option, + pub new_group: Observed, + pub control_group: Observed, +} + +/// The version `name` should run under `state`. +pub fn desired_for<'a>(state: &'a RolloutState, name: &str) -> &'a str { + match state.phase { + Phase::Complete => &state.target, + Phase::Halted => &state.previous, + Phase::Canary | Phase::Wave => { + if state.members.iter().any(|m| m == name) { + &state.target + } else { + &state.previous + } + } + } +} + +fn observe<'a>(outcomes: &'a [Outcome], pick: impl Fn(&Outcome) -> bool + 'a) -> Observed { + let mut o = Observed::default(); + for x in outcomes.iter().filter(|x| pick(x)) { + match x.kind { + OutcomeKind::Ok => { + o.jobs += 1; + o.ok += 1; + } + OutcomeKind::HostFailure => { + o.jobs += 1; + o.host_failures += 1; + } + OutcomeKind::Other => o.jobs += 1, + } + } + o +} + +fn compare(s: &RolloutSettings, label: &str, new: &Observed, control: &Observed) -> Vec { + let have_control = control.jobs >= s.min_baseline_jobs; + let host_limit = if have_control { + control.host_pm() + s.max_host_failure_increase_pm + } else { + s.absolute_host_failure_pm + }; + let mut checks = vec![Check { + name: "host_failures", + ok: new.host_pm() <= host_limit, + detail: format!( + "{label}: {} host failures in {} jobs ({} permille, limit {host_limit}{})", + new.host_failures, + new.jobs, + new.host_pm(), + if have_control { + ", vs control" + } else { + ", absolute" + } + ), + }]; + if have_control && new.jobs > 0 { + let floor = control.success_pm().saturating_sub(s.max_success_drop_pm); + checks.push(Check { + name: "success_rate", + ok: new.success_pm() >= floor, + detail: format!( + "{label}: success {} permille, floor {floor} (control {})", + new.success_pm(), + control.success_pm() + ), + }); + } + checks +} + +/// Advance the rollout by one pass. `state` is the persisted state (or `None` +/// before the first pass for this target). +#[allow( + clippy::too_many_lines, + reason = "one linear state machine; splitting it would scatter the stage rules" +)] +pub fn advance( + now: i64, + s: &RolloutSettings, + previous: &str, + target: &str, + state: Option, + drones: &[DroneView], + outcomes: &[Outcome], +) -> Advance { + let mut events = Vec::new(); + let mut st = match state { + Some(st) if st.target == target => st, + _ => { + events.push(format!("rollout of {target} started (from {previous})")); + RolloutState { + target: target.to_owned(), + previous: previous.to_owned(), + phase: Phase::Canary, + canaries: Vec::new(), + members: Vec::new(), + wave: 0, + started_at: now, + stage_started_at: now, + halted_reason: None, + } + } + }; + let done = |st: RolloutState, events: Vec, waiting: Option| Advance { + state: st, + events, + checks: Vec::new(), + waiting, + new_group: Observed::default(), + control_group: Observed::default(), + }; + if matches!(st.phase, Phase::Complete | Phase::Halted) { + return done(st, events, None); + } + + // Canary selection: one fresh, schedulable host per OS, kept once chosen. + if st.phase == Phase::Canary && st.canaries.len() < s.canary_oses.len() { + let chosen: BTreeSet = st.canaries.iter().cloned().collect(); + let mut picked = st.canaries.clone(); + for os in s.canary_oses.iter().skip(st.canaries.len()) { + let pick = drones + .iter() + .filter(|d| d.fresh && !d.held && d.os == *os && !chosen.contains(&d.name)) + .map(|d| d.name.clone()) + .min(); + match pick { + Some(name) => { + events.push(format!("canary for {os}: {name}")); + picked.push(name); + } + None => { + st.canaries = picked; + st.members = st.canaries.clone(); + let why = format!("waiting for a fresh {os} host to use as canary"); + return done(st, events, Some(why)); + } + } + } + st.canaries = picked; + st.members = st.canaries.clone(); + st.stage_started_at = now; + } + + let by_name = |n: &str| drones.iter().find(|d| d.name == n); + let members: BTreeSet<&str> = st.members.iter().map(String::as_str).collect(); + let on_target = |d: &DroneView| d.version.as_deref() == Some(target); + let since_of = |n: &str| by_name(n).and_then(|d| d.version_since).unwrap_or(i64::MAX); + let timeout = i64::try_from(s.update_timeout_secs).unwrap_or(i64::MAX); + let overdue = now.saturating_sub(st.stage_started_at) > timeout; + + // Observations: members on the target, since they got it (and since the + // stage began for wave gating); control = everyone not told the target. + let stage_start = st.stage_started_at; + let new_group = observe(outcomes, |o| { + members.contains(o.drone.as_str()) && o.finished_at >= since_of(&o.drone) + }); + let control_group = observe(outcomes, |o| { + !members.contains(o.drone.as_str()) && o.finished_at >= stage_start + }); + let stage_jobs = observe(outcomes, |o| { + members.contains(o.drone.as_str()) && o.finished_at >= since_of(&o.drone).max(stage_start) + }); + + let mut checks = Vec::new(); + let mut halt: Option = None; + let mut wait: Option = None; + + for n in &st.members { + let Some(d) = by_name(n) else { + continue; + }; + if !on_target(d) { + if overdue { + halt.get_or_insert(format!( + "{n} did not report {target} within {}s of being told", + s.update_timeout_secs + )); + } else { + wait.get_or_insert(format!("waiting for {n} to update to {target}")); + } + continue; + } + if d.held { + halt.get_or_insert(format!("{n} is quarantined on {target}")); + } else if !d.fresh { + if overdue { + halt.get_or_insert(format!("{n} stopped heartbeating on {target}")); + } else { + wait.get_or_insert(format!("waiting for {n} to heartbeat on {target}")); + } + } + } + checks.push(Check { + name: "update_and_heartbeat", + ok: halt.is_none(), + detail: halt + .clone() + .or_else(|| wait.clone()) + .unwrap_or_else(|| "all told drones run the target and heartbeat".into()), + }); + + checks.extend(compare(s, "new version", &new_group, &control_group)); + if st.phase == Phase::Canary { + // Independence: each canary host must pass on its own. + for n in &st.canaries { + let own = observe(outcomes, |o| o.drone == *n && o.finished_at >= since_of(n)); + let own_checks = compare(s, n, &own, &control_group); + for c in own_checks { + if !c.ok && own.jobs >= 3 { + let others_ok = st.canaries.len() > 1; + halt.get_or_insert(format!( + "canary {n} failed {}{}: {}", + c.name, + if others_ok { + " while another canary did not" + } else { + "" + }, + c.detail + )); + } + checks.push(c); + } + } + } + if let Some(bad) = checks.iter().find(|c| !c.ok) { + halt.get_or_insert(format!("{}: {}", bad.name, bad.detail)); + } + + if let Some(reason) = halt { + st.phase = Phase::Halted; + st.halted_reason = Some(reason.clone()); + events.push(format!( + "rollout of {target} HALTED, rolled back to {}: {reason}", + st.previous + )); + return Advance { + state: st, + events, + checks, + waiting: None, + new_group, + control_group, + }; + } + + // Promotion. + let promote = match st.phase { + Phase::Canary => { + let enough = st.canaries.iter().all(|n| { + observe(outcomes, |o| o.drone == *n && o.finished_at >= since_of(n)).ok + >= s.canary_jobs + }); + let all_on = wait.is_none(); + if !enough { + wait.get_or_insert(format!( + "observing canaries: need {} successful jobs on each", + s.canary_jobs + )); + } + enough && all_on + } + Phase::Wave => { + let soaked = now.saturating_sub(st.stage_started_at) + >= i64::try_from(s.soak_secs).unwrap_or(i64::MAX); + let enough = stage_jobs.jobs >= s.wave_jobs; + if wait.is_none() && !soaked { + wait = Some(format!("soaking wave {} for {}s", st.wave + 1, s.soak_secs)); + } else if wait.is_none() && !enough { + wait = Some(format!( + "wave {}: {} of {} jobs observed", + st.wave + 1, + stage_jobs.jobs, + s.wave_jobs + )); + } + soaked && enough && wait.is_none() + } + _ => false, + }; + if promote { + let next = if st.phase == Phase::Canary { + 0 + } else { + st.wave + 1 + }; + if next >= s.waves_pct.len() { + st.phase = Phase::Complete; + events.push(format!("rollout of {target} complete")); + } else { + let eligible: Vec<&DroneView> = drones.iter().filter(|d| !d.held).collect(); + let pct = usize::try_from(s.waves_pct[next]).unwrap_or(100); + let want = (eligible.len() * pct).div_ceil(100); + let mut extra: Vec<&str> = eligible + .iter() + .map(|d| d.name.as_str()) + .filter(|n| !members.contains(n)) + .collect(); + extra.sort_unstable(); + for n in extra + .into_iter() + .take(want.saturating_sub(st.members.len())) + { + st.members.push(n.to_owned()); + } + st.phase = Phase::Wave; + st.wave = next; + st.stage_started_at = now; + events.push(format!( + "rollout of {target}: wave {} ({pct}%), {} drone(s) told", + next + 1, + st.members.len() + )); + } + wait = None; + } + Advance { + state: st, + events, + checks, + waiting: wait, + new_group, + control_group, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const NEW: &str = "2.0.0"; + const OLD: &str = "1.0.0"; + + fn settings() -> RolloutSettings { + RolloutSettings { + previous_version: Some(OLD.into()), + canary_jobs: 3, + wave_jobs: 4, + soak_secs: 100, + update_timeout_secs: 60, + min_baseline_jobs: 4, + ..RolloutSettings::default() + } + } + + fn drone(name: &str, os: &str, version: &str, since: i64) -> DroneView { + DroneView { + name: name.into(), + os: os.into(), + fresh: true, + held: false, + version: Some(version.into()), + version_since: Some(since), + } + } + + fn job(drone: &str, at: i64, kind: OutcomeKind) -> Outcome { + Outcome { + drone: drone.into(), + lease_id: format!("{drone}-{at}"), + repo: "o/r".into(), + run_id: u64::try_from(at).unwrap_or(0), + job_name: "j".into(), + finished_at: at, + kind, + } + } + + /// Fleet: linux l1..l3, macos m1, m2, all on OLD. + fn fleet() -> Vec { + vec![ + drone("l1", "linux", OLD, 0), + drone("l2", "linux", OLD, 0), + drone("l3", "linux", OLD, 0), + drone("m1", "macos", OLD, 0), + drone("m2", "macos", OLD, 0), + ] + } + + fn step(now: i64, st: Option, fleet: &[DroneView], jobs: &[Outcome]) -> Advance { + advance(now, &settings(), OLD, NEW, st, fleet, jobs) + } + + fn set_version(fleet: &mut [DroneView], names: &[&str], version: &str, since: i64) { + for d in fleet + .iter_mut() + .filter(|d| names.contains(&d.name.as_str())) + { + d.version = Some(version.into()); + d.version_since = Some(since); + } + } + + #[test] + fn the_first_pass_picks_one_canary_per_os_and_tells_only_them() { + let a = step(100, None, &fleet(), &[]); + assert_eq!(a.state.canaries, ["l1", "m1"]); + assert_eq!(a.state.phase, Phase::Canary); + assert_eq!(desired_for(&a.state, "l1"), NEW); + assert_eq!(desired_for(&a.state, "m1"), NEW); + assert_eq!(desired_for(&a.state, "l2"), OLD); + assert!(a.events[0].contains("started")); + // Waiting for the canaries to update. + assert!(a.waiting.as_deref().unwrap().contains("update")); + } + + #[test] + fn without_a_macos_host_the_canary_stage_waits_and_keeps_the_linux_pick() { + let f: Vec = fleet().into_iter().filter(|d| d.os == "linux").collect(); + let a = step(100, None, &f, &[]); + assert!(a.waiting.as_deref().unwrap().contains("macos")); + assert_eq!(a.state.canaries, ["l1"]); + // Later a macOS host appears: the stage proceeds with the same Linux pick. + let b = step(110, Some(a.state), &fleet(), &[]); + assert_eq!(b.state.canaries, ["l1", "m1"]); + } + + #[test] + fn canaries_promote_to_the_first_wave_only_after_enough_successful_jobs_on_each() { + let mut f = fleet(); + let a = step(100, None, &f, &[]); + set_version(&mut f, &["l1", "m1"], NEW, 105); + // Two successes on l1, none on m1: not enough. + let mut jobs: Vec = (0..3) + .map(|i| job("l1", 110 + i, OutcomeKind::Ok)) + .collect(); + jobs.push(job("m1", 120, OutcomeKind::Ok)); + let b = step(130, Some(a.state), &f, &jobs); + assert_eq!(b.state.phase, Phase::Canary); + assert!(b.waiting.as_deref().unwrap().contains("observing canaries")); + // Jobs from before the drone got the new version do not count. + jobs.push(job("m1", 50, OutcomeKind::Ok)); + jobs.push(job("m1", 51, OutcomeKind::Ok)); + let c = step(131, Some(b.state.clone()), &f, &jobs); + assert_eq!(c.state.phase, Phase::Canary); + jobs.push(job("m1", 121, OutcomeKind::Ok)); + jobs.push(job("m1", 122, OutcomeKind::Ok)); + let d = step(140, Some(b.state), &f, &jobs); + assert_eq!(d.state.phase, Phase::Wave); + assert_eq!(d.state.wave, 0); + // 25% of 5 eligible drones, rounded up: 2 drones, the two canaries. + assert_eq!(d.state.members, ["l1", "m1"]); + } + + fn promoted_to_wave() -> (Vec, RolloutState) { + let mut f = fleet(); + let a = step(100, None, &f, &[]); + set_version(&mut f, &["l1", "m1"], NEW, 105); + let jobs: Vec = ["l1", "m1"] + .iter() + .flat_map(|d| (0..3).map(move |i| job(d, 110 + i, OutcomeKind::Ok))) + .collect(); + let b = step(140, Some(a.state), &f, &jobs); + assert_eq!(b.state.phase, Phase::Wave); + (f, b.state) + } + + #[test] + fn waves_need_a_soak_and_jobs_then_grow_the_set_and_finally_complete() { + let (mut f, st) = promoted_to_wave(); + let jobs: Vec = (0..6) + .map(|i| job("l1", 150 + i, OutcomeKind::Ok)) + .collect(); + // Soak (100s from 140) not over yet. + let a = step(200, Some(st.clone()), &f, &jobs); + assert_eq!(a.state.wave, 0); + assert!(a.waiting.as_deref().unwrap().contains("soaking")); + // Soaked and enough jobs: wave 2 (50% of 5 = 3 drones). + let b = step(241, Some(st), &f, &jobs); + assert_eq!((b.state.phase, b.state.wave), (Phase::Wave, 1)); + assert_eq!(b.state.members, ["l1", "m1", "l2"]); + assert_eq!(desired_for(&b.state, "l2"), NEW); + assert_eq!(desired_for(&b.state, "l3"), OLD); + // Wave 3 is 100%: everyone is told. + set_version(&mut f, &["l2"], NEW, 250); + let jobs2: Vec = (0..6) + .map(|i| job("l2", 260 + i, OutcomeKind::Ok)) + .collect(); + let c = step(400, Some(b.state), &f, &jobs2); + assert_eq!((c.state.phase, c.state.wave), (Phase::Wave, 2)); + assert_eq!(c.state.members.len(), 5); + // After the last wave soaks and passes: complete, new drones included. + set_version(&mut f, &["l3", "m2"], NEW, 410); + let jobs3: Vec = (0..6) + .map(|i| job("l3", 420 + i, OutcomeKind::Ok)) + .collect(); + let d = step(520, Some(c.state), &f, &jobs3); + assert_eq!(d.state.phase, Phase::Complete); + assert_eq!(desired_for(&d.state, "brand-new"), NEW); + } + + #[test] + fn a_host_failure_spike_on_the_new_version_halts_and_rolls_everyone_back() { + let (f, st) = promoted_to_wave(); + let mut jobs: Vec = (0..4) + .map(|i| job("l1", 150 + i, OutcomeKind::Ok)) + .collect(); + jobs.extend((0..3).map(|i| job("m1", 150 + i, OutcomeKind::HostFailure))); + let a = step(160, Some(st), &f, &jobs); + assert_eq!(a.state.phase, Phase::Halted); + assert!( + a.state + .halted_reason + .as_deref() + .unwrap() + .contains("host_failures") + ); + assert_eq!(desired_for(&a.state, "l1"), OLD); + assert_eq!(desired_for(&a.state, "l3"), OLD); + assert!(a.events.last().unwrap().contains("HALTED")); + // A halted rollout stays halted, whatever happens next. + let b = step(10_000, Some(a.state.clone()), &f, &[]); + assert_eq!(b.state, a.state); + } + + #[test] + fn the_control_group_makes_the_comparison_fair() { + // The whole fleet is failing equally (a bad day, not a bad version): + // the new version is no worse than the control, so it is not halted. + let (f, st) = promoted_to_wave(); + let mut jobs = Vec::new(); + for d in ["l1", "m1", "l2", "l3", "m2"] { + jobs.push(job(d, 150, OutcomeKind::Ok)); + jobs.push(job(d, 151, OutcomeKind::Ok)); + jobs.push(job(d, 152, OutcomeKind::Other)); + jobs.push(job(d, 153, OutcomeKind::Other)); + } + let a = step(160, Some(st.clone()), &f, &jobs); + assert_eq!(a.state.phase, Phase::Wave); + // The same failures only on the new version do halt it (success drop). + let mut jobs: Vec = Vec::new(); + for d in ["l1", "m1"] { + jobs.extend((0..4).map(|i| job(d, 150 + i, OutcomeKind::Other))); + } + for d in ["l2", "l3", "m2"] { + jobs.extend((0..4).map(|i| job(d, 150 + i, OutcomeKind::Ok))); + } + let b = step(160, Some(st), &f, &jobs); + assert_eq!(b.state.phase, Phase::Halted); + assert!( + b.state + .halted_reason + .as_deref() + .unwrap() + .contains("success_rate") + ); + } + + #[test] + fn one_clean_canary_cannot_vouch_for_a_failing_one() { + let mut f = fleet(); + let a = step(100, None, &f, &[]); + set_version(&mut f, &["l1", "m1"], NEW, 105); + let mut jobs: Vec = (0..8) + .map(|i| job("l1", 110 + i, OutcomeKind::Ok)) + .collect(); + // m1: 3 host failures out of 3 jobs; aggregate 3/11 = 272 permille. + jobs.extend((0..3).map(|i| job("m1", 110 + i, OutcomeKind::HostFailure))); + let b = step(130, Some(a.state), &f, &jobs); + assert_eq!(b.state.phase, Phase::Halted); + assert!( + b.state.halted_reason.as_deref().unwrap().contains("m1"), + "{:?}", + b.state.halted_reason + ); + } + + #[test] + fn a_canary_that_never_updates_or_goes_silent_halts_after_the_timeout() { + let f = fleet(); + let a = step(100, None, &f, &[]); + // Within the timeout (60s): wait. + let b = step(150, Some(a.state.clone()), &f, &[]); + assert_eq!(b.state.phase, Phase::Canary); + assert!(b.waiting.as_deref().unwrap().contains("update")); + // Past it: halt. + let c = step(170, Some(a.state.clone()), &f, &[]); + assert_eq!(c.state.phase, Phase::Halted); + assert!( + c.state + .halted_reason + .as_deref() + .unwrap() + .contains("did not report") + ); + // A quarantined canary halts immediately. + let mut g = fleet(); + set_version(&mut g, &["l1", "m1"], NEW, 105); + g[0].held = true; + let d = step(110, Some(a.state), &g, &[]); + assert_eq!(d.state.phase, Phase::Halted); + assert!( + d.state + .halted_reason + .as_deref() + .unwrap() + .contains("quarantined") + ); + } + + #[test] + fn a_new_target_starts_a_fresh_rollout_even_after_a_halt() { + let f = fleet(); + let a = step(100, None, &f, &[]); + let halted = step(500, Some(a.state), &f, &[]); + assert_eq!(halted.state.phase, Phase::Halted); + let b = advance(600, &settings(), OLD, "3.0.0", Some(halted.state), &f, &[]); + assert_eq!( + (b.state.phase, b.state.target.as_str()), + (Phase::Canary, "3.0.0") + ); + } + + #[test] + fn the_pass_is_deterministic_and_settings_validate() { + let f = fleet(); + assert_eq!(step(100, None, &f, &[]), step(100, None, &f, &[])); + assert!(RolloutSettings::default().validate().is_ok()); + for bad in [ + RolloutSettings { + waves_pct: vec![], + ..RolloutSettings::default() + }, + RolloutSettings { + waves_pct: vec![50, 25, 100], + ..RolloutSettings::default() + }, + RolloutSettings { + waves_pct: vec![25, 50], + ..RolloutSettings::default() + }, + RolloutSettings { + canary_oses: vec![], + ..RolloutSettings::default() + }, + RolloutSettings { + canary_jobs: 0, + ..RolloutSettings::default() + }, + ] { + assert!(bad.validate().is_err()); + } + } +} diff --git a/src/run.rs b/src/run.rs new file mode 100644 index 0000000..4f4dde7 --- /dev/null +++ b/src/run.rs @@ -0,0 +1,500 @@ +//! `borg run -- `: run a command where capacity is free, or here. +//! +//! The client resolves the tree it is asked to run against, submits a work unit +//! to the queen (`POST /v1/units`, see [`crate::units`]), streams the unit's log +//! to this process's stdout and exits with the unit's exit code. It never makes +//! things worse than running the command directly: when the queen is not +//! configured, unreachable, refuses the unit, or no drone picks it up within +//! `wait_secs`, the command runs locally with this process's own stdio. +//! +//! * The tree is `git write-tree` of the index (`--index`, for pre-commit +//! hooks) or `HEAD^{tree}` (default), or an explicit SHA. +//! * An identical unit that is already running is joined (its output is +//! streamed); one that already succeeded is not run at all: the receipt is +//! reported and the exit code is 0. +//! * A unit that fails for infrastructure reasons (cancelled, timed out, out of +//! attempts, or the queen lost track of it) is run locally rather than +//! reported as the command's failure. A command that ran and exited nonzero +//! on a drone is a real failure: its code is returned as is. +//! * The drone must be able to get the tree. A tree from a pushed commit is in +//! its git mirror. Otherwise `--publish` uploads a `git bundle` of just the +//! missing objects (everything since the merge base with `origin/main`) to the +//! queen, which hands it to the drone; nothing is pushed to any remote. Without +//! `--publish`, an unpublished tree runs locally. + +use std::io::Write; +use std::path::{Path, PathBuf}; +use std::process::Command as StdCommand; +use std::time::{Duration, Instant}; + +use serde::Deserialize; + +use crate::secret::Secret; + +/// Where the tree comes from. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TreeSource { + /// `git write-tree`: the staged content (pre-commit hooks). + Index, + /// `HEAD^{tree}`. + Head, + Explicit(String), +} + +#[derive(Debug, Clone)] +pub struct RunOptions { + /// Queen base URL; `None` runs locally. + pub queen: Option, + pub token: Option, + pub dir: PathBuf, + pub tree: TreeSource, + /// `owner/name`; default derived from `origin`. + pub repo: Option, + pub platform: String, + pub class: String, + pub toolchain: String, + pub timeout_secs: Option, + pub force: bool, + /// How long to wait for a drone to pick the unit up before running locally. + pub wait_secs: u64, + pub publish: bool, + pub argv: Vec, +} + +/// `-` of this machine in the queen's notation. +pub fn local_platform() -> String { + let arch = match std::env::consts::ARCH { + "arm64" => "aarch64", + "amd64" => "x86_64", + other => other, + }; + format!("{}-{arch}", std::env::consts::OS) +} + +fn git(dir: &Path, args: &[&str]) -> Result { + #[allow( + clippy::disallowed_methods, + reason = "fixed git subcommands that resolve the tree and repository for the user's own checkout" + )] + let out = StdCommand::new("git") + .args(args) + .current_dir(dir) + .output() + .map_err(|e| format!("git: {e}"))?; + if out.status.success() { + Ok(String::from_utf8_lossy(&out.stdout).trim().to_owned()) + } else { + Err(format!( + "git {}: {}", + args.join(" "), + String::from_utf8_lossy(&out.stderr).trim() + )) + } +} + +/// The tree object id for `source`. +pub fn resolve_tree(dir: &Path, source: &TreeSource) -> Result { + match source { + TreeSource::Index => git(dir, &["write-tree"]), + TreeSource::Head => git(dir, &["rev-parse", "HEAD^{tree}"]), + TreeSource::Explicit(sha) => { + git(dir, &["rev-parse", "--verify", &format!("{sha}^{{tree}}")]) + } + } +} + +/// The channel a `rust-toolchain.toml` (or a legacy one-line `rust-toolchain`) +/// names. +pub fn parse_toolchain_channel(text: &str) -> Option { + let valid = |c: &str| { + !c.is_empty() + && c.len() <= 64 + && c.bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_')) + }; + let mut lines = text + .lines() + .map(str::trim) + .filter(|l| !l.is_empty() && !l.starts_with('#')); + let first = lines.clone().next()?; + if !first.starts_with('[') && !first.contains('=') { + return Some(first.to_owned()).filter(|c| valid(c)); + } + lines + .find_map(|l| { + let (k, v) = l.split_once('=')?; + (k.trim() == "channel").then(|| { + v.split('#') + .next() + .unwrap_or("") + .trim() + .trim_matches('"') + .to_owned() + }) + }) + .filter(|c| valid(c)) +} + +/// The Rust toolchain the tree asks for, looked for at the repository root and +/// under `rust/`. +pub fn tree_rust_toolchain(dir: &Path, tree: &str) -> Option { + [ + "rust-toolchain.toml", + "rust/rust-toolchain.toml", + "rust-toolchain", + ] + .iter() + .find_map(|p| git(dir, &["show", &format!("{tree}:{p}")]).ok()) + .and_then(|t| parse_toolchain_channel(&t)) +} + +/// `owner/name` from a remote URL (`git@host:owner/name.git`, +/// `https://host/owner/name`). +pub fn repo_from_remote(url: &str) -> Option { + let path = url.trim().trim_end_matches('/').trim_end_matches(".git"); + let tail = path.rsplit([':', '/']).take(2).collect::>(); + match tail.as_slice() { + [name, owner] if !name.is_empty() && !owner.is_empty() => Some(format!("{owner}/{name}")), + _ => None, + } +} + +/// Whether the tree exists on a remote already: some commit reachable from a +/// remote-tracking ref has exactly this tree. Conservative: `false` when unsure. +fn tree_is_published(dir: &Path, tree: &str) -> bool { + git(dir, &["log", "--remotes", "--format=%T", "-n", "2000"]) + .map(|t| t.lines().any(|l| l == tree)) + .unwrap_or(false) +} + +fn git_bytes(dir: &Path, args: &[&str]) -> Result, String> { + #[allow( + clippy::disallowed_methods, + reason = "fixed git subcommands that bundle the user's own objects" + )] + let out = StdCommand::new("git") + .args(args) + .current_dir(dir) + .output() + .map_err(|e| format!("git: {e}"))?; + if out.status.success() { + Ok(out.stdout) + } else { + Err(format!( + "git {}: {}", + args.join(" "), + String::from_utf8_lossy(&out.stderr).trim() + )) + } +} + +/// The merge base of HEAD with the remote's default branch: what a drone's +/// mirror is certain to have. +fn bundle_base(dir: &Path) -> Result { + for upstream in ["origin/HEAD", "origin/main", "origin/master"] { + if let Ok(base) = git(dir, &["merge-base", "HEAD", upstream]) { + return Ok(base); + } + } + Err("no merge base with origin/main".into()) +} + +/// A git bundle holding `tree`: a synthetic commit whose parent is the merge +/// base with the remote default branch, bundled with everything the merge base +/// already reaches excluded. The bundle's single ref is `refs/borg/unit/`. +/// Nothing is pushed anywhere. +pub fn build_bundle(dir: &Path, tree: &str) -> Result, String> { + let base = bundle_base(dir)?; + let commit = git( + dir, + &[ + "-c", + "user.name=borg", + "-c", + "user.email=borg@localhost", + "commit-tree", + tree, + "-p", + &base, + "-m", + "borg unit tree", + ], + )?; + let name = format!("refs/borg/unit/{tree}"); + git(dir, &["update-ref", &name, &commit])?; + let bundle = git_bytes(dir, &["bundle", "create", "-", &name, &format!("^{base}")]); + // Never leave the temporary ref behind, whatever the bundle did. + let _ = git(dir, &["update-ref", "-d", &name]); + bundle +} + +/// Run the command here with this process's stdio; returns its exit code. +pub fn run_local(argv: &[String], dir: &Path) -> i32 { + let Some((prog, args)) = argv.split_first() else { + return 2; + }; + #[allow( + clippy::disallowed_methods, + reason = "the fallback runs the command the user typed after --, on their own machine, with their own stdio" + )] + let status = StdCommand::new(prog).args(args).current_dir(dir).status(); + match status { + Ok(s) => s.code().unwrap_or(128 + signal_of(&s)), + Err(e) => { + eprintln!("borg run: cannot start {prog}: {e}"); + 127 + } + } +} + +#[cfg(unix)] +fn signal_of(s: &std::process::ExitStatus) -> i32 { + use std::os::unix::process::ExitStatusExt; + s.signal().unwrap_or(0) +} + +#[cfg(not(unix))] +fn signal_of(_: &std::process::ExitStatus) -> i32 { + 0 +} + +#[derive(Deserialize)] +struct Submitted { + outcome: String, + unit_id: String, + receipt: Option, +} + +struct Remote<'a> { + http: reqwest::Client, + base: &'a str, + token: &'a Secret, +} + +impl Remote<'_> { + fn url(&self, path: &str) -> String { + format!("{}{path}", self.base.trim_end_matches('/')) + } +} + +async fn upload_bundle(q: &Remote<'_>, dir: &Path, tree: &str, unit: &str) -> Result<(), String> { + let (dir, tree_id) = (dir.to_path_buf(), tree.to_owned()); + let bundle = tokio::task::spawn_blocking(move || build_bundle(&dir, &tree_id)) + .await + .map_err(|e| format!("bundling failed ({e})"))? + .map_err(|e| format!("bundling the tree failed ({e})"))?; + let resp = q + .http + .put(q.url(&format!("/v1/units/{unit}/bundle"))) + .bearer_auth(q.token.expose()) + .timeout(Duration::from_secs(120)) + .body(bundle) + .send() + .await + .map_err(|e| format!("uploading the bundle failed ({e})"))?; + if resp.status().is_success() { + Ok(()) + } else { + Err(format!("the queen refused the bundle ({})", resp.status())) + } +} + +/// Why a remote attempt did not produce the command's result. +enum Fallback { + /// Run it here, with this reason on stderr. + Local(String), +} + +/// Run `opts.argv`: remotely when possible, otherwise here. Returns the exit +/// code to exit with. Remote log bytes go to `out`. +pub async fn run(opts: &RunOptions, out: &mut dyn Write) -> i32 { + match remote(opts, out).await { + Ok(code) => code, + Err(Fallback::Local(why)) => { + eprintln!("borg run: {why}; running locally"); + run_local(&opts.argv, &opts.dir) + } + } +} + +#[allow( + clippy::too_many_lines, + reason = "one linear submit, wait, stream, finish flow" +)] +async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result { + let local = |why: &str| Fallback::Local(why.to_owned()); + let (Some(base), Some(token)) = (&opts.queen, &opts.token) else { + return Err(local("no queen configured")); + }; + let tree = resolve_tree(&opts.dir, &opts.tree).map_err(Fallback::Local)?; + let repo = match &opts.repo { + Some(r) => r.clone(), + None => git(&opts.dir, &["remote", "get-url", "origin"]) + .ok() + .and_then(|u| repo_from_remote(&u)) + .ok_or_else(|| local("cannot tell the repository (no origin remote)"))?, + }; + let needs_bundle = !tree_is_published(&opts.dir, &tree); + if needs_bundle && !opts.publish { + return Err(local("the tree is not on a remote (use --publish)")); + } + let http = reqwest::Client::builder() + .timeout(Duration::from_secs(40)) + .build() + .map_err(|e| Fallback::Local(e.to_string()))?; + let q = Remote { http, base, token }; + + let body = serde_json::json!({ + "repo": repo, "tree_sha": tree, "argv": opts.argv, "platform": opts.platform, + "toolchain": opts.toolchain, "class": opts.class, + "timeout_secs": opts.timeout_secs, "force": opts.force, "bundle": needs_bundle, + "rust_toolchain": tree_rust_toolchain(&opts.dir, &tree), + }); + let resp = q + .http + .post(q.url("/v1/units")) + .bearer_auth(q.token.expose()) + .timeout(Duration::from_secs(10)) + .json(&body) + .send() + .await + .map_err(|e| Fallback::Local(format!("queen unreachable ({e})")))?; + if !resp.status().is_success() { + return Err(Fallback::Local(format!( + "queen refused the unit ({})", + resp.status() + ))); + } + let sub: Submitted = resp + .json() + .await + .map_err(|e| Fallback::Local(format!("unreadable queen answer ({e})")))?; + if sub.outcome == "receipt" { + let drone = sub + .receipt + .as_ref() + .and_then(|r| r["drone"].as_str()) + .unwrap_or("a drone"); + eprintln!( + "borg run: tree {} already passed on {drone} (unit {}); not running again", + &tree[..tree.len().min(12)], + sub.unit_id + ); + return Ok(0); + } + let created = sub.outcome == "new"; + if created && needs_bundle { + let sent = upload_bundle(&q, &opts.dir, &tree, &sub.unit_id).await; + if let Err(why) = sent { + let _ = q + .http + .post(q.url(&format!("/v1/units/{}/cancel", sub.unit_id))) + .bearer_auth(q.token.expose()) + .send() + .await; + return Err(Fallback::Local(why)); + } + } + eprintln!( + "borg run: unit {} {} on the fleet", + sub.unit_id, + if created { "queued" } else { "joined" } + ); + + let started = Instant::now(); + let mut offset: u64 = 0; + let mut running = false; + let mut unreachable_since: Option = None; + loop { + let poll = q + .http + .get(q.url(&format!( + "/v1/units/{}/log?from={offset}&wait=10", + sub.unit_id + ))) + .bearer_auth(q.token.expose()) + .timeout(Duration::from_secs(25)); + let resp = tokio::select! { + r = poll.send() => r, + _ = tokio::signal::ctrl_c() => { + if created { + let _ = q.http.post(q.url(&format!("/v1/units/{}/cancel", sub.unit_id))) + .bearer_auth(q.token.expose()).send().await; + } + eprintln!("borg run: interrupted"); + return Ok(130); + } + }; + let resp = match resp { + Ok(r) if r.status().is_success() => { + unreachable_since = None; + r + } + other => { + let since = *unreachable_since.get_or_insert_with(Instant::now); + if since.elapsed() > Duration::from_secs(60) { + return Err(Fallback::Local(match other { + Ok(r) => format!("queen answered {} for 60s", r.status()), + Err(e) => format!("lost the queen ({e})"), + })); + } + tokio::time::sleep(Duration::from_secs(2)).await; + continue; + } + }; + let hdr = |k: &str| { + resp.headers() + .get(k) + .and_then(|v| v.to_str().ok()) + .map(str::to_owned) + }; + let state = hdr("x-borg-state").unwrap_or_default(); + let next = hdr("x-borg-next").and_then(|v| v.parse::().ok()); + let exit = hdr("x-borg-exit").and_then(|v| v.parse::().ok()); + let failure = hdr("x-borg-failure"); + let bytes = resp + .bytes() + .await + .map_err(|e| Fallback::Local(format!("log read failed ({e})")))?; + if !bytes.is_empty() { + let _ = out.write_all(&bytes); + let _ = out.flush(); + } + if let Some(n) = next { + offset = n; + } + if state == "running" || state == "done" { + running = true; + } + match state.as_str() { + "done" => { + if let Some(code) = exit { + return Ok(code); + } + } + "failed" => { + let why = failure.unwrap_or_else(|| "failed".into()); + return Err(Fallback::Local(format!( + "the unit did not run to completion ({why})" + ))); + } + "queued" | "leased" + if !running && started.elapsed() >= Duration::from_secs(opts.wait_secs) => + { + if created { + let _ = q + .http + .post(q.url(&format!("/v1/units/{}/cancel", sub.unit_id))) + .bearer_auth(q.token.expose()) + .send() + .await; + } + return Err(Fallback::Local(format!( + "no drone took it within {}s", + opts.wait_secs + ))); + } + _ => {} + } + } +} diff --git a/src/store.rs b/src/store.rs index 751eab3..63adc6c 100644 --- a/src/store.rs +++ b/src/store.rs @@ -12,10 +12,15 @@ use serde::de::DeserializeOwned; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; use sqlx::{Row, SqlitePool}; -use crate::allocator::HostCaps; +use crate::allocator::{DroneState, HostCaps}; use crate::api::CacheStat; +use crate::bid::DroneBid; use crate::classes::{Class, Scope}; +use crate::envcontract::EnvFacts; +use crate::health::{Action, Hold}; use crate::probe::Capacity; +use crate::reclaim::{DeletionRecord, ReclaimStats}; +use crate::rollout::RolloutState; use crate::secret::Secret; const SCHEMA: &str = " @@ -30,7 +35,93 @@ CREATE TABLE IF NOT EXISTS drones ( cache_stats TEXT, quarantined INTEGER NOT NULL DEFAULT 0, quarantine_reason TEXT, - version TEXT + version TEXT, + extra_labels TEXT, + bids TEXT, + quarantine_source TEXT, + quarantined_at INTEGER, + probation_since INTEGER, + quarantine_strikes INTEGER NOT NULL DEFAULT 0, + health_cleared_at INTEGER, + version_since INTEGER, + reclaim TEXT, + env_facts TEXT +); +CREATE TABLE IF NOT EXISTS reclaim_log ( + drone TEXT NOT NULL, + seq INTEGER NOT NULL, + at INTEGER NOT NULL, + path TEXT NOT NULL, + kind TEXT NOT NULL, + size_bytes INTEGER NOT NULL, + reason TEXT NOT NULL, + outcome TEXT NOT NULL, + error TEXT, + PRIMARY KEY (drone, seq) +); +CREATE INDEX IF NOT EXISTS reclaim_log_at ON reclaim_log(at); +CREATE TABLE IF NOT EXISTS clients ( + name TEXT PRIMARY KEY, + token_hash TEXT NOT NULL UNIQUE, + created_at INTEGER NOT NULL +); +CREATE TABLE IF NOT EXISTS units ( + id TEXT PRIMARY KEY, + key TEXT NOT NULL, + repo TEXT NOT NULL, + tree_sha TEXT NOT NULL, + argv TEXT NOT NULL, + platform TEXT NOT NULL, + toolchain TEXT NOT NULL, + rust_toolchain TEXT, + class TEXT NOT NULL, + timeout_secs INTEGER NOT NULL, + client TEXT NOT NULL, + state TEXT NOT NULL, + drone TEXT, + attempt INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + leased_at INTEGER, + started_at INTEGER, + finished_at INTEGER, + exit_code INTEGER, + duration_secs INTEGER, + cancel_requested INTEGER NOT NULL DEFAULT 0, + failure TEXT, + log_bytes INTEGER NOT NULL DEFAULT 0, + awaiting_bundle INTEGER NOT NULL DEFAULT 0, + bundle_bytes INTEGER NOT NULL DEFAULT 0 +); +CREATE TABLE IF NOT EXISTS unit_bundles ( + unit_id TEXT PRIMARY KEY, + data BLOB NOT NULL, + created_at INTEGER NOT NULL +); +CREATE UNIQUE INDEX IF NOT EXISTS units_active_key ON units(key) WHERE state IN ('queued', 'leased', 'running'); +CREATE INDEX IF NOT EXISTS units_state ON units(state, created_at); +CREATE TABLE IF NOT EXISTS unit_logs ( + unit_id TEXT NOT NULL, + offset INTEGER NOT NULL, + data BLOB NOT NULL, + PRIMARY KEY (unit_id, offset) +); +CREATE TABLE IF NOT EXISTS rollout_state ( + id INTEGER PRIMARY KEY CHECK (id = 1), + state TEXT NOT NULL, + updated_at INTEGER NOT NULL +); +CREATE TABLE IF NOT EXISTS rollout_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + at INTEGER NOT NULL, + event TEXT NOT NULL +); +CREATE TABLE IF NOT EXISTS health_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + at INTEGER NOT NULL, + drone TEXT NOT NULL, + event TEXT NOT NULL, + reason TEXT, + evidence TEXT ); CREATE TABLE IF NOT EXISTS leases ( id TEXT PRIMARY KEY, @@ -127,13 +218,94 @@ pub struct DroneRow { pub quarantine_reason: Option, /// The borg version the drone last reported, if it reports one. pub version: Option, + pub extra_labels: Vec, + /// The drone's last reported bids (may have expired; see `expires_at`). + pub bids: Vec, + /// `operator` or `auto` while quarantined. + pub quarantine_source: Option, + pub quarantined_at: Option, + /// Set while an auto-quarantined drone is on probation. + pub probation_since: Option, + pub quarantine_strikes: u32, + /// Evidence at or before this time was cleared by a return to service. + pub health_cleared_at: Option, + /// When the queen first saw the drone report its current `version`. + pub version_since: Option, + /// The drone's last reported build-box reclaim stats (audit entries are + /// in `reclaim_log`, not here). + pub reclaim: Option, + /// The build environment the drone last reported. + pub env_facts: Option, } impl DroneRow { + /// Quarantined and not on probation: no new leases at all. + pub fn held(&self) -> bool { + self.quarantined && self.probation_since.is_none() + } + + /// On probation after an automatic quarantine: one lease at a time. + pub fn on_probation(&self) -> bool { + self.quarantined && self.probation_since.is_some() + } + /// Whether the allocator may hand this drone new leases. pub fn schedulable(&self) -> bool { - self.accepting && !self.quarantined + self.accepting && !self.held() + } + + /// Why the drone is held, for the health evaluation. + pub fn hold(&self) -> Hold { + if !self.quarantined { + return Hold::None; + } + if self.quarantine_source.as_deref() != Some("auto") { + return Hold::Operator; + } + let strikes = self.quarantine_strikes; + match (self.probation_since, self.quarantined_at) { + (Some(since), _) => Hold::Probation { since, strikes }, + (None, since) => Hold::Auto { + since: since.unwrap_or(0), + strikes, + }, + } } + + /// The allocator's view of a fresh drone with a known capacity. + pub fn allocator_state(&self, cap: &Capacity, now: i64) -> DroneState { + let mut s = DroneState::from_capacity(self.name.clone(), cap, &self.caps); + s.accepting = self.schedulable(); + if self.on_probation() { + s.max_leases = Some(s.max_leases.map_or(1, |m| m.min(1))); + } + s.bids = self + .bids + .iter() + .filter(|b| b.expires_at > now) + .copied() + .collect(); + s + } +} + +/// One stored build-box reclaim audit entry. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct ReclaimLogRow { + pub drone: String, + #[serde(flatten)] + pub record: DeletionRecord, +} + +/// One entry of the quarantine and health audit trail. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct HealthLogRow { + pub at: i64, + pub drone: String, + pub event: String, + pub reason: Option, + /// The failures that justified an automatic action (JSON). + pub evidence: Option, } #[derive(Debug, Clone)] @@ -210,6 +382,16 @@ impl Store { "ALTER TABLE drones ADD COLUMN quarantined INTEGER NOT NULL DEFAULT 0", "ALTER TABLE drones ADD COLUMN quarantine_reason TEXT", "ALTER TABLE drones ADD COLUMN version TEXT", + "ALTER TABLE drones ADD COLUMN extra_labels TEXT", + "ALTER TABLE drones ADD COLUMN bids TEXT", + "ALTER TABLE drones ADD COLUMN quarantine_source TEXT", + "ALTER TABLE drones ADD COLUMN quarantined_at INTEGER", + "ALTER TABLE drones ADD COLUMN probation_since INTEGER", + "ALTER TABLE drones ADD COLUMN quarantine_strikes INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE drones ADD COLUMN health_cleared_at INTEGER", + "ALTER TABLE drones ADD COLUMN version_since INTEGER", + "ALTER TABLE drones ADD COLUMN reclaim TEXT", + "ALTER TABLE drones ADD COLUMN env_facts TEXT", ] { // sqlx-runtime-query: one-off additive migration, no data involved if let Err(e) = sqlx::query(ddl).execute(&pool).await @@ -272,10 +454,56 @@ impl Store { Ok(()) } + /// Labels the drone asks to have added to its runners. + pub async fn set_extra_labels(&self, name: &str, labels: &[String]) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE drones SET extra_labels = ? WHERE name = ?") + .bind(to_text(&labels)?) + .bind(name) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Delete a drone (its token hash with it) and its lease records. Returns + /// the number of unfinished leases dropped, or `None` if there is no such + /// drone. With the lease rows gone the queen no longer treats those + /// runners as owned, so it deregisters any that are still offline. + pub async fn remove_drone(&self, name: &str) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let exists = sqlx::query("SELECT 1 FROM drones WHERE name = ?") + .bind(name) + .fetch_optional(&self.pool) + .await? + .is_some(); + if !exists { + return Ok(None); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let unfinished: i64 = sqlx::query( + "SELECT COUNT(*) AS n FROM leases WHERE drone = ? AND state IN ('pending', 'running')", + ) + .bind(name) + .fetch_one(&self.pool) + .await? + .get("n"); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("DELETE FROM leases WHERE drone = ?") + .bind(name) + .execute(&self.pool) + .await?; + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("DELETE FROM drones WHERE name = ?") + .bind(name) + .execute(&self.pool) + .await?; + Ok(Some(u32::try_from(unfinished).unwrap_or(u32::MAX))) + } + pub async fn drones(&self) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let rows = sqlx::query( - "SELECT name, last_heartbeat, capacity, caps, accepting, cache_stats, quarantined, quarantine_reason, version FROM drones ORDER BY name", + "SELECT name, last_heartbeat, capacity, caps, accepting, cache_stats, quarantined, quarantine_reason, version, extra_labels, bids, quarantine_source, quarantined_at, probation_since, quarantine_strikes, health_cleared_at, version_since, reclaim, env_facts FROM drones ORDER BY name", ) .fetch_all(&self.pool) .await?; @@ -284,6 +512,7 @@ impl Store { let capacity: Option = r.get("capacity"); let caps: Option = r.get("caps"); let stats: Option = r.get("cache_stats"); + let labels: Option = r.get("extra_labels"); Ok(DroneRow { name: r.get("name"), last_heartbeat: r.get("last_heartbeat"), @@ -302,33 +531,294 @@ impl Store { quarantined: r.get::("quarantined") != 0, quarantine_reason: r.get("quarantine_reason"), version: r.get("version"), + extra_labels: labels + .as_deref() + .map(from_text) + .transpose()? + .unwrap_or_default(), + quarantine_source: r.get("quarantine_source"), + quarantined_at: r.get("quarantined_at"), + probation_since: r.get("probation_since"), + quarantine_strikes: u32::try_from(r.get::("quarantine_strikes")) + .unwrap_or(0), + health_cleared_at: r.get("health_cleared_at"), + version_since: r.get("version_since"), + env_facts: r + .get::, _>("env_facts") + .as_deref() + .map(from_text) + .transpose()?, + reclaim: r + .get::, _>("reclaim") + .as_deref() + .map(from_text) + .transpose()?, + bids: r + .get::, _>("bids") + .as_deref() + .map(from_text) + .transpose()? + .unwrap_or_default(), }) }) .collect() } - /// Quarantine (`Some(reason)`) or release (`None`) a drone. False if no - /// such drone is enrolled. + /// Quarantine (`Some(reason)`) or release (`None`) a drone on an + /// operator's say-so. A release forgets earlier health evidence. False if + /// no such drone is enrolled. pub async fn set_quarantine( &self, name: &str, reason: Option<&str>, ) -> Result { - const SQL: &str = "UPDATE drones SET quarantined = ?, quarantine_reason = ? WHERE name = ?"; + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + const HOLD: &str = "UPDATE drones SET quarantined = 1, quarantine_reason = ?, quarantine_source = 'operator', quarantined_at = ?, probation_since = NULL, quarantine_strikes = 0 WHERE name = ?"; + const FREE: &str = "UPDATE drones SET quarantined = 0, quarantine_reason = NULL, quarantine_source = NULL, quarantined_at = NULL, probation_since = NULL, quarantine_strikes = 0, health_cleared_at = ? WHERE name = ?"; + let r = match reason { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + Some(why) => sqlx::query(HOLD).bind(why).bind(now).bind(name), + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + None => sqlx::query(FREE).bind(now).bind(name), + } + .execute(&self.pool) + .await?; + if r.rows_affected() == 0 { + return Ok(false); + } + let event = if reason.is_some() { + "manual_quarantine" + } else { + "manual_release" + }; + self.log_health(now, name, event, reason, None).await?; + Ok(true) + } + + async fn log_health( + &self, + at: i64, + drone: &str, + event: &str, + reason: Option<&str>, + evidence: Option<&str>, + ) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT INTO health_log (at, drone, event, reason, evidence) VALUES (?, ?, ?, ?, ?)", + ) + .bind(at) + .bind(drone) + .bind(event) + .bind(reason) + .bind(evidence) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Apply one automatic health action and record it in the audit trail + /// (which is never pruned). + pub async fn apply_health_action(&self, action: &Action, now: i64) -> Result<(), StoreError> { + const HOLD: &str = "UPDATE drones SET quarantined = 1, quarantine_reason = ?, quarantine_source = 'auto', quarantined_at = ?, probation_since = NULL, quarantine_strikes = ? WHERE name = ?"; + const PROBATION: &str = + "UPDATE drones SET probation_since = ? WHERE name = ? AND quarantine_source = 'auto'"; + const FREE: &str = "UPDATE drones SET quarantined = 0, quarantine_reason = NULL, quarantine_source = NULL, quarantined_at = NULL, probation_since = NULL, quarantine_strikes = 0, health_cleared_at = ? WHERE name = ? AND quarantine_source = 'auto'"; + let (reason, evidence) = match action { + Action::Quarantine { + drone, + reason, + evidence, + strikes, + } + | Action::Requarantine { + drone, + reason, + evidence, + strikes, + } => { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query(HOLD) + .bind(reason) + .bind(now) + .bind(i64::from(*strikes)) + .bind(drone) + .execute(&self.pool) + .await?; + (Some(reason.as_str()), Some(to_text(evidence)?)) + } + Action::BeginProbation { drone } => { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query(PROBATION) + .bind(now) + .bind(drone) + .execute(&self.pool) + .await?; + (None, None) + } + Action::Release { drone } => { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query(FREE) + .bind(now) + .bind(drone) + .execute(&self.pool) + .await?; + (None, None) + } + }; + self.log_health( + now, + action.drone(), + action.event(), + reason, + evidence.as_deref(), + ) + .await + } + + /// The newest `limit` health audit entries, newest first. + pub async fn health_log(&self, limit: i64) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT at, drone, event, reason, evidence FROM health_log ORDER BY id DESC LIMIT ?", + ) + .bind(limit) + .fetch_all(&self.pool) + .await?; + rows.iter() + .map(|r| { + let evidence: Option = r.get("evidence"); + Ok(HealthLogRow { + at: r.get("at"), + drone: r.get("drone"), + event: r.get("event"), + reason: r.get("reason"), + evidence: evidence.as_deref().map(from_text).transpose()?, + }) + }) + .collect() + } + + /// Store a drone's reclaim stats and append its new audit entries. Entries + /// already stored (same drone and `seq`) are ignored, so a repeated + /// heartbeat cannot duplicate them. Text is bounded. + pub async fn record_reclaim(&self, name: &str, stats: &ReclaimStats) -> Result<(), StoreError> { + const MAX_TEXT: usize = 500; + const MAX_ENTRIES: usize = 500; + let clip = |s: &str| { + s.chars() + .filter(|c| !c.is_control()) + .take(MAX_TEXT) + .collect::() + }; + for r in stats.recent.iter().take(MAX_ENTRIES) { + let size = i64::try_from(r.size_bytes).unwrap_or(i64::MAX); + let seq = i64::try_from(r.seq).unwrap_or(i64::MAX); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("INSERT OR IGNORE INTO reclaim_log (drone, seq, at, path, kind, size_bytes, reason, outcome, error) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)") + .bind(name) + .bind(seq) + .bind(r.at) + .bind(clip(&r.path)) + .bind(to_text(&r.kind)?.trim_matches('"').to_owned()) + .bind(size) + .bind(clip(&r.reason)) + .bind(to_text(&r.outcome)?.trim_matches('"').to_owned()) + .bind(r.error.as_deref().map(clip)) + .execute(&self.pool) + .await?; + } + let mut kept = stats.clone(); + kept.recent.clear(); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE drones SET reclaim = ? WHERE name = ?") + .bind(to_text(&kept)?) + .bind(name) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// The newest reclaim audit entries across drones, newest first. + pub async fn reclaim_log(&self, limit: i64) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT drone, seq, at, path, kind, size_bytes, reason, outcome, error FROM reclaim_log ORDER BY at DESC, seq DESC LIMIT ?") + .bind(limit) + .fetch_all(&self.pool) + .await?; + rows.iter() + .map(|r| { + Ok(ReclaimLogRow { + drone: r.get("drone"), + record: DeletionRecord { + seq: u64::try_from(r.get::("seq")).unwrap_or(0), + at: r.get("at"), + path: r.get("path"), + kind: enum_from(&r.get::("kind"))?, + size_bytes: u64::try_from(r.get::("size_bytes")).unwrap_or(0), + reason: r.get("reason"), + outcome: enum_from(&r.get::("outcome"))?, + error: r.get("error"), + }, + }) + }) + .collect() + } + + /// Total bytes recorded as really deleted, across all drones and time. + pub async fn reclaim_bytes_logged(&self) -> Result { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database - let r = sqlx::query(SQL) - .bind(i64::from(reason.is_some())) - .bind(reason) + let row = sqlx::query( + "SELECT COALESCE(SUM(size_bytes), 0) AS n FROM reclaim_log WHERE outcome = 'deleted'", + ) + .fetch_one(&self.pool) + .await?; + Ok(u64::try_from(row.get::("n")).unwrap_or(0)) + } + + pub async fn prune_reclaim_log(&self, before: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let r = sqlx::query("DELETE FROM reclaim_log WHERE at < ?") + .bind(before) + .execute(&self.pool) + .await?; + Ok(r.rows_affected()) + } + + /// Store the build environment a drone last reported. + pub async fn set_env_facts(&self, name: &str, facts: &EnvFacts) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE drones SET env_facts = ? WHERE name = ?") + .bind(to_text(facts)?) + .bind(name) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Replace the bids a drone last reported. + pub async fn set_bids(&self, name: &str, bids: &[DroneBid]) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE drones SET bids = ? WHERE name = ?") + .bind(to_text(&bids)?) .bind(name) .execute(&self.pool) .await?; - Ok(r.rows_affected() > 0) + Ok(()) } /// Record the borg version a drone reported in its heartbeat. - pub async fn set_version(&self, name: &str, version: Option<&str>) -> Result<(), StoreError> { + pub async fn set_version( + &self, + name: &str, + version: Option<&str>, + now: i64, + ) -> Result<(), StoreError> { + // `version_since` moves only when the reported version changes. // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database - sqlx::query("UPDATE drones SET version = ? WHERE name = ?") + sqlx::query("UPDATE drones SET version_since = CASE WHEN version IS ? THEN version_since ELSE ? END, version = ? WHERE name = ?") + .bind(version) + .bind(now) .bind(version) .bind(name) .execute(&self.pool) @@ -336,6 +826,57 @@ impl Store { Ok(()) } + /// The persisted state of the current rollout, if any. + pub async fn rollout_state(&self) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT state FROM rollout_state WHERE id = 1") + .fetch_optional(&self.pool) + .await?; + row.map(|r| from_text(&r.get::("state"))) + .transpose() + } + + pub async fn save_rollout_state(&self, st: &RolloutState, now: i64) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT OR REPLACE INTO rollout_state (id, state, updated_at) VALUES (1, ?, ?)", + ) + .bind(to_text(st)?) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Forget the current rollout; the next pass starts it over. + pub async fn clear_rollout_state(&self) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("DELETE FROM rollout_state WHERE id = 1") + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn log_rollout(&self, at: i64, event: &str) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("INSERT INTO rollout_log (at, event) VALUES (?, ?)") + .bind(at) + .bind(event) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// The newest rollout audit lines, newest first. + pub async fn rollout_log(&self, limit: i64) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT at, event FROM rollout_log ORDER BY id DESC LIMIT ?") + .bind(limit) + .fetch_all(&self.pool) + .await?; + Ok(rows.iter().map(|r| (r.get("at"), r.get("event"))).collect()) + } + pub async fn insert_lease(&self, l: &LeaseRow) -> Result<(), StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database sqlx::query( @@ -590,7 +1131,12 @@ mod tests { let s = store_with("d1").await; assert!(!s.set_quarantine("nope", Some("x")).await.unwrap()); assert!(s.set_quarantine("d1", Some("bad disk")).await.unwrap()); - s.set_version("d1", Some("1.2.3")).await.unwrap(); + s.set_version("d1", Some("1.2.3"), 50).await.unwrap(); + s.set_version("d1", Some("1.2.3"), 99).await.unwrap(); + assert_eq!(s.drones().await.unwrap()[0].version_since, Some(50)); + s.set_version("d1", Some("1.2.4"), 120).await.unwrap(); + assert_eq!(s.drones().await.unwrap()[0].version_since, Some(120)); + s.set_version("d1", Some("1.2.3"), 130).await.unwrap(); let d = &s.drones().await.unwrap()[0]; assert!(d.accepting && d.quarantined && !d.schedulable()); assert_eq!(d.quarantine_reason.as_deref(), Some("bad disk")); @@ -601,6 +1147,71 @@ mod tests { assert_eq!(d.quarantine_reason, None); } + #[tokio::test] + async fn probation_caps_a_drone_at_one_lease_and_a_hold_makes_it_unschedulable() { + use crate::health::{Action, Hold}; + let s = store_with("d1").await; + let cap = crate::probe::Capacity { + hostname: "h".into(), + os: "linux".into(), + arch: "x86_64".into(), + kind: crate::probe::HostKind::Vm, + vcpu: 8, + mem_mib: 16384, + free_disk_mib: None, + docker_usable: true, + mise_version: None, + git_version: None, + }; + let row = || async { s.drones().await.unwrap().remove(0) }; + assert!(row().await.schedulable()); + assert_eq!(row().await.hold(), Hold::None); + + let hold = Action::Quarantine { + drone: "d1".into(), + reason: "auto: test".into(), + evidence: vec![], + strikes: 2, + }; + s.apply_health_action(&hold, 100).await.unwrap(); + let r = row().await; + assert!(!r.schedulable() && r.held()); + assert_eq!( + r.hold(), + Hold::Auto { + since: 100, + strikes: 2 + } + ); + + s.apply_health_action(&Action::BeginProbation { drone: "d1".into() }, 200) + .await + .unwrap(); + let r = row().await; + assert!(r.schedulable() && r.on_probation()); + assert_eq!( + r.hold(), + Hold::Probation { + since: 200, + strikes: 2 + } + ); + assert_eq!(r.allocator_state(&cap, 0).max_leases, Some(1)); + + s.apply_health_action(&Action::Release { drone: "d1".into() }, 300) + .await + .unwrap(); + let r = row().await; + assert_eq!( + (r.hold(), r.health_cleared_at, r.quarantine_strikes), + (Hold::None, Some(300), 0) + ); + assert_eq!(r.allocator_state(&cap, 0).max_leases, None); + let log = s.health_log(10).await.unwrap(); + let events: Vec<&str> = log.iter().map(|e| e.event.as_str()).collect(); + assert_eq!(events, ["release", "probation_start", "auto_quarantine"]); + } + #[tokio::test] async fn heartbeat_round_trips_capacity() { let s = store_with("d1").await; @@ -635,9 +1246,13 @@ mod tests { ) .await .unwrap(); + s.set_extra_labels("d1", &["on-prem".to_owned()]) + .await + .unwrap(); let d = &s.drones().await.unwrap()[0]; assert_eq!((d.last_heartbeat, d.accepting), (Some(123), false)); assert_eq!(d.cache_stats, [stat]); + assert_eq!(d.extra_labels, ["on-prem"]); assert_eq!(d.capacity.as_ref(), Some(&cap)); } } diff --git a/src/tailnet.rs b/src/tailnet.rs new file mode 100644 index 0000000..1e8db24 --- /dev/null +++ b/src/tailnet.rs @@ -0,0 +1,376 @@ +//! Zero-touch drone enrollment for hosts on a Tailscale tailnet. +//! +//! `POST /v1/enroll` is disabled unless `tailnet_enroll` is configured. When it +//! is, the queen asks an [`IdentityVerifier`] who owns the caller's network +//! address. The production verifier ([`TailscaledWhois`]) asks the local +//! `tailscaled` over its Unix socket (`/localapi/v0/whois`), which answers from +//! the WireGuard keys of the tailnet, so the identity cannot be forged by the +//! caller. If the node carries an allowed tag, or (for an untagged node) its +//! user is an allowed user, the queen enrolls a drone named after the node and +//! returns a fresh token, exactly like `borg queen enroll`. +//! +//! Requirements and limits: +//! * The caller's address must be the one the queen sees. Bind the queen to its +//! tailnet address and do not put a proxy in front of `/v1/enroll`: a proxy's +//! address is the proxy's identity. +//! * The queen's service user must be allowed to use the tailscaled socket +//! (`tailscale set --operator=`). +//! * A node name that is already enrolled is refused (409). There is no rotate +//! yet; delete the drone first. +//! * A tagged node is authorized by its tags only; its "user" is the tag owner +//! pseudo-user and is ignored. + +use std::collections::HashMap; +use std::future::Future; +use std::net::{IpAddr, SocketAddr}; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::Arc; + +use serde::Deserialize; + +use crate::enroll::valid_drone_name; + +pub type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// Who a tailnet address belongs to. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TailnetIdentity { + /// The node's MagicDNS name, e.g. `host-a.example.invalid.`. + pub node_name: String, + /// ACL tags on the node, e.g. `tag:ci`. + pub tags: Vec, + /// Login of the owning user; meaningful only for untagged nodes. + pub user: Option, +} + +#[derive(Debug, thiserror::Error)] +pub enum WhoisError { + #[error("tailscaled socket: {0}")] + Io(String), + #[error("tailscaled answered {0}")] + Status(u16), + #[error("tailscaled sent an unusable answer: {0}")] + Decode(String), +} + +/// Resolves a caller's address to a tailnet identity. Pluggable so tests can +/// use [`StaticIdentities`]. +pub trait IdentityVerifier: Send + Sync { + /// `Ok(None)`: the address is not a known tailnet node. + fn whois(&self, addr: SocketAddr) + -> BoxFuture<'_, Result, WhoisError>>; +} + +/// A fixed table, for tests. +#[derive(Debug, Default, Clone)] +pub struct StaticIdentities(pub HashMap); + +impl IdentityVerifier for StaticIdentities { + fn whois( + &self, + addr: SocketAddr, + ) -> BoxFuture<'_, Result, WhoisError>> { + let found = self.0.get(&addr.ip()).cloned(); + Box::pin(async move { Ok(found) }) + } +} + +/// Allowlist from the queen config. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TailnetEnrollSettings { + #[serde(default = "default_socket")] + pub socket_path: PathBuf, + /// Tags whose nodes may enroll, e.g. `tag:ci`. + #[serde(default)] + pub allowed_tags: Vec, + /// Logins of users whose untagged nodes may enroll. + #[serde(default)] + pub allowed_users: Vec, +} + +fn default_socket() -> PathBuf { + PathBuf::from("/var/run/tailscale/tailscaled.sock") +} + +impl TailnetEnrollSettings { + pub fn validate(&self) -> Result<(), String> { + if self.allowed_tags.is_empty() && self.allowed_users.is_empty() { + return Err("tailnet_enroll needs allowed_tags or allowed_users".to_owned()); + } + if self.allowed_tags.iter().any(|t| !t.starts_with("tag:")) { + return Err("tailnet_enroll allowed_tags entries look like tag:name".to_owned()); + } + Ok(()) + } + + /// Whether `id` is on the allowlist. + pub fn allows(&self, id: &TailnetIdentity) -> bool { + if !id.tags.is_empty() { + return id.tags.iter().any(|t| self.allowed_tags.contains(t)); + } + id.user + .as_ref() + .is_some_and(|u| self.allowed_users.iter().any(|a| a.eq_ignore_ascii_case(u))) + } +} + +/// The drone name for a node: its first DNS label, lowercased, with anything +/// outside the drone-name alphabet turned into `-`. `None` if nothing usable +/// is left. +pub fn drone_name_for(node_name: &str) -> Option { + let label = node_name.trim_end_matches('.').split('.').next()?; + let mut name: String = label + .chars() + .map(|c| { + let c = c.to_ascii_lowercase(); + if c.is_ascii_lowercase() || c.is_ascii_digit() { + c + } else { + '-' + } + }) + .collect(); + name = name.trim_start_matches('-').to_owned(); + name.truncate(40); + (valid_drone_name(&name)).then_some(name) +} + +/// What the queen needs to serve `/v1/enroll`. +#[derive(Clone)] +pub struct TailnetEnroll { + pub verifier: Arc, + pub settings: TailnetEnrollSettings, +} + +// ---------------------------------------------------------------- tailscaled + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +struct WhoisReply { + node: Option, + user_profile: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +struct WhoisNode { + #[serde(default)] + name: String, + #[serde(default)] + tags: Option>, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +struct WhoisUser { + #[serde(default)] + login_name: String, +} + +/// Parse the body of `GET /localapi/v0/whois`. +pub fn parse_whois(body: &[u8]) -> Result, WhoisError> { + let reply: WhoisReply = + serde_json::from_slice(body).map_err(|e| WhoisError::Decode(e.to_string()))?; + let Some(node) = reply.node else { + return Ok(None); + }; + if node.name.is_empty() { + return Ok(None); + } + Ok(Some(TailnetIdentity { + node_name: node.name, + tags: node.tags.unwrap_or_default(), + user: reply + .user_profile + .map(|u| u.login_name) + .filter(|l| !l.is_empty()), + })) +} + +/// Asks the local tailscaled over its Unix socket. +pub struct TailscaledWhois { + socket: PathBuf, +} + +impl TailscaledWhois { + pub fn new(socket: impl Into) -> Self { + Self { + socket: socket.into(), + } + } +} + +const MAX_REPLY_BYTES: u64 = 1 << 20; +const WHOIS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); + +impl IdentityVerifier for TailscaledWhois { + fn whois( + &self, + addr: SocketAddr, + ) -> BoxFuture<'_, Result, WhoisError>> { + Box::pin(async move { + tokio::time::timeout(WHOIS_TIMEOUT, self.query(addr)) + .await + .map_err(|_| WhoisError::Io("timed out".to_owned()))? + }) + } +} + +impl TailscaledWhois { + #[cfg(unix)] + async fn query(&self, addr: SocketAddr) -> Result, WhoisError> { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let io = |e: std::io::Error| WhoisError::Io(e.to_string()); + let mut stream = tokio::net::UnixStream::connect(&self.socket) + .await + .map_err(io)?; + // HTTP/1.0 so the reply is neither chunked nor kept alive. + let request = format!( + "GET /localapi/v0/whois?addr={addr} HTTP/1.0\r\nHost: local-tailscaled.sock\r\nSec-Tailscale: localapi\r\n\r\n" + ); + stream.write_all(request.as_bytes()).await.map_err(io)?; + let mut raw = Vec::new(); + stream + .take(MAX_REPLY_BYTES) + .read_to_end(&mut raw) + .await + .map_err(io)?; + let split = raw + .windows(4) + .position(|w| w == b"\r\n\r\n") + .ok_or_else(|| WhoisError::Decode("no header end".to_owned()))?; + let head = String::from_utf8_lossy(&raw[..split]); + let status: u16 = head + .lines() + .next() + .and_then(|l| l.split_whitespace().nth(1)) + .and_then(|s| s.parse().ok()) + .ok_or_else(|| WhoisError::Decode("bad status line".to_owned()))?; + match status { + 200 => parse_whois(&raw[split + 4..]), + // tailscaled answers 404 for an address that is not a tailnet node. + 404 => Ok(None), + other => Err(WhoisError::Status(other)), + } + } + + #[cfg(not(unix))] + async fn query(&self, _addr: SocketAddr) -> Result, WhoisError> { + let _ = (&self.socket, MAX_REPLY_BYTES); + Err(WhoisError::Io("unix sockets are not available".to_owned())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn settings(tags: &[&str], users: &[&str]) -> TailnetEnrollSettings { + TailnetEnrollSettings { + socket_path: default_socket(), + allowed_tags: tags.iter().map(|s| (*s).to_owned()).collect(), + allowed_users: users.iter().map(|s| (*s).to_owned()).collect(), + } + } + + fn id(tags: &[&str], user: Option<&str>) -> TailnetIdentity { + TailnetIdentity { + node_name: "host-a.example.invalid.".to_owned(), + tags: tags.iter().map(|s| (*s).to_owned()).collect(), + user: user.map(str::to_owned), + } + } + + #[test] + fn tags_authorize_tagged_nodes_and_users_only_untagged_ones() { + let s = settings(&["tag:ci"], &["dev@example.test"]); + assert!(s.allows(&id(&["tag:ci"], None))); + assert!(s.allows(&id(&["tag:other", "tag:ci"], None))); + assert!(!s.allows(&id(&["tag:other"], Some("dev@example.test")))); + assert!(s.allows(&id(&[], Some("Dev@Example.test")))); + assert!(!s.allows(&id(&[], Some("other@example.test")))); + assert!(!s.allows(&id(&[], None))); + } + + #[test] + fn settings_need_an_allowlist_of_well_formed_tags() { + assert!(settings(&[], &[]).validate().is_err()); + assert!(settings(&["ci"], &[]).validate().is_err()); + assert!(settings(&["tag:ci"], &[]).validate().is_ok()); + assert!(settings(&[], &["a@b.test"]).validate().is_ok()); + } + + #[test] + fn drone_names_come_from_the_first_dns_label() { + assert_eq!( + drone_name_for("Host_A.tail.example.invalid.").as_deref(), + Some("host-a") + ); + assert_eq!(drone_name_for("host-a").as_deref(), Some("host-a")); + assert_eq!(drone_name_for("_x.example.invalid.").as_deref(), Some("x")); + assert_eq!(drone_name_for("---.example.invalid."), None); + assert_eq!(drone_name_for(""), None); + assert_eq!( + drone_name_for(&format!("{}.example.invalid", "a".repeat(80))).map(|n| n.len()), + Some(40) + ); + } + + #[test] + fn whois_replies_parse_for_tagged_user_and_unknown_nodes() { + let tagged = br#"{"Node":{"Name":"host-a.example.invalid.","Tags":["tag:ci"]}, + "UserProfile":{"LoginName":"tagged-devices"}}"#; + let t = parse_whois(tagged).unwrap().unwrap(); + assert_eq!( + (t.node_name.as_str(), t.tags.as_slice()), + ("host-a.example.invalid.", ["tag:ci".to_owned()].as_slice()) + ); + let user = br#"{"Node":{"Name":"laptop.example.invalid.","Tags":null}, + "UserProfile":{"LoginName":"dev@example.test"}}"#; + let u = parse_whois(user).unwrap().unwrap(); + assert!(u.tags.is_empty()); + assert_eq!(u.user.as_deref(), Some("dev@example.test")); + assert_eq!(parse_whois(b"{}").unwrap(), None); + assert!(parse_whois(b"not json").is_err()); + } + + #[cfg(unix)] + #[tokio::test] + async fn the_tailscaled_client_speaks_http_over_the_unix_socket() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let dir = tempfile::tempdir().unwrap(); + let sock = dir.path().join("ts.sock"); + let listener = tokio::net::UnixListener::bind(&sock).unwrap(); + tokio::spawn(async move { + for answer in ["ok", "unknown", "broken"] { + let (mut conn, _) = listener.accept().await.unwrap(); + let mut buf = vec![0u8; 4096]; + let n = conn.read(&mut buf).await.unwrap(); + let req = String::from_utf8_lossy(&buf[..n]).into_owned(); + assert!( + req.starts_with("GET /localapi/v0/whois?addr=100.64.0.7:5555 HTTP/1.0\r\n"), + "{req}" + ); + let reply = match answer { + "ok" => "HTTP/1.0 200 OK\r\nContent-Type: application/json\r\n\r\n{\"Node\":{\"Name\":\"n.example.invalid.\",\"Tags\":[\"tag:ci\"]}}".to_owned(), + "unknown" => "HTTP/1.0 404 Not Found\r\n\r\nno match".to_owned(), + _ => "HTTP/1.0 500 Oops\r\n\r\n".to_owned(), + }; + conn.write_all(reply.as_bytes()).await.unwrap(); + } + }); + let v = TailscaledWhois::new(&sock); + let addr: SocketAddr = "100.64.0.7:5555".parse().unwrap(); + let found = v.whois(addr).await.unwrap().unwrap(); + assert_eq!(found.tags, ["tag:ci"]); + assert_eq!(v.whois(addr).await.unwrap(), None); + assert!(matches!(v.whois(addr).await, Err(WhoisError::Status(500)))); + // No daemon listening at all is an error, not "unknown node". + let gone = TailscaledWhois::new(dir.path().join("missing.sock")); + assert!(matches!(gone.whois(addr).await, Err(WhoisError::Io(_)))); + } +} diff --git a/src/units.rs b/src/units.rs new file mode 100644 index 0000000..cf734a7 --- /dev/null +++ b/src/units.rs @@ -0,0 +1,1238 @@ +//! Work units: whole commands (gate and hook runs) that clients submit to the +//! queen and drones execute, with single-flight and receipts. +//! +//! A *unit* is `(repo, tree SHA, argv, platform, toolchain digest)`. Clients +//! (`borg run`, the pre-commit hook, `mq`) submit one with `POST /v1/units`: +//! +//! * **Single-flight.** The unit key is a SHA-256 over those fields. A +//! submission whose key is already queued, leased or running *joins* that +//! unit; one whose key has a successful (exit 0) result younger than +//! `receipt_ttl_secs` returns that *receipt* without running anything. Only +//! successes are reused, and only for the exact tree SHA, so a stale pass can +//! never stand in for a different tree. `force` skips reuse. +//! * **Placement.** Queued units are placed by the same allocator that places +//! runner leases (class capacity, owner caps, bids, learned profiles, +//! quarantine and probation), on drones of the unit's platform, after CI +//! runner demand has been served. +//! * **Execution** is the drone's: it is told about a unit in its heartbeat +//! response (`units`), acknowledges it with `active_units`, appends log bytes +//! with `POST /v1/units/{id}/log?offset=N`, and ends it with +//! `POST /v1/units/{id}/result`. A unit that was never acknowledged is +//! re-queued after `lease_ttl_secs`; one on a silent drone is re-queued (up to +//! `max_attempts`); one over its timeout fails. +//! * **Logs** are stored in SQLite up to `max_log_bytes` and read by clients +//! with `GET /v1/units/{id}/log?from=N&wait=S` (long poll). +//! +//! Clients authenticate with a client token (`borg queen enroll-client`); +//! drones with their drone token. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use sqlx::Row; + +use crate::classes::Class; +use crate::store::{Store, StoreError}; + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct UnitSettings { + /// Off unless the queen opts in; the endpoints answer 404 otherwise. + pub enabled: bool, + /// A successful result is reused for this long. + pub receipt_ttl_secs: i64, + /// A placed unit the drone has not acknowledged within this is re-queued. + pub lease_ttl_secs: i64, + /// Time limit when the client gives none. + pub default_timeout_secs: u64, + pub max_timeout_secs: u64, + /// Placement attempts before a unit fails. + pub max_attempts: u32, + /// Units queued or in flight, per client. + pub max_queued_per_client: u32, + pub max_log_bytes: u64, + /// Largest git bundle a client may upload for a unit. + pub max_bundle_bytes: u64, + /// A unit announced with a bundle waits this long for it, then fails. + pub bundle_wait_secs: i64, + /// Finished units and their logs are deleted after this long. + pub keep_secs: i64, +} + +impl Default for UnitSettings { + fn default() -> Self { + Self { + enabled: false, + receipt_ttl_secs: 7 * 86_400, + lease_ttl_secs: 120, + default_timeout_secs: 3600, + max_timeout_secs: 4 * 3600, + max_attempts: 3, + max_queued_per_client: 200, + max_log_bytes: 8 << 20, + max_bundle_bytes: 64 << 20, + bundle_wait_secs: 60, + keep_secs: 7 * 86_400, + } + } +} + +impl UnitSettings { + pub fn validate(&self) -> Result<(), String> { + if self.receipt_ttl_secs <= 0 + || self.lease_ttl_secs <= 0 + || self.keep_secs <= 0 + || self.default_timeout_secs == 0 + || self.max_timeout_secs < self.default_timeout_secs + || self.max_attempts == 0 + || self.max_queued_per_client == 0 + || self.max_log_bytes == 0 + || self.max_bundle_bytes == 0 + || self.bundle_wait_secs <= 0 + { + return Err( + "units settings must be positive and max_timeout_secs >= default_timeout_secs" + .into(), + ); + } + Ok(()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum UnitState { + Queued, + /// Placed on a drone, not yet acknowledged. + Leased, + Running, + Done, + /// Ended without a result: timed out, cancelled or out of attempts. + Failed, +} + +impl UnitState { + pub const ALL: [Self; 5] = [ + Self::Queued, + Self::Leased, + Self::Running, + Self::Done, + Self::Failed, + ]; + + pub fn as_str(self) -> &'static str { + match self { + Self::Queued => "queued", + Self::Leased => "leased", + Self::Running => "running", + Self::Done => "done", + Self::Failed => "failed", + } + } + + fn parse(s: &str) -> Result { + Self::ALL + .into_iter() + .find(|v| v.as_str() == s) + .ok_or_else(|| StoreError::Corrupt(format!("unit state {s}"))) + } + + pub fn is_active(self) -> bool { + matches!(self, Self::Queued | Self::Leased | Self::Running) + } +} + +/// What a client submits. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SubmitUnit { + pub repo: String, + pub tree_sha: String, + pub argv: Vec, + /// `-`, e.g. `linux-x86_64`. + pub platform: String, + /// Digest of the toolchain the result depends on; part of the key. + #[serde(default)] + pub toolchain: String, + #[serde(default = "default_class")] + pub class: Class, + #[serde(default)] + pub timeout_secs: Option, + /// Run even if a receipt exists. + #[serde(default)] + pub force: bool, + /// The Rust toolchain the command needs (from `rust-toolchain.toml` at the + /// tree); the unit is only placed on a drone that has it. + #[serde(default)] + pub rust_toolchain: Option, + /// The client will upload a git bundle holding the tree's objects + /// (`PUT /v1/units/{id}/bundle`); the unit is not placed until it arrives. + #[serde(default)] + pub bundle: bool, +} + +fn default_class() -> Class { + Class::Medium +} + +impl SubmitUnit { + pub fn validate(&self) -> Result<(), String> { + let hex_ok = + |s: &str| matches!(s.len(), 40 | 64) && s.bytes().all(|b| b.is_ascii_hexdigit()); + if self.repo.split('/').filter(|p| !p.is_empty()).count() != 2 + || self.repo.len() > 200 + || self.repo.contains(char::is_whitespace) + { + return Err("repo must look like owner/name".into()); + } + if !hex_ok(&self.tree_sha) { + return Err("tree_sha must be a 40 or 64 digit hex object id".into()); + } + if self.argv.is_empty() + || self.argv.len() > 64 + || self.argv.iter().any(|a| a.len() > 4096 || a.contains('\0')) + { + return Err("argv must have 1 to 64 arguments of at most 4096 bytes".into()); + } + let plat_ok = !self.platform.is_empty() + && self.platform.len() <= 40 + && self + .platform + .bytes() + .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-' || b == b'_'); + if !plat_ok { + return Err("platform must look like linux-x86_64".into()); + } + let tc_ok = |t: &str| { + !t.is_empty() + && t.len() <= 64 + && t.bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'_')) + }; + if self.rust_toolchain.as_deref().is_some_and(|t| !tc_ok(t)) { + return Err("rust_toolchain must look like 1.95.0 or nightly-2026-01-01".into()); + } + if self.toolchain.len() > 200 || self.toolchain.chars().any(char::is_control) { + return Err("toolchain must be at most 200 printable characters".into()); + } + Ok(()) + } + + /// The single-flight key. Length-prefixed fields, so no two different + /// submissions can collide by concatenation. + pub fn key(&self) -> String { + let mut h = Sha256::new(); + let mut field = |s: &str| { + h.update(u64::try_from(s.len()).unwrap_or(0).to_be_bytes()); + h.update(s.as_bytes()); + }; + field(&self.repo); + field(&self.tree_sha.to_lowercase()); + field(&self.platform); + field(&self.toolchain); + field(self.rust_toolchain.as_deref().unwrap_or("")); + for a in &self.argv { + field(a); + } + hex::encode(h.finalize()) + } +} + +/// A unit as stored and shown. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct Unit { + pub id: String, + pub key: String, + pub repo: String, + pub tree_sha: String, + pub argv: Vec, + pub platform: String, + pub toolchain: String, + pub rust_toolchain: Option, + pub class: Class, + pub timeout_secs: u64, + pub client: String, + pub state: UnitState, + pub drone: Option, + pub attempt: u32, + pub created_at: i64, + pub leased_at: Option, + pub started_at: Option, + pub finished_at: Option, + pub exit_code: Option, + pub duration_secs: Option, + pub cancel_requested: bool, + pub failure: Option, + pub log_bytes: u64, + /// Waiting for the client's bundle upload; not placeable yet. + pub awaiting_bundle: bool, + /// Size of the uploaded bundle, 0 if none. + pub bundle_bytes: u64, +} + +/// The proof a unit succeeded for exactly this tree. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Receipt { + pub unit_id: String, + pub repo: String, + pub tree_sha: String, + pub argv: Vec, + pub platform: String, + pub toolchain: String, + pub drone: String, + pub exit_code: i32, + pub finished_at: i64, + pub duration_secs: u64, + /// SHA-256 over the fields above; lets a client check it was not altered in + /// transit or in a log. + pub digest: String, +} + +impl Unit { + pub fn receipt(&self) -> Option { + if self.state != UnitState::Done || self.exit_code != Some(0) { + return None; + } + let mut r = Receipt { + unit_id: self.id.clone(), + repo: self.repo.clone(), + tree_sha: self.tree_sha.clone(), + argv: self.argv.clone(), + platform: self.platform.clone(), + toolchain: self.toolchain.clone(), + drone: self.drone.clone()?, + exit_code: 0, + finished_at: self.finished_at?, + duration_secs: self.duration_secs.unwrap_or(0), + digest: String::new(), + }; + r.digest = r.compute_digest(); + Some(r) + } +} + +impl Receipt { + pub fn compute_digest(&self) -> String { + let mut h = Sha256::new(); + for s in [ + self.unit_id.as_str(), + &self.repo, + &self.tree_sha, + &self.platform, + &self.toolchain, + &self.drone, + ] { + h.update(u64::try_from(s.len()).unwrap_or(0).to_be_bytes()); + h.update(s.as_bytes()); + } + for a in &self.argv { + h.update(u64::try_from(a.len()).unwrap_or(0).to_be_bytes()); + h.update(a.as_bytes()); + } + h.update(self.exit_code.to_be_bytes()); + h.update(self.finished_at.to_be_bytes()); + h.update(self.duration_secs.to_be_bytes()); + hex::encode(h.finalize()) + } + + pub fn verifies(&self) -> bool { + self.digest == self.compute_digest() + } +} + +/// What the drone is told to run. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct UnitGrant { + pub unit_id: String, + pub repo: String, + pub tree_sha: String, + pub argv: Vec, + pub class: Class, + pub toolchain: String, + pub timeout_secs: u64, + pub attempt: u32, + /// Size of the git bundle to fetch with `GET /v1/units/{id}/bundle` when + /// the tree is not in the drone's mirror; 0 if there is none. + #[serde(default)] + pub bundle_bytes: u64, +} + +/// The result of storing a bundle. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BundleStored { + Stored, + /// No such unit, or it belongs to another client. + NotFound, + /// The unit has already been placed or ended. + TooLate, +} + +/// The outcome of a submission. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Submitted { + /// A new unit was queued. + New(Unit), + /// An identical unit is already queued or running. + Joined(Unit), + /// An identical unit already succeeded recently. + Receipt(Unit), +} + +fn to_i64(v: u64) -> i64 { + i64::try_from(v).unwrap_or(i64::MAX) +} + +fn unit_from_row(r: &sqlx::sqlite::SqliteRow) -> Result { + let argv: String = r.get("argv"); + let class: String = r.get("class"); + Ok(Unit { + id: r.get("id"), + key: r.get("key"), + repo: r.get("repo"), + tree_sha: r.get("tree_sha"), + argv: serde_json::from_str(&argv).map_err(|e| StoreError::Corrupt(e.to_string()))?, + platform: r.get("platform"), + toolchain: r.get("toolchain"), + rust_toolchain: r.get("rust_toolchain"), + class: serde_json::from_str(&format!("\"{class}\"")) + .map_err(|e| StoreError::Corrupt(e.to_string()))?, + timeout_secs: u64::try_from(r.get::("timeout_secs")).unwrap_or(0), + client: r.get("client"), + state: UnitState::parse(&r.get::("state"))?, + drone: r.get("drone"), + attempt: u32::try_from(r.get::("attempt")).unwrap_or(0), + created_at: r.get("created_at"), + leased_at: r.get("leased_at"), + started_at: r.get("started_at"), + finished_at: r.get("finished_at"), + exit_code: r + .get::, _>("exit_code") + .and_then(|v| i32::try_from(v).ok()), + duration_secs: r + .get::, _>("duration_secs") + .map(|v| u64::try_from(v).unwrap_or(0)), + cancel_requested: r.get::("cancel_requested") != 0, + failure: r.get("failure"), + log_bytes: u64::try_from(r.get::("log_bytes")).unwrap_or(0), + awaiting_bundle: r.get::("awaiting_bundle") != 0, + bundle_bytes: u64::try_from(r.get::("bundle_bytes")).unwrap_or(0), + }) +} + +const ACTIVE: &str = "('queued', 'leased', 'running')"; + +impl Store { + pub async fn enroll_client( + &self, + name: &str, + token_hash: &str, + now: i64, + ) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let r = sqlx::query("INSERT INTO clients (name, token_hash, created_at) VALUES (?, ?, ?)") + .bind(name) + .bind(token_hash) + .bind(now) + .execute(&self.pool) + .await; + match r { + Ok(_) => Ok(()), + Err(sqlx::Error::Database(e)) if e.is_unique_violation() => { + Err(StoreError::AlreadyEnrolled(name.to_owned())) + } + Err(e) => Err(e.into()), + } + } + + pub async fn client_for_token_hash( + &self, + token_hash: &str, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT name FROM clients WHERE token_hash = ?") + .bind(token_hash) + .fetch_optional(&self.pool) + .await?; + Ok(row.map(|r| r.get("name"))) + } + + pub async fn unit(&self, id: &str) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT * FROM units WHERE id = ?") + .bind(id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(unit_from_row).transpose() + } + + /// Units in `states`, oldest first. + pub async fn units_in(&self, states: &[UnitState]) -> Result, StoreError> { + let list = states + .iter() + .map(|s| format!("'{}'", s.as_str())) + .collect::>() + .join(", "); + let sql = format!("SELECT * FROM units WHERE state IN ({list}) ORDER BY created_at, id"); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query(&sql).fetch_all(&self.pool).await?; + rows.iter().map(unit_from_row).collect() + } + + pub async fn recent_units(&self, limit: i64) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT * FROM units ORDER BY created_at DESC, id DESC LIMIT ?") + .bind(limit) + .fetch_all(&self.pool) + .await?; + rows.iter().map(unit_from_row).collect() + } + + /// Submit with single-flight. The caller has validated `s`. + pub async fn submit_unit( + &self, + s: &SubmitUnit, + client: &str, + id: &str, + timeout_secs: u64, + receipt_ttl: i64, + now: i64, + ) -> Result { + let key = s.key(); + if let Some(existing) = self.find_reusable(&key, s.force, receipt_ttl, now).await? { + return Ok(existing); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let inserted = sqlx::query("INSERT INTO units (id, key, repo, tree_sha, argv, platform, toolchain, rust_toolchain, class, timeout_secs, client, state, attempt, created_at, awaiting_bundle) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?, ?)") + .bind(id) + .bind(&key) + .bind(&s.repo) + .bind(s.tree_sha.to_lowercase()) + .bind(serde_json::to_string(&s.argv).map_err(|e| StoreError::Corrupt(e.to_string()))?) + .bind(&s.platform) + .bind(&s.toolchain) + .bind(s.rust_toolchain.as_deref()) + .bind(s.class.as_str()) + .bind(to_i64(timeout_secs)) + .bind(client) + .bind(now) + .bind(i64::from(s.bundle)) + .execute(&self.pool) + .await; + match inserted { + Ok(_) => {} + // Another submission won the race for the same key: join it. + Err(sqlx::Error::Database(e)) if e.is_unique_violation() => { + if let Some(existing) = self.find_reusable(&key, s.force, receipt_ttl, now).await? { + return Ok(existing); + } + } + Err(e) => return Err(e.into()), + } + let unit = self + .unit(id) + .await? + .ok_or_else(|| StoreError::Corrupt("unit vanished".into()))?; + Ok(Submitted::New(unit)) + } + + /// An identical unit that is running or succeeded recently, if any (what a + /// submission would join or reuse). + pub async fn find_reusable_unit( + &self, + s: &SubmitUnit, + receipt_ttl: i64, + now: i64, + ) -> Result, StoreError> { + self.find_reusable(&s.key(), s.force, receipt_ttl, now) + .await + } + + async fn find_reusable( + &self, + key: &str, + force: bool, + receipt_ttl: i64, + now: i64, + ) -> Result, StoreError> { + let sql = format!("SELECT * FROM units WHERE key = ? AND state IN {ACTIVE} LIMIT 1"); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let active = sqlx::query(&sql) + .bind(key) + .fetch_optional(&self.pool) + .await?; + if let Some(r) = active { + return Ok(Some(Submitted::Joined(unit_from_row(&r)?))); + } + if force { + return Ok(None); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let done = sqlx::query("SELECT * FROM units WHERE key = ? AND state = 'done' AND exit_code = 0 AND finished_at >= ? ORDER BY finished_at DESC LIMIT 1") + .bind(key) + .bind(now - receipt_ttl) + .fetch_optional(&self.pool) + .await?; + done.as_ref() + .map(|r| unit_from_row(r).map(Submitted::Receipt)) + .transpose() + } + + /// Number of a client's units that are queued or in flight. + pub async fn client_active_units(&self, client: &str) -> Result { + let sql = format!("SELECT COUNT(*) AS n FROM units WHERE client = ? AND state IN {ACTIVE}"); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query(&sql).bind(client).fetch_one(&self.pool).await?; + Ok(u32::try_from(row.get::("n")).unwrap_or(u32::MAX)) + } + + /// Place a queued unit on a drone. + pub async fn lease_unit(&self, id: &str, drone: &str, now: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let r = sqlx::query("UPDATE units SET state = 'leased', drone = ?, leased_at = ?, attempt = attempt + 1 WHERE id = ? AND state = 'queued'") + .bind(drone) + .bind(now) + .bind(id) + .execute(&self.pool) + .await?; + Ok(r.rows_affected() > 0) + } + + /// Units placed on `drone` that it has not acknowledged yet. + pub async fn unacked_units(&self, drone: &str) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT * FROM units WHERE drone = ? AND state = 'leased' AND cancel_requested = 0 ORDER BY created_at, id") + .bind(drone) + .fetch_all(&self.pool) + .await?; + rows.iter().map(unit_from_row).collect() + } + + /// The drone acknowledges units; leased ones start running. + pub async fn ack_units(&self, drone: &str, ids: &[String], now: i64) -> Result<(), StoreError> { + for id in ids { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE units SET state = 'running', started_at = ? WHERE id = ? AND drone = ? AND state = 'leased'") + .bind(now) + .bind(id) + .bind(drone) + .execute(&self.pool) + .await?; + } + Ok(()) + } + + /// Ids on `drone` that a client asked to cancel while queued, leased or + /// running. + pub async fn cancel_requests(&self, drone: &str) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT id FROM units WHERE drone = ? AND cancel_requested = 1 AND state IN ('leased', 'running') ORDER BY id") + .bind(drone) + .fetch_all(&self.pool) + .await?; + Ok(rows.iter().map(|r| r.get("id")).collect()) + } + + /// Client cancel: queued and leased-but-unstarted units end at once; + /// running ones are flagged for the drone. False if not found or already + /// finished. + pub async fn cancel_unit(&self, id: &str, now: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let ended = sqlx::query("UPDATE units SET state = 'failed', failure = 'cancelled', finished_at = ?, cancel_requested = 1 WHERE id = ? AND state IN ('queued', 'leased')") + .bind(now) + .bind(id) + .execute(&self.pool) + .await?; + if ended.rows_affected() > 0 { + self.delete_unit_bundle(id).await?; + return Ok(true); + } + const FLAG: &str = + "UPDATE units SET cancel_requested = 1 WHERE id = ? AND state = 'running'"; + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let flagged = sqlx::query(FLAG).bind(id).execute(&self.pool).await?; + Ok(flagged.rows_affected() > 0) + } + + /// The drone reports a unit's end. Idempotent; false if the unit is not + /// this drone's or already finished with a different outcome. + pub async fn finish_unit( + &self, + id: &str, + drone: &str, + exit_code: i32, + duration_secs: u64, + now: i64, + ) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let owned = sqlx::query("SELECT state FROM units WHERE id = ? AND drone = ?") + .bind(id) + .bind(drone) + .fetch_optional(&self.pool) + .await?; + let Some(row) = owned else { + return Ok(false); + }; + let state: String = row.get("state"); + if state == "done" { + return Ok(true); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE units SET state = 'done', exit_code = ?, duration_secs = ?, finished_at = ? WHERE id = ? AND drone = ? AND state IN ('leased', 'running', 'failed')") + .bind(i64::from(exit_code)) + .bind(to_i64(duration_secs)) + .bind(now) + .bind(id) + .bind(drone) + .execute(&self.pool) + .await?; + Ok(true) + } + + /// Running units on `drone` that its heartbeat no longer lists (started + /// before `cutoff`) are lost: re-queue them, or fail them once out of + /// attempts. Returns how many were affected. + pub async fn requeue_lost_units( + &self, + drone: &str, + active: &[String], + cutoff: i64, + max_attempts: u32, + now: i64, + ) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT id, attempt FROM units WHERE drone = ? AND state = 'running' AND started_at < ?") + .bind(drone) + .bind(cutoff) + .fetch_all(&self.pool) + .await?; + let mut n = 0; + for r in rows { + let id: String = r.get("id"); + if active.contains(&id) { + continue; + } + let attempt = u32::try_from(r.get::("attempt")).unwrap_or(0); + self.retry_or_fail(&id, attempt, max_attempts, "drone lost the unit", now) + .await?; + n += 1; + } + Ok(n) + } + + /// Re-queue a unit, or fail it if it has used all its attempts. + pub async fn retry_or_fail( + &self, + id: &str, + attempt: u32, + max_attempts: u32, + why: &str, + now: i64, + ) -> Result<(), StoreError> { + if attempt >= max_attempts { + self.fail_unit(id, &format!("{why}; out of attempts"), now) + .await + } else { + self.requeue_unit(id).await + } + } + + /// Put a unit back in the queue (its drone vanished or never started). + pub async fn requeue_unit(&self, id: &str) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE units SET state = 'queued', drone = NULL, leased_at = NULL, started_at = NULL WHERE id = ? AND state IN ('leased', 'running')") + .bind(id) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn fail_unit(&self, id: &str, why: &str, now: i64) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE units SET state = 'failed', failure = ?, finished_at = ? WHERE id = ? AND state IN ('queued', 'leased', 'running')") + .bind(why) + .bind(now) + .bind(id) + .execute(&self.pool) + .await?; + self.delete_unit_bundle(id).await + } + + /// Store the client's bundle for a unit that is still waiting to be + /// placed. The bundle may replace an earlier upload. + pub async fn put_unit_bundle( + &self, + id: &str, + client: &str, + data: &[u8], + now: i64, + ) -> Result { + let Some(unit) = self.unit(id).await? else { + return Ok(BundleStored::NotFound); + }; + if unit.client != client { + return Ok(BundleStored::NotFound); + } + if unit.state != UnitState::Queued { + return Ok(BundleStored::TooLate); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT OR REPLACE INTO unit_bundles (unit_id, data, created_at) VALUES (?, ?, ?)", + ) + .bind(id) + .bind(data) + .bind(now) + .execute(&self.pool) + .await?; + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE units SET awaiting_bundle = 0, bundle_bytes = ? WHERE id = ?") + .bind(to_i64(u64::try_from(data.len()).unwrap_or(0))) + .bind(id) + .execute(&self.pool) + .await?; + Ok(BundleStored::Stored) + } + + pub async fn unit_bundle(&self, id: &str) -> Result>, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT data FROM unit_bundles WHERE unit_id = ?") + .bind(id) + .fetch_optional(&self.pool) + .await?; + Ok(row.map(|r| r.get("data"))) + } + + pub async fn delete_unit_bundle(&self, id: &str) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("DELETE FROM unit_bundles WHERE unit_id = ?") + .bind(id) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Append drone log bytes at `offset`. Idempotent for a repeated chunk; + /// an offset beyond the stored end (a gap) is rejected so the log never + /// has holes. Returns the new length, or `None` for a gap. + pub async fn append_unit_log( + &self, + id: &str, + offset: u64, + data: &[u8], + max_bytes: u64, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT log_bytes FROM units WHERE id = ?") + .bind(id) + .fetch_optional(&self.pool) + .await?; + let Some(row) = row else { + return Ok(None); + }; + let have = u64::try_from(row.get::("log_bytes")).unwrap_or(0); + if offset > have { + return Ok(None); + } + let skip = usize::try_from(have - offset) + .unwrap_or(usize::MAX) + .min(data.len()); + let fresh = &data[skip..]; + let room = usize::try_from(max_bytes.saturating_sub(have)).unwrap_or(usize::MAX); + let fresh = &fresh[..fresh.len().min(room)]; + if fresh.is_empty() { + return Ok(Some(have)); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("INSERT OR IGNORE INTO unit_logs (unit_id, offset, data) VALUES (?, ?, ?)") + .bind(id) + .bind(to_i64(have)) + .bind(fresh) + .execute(&self.pool) + .await?; + let new_len = have + u64::try_from(fresh.len()).unwrap_or(0); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE units SET log_bytes = ? WHERE id = ?") + .bind(to_i64(new_len)) + .bind(id) + .execute(&self.pool) + .await?; + Ok(Some(new_len)) + } + + /// Log bytes from `from`, at most `limit`. + pub async fn read_unit_log( + &self, + id: &str, + from: u64, + limit: usize, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT offset, data FROM unit_logs WHERE unit_id = ? AND offset + length(data) > ? ORDER BY offset") + .bind(id) + .bind(to_i64(from)) + .fetch_all(&self.pool) + .await?; + let mut out = Vec::new(); + for r in rows { + let off = u64::try_from(r.get::("offset")).unwrap_or(0); + let data: Vec = r.get("data"); + let skip = usize::try_from(from.saturating_sub(off)) + .unwrap_or(usize::MAX) + .min(data.len()); + out.extend_from_slice(&data[skip..]); + if out.len() >= limit { + out.truncate(limit); + break; + } + } + Ok(out) + } + + /// Delete finished units (and their logs) older than `before`. + pub async fn prune_units(&self, before: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("DELETE FROM unit_bundles WHERE unit_id IN (SELECT id FROM units WHERE state IN ('done', 'failed') AND finished_at < ?)") + .bind(before) + .execute(&self.pool) + .await?; + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("DELETE FROM unit_logs WHERE unit_id IN (SELECT id FROM units WHERE state IN ('done', 'failed') AND finished_at < ?)") + .bind(before) + .execute(&self.pool) + .await?; + const DROP: &str = + "DELETE FROM units WHERE state IN ('done', 'failed') AND finished_at < ?"; + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let r = sqlx::query(DROP).bind(before).execute(&self.pool).await?; + Ok(r.rows_affected()) + } + + /// Counts by state. + pub async fn unit_counts(&self) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT state, COUNT(*) AS n FROM units GROUP BY state") + .fetch_all(&self.pool) + .await?; + let mut out: BTreeMap<&'static str, u32> = + UnitState::ALL.iter().map(|s| (s.as_str(), 0)).collect(); + for r in rows { + let s = UnitState::parse(&r.get::("state"))?; + out.insert( + s.as_str(), + u32::try_from(r.get::("n")).unwrap_or(u32::MAX), + ); + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn submit() -> SubmitUnit { + SubmitUnit { + repo: "o/r".into(), + tree_sha: "a".repeat(40), + argv: vec!["cargo".into(), "check".into()], + platform: "linux-x86_64".into(), + toolchain: "rust-1.95".into(), + class: Class::Medium, + timeout_secs: None, + force: false, + rust_toolchain: None, + bundle: false, + } + } + + #[test] + fn the_key_covers_every_field_and_cannot_collide_by_concatenation() { + let base = submit(); + let k = base.key(); + assert_eq!(k, submit().key()); + for changed in [ + SubmitUnit { + repo: "o/x".into(), + ..submit() + }, + SubmitUnit { + tree_sha: "b".repeat(40), + ..submit() + }, + SubmitUnit { + platform: "macos-aarch64".into(), + ..submit() + }, + SubmitUnit { + toolchain: "rust-1.96".into(), + ..submit() + }, + SubmitUnit { + argv: vec!["cargo".into(), "test".into()], + ..submit() + }, + SubmitUnit { + argv: vec!["cargo check".into()], + ..submit() + }, + SubmitUnit { + argv: vec!["cargoc".into(), "heck".into()], + ..submit() + }, + ] { + assert_ne!(changed.key(), k, "{changed:?}"); + } + // Case of the object id is not part of the identity; class and force are not either. + let upper = SubmitUnit { + tree_sha: "A".repeat(40), + ..submit() + }; + assert_eq!(upper.key(), k); + assert_eq!( + SubmitUnit { + class: Class::Heavy, + force: true, + ..submit() + } + .key(), + k + ); + } + + #[test] + fn submissions_are_validated() { + assert!(submit().validate().is_ok()); + for bad in [ + SubmitUnit { + repo: "nope".into(), + ..submit() + }, + SubmitUnit { + tree_sha: "xyz".into(), + ..submit() + }, + SubmitUnit { + argv: vec![], + ..submit() + }, + SubmitUnit { + argv: vec!["a\0b".into()], + ..submit() + }, + SubmitUnit { + platform: "Linux X".into(), + ..submit() + }, + SubmitUnit { + toolchain: "x\n".into(), + ..submit() + }, + ] { + assert!(bad.validate().is_err(), "{bad:?}"); + } + } + + async fn store() -> Store { + Store::open_in_memory().await.unwrap() + } + + #[tokio::test] + async fn single_flight_joins_running_units_and_reuses_only_fresh_successes() { + let s = store().await; + let first = s + .submit_unit(&submit(), "c", "u1", 60, 1000, 100) + .await + .unwrap(); + assert!(matches!(first, Submitted::New(_))); + // Queued, then running: still the same unit. + let again = s + .submit_unit(&submit(), "c2", "u2", 60, 1000, 101) + .await + .unwrap(); + let Submitted::Joined(u) = again else { + panic!("{again:?}") + }; + assert_eq!(u.id, "u1"); + + s.enroll("d1", "h", 0).await.unwrap(); + assert!(s.lease_unit("u1", "d1", 102).await.unwrap()); + s.ack_units("d1", &["u1".into()], 103).await.unwrap(); + assert!(matches!( + s.submit_unit(&submit(), "c", "u3", 60, 1000, 104) + .await + .unwrap(), + Submitted::Joined(_) + )); + + // A failure is not reused: a new unit runs. + s.finish_unit("u1", "d1", 1, 5, 110).await.unwrap(); + let after_fail = s + .submit_unit(&submit(), "c", "u4", 60, 1000, 111) + .await + .unwrap(); + assert!(matches!(after_fail, Submitted::New(_)), "{after_fail:?}"); + + // A success is reused until its ttl ends, never with force. + s.lease_unit("u4", "d1", 112).await.unwrap(); + s.finish_unit("u4", "d1", 0, 7, 120).await.unwrap(); + let hit = s + .submit_unit(&submit(), "c", "u5", 60, 1000, 500) + .await + .unwrap(); + let Submitted::Receipt(u) = hit else { + panic!("{hit:?}") + }; + let r = u.receipt().unwrap(); + assert!(r.verifies()); + assert_eq!( + (r.tree_sha.as_str(), r.exit_code, r.drone.as_str()), + (&*"a".repeat(40), 0, "d1") + ); + let forced = SubmitUnit { + force: true, + ..submit() + }; + assert!(matches!( + s.submit_unit(&forced, "c", "u6", 60, 1000, 500) + .await + .unwrap(), + Submitted::New(_) + )); + // Expired receipt: a new unit. (u6 is active, so first end it.) + s.cancel_unit("u6", 501).await.unwrap(); + let late = s + .submit_unit(&submit(), "c", "u7", 60, 1000, 5000) + .await + .unwrap(); + assert!(matches!(late, Submitted::New(_))); + // A different tree never matches. + let other = SubmitUnit { + tree_sha: "c".repeat(40), + ..submit() + }; + assert!(matches!( + s.submit_unit(&other, "c", "u8", 60, 1000, 130) + .await + .unwrap(), + Submitted::New(_) + )); + } + + #[tokio::test] + async fn receipts_detect_tampering() { + let s = store().await; + s.enroll("d1", "h", 0).await.unwrap(); + s.submit_unit(&submit(), "c", "u1", 60, 1000, 100) + .await + .unwrap(); + s.lease_unit("u1", "d1", 101).await.unwrap(); + s.finish_unit("u1", "d1", 0, 3, 110).await.unwrap(); + let mut r = s.unit("u1").await.unwrap().unwrap().receipt().unwrap(); + assert!(r.verifies()); + r.tree_sha = "b".repeat(40); + assert!(!r.verifies()); + // Nonzero exits have no receipt. + s.submit_unit( + &SubmitUnit { + force: true, + ..submit() + }, + "c", + "u2", + 60, + 1000, + 120, + ) + .await + .unwrap(); + s.lease_unit("u2", "d1", 121).await.unwrap(); + s.finish_unit("u2", "d1", 2, 3, 130).await.unwrap(); + assert!(s.unit("u2").await.unwrap().unwrap().receipt().is_none()); + } + + #[tokio::test] + async fn logs_append_idempotently_without_gaps_and_read_from_any_offset() { + let s = store().await; + s.submit_unit(&submit(), "c", "u1", 60, 1000, 100) + .await + .unwrap(); + assert_eq!( + s.append_unit_log("u1", 0, b"hello ", 100).await.unwrap(), + Some(6) + ); + // The same chunk again, and an overlapping one. + assert_eq!( + s.append_unit_log("u1", 0, b"hello ", 100).await.unwrap(), + Some(6) + ); + assert_eq!( + s.append_unit_log("u1", 3, b"lo world", 100).await.unwrap(), + Some(11) + ); + // A gap is refused. + assert_eq!(s.append_unit_log("u1", 50, b"x", 100).await.unwrap(), None); + assert_eq!( + s.read_unit_log("u1", 0, 1000).await.unwrap(), + b"hello world" + ); + assert_eq!(s.read_unit_log("u1", 4, 1000).await.unwrap(), b"o world"); + assert_eq!(s.read_unit_log("u1", 4, 3).await.unwrap(), b"o w"); + assert!(s.read_unit_log("u1", 11, 10).await.unwrap().is_empty()); + // The cap truncates. + assert_eq!( + s.append_unit_log("u1", 11, b"0123456789", 15) + .await + .unwrap(), + Some(15) + ); + assert_eq!(s.append_unit_log("nope", 0, b"x", 15).await.unwrap(), None); + } + + #[tokio::test] + async fn cancel_requeue_fail_and_prune() { + let s = store().await; + s.enroll("d1", "h", 0).await.unwrap(); + s.submit_unit(&submit(), "c", "u1", 60, 1000, 100) + .await + .unwrap(); + assert!(s.cancel_unit("u1", 101).await.unwrap()); + let u = s.unit("u1").await.unwrap().unwrap(); + assert_eq!( + (u.state, u.failure.as_deref()), + (UnitState::Failed, Some("cancelled")) + ); + assert!(!s.cancel_unit("u1", 102).await.unwrap()); + + s.submit_unit(&submit(), "c", "u2", 60, 1000, 110) + .await + .unwrap(); + s.lease_unit("u2", "d1", 111).await.unwrap(); + s.requeue_unit("u2").await.unwrap(); + let u = s.unit("u2").await.unwrap().unwrap(); + assert_eq!((u.state, u.drone, u.attempt), (UnitState::Queued, None, 1)); + s.lease_unit("u2", "d1", 112).await.unwrap(); + s.ack_units("d1", &["u2".into()], 113).await.unwrap(); + assert!(s.cancel_unit("u2", 114).await.unwrap()); + assert_eq!(s.cancel_requests("d1").await.unwrap(), ["u2"]); + s.fail_unit("u2", "timed out", 200).await.unwrap(); + assert_eq!(s.client_active_units("c").await.unwrap(), 0); + assert_eq!(s.prune_units(150).await.unwrap(), 1); + assert_eq!(s.prune_units(300).await.unwrap(), 1); + assert!(s.unit_counts().await.unwrap().values().all(|&n| n == 0)); + } + + #[tokio::test] + async fn clients_enroll_once() { + let s = store().await; + s.enroll_client("agent", "hash", 1).await.unwrap(); + assert!(matches!( + s.enroll_client("agent", "other", 2).await, + Err(StoreError::AlreadyEnrolled(_)) + )); + assert_eq!( + s.client_for_token_hash("hash").await.unwrap().as_deref(), + Some("agent") + ); + assert!(s.client_for_token_hash("nope").await.unwrap().is_none()); + } +} diff --git a/src/units_api.rs b/src/units_api.rs new file mode 100644 index 0000000..bf9ce60 --- /dev/null +++ b/src/units_api.rs @@ -0,0 +1,335 @@ +//! HTTP endpoints for work units; see [`crate::units`]. +//! +//! Clients (bearer: client token): `POST /v1/units`, `GET /v1/units/{id}`, +//! `GET /v1/units/{id}/log?from=N&wait=S`, `POST /v1/units/{id}/cancel`. +//! Drones (bearer: drone token): `POST /v1/units/{id}/log?offset=N` and +//! `POST /v1/units/{id}/result`; units reach a drone in its heartbeat +//! response. + +use std::sync::Arc; +use std::time::Duration; + +use axum::Json; +use axum::Router; +use axum::body::Body; +use axum::body::Bytes; +use axum::extract::{Path, Query, State}; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post, put}; +use serde::{Deserialize, Serialize}; +use time::OffsetDateTime; + +use crate::api::{ApiError, authenticate}; +use crate::enroll::hash_token; +use crate::queen::Queen; +use crate::units::{BundleStored, Receipt, SubmitUnit, Submitted, Unit}; + +/// Longest a log long-poll waits. +const MAX_WAIT_SECS: u64 = 30; +/// Most log bytes returned by one read. +const READ_CHUNK: usize = 256 * 1024; + +pub fn routes() -> Router> { + Router::new() + .route("/v1/units", post(submit)) + .route("/v1/units/{id}", get(show)) + .route("/v1/units/{id}/log", get(read_log).post(append_log)) + .route("/v1/units/{id}/bundle", put(put_bundle).get(get_bundle)) + .route("/v1/units/{id}/cancel", post(cancel)) + .route("/v1/units/{id}/result", post(result)) +} + +fn now() -> i64 { + OffsetDateTime::now_utc().unix_timestamp() +} + +fn enabled(queen: &Queen) -> Result<(), ApiError> { + if queen.config().units.enabled { + Ok(()) + } else { + Err(ApiError::NotFound) + } +} + +async fn authenticate_client(queen: &Queen, headers: &HeaderMap) -> Result { + let presented = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.strip_prefix("Bearer ")) + .ok_or(ApiError::Unauthorized)?; + queen + .store() + .client_for_token_hash(&hash_token(presented)) + .await? + .ok_or(ApiError::Unauthorized) +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct SubmitResponse { + /// `new`, `joined` or `receipt`. + pub outcome: String, + pub unit_id: String, + pub state: String, + pub receipt: Option, +} + +async fn submit( + State(queen): State>, + headers: HeaderMap, + Json(req): Json, +) -> Result<(StatusCode, Json), ApiError> { + enabled(&queen)?; + let client = authenticate_client(&queen, &headers).await?; + req.validate().map_err(ApiError::BadRequest)?; + let cfg = &queen.config().units; + let timeout = req + .timeout_secs + .unwrap_or(cfg.default_timeout_secs) + .clamp(1, cfg.max_timeout_secs); + let store = queen.store(); + let t = now(); + // Joining or reusing a result costs nothing, so it is allowed even when the + // client is at its limit of queued units. + if store.client_active_units(&client).await? >= cfg.max_queued_per_client + && store + .find_reusable_unit(&req, cfg.receipt_ttl_secs, t) + .await? + .is_none() + { + return Err(ApiError::TooMany); + } + let id = crate::queen::new_unit_id(); + let submitted = store + .submit_unit(&req, &client, &id, timeout, cfg.receipt_ttl_secs, t) + .await?; + let (outcome, unit, status) = match submitted { + Submitted::New(u) => ("new", u, StatusCode::CREATED), + Submitted::Joined(u) => { + queen.count_unit_join(); + ("joined", u, StatusCode::OK) + } + Submitted::Receipt(u) => { + queen.count_unit_receipt(); + ("receipt", u, StatusCode::OK) + } + }; + tracing::info!( + "unit {} {outcome} for {} {} by {client}", + unit.id, + unit.repo, + &unit.tree_sha[..unit.tree_sha.len().min(12)] + ); + Ok(( + status, + Json(SubmitResponse { + outcome: outcome.to_owned(), + unit_id: unit.id.clone(), + state: unit.state.as_str().to_owned(), + receipt: unit.receipt(), + }), + )) +} + +#[derive(Debug, Serialize)] +pub struct UnitView { + #[serde(flatten)] + pub unit: Unit, + pub receipt: Option, +} + +async fn show( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, +) -> Result, ApiError> { + enabled(&queen)?; + authenticate_client(&queen, &headers).await?; + let unit = queen.store().unit(&id).await?.ok_or(ApiError::NotFound)?; + let receipt = unit.receipt(); + Ok(Json(UnitView { unit, receipt })) +} + +async fn cancel( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, +) -> Result { + enabled(&queen)?; + authenticate_client(&queen, &headers).await?; + if queen.store().cancel_unit(&id, now()).await? { + Ok(StatusCode::NO_CONTENT) + } else { + Err(ApiError::NotFound) + } +} + +#[derive(Debug, Deserialize)] +struct LogQuery { + from: Option, + wait: Option, +} + +/// Log bytes from `from`. Waits up to `wait` seconds (at most 30) for new +/// bytes while the unit is still active. Headers: `x-borg-next` (offset to +/// ask for next), `x-borg-state`, and `x-borg-exit` once finished. +async fn read_log( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, + Query(q): Query, +) -> Result { + enabled(&queen)?; + authenticate_client(&queen, &headers).await?; + let from = q.from.unwrap_or(0); + let deadline = + tokio::time::Instant::now() + Duration::from_secs(q.wait.unwrap_or(0).min(MAX_WAIT_SECS)); + loop { + let unit = queen.store().unit(&id).await?.ok_or(ApiError::NotFound)?; + let data = queen.store().read_unit_log(&id, from, READ_CHUNK).await?; + let finished = !unit.state.is_active(); + // Fully drained and finished, or data to hand over, or out of time. + if !data.is_empty() || finished || tokio::time::Instant::now() >= deadline { + let next = from + u64::try_from(data.len()).unwrap_or(0); + let mut resp = (StatusCode::OK, data).into_response(); + let h = resp.headers_mut(); + let mut put = |k: &'static str, v: String| { + if let Ok(v) = v.parse() { + h.insert(k, v); + } + }; + put("x-borg-next", next.to_string()); + put("x-borg-state", unit.state.as_str().to_owned()); + // Only report the end once every byte has been handed out. + if finished && next >= unit.log_bytes { + if let Some(code) = unit.exit_code { + put("x-borg-exit", code.to_string()); + } + if let Some(why) = &unit.failure { + put("x-borg-failure", why.clone()); + } + } + return Ok(resp); + } + tokio::time::sleep(Duration::from_millis(250)).await; + } +} + +#[derive(Debug, Deserialize)] +struct AppendQuery { + offset: u64, +} + +async fn drone_owns(queen: &Queen, drone: &str, id: &str) -> Result { + let unit = queen.store().unit(id).await?.ok_or(ApiError::NotFound)?; + if unit.drone.as_deref() == Some(drone) { + Ok(unit) + } else { + Err(ApiError::NotFound) + } +} + +async fn append_log( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, + Query(q): Query, + body: Bytes, +) -> Result { + enabled(&queen)?; + let drone = authenticate(&queen, &headers).await?; + drone_owns(&queen, &drone, &id).await?; + let max = queen.config().units.max_log_bytes; + match queen + .store() + .append_unit_log(&id, q.offset, &body, max) + .await? + { + Some(len) => Ok((StatusCode::OK, len.to_string()).into_response()), + // The drone is ahead of what the queen has: it must resend from the + // length the queen holds. + None => Ok(StatusCode::CONFLICT.into_response()), + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct UnitResult { + pub exit_code: i32, + pub duration_secs: u64, +} + +async fn result( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, + Json(req): Json, +) -> Result { + enabled(&queen)?; + let drone = authenticate(&queen, &headers).await?; + drone_owns(&queen, &drone, &id).await?; + if queen + .store() + .finish_unit(&id, &drone, req.exit_code, req.duration_secs, now()) + .await? + { + Ok(StatusCode::NO_CONTENT) + } else { + Err(ApiError::NotFound) + } +} + +/// `PUT /v1/units/{id}/bundle`: the client uploads a git bundle holding the +/// objects of a tree that is not on any remote (a unit submitted with +/// `bundle: true`). Authenticated as the unit's client and capped at +/// `units.max_bundle_bytes`; deleted with the unit when its receipt expires (or +/// at once if the unit fails or is cancelled). +async fn put_bundle( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, + body: Body, +) -> Result { + enabled(&queen)?; + let client = authenticate_client(&queen, &headers).await?; + let cap = queen.config().units.max_bundle_bytes; + let declared = headers + .get("content-length") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()); + if declared.is_some_and(|n| n > cap) { + return Err(ApiError::TooLarge); + } + let data = axum::body::to_bytes(body, usize::try_from(cap).unwrap_or(usize::MAX)) + .await + .map_err(|_| ApiError::TooLarge)?; + if data.is_empty() { + return Err(ApiError::BadRequest("empty bundle".into())); + } + match queen + .store() + .put_unit_bundle(&id, &client, &data, now()) + .await? + { + BundleStored::Stored => Ok(StatusCode::NO_CONTENT), + BundleStored::NotFound => Err(ApiError::NotFound), + BundleStored::TooLate => Err(ApiError::Conflict), + } +} + +/// `GET /v1/units/{id}/bundle`: the drone running the unit fetches the bundle +/// when the tree is not in its git mirror. +async fn get_bundle( + State(queen): State>, + headers: HeaderMap, + Path(id): Path, +) -> Result { + enabled(&queen)?; + let drone = authenticate(&queen, &headers).await?; + drone_owns(&queen, &drone, &id).await?; + let data = queen + .store() + .unit_bundle(&id) + .await? + .ok_or(ApiError::NotFound)?; + Ok(([("content-type", "application/octet-stream")], data).into_response()) +} diff --git a/src/update/cli.rs b/src/update/cli.rs new file mode 100644 index 0000000..c15ff42 --- /dev/null +++ b/src/update/cli.rs @@ -0,0 +1,168 @@ +//! The `borg update` command line: configuration, manifest fetch, the real +//! service manager and health file, and the request/result hand-off with a +//! running drone. + +use std::path::PathBuf; +use std::time::Duration; + +use anyhow::{Context, Result, bail}; + +use crate::drone::exec::SystemExec; +use crate::update::config::UpdateConfig; +use crate::update::health::FileHealth; +use crate::update::manifest; +use crate::update::service::{Launchd, ServiceManager, Systemd}; +use crate::update::{Options, Outcome, Request, UpdateResult, request_path, result_path}; + +pub struct Args { + pub config: PathBuf, + pub version: Option, + pub check: bool, + pub runner_only: bool, + /// Read the target version from the drone's request file and leave a + /// result for it (how the systemd path unit invokes the updater). + pub from_request: bool, + pub json: bool, +} + +fn unix_now() -> i64 { + time::OffsetDateTime::now_utc().unix_timestamp() +} + +/// Run an update; returns the process exit code (0 ok, 3 rolled back). +pub async fn run(args: Args) -> Result { + let cfg = UpdateConfig::load(&args.config)?; + let mut version = args.version.clone(); + if args.from_request { + let path = request_path(&cfg.state_dir); + let text = + std::fs::read_to_string(&path).with_context(|| format!("read {}", path.display()))?; + let request: Request = serde_json::from_str(&text).context("parse update request")?; + if !manifest::valid_version(&request.version) { + bail!("the update request names an invalid version"); + } + version = Some(request.version); + } + + let outcome = attempt(&cfg, &args, version.as_deref()).await; + if args.from_request { + let result = + UpdateResult::from_outcome(version.as_deref().unwrap_or(""), &outcome, unix_now()); + std::fs::write(result_path(&cfg.state_dir), serde_json::to_vec(&result)?) + .context("write update result")?; + let _ = std::fs::remove_file(request_path(&cfg.state_dir)); + } + let outcome = outcome?; + if args.json { + println!("{}", serde_json::to_string_pretty(&outcome)?); + } else { + println!("{}", describe(&outcome)); + } + Ok(i32::from(matches!(outcome, Outcome::RolledBack { .. })) * 3) +} + +async fn attempt( + cfg: &UpdateConfig, + args: &Args, + version: Option<&str>, +) -> Result { + let manifest = manifest::fetch(&cfg.manifest_url_for(version), cfg.allow_insecure).await?; + let exec = SystemExec; + let service: Box = if cfg!(target_os = "macos") { + Box::new(Launchd { + label: cfg.service.clone(), + plist: cfg.launchd_plist.clone().unwrap_or_default(), + exec: &exec, + }) + } else { + Box::new(Systemd { + unit: cfg.service.clone(), + exec: &exec, + }) + }; + let opts = Options { + bin: &cfg.bin, + current_version: env!("CARGO_PKG_VERSION"), + requested_version: version, + check: args.check, + runner_only: args.runner_only, + allow_insecure: cfg.allow_insecure, + health_timeout: Duration::from_secs(cfg.health_timeout_secs), + poll: Duration::from_secs(2), + host_os: std::env::consts::OS, + host_arch: std::env::consts::ARCH, + runner_dir: cfg.runner_dir.as_deref(), + runner_owner: cfg.runner_owner.as_deref(), + }; + crate::update::run( + &opts, + &manifest, + service.as_ref(), + &FileHealth(cfg.health_file.clone()), + &exec, + ) + .await +} + +pub fn describe(outcome: &Outcome) -> String { + match outcome { + Outcome::UpToDate { version } => format!("up to date ({version})"), + Outcome::Available { + current, + target, + runner, + } => match runner { + Some(r) => format!("update available: {current} -> {target}, runner {r}"), + None => format!("update available: {current} -> {target}"), + }, + Outcome::Updated { + from, + to, + runner_updated, + } => { + format!( + "updated {from} -> {to}{}", + if *runner_updated { + " (runner updated)" + } else { + "" + } + ) + } + Outcome::RunnerUpdated { version } => format!("runner updated to {version}"), + Outcome::RolledBack { from, to, reason } => { + format!("rolled back to {from}: {to} failed: {reason}") + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn outcomes_read_clearly() { + assert_eq!( + describe(&Outcome::UpToDate { + version: "1".into() + }), + "up to date (1)" + ); + assert!( + describe(&Outcome::Available { + current: "1".into(), + target: "2".into(), + runner: Some("2.3".into()) + }) + .contains("runner 2.3") + ); + assert!( + describe(&Outcome::RolledBack { + from: "1".into(), + to: "2".into(), + reason: "x".into() + }) + .starts_with("rolled back to 1") + ); + } +} diff --git a/src/update/config.rs b/src/update/config.rs new file mode 100644 index 0000000..2d6ab2c --- /dev/null +++ b/src/update/config.rs @@ -0,0 +1,171 @@ +//! Update configuration. It is read by the privileged updater, so it must live +//! somewhere only root can change (`/etc/borg-update`, not the drone's own +//! `/etc/borg`): otherwise a compromised drone could point the updater at its +//! own "release" and have it installed as root. + +use std::path::{Path, PathBuf}; + +use serde::{Deserialize, Serialize}; + +pub const DEFAULT_DRONE_CONFIG: &str = "/etc/borg-update/update.json"; +pub const DEFAULT_QUEEN_CONFIG: &str = "/etc/borg-update/queen.json"; +/// The projected public repository's release manifest, used when the installer +/// is not told otherwise. +pub const DEFAULT_MANIFEST_URL: &str = + "https://github.com/dx-corp/borg/releases/latest/download/borg-manifest.json"; +pub const DEFAULT_VERSIONED_MANIFEST_URL: &str = + "https://github.com/dx-corp/borg/releases/download/v{version}/borg-manifest.json"; + +#[derive(Debug, thiserror::Error)] +pub enum UpdateConfigError { + #[error("read {path}: {source}")] + Read { + path: PathBuf, + source: std::io::Error, + }, + #[error("parse update config: {0}")] + Parse(#[from] serde_json::Error), + #[error("{0}")] + Invalid(String), +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UpdateConfig { + /// Manifest of the newest release. + pub manifest_url: String, + /// Manifest of a specific release; `{version}` is replaced. Used when a + /// version is requested (`--version`, or by the queen's rollout). + #[serde(default)] + pub versioned_manifest_url: Option, + /// Permit plain `http://` URLs (trusted network, tests). + #[serde(default)] + pub allow_insecure: bool, + /// systemd unit (Linux) or launchd label (macOS) of the service. + pub service: String, + /// launchd only: the plist path used to start the service again. + #[serde(default)] + pub launchd_plist: Option, + #[serde(default = "default_bin")] + pub bin: PathBuf, + /// The record the daemon writes after each successful heartbeat/pass. + pub health_file: PathBuf, + /// Directory of the drone's `update-request` and `update-result.json`. + pub state_dir: PathBuf, + /// The pinned actions-runner directory and its owner, if this host has one. + #[serde(default)] + pub runner_dir: Option, + #[serde(default)] + pub runner_owner: Option, + #[serde(default = "default_health_timeout")] + pub health_timeout_secs: u64, +} + +fn default_bin() -> PathBuf { + "/usr/local/bin/borg".into() +} +fn default_health_timeout() -> u64 { + 120 +} + +impl UpdateConfig { + pub fn load(path: &Path) -> Result { + let text = std::fs::read_to_string(path).map_err(|source| UpdateConfigError::Read { + path: path.to_owned(), + source, + })?; + Self::from_json(&text) + } + + pub fn from_json(json: &str) -> Result { + let cfg: Self = serde_json::from_str(json)?; + if cfg.service.is_empty() || cfg.manifest_url.is_empty() { + return Err(UpdateConfigError::Invalid( + "service and manifest_url are required".to_owned(), + )); + } + if let Some(t) = &cfg.versioned_manifest_url + && !t.contains("{version}") + { + return Err(UpdateConfigError::Invalid( + "versioned_manifest_url must contain {version}".to_owned(), + )); + } + if cfg.health_timeout_secs == 0 { + return Err(UpdateConfigError::Invalid( + "health_timeout_secs must be positive".to_owned(), + )); + } + Ok(cfg) + } + + /// The manifest to read for `version` (the newest release if `None`). + pub fn manifest_url_for(&self, version: Option<&str>) -> String { + match (version, &self.versioned_manifest_url) { + (Some(v), Some(template)) => template.replace("{version}", v), + _ => self.manifest_url.clone(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const BASE: &str = r#"{"manifest_url":"https://m.example.test/latest.json","service":"borg-drone.service","health_file":"/var/lib/borg/health.json","state_dir":"/var/lib/borg"}"#; + + #[test] + fn defaults_and_url_selection() { + let c = UpdateConfig::from_json(BASE).unwrap(); + assert_eq!( + (c.health_timeout_secs, c.bin.as_path()), + (120, Path::new("/usr/local/bin/borg")) + ); + assert_eq!( + c.manifest_url_for(None), + "https://m.example.test/latest.json" + ); + assert_eq!( + c.manifest_url_for(Some("1.2.3")), + "https://m.example.test/latest.json", + "no template: newest" + ); + + let t = BASE.replace( + '}', + r#","versioned_manifest_url":"https://m.example.test/v{version}.json"}"#, + ); + let c = UpdateConfig::from_json(&t).unwrap(); + assert_eq!( + c.manifest_url_for(Some("1.2.3")), + "https://m.example.test/v1.2.3.json" + ); + assert_eq!( + c.manifest_url_for(None), + "https://m.example.test/latest.json" + ); + } + + #[test] + fn rejects_bad_config() { + for bad in [ + BASE.replace( + '}', + r#","versioned_manifest_url":"https://m.example.test/v.json"}"#, + ), + BASE.replace('}', r#","health_timeout_secs":0}"#), + BASE.replace('}', r#","bogus":1}"#), + BASE.replace("borg-drone.service", ""), + ] { + assert!(UpdateConfig::from_json(&bad).is_err(), "{bad}"); + } + } + + #[test] + fn default_urls_point_at_the_public_release_manifest() { + assert!( + DEFAULT_MANIFEST_URL.starts_with("https://") + && DEFAULT_VERSIONED_MANIFEST_URL.contains("{version}") + ); + } +} diff --git a/src/update/health.rs b/src/update/health.rs new file mode 100644 index 0000000..a2cc603 --- /dev/null +++ b/src/update/health.rs @@ -0,0 +1,65 @@ +//! Health evidence after an update. +//! +//! A running daemon writes a small record after every successful heartbeat +//! (drone) or reconciliation pass (queen). The updater polls it to learn that +//! the new binary came up and is working; no new endpoint is needed and the +//! updater needs no credentials. + +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct HealthRecord { + pub version: String, + /// Unix seconds of the last success. + pub at: i64, + pub pid: u32, +} + +/// Write the record atomically (temporary file, then rename). +pub fn write(path: &Path, record: &HealthRecord) -> io::Result<()> { + let tmp = path.with_extension("tmp"); + fs::write(&tmp, serde_json::to_vec(record).map_err(io::Error::other)?)?; + fs::rename(&tmp, path) +} + +pub fn read(path: &Path) -> Option { + serde_json::from_slice(&fs::read(path).ok()?).ok() +} + +pub trait HealthProbe { + /// The newest record the daemon has written, if any. + fn current(&self) -> Option; +} + +pub struct FileHealth(pub PathBuf); + +impl HealthProbe for FileHealth { + fn current(&self) -> Option { + read(&self.0) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn records_round_trip_and_garbage_reads_as_none() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("health.json"); + assert_eq!(FileHealth(path.clone()).current(), None); + let rec = HealthRecord { + version: "1.2.3".into(), + at: 100, + pid: 7, + }; + write(&path, &rec).unwrap(); + assert_eq!(FileHealth(path.clone()).current(), Some(rec)); + fs::write(&path, "not json").unwrap(); + assert_eq!(read(&path), None); + } +} diff --git a/src/update/manifest.rs b/src/update/manifest.rs new file mode 100644 index 0000000..fd336c5 --- /dev/null +++ b/src/update/manifest.rs @@ -0,0 +1,260 @@ +//! The release manifest: what version is current and where to get it. +//! +//! A manifest is a small JSON document published next to a release (for +//! example as a GitHub release asset of the public borg repository): +//! +//! ```json +//! { +//! "schema": 1, +//! "version": "0.5.0", +//! "targets": { +//! "linux-x64-musl": { "url": "https://.../borg-linux-x64-musl", "sha256": "<64 hex>" }, +//! "linux-arm64-musl": { "url": "https://.../borg-linux-arm64-musl", "sha256": "<64 hex>" }, +//! "darwin-arm64": { "url": "https://.../borg-darwin-arm64", "sha256": "<64 hex>" } +//! }, +//! "runner": { "version": "2.337.0", "sha256": { "linux-x64": "<64 hex>", "osx-arm64": "<64 hex>" } } +//! } +//! ``` +//! +//! sha256 pinning against the manifest is mandatory: an artifact is never +//! installed unless its digest matches. The manifest itself is trusted by the +//! transport (HTTPS to the configured URL); signature verification is not +//! implemented. + +use std::collections::BTreeMap; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; + +use crate::drone::pins::RunnerPin; + +pub const SCHEMA: u32 = 1; + +#[derive(Debug, thiserror::Error)] +pub enum ManifestError { + #[error("fetch manifest {url}: {source}")] + Fetch { url: String, source: reqwest::Error }, + #[error("manifest {url} returned {status}")] + Status { url: String, status: u16 }, + #[error("parse manifest: {0}")] + Parse(#[from] serde_json::Error), + #[error("invalid manifest: {0}")] + Invalid(String), + #[error("manifest URL must be https:// (set allow_insecure for a trusted network): {0}")] + Insecure(String), +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Artifact { + pub url: String, + pub sha256: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RunnerRelease { + pub version: String, + /// actions-runner asset platform (`linux-x64`, `linux-arm64`, `osx-x64`, + /// `osx-arm64`) to the SHA-256 of that tarball. + pub sha256: BTreeMap, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Manifest { + pub schema: u32, + pub version: String, + pub targets: BTreeMap, + #[serde(default)] + pub runner: Option, +} + +fn is_sha256(s: &str) -> bool { + s.len() == 64 + && s.bytes() + .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) +} + +/// A version is a short, printable string without path or shell characters. +pub fn valid_version(v: &str) -> bool { + !v.is_empty() + && v.len() <= 64 + && v.bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'-' | b'+' | b'_')) +} + +impl Manifest { + pub fn parse(json: &str) -> Result { + let m: Self = serde_json::from_str(json)?; + m.validate()?; + Ok(m) + } + + pub fn validate(&self) -> Result<(), ManifestError> { + let bad = |m: String| Err(ManifestError::Invalid(m)); + if self.schema != SCHEMA { + return bad(format!( + "unsupported schema {} (expected {SCHEMA})", + self.schema + )); + } + if !valid_version(&self.version) { + return bad(format!("bad version {:?}", self.version)); + } + if self.targets.is_empty() { + return bad("no targets".to_owned()); + } + for (name, a) in &self.targets { + if !is_sha256(&a.sha256) { + return bad(format!( + "target {name}: sha256 must be 64 lowercase hex characters" + )); + } + if !(a.url.starts_with("https://") || a.url.starts_with("http://")) { + return bad(format!("target {name}: url must be http(s)")); + } + } + if let Some(r) = &self.runner { + if !valid_version(&r.version) { + return bad(format!("bad runner version {:?}", r.version)); + } + if let Some((p, _)) = r.sha256.iter().find(|(_, s)| !is_sha256(s)) { + return bad(format!( + "runner {p}: sha256 must be 64 lowercase hex characters" + )); + } + } + Ok(()) + } + + pub fn artifact_for(&self, target: &str) -> Option<&Artifact> { + self.targets.get(target) + } + + /// The actions-runner to install on a host, if the manifest pins one for it. + pub fn runner_pin_for(&self, os: &str, arch: &str) -> Option { + let release = self.runner.as_ref()?; + let platform = crate::drone::pins::pin_for(os, arch)?.platform; + let sha256 = release.sha256.get(&platform)?.clone(); + Some(RunnerPin { + version: release.version.clone(), + platform, + sha256, + }) + } +} + +/// The manifest key for a host's borg binary. Linux gets the static musl +/// build, which runs on any glibc. +pub fn host_target(os: &str, arch: &str) -> Option<&'static str> { + match (os, arch) { + ("linux", "x86_64") => Some("linux-x64-musl"), + ("linux", "aarch64") => Some("linux-arm64-musl"), + ("macos", "aarch64") => Some("darwin-arm64"), + _ => None, + } +} + +/// Fetch and validate a manifest. `https://` is required unless `allow_insecure`. +pub async fn fetch(url: &str, allow_insecure: bool) -> Result { + let scheme_ok = url.starts_with("https://") || (allow_insecure && url.starts_with("http://")); + if !scheme_ok { + return Err(ManifestError::Insecure(url.to_owned())); + } + let client = reqwest::Client::builder() + .user_agent("borg-update") + .timeout(Duration::from_secs(30)) + .build() + .map_err(|source| ManifestError::Fetch { + url: url.to_owned(), + source, + })?; + let resp = client + .get(url) + .send() + .await + .map_err(|source| ManifestError::Fetch { + url: url.to_owned(), + source, + })?; + if !resp.status().is_success() { + return Err(ManifestError::Status { + url: url.to_owned(), + status: resp.status().as_u16(), + }); + } + let text = resp.text().await.map_err(|source| ManifestError::Fetch { + url: url.to_owned(), + source, + })?; + Manifest::parse(&text) +} + +#[cfg(test)] +mod tests { + use super::*; + + const SHA: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + + fn json(version: &str, sha: &str) -> String { + format!( + r#"{{"schema":1,"version":"{version}","targets":{{"linux-x64-musl":{{"url":"https://example.test/borg","sha256":"{sha}"}}}}, + "runner":{{"version":"2.337.0","sha256":{{"linux-x64":"{SHA}"}}}}}}"# + ) + } + + #[test] + fn parses_and_selects_targets_and_the_runner_pin() { + let m = Manifest::parse(&json("0.5.0", SHA)).unwrap(); + assert_eq!(m.version, "0.5.0"); + assert_eq!(m.artifact_for("linux-x64-musl").unwrap().sha256, SHA); + assert!(m.artifact_for("darwin-arm64").is_none()); + let pin = m.runner_pin_for("linux", "x86_64").unwrap(); + assert_eq!( + ( + pin.version.as_str(), + pin.platform.as_str(), + pin.sha256.as_str() + ), + ("2.337.0", "linux-x64", SHA) + ); + assert!( + m.runner_pin_for("linux", "aarch64").is_none(), + "no digest for that platform" + ); + assert!(m.runner_pin_for("plan9", "mips").is_none()); + } + + #[test] + fn rejects_malformed_manifests() { + for bad in [ + json("0.5.0", "NOTHEX"), + json("0.5.0", &SHA.to_uppercase()), + json("../evil", SHA), + json("", SHA), + json("0.5.0", SHA).replace("\"schema\":1", "\"schema\":2"), + json("0.5.0", SHA).replace("https://example.test/borg", "ftp://example.test/borg"), + json("0.5.0", SHA).replace("\"targets\":{", "\"extra\":1,\"targets\":{"), + r#"{"schema":1,"version":"1","targets":{}}"#.to_owned(), + ] { + assert!(Manifest::parse(&bad).is_err(), "{bad}"); + } + } + + #[test] + fn host_targets() { + assert_eq!(host_target("linux", "x86_64"), Some("linux-x64-musl")); + assert_eq!(host_target("linux", "aarch64"), Some("linux-arm64-musl")); + assert_eq!(host_target("macos", "aarch64"), Some("darwin-arm64")); + assert_eq!(host_target("macos", "x86_64"), None); + } + + #[tokio::test] + async fn plain_http_needs_the_override() { + assert!(matches!( + fetch("http://example.test/m.json", false).await, + Err(ManifestError::Insecure(_)) + )); + } +} diff --git a/src/update/mod.rs b/src/update/mod.rs new file mode 100644 index 0000000..ce379b1 --- /dev/null +++ b/src/update/mod.rs @@ -0,0 +1,266 @@ +//! Self-update for the drone and the queen: `borg update`. +//! +//! The binary comes from a release [`manifest`] whose SHA-256 pins are +//! mandatory. The update drains the service (the service manager's stop is the +//! drain), swaps the binary atomically keeping `borg.prev`, restarts it, waits +//! for health evidence from the new version and rolls back to `borg.prev` if it +//! does not appear. The same machinery installs a newer pinned actions-runner, +//! on its own with `--runner-only` or together with a binary update. + +pub mod cli; +pub mod config; +pub mod health; +pub mod manifest; +pub mod runner; +pub mod service; +pub mod swap; + +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; + +use crate::drone::exec::Exec; +use health::HealthProbe; +use manifest::Manifest; +use service::ServiceManager; + +#[derive(Debug, thiserror::Error)] +pub enum UpdateError { + #[error(transparent)] + Manifest(#[from] manifest::ManifestError), + #[error(transparent)] + Swap(#[from] swap::SwapError), + #[error(transparent)] + Runner(#[from] runner::RunnerError), + #[error("the manifest is version {manifest}, not the requested {requested}")] + VersionMismatch { manifest: String, requested: String }, + #[error("the manifest has no binary for {0}")] + NoTarget(String), + #[error("this platform ({0}) has no released borg binary")] + UnsupportedHost(String), + #[error("could not drain and stop the service; nothing was changed: {0}")] + Drain(String), + #[error("could not restart the service and could not roll back: {0}")] + Unrecoverable(String), +} + +/// What an update run did (or, with `check`, would do). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "result", rename_all = "snake_case")] +pub enum Outcome { + UpToDate { + version: String, + }, + Available { + current: String, + target: String, + runner: Option, + }, + Updated { + from: String, + to: String, + runner_updated: bool, + }, + RunnerUpdated { + version: String, + }, + RolledBack { + from: String, + to: String, + reason: String, + }, +} + +pub struct Options<'a> { + pub bin: &'a Path, + /// The version of the installed binary (the one running this update). + pub current_version: &'a str, + pub requested_version: Option<&'a str>, + pub check: bool, + pub runner_only: bool, + pub allow_insecure: bool, + pub health_timeout: Duration, + pub poll: Duration, + pub host_os: &'a str, + pub host_arch: &'a str, + /// Where the pinned runner lives and who owns it; `None` skips runner updates. + pub runner_dir: Option<&'a Path>, + pub runner_owner: Option<&'a str>, +} + +/// Run an update against an already fetched manifest. +pub async fn run( + opts: &Options<'_>, + manifest: &Manifest, + service: &dyn ServiceManager, + health: &dyn HealthProbe, + exec: &dyn Exec, +) -> Result { + if let Some(requested) = opts.requested_version + && requested != manifest.version + { + return Err(UpdateError::VersionMismatch { + manifest: manifest.version.clone(), + requested: requested.to_owned(), + }); + } + let runner_pin = match opts.runner_dir { + Some(dir) => manifest + .runner_pin_for(opts.host_os, opts.host_arch) + .filter(|pin| !runner::is_installed(dir, pin)), + None => None, + }; + let binary_needed = !opts.runner_only && manifest.version != opts.current_version; + + if opts.check { + return Ok(if binary_needed || runner_pin.is_some() { + Outcome::Available { + current: opts.current_version.to_owned(), + target: manifest.version.clone(), + runner: runner_pin.map(|p| p.version), + } + } else { + Outcome::UpToDate { + version: opts.current_version.to_owned(), + } + }); + } + if !binary_needed && runner_pin.is_none() { + return Ok(Outcome::UpToDate { + version: opts.current_version.to_owned(), + }); + } + + // Everything that can fail without touching the running service comes first. + let new_bin = swap::new_path(opts.bin); + if binary_needed { + let target = manifest::host_target(opts.host_os, opts.host_arch).ok_or_else(|| { + UpdateError::UnsupportedHost(format!("{}-{}", opts.host_os, opts.host_arch)) + })?; + let artifact = manifest + .artifact_for(target) + .ok_or_else(|| UpdateError::NoTarget(target.to_owned()))?; + swap::download_verified( + &artifact.url, + &artifact.sha256, + &new_bin, + opts.allow_insecure, + ) + .await?; + } + let mut runner_updated = false; + if let (Some(pin), Some(dir)) = (&runner_pin, opts.runner_dir) { + if let Err(e) = runner::install(exec, pin, dir, opts.runner_owner) { + let _ = std::fs::remove_file(&new_bin); + return Err(e.into()); + } + runner_updated = true; + } + if !binary_needed { + return Ok(Outcome::RunnerUpdated { + version: runner_pin.map(|p| p.version).unwrap_or_default(), + }); + } + + // Drain: no new leases, in-flight jobs finish, then the service exits. + if let Err(e) = service.stop() { + let _ = std::fs::remove_file(&new_bin); + return Err(UpdateError::Drain(e)); + } + swap::install(opts.bin, &new_bin)?; + let started = service.start(); + let healthy = started.is_ok() + && wait_healthy(health, &manifest.version, opts.health_timeout, opts.poll).await; + if healthy { + return Ok(Outcome::Updated { + from: opts.current_version.to_owned(), + to: manifest.version.clone(), + runner_updated, + }); + } + + let reason = match started { + Err(e) => format!("the new version did not start: {e}"), + Ok(()) => format!( + "no health evidence from {} within {}s", + manifest.version, + opts.health_timeout.as_secs() + ), + }; + // Roll back: stop the failed version, restore the previous binary, restart. + let _ = service.stop(); + swap::rollback(opts.bin).map_err(|e| UpdateError::Unrecoverable(e.to_string()))?; + service.start().map_err(UpdateError::Unrecoverable)?; + Ok(Outcome::RolledBack { + from: opts.current_version.to_owned(), + to: manifest.version.clone(), + reason, + }) +} + +/// Wait until the daemon reports `version` as healthy, or the timeout passes. +async fn wait_healthy( + health: &dyn HealthProbe, + version: &str, + timeout: Duration, + poll: Duration, +) -> bool { + let deadline = tokio::time::Instant::now() + timeout; + loop { + if health.current().is_some_and(|r| r.version == version) { + return true; + } + if tokio::time::Instant::now() >= deadline { + return false; + } + tokio::time::sleep(poll).await; + } +} + +/// The path the drone writes an update request to and the updater reads. +pub fn request_path(state_dir: &Path) -> PathBuf { + state_dir.join("update-request") +} + +/// Where the updater leaves its result for the drone to report. +pub fn result_path(state_dir: &Path) -> PathBuf { + state_dir.join("update-result.json") +} + +/// A request for the privileged updater, written by the drone. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Request { + pub version: String, +} + +/// What the updater leaves for the drone to report to the queen. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct UpdateResult { + pub target: String, + /// `installed`, `runner_updated`, `up_to_date`, `rolled_back` or `failed`. + pub state: String, + pub detail: String, + pub at: i64, +} + +impl UpdateResult { + pub fn from_outcome(target: &str, outcome: &Result, at: i64) -> Self { + let (state, detail) = match outcome { + Ok(Outcome::Updated { to, .. }) => ("installed", format!("updated to {to}")), + Ok(Outcome::RunnerUpdated { version }) => { + ("runner_updated", format!("runner {version}")) + } + Ok(Outcome::UpToDate { version }) => ("up_to_date", version.clone()), + Ok(Outcome::Available { target, .. }) => ("up_to_date", format!("available {target}")), + Ok(Outcome::RolledBack { reason, .. }) => ("rolled_back", reason.clone()), + Err(e) => ("failed", e.to_string()), + }; + Self { + target: target.to_owned(), + state: state.to_owned(), + detail, + at, + } + } +} diff --git a/src/update/runner.rs b/src/update/runner.rs new file mode 100644 index 0000000..2271999 --- /dev/null +++ b/src/update/runner.rs @@ -0,0 +1,185 @@ +//! Installing the pinned actions-runner from a manifest pin: download, +//! verify the SHA-256, unpack, swap the directory, hand it to the job user and +//! make it read-only (leases hard-link its files). +//! +//! Leases clone the runner directory when they start, so replacing it while +//! jobs run is safe: running leases keep their own links. + +use std::fs; +use std::path::{Path, PathBuf}; + +use sha2::{Digest, Sha256}; + +use crate::drone::exec::Exec; +use crate::drone::pins::RunnerPin; + +#[derive(Debug, thiserror::Error)] +pub enum RunnerError { + #[error("`{0}` failed")] + Command(String), + #[error("checksum mismatch for {file}: expected {expected}, got {actual}")] + Checksum { + file: String, + expected: String, + actual: String, + }, + #[error("{path}: {source}")] + Io { + path: PathBuf, + source: std::io::Error, + }, +} + +fn io(path: &Path) -> impl FnOnce(std::io::Error) -> RunnerError + '_ { + move |source| RunnerError::Io { + path: path.to_owned(), + source, + } +} + +fn must(exec: &dyn Exec, program: &str, args: &[&str]) -> Result<(), RunnerError> { + match exec.run(program, args) { + Ok(o) if o.success => Ok(()), + _ => Err(RunnerError::Command(format!( + "{program} {}", + args.join(" ") + ))), + } +} + +/// What the pin installs, as recorded in `/.borg-runner-version`. +pub fn marker(pin: &RunnerPin) -> String { + format!("{} {}\n", pin.version, pin.sha256) +} + +/// Whether `runner_dir` already holds exactly this pin. +pub fn is_installed(runner_dir: &Path, pin: &RunnerPin) -> bool { + fs::read_to_string(runner_dir.join(".borg-runner-version")).is_ok_and(|m| m == marker(pin)) +} + +/// Install `pin` into `runner_dir`, owned by `owner` (when given). +pub fn install( + exec: &dyn Exec, + pin: &RunnerPin, + runner_dir: &Path, + owner: Option<&str>, +) -> Result<(), RunnerError> { + let parent = runner_dir.parent().unwrap_or(Path::new(".")); + let staging = parent.join("runner.new"); + let tarball = parent.join(pin.filename()); + let _ = fs::remove_dir_all(&staging); + fs::create_dir_all(&staging).map_err(io(&staging))?; + + must( + exec, + "curl", + &["-fsSL", "-o", &tarball.to_string_lossy(), &pin.url()], + )?; + let actual = hex::encode(Sha256::digest(fs::read(&tarball).map_err(io(&tarball))?)); + if actual != pin.sha256 { + let _ = fs::remove_file(&tarball); + let _ = fs::remove_dir_all(&staging); + return Err(RunnerError::Checksum { + file: pin.filename(), + expected: pin.sha256.clone(), + actual, + }); + } + must( + exec, + "tar", + &[ + "-xzf", + &tarball.to_string_lossy(), + "-C", + &staging.to_string_lossy(), + ], + )?; + let _ = fs::remove_file(&tarball); + fs::write(staging.join(".borg-runner-version"), marker(pin)).map_err(io(&staging))?; + + if runner_dir.exists() { + // The installed tree is read-only; restore write access to delete it. + must(exec, "chmod", &["-R", "u+w", &runner_dir.to_string_lossy()])?; + fs::remove_dir_all(runner_dir).map_err(io(runner_dir))?; + } + fs::rename(&staging, runner_dir).map_err(io(runner_dir))?; + if let Some(owner) = owner { + must( + exec, + "chown", + &[ + "-R", + &format!("{owner}:{owner}"), + &runner_dir.to_string_lossy(), + ], + )?; + } + must(exec, "chmod", &["-R", "a-w", &runner_dir.to_string_lossy()]) +} + +#[cfg(test)] +mod tests { + use std::io; + + use super::*; + use crate::drone::cleanup::testing::FakeExec; + use crate::drone::exec::Output; + + const BODY: &[u8] = b"fake runner tarball"; + + struct Hooked(FakeExec); + impl Exec for Hooked { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + match program { + "curl" => fs::write(args[2], BODY)?, + "tar" => fs::write(Path::new(args[3]).join("run.sh"), "#!/bin/sh\n")?, + _ => {} + } + self.0.run(program, args) + } + } + + fn pin(sha: &str, version: &str) -> RunnerPin { + RunnerPin { + version: version.into(), + platform: "linux-x64".into(), + sha256: sha.into(), + } + } + + #[test] + fn installs_verified_runners_and_is_idempotent_by_marker() { + let dir = tempfile::tempdir().unwrap(); + let runner = dir.path().join("runner"); + let p = pin(&hex::encode(Sha256::digest(BODY)), "2.337.0"); + assert!(!is_installed(&runner, &p)); + let x = Hooked(FakeExec::default()); + install(&x, &p, &runner, Some("borg-job")).unwrap(); + assert!(is_installed(&runner, &p)); + assert!(runner.join("run.sh").exists()); + let calls = x.0.calls(); + assert!( + calls + .iter() + .any(|c| c.starts_with("chown -R borg-job:borg-job")), + "{calls:?}" + ); + assert!(calls.last().unwrap().starts_with("chmod -R a-w")); + + // A newer pin replaces it. + let newer = pin(&p.sha256, "2.338.0"); + install(&x, &newer, &runner, None).unwrap(); + assert!(is_installed(&runner, &newer) && !is_installed(&runner, &p)); + } + + #[test] + fn a_bad_checksum_installs_nothing() { + let dir = tempfile::tempdir().unwrap(); + let runner = dir.path().join("runner"); + let x = Hooked(FakeExec::default()); + let err = install(&x, &pin(&"0".repeat(64), "2.337.0"), &runner, None).unwrap_err(); + assert!(matches!(err, RunnerError::Checksum { .. }), "{err}"); + assert!(!runner.exists() && !dir.path().join("runner.new").exists()); + } +} diff --git a/src/update/service.rs b/src/update/service.rs new file mode 100644 index 0000000..c738f41 --- /dev/null +++ b/src/update/service.rs @@ -0,0 +1,107 @@ +//! The service manager the update restarts the daemon through. +//! +//! Stopping is the drain: the unit sends `SIGTERM` to the daemon, which stops +//! taking leases and lets in-flight jobs finish (up to its drain timeout, which +//! the unit's stop timeout exceeds), and `stop` returns only once it has exited. + +use crate::drone::exec::Exec; + +pub trait ServiceManager { + /// Stop the service, returning only after it has fully stopped. + fn stop(&self) -> Result<(), String>; + fn start(&self) -> Result<(), String>; +} + +fn run(exec: &dyn Exec, program: &str, args: &[&str]) -> Result<(), String> { + match exec.run(program, args) { + Ok(out) if out.success => Ok(()), + Ok(_) => Err(format!("`{program} {}` failed", args.join(" "))), + Err(e) => Err(format!("{program}: {e}")), + } +} + +pub struct Systemd<'a> { + pub unit: String, + pub exec: &'a dyn Exec, +} + +impl ServiceManager for Systemd<'_> { + fn stop(&self) -> Result<(), String> { + run(self.exec, "systemctl", &["stop", &self.unit]) + } + + fn start(&self) -> Result<(), String> { + run(self.exec, "systemctl", &["start", &self.unit]) + } +} + +pub struct Launchd<'a> { + pub label: String, + pub plist: String, + pub exec: &'a dyn Exec, +} + +impl ServiceManager for Launchd<'_> { + fn stop(&self) -> Result<(), String> { + run( + self.exec, + "launchctl", + &["bootout", &format!("system/{}", self.label)], + ) + } + + fn start(&self) -> Result<(), String> { + run( + self.exec, + "launchctl", + &["bootstrap", "system", &self.plist], + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::drone::cleanup::testing::FakeExec; + + #[test] + fn systemd_and_launchd_issue_the_expected_commands() { + let x = FakeExec::default(); + let s = Systemd { + unit: "borg-drone.service".into(), + exec: &x, + }; + s.stop().unwrap(); + s.start().unwrap(); + let l = Launchd { + label: "org.borg.drone".into(), + plist: "/Library/LaunchDaemons/org.borg.drone.plist".into(), + exec: &x, + }; + l.stop().unwrap(); + l.start().unwrap(); + assert_eq!( + x.calls(), + [ + "systemctl stop borg-drone.service", + "systemctl start borg-drone.service", + "launchctl bootout system/org.borg.drone", + "launchctl bootstrap system /Library/LaunchDaemons/org.borg.drone.plist", + ] + ); + } + + #[test] + fn a_failing_stop_is_an_error() { + let x = FakeExec::default(); + x.fail("systemctl stop u"); + assert!( + Systemd { + unit: "u".into(), + exec: &x + } + .stop() + .is_err() + ); + } +} diff --git a/src/update/swap.rs b/src/update/swap.rs new file mode 100644 index 0000000..6da0056 --- /dev/null +++ b/src/update/swap.rs @@ -0,0 +1,213 @@ +//! Downloading a verified binary and swapping it in atomically. +//! +//! Streaming the body into a temporary file next to the destination and +//! renaming it over the target is ported from flotilla's `put_file` +//! (`crates/flotillad/src/server.rs`): a running binary keeps its old inode and +//! readers never see a partial file. The digest is checked before anything is +//! renamed into place. + +use std::fs; +use std::io; +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use sha2::{Digest, Sha256}; +use tokio::io::AsyncWriteExt; + +#[derive(Debug, thiserror::Error)] +pub enum SwapError { + #[error("download {url}: {source}")] + Download { url: String, source: reqwest::Error }, + #[error("download {url} returned {status}")] + Status { url: String, status: u16 }, + #[error("checksum mismatch for {url}: expected {expected}, got {actual}")] + Checksum { + url: String, + expected: String, + actual: String, + }, + #[error("{path}: {source}")] + Io { path: PathBuf, source: io::Error }, + #[error("there is no previous binary to roll back to at {0}")] + NoPrevious(PathBuf), +} + +fn io_err(path: &Path) -> impl FnOnce(io::Error) -> SwapError + '_ { + move |source| SwapError::Io { + path: path.to_owned(), + source, + } +} + +/// `/borg` -> `/borg.new` +pub fn new_path(bin: &Path) -> PathBuf { + sibling(bin, "new") +} + +/// `/borg` -> `/borg.prev` +pub fn prev_path(bin: &Path) -> PathBuf { + sibling(bin, "prev") +} + +fn sibling(bin: &Path, suffix: &str) -> PathBuf { + let name = bin + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_default(); + bin.with_file_name(format!("{name}.{suffix}")) +} + +/// Download `url` to `dest` (mode 0755), verifying its SHA-256 against +/// `expected_sha256`. Nothing is left at `dest` unless the digest matches. +pub async fn download_verified( + url: &str, + expected_sha256: &str, + dest: &Path, + allow_insecure: bool, +) -> Result<(), SwapError> { + let scheme_ok = url.starts_with("https://") || (allow_insecure && url.starts_with("http://")); + if !scheme_ok { + return Err(SwapError::Status { + url: url.to_owned(), + status: 0, + }); + } + let client = reqwest::Client::builder() + .user_agent("borg-update") + .timeout(Duration::from_secs(900)) + .build() + .map_err(|source| SwapError::Download { + url: url.to_owned(), + source, + })?; + let mut resp = client + .get(url) + .send() + .await + .map_err(|source| SwapError::Download { + url: url.to_owned(), + source, + })?; + if !resp.status().is_success() { + return Err(SwapError::Status { + url: url.to_owned(), + status: resp.status().as_u16(), + }); + } + let tmp = sibling(dest, "download"); + let result: Result = async { + let mut file = tokio::fs::File::create(&tmp).await.map_err(io_err(&tmp))?; + let mut hasher = Sha256::new(); + while let Some(chunk) = resp.chunk().await.map_err(|source| SwapError::Download { + url: url.to_owned(), + source, + })? { + hasher.update(&chunk); + file.write_all(&chunk).await.map_err(io_err(&tmp))?; + } + file.flush().await.map_err(io_err(&tmp))?; + Ok(hex::encode(hasher.finalize())) + } + .await; + let actual = match result { + Ok(actual) => actual, + Err(e) => { + let _ = fs::remove_file(&tmp); + return Err(e); + } + }; + if actual != expected_sha256 { + let _ = fs::remove_file(&tmp); + return Err(SwapError::Checksum { + url: url.to_owned(), + expected: expected_sha256.to_owned(), + actual, + }); + } + fs::set_permissions(&tmp, fs::Permissions::from_mode(0o755)).map_err(io_err(&tmp))?; + fs::rename(&tmp, dest).map_err(io_err(dest)) +} + +/// Make `new` the installed binary, keeping the current one as `.prev`. +/// Both steps are renames/links in one directory, so the installed path always +/// names a complete executable. +pub fn install(bin: &Path, new: &Path) -> Result<(), SwapError> { + let prev = prev_path(bin); + if bin.exists() { + let _ = fs::remove_file(&prev); + fs::hard_link(bin, &prev) + .or_else(|_| fs::copy(bin, &prev).map(|_| ())) + .map_err(io_err(&prev))?; + } + fs::rename(new, bin).map_err(io_err(bin)) +} + +/// Put `.prev` back as the installed binary. +pub fn rollback(bin: &Path) -> Result<(), SwapError> { + let prev = prev_path(bin); + if !prev.exists() { + return Err(SwapError::NoPrevious(prev)); + } + fs::rename(&prev, bin).map_err(io_err(bin)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn install_keeps_the_previous_binary_and_rollback_restores_it() { + let dir = tempfile::tempdir().unwrap(); + let bin = dir.path().join("borg"); + fs::write(&bin, "v1").unwrap(); + let new = new_path(&bin); + fs::write(&new, "v2").unwrap(); + + install(&bin, &new).unwrap(); + assert_eq!(fs::read_to_string(&bin).unwrap(), "v2"); + assert_eq!(fs::read_to_string(prev_path(&bin)).unwrap(), "v1"); + assert!(!new.exists()); + + rollback(&bin).unwrap(); + assert_eq!(fs::read_to_string(&bin).unwrap(), "v1"); + assert!( + matches!(rollback(&bin), Err(SwapError::NoPrevious(_))), + "nothing further to restore" + ); + } + + #[test] + fn a_second_install_replaces_the_old_prev() { + let dir = tempfile::tempdir().unwrap(); + let bin = dir.path().join("borg"); + fs::write(&bin, "v1").unwrap(); + for (version, prev) in [("v2", "v1"), ("v3", "v2")] { + fs::write(new_path(&bin), version).unwrap(); + install(&bin, &new_path(&bin)).unwrap(); + assert_eq!(fs::read_to_string(prev_path(&bin)).unwrap(), prev); + } + } + + #[test] + fn installing_over_nothing_works_for_a_first_install() { + let dir = tempfile::tempdir().unwrap(); + let bin = dir.path().join("borg"); + fs::write(new_path(&bin), "v1").unwrap(); + install(&bin, &new_path(&bin)).unwrap(); + assert_eq!(fs::read_to_string(&bin).unwrap(), "v1"); + assert!(!prev_path(&bin).exists()); + } + + #[test] + fn sibling_paths() { + assert_eq!( + new_path(Path::new("/usr/local/bin/borg")), + Path::new("/usr/local/bin/borg.new") + ); + assert_eq!( + prev_path(Path::new("/usr/local/bin/borg")), + Path::new("/usr/local/bin/borg.prev") + ); + } +} diff --git a/src/wake.rs b/src/wake.rs new file mode 100644 index 0000000..95e9340 --- /dev/null +++ b/src/wake.rs @@ -0,0 +1,533 @@ +//! Wake-on-LAN for idle hosts, driven by unmet demand. +//! +//! When demand for a class stays unmet for `after_secs` (the fleet has no +//! capacity for it), the queen sends a magic packet to a configured host that +//! can serve the class and is not awake, rate limited three ways: at most one +//! wake per `min_interval_secs` anywhere, at most `max_per_hour` per hour, and +//! at most one per host per `host_cooldown_secs` (a host that was woken and +//! never came up is retried only after that). A host counts as awake when a +//! drone of the same `name` is fresh. +//! +//! The planner ([`WakePlanner`]) is pure and takes the clock as an argument; the +//! sender is a trait ([`WakeSender`]) so tests use a fake and production sends +//! UDP. The packet code ([`parse_mac`], [`magic_packet`], the broadcast +//! address math and the port pair 9 and 7) is ported from the wake sender in +//! haasonsaas/flotilla (`crates/flotilla-core/src/api.rs` and +//! `crates/flotillad/src/server.rs`, MIT, same author), without its interface +//! discovery: broadcast targets are configuration here. +//! +//! The magic packet must originate on the target's LAN; run the queen on (or +//! route directed broadcasts to) the network the hosts sleep on. + +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::future::Future; +use std::net::{Ipv4Addr, SocketAddrV4}; +use std::pin::Pin; + +use serde::{Deserialize, Serialize}; + +use crate::classes::Class; + +pub type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// Parse `aa:bb:cc:dd:ee:ff`, `aa-bb-...` or `aabbccddeeff`. +/// (Ported from flotilla-core `parse_mac`.) +pub fn parse_mac(s: &str) -> Option<[u8; 6]> { + let hex: String = s.chars().filter(char::is_ascii_hexdigit).collect(); + if hex.len() != 12 { + return None; + } + let mut out = [0u8; 6]; + for (i, b) in out.iter_mut().enumerate() { + *b = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16).ok()?; + } + Some(out) +} + +/// A wake-on-LAN magic packet: six 0xFF bytes, then the MAC sixteen times. +/// (Ported from flotilla-core `magic_packet`.) +pub fn magic_packet(mac: [u8; 6]) -> Vec { + let mut p = vec![0xFFu8; 6]; + for _ in 0..16 { + p.extend_from_slice(&mac); + } + p +} + +/// The directed broadcast address of `ip/prefix`. (The mask math is ported +/// from the interface loop in flotillad `post_wake`.) +pub fn directed_broadcast(ip: Ipv4Addr, prefix: u8) -> Ipv4Addr { + let prefix = u32::from(prefix.min(32)); + let mask = if prefix == 0 { + 0 + } else { + u32::MAX << (32 - prefix) + }; + Ipv4Addr::from(u32::from(ip) | !mask) +} + +/// A broadcast target from configuration: a plain IPv4 address (used as is) or +/// `a.b.c.d/len` (its directed broadcast). +pub fn parse_target(s: &str) -> Option { + match s.split_once('/') { + Some((ip, len)) => Some(directed_broadcast(ip.parse().ok()?, len.parse().ok()?)), + None => s.parse().ok(), + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct WakeHost { + /// Drone name this machine runs under; the host counts as awake while that + /// drone is fresh. Also the label in logs and metrics. + pub name: String, + pub mac: String, + /// Classes the host can serve; all when empty. + #[serde(default)] + pub classes: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct WakeSettings { + pub enabled: bool, + pub hosts: Vec, + /// Demand must stay unmet this long before anything is woken. + pub after_secs: u64, + pub min_interval_secs: u64, + pub max_per_hour: u32, + pub host_cooldown_secs: u64, + /// IPv4 addresses or `addr/len` networks to send the packet to. + pub broadcast: Vec, + pub ports: Vec, +} + +impl Default for WakeSettings { + fn default() -> Self { + Self { + enabled: false, + hosts: Vec::new(), + after_secs: 300, + min_interval_secs: 300, + max_per_hour: 6, + host_cooldown_secs: 900, + broadcast: vec!["255.255.255.255".into()], + ports: vec![9, 7], + } + } +} + +impl WakeSettings { + pub fn validate(&self) -> Result<(), String> { + if self.after_secs == 0 || self.max_per_hour == 0 || self.ports.is_empty() { + return Err( + "wake after_secs and max_per_hour must be positive and ports non-empty".into(), + ); + } + if self.enabled && self.hosts.is_empty() { + return Err("wake is enabled but lists no hosts".into()); + } + let mut names = BTreeSet::new(); + for h in &self.hosts { + if parse_mac(&h.mac).is_none() { + return Err(format!("wake host {}: bad mac", h.name)); + } + if h.name.is_empty() || !names.insert(&h.name) { + return Err("wake host names must be non-empty and unique".into()); + } + } + if self.broadcast.iter().any(|b| parse_target(b).is_none()) { + return Err("wake broadcast entries are IPv4 addresses or addr/len".into()); + } + Ok(()) + } + + /// Where the packet goes: every target on every port. + pub fn destinations(&self) -> Vec { + let mut out = Vec::new(); + for ip in self.broadcast.iter().filter_map(|b| parse_target(b)) { + for &port in &self.ports { + let d = SocketAddrV4::new(ip, port); + if !out.contains(&d) { + out.push(d); + } + } + } + out + } +} + +/// One wake the planner decided on. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct WakeEvent { + pub at: i64, + pub host: String, + pub class: Class, + pub reason: String, +} + +#[derive(Debug, Default)] +pub struct WakePlanner { + /// Since when each class has had unmet demand without a break. + unmet_since: BTreeMap, + last_wake: Option, + per_host: BTreeMap, + recent: VecDeque, + attempts: BTreeMap, + events: VecDeque, +} + +/// What `/status` shows. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)] +pub struct WakeView { + pub enabled: bool, + /// Seconds each class has had unmet demand. + pub unmet_secs: BTreeMap, + pub wakes_last_hour: u32, + /// Wake attempts per host since the queen started. + pub attempts: BTreeMap, + /// The newest wakes, newest first. + pub events: Vec, +} + +impl WakePlanner { + pub fn new() -> Self { + Self::default() + } + + /// Decide wakes for this pass. `unmet` are classes with demand the fleet + /// could not place; `awake` are names of drones that are fresh. + pub fn step( + &mut self, + now: i64, + s: &WakeSettings, + unmet: &BTreeSet, + awake: &BTreeSet<&str>, + ) -> Vec<(usize, WakeEvent)> { + self.unmet_since.retain(|c, _| unmet.contains(c)); + for c in unmet { + self.unmet_since.entry(*c).or_insert(now); + } + while self.recent.front().is_some_and(|&t| now - t >= 3600) { + self.recent.pop_front(); + } + let secs = |v: u64| i64::try_from(v).unwrap_or(i64::MAX); + let mut out: Vec<(usize, WakeEvent)> = Vec::new(); + if !s.enabled { + return out; + } + // One wake per pass: the next one needs the min interval to have + // passed, which gives the woken host time to boot and heartbeat. + if self + .last_wake + .is_some_and(|t| now - t < secs(s.min_interval_secs)) + || u32::try_from(self.recent.len()).unwrap_or(u32::MAX) >= s.max_per_hour + { + return out; + } + for (&class, &since) in &self.unmet_since { + if now - since < secs(s.after_secs) { + continue; + } + let pick = s.hosts.iter().enumerate().find(|(_, h)| { + (h.classes.is_empty() || h.classes.contains(&class)) + && !awake.contains(h.name.as_str()) + && self + .per_host + .get(&h.name) + .is_none_or(|&t| now - t >= secs(s.host_cooldown_secs)) + }); + if let Some((i, h)) = pick { + self.last_wake = Some(now); + self.per_host.insert(h.name.clone(), now); + self.recent.push_back(now); + *self.attempts.entry(h.name.clone()).or_default() += 1; + let event = WakeEvent { + at: now, + host: h.name.clone(), + class, + reason: format!( + "{class} demand unmet for {}s (>= {}s)", + now - since, + s.after_secs + ), + }; + self.events.push_front(event.clone()); + self.events.truncate(20); + out.push((i, event)); + break; + } + } + out + } + + pub fn view(&self, now: i64, enabled: bool) -> WakeView { + WakeView { + enabled, + unmet_secs: self + .unmet_since + .iter() + .map(|(c, &t)| (*c, u64::try_from(now - t).unwrap_or(0))) + .collect(), + wakes_last_hour: u32::try_from(self.recent.len()).unwrap_or(u32::MAX), + attempts: self.attempts.clone(), + events: self.events.iter().cloned().collect(), + } + } +} + +/// Sends magic packets. Returns the destinations it reached. +pub trait WakeSender: Send + Sync { + fn send<'a>( + &'a self, + mac: [u8; 6], + to: &'a [SocketAddrV4], + ) -> BoxFuture<'a, Result, String>>; +} + +/// Sends the packet over UDP with broadcast enabled. +pub struct UdpWakeSender; + +impl WakeSender for UdpWakeSender { + fn send<'a>( + &'a self, + mac: [u8; 6], + to: &'a [SocketAddrV4], + ) -> BoxFuture<'a, Result, String>> { + Box::pin(async move { + let packet = magic_packet(mac); + let sock = tokio::net::UdpSocket::bind("0.0.0.0:0") + .await + .map_err(|e| e.to_string())?; + sock.set_broadcast(true).map_err(|e| e.to_string())?; + let mut sent = Vec::new(); + for d in to { + if sock.send_to(&packet, d).await.is_ok() { + sent.push(*d); + } + } + if sent.is_empty() { + Err("no destination accepted the packet".into()) + } else { + Ok(sent) + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn host(name: &str, mac: &str, classes: &[Class]) -> WakeHost { + WakeHost { + name: name.into(), + mac: mac.into(), + classes: classes.to_vec(), + } + } + + fn settings() -> WakeSettings { + WakeSettings { + enabled: true, + hosts: vec![ + host("a", "aa:bb:cc:00:00:01", &[Class::Heavy]), + host("b", "aa:bb:cc:00:00:02", &[]), + ], + after_secs: 300, + min_interval_secs: 600, + max_per_hour: 2, + host_cooldown_secs: 1200, + ..WakeSettings::default() + } + } + + fn unmet(c: &[Class]) -> BTreeSet { + c.iter().copied().collect() + } + + fn step(p: &mut WakePlanner, now: i64, u: &[Class], awake: &[&str]) -> Vec { + let awake: BTreeSet<&str> = awake.iter().copied().collect(); + p.step(now, &settings(), &unmet(u), &awake) + .into_iter() + .map(|(_, e)| e.host) + .collect() + } + + #[test] + fn ported_packet_helpers_match_flotilla() { + let m = parse_mac("A4:83:E7:12:34:56").unwrap(); + assert_eq!(m, [0xa4, 0x83, 0xe7, 0x12, 0x34, 0x56]); + assert_eq!(parse_mac("a4-83-e7-12-34-56"), Some(m)); + assert_eq!(parse_mac("a483e7123456"), Some(m)); + assert!(parse_mac("nope").is_none()); + let p = magic_packet(m); + assert_eq!(p.len(), 102); + assert_eq!(&p[..6], &[0xff; 6]); + assert_eq!(&p[6..12], &m); + assert_eq!(&p[96..], &m); + } + + #[test] + fn broadcast_math_and_targets() { + assert_eq!( + directed_broadcast("192.0.2.77".parse().unwrap(), 24).to_string(), + "192.0.2.255" + ); + assert_eq!( + directed_broadcast("198.51.100.7".parse().unwrap(), 16).to_string(), + "198.51.255.255" + ); + assert_eq!( + directed_broadcast("198.51.100.7".parse().unwrap(), 0).to_string(), + "255.255.255.255" + ); + assert_eq!( + directed_broadcast("198.51.100.7".parse().unwrap(), 32).to_string(), + "198.51.100.7" + ); + assert_eq!( + parse_target("192.0.2.0/24").unwrap().to_string(), + "192.0.2.255" + ); + assert_eq!( + parse_target("255.255.255.255").unwrap().to_string(), + "255.255.255.255" + ); + assert!(parse_target("nope").is_none() && parse_target("1.2.3.4/x").is_none()); + let s = WakeSettings { + broadcast: vec![ + "255.255.255.255".into(), + "192.0.2.0/24".into(), + "192.0.2.255".into(), + ], + ..WakeSettings::default() + }; + // Both spellings of the same broadcast address are sent once per port. + let d: Vec = s.destinations().iter().map(ToString::to_string).collect(); + assert_eq!( + d, + [ + "255.255.255.255:9", + "255.255.255.255:7", + "192.0.2.255:9", + "192.0.2.255:7" + ] + ); + } + + #[test] + fn demand_must_stay_unmet_for_the_whole_period() { + let mut p = WakePlanner::new(); + assert!(step(&mut p, 0, &[Class::Medium], &[]).is_empty()); + assert!(step(&mut p, 299, &[Class::Medium], &[]).is_empty()); + // A gap resets the clock. + assert!(step(&mut p, 300, &[], &[]).is_empty()); + assert!(step(&mut p, 310, &[Class::Medium], &[]).is_empty()); + assert!(step(&mut p, 609, &[Class::Medium], &[]).is_empty()); + assert_eq!(step(&mut p, 610, &[Class::Medium], &[]), ["b"]); + } + + #[test] + fn only_a_sleeping_host_that_can_serve_the_class_is_woken() { + // Heavy demand: "a" is listed first and serves heavy. + let mut p = WakePlanner::new(); + step(&mut p, 0, &[Class::Heavy], &[]); + assert_eq!(step(&mut p, 300, &[Class::Heavy], &[]), ["a"]); + // "a" cannot serve light; "b" can. + let mut p = WakePlanner::new(); + step(&mut p, 0, &[Class::Light], &[]); + assert_eq!(step(&mut p, 300, &[Class::Light], &[]), ["b"]); + // If "a" is awake (fresh drone), heavy demand falls to "b". + let mut p = WakePlanner::new(); + step(&mut p, 0, &[Class::Heavy], &["a"]); + assert_eq!(step(&mut p, 300, &[Class::Heavy], &["a"]), ["b"]); + // If everything that could serve it is awake, nothing is sent. + let mut p = WakePlanner::new(); + step(&mut p, 0, &[Class::Heavy], &["a", "b"]); + assert!(step(&mut p, 900, &[Class::Heavy], &["a", "b"]).is_empty()); + } + + #[test] + fn the_global_interval_the_hourly_cap_and_the_host_cooldown_rate_limit_wakes() { + let mut p = WakePlanner::new(); + step(&mut p, 0, &[Class::Light], &[]); + assert_eq!(step(&mut p, 300, &[Class::Light], &[]), ["b"]); + // Inside the 600s interval: nothing, even though still unmet. + assert!(step(&mut p, 599, &[Class::Light], &[]).is_empty()); + // After it, "b" is still cooling down (1200s) and "a" cannot serve light. + assert!(step(&mut p, 950, &[Class::Light], &[]).is_empty()); + // Cooldown over: "b" is retried (a host that never came up). + assert_eq!(step(&mut p, 1500, &[Class::Light], &[]), ["b"]); + // Two wakes this hour is the cap. + assert!(step(&mut p, 3000, &[Class::Light], &[]).is_empty()); + // An hour after the first wake the window slides. + assert_eq!(step(&mut p, 3901, &[Class::Light], &[]), ["b"]); + } + + #[test] + fn one_wake_per_pass_even_with_several_classes_unmet() { + let mut p = WakePlanner::new(); + step(&mut p, 0, &[Class::Light, Class::Heavy], &[]); + let got = step(&mut p, 300, &[Class::Light, Class::Heavy], &[]); + assert_eq!(got.len(), 1); + } + + #[test] + fn disabled_planner_tracks_but_never_wakes_and_settings_validate() { + let mut p = WakePlanner::new(); + let off = WakeSettings { + enabled: false, + ..settings() + }; + let awake = BTreeSet::new(); + assert!(p.step(0, &off, &unmet(&[Class::Light]), &awake).is_empty()); + assert!( + p.step(10_000, &off, &unmet(&[Class::Light]), &awake) + .is_empty() + ); + + assert!(settings().validate().is_ok()); + for bad in [ + WakeSettings { + hosts: vec![], + ..settings() + }, + WakeSettings { + hosts: vec![host("a", "zz", &[])], + ..settings() + }, + WakeSettings { + hosts: vec![ + host("a", "aa:bb:cc:00:00:01", &[]), + host("a", "aa:bb:cc:00:00:02", &[]), + ], + ..settings() + }, + WakeSettings { + broadcast: vec!["x".into()], + ..settings() + }, + WakeSettings { + max_per_hour: 0, + ..settings() + }, + ] { + assert!(bad.validate().is_err()); + } + } + + #[tokio::test] + async fn the_udp_sender_delivers_a_magic_packet() { + let rx = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let std::net::SocketAddr::V4(addr) = rx.local_addr().unwrap() else { + panic!("ipv4 expected"); + }; + let mac = parse_mac("aa:bb:cc:dd:ee:ff").unwrap(); + let sent = UdpWakeSender.send(mac, &[addr]).await.unwrap(); + assert_eq!(sent, [addr]); + let mut buf = [0u8; 256]; + let (n, _) = rx.recv_from(&mut buf).await.unwrap(); + assert_eq!(&buf[..n], magic_packet(mac).as_slice()); + } +} diff --git a/tests/admin_cli.rs b/tests/admin_cli.rs new file mode 100644 index 0000000..158aea4 --- /dev/null +++ b/tests/admin_cli.rs @@ -0,0 +1,113 @@ +//! The `borg` binary's queen administration commands and `--token-file`. + +use std::fs; +use std::path::Path; + +use borg::drone::exec::{Exec, SystemExec}; + +/// Cargo supplies the binary path at compile time for `cargo test`; an +/// all-target check can compile this file without it, so fall back to the +/// binary beside the test's deps directory. +fn borg_bin() -> String { + option_env!("CARGO_BIN_EXE_borg") + .map(str::to_owned) + .unwrap_or_else(|| { + let exe = std::env::current_exe().expect("test executable path"); + let dir = exe + .parent() + .and_then(Path::parent) + .expect("target profile directory"); + dir.join("borg").to_string_lossy().into_owned() + }) +} + +fn run(args: &[&str]) -> (bool, String) { + let o = SystemExec.run(&borg_bin(), args).unwrap(); + (o.success, o.stdout) +} + +fn write_config(dir: &Path) -> String { + let path = dir.join("queen.json"); + let db = dir.join("queen.db"); + fs::write( + &path, + format!( + r#"{{"github":{{"org":"example-org","app_id":"1"}},"repos":["example-org/app"],"db_path":"{}"}}"#, + db.display() + ), + ) + .unwrap(); + path.to_string_lossy().into_owned() +} + +#[test] +fn enroll_prints_only_the_token_and_list_and_remove_manage_drones() { + let dir = tempfile::tempdir().unwrap(); + let cfg = write_config(dir.path()); + + let (ok, token) = run(&["queen", "enroll", "--config", &cfg, "host-a"]); + assert!(ok); + assert!( + token.starts_with("borg_") + && token.len() == 5 + 64 + && token + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_'), + "stdout is exactly the token: {token:?}" + ); + + let (ok, list) = run(&["queen", "drone", "list", "--config", &cfg]); + assert!(ok); + assert!(list.contains("host-a") && list.contains("never"), "{list}"); + assert!(!list.contains(&token), "the token is never shown again"); + + let (ok, removed) = run(&["queen", "drone", "remove", "--config", &cfg, "host-a"]); + assert!(ok && removed.contains("removed drone host-a"), "{removed}"); + let (_, list) = run(&["queen", "drone", "list", "--config", &cfg]); + assert!(!list.contains("host-a"), "{list}"); + let (ok, _) = run(&["queen", "drone", "remove", "--config", &cfg, "host-a"]); + assert!(!ok, "removing an unknown drone fails"); +} + +#[test] +fn install_reads_the_token_from_a_file_and_never_needs_it_on_the_command_line() { + let dir = tempfile::tempdir().unwrap(); + let token = dir.path().join("token"); + fs::write(&token, "borg_from_a_file\n").unwrap(); + let me = SystemExec.run("id", &["-un"]).unwrap().stdout; + let root = dir.path().join("root"); + // Linux hosts must be running systemd; stand in for that under --root. + fs::create_dir_all(root.join("run/systemd/system")).unwrap(); + let (ok, plan) = run(&[ + "drone", + "install", + "--queen", + "https://queen.example.test", + "--token-file", + &token.to_string_lossy(), + "--dry-run", + "--no-job-user", + "--run-as", + &me, + "--root", + &root.to_string_lossy(), + ]); + assert!(ok, "{plan}"); + assert!(plan.contains("drone.token"), "{plan}"); + assert!( + !plan.contains("borg_from_a_file"), + "the token is never printed" + ); + + let (ok, _) = run(&[ + "drone", + "install", + "--queen", + "https://queen.example.test", + "--dry-run", + "--no-job-user", + "--run-as", + &me, + ]); + assert!(!ok, "a token is required"); +} diff --git a/tests/capacity.rs b/tests/capacity.rs new file mode 100644 index 0000000..3b20130 --- /dev/null +++ b/tests/capacity.rs @@ -0,0 +1,129 @@ +mod support; + +use serde_json::{Value, json}; +use support::Harness; +use time::OffsetDateTime; + +fn add_static(h: &Harness, id: u64, label: &str, busy: bool) { + h.state().runners.push(json!({ + "id": id, "name": format!("static-{id}"), "status": "online", "busy": busy, + "labels": [{"name": "self-hosted"}, {"name": label}] + })); +} + +fn drone<'a>(status: &'a Value, name: &str) -> &'a Value { + status["drones"] + .as_array() + .unwrap() + .iter() + .find(|d| d["name"] == name) + .unwrap() +} + +#[tokio::test] +async fn status_reports_free_capacity_next_to_theoretical_capacity() { + let h = Harness::start(json!({})).await; + let t = h.enroll("d1").await; + h.heartbeat(&t, &[]).await; + + // An empty fleet: free equals theoretical (28 GiB budget, 16 vCPU). + let s = h.status().await; + let d = drone(&s, "d1"); + assert_eq!(d["budget_mem_mib"], 28 * 1024); + assert_eq!(d["free_mem_mib"], 28 * 1024); + assert_eq!(d["free_slots"]["heavy"], d["theoretical_slots"]["heavy"]); + assert_eq!(d["free_slots"]["heavy"], 1); + assert_eq!(s["classes"]["heavy"]["free_slots"], 1); + + // One heavy lease (24 GiB, 8 vCPU) eats the budget: the theoretical heavy + // slot is still 1, but nothing heavy is free and the memory shows it. + h.queue_job(1, "borg-ci-heavy-pr", 5); + h.tick().await; + let s = h.status().await; + let d = drone(&s, "d1"); + assert_eq!(d["theoretical_slots"]["heavy"], 1); + assert_eq!(d["free_slots"]["heavy"], 0); + assert_eq!(d["free_mem_mib"], 4 * 1024); + assert_eq!(d["free_vcpu"], 8); + assert_eq!(d["free_slots"]["light"], 1); + assert_eq!(s["classes"]["heavy"]["slots"], 1); + assert_eq!(s["classes"]["heavy"]["free_slots"], 0); + assert_eq!(s["classes"]["heavy"]["leases_pending"], 1); + + let m = h.metrics().await; + for line in [ + "borg_drone_free_mem_mib{drone=\"d1\"} 4096", + "borg_drone_free_slots{drone=\"d1\",class=\"heavy\"} 0", + "borg_class_free_slots{class=\"heavy\"} 0", + ] { + assert!(m.contains(line), "missing {line}\n{m}"); + } +} + +#[tokio::test] +async fn demand_is_split_into_what_non_borg_runners_will_take_and_what_borg_must_place() { + let h = Harness::start(json!({})).await; + // No drones at all: every queued job is unmet. + for i in 0..5 { + h.queue_job(i, "borg-ci-light-pr", 10); + } + for id in 901..904 { + add_static(&h, id, "borg-ci-light-pr", false); + } + add_static(&h, 904, "borg-ci-light-pr", true); // busy: online, not idle + add_static(&h, 905, "borg-ci-light", false); // main label: not for these jobs + h.tick().await; + + let s = h.status().await; + let pr = &s["classes"]["light"]["demand_pr"]; + assert_eq!(pr["queued"], 5); + assert_eq!(pr["static_online"], 4); + assert_eq!(pr["static_idle"], 3); + assert_eq!(pr["claimable_by_static_estimate"], 3); + assert_eq!(pr["borg_placeable"], 2); + assert_eq!(s["classes"]["light"]["demand_main"]["static_idle"], 1); + // The allocator's own unmet still says 5; the placeable view says 2. + assert_eq!(s["unmet"][0]["count"], 5); + assert_eq!(s["unmet_placeable"][0]["count"], 2); + let m = h.metrics().await; + assert!( + m.contains("borg_queued_borg_placeable{class=\"light\",scope=\"pr\"} 2"), + "{m}" + ); + assert!(m.contains("borg_queued_claimable_by_static{class=\"light\",scope=\"pr\"} 3")); +} + +#[tokio::test] +async fn wake_on_lan_ignores_demand_that_non_borg_runners_will_take() { + let cfg = json!({ + "drone_stale_secs": 1_000_000, + "wake": {"enabled": true, "after_secs": 300, "hosts": [ + {"name": "sleeper", "mac": "aa:bb:cc:00:00:01"}]} + }); + let base = OffsetDateTime::now_utc(); + let at = |s: i64| base + time::Duration::seconds(s); + + // Three queued jobs, three idle static runners: nothing for borg to wake for. + let h = Harness::start(cfg.clone()).await; + for i in 0..3 { + h.queue_job(i, "borg-ci-light-pr", 10); + add_static(&h, 900 + i, "borg-ci-light-pr", false); + } + for s in [0, 310, 700] { + h.queen.tick(at(s)).await.unwrap(); + } + assert!(h.wakes.0.lock().unwrap().is_empty()); + + // Two static runners for three jobs: one is borg's, and it wakes the host. + let h = Harness::start(cfg).await; + for i in 0..3 { + h.queue_job(i, "borg-ci-light-pr", 10); + } + for i in 0..2 { + add_static(&h, 900 + i, "borg-ci-light-pr", false); + } + for s in [0, 310] { + h.queen.tick(at(s)).await.unwrap(); + } + assert_eq!(h.wakes.0.lock().unwrap().len(), 1); +} diff --git a/tests/contract.rs b/tests/contract.rs new file mode 100644 index 0000000..fe18dca --- /dev/null +++ b/tests/contract.rs @@ -0,0 +1,287 @@ +mod support; + +use borg::enroll; +use serde_json::{Value, json}; +use support::Harness; +use time::OffsetDateTime; + +fn now() -> i64 { + OffsetDateTime::now_utc().unix_timestamp() +} + +fn good(toolchain: &str) -> Value { + json!({ + "observed_at": now(), + "rustc_wrapper": "/usr/local/bin/sccache", "sccache_dir": "/var/cache/sccache", + "sccache_server_reachable": true, "sccache_stats_ok": true, + "cache_requests": 1000, "cache_hits": 500, "cache_errors": 0, + "cache_dir_free_bytes": 100u64 << 30, "cache_write_ok": true, + "toolchains": [{"channel": toolchain, "rustc_version": toolchain}], + "mem_available_mib": 30_000, "cpu_psi_some_avg10_pm": 50 + }) +} + +fn cfg(extra: Value) -> Value { + let mut c = json!({ + "drone_stale_secs": 1_000_000, + "env_contract": { + "enabled": true, + "checks": {"compiler_cache": "required", "toolchain": "required", + "cache_store": "required", "resources": "required"} + }, + "units": {"enabled": true} + }); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + c[k] = v; + } + c +} + +async fn beat(h: &Harness, token: &str, facts: Option) -> Value { + match facts { + Some(f) => h.heartbeat_with(token, json!({"env_facts": f})).await, + None => h.heartbeat_with(token, json!({})).await, + } +} + +fn drone<'a>(status: &'a Value, name: &str) -> &'a Value { + status["drones"] + .as_array() + .unwrap() + .iter() + .find(|d| d["name"] == name) + .unwrap() +} + +#[tokio::test] +async fn a_failing_required_check_keeps_ci_leases_off_that_drone() { + let h = Harness::start(cfg(json!({}))).await; + // "a-bad" would win every name tie; it has no compiler cache wrapper. + let bad = h.enroll("a-bad").await; + let ok = h.enroll("b-good").await; + let mut broken = good("1.95.0"); + broken["rustc_wrapper"] = Value::Null; + beat(&h, &bad, Some(broken)).await; + beat(&h, &ok, Some(good("1.95.0"))).await; + + h.queue_job(1, "borg-ci-light-pr", 5); + h.tick().await; + assert!( + beat(&h, &bad, None).await["leases"] + .as_array() + .unwrap() + .is_empty() + ); + assert_eq!( + beat(&h, &ok, None).await["leases"] + .as_array() + .unwrap() + .len(), + 1 + ); + + let s = h.status().await; + let c = &drone(&s, "a-bad")["contract"]; + assert_eq!(c["admission"]["light"]["admit"], false); + assert!( + c["admission"]["light"]["why"] + .as_str() + .unwrap() + .contains("RUSTC_WRAPPER is not set") + ); + let cc = c["checks"] + .as_array() + .unwrap() + .iter() + .find(|x| x["kind"] == "compiler_cache") + .unwrap(); + assert_eq!( + (cc["state"].as_str(), cc["mode"].as_str()), + (Some("fail"), Some("required")) + ); + assert_eq!( + drone(&s, "b-good")["contract"]["admission"]["heavy"]["admit"], + true + ); + // A blocked drone shows no free slots for the class it is not admitted to. + assert_eq!(drone(&s, "a-bad")["free_slots"]["light"], 0); + + let m = h.metrics().await; + for line in [ + "borg_env_contract_check{drone=\"a-bad\",check=\"compiler_cache\",state=\"fail\"} 1", + "borg_env_contract_check{drone=\"b-good\",check=\"compiler_cache\",state=\"pass\"} 1", + "borg_env_contract_admitted{drone=\"a-bad\",class=\"light\"} 0", + "borg_env_contract_admitted{drone=\"b-good\",class=\"light\"} 1", + ] { + assert!(m.contains(line), "missing {line}\n{m}"); + } +} + +#[tokio::test] +async fn units_are_placed_only_on_drones_with_the_toolchain_the_tree_asks_for() { + let h = Harness::start(cfg(json!({}))).await; + let client = enroll::new_token(); + h.queen + .store() + .enroll_client("agent", &enroll::hash_token(&client), 0) + .await + .unwrap(); + let old = h.enroll("a-old").await; + let new = h.enroll("b-new").await; + beat(&h, &old, Some(good("1.90.0"))).await; + beat(&h, &new, Some(good("1.95.0"))).await; + + let r = h + .http + .post(format!("{}/v1/units", h.url)) + .bearer_auth(&client) + .json(&json!({"repo": "example-org/app", "tree_sha": "a".repeat(40), "argv": ["cargo", "check"], + "platform": "linux-x86_64", "class": "light", "rust_toolchain": "1.95.0"})) + .send() + .await + .unwrap(); + assert_eq!(r.status(), 201); + h.tick().await; + assert!( + beat(&h, &old, None).await["units"] + .as_array() + .unwrap() + .is_empty() + ); + assert_eq!( + beat(&h, &new, None).await["units"] + .as_array() + .unwrap() + .len(), + 1 + ); + + // Nobody has 1.99.0: the unit waits, and the status says why for each drone. + let r = h + .http + .post(format!("{}/v1/units", h.url)) + .bearer_auth(&client) + .json(&json!({"repo": "example-org/app", "tree_sha": "b".repeat(40), "argv": ["cargo", "check"], + "platform": "linux-x86_64", "class": "light", "rust_toolchain": "1.99.0"})) + .send() + .await + .unwrap(); + let id = r.json::().await.unwrap()["unit_id"] + .as_str() + .unwrap() + .to_owned(); + h.tick().await; + assert_eq!( + h.queen + .store() + .unit(&id) + .await + .unwrap() + .unwrap() + .state + .as_str(), + "queued" + ); + let status = h.status().await; + let blocks = status["units"]["not_admitted"].as_array().unwrap(); + assert_eq!(blocks.len(), 2, "{blocks:?}"); + assert!(blocks.iter().all(|b| b["unit_id"] == id.as_str())); + assert!( + blocks[0]["why"] + .as_str() + .unwrap() + .contains("toolchain 1.99.0 is not installed (have 1.90.0)") + ); +} + +#[tokio::test] +async fn memory_and_cpu_pressure_are_checked_against_the_class_needs() { + let h = Harness::start(cfg(json!({}))).await; + let t = h.enroll("d1").await; + let mut f = good("1.95.0"); + f["mem_available_mib"] = json!(10_000); // heavy needs 24 GiB, medium 12 GiB + beat(&h, &t, Some(f)).await; + h.tick().await; + let a = &drone(&h.status().await, "d1")["contract"]["admission"].clone(); + assert_eq!( + ( + a["light"]["admit"].as_bool(), + a["medium"]["admit"].as_bool(), + a["heavy"]["admit"].as_bool() + ), + (Some(true), Some(false), Some(false)) + ); + assert!( + a["heavy"]["why"] + .as_str() + .unwrap() + .contains("10000 MiB available") + ); + + // CPU pressure over the ceiling blocks everything. + let mut f = good("1.95.0"); + f["cpu_psi_some_avg10_pm"] = json!(950); + beat(&h, &t, Some(f)).await; + let a = &drone(&h.status().await, "d1")["contract"]["admission"].clone(); + assert_eq!(a["light"]["admit"], false); +} + +#[tokio::test] +async fn drones_that_never_reported_are_admitted_by_default_and_stale_facts_are_unknown() { + let h = Harness::start(cfg(json!({}))).await; + let t = h.enroll("old-drone").await; + h.heartbeat(&t, &[]).await; + h.queue_job(1, "borg-ci-light-pr", 5); + h.tick().await; + // Predates the contract: still gets work. + assert_eq!( + beat(&h, &t, None).await["leases"].as_array().unwrap().len(), + 1 + ); + let s = h.status().await; + assert_eq!(drone(&s, "old-drone")["contract"]["reported"], false); + + // Facts that go stale are unknown, which a required check treats as blocking. + let t2 = h.enroll("z-drone").await; + let mut f = good("1.95.0"); + f["observed_at"] = json!(now() - 3600); + beat(&h, &t2, Some(f)).await; + let s = h.status().await; + let c = &drone(&s, "z-drone")["contract"]; + assert_eq!(c["admission"]["light"]["admit"], false); + assert_eq!(c["checks"][0]["state"], "unknown"); +} + +#[tokio::test] +async fn the_contract_is_off_by_default_and_require_report_is_opt_in() { + // Off: even terrible facts block nothing and no contract appears. + let h = Harness::start(json!({"drone_stale_secs": 1_000_000})).await; + let t = h.enroll("d1").await; + let mut f = good("1.95.0"); + f["rustc_wrapper"] = Value::Null; + f["mem_available_mib"] = json!(1); + beat(&h, &t, Some(f)).await; + h.queue_job(1, "borg-ci-light-pr", 5); + h.tick().await; + assert_eq!( + beat(&h, &t, None).await["leases"].as_array().unwrap().len(), + 1 + ); + assert!(drone(&h.status().await, "d1")["contract"].is_null()); + + // require_report: a drone with no facts gets nothing. + let h = Harness::start(cfg( + json!({"env_contract": {"enabled": true, "require_report": true}}), + )) + .await; + let t = h.enroll("d1").await; + h.heartbeat(&t, &[]).await; + h.queue_job(1, "borg-ci-light-pr", 5); + h.tick().await; + assert!( + beat(&h, &t, None).await["leases"] + .as_array() + .unwrap() + .is_empty() + ); +} diff --git a/tests/drone.rs b/tests/drone.rs index c64e01c..992de54 100644 --- a/tests/drone.rs +++ b/tests/drone.rs @@ -67,7 +67,11 @@ struct Rig { impl Rig { async fn new() -> Self { - let h = Harness::start(json!({})).await; + Self::with_queen(json!({})).await + } + + async fn with_queen(queen: Value) -> Self { + let h = Harness::start(queen).await; let token = h.enroll("d1").await; Self { h, @@ -310,6 +314,19 @@ impl Exec for GitOnlyExec { } } +/// Runs everything except docker, which is only recorded. +struct AllButDocker(RecordingExec); + +impl Exec for AllButDocker { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + if program == "docker" { + self.0.run(program, args) + } else { + SystemExec.run(program, args) + } + } +} + fn git(args: &[&str]) -> String { let out = SystemExec.run("git", args).unwrap(); assert!(out.success, "git {args:?} failed"); @@ -471,6 +488,7 @@ async fn the_mirror_timer_picks_up_new_commits() { let caches = Arc::new(borg::drone::caches::Caches::new( cfg, Arc::new(SystemExec), + Arc::new(borg::drone::fsops::LocalFs), Arc::new(FixedHost), )); let timer = caches.spawn_background(); @@ -551,7 +569,7 @@ async fn leases_get_distinct_slots_while_concurrent_and_the_same_path_when_a_slo v }; let first_pwds: Vec = read_all("pwd-").into_iter().map(|(_, p)| p).collect(); - let work = rig.path("work"); + let work = fs::canonicalize(rig.path("work")).unwrap(); let mut want = vec![ work.join("slot-0").display().to_string(), work.join("slot-1").display().to_string(), @@ -577,3 +595,200 @@ async fn leases_get_distinct_slots_while_concurrent_and_the_same_path_when_a_slo ); assert_eq!(fs::read_dir(rig.path("work")).unwrap().count(), 0); } + +/// A stand-in for `setpriv` that records its arguments, then runs the command +/// after `--` as the current user. +fn fake_setpriv(dir: &Path, log: &Path) -> PathBuf { + let path = dir.join("fake-setpriv"); + fs::write( + &path, + format!( + "#!/bin/sh\necho \"$@\" >> {}\nwhile [ \"$1\" != \"--\" ]; do shift; done\nshift\nexec \"$@\"\n", + log.display() + ), + ) + .unwrap(); + fs::set_permissions(&path, fs::Permissions::from_mode(0o755)).unwrap(); + path +} + +#[tokio::test] +async fn with_a_job_user_the_runner_and_its_file_operations_go_through_setpriv() { + let rig = Rig::new().await; + let log = rig.path("setpriv.log"); + let setpriv = fake_setpriv(rig.dir.path(), &log); + rig.install_runner( + "printf '%s' \"$ACTIONS_RUNNER_INPUT_JITCONFIG\" > $OUT/jit.txt\necho \"$HOME|$USER\" > $OUT/who.txt\necho \"$@\" > $OUT/args.txt", + ); + let drone = rig.drone_with( + json!({"job_user": { + "name": "borg-job", + "home": rig.path("job-home"), + "priv_bin": setpriv, + }}), + Arc::new(AllButDocker(RecordingExec::default())), + ); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + rig.h.queue_job(1, "borg-ci-light", 1); + rig.h.tick().await; + assert_eq!(drone.heartbeat_once(true, &mut tasks).await.unwrap(), 1); + drain(&mut tasks).await; + + let calls = fs::read_to_string(&log).unwrap(); + let prefix = "--reuid borg-job --regid borg-job --init-groups --inh-caps=-all --ambient-caps=-all --bounding-set=-all --"; + // File contents passed as arguments can span lines; only count call starts. + let lines: Vec<&str> = calls.lines().collect(); + let starts: Vec<&&str> = lines.iter().filter(|l| l.starts_with("--reuid")).collect(); + assert!(starts.len() >= 5, "{calls}"); + assert!( + starts.iter().all(|l| l.starts_with(prefix)), + "every call drops to the job user: {calls}" + ); + let has = |needle: &str| lines.iter().any(|l| l.contains(needle)); + assert!( + has("cp -al"), + "the runner install is cloned as the job user: {calls}" + ); + assert!( + has("/work/slot-0/run.sh"), + "run.sh is started through setpriv: {calls}" + ); + assert!( + has("rm -rf -- ") && has("/work/slot-0"), + "cleanup removes the slot as the job user: {calls}" + ); + + let name = rig.h.state().mints[0]["name"].as_str().unwrap().to_owned(); + assert_eq!( + fs::read_to_string(rig.path("out/jit.txt")).unwrap(), + format!("jit-{name}") + ); + assert_eq!( + fs::read_to_string(rig.path("out/args.txt")).unwrap().trim(), + "" + ); + assert_eq!( + fs::read_to_string(rig.path("out/who.txt")).unwrap().trim(), + format!("{}|borg-job", rig.path("job-home").display()), + "the job sees the job user's home" + ); + assert!(!rig.path("work/slot-0").exists()); +} + +#[tokio::test] +async fn extra_labels_reach_the_minted_runner_but_class_labels_are_dropped() { + let rig = Rig::new().await; + rig.install_runner("true"); + let drone = rig.drone(json!({"extra_labels": ["on-prem", "legacy-lane"]})); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + rig.h.queue_job(1, "borg-ci-light-pr", 1); + rig.h.tick().await; + let labels: Vec = rig.h.state().mints[0]["labels"] + .as_array() + .unwrap() + .iter() + .map(|l| l.as_str().unwrap().to_owned()) + .collect(); + assert_eq!( + labels, + [ + "self-hosted", + "Linux", + "X64", + "on-prem", + "legacy-lane", + "borg-ci-light", + "borg-ci-light-pr" + ] + ); + drain(&mut tasks).await; + + // A drone cannot smuggle a class label past the queen: a PR runner stays PR-scoped. + let req = json!({"capacity": support::capacity(), "extra_labels": ["borg-ci-light", "borg-ci-heavy-pr", "ok-label"]}); + let r = rig + .h + .http + .post(format!("{}/v1/heartbeat", rig.h.url)) + .bearer_auth(&rig.token) + .json(&req) + .send() + .await + .unwrap(); + assert!(r.status().is_success()); + rig.h.state().jobs.clear(); + rig.h.state().mints.clear(); + rig.h.queue_job(2, "borg-ci-medium-pr", 1); + rig.h.tick().await; + let labels: Vec = rig.h.state().mints[0]["labels"] + .as_array() + .unwrap() + .iter() + .map(|l| l.as_str().unwrap().to_owned()) + .collect(); + assert_eq!( + labels, + [ + "self-hosted", + "Linux", + "X64", + "ok-label", + "borg-ci-medium", + "borg-ci-medium-pr" + ] + ); +} + +#[tokio::test] +async fn runner_env_and_hooks_are_exported_to_every_lease() { + let rig = Rig::new().await; + rig.install_runner( + "echo \"$HTTP_PROXY|$ACTIONS_RUNNER_HOOK_JOB_STARTED|$ACTIONS_RUNNER_HOOK_JOB_COMPLETED\" > $OUT/env.txt", + ); + let drone = rig.drone(json!({ + "runner_env": {"HTTP_PROXY": "http://proxy.example.test:3128"}, + "job_started_hook": "/etc/borg/hooks/start.sh", + "job_completed_hook": "/etc/borg/hooks/done.sh", + })); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + rig.h.queue_job(1, "borg-ci-light", 1); + rig.h.tick().await; + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + drain(&mut tasks).await; + assert_eq!( + fs::read_to_string(rig.path("out/env.txt")).unwrap().trim(), + "http://proxy.example.test:3128|/etc/borg/hooks/start.sh|/etc/borg/hooks/done.sh" + ); +} + +fn request_version(dir: &Path) -> Option { + let v: Value = serde_json::from_slice(&fs::read(dir.join("update-request")).ok()?).ok()?; + v["version"].as_str().map(str::to_owned) +} + +#[tokio::test] +async fn an_idle_drone_asks_for_the_version_the_queen_wants_and_takes_no_lease() { + let rig = Rig::with_queen(json!({"rollout_target_version": "9.9.9"})).await; + rig.install_runner("exit 0"); + let state = rig.path("state"); + fs::create_dir_all(&state).unwrap(); + let drone = rig.drone(json!({"state_dir": state})); + let mut tasks = JoinSet::new(); + + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + assert_eq!(request_version(&state).as_deref(), Some("9.9.9")); + assert!( + state.join("health.json").exists(), + "health evidence is written" + ); + + rig.h.queue_job(1, "borg-ci-light-pr", 3); + rig.h.tick().await; + assert_eq!( + drone.heartbeat_once(true, &mut tasks).await.unwrap(), + 0, + "no new lease while an update is pending" + ); +} diff --git a/tests/health.rs b/tests/health.rs new file mode 100644 index 0000000..60ec8ba --- /dev/null +++ b/tests/health.rs @@ -0,0 +1,318 @@ +mod support; + +use borg::classes::{Class, Scope}; +use borg::store::{LeaseRow, LeaseStatus}; +use borg::telemetry::{Conclusion, FailureClass, JobTelemetry}; +use serde_json::{Value, json}; +use support::Harness; +use time::OffsetDateTime; + +fn at(t: i64) -> OffsetDateTime { + OffsetDateTime::from_unix_timestamp(t).unwrap() +} + +struct Fleet { + h: Harness, + tokens: Vec<(String, String)>, + base: i64, + n: u64, +} + +impl Fleet { + async fn new(drones: &[&str], extra: Value) -> Self { + let mut cfg = json!({"drone_stale_secs": 1_000_000_000}); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + cfg[k] = v; + } + let h = Harness::start(cfg).await; + let mut tokens = Vec::new(); + for d in drones { + let token = h.enroll(d).await; + h.heartbeat(&token, &[]).await; + tokens.push(((*d).to_owned(), token)); + } + Self { + h, + tokens, + base: OffsetDateTime::now_utc().unix_timestamp(), + n: 0, + } + } + + /// Record a finished job on `drone` at `base + offset`. + async fn job(&mut self, drone: &str, offset: i64, conclusion: Conclusion, class: FailureClass) { + self.n += 1; + let lease = LeaseRow { + id: format!("lease-{}", self.n), + drone: drone.to_owned(), + class: Class::Light, + scope: Scope::Pr, + labels: vec![], + runner_name: "r".into(), + runner_id: 1, + jit_config: None, + state: LeaseStatus::Done, + created_at: 0, + expires_at: 0, + }; + let t = JobTelemetry { + job_name: "job".into(), + repo: "example-org/app".into(), + run_id: 1000 + self.n, + conclusion, + failure_class: class, + queue_wait_secs: None, + steps: vec![], + }; + self.h + .queen + .store() + .record_telemetry(&lease, &t, 10, self.base + offset) + .await + .unwrap(); + } + + /// A successful job on `drone` that took `secs`. + async fn job_taking(&mut self, drone: &str, offset: i64, secs: u64) { + self.n += 1; + let lease = LeaseRow { + id: format!("lease-{}", self.n), + drone: drone.to_owned(), + class: Class::Light, + scope: Scope::Pr, + labels: vec![], + runner_name: "r".into(), + runner_id: 1, + jit_config: None, + state: LeaseStatus::Done, + created_at: 0, + expires_at: 0, + }; + let t = JobTelemetry { + job_name: format!("job (shard {offset})"), + repo: "example-org/app".into(), + run_id: 5000 + self.n, + conclusion: Conclusion::Success, + failure_class: FailureClass::Unknown, + queue_wait_secs: None, + steps: vec![], + }; + self.h + .queen + .store() + .record_telemetry(&lease, &t, secs, self.base + offset) + .await + .unwrap(); + } + + async fn host_failures(&mut self, drone: &str, count: usize, offset: i64) { + for _ in 0..count { + self.job(drone, offset, Conclusion::Failure, FailureClass::Host) + .await; + } + } + + async fn tick(&self, offset: i64) { + self.h.queen.tick(at(self.base + offset)).await.unwrap(); + } + + async fn health(&self, drone: &str) -> Value { + let status = self.h.status().await; + status["drones"] + .as_array() + .unwrap() + .iter() + .find(|d| d["name"] == drone) + .unwrap()["health"] + .clone() + } + + fn token(&self, drone: &str) -> &str { + &self.tokens.iter().find(|(n, _)| n == drone).unwrap().1 + } +} + +#[tokio::test] +async fn corroborated_host_failures_quarantine_a_drone_and_leases_move_away() { + let mut f = Fleet::new(&["d1", "d2"], json!({})).await; + f.host_failures("d1", 3, 0).await; + f.tick(10).await; + + let h = f.health("d1").await; + assert_eq!(h["state"], "quarantined", "{h}"); + assert_eq!(h["source"], "auto"); + assert_eq!(h["corroborating_runs"], 3); + assert_eq!(f.health("d2").await["state"], "healthy"); + + // The audit trail keeps the evidence. + let events = f.h.status().await["health_events"].clone(); + assert_eq!(events[0]["event"], "auto_quarantine"); + assert_eq!(events[0]["evidence"].as_array().unwrap().len(), 3); + + // Ties go to d1 by name, so a lease landing on d2 shows d1 is held back. + f.h.queue_job(1, "borg-ci-light-pr", 5); + f.tick(11).await; + assert!(f.h.leases(f.token("d1"), &[]).await.is_empty()); + assert_eq!(f.h.leases(f.token("d2"), &[]).await.len(), 1); + + let metrics = f.h.metrics().await; + assert!( + metrics.contains("borg_drone_health_state{drone=\"d1\",state=\"quarantined\"} 1"), + "{metrics}" + ); + assert!(metrics.contains("borg_drones{state=\"quarantined\"} 1")); +} + +#[tokio::test] +async fn code_failures_and_a_single_host_failure_never_suppress_a_host() { + let mut f = Fleet::new(&["d1", "d2"], json!({})).await; + for _ in 0..6 { + f.job("d1", 0, Conclusion::Failure, FailureClass::Code) + .await; + } + f.host_failures("d1", 1, 0).await; + f.tick(10).await; + assert_eq!(f.health("d1").await["state"], "healthy"); + f.host_failures("d1", 1, 0).await; + f.tick(11).await; + // Two corroborated failures only recommend. + assert_eq!(f.health("d1").await["state"], "recommended"); + assert!( + !f.h.status().await["drones"][0]["quarantined"] + .as_bool() + .unwrap() + ); +} + +#[tokio::test] +async fn cooldown_probation_and_a_health_proof_return_a_host_to_service() { + let mut f = Fleet::new(&["d1", "d2"], json!({})).await; + f.host_failures("d1", 3, 0).await; + f.tick(10).await; + assert_eq!(f.health("d1").await["state"], "quarantined"); + + // Inside the cooldown (default 1800s) nothing changes. + f.tick(1000).await; + assert_eq!(f.health("d1").await["state"], "quarantined"); + + // After it, probation: d1 is schedulable again, one lease at a time. + f.tick(1900).await; + f.tick(1901).await; + assert_eq!(f.health("d1").await["state"], "probation"); + f.h.queue_job(1, "borg-ci-light-pr", 5); + f.tick(1902).await; + assert_eq!(f.h.leases(f.token("d1"), &[]).await.len(), 1); + + // Two successful jobs after probation began prove it healthy. + f.job("d1", 1950, Conclusion::Success, FailureClass::Unknown) + .await; + f.tick(2000).await; + assert_eq!(f.health("d1").await["state"], "probation"); + f.job("d1", 2010, Conclusion::Success, FailureClass::Unknown) + .await; + f.tick(2020).await; + f.tick(2021).await; + assert_eq!(f.health("d1").await["state"], "healthy"); + let events: Vec = f.h.status().await["health_events"] + .as_array() + .unwrap() + .iter() + .map(|e| e["event"].as_str().unwrap().to_owned()) + .collect(); + assert_eq!(events, ["release", "probation_start", "auto_quarantine"]); + + // The failures that caused the first quarantine do not count again. + f.tick(2100).await; + assert_eq!(f.health("d1").await["state"], "healthy"); +} + +#[tokio::test] +async fn a_host_failure_on_probation_requarantines_with_a_longer_cooldown() { + let mut f = Fleet::new(&["d1", "d2"], json!({})).await; + f.host_failures("d1", 3, 0).await; + f.tick(10).await; + f.tick(1900).await; + f.tick(1901).await; + assert_eq!(f.health("d1").await["state"], "probation"); + + f.host_failures("d1", 1, 1950).await; + f.tick(2000).await; + f.tick(2001).await; + let h = f.health("d1").await; + assert_eq!( + (h["state"].as_str(), h["strikes"].as_u64()), + (Some("quarantined"), Some(1)), + "{h}" + ); + // The second cooldown is twice the first (3600s from the requarantine). + assert_eq!(h["probation_at"], f.base + 2000 + 3600); +} + +#[tokio::test] +async fn operator_quarantine_is_never_lifted_by_the_cooldown() { + let f = Fleet::new(&["d1", "d2"], json!({})).await; + f.h.queen + .store() + .set_quarantine("d1", Some("rack move")) + .await + .unwrap(); + f.tick(1_000_000).await; + let h = f.health("d1").await; + assert_eq!( + (h["state"].as_str(), h["source"].as_str()), + (Some("quarantined"), Some("operator")) + ); +} + +#[tokio::test] +async fn fleet_wide_host_failures_only_recommend() { + let mut f = Fleet::new(&["d1", "d2", "d3", "d4"], json!({})).await; + f.host_failures("d1", 3, 0).await; + f.host_failures("d2", 1, 0).await; + f.job("d3", 0, Conclusion::Success, FailureClass::Unknown) + .await; + f.tick(10).await; + assert_eq!(f.health("d1").await["state"], "recommended"); + assert_eq!(f.h.status().await["health_systemic"], true); + assert!(f.h.metrics().await.contains("borg_health_systemic 1")); +} + +#[tokio::test] +async fn automatic_quarantine_can_be_switched_off() { + let mut f = Fleet::new(&["d1", "d2"], json!({"health": {"auto_quarantine": false}})).await; + f.host_failures("d1", 4, 0).await; + f.tick(10).await; + let h = f.health("d1").await; + assert_eq!(h["state"], "recommended"); + assert!(h["note"].as_str().unwrap().contains("off")); +} + +#[tokio::test] +async fn learned_profiles_steer_new_leases_to_the_faster_host() { + // d1 would win every tie by name; d2 has proven faster on the same family. + let mut f = Fleet::new(&["d1", "d2"], json!({})).await; + for i in 0..8 { + f.job_taking("d1", i, 200 + i.unsigned_abs()).await; + f.job_taking("d2", i, 100 + i.unsigned_abs()).await; + } + f.tick(10).await; + let status = f.h.status().await; + let p2 = status["profiles"] + .as_array() + .unwrap() + .iter() + .find(|p| p["drone"] == "d2" && p["class"] == "light") + .unwrap(); + assert!(p2["speed_permille"].as_u64().unwrap() > 1100, "{p2}"); + assert_eq!(p2["families"][0]["family"], "job"); + assert!( + f.h.metrics() + .await + .contains("borg_drone_profile_speed{drone=\"d2\",class=\"light\"}") + ); + + f.h.queue_job(1, "borg-ci-light-pr", 5); + f.tick(20).await; + assert!(f.h.leases(f.token("d1"), &[]).await.is_empty()); + assert_eq!(f.h.leases(f.token("d2"), &[]).await.len(), 1); +} diff --git a/tests/job_user.rs b/tests/job_user.rs new file mode 100644 index 0000000..2b8b722 --- /dev/null +++ b/tests/job_user.rs @@ -0,0 +1,94 @@ +//! Real privilege separation. Needs root (to change user) and util-linux +//! `setpriv`; skipped otherwise. Run it with, for example, +//! `sudo -E cargo test --test job_user -- --nocapture`. + +use std::fs; +use std::os::unix::fs::PermissionsExt; +use std::sync::Arc; + +use borg::drone::exec::{Exec, SystemExec}; +use borg::drone::jobuser::{JobExec, JobUser}; + +/// Numeric id of `nobody`: present on every Linux system, unprivileged, and +/// in no interesting group. +const JOB_ID: &str = "65534"; + +fn out(exec: &dyn Exec, program: &str, args: &[&str]) -> (bool, String) { + let o = exec.run(program, args).unwrap(); + (o.success, o.stdout) +} + +#[test] +fn a_lease_runs_as_a_different_uid_without_access_to_the_drone_token_or_any_capability() { + let (_, uid) = out(&SystemExec, "id", &["-u"]); + if uid != "0" + || !SystemExec + .run("setpriv", &["--version"]) + .is_ok_and(|o| o.success) + { + println!("skipping: needs root and setpriv"); + return; + } + let drone_uid = uid; + + // Shaped like /etc/borg: a directory only the drone's user can enter. + let dir = tempfile::tempdir().unwrap(); + let config_dir = dir.path().join("etc-borg"); + fs::create_dir(&config_dir).unwrap(); + fs::set_permissions(&config_dir, fs::Permissions::from_mode(0o700)).unwrap(); + let token = config_dir.join("drone.token"); + fs::write(&token, "borg_secret_token\n").unwrap(); + fs::set_permissions(&token, fs::Permissions::from_mode(0o600)).unwrap(); + // Let the job user at least reach the temporary parent. + fs::set_permissions(dir.path(), fs::Permissions::from_mode(0o755)).unwrap(); + + let job = JobExec::new( + Arc::new(SystemExec), + JobUser { + name: JOB_ID.to_owned(), + home: "/nonexistent".into(), + priv_bin: "setpriv".to_owned(), + }, + ); + + let (ok, lease_uid) = out(&job, "id", &["-u"]); + assert!(ok, "the lease command ran"); + assert_eq!(lease_uid, JOB_ID); + assert_ne!( + lease_uid, drone_uid, + "the lease runs as a different uid than the drone" + ); + + let token_path = token.to_string_lossy().into_owned(); + let (read_ok, contents) = out(&job, "cat", &[&token_path]); + assert!( + !read_ok && !contents.contains("secret"), + "the job user cannot read the token" + ); + let (listed, _) = out(&job, "ls", &[&config_dir.to_string_lossy()]); + assert!(!listed, "nor list the drone's config directory"); + // The drone itself can, which is the point of the split. + assert_eq!( + fs::read_to_string(&token).unwrap().trim(), + "borg_secret_token" + ); + + // No capability survives into the job. + let (_, caps) = out( + &job, + "sh", + &[ + "-c", + "grep -E '^(CapEff|CapPrm|CapAmb|CapBnd):' /proc/self/status", + ], + ); + for line in caps.lines() { + assert!( + line.ends_with("0000000000000000"), + "capabilities remain: {line}" + ); + } + println!( + "lease uid {lease_uid} vs drone uid {drone_uid}; token unreadable; capabilities dropped" + ); +} diff --git a/tests/queen.rs b/tests/queen.rs index 201db54..ab41fb6 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -548,3 +548,144 @@ async fn old_telemetry_is_pruned_by_the_tick() { .is_empty() ); } + +#[tokio::test] +async fn bids_gate_leases_on_published_pressure_and_expired_bids_are_ignored() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + let shy = |expires_at: i64| { + json!({"bids": [{"class": "light", "scope": "pr", "available_slots": 4, + "min_pressure": 1000, "expires_at": expires_at}]}) + }; + let far = OffsetDateTime::now_utc().unix_timestamp() + 3600; + + h.heartbeat_with(&token, shy(far)).await; + h.queue_job(1, "borg-ci-light-pr", 5); + h.tick().await; + // One queued job against a mostly idle host is low pressure: the shy bid + // stays out, and the queen publishes the pressure it computed. + let resp = h.heartbeat_with(&token, shy(far)).await; + assert!(resp["leases"].as_array().unwrap().is_empty(), "{resp}"); + let light = resp["pressure"] + .as_array() + .unwrap() + .iter() + .find(|p| p["class"] == "light") + .unwrap(); + assert!(light["pressure"].as_u64().unwrap() < 1000, "{resp}"); + assert!( + h.metrics() + .await + .contains("borg_demand_pressure{class=\"light\"}") + ); + + // The same bid, expired, is ignored and the drone is a legacy drone. + h.heartbeat_with(&token, shy(1)).await; + h.tick().await; + let resp = h.heartbeat_with(&token, shy(1)).await; + assert_eq!(resp["leases"].as_array().unwrap().len(), 1, "{resp}"); +} + +#[tokio::test] +async fn homeostasis_raises_warm_floors_under_sustained_queueing_and_stays_off_by_default() { + let cfg = json!({"drone_stale_secs": 100_000_000, "homeostasis": { + "enabled": true, "target_secs": 60, "raise_after_secs": 100, "cooldown_secs": 10, + "max_floor": {"light": 3}}}); + let h = Harness::start(cfg).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let base = OffsetDateTime::now_utc(); + // One job queued for 5 minutes at the start: far over the 60s target. + h.queue_job(1, "borg-ci-light-pr", 300); + + for secs in [0, 50, 100, 110, 220] { + h.queen + .tick(base + time::Duration::seconds(secs)) + .await + .unwrap(); + } + let status = h.status().await; + let light_pr = status["homeostasis"] + .as_array() + .unwrap() + .iter() + .find(|r| r["class"] == "light" && r["scope"] == "pr") + .unwrap(); + assert_eq!(light_pr["level"], 2, "{light_pr}"); + assert_eq!(light_pr["effective_floor"], 2); + // Two floors against one queued job: the second is an extra warm runner. + assert_eq!(h.state().mints.len(), 2); + let metrics = h.metrics().await; + assert!( + metrics.contains("borg_warm_floor_dynamic{class=\"light\",scope=\"pr\"} 2"), + "{metrics}" + ); + assert!(metrics.contains("borg_warm_floor{class=\"light\",scope=\"pr\"} 2")); + + // Off by default: the same load changes nothing. + let off = Harness::start(json!({"drone_stale_secs": 100_000_000})).await; + let t = off.enroll("d1").await; + off.heartbeat(&t, &[]).await; + off.queue_job(1, "borg-ci-light-pr", 300); + for secs in [0, 130, 260] { + off.queen + .tick(base + time::Duration::seconds(secs)) + .await + .unwrap(); + } + assert_eq!(off.state().mints.len(), 1); + assert!( + off.status().await["homeostasis"] + .as_array() + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn reclaim_reports_are_audited_once_and_shown_in_status_and_metrics() { + let h = Harness::start(json!({})).await; + let token = h.enroll("devbox").await; + let stats = |entries: Value, total: u64| { + json!({"reclaim": { + "enabled": true, "dry_run": false, "reclaimed_bytes_total": total, + "last_run_at": 100, "last_run_reclaimed_bytes": total, + "last_run_would_reclaim_bytes": 0, "last_run_errors": 1, + "usage_pct": 91, "last_run_mode": "pressure", + "kept": {"busy": 2}, "recent": entries + }}) + }; + let entries = json!([ + {"seq": 1, "at": 100, "path": "/srv/devdata/cargo-target/old-branch", "kind": "cargo_target", + "size_bytes": 5_368_709_120u64, "reason": "disk 91% >= 85% high water", "outcome": "deleted"}, + {"seq": 2, "at": 101, "path": "/srv/devdata/worktrees/x", "kind": "worktree", + "size_bytes": 1_073_741_824u64, "reason": "r", "outcome": "failed", "error": "permission denied"} + ]); + h.heartbeat_with(&token, stats(entries.clone(), 5_368_709_120)) + .await; + // The same entries again (a repeated heartbeat) must not duplicate. + h.heartbeat_with(&token, stats(entries, 5_368_709_120)) + .await; + + let status = h.status().await; + assert_eq!( + status["reclaim"]["reclaimed_bytes_logged"], + 5_368_709_120u64 + ); + let events = status["reclaim"]["events"].as_array().unwrap(); + assert_eq!(events.len(), 2, "{events:?}"); + assert_eq!(events[0]["drone"], "devbox"); + assert_eq!(events[1]["path"], "/srv/devdata/cargo-target/old-branch"); + assert_eq!(events[1]["outcome"], "deleted"); + assert_eq!(status["drones"][0]["reclaim"]["usage_pct"], 91); + assert_eq!(status["drones"][0]["reclaim"]["recent"], json!([])); + + let metrics = h.metrics().await; + for line in [ + "borg_reclaim_bytes_total{drone=\"devbox\"} 5368709120", + "borg_reclaim_disk_usage_percent{drone=\"devbox\"} 91", + "borg_reclaim_errors{drone=\"devbox\"} 1", + ] { + assert!(metrics.contains(line), "missing {line}\n{metrics}"); + } +} diff --git a/tests/rollout.rs b/tests/rollout.rs new file mode 100644 index 0000000..1f1fe44 --- /dev/null +++ b/tests/rollout.rs @@ -0,0 +1,229 @@ +mod support; + +use borg::classes::{Class, Scope}; +use borg::store::{LeaseRow, LeaseStatus}; +use borg::telemetry::{Conclusion, FailureClass, JobTelemetry}; +use serde_json::{Value, json}; +use support::Harness; +use time::OffsetDateTime; + +struct Rig { + h: Harness, + tokens: Vec<(String, String)>, + base: i64, + n: u64, +} + +fn gated() -> Value { + json!({ + "drone_stale_secs": 1_000_000_000, + "rollout_target_version": "2.0.0", + "rollout": { + "previous_version": "1.0.0", "canary_jobs": 2, "waves_pct": [50, 100], + "wave_jobs": 2, "soak_secs": 60, "update_timeout_secs": 300, + "min_baseline_jobs": 100 + } + }) +} + +impl Rig { + async fn new(cfg: Value) -> Self { + let h = Harness::start(cfg).await; + let mut tokens = Vec::new(); + for d in ["l1", "l2", "m1"] { + tokens.push((d.to_owned(), h.enroll(d).await)); + } + let rig = Self { + h, + tokens, + base: OffsetDateTime::now_utc().unix_timestamp(), + n: 0, + }; + for d in ["l1", "l2", "m1"] { + rig.beat(d, "1.0.0").await; + } + rig + } + + fn token(&self, d: &str) -> &str { + &self.tokens.iter().find(|(n, _)| n == d).unwrap().1 + } + + /// Heartbeat as `d` running `version`; returns the response. + async fn beat(&self, d: &str, version: &str) -> Value { + let mut cap = serde_json::to_value(support::capacity()).unwrap(); + if d.starts_with('m') { + cap["os"] = json!("macos"); + } + self.h + .heartbeat_with(self.token(d), json!({"capacity": cap, "version": version})) + .await + } + + async fn desired(&self, d: &str, running: &str) -> Option { + self.beat(d, running).await["desired_version"] + .as_str() + .map(str::to_owned) + } + + async fn job(&mut self, d: &str, offset: i64, ok: bool) { + self.n += 1; + let lease = LeaseRow { + id: format!("lease-{}", self.n), + drone: d.to_owned(), + class: Class::Light, + scope: Scope::Pr, + labels: vec![], + runner_name: "r".into(), + runner_id: 1, + jit_config: None, + state: LeaseStatus::Done, + created_at: 0, + expires_at: 0, + }; + let t = JobTelemetry { + job_name: "job".into(), + repo: "example-org/app".into(), + run_id: 9000 + self.n, + conclusion: if ok { + Conclusion::Success + } else { + Conclusion::Failure + }, + failure_class: if ok { + FailureClass::Unknown + } else { + FailureClass::Host + }, + queue_wait_secs: None, + steps: vec![], + }; + self.h + .queen + .store() + .record_telemetry(&lease, &t, 10, self.base + offset) + .await + .unwrap(); + } + + async fn tick(&self, offset: i64) { + self.h + .queen + .tick(OffsetDateTime::from_unix_timestamp(self.base + offset).unwrap()) + .await + .unwrap(); + } + + async fn phase(&self) -> String { + self.h.status().await["rollout"]["gate"]["state"]["phase"] + .as_str() + .unwrap_or("none") + .to_owned() + } +} + +#[tokio::test] +async fn a_rollout_goes_canary_then_waves_then_complete_with_the_right_drones_told() { + let mut r = Rig::new(gated()).await; + // Before the first pass the queen gives no instruction. + assert_eq!(r.desired("l1", "1.0.0").await, None); + + r.tick(0).await; + assert_eq!(r.phase().await, "canary"); + assert_eq!(r.desired("l1", "1.0.0").await.as_deref(), Some("2.0.0")); + assert_eq!(r.desired("m1", "1.0.0").await.as_deref(), Some("2.0.0")); + assert_eq!(r.desired("l2", "1.0.0").await.as_deref(), Some("1.0.0")); + + // The canaries update and each completes two jobs. + r.beat("l1", "2.0.0").await; + r.beat("m1", "2.0.0").await; + for d in ["l1", "m1"] { + r.job(d, 10, true).await; + r.job(d, 11, true).await; + } + r.tick(20).await; + assert_eq!(r.phase().await, "wave"); + + // Soak not over: still wave 1. + r.job("l1", 30, true).await; + r.job("m1", 31, true).await; + r.tick(50).await; + let st = r.h.status().await; + assert_eq!(st["rollout"]["gate"]["state"]["wave"], 0); + assert!( + st["rollout"]["gate"]["waiting"] + .as_str() + .unwrap() + .contains("soaking") + ); + + // Soaked, enough jobs: wave 2 tells everyone. + r.tick(100).await; + assert_eq!(r.desired("l2", "1.0.0").await.as_deref(), Some("2.0.0")); + r.beat("l2", "2.0.0").await; + r.job("l2", 110, true).await; + r.job("l2", 111, true).await; + r.tick(200).await; + assert_eq!(r.phase().await, "complete"); + let log = r.h.status().await["rollout"]["log"].clone(); + assert!(log[0][1].as_str().unwrap().contains("complete"), "{log}"); + let metrics = r.h.metrics().await; + assert!( + metrics.contains("borg_rollout_phase{phase=\"complete\"} 1"), + "{metrics}" + ); + + // The state survives a restart: a fresh pass keeps it complete. + r.tick(300).await; + assert_eq!(r.phase().await, "complete"); +} + +#[tokio::test] +async fn host_failures_on_a_canary_halt_the_rollout_and_roll_every_drone_back() { + let mut r = Rig::new(gated()).await; + r.tick(0).await; + r.beat("l1", "2.0.0").await; + r.beat("m1", "2.0.0").await; + for i in 0..3 { + r.job("l1", 10 + i, true).await; + r.job("m1", 10 + i, false).await; + } + r.tick(30).await; + assert_eq!(r.phase().await, "halted"); + for d in ["l1", "m1", "l2"] { + assert_eq!(r.desired(d, "2.0.0").await.as_deref(), Some("1.0.0"), "{d}"); + } + let st = r.h.status().await; + assert!( + st["rollout"]["gate"]["state"]["halted_reason"] + .as_str() + .unwrap() + .contains("m1") + ); + assert!( + r.h.metrics() + .await + .contains("borg_rollout_phase{phase=\"halted\"} 1") + ); + + // Halted stays halted, even with a clean later history. + r.job("m1", 40, true).await; + r.tick(100).await; + assert_eq!(r.phase().await, "halted"); +} + +#[tokio::test] +async fn without_a_previous_version_the_target_is_only_reported_and_everyone_is_told_it() { + let r = Rig::new(json!({ + "drone_stale_secs": 1_000_000_000, + "rollout_target_version": "2.0.0" + })) + .await; + r.tick(0).await; + assert_eq!(r.desired("l2", "1.0.0").await.as_deref(), Some("2.0.0")); + assert_eq!(r.phase().await, "none"); + + let none = Rig::new(json!({"drone_stale_secs": 1_000_000_000})).await; + none.tick(0).await; + assert_eq!(none.desired("l1", "1.0.0").await, None); +} diff --git a/tests/run.rs b/tests/run.rs new file mode 100644 index 0000000..abf5325 --- /dev/null +++ b/tests/run.rs @@ -0,0 +1,449 @@ +//! `borg run` against a fake queen and a real temporary git repository. + +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::{Arc, Mutex}; + +use axum::extract::{Path as AxPath, State}; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use borg::run::{ + RunOptions, TreeSource, parse_toolchain_channel, repo_from_remote, resolve_tree, run, + tree_rust_toolchain, +}; +use serde_json::{Value, json}; + +#[derive(Clone, Copy, PartialEq)] +enum Mode { + Receipt, + Run, + Failed, + Queued, + Refuse, + /// Runs like `Run` but rejects the bundle upload. + RefuseBundle, +} + +struct Fake { + mode: Mode, + submitted: Mutex>, + cancelled: Mutex, + bundles: Mutex>>, +} + +type Shared = Arc; + +async fn submit(State(f): State, Json(b): Json) -> Response { + f.submitted.lock().unwrap().push(b); + match f.mode { + Mode::Refuse => StatusCode::INTERNAL_SERVER_ERROR.into_response(), + Mode::Receipt => Json( + json!({"outcome": "receipt", "unit_id": "u-1", "state": "done", + "receipt": {"drone": "d9", "exit_code": 0}}), + ) + .into_response(), + _ => ( + StatusCode::CREATED, + Json(json!({"outcome": "new", "unit_id": "u-1", "state": "queued"})), + ) + .into_response(), + } +} + +async fn log(State(f): State, AxPath(_id): AxPath) -> Response { + let mut r = match f.mode { + Mode::Run | Mode::RefuseBundle => "hello from the fleet\n".into_response(), + _ => { + tokio::time::sleep(std::time::Duration::from_millis(150)).await; + "".into_response() + } + }; + let h = r.headers_mut(); + let (state, next) = match f.mode { + Mode::Run | Mode::RefuseBundle => ("done", "21"), + Mode::Failed => ("failed", "0"), + _ => ("queued", "0"), + }; + h.insert("x-borg-state", state.parse().unwrap()); + h.insert("x-borg-next", next.parse().unwrap()); + if matches!(f.mode, Mode::Run | Mode::RefuseBundle) { + h.insert("x-borg-exit", "3".parse().unwrap()); + } + if f.mode == Mode::Failed { + h.insert("x-borg-failure", "timed out".parse().unwrap()); + } + r +} + +async fn cancel(State(f): State, AxPath(_id): AxPath) -> StatusCode { + *f.cancelled.lock().unwrap() += 1; + StatusCode::NO_CONTENT +} + +async fn bundle( + State(f): State, + AxPath(_id): AxPath, + body: axum::body::Bytes, +) -> StatusCode { + if f.mode == Mode::RefuseBundle { + return StatusCode::PAYLOAD_TOO_LARGE; + } + f.bundles.lock().unwrap().push(body.to_vec()); + StatusCode::NO_CONTENT +} + +async fn fake_queen(mode: Mode) -> (String, Shared) { + let f = Arc::new(Fake { + mode, + submitted: Mutex::new(Vec::new()), + cancelled: Mutex::new(0), + bundles: Mutex::new(Vec::new()), + }); + let app = Router::new() + .route("/v1/units", post(submit)) + .route("/v1/units/{id}/log", get(log)) + .route("/v1/units/{id}/bundle", axum::routing::put(bundle)) + .route("/v1/units/{id}/cancel", post(cancel)) + .with_state(Arc::clone(&f)); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", l.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(l, app).await }); + (url, f) +} + +fn git(dir: &Path, args: &[&str]) -> String { + #[allow( + clippy::disallowed_methods, + reason = "test fixture: scratch git repositories built with the real git binary" + )] + let out = Command::new("git") + .args([ + "-c", + "user.name=t", + "-c", + "user.email=t@example.test", + "-c", + "commit.gpgsign=false", + ]) + .args(args) + .current_dir(dir) + .output() + .unwrap(); + assert!( + out.status.success(), + "git {args:?}: {}", + String::from_utf8_lossy(&out.stderr) + ); + String::from_utf8_lossy(&out.stdout).trim().to_owned() +} + +fn git_succeeds(dir: &Path, args: &[&str]) -> bool { + #[allow( + clippy::disallowed_methods, + reason = "test fixture: asks the real git binary whether an object exists" + )] + let status = Command::new("git").args(args).current_dir(dir).status(); + status.is_ok_and(|s| s.success()) +} + +/// A repository with one pushed commit and its bare origin. +struct Repo { + _tmp: tempfile::TempDir, + dir: PathBuf, + origin: PathBuf, +} + +fn repo() -> Repo { + let tmp = tempfile::tempdir().unwrap(); + let root = std::fs::canonicalize(tmp.path()).unwrap(); + let dir = root.join("work"); + let origin = root.join("origin.git"); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::create_dir_all(&origin).unwrap(); + git(&origin, &["init", "-q", "--bare", "-b", "main"]); + git(&dir, &["init", "-q", "-b", "main"]); + std::fs::write(dir.join("a.txt"), "one\n").unwrap(); + git(&dir, &["add", "."]); + git(&dir, &["commit", "-q", "-m", "first"]); + git(&dir, &["remote", "add", "origin", origin.to_str().unwrap()]); + git(&dir, &["push", "-q", "origin", "main"]); + Repo { + _tmp: tmp, + dir, + origin, + } +} + +fn opts(r: &Repo, queen: Option, marker: &str) -> RunOptions { + RunOptions { + token: queen + .as_ref() + .map(|_| borg::secret::Secret::new("borg_client")), + queen, + dir: r.dir.clone(), + tree: TreeSource::Head, + repo: Some("example-org/app".into()), + platform: "linux-x86_64".into(), + class: "medium".into(), + toolchain: "tc".into(), + timeout_secs: None, + force: false, + wait_secs: 1, + publish: false, + argv: vec!["sh".into(), "-c".into(), format!("touch {marker}; exit 7")], + } +} + +fn marker(r: &Repo) -> String { + r.dir.join("ran-locally").display().to_string() +} + +#[test] +fn trees_and_repository_names_are_resolved_from_git() { + let r = repo(); + let head = resolve_tree(&r.dir, &TreeSource::Head).unwrap(); + assert_eq!(head, git(&r.dir, &["rev-parse", "HEAD^{tree}"])); + assert_eq!(resolve_tree(&r.dir, &TreeSource::Index).unwrap(), head); + // Staging a change moves the index tree, not HEAD's. + std::fs::write(r.dir.join("a.txt"), "two\n").unwrap(); + git(&r.dir, &["add", "a.txt"]); + let staged = resolve_tree(&r.dir, &TreeSource::Index).unwrap(); + assert_ne!(staged, head); + assert_eq!(resolve_tree(&r.dir, &TreeSource::Head).unwrap(), head); + // A commit id resolves to its tree. + let commit = git(&r.dir, &["rev-parse", "HEAD"]); + assert_eq!( + resolve_tree(&r.dir, &TreeSource::Explicit(commit)).unwrap(), + head + ); + assert!(resolve_tree(&r.dir, &TreeSource::Explicit("nope".into())).is_err()); + + for (url, want) in [ + ("git@github.com:dx-corp/mono.git", Some("dx-corp/mono")), + ("https://github.com/dx-corp/mono", Some("dx-corp/mono")), + ("https://github.com/dx-corp/mono.git/", Some("dx-corp/mono")), + ("mono", None), + ] { + assert_eq!(repo_from_remote(url).as_deref(), want, "{url}"); + } +} + +#[tokio::test] +async fn without_a_queen_the_command_just_runs_here() { + let r = repo(); + let code = run(&opts(&r, None, &marker(&r)), &mut Vec::new()).await; + assert_eq!(code, 7); + assert!(Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn a_receipt_means_nothing_runs_and_the_exit_code_is_zero() { + let r = repo(); + let (url, fake) = fake_queen(Mode::Receipt).await; + let code = run(&opts(&r, Some(url), &marker(&r)), &mut Vec::new()).await; + assert_eq!(code, 0); + assert!(!Path::new(&marker(&r)).exists()); + let sent = fake.submitted.lock().unwrap()[0].clone(); + assert_eq!( + sent["tree_sha"], + resolve_tree(&r.dir, &TreeSource::Head).unwrap() + ); + assert_eq!(sent["repo"], "example-org/app"); + assert_eq!(sent["argv"][0], "sh"); + assert_eq!(sent["platform"], "linux-x86_64"); +} + +#[tokio::test] +async fn a_remote_run_streams_the_log_and_returns_the_units_exit_code() { + let r = repo(); + let (url, _fake) = fake_queen(Mode::Run).await; + let mut out = Vec::new(); + let code = run(&opts(&r, Some(url), &marker(&r)), &mut out).await; + assert_eq!(code, 3, "the unit's code, not the local command's"); + assert_eq!(String::from_utf8(out).unwrap(), "hello from the fleet\n"); + assert!(!Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn the_index_tree_is_what_a_hook_submits_and_travels_as_a_small_bundle() { + let r = repo(); + std::fs::write(r.dir.join("a.txt"), "staged\n").unwrap(); + git(&r.dir, &["add", "a.txt"]); + let tree = git(&r.dir, &["write-tree"]); + let (url, fake) = fake_queen(Mode::Run).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.tree = TreeSource::Index; + // The staged tree is not on any remote: without --publish this runs locally. + assert_eq!(run(&o, &mut Vec::new()).await, 7); + assert!(fake.submitted.lock().unwrap().is_empty()); + + // With --publish it is submitted with a bundle, and nothing is pushed. + o.publish = true; + assert_eq!(run(&o, &mut Vec::new()).await, 3); + let sent = fake.submitted.lock().unwrap()[0].clone(); + assert_eq!(sent["tree_sha"], tree.as_str()); + assert_eq!(sent["bundle"], true); + let bundle = fake.bundles.lock().unwrap()[0].clone(); + assert!(bundle.starts_with(b"# v2 git bundle")); + assert!( + git(&r.origin, &["for-each-ref"]) + .lines() + .all(|l| !l.contains("refs/borg")) + ); + assert!( + git(&r.dir, &["for-each-ref"]) + .lines() + .all(|l| !l.contains("refs/borg")) + ); + + // What a drone does: a clone of the remote (its mirror) unbundles it and + // then has the tree, though it never had the staged content. + let mirror = r.dir.parent().unwrap().join("mirror"); + git( + &r.dir, + &[ + "clone", + "-q", + "--bare", + r.origin.to_str().unwrap(), + mirror.to_str().unwrap(), + ], + ); + assert!(!git_succeeds(&mirror, &["cat-file", "-e", &tree])); + let file = mirror.join("unit.bundle"); + std::fs::write(&file, &bundle).unwrap(); + git(&mirror, &["bundle", "unbundle", file.to_str().unwrap()]); + assert_eq!(git(&mirror, &["cat-file", "-t", &tree]), "tree"); + assert_eq!(git(&mirror, &["show", &format!("{tree}:a.txt")]), "staged"); +} + +#[tokio::test] +async fn a_pushed_tree_sends_no_bundle_and_a_refused_bundle_falls_back() { + let r = repo(); + let (url, fake) = fake_queen(Mode::Run).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.publish = true; + assert_eq!(run(&o, &mut Vec::new()).await, 3); + assert_eq!(fake.submitted.lock().unwrap()[0]["bundle"], false); + assert!(fake.bundles.lock().unwrap().is_empty()); + + // The queen refuses the bundle (too large): the unit is cancelled and the + // command runs here. + std::fs::write(r.dir.join("a.txt"), "staged\n").unwrap(); + git(&r.dir, &["add", "a.txt"]); + let (url, fake) = fake_queen(Mode::RefuseBundle).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.tree = TreeSource::Index; + o.publish = true; + assert_eq!(run(&o, &mut Vec::new()).await, 7); + assert_eq!(*fake.cancelled.lock().unwrap(), 1); + assert!(Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn an_unreachable_or_refusing_queen_falls_back_to_running_here() { + let r = repo(); + // Nothing listens on this port. + let dead = { + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + format!("http://{}", l.local_addr().unwrap()) + }; + assert_eq!( + run(&opts(&r, Some(dead), &marker(&r)), &mut Vec::new()).await, + 7 + ); + assert!(Path::new(&marker(&r)).exists()); + std::fs::remove_file(marker(&r)).unwrap(); + + let (url, _) = fake_queen(Mode::Refuse).await; + assert_eq!( + run(&opts(&r, Some(url), &marker(&r)), &mut Vec::new()).await, + 7 + ); + assert!(Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn no_drone_within_the_wait_cancels_the_unit_and_runs_here() { + let r = repo(); + let (url, fake) = fake_queen(Mode::Queued).await; + let started = std::time::Instant::now(); + assert_eq!( + run(&opts(&r, Some(url), &marker(&r)), &mut Vec::new()).await, + 7 + ); + assert!(started.elapsed() < std::time::Duration::from_secs(10)); + assert_eq!(*fake.cancelled.lock().unwrap(), 1); + assert!(Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn a_unit_that_failed_for_infrastructure_reasons_is_run_here_instead() { + let r = repo(); + let (url, _) = fake_queen(Mode::Failed).await; + assert_eq!( + run(&opts(&r, Some(url), &marker(&r)), &mut Vec::new()).await, + 7 + ); + assert!(Path::new(&marker(&r)).exists()); +} + +#[test] +fn the_rust_toolchain_is_read_from_the_tree() { + for (text, want) in [ + ( + "[toolchain]\nchannel = \"1.95.0\"\ncomponents = [\"clippy\"]\n", + Some("1.95.0"), + ), + ( + "# pinned\n[toolchain]\nchannel=\"nightly-2026-01-01\" # nightly\n", + Some("nightly-2026-01-01"), + ), + ("stable\n", Some("stable")), + ("[toolchain]\ncomponents = []\n", None), + ("[toolchain]\nchannel = \"bad toolchain!\"\n", None), + ("", None), + ] { + assert_eq!(parse_toolchain_channel(text).as_deref(), want, "{text:?}"); + } + let r = repo(); + assert_eq!(tree_rust_toolchain(&r.dir, "HEAD^{tree}"), None); + std::fs::create_dir_all(r.dir.join("rust")).unwrap(); + std::fs::write( + r.dir.join("rust/rust-toolchain.toml"), + "[toolchain]\nchannel = \"1.95.0\"\n", + ) + .unwrap(); + git(&r.dir, &["add", "."]); + git(&r.dir, &["commit", "-q", "-m", "toolchain"]); + assert_eq!( + tree_rust_toolchain(&r.dir, "HEAD^{tree}").as_deref(), + Some("1.95.0") + ); + // The tree asked about decides, not the working directory. + assert_eq!(tree_rust_toolchain(&r.dir, "HEAD~1^{tree}"), None); +} + +#[tokio::test] +async fn the_submission_carries_the_trees_toolchain() { + let r = repo(); + std::fs::write( + r.dir.join("rust-toolchain.toml"), + "[toolchain]\nchannel = \"1.95.0\"\n", + ) + .unwrap(); + git(&r.dir, &["add", "."]); + git(&r.dir, &["commit", "-q", "-m", "toolchain"]); + git(&r.dir, &["push", "-q", "origin", "main"]); + let (url, fake) = fake_queen(Mode::Receipt).await; + assert_eq!( + run(&opts(&r, Some(url), &marker(&r)), &mut Vec::new()).await, + 0 + ); + assert_eq!( + fake.submitted.lock().unwrap()[0]["rust_toolchain"], + "1.95.0" + ); +} diff --git a/tests/support/mod.rs b/tests/support/mod.rs index 89952f0..d5b2d90 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -272,7 +272,24 @@ async fn delete_runner( Ok(StatusCode::NO_CONTENT) } +/// Records wake-on-LAN sends instead of using the network. +#[derive(Default)] +pub struct RecordingWake(pub Mutex)>>); + +impl borg::wake::WakeSender for RecordingWake { + fn send<'a>( + &'a self, + mac: [u8; 6], + to: &'a [std::net::SocketAddrV4], + ) -> borg::wake::BoxFuture<'a, Result, String>> { + self.0.lock().unwrap().push((mac, to.to_vec())); + let sent = to.to_vec(); + Box::pin(async move { Ok(sent) }) + } +} + pub struct Harness { + pub wakes: Arc, pub mock: Shared, pub queen: Arc, pub url: String, @@ -298,6 +315,12 @@ pub fn capacity() -> Capacity { impl Harness { /// `extra` is merged into the queen config; warm floors default to none. pub async fn start(extra: Value) -> Self { + Self::start_with(extra, borg::tailnet::StaticIdentities::default()).await + } + + /// Like `start`, with a fake tailnet identity table (callers all connect + /// from 127.0.0.1) behind `tailnet_enroll` when the config enables it. + pub async fn start_with(extra: Value, identities: borg::tailnet::StaticIdentities) -> Self { let mock: Shared = Arc::default(); let (private_pem, public_pem) = generate_key(); lock(&mock).public_pem = public_pem; @@ -340,17 +363,31 @@ impl Harness { borg::janitor::Janitor::new(&gh_url, borg::secret::Secret::new(JANITOR_TOKEN), settings) .unwrap() }); - let mut queen = Queen::new(cfg, store, github); + let wakes = Arc::new(RecordingWake::default()); + let mut queen = Queen::new(cfg, store, github).with_wake_sender(wakes.clone()); if let Some(j) = janitor { queen = queen.with_janitor(j); } + if let Some(settings) = queen.config().tailnet_enroll.clone() { + queen = queen.with_tailnet_enroll(borg::tailnet::TailnetEnroll { + verifier: Arc::new(identities), + settings, + }); + } let queen = Arc::new(queen); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let url = format!("http://{}", listener.local_addr().unwrap()); let router = api::router(Arc::clone(&queen)); - tokio::spawn(async move { axum::serve(listener, router).await }); + tokio::spawn(async move { + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .await + }); Self { + wakes, mock, queen, url, @@ -394,6 +431,24 @@ impl Harness { .unwrap() } + /// Heartbeat with extra request fields (e.g. `bids`) merged in. + pub async fn heartbeat_with(&self, token: &str, extra: Value) -> Value { + let mut body = json!({"capacity": capacity()}); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + body[k] = v; + } + self.http + .post(format!("{}/v1/heartbeat", self.url)) + .bearer_auth(token) + .json(&body) + .send() + .await + .unwrap() + .json() + .await + .unwrap() + } + pub async fn metrics(&self) -> String { self.http .get(format!("{}/metrics", self.url)) diff --git a/tests/tailnet.rs b/tests/tailnet.rs new file mode 100644 index 0000000..6942e05 --- /dev/null +++ b/tests/tailnet.rs @@ -0,0 +1,116 @@ +mod support; + +use std::net::{IpAddr, Ipv4Addr}; + +use borg::tailnet::{StaticIdentities, TailnetIdentity}; +use serde_json::{Value, json}; +use support::Harness; + +fn identities(node: &str, tags: &[&str], user: Option<&str>) -> StaticIdentities { + let mut m = StaticIdentities::default(); + m.0.insert( + IpAddr::V4(Ipv4Addr::LOCALHOST), + TailnetIdentity { + node_name: node.to_owned(), + tags: tags.iter().map(|t| (*t).to_owned()).collect(), + user: user.map(str::to_owned), + }, + ); + m +} + +fn enabled() -> Value { + json!({"tailnet_enroll": {"allowed_tags": ["tag:ci"], "allowed_users": ["dev@example.test"]}}) +} + +async fn enroll(h: &Harness) -> reqwest::Response { + h.http + .post(format!("{}/v1/enroll", h.url)) + .send() + .await + .unwrap() +} + +#[tokio::test] +async fn an_allowlisted_tagged_node_gets_a_working_drone_token() { + let h = Harness::start_with( + enabled(), + identities("Host_A.example.invalid.", &["tag:ci"], None), + ) + .await; + let resp = enroll(&h).await; + assert_eq!(resp.status(), 200); + let body: Value = resp.json().await.unwrap(); + assert_eq!(body["name"], "host-a"); + let token = body["token"].as_str().unwrap().to_owned(); + assert!(token.starts_with("borg_")); + + // The issued token authenticates as that drone. + assert_eq!(h.heartbeat(&token, &[]).await.status(), 200); + assert_eq!(h.status().await["drones"][0]["name"], "host-a"); +} + +#[tokio::test] +async fn an_untagged_node_enrolls_only_for_an_allowed_user() { + let h = Harness::start_with( + enabled(), + identities("laptop.example.invalid.", &[], Some("dev@example.test")), + ) + .await; + assert_eq!(enroll(&h).await.status(), 200); + + let h = Harness::start_with( + enabled(), + identities("laptop.example.invalid.", &[], Some("eve@example.test")), + ) + .await; + assert_eq!(enroll(&h).await.status(), 403); +} + +#[tokio::test] +async fn nodes_outside_the_allowlist_or_the_tailnet_are_refused() { + let h = Harness::start_with( + enabled(), + identities("host-a.example.invalid.", &["tag:other"], None), + ) + .await; + assert_eq!(enroll(&h).await.status(), 403); + assert!(h.status().await["drones"].as_array().unwrap().is_empty()); + + // Not in the verifier's table at all: not a tailnet node. + let h = Harness::start_with(enabled(), StaticIdentities::default()).await; + assert_eq!(enroll(&h).await.status(), 403); +} + +#[tokio::test] +async fn a_second_enrollment_of_the_same_node_is_a_conflict_and_keeps_the_first_token() { + let h = Harness::start_with( + enabled(), + identities("host-a.example.invalid.", &["tag:ci"], None), + ) + .await; + let first: Value = enroll(&h).await.json().await.unwrap(); + assert_eq!(enroll(&h).await.status(), 409); + let token = first["token"].as_str().unwrap(); + assert_eq!(h.heartbeat(token, &[]).await.status(), 200); +} + +#[tokio::test] +async fn an_unusable_node_name_is_unprocessable() { + let h = Harness::start_with( + enabled(), + identities("---.example.invalid.", &["tag:ci"], None), + ) + .await; + assert_eq!(enroll(&h).await.status(), 422); +} + +#[tokio::test] +async fn the_endpoint_does_not_exist_unless_configured() { + let h = Harness::start_with( + json!({}), + identities("host-a.example.invalid.", &["tag:ci"], None), + ) + .await; + assert_eq!(enroll(&h).await.status(), 404); +} diff --git a/tests/units.rs b/tests/units.rs new file mode 100644 index 0000000..0dc87e2 --- /dev/null +++ b/tests/units.rs @@ -0,0 +1,599 @@ +mod support; + +use borg::enroll; +use serde_json::{Value, json}; +use support::Harness; +use time::OffsetDateTime; + +const TREE: &str = "0123456789abcdef0123456789abcdef01234567"; + +fn cfg() -> Value { + json!({"drone_stale_secs": 1_000_000, "units": {"enabled": true, "lease_ttl_secs": 100, "max_attempts": 2}}) +} + +struct Rig { + h: Harness, + client: String, + base: OffsetDateTime, +} + +impl Rig { + async fn new(extra: Value) -> Self { + let mut c = cfg(); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + c[k] = v; + } + let h = Harness::start(c).await; + let client = enroll::new_token(); + h.queen + .store() + .enroll_client("agent", &enroll::hash_token(&client), 0) + .await + .unwrap(); + Self { + h, + client, + base: OffsetDateTime::now_utc(), + } + } + + async fn tick(&self, secs: i64) { + self.h + .queen + .tick(self.base + time::Duration::seconds(secs)) + .await + .unwrap(); + } + + fn unit(&self, extra: Value) -> Value { + let mut u = json!({"repo": "example-org/app", "tree_sha": TREE, "argv": ["cargo", "check"], + "platform": "linux-x86_64", "toolchain": "rust-1.95", "class": "light"}); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + u[k] = v; + } + u + } + + async fn submit_as(&self, token: &str, body: &Value) -> (u16, Value) { + let r = self + .h + .http + .post(format!("{}/v1/units", self.h.url)) + .bearer_auth(token) + .json(body) + .send() + .await + .unwrap(); + (r.status().as_u16(), r.json().await.unwrap_or(Value::Null)) + } + + async fn submit(&self, body: &Value) -> (u16, Value) { + self.submit_as(&self.client, body).await + } + + async fn get(&self, path: &str) -> reqwest::Response { + self.h + .http + .get(format!("{}{path}", self.h.url)) + .bearer_auth(&self.client) + .send() + .await + .unwrap() + } + + async fn drone(&self, name: &str) -> String { + let t = self.h.enroll(name).await; + self.h.heartbeat(&t, &[]).await; + t + } + + async fn beat(&self, token: &str, active: &[&str]) -> Value { + self.h + .heartbeat_with(token, json!({"active_units": active})) + .await + } + + async fn post_log(&self, token: &str, id: &str, offset: u64, data: &str) -> u16 { + self.h + .http + .post(format!("{}/v1/units/{id}/log?offset={offset}", self.h.url)) + .bearer_auth(token) + .body(data.to_owned()) + .send() + .await + .unwrap() + .status() + .as_u16() + } + + async fn post_result(&self, token: &str, id: &str, code: i32) -> u16 { + self.h + .http + .post(format!("{}/v1/units/{id}/result", self.h.url)) + .bearer_auth(token) + .json(&json!({"exit_code": code, "duration_secs": 12})) + .send() + .await + .unwrap() + .status() + .as_u16() + } +} + +#[tokio::test] +async fn a_unit_is_placed_run_logged_finished_and_then_served_as_a_receipt() { + let r = Rig::new(json!({})).await; + let d1 = r.drone("d1").await; + + let (code, sub) = r.submit(&r.unit(json!({}))).await; + assert_eq!((code, sub["outcome"].as_str()), (201, Some("new")), "{sub}"); + let id = sub["unit_id"].as_str().unwrap().to_owned(); + + // Placement happens on the tick; the drone then hears about it. + assert!( + r.beat(&d1, &[]).await["units"] + .as_array() + .unwrap() + .is_empty() + ); + r.tick(5).await; + let hb = r.beat(&d1, &[]).await; + let grant = &hb["units"][0]; + assert_eq!(grant["unit_id"], id.as_str()); + assert_eq!(grant["tree_sha"], TREE); + assert_eq!(grant["argv"], json!(["cargo", "check"])); + // A placed unit holds capacity: the drone's free memory in /status drops by + // the unit's class (light: 2.5 GiB of the 28 GiB budget). + let status = r.h.status().await; + let d = status["drones"] + .as_array() + .unwrap() + .iter() + .find(|d| d["name"] == "d1") + .unwrap(); + assert_eq!(d["free_mem_mib"], 28 * 1024 - 2560); + // Resent until acknowledged; acknowledging starts it. + assert_eq!(r.beat(&d1, &[]).await["units"].as_array().unwrap().len(), 1); + assert!( + r.beat(&d1, &[&id]).await["units"] + .as_array() + .unwrap() + .is_empty() + ); + assert_eq!( + r.get(&format!("/v1/units/{id}")) + .await + .json::() + .await + .unwrap()["state"], + "running" + ); + + // The drone streams the log; the client reads it from any offset. + assert_eq!(r.post_log(&d1, &id, 0, "compiling a\n").await, 200); + assert_eq!(r.post_log(&d1, &id, 12, "compiling b\n").await, 200); + assert_eq!(r.post_log(&d1, &id, 99, "gap").await, 409); + let log = r.get(&format!("/v1/units/{id}/log?from=0")).await; + assert_eq!(log.headers()["x-borg-next"], "24"); + assert_eq!(log.headers()["x-borg-state"], "running"); + assert!(log.headers().get("x-borg-exit").is_none()); + assert_eq!(log.text().await.unwrap(), "compiling a\ncompiling b\n"); + let tail = r.get(&format!("/v1/units/{id}/log?from=12")).await; + assert_eq!(tail.text().await.unwrap(), "compiling b\n"); + + assert_eq!(r.post_result(&d1, &id, 0).await, 204); + let done = r.get(&format!("/v1/units/{id}/log?from=24")).await; + assert_eq!(done.headers()["x-borg-exit"], "0"); + assert_eq!(done.headers()["x-borg-state"], "done"); + let view = r + .get(&format!("/v1/units/{id}")) + .await + .json::() + .await + .unwrap(); + assert_eq!(view["receipt"]["tree_sha"], TREE); + assert_eq!(view["receipt"]["drone"], "d1"); + + // The same work again is answered from the receipt: nothing runs. + let (code, again) = r.submit(&r.unit(json!({}))).await; + assert_eq!((code, again["outcome"].as_str()), (200, Some("receipt"))); + assert_eq!(again["unit_id"], id.as_str()); + assert_eq!(again["receipt"]["exit_code"], 0); + // A different tree is new work. + let other = r.unit(json!({"tree_sha": "f".repeat(40)})); + assert_eq!(r.submit(&other).await.1["outcome"], "new"); + + let metrics = r.h.metrics().await; + assert!( + metrics.contains("borg_unit_receipt_hits_total 1"), + "{metrics}" + ); + assert!(metrics.contains("borg_units{state=\"done\"} 1")); + assert_eq!(r.h.status().await["units"]["receipt_hits"], 1); +} + +#[tokio::test] +async fn identical_submissions_join_the_running_unit() { + let r = Rig::new(json!({})).await; + r.drone("d1").await; + let (_, a) = r.submit(&r.unit(json!({}))).await; + let (code, b) = r.submit(&r.unit(json!({}))).await; + assert_eq!((code, b["outcome"].as_str()), (200, Some("joined"))); + assert_eq!(a["unit_id"], b["unit_id"]); + assert_eq!( + r.h.metrics() + .await + .matches("borg_unit_joins_total 1") + .count(), + 1 + ); +} + +#[tokio::test] +async fn placement_honors_platform_quarantine_and_class_capacity() { + let r = Rig::new(json!({})).await; + let d1 = r.drone("d1").await; + // A macOS unit has no macOS drone: it waits. + let (_, mac) = r + .submit(&r.unit(json!({"platform": "macos-aarch64"}))) + .await; + r.tick(1).await; + let id = mac["unit_id"].as_str().unwrap(); + assert_eq!( + r.h.queen + .store() + .unit(id) + .await + .unwrap() + .unwrap() + .state + .as_str(), + "queued" + ); + assert!( + r.beat(&d1, &[]).await["units"] + .as_array() + .unwrap() + .is_empty() + ); + // A macOS drone appears (the harness sends linux unless overridden). + let m1 = r.h.enroll("m1").await; + let mut cap = serde_json::to_value(support::capacity()).unwrap(); + cap["os"] = json!("macos"); + cap["arch"] = json!("arm64"); + r.h.heartbeat_with(&m1, json!({"capacity": cap})).await; + r.tick(2).await; + let hb = r.h.heartbeat_with(&m1, json!({"capacity": cap})).await; + assert_eq!(hb["units"][0]["unit_id"], id); + + // A quarantined drone gets nothing. + r.h.queen + .store() + .set_quarantine("d1", Some("test")) + .await + .unwrap(); + r.submit(&r.unit(json!({"tree_sha": "e".repeat(40)}))).await; + r.tick(3).await; + assert!( + r.beat(&d1, &[]).await["units"] + .as_array() + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn units_that_are_never_acknowledged_or_whose_drone_goes_silent_are_retried_then_fail() { + let r = Rig::new(json!({})).await; + let d1 = r.drone("d1").await; + let (_, sub) = r.submit(&r.unit(json!({}))).await; + let id = sub["unit_id"].as_str().unwrap().to_owned(); + let state = || async { r.h.queen.store().unit(&id).await.unwrap().unwrap() }; + + r.tick(1).await; // placed (attempt 1); drone heartbeats stay fresh in this fake time + assert_eq!(state().await.state.as_str(), "leased"); + // Never acknowledged within lease_ttl (100s): back in the queue, then placed again. + r.h.heartbeat(&d1, &[]).await; + r.tick(150).await; + let u = state().await; + assert_eq!((u.attempt, u.state.as_str()), (2, "leased"), "{u:?}"); + // Out of attempts (2): the next expiry fails it. + r.tick(400).await; + let u = state().await; + assert_eq!(u.state.as_str(), "failed"); + assert!(u.failure.unwrap().contains("out of attempts")); + // Failed units are not receipts. + let (_, again) = r.submit(&r.unit(json!({}))).await; + assert_eq!(again["outcome"], "new"); +} + +#[tokio::test] +async fn running_units_time_out_and_cancel_reaches_the_drone() { + let r = Rig::new(json!({})).await; + let d1 = r.drone("d1").await; + let (_, sub) = r.submit(&r.unit(json!({"timeout_secs": 30}))).await; + let id = sub["unit_id"].as_str().unwrap().to_owned(); + r.tick(1).await; + r.beat(&d1, &[]).await; + r.beat(&d1, &[&id]).await; + + // Cancelling a running unit flags it for the drone. + let c = + r.h.http + .post(format!("{}/v1/units/{id}/cancel", r.h.url)) + .bearer_auth(&r.client) + .send() + .await + .unwrap(); + assert_eq!(c.status(), 204); + assert_eq!(r.beat(&d1, &[&id]).await["cancel_units"], json!([id])); + + // Past its 30 second limit it fails by itself. + r.tick(60).await; + let u = r.h.queen.store().unit(&id).await.unwrap().unwrap(); + assert_eq!(u.state.as_str(), "failed"); + assert!(u.failure.unwrap().contains("timed out")); + // A queued unit cancels at once. + let (_, q) = r.submit(&r.unit(json!({"tree_sha": "d".repeat(40)}))).await; + let qid = q["unit_id"].as_str().unwrap(); + r.h.http + .post(format!("{}/v1/units/{qid}/cancel", r.h.url)) + .bearer_auth(&r.client) + .send() + .await + .unwrap(); + assert_eq!( + r.h.queen + .store() + .unit(qid) + .await + .unwrap() + .unwrap() + .failure + .as_deref(), + Some("cancelled") + ); +} + +#[tokio::test] +async fn auth_validation_limits_and_the_disabled_default() { + let r = Rig::new(json!({"units": {"enabled": true, "max_queued_per_client": 1}})).await; + let d1 = r.drone("d1").await; + // No token, a drone token and a junk token cannot submit. + for token in ["", d1.as_str(), "borg_nope"] { + assert_eq!( + r.submit_as(token, &r.unit(json!({}))).await.0, + 401, + "{token:?}" + ); + } + // Bad bodies are 400. + for bad in [ + json!({"tree_sha": "xyz"}), + json!({"argv": []}), + json!({"repo": "x"}), + ] { + assert_eq!(r.submit(&r.unit(bad)).await.0, 400); + } + // The per-client limit: a second distinct unit is refused, but joining is not. + assert_eq!(r.submit(&r.unit(json!({}))).await.0, 201); + assert_eq!( + r.submit(&r.unit(json!({"tree_sha": "c".repeat(40)}))) + .await + .0, + 429 + ); + assert_eq!(r.submit(&r.unit(json!({}))).await.1["outcome"], "joined"); + + // A drone cannot write another drone's unit, and clients cannot post results. + let (_, sub) = r.submit(&r.unit(json!({}))).await; + let id = sub["unit_id"].as_str().unwrap(); + r.tick(1).await; + let d2 = r.drone("d2").await; + assert_eq!(r.post_log(&d2, id, 0, "x").await, 404); + assert_eq!(r.post_result(&r.client, id, 0).await, 401); + assert_eq!(r.get(&format!("/v1/units/{id}")).await.status(), 200); + + // Off by default. + let off = Harness::start(json!({})).await; + let resp = off + .http + .post(format!("{}/v1/units", off.url)) + .bearer_auth("x") + .json(&r.unit(json!({}))) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 404); +} + +#[tokio::test] +async fn a_log_long_poll_returns_as_soon_as_bytes_arrive() { + let r = Rig::new(json!({})).await; + let d1 = r.drone("d1").await; + let (_, sub) = r.submit(&r.unit(json!({}))).await; + let id = sub["unit_id"].as_str().unwrap().to_owned(); + r.tick(1).await; + r.beat(&d1, &[]).await; + r.beat(&d1, &[&id]).await; + + let url = format!("{}/v1/units/{id}/log?from=0&wait=10", r.h.url); + let http = r.h.http.clone(); + let client = r.client.clone(); + let poll = tokio::spawn(async move { + http.get(url) + .bearer_auth(client) + .send() + .await + .unwrap() + .text() + .await + .unwrap() + }); + tokio::time::sleep(std::time::Duration::from_millis(400)).await; + r.post_log(&d1, &id, 0, "late bytes").await; + let started = std::time::Instant::now(); + assert_eq!(poll.await.unwrap(), "late bytes"); + assert!(started.elapsed() < std::time::Duration::from_secs(5)); +} + +#[tokio::test] +async fn units_on_a_drone_that_goes_silent_return_to_the_queue() { + let r = Rig::new(json!({"drone_stale_secs": 60})).await; + let d1 = r.drone("d1").await; + let (_, sub) = r.submit(&r.unit(json!({}))).await; + let id = sub["unit_id"].as_str().unwrap().to_owned(); + r.tick(1).await; + r.beat(&d1, &[]).await; + r.beat(&d1, &[&id]).await; + assert_eq!( + r.h.queen + .store() + .unit(&id) + .await + .unwrap() + .unwrap() + .state + .as_str(), + "running" + ); + // No heartbeat for longer than drone_stale_secs: the unit is requeued and, + // with no fresh drone, waits. + r.tick(200).await; + let u = r.h.queen.store().unit(&id).await.unwrap().unwrap(); + assert_eq!((u.state.as_str(), u.drone, u.attempt), ("queued", None, 1)); +} + +impl Rig { + async fn put_bundle(&self, token: &str, id: &str, data: Vec) -> u16 { + self.h + .http + .put(format!("{}/v1/units/{id}/bundle", self.h.url)) + .bearer_auth(token) + .body(data) + .send() + .await + .unwrap() + .status() + .as_u16() + } + + async fn get_bundle(&self, token: &str, id: &str) -> (u16, Vec) { + let r = self + .h + .http + .get(format!("{}/v1/units/{id}/bundle", self.h.url)) + .bearer_auth(token) + .send() + .await + .unwrap(); + (r.status().as_u16(), r.bytes().await.unwrap().to_vec()) + } +} + +#[tokio::test] +async fn a_unit_waits_for_its_bundle_then_the_drone_fetches_it() { + let r = Rig::new(json!({"units": {"enabled": true, "max_bundle_bytes": 1000}})).await; + let d1 = r.drone("d1").await; + let (_, sub) = r.submit(&r.unit(json!({"bundle": true}))).await; + let id = sub["unit_id"].as_str().unwrap().to_owned(); + + // Not placeable until the bundle arrives. + r.tick(1).await; + assert!( + r.beat(&d1, &[]).await["units"] + .as_array() + .unwrap() + .is_empty() + ); + let u = r + .get(&format!("/v1/units/{id}")) + .await + .json::() + .await + .unwrap(); + assert_eq!( + (u["state"].as_str(), u["awaiting_bundle"].as_bool()), + (Some("queued"), Some(true)) + ); + + // Size cap, a non-owner, and an empty body are refused. + assert_eq!(r.put_bundle(&r.client, &id, vec![7; 2000]).await, 413); + let other = enroll::new_token(); + r.h.queen + .store() + .enroll_client("other", &enroll::hash_token(&other), 0) + .await + .unwrap(); + assert_eq!(r.put_bundle(&other, &id, b"x".to_vec()).await, 404); + assert_eq!(r.put_bundle(&r.client, &id, vec![]).await, 400); + assert_eq!(r.put_bundle("", &id, b"x".to_vec()).await, 401); + + let bundle = b"# v2 git bundle\nfake".to_vec(); + assert_eq!(r.put_bundle(&r.client, &id, bundle.clone()).await, 204); + r.tick(2).await; + let grant = &r.beat(&d1, &[]).await["units"][0]; + assert_eq!(grant["unit_id"], id.as_str()); + assert_eq!(grant["bundle_bytes"], bundle.len()); + + // Only the unit's drone can fetch it. + assert_eq!(r.get_bundle(&d1, &id).await, (200, bundle.clone())); + let d2 = r.drone("d2").await; + assert_eq!(r.get_bundle(&d2, &id).await.0, 404); + assert_eq!(r.get_bundle(&r.client, &id).await.0, 401); + // Once placed, a late upload is refused. + assert_eq!(r.put_bundle(&r.client, &id, bundle).await, 409); +} + +#[tokio::test] +async fn a_bundle_that_never_arrives_fails_the_unit_and_bundles_go_with_their_unit() { + let r = Rig::new(json!({})).await; + let d1 = r.drone("d1").await; + let (_, a) = r.submit(&r.unit(json!({"bundle": true}))).await; + let a = a["unit_id"].as_str().unwrap().to_owned(); + r.tick(30).await; + assert_eq!( + r.h.queen + .store() + .unit(&a) + .await + .unwrap() + .unwrap() + .state + .as_str(), + "queued" + ); + r.tick(100).await; // bundle_wait_secs is 60 + let u = r.h.queen.store().unit(&a).await.unwrap().unwrap(); + assert_eq!(u.state.as_str(), "failed"); + assert!(u.failure.unwrap().contains("bundle never arrived")); + + // Cancelling deletes the bundle at once. + let (_, b) = r + .submit(&r.unit(json!({"bundle": true, "tree_sha": "e".repeat(40)}))) + .await; + let b = b["unit_id"].as_str().unwrap().to_owned(); + assert_eq!(r.put_bundle(&r.client, &b, b"bundle".to_vec()).await, 204); + r.h.http + .post(format!("{}/v1/units/{b}/cancel", r.h.url)) + .bearer_auth(&r.client) + .send() + .await + .unwrap(); + assert!(r.h.queen.store().unit_bundle(&b).await.unwrap().is_none()); + + // A finished unit keeps its bundle for the life of its receipt, then both go. + let (_, c) = r + .submit(&r.unit(json!({"bundle": true, "tree_sha": "d".repeat(40)}))) + .await; + let c = c["unit_id"].as_str().unwrap().to_owned(); + r.put_bundle(&r.client, &c, b"bundle".to_vec()).await; + r.tick(101).await; + r.beat(&d1, &[]).await; + r.beat(&d1, &[&c]).await; + assert_eq!(r.post_result(&d1, &c, 0).await, 204); + assert!(r.h.queen.store().unit_bundle(&c).await.unwrap().is_some()); + r.tick(8 * 86_400).await; // past keep_secs (7 days) + assert!(r.h.queen.store().unit(&c).await.unwrap().is_none()); + assert!(r.h.queen.store().unit_bundle(&c).await.unwrap().is_none()); +} diff --git a/tests/update.rs b/tests/update.rs new file mode 100644 index 0000000..726bfb2 --- /dev/null +++ b/tests/update.rs @@ -0,0 +1,640 @@ +//! `borg update` against a fake manifest/artifact server, a fake service +//! manager and a fake health probe. + +use std::collections::BTreeMap; +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::Router; +use axum::extract::State; +use axum::http::StatusCode; +use axum::routing::get; +use borg::drone::exec::{Exec, Output}; +use borg::update::health::{HealthProbe, HealthRecord}; +use borg::update::manifest::{self, Artifact, Manifest, ManifestError, RunnerRelease}; +use borg::update::service::ServiceManager; +use borg::update::swap::{self, SwapError}; +use borg::update::{Options, Outcome, UpdateError}; +use sha2::{Digest, Sha256}; + +const NEW_BINARY: &[u8] = b"#!/bin/sh\necho new borg\n"; +const RUNNER_TARBALL: &[u8] = b"fake runner tarball"; + +fn sha(bytes: &[u8]) -> String { + hex::encode(Sha256::digest(bytes)) +} + +#[derive(Clone)] +struct Served { + hits: Arc, +} + +async fn binary(State(s): State) -> Vec { + s.hits.fetch_add(1, Ordering::SeqCst); + NEW_BINARY.to_vec() +} + +/// Serves `/borg` (the new binary) and `/missing` (404). +async fn serve() -> (String, Served) { + let served = Served { + hits: Arc::default(), + }; + let app = Router::new() + .route("/borg", get(binary)) + .route("/missing", get(|| async { StatusCode::NOT_FOUND })) + .with_state(served.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(listener, app).await }); + (url, served) +} + +fn manifest(base: &str, version: &str, binary_sha: &str, runner: Option<(&str, &str)>) -> Manifest { + let mut targets = BTreeMap::new(); + targets.insert( + "linux-x64-musl".to_owned(), + Artifact { + url: format!("{base}/borg"), + sha256: binary_sha.to_owned(), + }, + ); + Manifest { + schema: manifest::SCHEMA, + version: version.to_owned(), + targets, + runner: runner.map(|(v, s)| RunnerRelease { + version: v.to_owned(), + sha256: BTreeMap::from([("linux-x64".to_owned(), s.to_owned())]), + }), + } +} + +/// Records what the service was asked to do, and what the installed binary +/// looked like at that moment. +struct FakeService { + bin: PathBuf, + events: Arc>>, + fail_stop: bool, + /// Make the nth start (1-based) fail. + fail_start_number: Option, + starts: AtomicU32, +} + +impl FakeService { + fn log(&self, what: &str) { + let content = fs::read_to_string(&self.bin).unwrap_or_default(); + self.events + .lock() + .unwrap() + .push(format!("{what}:{}", content.trim())); + } +} + +impl ServiceManager for FakeService { + fn stop(&self) -> Result<(), String> { + self.log("stop"); + if self.fail_stop { + Err("unit refused to stop".into()) + } else { + Ok(()) + } + } + fn start(&self) -> Result<(), String> { + self.log("start"); + let nth = self.starts.fetch_add(1, Ordering::SeqCst) as usize + 1; + if self.fail_start_number == Some(nth) { + Err("failed to start".into()) + } else { + Ok(()) + } + } +} + +/// Healthy for `version` only if `healthy_for` says so. +struct FakeHealth { + healthy_for: Option, +} + +impl HealthProbe for FakeHealth { + fn current(&self) -> Option { + self.healthy_for.clone().map(|version| HealthRecord { + version, + at: 1, + pid: 1, + }) + } +} + +struct Rig { + dir: tempfile::TempDir, + bin: PathBuf, + events: Arc>>, +} + +impl Rig { + fn new() -> Self { + let dir = tempfile::tempdir().unwrap(); + let bin = dir.path().join("borg"); + fs::write(&bin, "v1\n").unwrap(); + Self { + dir, + bin, + events: Arc::default(), + } + } + + fn service(&self) -> FakeService { + FakeService { + bin: self.bin.clone(), + events: Arc::clone(&self.events), + fail_stop: false, + fail_start_number: None, + starts: AtomicU32::new(0), + } + } + + fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } + + fn options<'a>(&'a self, current: &'a str) -> Options<'a> { + Options { + bin: &self.bin, + current_version: current, + requested_version: None, + check: false, + runner_only: false, + allow_insecure: true, + health_timeout: Duration::from_millis(300), + poll: Duration::from_millis(10), + host_os: "linux", + host_arch: "x86_64", + runner_dir: None, + runner_owner: None, + } + } + + fn leftovers(&self) -> Vec { + let mut v: Vec = fs::read_dir(self.dir.path()) + .unwrap() + .filter_map(Result::ok) + .map(|e| e.file_name().to_string_lossy().into_owned()) + .filter(|n| n != "borg" && n != "borg.prev" && n != "jobhome") + .collect(); + v.sort(); + v + } +} + +struct NoExec; +impl Exec for NoExec { + fn run(&self, program: &str, _: &[&str]) -> io::Result { + panic!("unexpected command {program}"); + } +} + +#[tokio::test] +async fn updates_drain_first_then_swap_then_confirm_health_and_keep_the_previous_binary() { + let (base, served) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + let out = borg::update::run( + &rig.options("0.4.0"), + &m, + &rig.service(), + &FakeHealth { + healthy_for: Some("0.5.0".into()), + }, + &NoExec, + ) + .await + .unwrap(); + + assert_eq!( + out, + Outcome::Updated { + from: "0.4.0".into(), + to: "0.5.0".into(), + runner_updated: false + } + ); + // The drain (stop) saw the old binary; the restart saw the new one. + assert_eq!(rig.events(), ["stop:v1", "start:#!/bin/sh\necho new borg"]); + assert_eq!(fs::read(&rig.bin).unwrap(), NEW_BINARY); + assert_eq!( + fs::read_to_string(swap::prev_path(&rig.bin)).unwrap(), + "v1\n", + "borg.prev keeps the old binary" + ); + assert_eq!(served.hits.load(Ordering::SeqCst), 1); + assert!(rig.leftovers().is_empty(), "{:?}", rig.leftovers()); +} + +#[tokio::test] +async fn a_checksum_mismatch_refuses_the_update_before_anything_is_drained() { + let (base, _) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(b"something else"), None); + let err = borg::update::run( + &rig.options("0.4.0"), + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &NoExec, + ) + .await + .unwrap_err(); + assert!( + matches!(err, UpdateError::Swap(SwapError::Checksum { .. })), + "{err}" + ); + assert!( + rig.events().is_empty(), + "the service was never touched: {:?}", + rig.events() + ); + assert_eq!(fs::read_to_string(&rig.bin).unwrap(), "v1\n"); + assert!( + rig.leftovers().is_empty(), + "no partial download is left: {:?}", + rig.leftovers() + ); +} + +#[tokio::test] +async fn a_failed_health_check_rolls_back_and_restarts_the_old_binary() { + let (base, _) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + // The new version never reports healthy. + let out = borg::update::run( + &rig.options("0.4.0"), + &m, + &rig.service(), + &FakeHealth { + healthy_for: Some("0.4.0".into()), + }, + &NoExec, + ) + .await + .unwrap(); + match out { + Outcome::RolledBack { from, to, reason } => { + assert_eq!((from.as_str(), to.as_str()), ("0.4.0", "0.5.0")); + assert!(reason.contains("no health evidence"), "{reason}"); + } + other => panic!("expected a rollback, got {other:?}"), + } + assert_eq!( + rig.events(), + [ + "stop:v1", + "start:#!/bin/sh\necho new borg", + "stop:#!/bin/sh\necho new borg", + "start:v1" + ], + "drain, swap, start, then stop the bad version, restore, start" + ); + assert_eq!(fs::read_to_string(&rig.bin).unwrap(), "v1\n"); +} + +#[tokio::test] +async fn a_service_that_will_not_start_is_rolled_back_too() { + let (base, _) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + let mut service = rig.service(); + service.fail_start_number = Some(1); // only the first start fails; the rollback's start works + let out = borg::update::run( + &rig.options("0.4.0"), + &m, + &service, + &FakeHealth { + healthy_for: Some("0.5.0".into()), + }, + &NoExec, + ) + .await + .unwrap(); + assert!( + matches!(&out, Outcome::RolledBack { reason, .. } if reason.contains("did not start")), + "{out:?}" + ); + assert_eq!(fs::read_to_string(&rig.bin).unwrap(), "v1\n"); +} + +#[tokio::test] +async fn a_drain_that_fails_changes_nothing() { + let (base, _) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + let mut service = rig.service(); + service.fail_stop = true; + let err = borg::update::run( + &rig.options("0.4.0"), + &m, + &service, + &FakeHealth { healthy_for: None }, + &NoExec, + ) + .await + .unwrap_err(); + assert!(matches!(err, UpdateError::Drain(_)), "{err}"); + assert_eq!(fs::read_to_string(&rig.bin).unwrap(), "v1\n"); + assert!(rig.leftovers().is_empty()); +} + +#[tokio::test] +async fn updating_to_the_current_version_does_nothing() { + let (base, served) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + let out = borg::update::run( + &rig.options("0.5.0"), + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &NoExec, + ) + .await + .unwrap(); + assert_eq!( + out, + Outcome::UpToDate { + version: "0.5.0".into() + } + ); + assert!(rig.events().is_empty()); + assert_eq!(served.hits.load(Ordering::SeqCst), 0, "no download"); +} + +#[tokio::test] +async fn check_reports_what_would_change_and_changes_nothing() { + let (base, served) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + let mut opts = rig.options("0.4.0"); + opts.check = true; + let out = borg::update::run( + &opts, + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &NoExec, + ) + .await + .unwrap(); + assert_eq!( + out, + Outcome::Available { + current: "0.4.0".into(), + target: "0.5.0".into(), + runner: None + } + ); + assert!(rig.events().is_empty() && served.hits.load(Ordering::SeqCst) == 0); + assert_eq!(fs::read_to_string(&rig.bin).unwrap(), "v1\n"); + + let up_to_date = borg::update::run( + &{ + let mut o = rig.options("0.5.0"); + o.check = true; + o + }, + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &NoExec, + ) + .await + .unwrap(); + assert_eq!( + up_to_date, + Outcome::UpToDate { + version: "0.5.0".into() + } + ); +} + +#[tokio::test] +async fn a_requested_version_must_match_the_manifest() { + let (base, _) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + let mut opts = rig.options("0.4.0"); + opts.requested_version = Some("0.4.9"); + let err = borg::update::run( + &opts, + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &NoExec, + ) + .await + .unwrap_err(); + assert!(matches!(err, UpdateError::VersionMismatch { .. }), "{err}"); + assert!(rig.events().is_empty()); +} + +#[tokio::test] +async fn a_host_without_a_released_binary_is_a_clear_error() { + let (base, _) = serve().await; + let rig = Rig::new(); + let m = manifest(&base, "0.5.0", &sha(NEW_BINARY), None); + let mut opts = rig.options("0.4.0"); + opts.host_os = "linux"; + opts.host_arch = "riscv64"; + let err = borg::update::run( + &opts, + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &NoExec, + ) + .await + .unwrap_err(); + assert!(matches!(err, UpdateError::UnsupportedHost(_)), "{err}"); +} + +struct Hooked { + calls: Mutex>, +} + +impl Exec for Hooked { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + self.calls + .lock() + .unwrap() + .push(format!("{program} {}", args.join(" "))); + match program { + "curl" => fs::write(args[2], RUNNER_TARBALL)?, + "tar" => fs::write(Path::new(args[3]).join("run.sh"), "#!/bin/sh\n")?, + _ => {} + } + Ok(Output { + success: true, + stdout: String::new(), + }) + } +} + +#[tokio::test] +async fn runner_only_updates_the_runner_without_touching_the_binary_or_the_service() { + let (base, served) = serve().await; + let rig = Rig::new(); + let runner_dir = rig.dir.path().join("jobhome/runner"); + // A newer binary is available too, but --runner-only leaves it alone. + let m = manifest( + &base, + "0.5.0", + &sha(NEW_BINARY), + Some(("2.338.0", &sha(RUNNER_TARBALL))), + ); + let mut opts = rig.options("0.4.0"); + opts.runner_only = true; + opts.runner_dir = Some(&runner_dir); + opts.runner_owner = Some("borg-job"); + let exec = Hooked { + calls: Mutex::default(), + }; + let out = borg::update::run( + &opts, + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &exec, + ) + .await + .unwrap(); + + assert_eq!( + out, + Outcome::RunnerUpdated { + version: "2.338.0".into() + } + ); + assert!( + rig.events().is_empty(), + "no drain or restart: {:?}", + rig.events() + ); + assert_eq!( + served.hits.load(Ordering::SeqCst), + 0, + "the borg binary was not downloaded" + ); + assert_eq!(fs::read_to_string(&rig.bin).unwrap(), "v1\n"); + let marker = fs::read_to_string(runner_dir.join(".borg-runner-version")).unwrap(); + assert_eq!(marker, format!("2.338.0 {}\n", sha(RUNNER_TARBALL))); + assert!( + exec.calls + .lock() + .unwrap() + .iter() + .any(|c| c.starts_with("chown -R borg-job:borg-job")) + ); + + // Idempotent: the same manifest again has nothing to do. + let again = borg::update::run( + &opts, + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &exec, + ) + .await + .unwrap(); + assert_eq!( + again, + Outcome::UpToDate { + version: "0.4.0".into() + } + ); +} + +#[tokio::test] +async fn a_binary_update_also_brings_the_runner_up_to_date() { + let (base, _) = serve().await; + let rig = Rig::new(); + let runner_dir = rig.dir.path().join("jobhome/runner"); + let m = manifest( + &base, + "0.5.0", + &sha(NEW_BINARY), + Some(("2.338.0", &sha(RUNNER_TARBALL))), + ); + let mut opts = rig.options("0.4.0"); + opts.runner_dir = Some(&runner_dir); + let out = borg::update::run( + &opts, + &m, + &rig.service(), + &FakeHealth { + healthy_for: Some("0.5.0".into()), + }, + &Hooked { + calls: Mutex::default(), + }, + ) + .await + .unwrap(); + assert_eq!( + out, + Outcome::Updated { + from: "0.4.0".into(), + to: "0.5.0".into(), + runner_updated: true + } + ); + assert!(runner_dir.join("run.sh").exists()); +} + +#[tokio::test] +async fn a_bad_runner_checksum_aborts_before_the_service_is_drained() { + let (base, _) = serve().await; + let rig = Rig::new(); + let runner_dir = rig.dir.path().join("jobhome/runner"); + let m = manifest( + &base, + "0.5.0", + &sha(NEW_BINARY), + Some(("2.338.0", &"0".repeat(64))), + ); + let mut opts = rig.options("0.4.0"); + opts.runner_dir = Some(&runner_dir); + let err = borg::update::run( + &opts, + &m, + &rig.service(), + &FakeHealth { healthy_for: None }, + &Hooked { + calls: Mutex::default(), + }, + ) + .await + .unwrap_err(); + assert!(matches!(err, UpdateError::Runner(_)), "{err}"); + assert!(rig.events().is_empty()); + assert!( + rig.leftovers().is_empty(), + "the verified binary download is cleaned up too: {:?}", + rig.leftovers() + ); +} + +#[tokio::test] +async fn manifests_are_fetched_over_http_only_when_allowed() { + let (base, _) = serve().await; + assert!(matches!( + manifest::fetch(&format!("{base}/missing"), true).await, + Err(ManifestError::Status { status: 404, .. }) + )); + assert!(matches!( + manifest::fetch(&format!("{base}/borg"), false).await, + Err(ManifestError::Insecure(_)) + )); +} diff --git a/tests/wake.rs b/tests/wake.rs new file mode 100644 index 0000000..e638eb9 --- /dev/null +++ b/tests/wake.rs @@ -0,0 +1,89 @@ +mod support; + +use serde_json::{Value, json}; +use support::Harness; +use time::OffsetDateTime; + +fn cfg() -> Value { + json!({ + "drone_stale_secs": 1_000_000_000, + "wake": { + "enabled": true, + "after_secs": 300, + "min_interval_secs": 600, + "broadcast": ["192.0.2.0/24"], + "ports": [9], + "hosts": [ + {"name": "sleeper", "mac": "aa:bb:cc:00:00:01", "classes": ["heavy"]}, + {"name": "other", "mac": "aa:bb:cc:00:00:02"} + ] + } + }) +} + +fn at(base: OffsetDateTime, secs: i64) -> OffsetDateTime { + base + time::Duration::seconds(secs) +} + +#[tokio::test] +async fn sustained_unmet_demand_wakes_a_sleeping_host_once_and_status_shows_it() { + let h = Harness::start(cfg()).await; + let base = OffsetDateTime::now_utc(); + // No drone at all: a queued heavy job cannot be placed. + h.queue_job(1, "borg-ci-heavy-pr", 10); + + h.queen.tick(at(base, 0)).await.unwrap(); + h.queen.tick(at(base, 200)).await.unwrap(); + assert!(h.wakes.0.lock().unwrap().is_empty(), "too early"); + + h.queen.tick(at(base, 310)).await.unwrap(); + { + let sent = h.wakes.0.lock().unwrap(); + assert_eq!(sent.len(), 1); + assert_eq!(sent[0].0, [0xaa, 0xbb, 0xcc, 0, 0, 1]); + assert_eq!(sent[0].1[0].to_string(), "192.0.2.255:9"); + } + // Rate limited: the same unmet demand does not wake again within the interval. + h.queen.tick(at(base, 400)).await.unwrap(); + assert_eq!(h.wakes.0.lock().unwrap().len(), 1); + + let status = h.status().await; + assert_eq!(status["wake"]["wakes_last_hour"], 1); + assert_eq!(status["wake"]["events"][0]["host"], "sleeper"); + assert!(status["wake"]["unmet_secs"]["heavy"].as_u64().unwrap() >= 300); + let metrics = h.metrics().await; + assert!( + metrics.contains("borg_wake_attempts_total{host=\"sleeper\"} 1"), + "{metrics}" + ); +} + +#[tokio::test] +async fn nothing_is_woken_when_the_host_is_already_awake_or_the_feature_is_off() { + // The host is awake: a fresh drone named like it exists (but cannot hold + // a heavy job, so demand stays unmet). + let h = Harness::start(cfg()).await; + let t = h.enroll("sleeper").await; + let mut cap = serde_json::to_value(support::capacity()).unwrap(); + cap["mem_mib"] = json!(4096); + h.heartbeat_with(&t, json!({"capacity": cap})).await; + let base = OffsetDateTime::now_utc(); + h.queue_job(1, "borg-ci-heavy-pr", 10); + for secs in [0, 310, 700, 1500] { + h.queen.tick(at(base, secs)).await.unwrap(); + } + // "other" has no class restriction so it is woken instead of "sleeper". + let sent = h.wakes.0.lock().unwrap().clone(); + assert!( + sent.iter() + .all(|(mac, _)| *mac == [0xaa, 0xbb, 0xcc, 0, 0, 2]), + "{sent:?}" + ); + + let off = Harness::start(json!({"drone_stale_secs": 1_000_000_000})).await; + off.queue_job(1, "borg-ci-heavy-pr", 10); + for secs in [0, 310, 700] { + off.queen.tick(at(base, secs)).await.unwrap(); + } + assert!(off.wakes.0.lock().unwrap().is_empty()); +} From 20c167b1ce454ce1442eb41958b066cc716addba Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sun, 11 Oct 2026 04:54:30 +0000 Subject: [PATCH 07/27] chore: project borg from Mono b271bbbd34fa --- .repository-projection.json | 4 ++-- src/main.rs | 2 ++ tests/admin_cli.rs | 18 ++++++++++++++++-- 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index aa3ad6d..fd66e56 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "7c0e3ee59b046f3f95c09c34cb8e769a031a7fc7", + "sourceSha": "b271bbbd34fa67d4b0b7b0faee44783e1de7ffeb", "destinationRepository": "dx-corp/borg", "priorProjectedBase": "68b3777037221a85f3b0683621dc7a07f1a8e3a1", "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "aa7636d26420d00c729a3d64cc047b6d9ef57bf433e1c8756954b9b2b830da52", + "contentDigest": "edcf6ae3f55554b157b8080376d28bf135d66c62e3ec8f9b10e1b4f0cc1a34b7", "publicationEligible": true } diff --git a/src/main.rs b/src/main.rs index a752cde..4588e3a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -729,6 +729,7 @@ async fn dispatch(command: Command) -> Result<()> { Command::Report(a) => report(a).await, Command::Drone { command } => { tracing_subscriber::fmt() + .with_writer(std::io::stderr) .with_env_filter( tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| "info".into()), @@ -741,6 +742,7 @@ async fn dispatch(command: Command) -> Result<()> { } Command::Queen { command } => { tracing_subscriber::fmt() + .with_writer(std::io::stderr) .with_env_filter( tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| "info".into()), diff --git a/tests/admin_cli.rs b/tests/admin_cli.rs index 158aea4..101eac7 100644 --- a/tests/admin_cli.rs +++ b/tests/admin_cli.rs @@ -2,6 +2,7 @@ use std::fs; use std::path::Path; +use std::process::Command; use borg::drone::exec::{Exec, SystemExec}; @@ -45,8 +46,19 @@ fn enroll_prints_only_the_token_and_list_and_remove_manage_drones() { let dir = tempfile::tempdir().unwrap(); let cfg = write_config(dir.path()); - let (ok, token) = run(&["queen", "enroll", "--config", &cfg, "host-a"]); - assert!(ok); + // Emit SQL diagnostics even on a fast host: a slow-query warning in CI + // must never become part of the token returned on stdout. + #[allow( + clippy::disallowed_methods, + reason = "test fixture: run the borg binary with isolated logging settings and capture both streams" + )] + let output = Command::new(borg_bin()) + .args(["queen", "enroll", "--config", &cfg, "host-a"]) + .env("RUST_LOG", "sqlx::query=debug") + .output() + .unwrap(); + assert!(output.status.success()); + let token = String::from_utf8(output.stdout).unwrap().trim().to_owned(); assert!( token.starts_with("borg_") && token.len() == 5 + 64 @@ -55,6 +67,8 @@ fn enroll_prints_only_the_token_and_list_and_remove_manage_drones() { .all(|b| b.is_ascii_alphanumeric() || b == b'_'), "stdout is exactly the token: {token:?}" ); + let diagnostics = String::from_utf8(output.stderr).unwrap(); + assert!(diagnostics.contains("sqlx::query"), "{diagnostics}"); let (ok, list) = run(&["queen", "drone", "list", "--config", &cfg]); assert!(ok); From 00e3a20414b684041d4ddd5bd2d0091b5dde2d6f Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sun, 11 Oct 2026 07:12:21 +0000 Subject: [PATCH 08/27] chore: project borg from Mono 6efbce212025 --- .repository-projection.json | 8 +- Cargo.lock | 47 +++ Cargo.toml | 1 + README.md | 764 ++++++------------------------------ docs/operations.md | 707 +++++++++++++++++++++++++++++++++ src/config.rs | 133 ++++++- src/configfile.rs | 106 +++++ src/drone/config.rs | 37 +- src/lib.rs | 1 + src/main.rs | 36 +- 10 files changed, 1176 insertions(+), 664 deletions(-) create mode 100644 docs/operations.md create mode 100644 src/configfile.rs diff --git a/.repository-projection.json b/.repository-projection.json index fd66e56..89be64d 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "b271bbbd34fa67d4b0b7b0faee44783e1de7ffeb", + "sourceSha": "6efbce21202579ccfae91687c9aef688428d0164", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "68b3777037221a85f3b0683621dc7a07f1a8e3a1", - "definitionDigest": "0ede17fcc026faff85cd66bb9e941fa3db5f1792165aae945b7b76b57cfd9c9a", + "priorProjectedBase": "2e48c3e0173446001b8e84070574e9add75b0db9", + "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "edcf6ae3f55554b157b8080376d28bf135d66c62e3ec8f9b10e1b4f0cc1a34b7", + "contentDigest": "bfd698d6787a8485a0097877e0b43db394bffb29733a7c97c98abf27c724571a", "publicationEligible": true } diff --git a/Cargo.lock b/Cargo.lock index 70bf517..929492c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -212,6 +212,7 @@ dependencies = [ "thiserror", "time", "tokio", + "toml", "tracing", "tracing-subscriber", ] @@ -1483,6 +1484,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1883,6 +1893,37 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" +dependencies = [ + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow", +] + [[package]] name = "tower" version = "0.5.3" @@ -2268,6 +2309,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/Cargo.toml b/Cargo.toml index 2b40cc4..3e43b72 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ sqlx = { version = "0.8", default-features = false, features = ["runtime-tokio-r thiserror = "2.0" time = { version = "0.3", features = ["formatting", "macros", "parsing"] } tokio = { version = "1.48", features = ["io-util", "macros", "net", "rt-multi-thread", "process", "signal", "sync", "time"] } +toml = { version = "1", default-features = false, features = ["parse", "serde"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } diff --git a/README.md b/README.md index 9aface8..2a2946d 100644 --- a/README.md +++ b/README.md @@ -1,689 +1,153 @@ -# borg +# Borg -borg is a self-organizing CI fleet agent for self-hosted GitHub Actions -runners. Every machine that joins the fleet runs a drone. A single queen -watches the queue, decides how much of each machine to lend to which kind of -job, and hands drones one-shot (ephemeral, JIT) runner configurations. +Borg runs self-hosted GitHub Actions runners across Linux, WSL, and macOS. +A **queen** watches queued jobs and allocates capacity. A **drone** runs on +each host, starts an ephemeral runner for each lease, and cleans up when the +job finishes. Hosts keep their build caches between jobs. -The crate contains the pure core and the queen: +Runner classes (`light`, `medium`, and `heavy`) define CPU, memory, and Docker +requirements. Host owners can cap the resources they lend to the fleet. -| Module | Does | -| --- | --- | -| `probe` | Reports vCPU, memory, free disk, architecture, OS, host kind (WSL, LXC, VM, bare metal), and whether Docker, mise and git are usable. | -| `classes` | Defines the `light`, `medium` and `heavy` runner classes, their resource requirements, and the label scheme. All of it is configuration with defaults. | -| `allocator` | A pure function from queued demand, drone capacity, in-flight leases, warm floors and per-host caps to lease grants with exact label sets. | - -| `queen`, `github`, `store`, `api` | The coordinator: polls queued jobs, runs the allocator, mints single-use JIT runner configs, keeps lease bookkeeping in SQLite and serves the drone API plus `/status` and `/metrics`. | +## Install from source -| `drone` | The agent on every machine: heartbeats to the queen, runs each lease as an ephemeral runner in a private directory, cleans up after every job, and installs itself. | - -## Build +Requires Rust 1.95 or newer. Build and install from the public repository: -``` -cargo build --locked -cargo test --locked +```sh +git clone https://github.com/dx-corp/borg.git +cd borg +cargo install --path . --locked +borg probe ``` -## Probe +`borg probe` reports host capacity and available tools as JSON. Add `--compact` +for a single line. -``` -borg probe # pretty JSON -borg probe --compact # one line -``` +## Start a fleet -## Labels +### 1. Configure the queen -A job asks for a class with a runner label: `` for jobs from -the default branch or a merge queue, and `-pr` for pull -requests. The default prefix is `borg-ci-`; set `label_prefix` in the classes -config to change it (for example `example-ci-`). A runner minted for main -demand carries only the bare label, so GitHub can hand it nothing but main -jobs. A runner minted for pull-request demand carries both labels. +Create a GitHub App with organization self-hosted runners **read and write** +and Actions **read** access to the repositories it will serve. Install the +App on the organization and save its private key on the queen host. -## Classes config +Save this as `queen.toml`, replacing the organization, repository, App ID, +and key path: -`Classes::from_json` accepts partial JSON; omitted fields keep their defaults: +```toml +repos = ["example-org/app"] +db_path = "queen.db" +listen = "127.0.0.1:8080" -```json -{ - "label_prefix": "example-ci-", - "heavy": { "mem_mib": 20480, "vcpu": 6, "docker": true } -} +[github] +org = "example-org" +app_id = "12345" +private_key_path = "/etc/borg/app.pem" ``` -## Queen +Validate the configuration, enroll a host, and start the queen: -``` -borg queen enroll --config queen.json host-a # prints a one-time token -borg queen run --config queen.json +```sh +borg config check queen.toml +umask 077 +borg queen enroll --config queen.toml host-a > host-a.token +borg queen run --config queen.toml ``` -Example `queen.json` (every field except `github`, `repos` and `db_path` is -optional): - -```json -{ - "github": { "org": "example-org", "app_id": "12345", "private_key_path": "/etc/borg/app.pem" }, - "repos": ["example-org/app"], - "db_path": "/var/lib/borg/queen.db", - "listen": "127.0.0.1:8080", - "classes": { "label_prefix": "example-ci-" }, - "starvation_after_secs": 120 -} -``` +Enrollment prints the token once, to stdout; diagnostics go to stderr. Keep +`host-a.token` private. The queen stores its hash. -Operators can stop new leases to one drone with `borg queen quarantine ---config queen.json --reason "" ` and lift it with `borg queen -release --config queen.json `. Set `rollout_target_version` in the queen -config to have `/metrics` and `/status` report how many drones run that borg -version (drones report theirs in every heartbeat; the queen never forces an -update). - -`warm_floors` keep idle runners alive per class and scope; the optional -`warm_schedule` raises them in time-of-day windows (`{"class": "light", "scope": -"pr", "count": 4, "start": "08:00", "end": "20:00", "days": ["mon", "tue"]}`) -at a fixed `utc_offset_minutes`. A window never lowers a floor. - -### Job telemetry - -A drone may add a `telemetry` object to `POST /v1/leases/{id}/result`: - -```json -{ - "exit_code": 1, "duration_secs": 412, - "telemetry": { - "job_name": "rust-test (shard 2)", "repo": "example-org/app", "run_id": 123456789, - "conclusion": "failure", "failure_class": "host", "queue_wait_secs": 18, - "steps": [{ "name": "cargo test", "duration_secs": 401, "conclusion": "failure" }] - } -} -``` +The queen serves plain HTTP. Put its loopback listener behind an HTTPS +endpoint reachable by the drones. The drone requires HTTPS unless explicitly +configured for an insecure connection; use `--ca-cert` to pin a private CA. -`conclusion` is `success`, `failure`, `cancelled`, `skipped`, `timed_out` or -`unknown`; `failure_class` is `host` (machine or runner failure), `code` (the -job's own failure) or `unknown`. Unrecognized values read as `unknown`; at most -64 steps and 120 characters per name are kept; `job_name` and a `owner/name` -`repo` are required, otherwise the telemetry is dropped and the result still -recorded. The queen keeps it for `telemetry_retention_days` (default 30). -`borg report --queen ` (or `--config queen.json` on the queen host) and -`GET /report` summarize time by job and step, queue wait per class, the -cancellation share and host-caused failures per drone. - -#### What a drone-side parser should produce - -The queen validates but does not interpret these fields, so the drone owns -their meaning. Keep it consistent across drones: - -- **Identity.** `repo`, `run_id` and `job_name` come from the job the runner - executed (its job context: repository, run id, job display name), not from - log text a workflow can print. `job_name` is the name as GitHub shows it, - including matrix suffixes. -- **Steps.** One entry per step in execution order, `duration_secs` rounded - down to whole seconds, `conclusion` mapped from the step result (succeeded to - `success`, failed to `failure`, canceled to `cancelled`, skipped to - `skipped`, anything else to `unknown`). Report names only; never log lines, - environment values or step output. -- **Conclusion.** The job's result, mapped the same way; a job the drone killed - at `max_job_secs` is `timed_out`. -- **`queue_wait_secs`.** Seconds from the job being queued on GitHub to the - runner starting it. Omit it when it cannot be determined; do not send 0 as a - placeholder, because the report averages the values it receives. -- **`failure_class`.** `host` when the machine or the runner failed rather than - the job: the runner process was killed by a signal or the kernel (OOM), the - disk filled, Docker was unreachable, the runner lost its connection, or the - lease was torn down by the drone. `code` when the job's own step exited - non-zero on a healthy host. `unknown` when the drone cannot tell. Leave it - `unknown` for jobs that did not fail. -- **Failure behaviour.** Parsing is best effort and must never delay or fail - the result: if anything cannot be read, send the result without `telemetry` - (or with fewer fields), and send it in the same request as `exit_code` and - `duration_secs`. A retried result may be sent again; the queen keeps the - latest. - -### Drone bids and demand pressure - -A drone may add `bids` to its heartbeat: at most one per `(class, scope)` it is -willing to serve. The queen answers every heartbeat with `pressure`, the -demand pressure (0 to 1000) per class, so the drone can set its bids against it. - -```json -{ "class": "medium", "scope": "pr", "available_slots": 2, - "estimated_start_ms": 1500, "cache_affinity": 800, "resource_pressure": 150, - "estimated_cost": 1000, "confidence": 900, "min_pressure": 0, "expires_at": 1760000000 } -``` +### 2. Install a drone -`cache_affinity`, `resource_pressure`, `confidence` and `min_pressure` are -permille (0 to 1000; larger values are clamped); `estimated_cost` is relative -(1000 is fleet typical); `expires_at` is Unix seconds, after which the queen -ignores the bid. `min_pressure` is the drone's response threshold: it does not -want this work while the class's pressure is below it, so a warm idle host -reports 0, a cold host something higher, and a workstation whose owner is -active something high. Starving demand overrides every threshold. - -The allocator drops a bid with no `available_slots` or above the current -pressure, and scores the rest with integers: -`3*cache_affinity - 3*resource_pressure - 2*start - 1*(cost-1000)`, where -`start` is `estimated_start_ms` as permille of 60 s, all scaled by -`confidence/1000`. The highest score wins; ties go to the tightest memory fit, -then the drone name. A drone that sends no bids, or only expired ones, is -treated as before (always eligible, neutral score); a drone that bids but not -for a key declines that key. GitHub still chooses which queued job a runner -takes: bids steer only which host receives a freshly minted runner. - -### Wake-on-LAN - -Optional (`wake` block, off by default). When demand for a class stays unmet for -`after_secs`, the queen sends a magic packet to a listed host that can serve the -class and whose drone is not fresh, rate limited per pass, per hour and per host. -Targets are IPv4 addresses or `addr/len` networks. See `/status` `wake`. - -### Build environment contract - -Optional (`env_contract.enabled`, off by default). Before a CI runner lease or a -work unit is placed on a drone, the queen checks what the drone last reported -about its build environment. Four checks, each `required`, `advisory` or `off` -globally and per class (`env_contract.checks`, `env_contract.classes.`): - -* `compiler_cache`: `RUSTC_WRAPPER` is an allowed wrapper (`allowed_wrappers`, - by file name: `sccache` by default, plus any single-flight wrapper in front of - it), `SCCACHE_DIR` is set, the sccache server is reachable. -* `toolchain`: when the work declares a Rust toolchain, the drone has it. A unit - declares it (`rust_toolchain`, which `borg run` reads from `rust-toolchain.toml` - or `rust/rust-toolchain.toml` at the unit's tree); a class can declare it in - config (`classes..rust_toolchain`). A mismatch is not admitted. -* `cache_store`: stats work, errors per thousand requests are under - `max_cache_error_pm`, the cache directory has at least - `min_cache_dir_free_bytes` free, and a write works. A hit rate at or under - `hit_rate_warn_pm` after `min_requests_for_hit_rate` requests is a *warning* in - `/status`, never a block. -* `resources`: memory available covers the class's memory (`mem_available_pct`) - and CPU pressure is at most `max_cpu_psi_pm`. (Free capacity after active - leases is the allocator's own rule.) - -A required check that fails **or is unknown** makes the drone ineligible for the -class; nothing else changes about placement. Facts older than -`facts_max_age_secs` (180) count as unknown. A drone that has never reported is -skipped unless `require_report`, so enabling the contract does not cut off -drones that predate it. Defaults when enabled: `toolchain` and `resources` -required, `compiler_cache` and `cache_store` advisory. - -`/status` shows each drone's `contract` (the checks with state, reason, observed -values and warnings, and the admission per class with the reason) and -`units.not_admitted` (why a queued unit is not admitted on each matching drone). -Metrics: `borg_env_contract_check{drone,check,state}` and -`borg_env_contract_admitted{drone,class}`. - -**Drone probe spec** (the heartbeat field `env_facts`). The drone only measures; -verdicts and thresholds are the queen's. Every field is optional, and `null` / -absent means "not measured" (never "fine"). Probe as the **job user**, with the -environment a job will get, and cache results about 60 seconds (resend the cached -object with its original `observed_at`; the queen treats facts over 180 seconds -old as unknown, so probe at least that often): - -| Field | How | -| --- | --- | -| `observed_at` | Unix seconds the probes ran. | -| `rustc_wrapper` | `RUSTC_WRAPPER` as a job sees it, resolved to a path through `PATH` if it is a bare name; absent if unset or empty. | -| `sccache_dir` | `SCCACHE_DIR` as a job sees it; absent if unset. | -| `sccache_server_reachable` | Whether the sccache server answers (connect to `SCCACHE_SERVER_PORT` or its socket, 5 s timeout). Never start a server to find out. | -| `sccache_stats_ok` | `sccache --show-stats --stats-format json` exits 0 within 5 s. | -| `cache_requests`, `cache_hits`, `cache_errors` | From that JSON: `compile_requests`, the sum of the `cache_hits.counts` values and the sum of the `cache_errors.counts` values (check against the installed sccache's schema). | -| `cache_dir_free_bytes` | Available bytes (`f_bavail * f_frsize`) on the filesystem holding `SCCACHE_DIR`. | -| `cache_write_ok` | Create, write and remove a small file named `.borg-probe-` in `SCCACHE_DIR`. | -| `toolchains` | For each toolchain `rustup toolchain list` shows, `{channel, rustc_version}`: the channel as `rust-toolchain.toml` would name it (`1.95.0`, `stable`, `nightly-2026-01-01`; strip the host triple) and the `release:` line of `rustc + -vV`. `null` if rustup could not be queried. | -| `mem_available_mib` | Linux `MemAvailable` from `/proc/meminfo`; macOS (free + inactive + speculative pages) times the page size, in MiB. | -| `cpu_psi_some_avg10_pm` | Linux `/proc/pressure/cpu` line `some`, `avg10` times ten, rounded; absent where PSI does not exist. | - -A unit's toolchain (`rust_toolchain`) matches an installed `channel` or -`rustc_version` exactly. - -### `borg run` +Copy the host's token and CA certificate to that host, then install the service: -``` -borg run [--index | --tree ] [--publish] [--class medium] [--wait-secs 60] -- +```sh +sudo "$(command -v borg)" drone install \ + --queen https://queen.example.test:8443 \ + --token-file host-a.token \ + --ca-cert queen-ca.pem ``` -Runs a command on the fleet when a drone has capacity, and otherwise right here; -the exit code is the command's either way. With no `BORG_QUEEN` set it just runs -the command locally, so wrappers can call it unconditionally. - -* Environment: `BORG_QUEEN` (queen URL), `BORG_CLIENT_TOKEN` or - `BORG_CLIENT_TOKEN_FILE` (default `~/.config/borg/client-token`; the token comes - from `borg queen enroll-client`), `BORG_RUN_WAIT_SECS`, `BORG_PUBLISH`, - `BORG_TOOLCHAIN`. -* The tree is `HEAD^{tree}` by default, `git write-tree` of the index with - `--index` (pre-commit hooks), or `--tree `. The unit's identity is the repo, - tree, argv, platform and toolchain, so an identical command that already passed - for this tree is not run again (the receipt is reported, exit 0) and one that is - running is joined. -* The drone has to be able to get the tree. A tree from a pushed commit is in - its git mirror. Otherwise `--publish` uploads a `git bundle` to the queen - (`PUT /v1/units/{id}/bundle`): a synthetic commit of the tree whose parent is the - merge base with `origin/main`, bundled without everything that merge base - already reaches, so it is small. Nothing is pushed to any remote. Without - `--publish`, an unpublished tree runs locally. -* Falls back to running here when the queen is unset or unreachable, refuses the - unit, no drone takes it within `--wait-secs` (the unit is cancelled), the queen - is lost for 60 seconds mid-run, or the unit ends for infrastructure reasons - (cancelled, timed out, out of attempts). A command that ran on a drone and - exited nonzero is a real failure and its code is returned. Ctrl-C cancels a - unit this invocation created. -### Work units (gate and hook runs) - -Optional (`units.enabled`, off by default). Clients such as `borg run`, the -pre-commit hook and `mq` submit whole commands; drones run them. Enroll a client -with `borg queen enroll-client ` (prints a token once). - -Client API (bearer: client token): - -* `POST /v1/units` with `{repo, tree_sha, argv, platform, toolchain, class, timeout_secs?, force?}` - returns `201` `{outcome: "new", unit_id, ...}`, or `200` with `outcome: - "joined"` (an identical unit is queued or running) or `"receipt"` (an - identical unit succeeded recently; `receipt` has the proof). The key is - (repo, tree SHA, argv, platform, toolchain); only successes are reused, never - for another tree, and `force` skips reuse. `platform` is `-` - (`linux-x86_64`, `macos-aarch64`). -* `GET /v1/units/{id}` returns the unit (state `queued|leased|running|done|failed`, - `exit_code`, `receipt`). -* `GET /v1/units/{id}/log?from=N&wait=S` returns log bytes from offset `N`, - waiting up to `S` seconds (at most 30) for new bytes while the unit is active. - Headers: `x-borg-next` (the offset to ask for next), `x-borg-state`, and - `x-borg-exit` once finished and fully read. -* `POST /v1/units/{id}/cancel`. -* `PUT /v1/units/{id}/bundle` (the unit's own client; body is a git bundle, at most - `units.max_bundle_bytes`, 64 MiB by default). A unit submitted with - `"bundle": true` is not placed until the bundle arrives, and fails after - `units.bundle_wait_secs` (60) without it. Bundles are deleted with their unit - (7 days, the life of its receipt) or at once if the unit fails or is cancelled. - -Drone protocol (bearer: drone token): - -1. The heartbeat response carries `units: [{unit_id, repo, tree_sha, argv, class, - toolchain, timeout_secs, attempt}]` for units placed on this drone, and - `cancel_units: [id]`. Placed units are resent until acknowledged. -2. Acknowledge by listing the unit id in `active_units` in the heartbeat request; - that starts the unit. Keep listing it while it runs. A running unit the drone - stops listing for 60 seconds is treated as lost and re-queued (up to - `units.max_attempts`). -3. Run it: get `tree_sha` into the local git mirror (if `git cat-file -e - ` fails and the grant has `bundle_bytes > 0`, fetch - `GET /v1/units/{id}/bundle` and `git bundle unbundle` it into the mirror; its - prerequisite is a commit the mirror already has, and if that fails the unit - fails and the client runs locally), check the tree out into a slot directory - (`git clean` it first), run `argv` as the unprivileged job user with the same - isolation rules as CI jobs, and enforce `timeout_secs`. -4. Stream output with `POST /v1/units/{id}/log?offset=N` (raw bytes; `N` is the - byte offset of the first byte of the body). Repeating a chunk is harmless; - `409` means the queen holds fewer than `N` bytes, so resend from its length - (the `200` body of any append is the queen's length). -5. Finish with `POST /v1/units/{id}/result` `{exit_code, duration_secs}`. - Idempotent. Stop and report on a `cancel_units` entry. - -Units are placed by the runner allocator (class capacity, owner caps, bids, -profiles, quarantine and probation) on drones whose `-` matches, using -capacity left after CI runner demand. - -### Build-box reclaim - -A drone on a shared build host can reclaim abandoned git worktrees and cargo -`target/` and `incremental/` directories by disk pressure (`reclaim` block of the -drone config; off by default, dry-run by default). `reclaim::run_cycle` is the -policy; the drone implements `ReclaimHost` and reports `ReclaimStats` in its -heartbeat (`reclaim`), which the queen audits (`/status` `reclaim`). - -What a `ReclaimHost` implementation must guarantee, because the policy relies on -it: `probe` reports **busy** if any process has its working directory or an open -file inside the path (check `/proc/*/cwd` and `fd`, or `lsof +D` on macOS); -**dirty** if `git status --porcelain` is non-empty, untracked files included; -**unpushed** if `git rev-list HEAD --not --remotes` is non-empty (or HEAD is -detached on such commits); **protected** for the main checkout (first entry of -`git worktree list`), a locked worktree, or a repository with stashes; **not_git** -when the path is not a work tree; and **unknown** whenever any step failed or -timed out. `delete` removes the directory (for a worktree, `git worktree remove` -or removal followed by `git worktree prune`) and returns the bytes freed. `scan` -lists each child of a root, and for `cargo_target` roots also each -`incremental/` directory inside (with `within` set to its target directory). - -### Quorum rollout - -`rollout_target_version` plus `rollout.previous_version` turns on a gated rollout -of a new drone version: the queen tells one Linux and one macOS canary the -target through `desired_version` in the heartbeat response, promotes through -waves only while success rate, host-failure rate against the control group and -heartbeat health agree, and halts and tells every drone the previous version -otherwise. State is persisted; see `rollout` in `/status` and -`borg queen rollout halt|reset`. Drones update between jobs when -`desired_version` differs from what they run, and may downgrade. - -### Learned profiles - -The queen builds per-drone, per-class profiles from telemetry: speed against the -fleet median for the same job family (the job name without matrix suffixes) and -reliability from host-caused failures only. They add to the bid score in -proportion to the evidence, so proven faster and more reliable hosts are -preferred when placing a runner. See `profiles` in the queen config and -`profiles[]` in `/status`. - -### Warm-floor controller - -Optional (`homeostasis.enabled`). A damped feedback loop raises the warm floor of -a class and scope when queue latency stays above a target and idle room exists, -and lowers it gradually when demand subsides, with hysteresis, per-key cooldowns, -a step size, an hourly change limit and per-class ceilings. The clock schedule -(`warm_schedule`) remains a minimum. State appears in `/status` and the -`borg_warm_floor*` metrics. - -### Host health - -The queen scores each drone's host-caused failures (from telemetry with -`failure_class: host`) with a half-life decay. Code failures never count. It -recommends quarantine on moderate evidence, and quarantines automatically only -when failures from several distinct runs corroborate, the failures do not look -fleet-wide, and enough other hosts remain in service. A quarantined host returns -after a cooldown, on probation (one lease at a time), once a few jobs succeed. -Tune it with a `health` block in the queen config (`auto_quarantine`, -`half_life_secs`, `window_secs`, `corroboration_runs`, `cooldown_secs`, -`proof_jobs`, and the bounds); `/status` shows each drone's `health`. - -### Tailnet enrollment - -Optional and off by default. With a `tailnet_enroll` block -(`allowed_tags`, `allowed_users`, optional `socket_path` for the local -tailscaled), `POST /v1/enroll` identifies the caller through the local -tailscaled (`whois` by remote address) and, if the node carries an allowed tag -(or, for an untagged node, belongs to an allowed user), enrolls a drone named -after the node and returns `{"name", "token"}` once. The queen must see the -caller's tailnet address directly (bind it to the tailnet address, no proxy in -front), and its service user needs access to the tailscaled socket -(`tailscale set --operator=`). A node that is already enrolled gets 409. - -The optional supersede janitor cancels queued pull-request workflow runs whose -head commit is no longer the pull request's head. Add a `supersede_janitor` -block (`token_path` or `BORG_JANITOR_TOKEN` for a token with `actions: write`, -`dry_run` defaulting to true, `interval_secs`, `max_cancels_per_pass`) to turn -it on. It never cancels a run that has started. - -The GitHub App needs organization self-hosted runners: read and write, plus -read access to Actions on the listed repositories. Set -`BORG_APP_PRIVATE_KEY` instead of `private_key_path` to pass the PEM in the -environment. The queen speaks plain HTTP; bind it to a private interface or -put a TLS terminator in front. - -| Endpoint | Purpose | -| --- | --- | -| `GET /status` | Per-class demand, slots and leases; drones and heartbeat age; oldest queued age per label; `reserved_for`. | -| `GET /metrics` | The same numbers in Prometheus text format. | -| `POST /v1/heartbeat` | Drone API. Reports capacity; returns new leases until their JIT config is acknowledged. | -| `POST /v1/leases/{id}/result` | Drone API. Reports a lease's exit code, duration and peak RSS. | - -Manage drones with `borg queen drone list` and `borg queen drone remove ` -(which revokes the token). `borg queen enroll` prints only the token to stdout. +On Linux, the installer creates a daemon account and a separate job account, +installs the pinned GitHub runner, and writes a systemd unit. WSL needs +systemd enabled. On macOS, add `--run-as ` to select the +LaunchDaemon's account. See the [drone guide](docs/operations.md#drone) for +platform setup and isolation details. -Drone requests carry `Authorization: Bearer `. The queen stores only -the SHA-256 of each token. +Add `--dry-run` to preview installation. Use `--max-vcpu`, `--max-mem-mib`, +and `--max-leases` to limit what a host lends. -### Starvation guard +### 3. Route a workflow -When demand for a class has waited longer than `starvation_after_secs`, the -allocator places it first. If it still does not fit, it holds back the drone -that can free up soonest (fewest in-flight leases, ties by name) and grants no -smaller lease there until the starving lease fits. Held drones are listed in -`/status` under `reserved_for`. +With the default label prefix, use `borg-ci-` for default-branch and +merge-queue jobs, and `borg-ci--pr` for pull requests. Replace `main` +with your repository's default branch: -## Drone +```yaml +name: Fleet check +on: + push: + branches: [main] + pull_request: + merge_group: -``` -sudo borg drone install --queen https://queen.example.test:8443 --token --ca-cert queen-ca.pem -borg drone install ... --dry-run # list changes without applying them -borg drone run --config /etc/borg/drone.json +jobs: + check: + runs-on: ${{ github.event_name == 'pull_request' && 'borg-ci-light-pr' || 'borg-ci-light' }} + steps: + - run: uname -a ``` -`install` is idempotent. It prints one line per action (`created`, `updated`, -`unchanged`, `ran`, ...) and restarts the service only if something changed. -It sets up: - -* **Linux and WSL:** a `borg` system user (added to the `docker` group when - that group exists) and a systemd unit. The unit sends `SIGTERM` to the drone - only (`KillMode=mixed`), so in-flight jobs finish while it drains. WSL needs - systemd enabled: put `[boot]` and `systemd=true` in `/etc/wsl.conf`, run - `wsl --shutdown` from Windows, and start the distribution again. -* **macOS:** a LaunchDaemon (`org.borg.drone`) that runs as `--run-as `, - which must already exist. -* The pinned actions-runner (version and SHA-256 are in `src/drone/pins.rs`), - downloaded and checksum-verified before it is unpacked, a copy of this - binary, `drone.json` and the token file (mode 0600). - -The queen URL must be `https://`; `ca_cert_path` makes that CA the only trust -root. A plain `http://` URL is refused unless `allow_insecure` is set. - -### Separate job user (Linux) - -Leases run code from pull requests, so they must not run as the account that -holds the enrollment token. `borg drone install` creates two accounts: the -drone daemon (owns `/etc/borg`, mode 0700, with the config and token at 0600) -and an unprivileged job user (`job_user` in the config; default `borg-job`) -that runs every lease and owns the work and cache directories. The daemon -becomes the job user with `setpriv --reuid ... --init-groups`, which drops every -capability before the job starts; the systemd unit grants the daemon only -`CAP_SETUID`, `CAP_SETGID` and `CAP_SETPCAP` for that. File operations on -job-owned directories (cloning the runner, cleanup, cache collection) also run -as the job user. `tests/job_user.rs` checks, as root, that a lease's uid differs -from the drone's, that the token and config directory are unreadable to it, and -that it holds no capability. Without `job_user` (macOS, or `--no-job-user`) -jobs run as the drone's own account and the drone logs a warning. - -### Labels, environment, hooks - -* `extra_labels` are added to every runner the drone registers (for lanes that - still route by label). The queen drops any that equal a job class label, so - they can never change whether a runner serves main or pull-request jobs. -* `runner_env` is a map injected into every lease's environment (keys starting - `BORG_` or `ACTIONS_RUNNER_`, and a few others, are reserved). -* `job_started_hook` and `job_completed_hook` are absolute paths ending in `.sh`, - exported as `ACTIONS_RUNNER_HOOK_JOB_STARTED` and `ACTIONS_RUNNER_HOOK_JOB_COMPLETED`. -* `borg drone run --token-file F` and `borg drone install --token-file F` read - the token from a file instead of the command line. - -### Owner caps - -`--max-vcpu`, `--max-mem-mib` and `--max-leases` (config keys of the same -names) are sent with every heartbeat. The queen allocates within the capped -budget, so an owner can say "this PC gives at most 8 cores and 24 GiB". - -### What a lease does - -1. The pinned runner install is cloned into `/` (hard - links where possible) and `run.sh` starts in its own process group. The JIT - config is passed in `ACTIONS_RUNNER_INPUT_JITCONFIG`, never on the command - line, so it is not visible in `ps`; the runner masks it and removes it from its - environment. Jobs see `BORG_LEASE_ID`, `BORG_DOCKER_PREFIX`, - `COMPOSE_PROJECT_NAME` (both `borg-`) and `BORG_CACHE_ROOT`. -2. When the runner exits, or `max_job_secs` passes, the process group is - killed, and the work directory (including `_diag`) is deleted. -3. Docker containers, networks and volumes are removed only if `docker - inspect` shows they carry the label `borg.lease=`, the compose - project `borg-`, or a name that is exactly that prefix or starts - with it plus `-` or `_`. There is no global prune. Jobs that start their own - containers should name them with `$BORG_DOCKER_PREFIX` (docker compose does - this through `COMPOSE_PROJECT_NAME`). -4. Tool caches under `cache_root` are never touched. -5. The exit code, duration and peak RSS of the runner's process group (Linux) - are reported to the queen. - -### Host caches - -The drone owns one stable directory per cache kind under `cache_root`. The -path is identical for every lease on the host, is exported to every job, is -never touched by per-job cleanup, and is size-capped by least-recently-used -garbage collection that runs only while no lease is running. `cache_root` and -`work_root` must not contain one another. Everything is configured in the -`caches` section of the drone config; every kind except `git` is on by default. - -| Kind | Environment for jobs | Default cap | Collected | -| --- | --- | --- | --- | -| `sccache` | `SCCACHE_DIR`, `SCCACHE_CACHE_SIZE` (and `SCCACHE_SERVER_PORT` with the optional host server) | 25% of free disk at start, at most 100G | never by the drone; sccache enforces its own size | -| `cargo` | `CARGO_HOME` (registry and git), `CARGO_INSTALL_ROOT` per lease, `BORG_RUSTFLAGS` and `RUSTFLAGS` | 20G | crates, unpacked sources, git db and checkouts | -| `bazel` | `BAZELISK_HOME`; `--disk_cache` and `--repository_cache` via a managed block in `~/.bazelrc` | 50G | `disk/cas`, `disk/ac`, repository cache entries | -| `pnpm` | `PNPM_STORE_DIR`, `npm_config_store_dir` | 20G | store files (the index is kept) | -| `go` | `GOMODCACHE`, `GOCACHE` | 20G | build cache only; the module cache is read-only trees and is left alone | -| `uv` | `UV_CACHE_DIR` | 10G | entries of each versioned bucket | -| `mise` | `MISE_DATA_DIR` | 20G | installed tool versions and downloads; shims are kept | -| `git` | `BORG_GIT_MIRROR__` per mirror | none | never (objects only accumulate) | - -```json -"caches": { - "git": { "mirrors": ["example-org/app"], "remote_base": "https://github.com" }, - "sccache": { "max_size": "50G" }, - "bazel": { "max_size": "40G", "bazelrc": "/var/lib/borg/.bazelrc" }, - "uv": { "enabled": false } -} -``` +Class requirements and the label prefix are configurable. See +[labels](docs/operations.md#labels) and +[class configuration](docs/operations.md#classes-config). -Notes on each kind: - -* **Garbage collection** removes whole entries (a crate, a content-addressed - blob, an installed tool version), oldest use first, until the collectable data - is at most 90% of the cap. It takes a lock that every lease holds for its whole - life, so it never overlaps a job; under constant load it simply waits for an - idle moment, and caps are soft in the meantime. -* **cargo** shares only the registry and git data. `cargo install` is pointed at - a per-lease root so no installed binary ever leaks between leases. -* **bazel**: the drone writes a marked block into the service user's `~/.bazelrc` - (or `caches.bazel.bazelrc`) and leaves the rest of the file alone. The user rc - is read after the workspace rc, so these two flags win. Bazel's disk cache is - safe for concurrent builds: entries are written to a temporary file and renamed - into place, and a collision between processes is harmless. A collected entry is - a cache miss. -* **go**: the module cache is never collected; the build cache is. -* **uv** and **mise** collection is coarse: it removes entries those tools - re-create on demand. mise serializes installs of one tool version with a lock - file. -* **git mirrors**: for each `caches.git.mirrors` entry (`owner/name`) the drone - keeps a bare mirror of the branches and tags at - `/git//.git`, refreshed every - `caches.git.interval_secs` (default 600) and before a lease starts (at most - once per 15 seconds). The variable name is upper-cased with every other - character replaced by `_`, and it is set only while the mirror exists, so a job - can test it and fall back to a normal clone: - - ``` - if [ -n "${BORG_GIT_MIRROR_EXAMPLE_ORG_APP:-}" ]; then - git clone --reference "$BORG_GIT_MIRROR_EXAMPLE_ORG_APP" src - else - git clone src - fi - ``` - - Mirrors fetch from `remote_base` (default `https://github.com`) plus - `/.git`. The drone holds no GitHub credentials: for a private - repository give the service user a credential helper, or set `remote_base` to - an ssh base such as `git@host:` with a read-only deploy key. A failed refresh - keeps the previous mirror and never blocks a job. - -#### Slots, path remapping and compile-cache keys - -A lease runs in `/slot-` and takes the lowest free slot -(`slots` in the drone config, default `max_leases`, else 16; it is also the most -leases that run at once). The path never contains the lease id, concurrent -leases always get distinct slots, and a slot is fully cleaned before it is -reused, so the next lease in a slot sees exactly the same absolute paths. The -lease-to-slot mapping is logged (`lease X -> slot N`) and kept in drone state. -Jobs see `BORG_SLOT`. - -Jobs also get `BORG_RUSTFLAGS` and `RUSTFLAGS`: `--remap-path-prefix` flags that -map **every** slot directory onto `/borg/work`, and `CARGO_HOME`, its registry -and its git data onto fixed virtual prefixes. The string lists all slots, so it -is byte-identical for every lease (a per-slot flag would itself differ between -slots and change the cache key). The same flags are written to -`$CARGO_HOME/config.toml` as `[build] rustflags`, as a fallback. Two cautions: - -* Cargo ignores `build.rustflags` and `target.*.rustflags` from config files - whenever `RUSTFLAGS` is set. A workflow that sets its own `RUSTFLAGS` replaces - ours; it should append the remap instead, e.g. - `RUSTFLAGS: "-D warnings ${BORG_RUSTFLAGS}"`. Repositories that depend on - config-file rustflags can turn the export off with - `caches.cargo.export_rustflags: false`. -* The remap does **not** make workspace crates hit the compile cache across - different slots. sccache hashes the compile's working directory and every - `CARGO_*` variable, including `CARGO_MANIFEST_DIR` (read from sccache's - `generate_hash_key`, v0.16), and no rustc flag changes those. What slots buy is - that a *reused* slot recompiles at identical paths and hits, and that - registry dependencies, whose paths are in the shared `CARGO_HOME`, hit from any - slot. `tests/sccache_slots.rs` compiles one crate in slot 0, again in a cleaned - slot 0, and in slot 1, and prints the hit counts for each. - -sccache keeps hit and miss counters only in the server's memory, so after a -lease's own server exits there is nothing on disk to read. Set -`caches.sccache.server: {"port": 4226}` to have the drone run one long-lived -sccache server per host (its configuration then applies to all jobs, and a job -whose sccache version differs from the host binary's restarts it); the drone -reads its totals with `sccache --show-stats --stats-format json` and reports -hits and misses for `sccache` in heartbeats, `/status` and `borg_cache_hits`. -Without it those fields stay empty. - -Every heartbeat carries each cache's size, cap, evicted bytes and last -collection time (and, for git, how many leases were given a mirror and how many -ran without one). The queen shows them under each drone in `/status` and as -`borg_cache_*` metrics. Sizes are re-measured every -`caches.maintenance_interval_secs` (default 600) and after leases. - -`borg drone install` takes `--git-mirror owner/name` (repeatable) and -`--git-remote-base`; edit `/etc/borg/drone.json` for the rest. - -On start the drone deletes leftover lease directories and borg-labelled docker -resources from a previous run. On `SIGTERM` it stops accepting leases and waits -up to `drain_timeout_secs` for running jobs before stopping them. - -## Updating a drone - -A drone follows the queen's `desired_version` (see Quorum rollout) between -jobs, in either direction. When the version differs it stops accepting leases, -finishes the jobs it holds, and then asks the updater: it writes `update-request` -into its state directory (`state_dir` in `drone.json`). The drone itself cannot -replace its binary, because it runs unprivileged and stopping it is what drains -it. A root-owned path unit (`borg-update.path` on Linux, the `org.borg.update` -launchd daemon on macOS) starts `borg update --from-request`, which: - -1. fetches the release manifest and downloads every artifact it needs, checking - each SHA-256 against the manifest **before** anything is stopped; -2. stops the drone service (which drains running jobs), swaps the binary - atomically and keeps the previous one as `borg.prev`; -3. starts the service and waits for a fresh `health.json`, which the new drone - writes after each heartbeat the queen answered; -4. rolls back to `borg.prev` if the service does not start or does not report - healthy within `health_timeout_secs`. - -The outcome is left in `update-result.json` and sent to the queen in the next -heartbeat as `update_report`. A target that failed or rolled back on a host is -not retried there until the queen asks for a different version. A queen that is -down is retried with growing delays, and job results are retried for several -minutes, so a restart of the queen loses nothing. - -The updater reads `/etc/borg-update/update.json`, which only root can write, so -a compromised drone cannot point it at another release. The manifest is a JSON -document (`schema`, `version`, per-platform `url` and `sha256` for -`linux-x64-musl`, `linux-arm64-musl` and `darwin-arm64`, and an optional pinned -actions-runner). `https://` is required unless the config sets `allow_insecure`. -There are no signatures: the trust anchor is HTTPS to the configured manifest -URL plus the pinned digests. - -```text -borg update --check # what would change -borg update --version 0.5.0 # install exactly this version -borg update --runner-only # replace just the pinned actions-runner, no restart +## Operate the fleet + +```sh +borg queen drone list --config queen.toml +borg queen quarantine --config queen.toml --reason "maintenance" host-a +borg queen release --config queen.toml host-a +borg report --config queen.toml ``` -`borg drone install` writes the updater configuration and units; -`--update-manifest-url` points them at another manifest and `--no-self-update` -skips them. Drones installed before this feature have to be updated once by -rerunning `borg drone install` with the new binary. +The queen exposes fleet state at `/status`, Prometheus metrics at `/metrics`, +and job telemetry at `/report`. + +| Guide | Topics | +| --- | --- | +| [Configuration](docs/operations.md#config-files) | TOML and JSON, validation, defaults, and redacted secrets. | +| [Scheduling](docs/operations.md#drone-bids-and-demand-pressure) | Bids, capacity, starvation protection, host health, and warm floors. | +| [Work units](docs/operations.md#borg-run) | Run commands on the fleet with `borg run`, including local fallback. | +| [Host caches](docs/operations.md#host-caches) | Cache limits, collection, git mirrors, slots, and compile-cache behavior. | +| [Drone updates](docs/operations.md#updating-a-drone) | Canary rollout, update manifests, health checks, and rollback. | +| [Operations reference](docs/operations.md) | APIs, telemetry fields, enrollment, and optional controllers. | + +## Development + +From the public repository root: -## Contributing +```sh +cargo build --locked +cargo test --locked +cargo clippy --locked --all-targets -- -D warnings +``` + +| Source | Responsibility | +| --- | --- | +| `src/probe/` | Host capacity and tool detection. | +| `src/classes.rs`, `src/allocator.rs` | Runner requirements and lease placement. | +| `src/queen.rs`, `src/github.rs`, `src/store.rs`, `src/api.rs` | Coordination, GitHub integration, persistence, and APIs. | +| `src/drone/` | Heartbeats, runner execution, installation, and cleanup. | -Issues are welcome here. Code changes land in the source-of-truth repository -and are projected into this one; pull requests opened directly against this -repository will be overwritten by the next projection. +Issues are welcome in this repository. Source changes land upstream in +[dx-corp/mono](https://github.com/dx-corp/mono/tree/main/rust/tools/borg) +and are projected here. Maintainers make code and documentation changes +upstream so the next projection preserves them. ## License -Business Source License 1.1 (BUSL-1.1). See [LICENSE](LICENSE). +[Business Source License 1.1](LICENSE) (BUSL-1.1). diff --git a/docs/operations.md b/docs/operations.md new file mode 100644 index 0000000..04693c2 --- /dev/null +++ b/docs/operations.md @@ -0,0 +1,707 @@ +# Operations guide + +Configuration, scheduling, drone administration, and protocol details for +[Borg](../README.md). Start with the README for installation and fleet setup. + +## Contents + +- [Labels](#labels) +- [Classes config](#classes-config) +- [Config files](#config-files) +- [Queen](#queen) + - [Job telemetry](#job-telemetry) + - [Drone bids and demand pressure](#drone-bids-and-demand-pressure) + - [Wake-on-LAN](#wake-on-lan) + - [Build environment contract](#build-environment-contract) + - [borg run](#borg-run) + - [Work units (gate and hook runs)](#work-units-gate-and-hook-runs) + - [Build-box reclaim](#build-box-reclaim) + - [Quorum rollout](#quorum-rollout) + - [Learned profiles](#learned-profiles) + - [Warm-floor controller](#warm-floor-controller) + - [Host health](#host-health) + - [Tailnet enrollment](#tailnet-enrollment) + - [Starvation guard](#starvation-guard) +- [Drone](#drone) + - [Separate job user (Linux)](#separate-job-user-linux) + - [Labels, environment, hooks](#labels-environment-hooks) + - [Owner caps](#owner-caps) + - [What a lease does](#what-a-lease-does) + - [Host caches](#host-caches) +- [Updating a drone](#updating-a-drone) + +## Labels + +A job asks for a class with a runner label: `` for jobs from +the default branch or a merge queue, and `-pr` for pull +requests. The default prefix is `borg-ci-`; set `label_prefix` in the classes +config to change it (for example `example-ci-`). A runner minted for main +demand carries only the bare label. A runner minted for pull-request demand +carries both labels. With this workflow labeling convention, default-branch +and merge-queue jobs request the bare label; pull-request jobs request the +`-pr` label. Labels route jobs; they do not enforce an event restriction. + +## Classes config + +The `classes` table is partial; omitted fields keep their defaults: + +```toml +[classes] +label_prefix = "example-ci-" +heavy = { mem_mib = 20480, vcpu = 6, docker = true } +``` + +## Config files + +Queen and drone configs are TOML or JSON with the same keys, defaults and +validation. A `.toml` or `.json` extension picks the format; any other name is +read as JSON when it starts with `{` and as TOML otherwise. Unknown keys are +errors, reported with the file, line and key. TOML has no null, so settings +that JSON turns off with `null` take `false` in either format +(`starvation_after_secs = false`). + +```sh +borg config check queen.toml # validate and print the effective config +borg config check --drone drone.toml +``` + +`check` prints every setting with its defaults filled in; secrets print as +`[redacted]`. + +## Queen + +```sh +borg queen enroll --config queen.toml host-a # prints a one-time token +borg queen run --config queen.toml +``` + +Example `queen.toml` (every key except `github`, `repos` and `db_path` is +optional): + +```toml +repos = ["example-org/app"] +db_path = "/var/lib/borg/queen.db" +listen = "127.0.0.1:8080" +starvation_after_secs = 120 + +[github] +org = "example-org" +app_id = "12345" +private_key_path = "/etc/borg/app.pem" + +[classes] +label_prefix = "example-ci-" +``` + +Operators can stop new leases to one drone with `borg queen quarantine +--config queen.toml --reason "" ` and lift it with `borg queen +release --config queen.toml `. Set `rollout_target_version` in the queen +config to have `/metrics` and `/status` report how many drones run that borg +version. The queen sends that target as `desired_version`; drones update +between jobs. Configure `rollout.previous_version` to gate the rollout +through canaries and waves (see [Quorum rollout](#quorum-rollout)). + +`warm_floors` keep idle runners alive per class and scope; the optional +`warm_schedule` raises them in time-of-day windows (`{"class": "light", "scope": +"pr", "count": 4, "start": "08:00", "end": "20:00", "days": ["mon", "tue"]}`) +at a fixed `utc_offset_minutes`. A window never lowers a floor. + +### Job telemetry + +A drone may add a `telemetry` object to `POST /v1/leases/{id}/result`: + +```json +{ + "exit_code": 1, "duration_secs": 412, + "telemetry": { + "job_name": "rust-test (shard 2)", "repo": "example-org/app", "run_id": 123456789, + "conclusion": "failure", "failure_class": "host", "queue_wait_secs": 18, + "steps": [{ "name": "cargo test", "duration_secs": 401, "conclusion": "failure" }] + } +} +``` + +`conclusion` is `success`, `failure`, `cancelled`, `skipped`, `timed_out` or +`unknown`; `failure_class` is `host` (machine or runner failure), `code` (the +job's own failure) or `unknown`. Unrecognized values read as `unknown`; at most +64 steps and 120 characters per name are kept; `job_name` and a `owner/name` +`repo` are required, otherwise the telemetry is dropped and the result still +recorded. The queen keeps it for `telemetry_retention_days` (default 30). +`borg report --queen ` (or `--config queen.json` on the queen host) and +`GET /report` summarize time by job and step, queue wait per class, the +cancellation share and host-caused failures per drone. + +#### What a drone-side parser should produce + +The queen validates but does not interpret these fields, so the drone owns +their meaning. Keep it consistent across drones: + +- **Identity.** `repo`, `run_id` and `job_name` come from the job the runner + executed (its job context: repository, run id, job display name), not from + log text a workflow can print. `job_name` is the name as GitHub shows it, + including matrix suffixes. +- **Steps.** One entry per step in execution order, `duration_secs` rounded + down to whole seconds, `conclusion` mapped from the step result (succeeded to + `success`, failed to `failure`, canceled to `cancelled`, skipped to + `skipped`, anything else to `unknown`). Report names only; never log lines, + environment values or step output. +- **Conclusion.** The job's result, mapped the same way; a job the drone killed + at `max_job_secs` is `timed_out`. +- **`queue_wait_secs`.** Seconds from the job being queued on GitHub to the + runner starting it. Omit it when it cannot be determined; do not send 0 as a + placeholder, because the report averages the values it receives. +- **`failure_class`.** `host` when the machine or the runner failed rather than + the job: the runner process was killed by a signal or the kernel (OOM), the + disk filled, Docker was unreachable, the runner lost its connection, or the + lease was torn down by the drone. `code` when the job's own step exited + non-zero on a healthy host. `unknown` when the drone cannot tell. Leave it + `unknown` for jobs that did not fail. +- **Failure behaviour.** Parsing is best effort and must never delay or fail + the result: if anything cannot be read, send the result without `telemetry` + (or with fewer fields), and send it in the same request as `exit_code` and + `duration_secs`. A retried result may be sent again; the queen keeps the + latest. + +### Drone bids and demand pressure + +A drone may add `bids` to its heartbeat: at most one per `(class, scope)` it is +willing to serve. The queen answers every heartbeat with `pressure`, the +demand pressure (0 to 1000) per class, so the drone can set its bids against it. + +```json +{ "class": "medium", "scope": "pr", "available_slots": 2, + "estimated_start_ms": 1500, "cache_affinity": 800, "resource_pressure": 150, + "estimated_cost": 1000, "confidence": 900, "min_pressure": 0, "expires_at": 1760000000 } +``` + +`cache_affinity`, `resource_pressure`, `confidence` and `min_pressure` are +permille (0 to 1000; larger values are clamped); `estimated_cost` is relative +(1000 is fleet typical); `expires_at` is Unix seconds, after which the queen +ignores the bid. `min_pressure` is the drone's response threshold: it does not +want this work while the class's pressure is below it, so a warm idle host +reports 0, a cold host something higher, and a workstation whose owner is +active something high. Starving demand overrides every threshold. + +The allocator drops a bid with no `available_slots` or above the current +pressure, and scores the rest with integers: +`3*cache_affinity - 3*resource_pressure - 2*start - 1*(cost-1000)`, where +`start` is `estimated_start_ms` as permille of 60 s, all scaled by +`confidence/1000`. The highest score wins; ties go to the tightest memory fit, +then the drone name. A drone that sends no bids, or only expired ones, is +treated as before (always eligible, neutral score); a drone that bids but not +for a key declines that key. GitHub still chooses which queued job a runner +takes: bids steer only which host receives a freshly minted runner. + +### Wake-on-LAN + +Optional (`wake` block, off by default). When demand for a class stays unmet for +`after_secs`, the queen sends a magic packet to a listed host that can serve the +class and whose drone is not fresh, rate limited per pass, per hour and per host. +Targets are IPv4 addresses or `addr/len` networks. See `/status` `wake`. + +### Build environment contract + +Optional (`env_contract.enabled`, off by default). Before a CI runner lease or a +work unit is placed on a drone, the queen checks what the drone last reported +about its build environment. Four checks, each `required`, `advisory` or `off` +globally and per class (`env_contract.checks`, `env_contract.classes.`): + +* `compiler_cache`: `RUSTC_WRAPPER` is an allowed wrapper (`allowed_wrappers`, + by file name: `sccache` by default, plus any single-flight wrapper in front of + it), `SCCACHE_DIR` is set, the sccache server is reachable. +* `toolchain`: when the work declares a Rust toolchain, the drone has it. A unit + declares it (`rust_toolchain`, which `borg run` reads from `rust-toolchain.toml` + or `rust/rust-toolchain.toml` at the unit's tree); a class can declare it in + config (`classes..rust_toolchain`). A mismatch is not admitted. +* `cache_store`: stats work, errors per thousand requests are under + `max_cache_error_pm`, the cache directory has at least + `min_cache_dir_free_bytes` free, and a write works. A hit rate at or under + `hit_rate_warn_pm` after `min_requests_for_hit_rate` requests is a *warning* in + `/status`, never a block. +* `resources`: memory available covers the class's memory (`mem_available_pct`) + and CPU pressure is at most `max_cpu_psi_pm`. (Free capacity after active + leases is the allocator's own rule.) + +A required check that fails **or is unknown** makes the drone ineligible for the +class; nothing else changes about placement. Facts older than +`facts_max_age_secs` (180) count as unknown. A drone that has never reported is +skipped unless `require_report`, so enabling the contract does not cut off +drones that predate it. Defaults when enabled: `toolchain` and `resources` +required, `compiler_cache` and `cache_store` advisory. + +`/status` shows each drone's `contract` (the checks with state, reason, observed +values and warnings, and the admission per class with the reason) and +`units.not_admitted` (why a queued unit is not admitted on each matching drone). +Metrics: `borg_env_contract_check{drone,check,state}` and +`borg_env_contract_admitted{drone,class}`. + +**Drone probe spec** (the heartbeat field `env_facts`). The drone only measures; +verdicts and thresholds are the queen's. Every field is optional, and `null` / +absent means "not measured" (never "fine"). Probe as the **job user**, with the +environment a job will get, and cache results about 60 seconds (resend the cached +object with its original `observed_at`; the queen treats facts over 180 seconds +old as unknown, so probe at least that often): + +| Field | How | +| --- | --- | +| `observed_at` | Unix seconds the probes ran. | +| `rustc_wrapper` | `RUSTC_WRAPPER` as a job sees it, resolved to a path through `PATH` if it is a bare name; absent if unset or empty. | +| `sccache_dir` | `SCCACHE_DIR` as a job sees it; absent if unset. | +| `sccache_server_reachable` | Whether the sccache server answers (connect to `SCCACHE_SERVER_PORT` or its socket, 5 s timeout). Never start a server to find out. | +| `sccache_stats_ok` | `sccache --show-stats --stats-format json` exits 0 within 5 s. | +| `cache_requests`, `cache_hits`, `cache_errors` | From that JSON: `compile_requests`, the sum of the `cache_hits.counts` values and the sum of the `cache_errors.counts` values (check against the installed sccache's schema). | +| `cache_dir_free_bytes` | Available bytes (`f_bavail * f_frsize`) on the filesystem holding `SCCACHE_DIR`. | +| `cache_write_ok` | Create, write and remove a small file named `.borg-probe-` in `SCCACHE_DIR`. | +| `toolchains` | For each toolchain `rustup toolchain list` shows, `{channel, rustc_version}`: the channel as `rust-toolchain.toml` would name it (`1.95.0`, `stable`, `nightly-2026-01-01`; strip the host triple) and the `release:` line of `rustc + -vV`. `null` if rustup could not be queried. | +| `mem_available_mib` | Linux `MemAvailable` from `/proc/meminfo`; macOS (free + inactive + speculative pages) times the page size, in MiB. | +| `cpu_psi_some_avg10_pm` | Linux `/proc/pressure/cpu` line `some`, `avg10` times ten, rounded; absent where PSI does not exist. | + +A unit's toolchain (`rust_toolchain`) matches an installed `channel` or +`rustc_version` exactly. + +### `borg run` + +```sh +borg run [--index | --tree ] [--publish] [--class medium] [--wait-secs 60] -- +``` + +Runs a command on the fleet when a drone has capacity, and otherwise right here; +the exit code is the command's either way. With no `BORG_QUEEN` set it just runs +the command locally, so wrappers can call it unconditionally. + +* Environment: `BORG_QUEEN` (queen URL), `BORG_CLIENT_TOKEN` or + `BORG_CLIENT_TOKEN_FILE` (default `~/.config/borg/client-token`; the token comes + from `borg queen enroll-client`), `BORG_RUN_WAIT_SECS`, `BORG_PUBLISH`, + `BORG_TOOLCHAIN`. +* The tree is `HEAD^{tree}` by default, `git write-tree` of the index with + `--index` (pre-commit hooks), or `--tree `. The unit's identity is the repo, + tree, argv, platform and toolchain, so an identical command that already passed + for this tree is not run again (the receipt is reported, exit 0) and one that is + running is joined. +* The drone has to be able to get the tree. A tree from a pushed commit is in + its git mirror. Otherwise `--publish` uploads a `git bundle` to the queen + (`PUT /v1/units/{id}/bundle`): a synthetic commit of the tree whose parent is the + merge base with `origin/main`, bundled without everything that merge base + already reaches, so it is small. Nothing is pushed to any remote. Without + `--publish`, an unpublished tree runs locally. +* Falls back to running here when the queen is unset or unreachable, refuses the + unit, no drone takes it within `--wait-secs` (the unit is cancelled), the queen + is lost for 60 seconds mid-run, or the unit ends for infrastructure reasons + (cancelled, timed out, out of attempts). A command that ran on a drone and + exited nonzero is a real failure and its code is returned. Ctrl-C cancels a + unit this invocation created. + +### Work units (gate and hook runs) + +Optional (`units.enabled`, off by default). Clients such as `borg run`, the +pre-commit hook and `mq` submit whole commands; drones run them. Enroll a client +with `borg queen enroll-client ` (prints a token once). + +Client API (bearer: client token): + +* `POST /v1/units` with `{repo, tree_sha, argv, platform, toolchain, class, timeout_secs?, force?}` + returns `201` `{outcome: "new", unit_id, ...}`, or `200` with `outcome: + "joined"` (an identical unit is queued or running) or `"receipt"` (an + identical unit succeeded recently; `receipt` has the proof). The key is + (repo, tree SHA, argv, platform, toolchain); only successes are reused, never + for another tree, and `force` skips reuse. `platform` is `-` + (`linux-x86_64`, `macos-aarch64`). +* `GET /v1/units/{id}` returns the unit (state `queued|leased|running|done|failed`, + `exit_code`, `receipt`). +* `GET /v1/units/{id}/log?from=N&wait=S` returns log bytes from offset `N`, + waiting up to `S` seconds (at most 30) for new bytes while the unit is active. + Headers: `x-borg-next` (the offset to ask for next), `x-borg-state`, and + `x-borg-exit` once finished and fully read. +* `POST /v1/units/{id}/cancel`. +* `PUT /v1/units/{id}/bundle` (the unit's own client; body is a git bundle, at most + `units.max_bundle_bytes`, 64 MiB by default). A unit submitted with + `"bundle": true` is not placed until the bundle arrives, and fails after + `units.bundle_wait_secs` (60) without it. Bundles are deleted with their unit + (7 days, the life of its receipt) or at once if the unit fails or is cancelled. + +Drone protocol (bearer: drone token): + +1. The heartbeat response carries `units: [{unit_id, repo, tree_sha, argv, class, + toolchain, timeout_secs, attempt}]` for units placed on this drone, and + `cancel_units: [id]`. Placed units are resent until acknowledged. +2. Acknowledge by listing the unit id in `active_units` in the heartbeat request; + that starts the unit. Keep listing it while it runs. A running unit the drone + stops listing for 60 seconds is treated as lost and re-queued (up to + `units.max_attempts`). +3. Run it: get `tree_sha` into the local git mirror (if `git cat-file -e + ` fails and the grant has `bundle_bytes > 0`, fetch + `GET /v1/units/{id}/bundle` and `git bundle unbundle` it into the mirror; its + prerequisite is a commit the mirror already has, and if that fails the unit + fails and the client runs locally), check the tree out into a slot directory + (`git clean` it first), run `argv` as the unprivileged job user with the same + isolation rules as CI jobs, and enforce `timeout_secs`. +4. Stream output with `POST /v1/units/{id}/log?offset=N` (raw bytes; `N` is the + byte offset of the first byte of the body). Repeating a chunk is harmless; + `409` means the queen holds fewer than `N` bytes, so resend from its length + (the `200` body of any append is the queen's length). +5. Finish with `POST /v1/units/{id}/result` `{exit_code, duration_secs}`. + Idempotent. Stop and report on a `cancel_units` entry. + +Units are placed by the runner allocator (class capacity, owner caps, bids, +profiles, quarantine and probation) on drones whose `-` matches, using +capacity left after CI runner demand. + +### Build-box reclaim + +A drone on a shared build host can reclaim abandoned git worktrees and cargo +`target/` and `incremental/` directories by disk pressure (`reclaim` block of the +drone config; off by default, dry-run by default). `reclaim::run_cycle` is the +policy; the drone implements `ReclaimHost` and reports `ReclaimStats` in its +heartbeat (`reclaim`), which the queen audits (`/status` `reclaim`). + +What a `ReclaimHost` implementation must guarantee, because the policy relies on +it: `probe` reports **busy** if any process has its working directory or an open +file inside the path (check `/proc/*/cwd` and `fd`, or `lsof +D` on macOS); +**dirty** if `git status --porcelain` is non-empty, untracked files included; +**unpushed** if `git rev-list HEAD --not --remotes` is non-empty (or HEAD is +detached on such commits); **protected** for the main checkout (first entry of +`git worktree list`), a locked worktree, or a repository with stashes; **not_git** +when the path is not a work tree; and **unknown** whenever any step failed or +timed out. `delete` removes the directory (for a worktree, `git worktree remove` +or removal followed by `git worktree prune`) and returns the bytes freed. `scan` +lists each child of a root, and for `cargo_target` roots also each +`incremental/` directory inside (with `within` set to its target directory). + +### Quorum rollout + +`rollout_target_version` plus `rollout.previous_version` turns on a gated rollout +of a new drone version: the queen tells one Linux and one macOS canary the +target through `desired_version` in the heartbeat response, promotes through +waves only while success rate, host-failure rate against the control group and +heartbeat health agree, and halts and tells every drone the previous version +otherwise. State is persisted; see `rollout` in `/status` and +`borg queen rollout halt|reset`. Drones update between jobs when +`desired_version` differs from what they run, and may downgrade. + +### Learned profiles + +The queen builds per-drone, per-class profiles from telemetry: speed against the +fleet median for the same job family (the job name without matrix suffixes) and +reliability from host-caused failures only. They add to the bid score in +proportion to the evidence, so proven faster and more reliable hosts are +preferred when placing a runner. See `profiles` in the queen config and +`profiles[]` in `/status`. + +### Warm-floor controller + +Optional (`homeostasis.enabled`). A damped feedback loop raises the warm floor of +a class and scope when queue latency stays above a target and idle room exists, +and lowers it gradually when demand subsides, with hysteresis, per-key cooldowns, +a step size, an hourly change limit and per-class ceilings. The clock schedule +(`warm_schedule`) remains a minimum. State appears in `/status` and the +`borg_warm_floor*` metrics. + +### Host health + +The queen scores each drone's host-caused failures (from telemetry with +`failure_class: host`) with a half-life decay. Code failures never count. It +recommends quarantine on moderate evidence, and quarantines automatically only +when failures from several distinct runs corroborate, the failures do not look +fleet-wide, and enough other hosts remain in service. A quarantined host returns +after a cooldown, on probation (one lease at a time), once a few jobs succeed. +Tune it with a `health` block in the queen config (`auto_quarantine`, +`half_life_secs`, `window_secs`, `corroboration_runs`, `cooldown_secs`, +`proof_jobs`, and the bounds); `/status` shows each drone's `health`. + +### Tailnet enrollment + +Optional and off by default. With a `tailnet_enroll` block +(`allowed_tags`, `allowed_users`, optional `socket_path` for the local +tailscaled), `POST /v1/enroll` identifies the caller through the local +tailscaled (`whois` by remote address) and, if the node carries an allowed tag +(or, for an untagged node, belongs to an allowed user), enrolls a drone named +after the node and returns `{"name", "token"}` once. The queen must see the +caller's tailnet address directly (bind it to the tailnet address, no proxy in +front), and its service user needs access to the tailscaled socket +(`tailscale set --operator=`). A node that is already enrolled gets 409. + +The optional supersede janitor cancels queued pull-request workflow runs whose +head commit is no longer the pull request's head. Add a `supersede_janitor` +block (`token_path` or `BORG_JANITOR_TOKEN` for a token with `actions: write`, +`dry_run` defaulting to true, `interval_secs`, `max_cancels_per_pass`) to turn +it on. It never cancels a run that has started. + +The GitHub App needs organization self-hosted runners: read and write, plus +read access to Actions on the listed repositories. Set +`BORG_APP_PRIVATE_KEY` instead of `private_key_path` to pass the PEM in the +environment. The queen speaks plain HTTP; bind it to a private interface or +put a TLS terminator in front. + +| Endpoint | Purpose | +| --- | --- | +| `GET /status` | Per-class demand, slots and leases; drones and heartbeat age; oldest queued age per label; `reserved_for`. | +| `GET /metrics` | The same numbers in Prometheus text format. | +| `POST /v1/heartbeat` | Drone API. Reports capacity; returns new leases until their JIT config is acknowledged. | +| `POST /v1/leases/{id}/result` | Drone API. Reports a lease's exit code, duration and peak RSS. | + +Manage drones with `borg queen drone list` and `borg queen drone remove ` +(which revokes the token). `borg queen enroll` prints only the token to stdout. + +Drone requests carry `Authorization: Bearer `. The queen stores only +the SHA-256 of each token. + +### Starvation guard + +When demand for a class has waited longer than `starvation_after_secs`, the +allocator places it first. If it still does not fit, it holds back the drone +that can free up soonest (fewest in-flight leases, ties by name) and grants no +smaller lease there until the starving lease fits. Held drones are listed in +`/status` under `reserved_for`. + +## Drone + +```sh +sudo borg drone install --queen https://queen.example.test:8443 --token --ca-cert queen-ca.pem +borg drone install ... --dry-run # list changes without applying them +borg drone run --config /etc/borg/drone.json +``` + +`install` writes `drone.json`; a hand-written `drone.toml` with the same keys +works too. + +`install` is idempotent. It prints one line per action (`created`, `updated`, +`unchanged`, `ran`, ...) and restarts the service only if something changed. +It sets up: + +* **Linux and WSL:** a `borg` system user (added to the `docker` group when + that group exists) and a systemd unit. The unit sends `SIGTERM` to the drone + only (`KillMode=mixed`), so in-flight jobs finish while it drains. WSL needs + systemd enabled: put `[boot]` and `systemd=true` in `/etc/wsl.conf`, run + `wsl --shutdown` from Windows, and start the distribution again. +* **macOS:** a LaunchDaemon (`org.borg.drone`) that runs as `--run-as `, + which must already exist. +* The pinned actions-runner (version and SHA-256 are in `src/drone/pins.rs`), + downloaded and checksum-verified before it is unpacked, a copy of this + binary, `drone.json` and the token file (mode 0600). + +The queen URL must be `https://`; `ca_cert_path` makes that CA the only trust +root. A plain `http://` URL is refused unless `allow_insecure` is set. + +### Separate job user (Linux) + +Leases run code from pull requests, so they must not run as the account that +holds the enrollment token. `borg drone install` creates two accounts: the +drone daemon (owns `/etc/borg`, mode 0700, with the config and token at 0600) +and an unprivileged job user (`job_user` in the config; default `borg-job`) +that runs every lease and owns the work and cache directories. The daemon +becomes the job user with `setpriv --reuid ... --init-groups`, which drops every +capability before the job starts; the systemd unit grants the daemon only +`CAP_SETUID`, `CAP_SETGID` and `CAP_SETPCAP` for that. File operations on +job-owned directories (cloning the runner, cleanup, cache collection) also run +as the job user. `tests/job_user.rs` checks, as root, that a lease's uid differs +from the drone's, that the token and config directory are unreadable to it, and +that it holds no capability. Without `job_user` (macOS, or `--no-job-user`) +jobs run as the drone's own account and the drone logs a warning. + +### Labels, environment, hooks + +* `extra_labels` are added to every runner the drone registers (for lanes that + still route by label). The queen drops any that equal a job class label, so + they can never change whether a runner serves main or pull-request jobs. +* `runner_env` is a map injected into every lease's environment (keys starting + `BORG_` or `ACTIONS_RUNNER_`, and a few others, are reserved). +* `job_started_hook` and `job_completed_hook` are absolute paths ending in `.sh`, + exported as `ACTIONS_RUNNER_HOOK_JOB_STARTED` and `ACTIONS_RUNNER_HOOK_JOB_COMPLETED`. +* `borg drone run --token-file F` and `borg drone install --token-file F` read + the token from a file instead of the command line. + +### Owner caps + +`--max-vcpu`, `--max-mem-mib` and `--max-leases` (config keys of the same +names) are sent with every heartbeat. The queen allocates within the capped +budget, so an owner can say "this PC gives at most 8 vCPUs and 24 GiB". + +### What a lease does + +1. The pinned runner install is cloned into `/slot-` (hard + links where possible) and `run.sh` starts in its own process group. The JIT + config is passed in `ACTIONS_RUNNER_INPUT_JITCONFIG`, never on the command + line, so it is not visible in `ps`; the runner masks it and removes it from its + environment. Jobs see `BORG_LEASE_ID`, `BORG_DOCKER_PREFIX`, + `COMPOSE_PROJECT_NAME` (both `borg-`) and `BORG_CACHE_ROOT`. +2. When the runner exits, or `max_job_secs` passes, the process group is + killed, and the work directory (including `_diag`) is deleted. +3. Docker containers, networks and volumes are removed only if `docker + inspect` shows they carry the label `borg.lease=`, the compose + project `borg-`, or a name that is exactly that prefix or starts + with it plus `-` or `_`. There is no global prune. Jobs that start their own + containers should name them with `$BORG_DOCKER_PREFIX` (docker compose does + this through `COMPOSE_PROJECT_NAME`). +4. Per-job cleanup leaves tool caches under `cache_root` intact. +5. The exit code, duration and peak RSS of the runner's process group (Linux) + are reported to the queen. + +### Host caches + +The drone owns one stable directory per cache kind under `cache_root`. The +path is identical for every lease on the host, is exported to every job, is +never touched by per-job cleanup, and is size-capped by least-recently-used +garbage collection that runs only while no lease is running. `cache_root` and +`work_root` must not contain one another. Everything is configured in the +`caches` section of the drone config; every kind except `git` is on by default. + +| Kind | Environment for jobs | Default cap | Collected | +| --- | --- | --- | --- | +| `sccache` | `SCCACHE_DIR`, `SCCACHE_CACHE_SIZE` (and `SCCACHE_SERVER_PORT` with the optional host server) | 25% of free disk at start, at most 100G | never by the drone; sccache enforces its own size | +| `cargo` | `CARGO_HOME` (registry and git), `CARGO_INSTALL_ROOT` per lease, `BORG_RUSTFLAGS` and `RUSTFLAGS` | 20G | crates, unpacked sources, git db and checkouts | +| `bazel` | `BAZELISK_HOME`; `--disk_cache` and `--repository_cache` via a managed block in `~/.bazelrc` | 50G | `disk/cas`, `disk/ac`, repository cache entries | +| `pnpm` | `PNPM_STORE_DIR`, `npm_config_store_dir` | 20G | store files (the index is kept) | +| `go` | `GOMODCACHE`, `GOCACHE` | 20G | build cache only; the module cache is read-only trees and is left alone | +| `uv` | `UV_CACHE_DIR` | 10G | entries of each versioned bucket | +| `mise` | `MISE_DATA_DIR` | 20G | installed tool versions and downloads; shims are kept | +| `git` | `BORG_GIT_MIRROR__` per mirror | none | never (objects only accumulate) | + +```json +{ + "caches": { + "git": { "mirrors": ["example-org/app"], "remote_base": "https://github.com" }, + "sccache": { "max_size": "50G" }, + "bazel": { "max_size": "40G", "bazelrc": "/var/lib/borg/.bazelrc" }, + "uv": { "enabled": false } + } +} +``` + +Notes on each kind: + +* **Garbage collection** removes whole entries (a crate, a content-addressed + blob, an installed tool version), oldest use first, until the collectable data + is at most 90% of the cap. It takes a lock that every lease holds for its whole + life, so it never overlaps a job; under constant load it simply waits for an + idle moment, and caps are soft in the meantime. +* **cargo** shares only the registry and git data. `cargo install` is pointed at + a per-lease root so no installed binary ever leaks between leases. +* **bazel**: the drone writes a marked block into the service user's `~/.bazelrc` + (or `caches.bazel.bazelrc`) and leaves the rest of the file alone. The user rc + is read after the workspace rc, so these two flags win. Bazel's disk cache is + safe for concurrent builds: entries are written to a temporary file and renamed + into place, and a collision between processes is harmless. A collected entry is + a cache miss. +* **go**: the module cache is never collected; the build cache is. +* **uv** and **mise** collection is coarse: it removes entries those tools + re-create on demand. mise serializes installs of one tool version with a lock + file. +* **git mirrors**: for each `caches.git.mirrors` entry (`owner/name`) the drone + keeps a bare mirror of the branches and tags at + `/git//.git`, refreshed every + `caches.git.interval_secs` (default 600) and before a lease starts (at most + once per 15 seconds). The variable name is upper-cased with every other + character replaced by `_`, and it is set only while the mirror exists, so a job + can test it and fall back to a normal clone: + + ``` + if [ -n "${BORG_GIT_MIRROR_EXAMPLE_ORG_APP:-}" ]; then + git clone --reference "$BORG_GIT_MIRROR_EXAMPLE_ORG_APP" src + else + git clone src + fi + ``` + + Mirrors fetch from `remote_base` (default `https://github.com`) plus + `/.git`. The drone holds no GitHub credentials: for a private + repository give the service user a credential helper, or set `remote_base` to + an ssh base such as `git@host:` with a read-only deploy key. A failed refresh + keeps the previous mirror and never blocks a job. + +#### Slots, path remapping and compile-cache keys + +A lease runs in `/slot-` and takes the lowest free slot +(`slots` in the drone config, default `max_leases`, else 16; it is also the most +leases that run at once). The path never contains the lease id, concurrent +leases always get distinct slots, and a slot is fully cleaned before it is +reused, so the next lease in a slot sees exactly the same absolute paths. The +lease-to-slot mapping is logged (`lease X -> slot N`) and kept in drone state. +Jobs see `BORG_SLOT`. + +Jobs also get `BORG_RUSTFLAGS` and `RUSTFLAGS`: `--remap-path-prefix` flags that +map **every** slot directory onto `/borg/work`, and `CARGO_HOME`, its registry +and its git data onto fixed virtual prefixes. The string lists all slots, so it +is byte-identical for every lease (a per-slot flag would itself differ between +slots and change the cache key). The same flags are written to +`$CARGO_HOME/config.toml` as `[build] rustflags`, as a fallback. Two cautions: + +* Cargo ignores `build.rustflags` and `target.*.rustflags` from config files + whenever `RUSTFLAGS` is set. A workflow that sets its own `RUSTFLAGS` replaces + ours; it should append the remap instead, e.g. + `RUSTFLAGS: "-D warnings ${BORG_RUSTFLAGS}"`. Repositories that depend on + config-file rustflags can turn the export off with + `caches.cargo.export_rustflags: false`. +* The remap does **not** make workspace crates hit the compile cache across + different slots. sccache hashes the compile's working directory and every + `CARGO_*` variable, including `CARGO_MANIFEST_DIR` (read from sccache's + `generate_hash_key`, v0.16), and no rustc flag changes those. What slots buy is + that a *reused* slot recompiles at identical paths and hits, and that + registry dependencies, whose paths are in the shared `CARGO_HOME`, hit from any + slot. `tests/sccache_slots.rs` compiles one crate in slot 0, again in a cleaned + slot 0, and in slot 1, and prints the hit counts for each. + +sccache keeps hit and miss counters only in the server's memory, so after a +lease's own server exits there is nothing on disk to read. Set +`caches.sccache.server: {"port": 4226}` to have the drone run one long-lived +sccache server per host (its configuration then applies to all jobs, and a job +whose sccache version differs from the host binary's restarts it); the drone +reads its totals with `sccache --show-stats --stats-format json` and reports +hits and misses for `sccache` in heartbeats, `/status` and `borg_cache_hits`. +Without it those fields stay empty. + +Every heartbeat carries each cache's size, cap, evicted bytes and last +collection time (and, for git, how many leases were given a mirror and how many +ran without one). The queen shows them under each drone in `/status` and as +`borg_cache_*` metrics. Sizes are re-measured every +`caches.maintenance_interval_secs` (default 600) and after leases. + +`borg drone install` takes `--git-mirror owner/name` (repeatable) and +`--git-remote-base`; edit `/etc/borg/drone.json` for the rest. + +On start the drone deletes leftover lease directories and borg-labelled docker +resources from a previous run. On `SIGTERM` it stops accepting leases and waits +up to `drain_timeout_secs` for running jobs before stopping them. + +## Updating a drone + +A drone follows the queen's `desired_version` (see [Quorum rollout](#quorum-rollout)) between +jobs, in either direction. When the version differs it stops accepting leases, +finishes the jobs it holds, and then asks the updater: it writes `update-request` +into its state directory (`state_dir` in `drone.json`). The drone itself cannot +replace its binary, because it runs unprivileged and stopping it is what drains +it. A root-owned path unit (`borg-update.path` on Linux, the `org.borg.update` +launchd daemon on macOS) starts `borg update --from-request`, which: + +1. fetches the release manifest and downloads every artifact it needs, checking + each SHA-256 against the manifest **before** anything is stopped; +2. stops the drone service (which drains running jobs), swaps the binary + atomically and keeps the previous one as `borg.prev`; +3. starts the service and waits for a fresh `health.json`, which the new drone + writes after each heartbeat the queen answered; +4. rolls back to `borg.prev` if the service does not start or does not report + healthy within `health_timeout_secs`. + +The outcome is left in `update-result.json` and sent to the queen in the next +heartbeat as `update_report`. A target that failed or rolled back on a host is +not retried there until the queen asks for a different version. A queen that is +down is retried with growing delays, and job results are retried for several +minutes after a queen outage. + +The updater reads `/etc/borg-update/update.json`, which only root can write, so +a compromised drone cannot point it at another release. The manifest is a JSON +document (`schema`, `version`, per-platform `url` and `sha256` for +`linux-x64-musl`, `linux-arm64-musl` and `darwin-arm64`, and an optional pinned +actions-runner). `https://` is required unless the config sets `allow_insecure`. +There are no signatures: the trust anchor is HTTPS to the configured manifest +URL plus the pinned digests. + +```text +borg update --check # what would change +borg update --version 0.5.0 # install exactly this version +borg update --runner-only # replace just the pinned actions-runner, no restart +``` + +`borg drone install` writes the updater configuration and units; +`--update-manifest-url` points them at another manifest and `--no-self-update` +skips them. Drones installed before this feature have to be updated once by +rerunning `borg drone install` with the new binary. diff --git a/src/config.rs b/src/config.rs index b8a73a1..eae2f49 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,6 @@ -//! Queen configuration, read from a JSON file. Nothing org-specific is baked -//! in: the organization, repositories, label prefix and class sizes all come -//! from here. +//! Queen configuration, read from a TOML or JSON file (see +//! [`crate::configfile`]). Nothing org-specific is baked in: the organization, +//! repositories, label prefix and class sizes all come from here. use std::net::SocketAddr; use std::path::{Path, PathBuf}; @@ -28,6 +28,8 @@ pub enum ConfigError { }, #[error("parse config: {0}")] Parse(#[from] serde_json::Error), + #[error("parse config {0}")] + File(#[from] crate::configfile::ParseError), #[error("invalid classes: {0}")] Classes(#[from] ClassesError), #[error("{0}")] @@ -112,8 +114,9 @@ pub struct QueenConfig { /// A drone silent for longer than this gets no new leases. #[serde(default = "default_drone_stale")] pub drone_stale_secs: i64, - /// Demand older than this triggers the starvation guard; null disables. - #[serde(default = "default_starvation")] + /// Demand older than this triggers the starvation guard. `false` (or JSON + /// `null`) disables it; TOML has no null. + #[serde(default = "default_starvation", deserialize_with = "secs_or_off")] pub starvation_after_secs: Option, #[serde(default)] pub classes: Classes, @@ -195,13 +198,38 @@ fn default_starvation() -> Option { Some(120) } +/// Seconds, or `false`/`null` for off. +fn secs_or_off<'de, D: serde::Deserializer<'de>>(d: D) -> Result, D::Error> { + #[derive(Deserialize)] + #[serde(untagged)] + enum SecsOrOff { + Secs(u64), + Off(bool), + } + match Option::::deserialize(d)? { + None | Some(SecsOrOff::Off(false)) => Ok(None), + Some(SecsOrOff::Secs(s)) => Ok(Some(s)), + Some(SecsOrOff::Off(true)) => Err(serde::de::Error::custom( + "expected seconds, or false to disable", + )), + } +} + impl QueenConfig { pub fn load(path: &Path) -> Result { let text = std::fs::read_to_string(path).map_err(|source| ConfigError::Read { path: path.to_owned(), source, })?; - Self::from_json(&text) + let cfg: Self = crate::configfile::parse(path, &text)?; + cfg.validate()?; + Ok(cfg) + } + + pub fn from_toml(text: &str) -> Result { + let cfg: Self = crate::configfile::parse(Path::new("config.toml"), text)?; + cfg.validate()?; + Ok(cfg) } pub fn from_json(json: &str) -> Result { @@ -371,5 +399,98 @@ mod tests { ) .unwrap(); assert_eq!(c.starvation_after_secs, None); + let c = QueenConfig::from_json( + &MIN.replace("\"repos\"", "\"starvation_after_secs\":false,\"repos\""), + ) + .unwrap(); + assert_eq!(c.starvation_after_secs, None); + assert!( + QueenConfig::from_json( + &MIN.replace("\"repos\"", "\"starvation_after_secs\":true,\"repos\"") + ) + .is_err() + ); + } + + const FULL_JSON: &str = r#"{ + "github": {"org": "example-org", "app_id": "7", "private_key_path": "/etc/borg/app.pem", "runner_group_id": 3}, + "repos": ["example-org/app", "example-org/lib"], + "db_path": "/var/lib/borg/queen.db", + "listen": "192.0.2.5:8080", + "classes": {"label_prefix": "example-ci-", "heavy": {"mem_mib": 20480, "vcpu": 6, "docker": true}}, + "warm_floors": [{"class": "medium", "scope": "main", "count": 2}], + "starvation_after_secs": null, + "health": {"auto_quarantine": false}, + "units": {"enabled": true}, + "supersede_janitor": {"dry_run": false} + }"#; + + const FULL_TOML: &str = r#" +repos = ["example-org/app", "example-org/lib"] +db_path = "/var/lib/borg/queen.db" +listen = "192.0.2.5:8080" +starvation_after_secs = false + +[github] +org = "example-org" +app_id = "7" +private_key_path = "/etc/borg/app.pem" +runner_group_id = 3 + +[classes] +label_prefix = "example-ci-" +heavy = { mem_mib = 20480, vcpu = 6, docker = true } + +[[warm_floors]] +class = "medium" +scope = "main" +count = 2 + +[health] +auto_quarantine = false + +[units] +enabled = true + +[supersede_janitor] +dry_run = false +"#; + + #[test] + fn toml_and_json_give_the_same_config() { + let j = QueenConfig::from_json(FULL_JSON).unwrap(); + let t = QueenConfig::from_toml(FULL_TOML).unwrap(); + assert_eq!(format!("{j:?}"), format!("{t:?}")); + assert_eq!(t.repos.len(), 2); + assert_eq!(t.starvation_after_secs, None); + assert!(t.units.enabled); + } + + #[test] + fn load_picks_the_format_and_validates() { + let dir = tempfile::tempdir().unwrap(); + let toml_path = dir.path().join("queen.toml"); + std::fs::write(&toml_path, FULL_TOML).unwrap(); + let json_path = dir.path().join("queen.json"); + std::fs::write(&json_path, FULL_JSON).unwrap(); + let t = QueenConfig::load(&toml_path).unwrap(); + let j = QueenConfig::load(&json_path).unwrap(); + assert_eq!(format!("{j:?}"), format!("{t:?}")); + + let bad = dir.path().join("bad.toml"); + std::fs::write(&bad, FULL_TOML.replace("repos = [", "nope = 1\nrepos = [")).unwrap(); + let e = QueenConfig::load(&bad).unwrap_err().to_string(); + assert!(e.contains("bad.toml") && e.contains("nope"), "{e}"); + + let empty = dir.path().join("empty.toml"); + std::fs::write( + &empty, + FULL_TOML.replace("\"example-org/app\", \"example-org/lib\"", ""), + ) + .unwrap(); + assert!(matches!( + QueenConfig::load(&empty), + Err(ConfigError::Invalid(_)) + )); } } diff --git a/src/configfile.rs b/src/configfile.rs new file mode 100644 index 0000000..835efc4 --- /dev/null +++ b/src/configfile.rs @@ -0,0 +1,106 @@ +//! Reads configuration files written as TOML or JSON. The two are equivalent: +//! the same keys, defaults and validation apply to either. +//! +//! A `.toml` or `.json` extension picks the format. Any other name is read as +//! JSON when its first non-blank character is `{`, and as TOML otherwise. + +use std::path::{Path, PathBuf}; + +use serde::de::DeserializeOwned; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Format { + Toml, + Json, +} + +impl Format { + pub fn of(path: &Path, text: &str) -> Self { + match path.extension().and_then(|e| e.to_str()) { + Some(e) if e.eq_ignore_ascii_case("toml") => Self::Toml, + Some(e) if e.eq_ignore_ascii_case("json") => Self::Json, + _ if text.trim_start().starts_with('{') => Self::Json, + _ => Self::Toml, + } + } +} + +/// A parse failure, naming the file and the place in it. TOML errors carry +/// the line, column and offending key; JSON errors the line and column. +#[derive(Debug, thiserror::Error)] +pub enum ParseError { + #[error("{path}: {source}")] + Toml { + path: PathBuf, + source: Box, + }, + #[error("{path}:{line}:{column}: {source}")] + Json { + path: PathBuf, + line: usize, + column: usize, + source: serde_json::Error, + }, +} + +/// Parses `text`, read from `path`, in the format [`Format::of`] picks. +pub fn parse(path: &Path, text: &str) -> Result { + match Format::of(path, text) { + Format::Toml => toml::from_str(text).map_err(|source| ParseError::Toml { + path: path.to_owned(), + source: Box::new(source), + }), + Format::Json => serde_json::from_str(text).map_err(|source| ParseError::Json { + path: path.to_owned(), + line: source.line(), + column: source.column(), + source, + }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Debug, serde::Deserialize, PartialEq)] + #[serde(deny_unknown_fields)] + struct Sample { + name: String, + #[serde(default)] + count: u32, + } + + #[test] + fn extension_wins_then_content() { + assert_eq!(Format::of(Path::new("a.toml"), "{"), Format::Toml); + assert_eq!(Format::of(Path::new("a.JSON"), "x = 1"), Format::Json); + assert_eq!( + Format::of(Path::new("config"), " \n{\"a\":1}"), + Format::Json + ); + assert_eq!(Format::of(Path::new("config"), "a = 1"), Format::Toml); + assert_eq!(Format::of(Path::new("config"), ""), Format::Toml); + } + + #[test] + fn both_formats_parse_to_the_same_value() { + let t: Sample = parse(Path::new("s.toml"), "name = \"x\"\ncount = 2\n").unwrap(); + let j: Sample = parse(Path::new("s.json"), r#"{"name":"x","count":2}"#).unwrap(); + assert_eq!(t, j); + } + + #[test] + fn errors_name_the_file_and_place() { + let e = parse::(Path::new("/etc/borg/q.toml"), "name = \"x\"\nnope = 1\n") + .unwrap_err() + .to_string(); + assert!(e.starts_with("/etc/borg/q.toml: "), "{e}"); + assert!(e.contains("line 2"), "{e}"); + assert!(e.contains("nope"), "{e}"); + let e = parse::(Path::new("/etc/borg/q.json"), "{\n\"nope\":1}") + .unwrap_err() + .to_string(); + assert!(e.starts_with("/etc/borg/q.json:2:"), "{e}"); + } +} diff --git a/src/drone/config.rs b/src/drone/config.rs index e61d0b9..f599697 100644 --- a/src/drone/config.rs +++ b/src/drone/config.rs @@ -1,4 +1,5 @@ -//! Drone configuration, a JSON file written by `borg drone install`. +//! Drone configuration: the JSON file `borg drone install` writes, or a +//! hand-written TOML file with the same keys (see [`crate::configfile`]). use std::collections::BTreeMap; use std::path::{Path, PathBuf}; @@ -18,6 +19,8 @@ pub enum DroneConfigError { }, #[error("parse config: {0}")] Parse(#[from] serde_json::Error), + #[error("parse config {0}")] + File(#[from] crate::configfile::ParseError), #[error("{0}")] Invalid(String), } @@ -289,7 +292,9 @@ impl DroneConfig { path: path.to_owned(), source, })?; - Self::from_json(&text) + let cfg: Self = crate::configfile::parse(path, &text)?; + cfg.validate()?; + Ok(cfg) } pub fn from_json(json: &str) -> Result { @@ -507,6 +512,34 @@ mod tests { assert_eq!((caps.max_vcpu, caps.max_mem_mib), (Some(8), Some(24576))); } + #[test] + fn toml_drone_config_matches_json() { + let dir = tempfile::tempdir().unwrap(); + let toml_path = dir.path().join("drone.toml"); + std::fs::write( + &toml_path, + r#" +queen_url = "https://queen.example.test:8443" +runner_dir = "/r" +work_root = "/w" +cache_root = "/c" +max_vcpu = 8 +extra_labels = ["on-prem"] + +[runner_env] +FOO = "bar" +"#, + ) + .unwrap(); + let json = BASE.replace( + '}', + r#","max_vcpu":8,"extra_labels":["on-prem"],"runner_env":{"FOO":"bar"}}"#, + ); + let t = DroneConfig::load(&toml_path).unwrap(); + let j = DroneConfig::from_json(&json).unwrap(); + assert_eq!(format!("{t:?}"), format!("{j:?}")); + } + #[test] fn plain_http_needs_explicit_opt_in() { let http = BASE.replace("https://", "http://"); diff --git a/src/lib.rs b/src/lib.rs index 6c778b0..26cdffb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,7 @@ pub mod api; pub mod bid; pub mod classes; pub mod config; +pub mod configfile; pub mod drone; pub mod enroll; pub mod envcontract; diff --git a/src/main.rs b/src/main.rs index 4588e3a..db42149 100644 --- a/src/main.rs +++ b/src/main.rs @@ -59,6 +59,24 @@ enum Command { #[command(subcommand)] command: QueenCommand, }, + /// Work with configuration files. + Config { + #[command(subcommand)] + command: ConfigCommand, + }, +} + +#[derive(Subcommand)] +enum ConfigCommand { + /// Parse and validate a queen (default) or drone config, TOML or JSON, and + /// print the effective configuration with every default filled in. + Check { + /// The config file. + path: PathBuf, + /// Read it as a drone config instead of a queen config. + #[arg(long)] + drone: bool, + }, } #[derive(Args)] @@ -152,7 +170,7 @@ struct ReportArgs { #[derive(Args)] struct ConfigArg { - /// Path to the queen's JSON config file. + /// Path to the queen's config file (TOML or JSON). #[arg(long, env = "BORG_QUEEN_CONFIG")] config: PathBuf, } @@ -234,7 +252,7 @@ enum RolloutCommand { #[derive(Args)] struct DroneConfigArg { - /// Path to the drone's JSON config file. + /// Path to the drone's config file (TOML or JSON). #[arg(long, env = "BORG_DRONE_CONFIG")] config: PathBuf, /// File holding the enrollment token (overrides `token_file` in the config). @@ -727,6 +745,20 @@ async fn dispatch(command: Command) -> Result<()> { Ok(()) } Command::Report(a) => report(a).await, + Command::Config { + command: ConfigCommand::Check { path, drone }, + } => { + // Debug output redacts secrets (see `borg::secret::Secret`). + if drone { + let cfg = DroneConfig::load(&path)?; + println!("{cfg:#?}"); + } else { + let cfg = QueenConfig::load(&path)?; + println!("{cfg:#?}"); + } + eprintln!("{}: ok", path.display()); + Ok(()) + } Command::Drone { command } => { tracing_subscriber::fmt() .with_writer(std::io::stderr) From 3a4a5d61d8964048c8437dd480e82e33e165f899 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sun, 11 Oct 2026 17:13:49 +0000 Subject: [PATCH 09/27] chore: project borg from Mono 5af6e0de217f --- .repository-projection.json | 6 +- Cargo.lock | 38 ++-- Cargo.toml | 2 +- docs/operations.md | 29 +++ src/github.rs | 180 ++++++++++++++++--- src/queen.rs | 108 +++++++++++- src/store.rs | 16 +- src/telemetry.rs | 164 ++++++++++++++++- tests/queen.rs | 339 ++++++++++++++++++++++++++++++++++++ tests/sccache_slots.rs | 4 + tests/support/mod.rs | 55 +++++- 11 files changed, 889 insertions(+), 52 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 89be64d..72e1b00 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "6efbce21202579ccfae91687c9aef688428d0164", + "sourceSha": "5af6e0de217fed962c7f8c796a6f2882c8df6661", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "2e48c3e0173446001b8e84070574e9add75b0db9", + "priorProjectedBase": "34e8c83d71a6e93badf1fee8996cc0e616da340d", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "bfd698d6787a8485a0097877e0b43db394bffb29733a7c97c98abf27c724571a", + "contentDigest": "35a2309bedf9953a929deb366f5a403f12c75b5dededb129a2377bfbd9a6817e", "publicationEligible": true } diff --git a/Cargo.lock b/Cargo.lock index 929492c..65b6f37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1486,11 +1486,11 @@ dependencies = [ [[package]] name = "serde_spanned" -version = "1.1.1" +version = "0.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3" dependencies = [ - "serde_core", + "serde", ] [[package]] @@ -1895,32 +1895,35 @@ dependencies = [ [[package]] name = "toml" -version = "1.1.2+spec-1.1.0" +version = "0.8.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" +checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" dependencies = [ - "serde_core", + "serde", "serde_spanned", "toml_datetime", - "toml_parser", - "winnow", + "toml_edit", ] [[package]] name = "toml_datetime" -version = "1.1.1+spec-1.1.0" +version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c" dependencies = [ - "serde_core", + "serde", ] [[package]] -name = "toml_parser" -version = "1.1.2+spec-1.1.0" +name = "toml_edit" +version = "0.22.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" dependencies = [ + "indexmap", + "serde", + "serde_spanned", + "toml_datetime", "winnow", ] @@ -2311,9 +2314,12 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "winnow" -version = "1.0.3" +version = "0.7.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +dependencies = [ + "memchr", +] [[package]] name = "wit-bindgen" diff --git a/Cargo.toml b/Cargo.toml index 3e43b72..79672c5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,7 @@ sqlx = { version = "0.8", default-features = false, features = ["runtime-tokio-r thiserror = "2.0" time = { version = "0.3", features = ["formatting", "macros", "parsing"] } tokio = { version = "1.48", features = ["io-util", "macros", "net", "rt-multi-thread", "process", "signal", "sync", "time"] } -toml = { version = "1", default-features = false, features = ["parse", "serde"] } +toml = { version = "0.8", default-features = false, features = ["parse"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } diff --git a/docs/operations.md b/docs/operations.md index 04693c2..8afb081 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -131,6 +131,35 @@ recorded. The queen keeps it for `telemetry_retention_days` (default 30). `GET /report` summarize time by job and step, queue wait per class, the cancellation share and host-caused failures per drone. +#### GitHub job feedback + +The queen also collects outcomes directly from the GitHub job records it reads +while polling queued and in-progress workflows. It accepts a job only when both +its runner ID and runner name match a lease the queen minted. Completed jobs +supply their actual execution duration, conclusion, step timings and queue wait +to the existing telemetry store, reports and learned placement profiles. +Missing or reversed timestamps remain unavailable; they never become placeholder +zeroes. GitHub failures retain `failure_class: unknown`, because an API result +does not establish whether the host caused the failure. Drone reports can add +that evidence and take precedence over this fallback. + +Once a running job is observed, its repository and job ID are retained in the +queen's existing SQLite store until completion. This lets the next poll query +the job directly after its workflow leaves the active list, including after a +queen restart. These bindings expire with `telemetry_retention_days`; a deleted +job (HTTP 404) or a completed job without usable timing drops its binding. +Collection runs after grants and the status snapshot are published. A pass has +five seconds, with at most sixteen direct lookups and one second per lookup; +the starting job rotates each tick so a slow job cannot starve other bindings. +Collection errors are logged and do not fail queue admission. No additional credentials or permissions are required beyond +the existing Actions read access. + +This is polling coverage: a job whose entire workflow starts and finishes +between polls can be missed. Reports count only collected outcomes. The existing +profile minimum sample and confidence rules still apply, and profiles refresh +on `profiles.refresh_secs`. This change does not alter owner resource caps, +class sizes or warm-floor settings. + #### What a drone-side parser should produce The queen validates but does not interpret these fields, so the drone owns diff --git a/src/github.rs b/src/github.rs index 4aa6db7..70056fd 100644 --- a/src/github.rs +++ b/src/github.rs @@ -15,6 +15,99 @@ use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; use crate::secret::Secret; +use crate::telemetry::{Conclusion, FailureClass, JobTelemetry, StepTiming}; + +/// GitHub-owned job metadata; never contains logs, outputs or environment values. +#[derive(Debug, Clone, Deserialize)] +pub struct WorkflowJob { + pub id: u64, + #[serde(default)] + pub run_id: u64, + #[serde(default)] + pub name: String, + pub status: String, + pub conclusion: Option, + pub labels: Option>, + pub created_at: String, + pub started_at: Option, + pub completed_at: Option, + pub runner_id: Option, + pub runner_name: Option, + #[serde(default)] + pub steps: Vec, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct JobStep { + pub name: String, + pub number: u64, + pub conclusion: Option, + pub started_at: Option, + pub completed_at: Option, +} + +fn timestamp(value: &str) -> Option { + OffsetDateTime::parse(value, &Rfc3339) + .ok() + .map(|t| t.unix_timestamp()) +} + +fn elapsed(start: Option<&str>, end: Option<&str>) -> Option { + u64::try_from(timestamp(end?)?.checked_sub(timestamp(start?)?)?).ok() +} + +impl WorkflowJob { + pub fn queued(&self, repo: &str) -> Result, GithubError> { + if self.status != "queued" { + return Ok(None); + } + let created_at = OffsetDateTime::parse(&self.created_at, &Rfc3339) + .map_err(|e| GithubError::Decode(format!("job {} created_at: {e}", self.id)))?; + Ok(Some(QueuedJob { + id: self.id, + repo: repo.to_owned(), + labels: self.labels.clone().ok_or_else(|| { + GithubError::Decode(format!("job {} missing queued labels", self.id)) + })?, + created_at, + })) + } + + /// A completed job's own execution time, excluding runner startup and idle time. + /// Missing or invalid timing is unavailable evidence, never a zero-duration sample. + pub fn telemetry(&self, repo: &str) -> Option<(JobTelemetry, u64, i64)> { + if self.status != "completed" { + return None; + } + let duration = elapsed(self.started_at.as_deref(), self.completed_at.as_deref())?; + let finished = timestamp(self.completed_at.as_deref()?)?; + let mut steps: Vec<_> = self.steps.iter().collect(); + steps.sort_by_key(|step| step.number); + let telemetry = JobTelemetry { + job_name: self.name.clone(), + repo: repo.to_owned(), + run_id: self.run_id, + conclusion: self.conclusion.unwrap_or_default(), + failure_class: FailureClass::Unknown, + queue_wait_secs: elapsed(Some(&self.created_at), self.started_at.as_deref()), + steps: steps + .into_iter() + .filter_map(|step| { + Some(StepTiming { + name: step.name.clone(), + conclusion: step.conclusion.unwrap_or_default(), + duration_secs: elapsed( + step.started_at.as_deref(), + step.completed_at.as_deref(), + )?, + }) + }) + .collect(), + } + .sanitized()?; + Some((telemetry, duration, finished)) + } +} /// Upper bound on pages fetched per list call (100 items per page). const MAX_PAGES: u32 = 20; @@ -223,8 +316,8 @@ impl GithubClient { self.send(Method::GET, path, rb).await } - /// Jobs that are `queued` in `repo`, across queued and in-progress runs. - pub async fn queued_jobs(&self, repo: &str) -> Result, GithubError> { + /// Jobs in queued and in-progress workflows, including their completed jobs. + pub async fn workflow_jobs(&self, repo: &str) -> Result, GithubError> { #[derive(Deserialize)] struct Runs { workflow_runs: Vec, @@ -235,14 +328,7 @@ impl GithubClient { } #[derive(Deserialize)] struct Jobs { - jobs: Vec, - } - #[derive(Deserialize)] - struct Job { - id: u64, - status: String, - labels: Vec, - created_at: String, + jobs: Vec, } let mut run_ids = BTreeSet::new(); @@ -267,18 +353,11 @@ impl GithubClient { .get(&format!("/repos/{repo}/actions/runs/{run}/jobs?filter=latest&per_page=100&page={page}")) .await?; let n = jobs.jobs.len(); - for j in jobs.jobs.into_iter().filter(|j| j.status == "queued") { - let created_at = - OffsetDateTime::parse(&j.created_at, &Rfc3339).map_err(|e| { - GithubError::Decode(format!("job {} created_at: {e}", j.id)) - })?; - out.push(QueuedJob { - id: j.id, - repo: repo.to_owned(), - labels: j.labels, - created_at, - }); - } + out.extend(jobs.jobs.into_iter().map(|mut job| { + // The enclosing run is authoritative even if an older API omits run_id. + job.run_id = run; + job + })); if n < 100 { break; } @@ -287,6 +366,20 @@ impl GithubClient { Ok(out) } + /// Compatibility surface for consumers that only need queued demand. + pub async fn queued_jobs(&self, repo: &str) -> Result, GithubError> { + self.workflow_jobs(repo) + .await? + .iter() + .filter_map(|j| j.queued(repo).transpose()) + .collect() + } + + pub async fn job(&self, repo: &str, job_id: u64) -> Result { + self.get(&format!("/repos/{repo}/actions/jobs/{job_id}")) + .await + } + /// Every self-hosted runner registered to the organization. pub async fn list_runners(&self) -> Result, GithubError> { #[derive(Deserialize)] @@ -373,3 +466,46 @@ fn static_method(m: &Method) -> &'static str { _ => "REQUEST", } } + +#[cfg(test)] +mod job_tests { + use super::*; + use serde_json::json; + + #[test] + fn missing_and_reversed_timestamps_do_not_invent_timings() { + let mut value = json!({"id": 7, "run_id": 1, "name": "lint", "status": "completed", "conclusion": "success", + "created_at": "bad", "started_at": "2026-10-11T10:00:00Z", "completed_at": "2026-10-11T10:00:30Z", + "steps": [ + {"name": "second", "number": 2, "conclusion": "success", "started_at": "2026-10-11T10:00:10Z", "completed_at": "2026-10-11T10:00:20Z"}, + {"name": "first", "number": 1, "conclusion": "success", "started_at": "2026-10-11T10:00:00Z", "completed_at": "2026-10-11T10:00:10Z"}, + {"name": "unavailable", "number": 3, "conclusion": null, "started_at": null, "completed_at": null} + ]}); + let job: WorkflowJob = serde_json::from_value(value.clone()).unwrap(); + let (telemetry, duration, _) = job.telemetry("example-org/app").unwrap(); + assert_eq!(duration, 30); + assert_eq!(telemetry.queue_wait_secs, None); + assert_eq!( + telemetry + .steps + .iter() + .map(|s| s.name.as_str()) + .collect::>(), + ["first", "second"] + ); + value["completed_at"] = json!("2026-10-11T09:00:00Z"); + assert!( + serde_json::from_value::(value.clone()) + .unwrap() + .telemetry("example-org/app") + .is_none() + ); + value["started_at"] = serde_json::Value::Null; + assert!( + serde_json::from_value::(value) + .unwrap() + .telemetry("example-org/app") + .is_none() + ); + } +} diff --git a/src/queen.rs b/src/queen.rs index 0207e55..2ddf1dc 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -11,6 +11,7 @@ //! Any GitHub failure before minting aborts the tick without minting: acting //! on a partial view of the queue would over- or under-grant. +use std::borrow::Cow; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::sync::Arc; use std::sync::RwLock; @@ -29,7 +30,7 @@ use crate::bid::{ClassPressure, ScoreWeights}; use crate::classes::{Class, Scope}; use crate::config::QueenConfig; use crate::envcontract::blocked_classes; -use crate::github::{GithubClient, GithubError, QueuedJob, Runner}; +use crate::github::{GithubClient, GithubError, QueuedJob, Runner, WorkflowJob}; use crate::health::{DroneHealth, DroneInfo, Outcome, OutcomeKind, evaluate}; use crate::homeostasis::{Controller, KeyInput, KeyReport}; use crate::janitor::Janitor; @@ -349,8 +350,14 @@ impl Queen { self.reap(&ours, &active).await; let mut jobs = Vec::new(); + let mut observed = Vec::new(); for repo in &self.cfg.repos { - jobs.extend(self.github.queued_jobs(repo).await?); + for job in self.github.workflow_jobs(repo).await? { + if let Some(queued) = job.queued(repo)? { + jobs.push(queued); + } + observed.push((repo.as_str(), job)); + } } let (queue, demand, queue_ages) = self.summarize(&jobs, now); @@ -521,6 +528,18 @@ impl Queen { if let Err(e) = self.store.prune_telemetry(now_ts - keep).await { tracing::warn!("prune telemetry failed: {e}"); } + // Grants and the snapshot are published before optional feedback. Bound + // this pass so unavailable job history cannot stall subsequent ticks. + match tokio::time::timeout( + std::time::Duration::from_secs(5), + self.collect_job_telemetry(&observed, now_ts), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(e)) => tracing::warn!("GitHub job telemetry collection failed: {e}"), + Err(_) => tracing::warn!("GitHub job telemetry collection exceeded its deadline"), + } Ok(()) } @@ -819,6 +838,91 @@ impl Queen { .view(now, settings.enabled) } + /// Fill the existing telemetry owner from GitHub's job records. The durable + /// binding is needed because a workflow can finish between queue polls; a + /// transient cache would lose that final outcome on a queen restart. + async fn collect_job_telemetry( + &self, + observed: &[(&str, WorkflowJob)], + now: i64, + ) -> Result<(), QueenError> { + let prefix = self.runner_prefix(); + let by_job: HashMap<_, _> = observed + .iter() + .map(|(repo, job)| ((*repo, job.id), job)) + .collect(); + for (repo, job) in observed { + if (job.status == "completed" && job.telemetry(repo).is_none()) + || !job + .runner_name + .as_deref() + .is_some_and(|name| name.starts_with(&prefix)) + { + continue; + } + if let (Some(id), Some(name)) = (job.runner_id, job.runner_name.as_deref()) + && let Some(lease) = self.store.lease_for_runner(id, name).await? + { + self.store.remember_job(&lease, repo, job.id, now).await?; + } + } + let mut pending = self.store.pending_jobs().await?; + if !pending.is_empty() { + // Rotate the starting point so a slow old job cannot starve later ones. + let tick = + now.div_euclid(i64::try_from(self.cfg.poll_interval_secs).unwrap_or(i64::MAX)); + let offset = tick.rem_euclid(pending.len() as i64) as usize; + pending.rotate_left(offset); + } + let mut lookups = 0; + for (lease, repo, id) in pending { + let job = if let Some(job) = by_job.get(&(repo.as_str(), id)) { + Cow::Borrowed(*job) + } else { + if lookups >= 16 { + continue; + } + lookups += 1; + match tokio::time::timeout( + std::time::Duration::from_secs(1), + self.github.job(&repo, id), + ) + .await + { + Ok(Ok(job)) => Cow::Owned(job), + Ok(Err(GithubError::Status { status: 404, .. })) => { + self.store.forget_job(&lease.id).await?; + continue; + } + Ok(Err(e)) => { + tracing::warn!("job {repo}/{id} telemetry unavailable: {e}"); + continue; + } + Err(_) => { + tracing::warn!("job {repo}/{id} telemetry lookup timed out"); + continue; + } + } + }; + if job.id != id + || job.runner_id != Some(lease.runner_id) + || job.runner_name.as_deref() != Some(lease.runner_name.as_str()) + { + continue; + } + if let Some((telemetry, duration, finished)) = job.telemetry(&repo) { + self.store + .record_github_telemetry(&lease, &telemetry, duration, finished) + .await?; + } + if job.status == "completed" { + // A terminal response with unusable timing cannot yield a sample. + self.store.forget_job(&lease.id).await?; + } + } + Ok(()) + } + /// Learned drone profiles, recomputed from telemetry at most every /// `profiles.refresh_secs`. Empty when profiles are off. async fn profiles(&self, now: i64) -> Vec { diff --git a/src/store.rs b/src/store.rs index 63adc6c..3c0b897 100644 --- a/src/store.rs +++ b/src/store.rs @@ -1,8 +1,8 @@ //! SQLite bookkeeping: enrolled drones and lease history. //! -//! The file is small and disposable in the sense that the queen re-derives -//! truth from GitHub plus drone heartbeats after a restart. Enrollment hashes -//! are the one thing worth backing up. +//! Back up enrollment hashes and pending job bindings. The queen re-derives +//! active capacity from GitHub and drone heartbeats, but losing a binding can +//! lose the completion sample once its workflow leaves the active queue. use std::path::Path; use std::str::FromStr; @@ -141,6 +141,14 @@ CREATE TABLE IF NOT EXISTS leases ( peak_rss_mib INTEGER ); CREATE INDEX IF NOT EXISTS leases_state ON leases(state); +CREATE INDEX IF NOT EXISTS leases_runner ON leases(runner_id, runner_name); +-- Recover observed job identity after a queen restart, until GitHub reports completion. +CREATE TABLE IF NOT EXISTS lease_jobs ( + lease_id TEXT PRIMARY KEY REFERENCES leases(id) ON DELETE CASCADE, + repo TEXT NOT NULL, + job_id INTEGER NOT NULL, + observed_at INTEGER NOT NULL +); CREATE TABLE IF NOT EXISTS job_telemetry ( lease_id TEXT PRIMARY KEY, drone TEXT NOT NULL, @@ -899,7 +907,7 @@ impl Store { Ok(()) } - fn lease_from_row(r: &sqlx::sqlite::SqliteRow) -> Result { + pub(crate) fn lease_from_row(r: &sqlx::sqlite::SqliteRow) -> Result { let jit: Option = r.get("jit_config"); Ok(LeaseRow { id: r.get("id"), diff --git a/src/telemetry.rs b/src/telemetry.rs index ae0580e..3c7f0ad 100644 --- a/src/telemetry.rs +++ b/src/telemetry.rs @@ -167,6 +167,74 @@ fn to_u64(v: i64) -> u64 { } impl Store { + /// Match both GitHub runner coordinates; never attribute a static or reused name. + pub async fn lease_for_runner( + &self, + id: u64, + name: &str, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT * FROM leases WHERE runner_id = ? AND runner_name = ?") + .bind(to_i64(id)) + .bind(name) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(Self::lease_from_row).transpose() + } + + /// Persist one GitHub job per ephemeral lease. Existing drone telemetry wins. + pub async fn remember_job( + &self, + lease: &LeaseRow, + repo: &str, + job_id: u64, + now: i64, + ) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT OR IGNORE INTO lease_jobs (lease_id, repo, job_id, observed_at) + SELECT ?, ?, ?, ? WHERE NOT EXISTS (SELECT 1 FROM job_telemetry WHERE lease_id = ?)", + ) + .bind(&lease.id) + .bind(repo) + .bind(to_i64(job_id)) + .bind(now) + .bind(&lease.id) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn pending_jobs(&self) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT leases.*, lease_jobs.repo AS job_repo, lease_jobs.job_id AS observed_job_id + FROM lease_jobs JOIN leases ON leases.id = lease_jobs.lease_id + WHERE NOT EXISTS (SELECT 1 FROM job_telemetry WHERE job_telemetry.lease_id = leases.id) + ORDER BY lease_jobs.observed_at, lease_jobs.lease_id", + ) + .fetch_all(&self.pool) + .await?; + rows.iter() + .map(|row| { + Ok(( + Self::lease_from_row(row)?, + row.get("job_repo"), + to_u64(row.get("observed_job_id")), + )) + }) + .collect() + } + + pub async fn forget_job(&self, lease_id: &str) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("DELETE FROM lease_jobs WHERE lease_id = ?") + .bind(lease_id) + .execute(&self.pool) + .await?; + Ok(()) + } + /// Store (or replace) the telemetry of a finished lease. pub async fn record_telemetry( &self, @@ -174,15 +242,46 @@ impl Store { t: &JobTelemetry, duration_secs: u64, now: i64, + ) -> Result<(), StoreError> { + self.write_telemetry(lease, t, duration_secs, now, true) + .await + } + + /// Fill a missing report atomically; a concurrent drone report takes precedence. + pub async fn record_github_telemetry( + &self, + lease: &LeaseRow, + t: &JobTelemetry, + duration_secs: u64, + now: i64, + ) -> Result<(), StoreError> { + self.write_telemetry(lease, t, duration_secs, now, false) + .await + } + + async fn write_telemetry( + &self, + lease: &LeaseRow, + t: &JobTelemetry, + duration_secs: u64, + now: i64, + replace: bool, ) -> Result<(), StoreError> { let steps = serde_json::to_string(&t.steps).map_err(|e| StoreError::Corrupt(e.to_string()))?; // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database sqlx::query( - "INSERT OR REPLACE INTO job_telemetry + "INSERT INTO job_telemetry (lease_id, drone, class, scope, repo, run_id, job_name, conclusion, failure_class, duration_secs, queue_wait_secs, steps, finished_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(lease_id) DO UPDATE SET + drone=excluded.drone, class=excluded.class, scope=excluded.scope, + repo=excluded.repo, run_id=excluded.run_id, job_name=excluded.job_name, + conclusion=excluded.conclusion, failure_class=excluded.failure_class, + duration_secs=excluded.duration_secs, queue_wait_secs=excluded.queue_wait_secs, + steps=excluded.steps, finished_at=excluded.finished_at + WHERE ?", ) .bind(&lease.id) .bind(&lease.drone) @@ -197,6 +296,7 @@ impl Store { .bind(t.queue_wait_secs.map(to_i64)) .bind(steps) .bind(now) + .bind(replace) .execute(&self.pool) .await?; Ok(()) @@ -243,6 +343,14 @@ impl Store { /// Delete reports finished before `before`; returns how many. pub async fn prune_telemetry(&self, before: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "DELETE FROM lease_jobs WHERE observed_at < ? OR EXISTS + (SELECT 1 FROM job_telemetry WHERE job_telemetry.lease_id = lease_jobs.lease_id)", + ) + .bind(before) + .execute(&self.pool) + .await?; // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let r = sqlx::query("DELETE FROM job_telemetry WHERE finished_at < ?") .bind(before) @@ -755,7 +863,16 @@ mod tests { FailureClass::Host, ) .telemetry; + let mut fallback = t.clone(); + fallback.failure_class = FailureClass::Unknown; + s.record_github_telemetry(&lease, &fallback, 20, 450) + .await + .unwrap(); s.record_telemetry(&lease, &t, 77, 500).await.unwrap(); + s.record_github_telemetry(&lease, &fallback, 20, 550) + .await + .unwrap(); + assert_eq!(s.telemetry_since(0, 10).await.unwrap()[0].telemetry, t); s.record_telemetry(&lease, &t, 88, 600).await.unwrap(); let rows = s.telemetry_since(0, 10).await.unwrap(); assert_eq!(rows.len(), 1); @@ -766,4 +883,47 @@ mod tests { assert_eq!(s.prune_telemetry(601).await.unwrap(), 1); assert!(s.telemetry_since(0, 10).await.unwrap().is_empty()); } + #[tokio::test] + async fn pending_job_identity_survives_restart_and_expires_with_telemetry() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("queen.sqlite"); + let s = Store::open(&path).await.unwrap(); + s.enroll("d1", "h", 0).await.unwrap(); + let lease = LeaseRow { + id: "l1".into(), + drone: "d1".into(), + class: Class::Light, + scope: Scope::Pr, + labels: vec![], + runner_name: "borg-d1-l1".into(), + runner_id: 123, + jit_config: None, + state: crate::store::LeaseStatus::Running, + created_at: 1, + expires_at: 500, + }; + s.insert_lease(&lease).await.unwrap(); + s.remember_job(&lease, "example-org/app", 77, 100) + .await + .unwrap(); + s.remember_job(&lease, "other/repo", 88, 200).await.unwrap(); + s.pool.close().await; + let reopened = Store::open(&path).await.unwrap(); + let pending = reopened.pending_jobs().await.unwrap(); + assert_eq!(pending.len(), 1); + assert_eq!( + (&pending[0].0.id, pending[0].1.as_str(), pending[0].2), + (&lease.id, "example-org/app", 77) + ); + reopened.prune_telemetry(101).await.unwrap(); + assert!(reopened.pending_jobs().await.unwrap().is_empty()); + reopened + .remember_job(&lease, "example-org/app", 77, 200) + .await + .unwrap(); + assert_eq!(reopened.remove_drone("d1").await.unwrap(), Some(1)); + assert!(reopened.pending_jobs().await.unwrap().is_empty()); + assert!(reopened.lease(&lease.id).await.unwrap().is_none()); + assert!(reopened.drones().await.unwrap().is_empty()); + } } diff --git a/tests/queen.rs b/tests/queen.rs index ab41fb6..dac71ad 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -689,3 +689,342 @@ async fn reclaim_reports_are_audited_once_and_shown_in_status_and_metrics() { assert!(metrics.contains(line), "missing {line}\n{metrics}"); } } + +#[tokio::test] +async fn github_job_outcomes_feed_profiles_without_drone_telemetry() { + let h = Harness::start(json!({"profiles": {"refresh_secs": 1, "min_samples": 1}})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let lease_id = lease_for(&h, &token, "borg-ci-light-pr", 71).await; + let lease = h.queen.store().lease(&lease_id).await.unwrap().unwrap(); + let end = OffsetDateTime::now_utc(); + let start = end - time::Duration::seconds(30); + let created = start - time::Duration::seconds(12); + h.state().jobs = vec![json!({ + "id": 71, "run_id": 1, "name": "lint (shard 2)", + "status": "completed", "conclusion": "success", "labels": ["borg-ci-light-pr"], + "runner_id": lease.runner_id, "runner_name": lease.runner_name, + "created_at": created.format(&time::format_description::well_known::Rfc3339).unwrap(), + "started_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "completed_at": end.format(&time::format_description::well_known::Rfc3339).unwrap(), + "steps": [{"name": "cargo check", "number": 1, "conclusion": "success", + "started_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "completed_at": end.format(&time::format_description::well_known::Rfc3339).unwrap()}] + })]; + h.tick().await; + let rows = h.queen.store().telemetry_since(0, 100).await.unwrap(); + assert_eq!( + rows.len(), + 1, + "GitHub completion must close the learning loop" + ); + assert_eq!(rows[0].lease_id, lease_id); + assert_eq!(rows[0].duration_secs, 30); + assert_eq!(rows[0].telemetry.queue_wait_secs, Some(12)); + assert_eq!(rows[0].telemetry.steps[0].duration_secs, 30); + // The profile refresh has its own bounded cadence. + h.queen + .tick(end + time::Duration::seconds(2)) + .await + .unwrap(); + assert_eq!(h.queen.snapshot().profiles[0].jobs, 1); + h.tick().await; + assert_eq!( + h.queen.store().telemetry_since(0, 100).await.unwrap(), + rows, + "repeated observations must preserve the completion timestamp and one sample" + ); +} + +#[tokio::test] +async fn observed_running_job_is_collected_after_workflow_leaves_the_queue() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let lease_id = lease_for(&h, &token, "borg-ci-light-pr", 72).await; + let lease = h.queen.store().lease(&lease_id).await.unwrap().unwrap(); + let start = OffsetDateTime::now_utc() - time::Duration::seconds(45); + h.state().jobs = vec![json!({ + "id": 72, "run_id": 1, "name": "rust-test", "status": "in_progress", + "conclusion": null, "labels": ["borg-ci-light-pr"], + "runner_id": lease.runner_id, "runner_name": lease.runner_name, + "created_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "started_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "completed_at": null, "steps": [] + })]; + h.tick().await; + assert_eq!(h.queen.store().pending_jobs().await.unwrap().len(), 1); + assert!( + h.queen + .store() + .telemetry_since(0, 100) + .await + .unwrap() + .is_empty() + ); + h.state().queued_runs.clear(); + h.state().jobs[0]["status"] = json!("completed"); + h.state().jobs[0]["conclusion"] = json!("failure"); + h.state().jobs[0]["completed_at"] = json!( + OffsetDateTime::now_utc() + .format(&time::format_description::well_known::Rfc3339) + .unwrap() + ); + h.tick().await; + let rows = h.queen.store().telemetry_since(0, 100).await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!( + rows[0].telemetry.conclusion, + borg::telemetry::Conclusion::Failure + ); + assert_eq!( + rows[0].telemetry.failure_class, + borg::telemetry::FailureClass::Unknown, + "GitHub failure alone must never quarantine a host" + ); + assert!(h.queen.store().pending_jobs().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn github_telemetry_requires_exact_runner_identity_and_preserves_drone_evidence() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let lease_id = lease_for(&h, &token, "borg-ci-light-pr", 73).await; + let lease = h.queen.store().lease(&lease_id).await.unwrap().unwrap(); + let start = OffsetDateTime::now_utc() - time::Duration::seconds(60); + let end = start + time::Duration::seconds(30); + let mut job = json!({ + "id": 73, "run_id": 1, "name": "untrusted name", "status": "completed", "conclusion": "success", + "runner_id": lease.runner_id + 1, "runner_name": lease.runner_name, + "created_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "started_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "completed_at": end.format(&time::format_description::well_known::Rfc3339).unwrap() + }); + h.state().jobs = vec![job.clone()]; + h.tick().await; + assert!( + h.queen + .store() + .telemetry_since(0, 100) + .await + .unwrap() + .is_empty() + ); + job["runner_id"] = json!(lease.runner_id); + job["runner_name"] = json!("other-runner"); + h.state().jobs = vec![job.clone()]; + h.tick().await; + assert!( + h.queen + .store() + .telemetry_since(0, 100) + .await + .unwrap() + .is_empty() + ); + assert_eq!( + post_result( + &h, + &token, + &lease_id, + json!({ + "exit_code": 1, "duration_secs": 60, + "telemetry": {"job_name": "drone evidence", "repo": "example-org/app", "run_id": 1, + "conclusion": "failure", "failure_class": "host"} + }) + ) + .await, + reqwest::StatusCode::NO_CONTENT + ); + job["runner_name"] = json!(lease.runner_name); + h.state().jobs = vec![job]; + h.tick().await; + let rows = h.queen.store().telemetry_since(0, 100).await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].telemetry.job_name, "drone evidence"); + assert_eq!( + rows[0].telemetry.failure_class, + borg::telemetry::FailureClass::Host + ); +} + +fn running_feedback_job(lease: &borg::store::LeaseRow, id: u64) -> Value { + let start = OffsetDateTime::now_utc() - time::Duration::seconds(45); + json!({ + "id": id, "run_id": 1, "name": "rust-test", "status": "in_progress", + "conclusion": null, "runner_id": lease.runner_id, "runner_name": lease.runner_name, + "created_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "started_at": start.format(&time::format_description::well_known::Rfc3339).unwrap(), + "completed_at": null + }) +} + +#[tokio::test] +async fn slow_feedback_does_not_delay_admission_or_overwrite_concurrent_drone_evidence() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let lease_id = lease_for(&h, &token, "borg-ci-light-pr", 80).await; + let lease = h.queen.store().lease(&lease_id).await.unwrap().unwrap(); + h.state().jobs = vec![running_feedback_job(&lease, 80)]; + h.tick().await; + let gate = std::sync::Arc::new(tokio::sync::Notify::new()); + { + let mut state = h.state(); + state.jobs[0]["status"] = json!("completed"); + state.jobs[0]["conclusion"] = json!("success"); + state.jobs[0]["completed_at"] = json!( + OffsetDateTime::now_utc() + .format(&time::format_description::well_known::Rfc3339) + .unwrap() + ); + state.queued_runs.clear(); + state.job_gate = Some(gate.clone()); + } + // A second workflow is queued while the old one has left the active list. + h.queue_job(81, "borg-ci-medium-pr", 5); + // Completed history is absent from the active workflow's jobs response. + h.state().hide_completed_jobs = true; + let queen = h.queen.clone(); + let tick = tokio::spawn(async move { queen.tick(OffsetDateTime::now_utc()).await }); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while h.state().job_requests == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_eq!( + h.state().mints.len(), + 2, + "admission must precede a slow completion lookup" + ); + assert_eq!( + post_result( + &h, + &token, + &lease_id, + json!({ + "exit_code": 1, "duration_secs": 60, + "telemetry": {"job_name": "drone evidence", "repo": "example-org/app", "run_id": 1, + "conclusion": "failure", "failure_class": "host"} + }) + ) + .await, + reqwest::StatusCode::NO_CONTENT + ); + gate.notify_one(); + tick.await.unwrap().unwrap(); + let rows = h.queen.store().telemetry_since(0, 100).await.unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!( + rows[0].telemetry.failure_class, + borg::telemetry::FailureClass::Host + ); + assert_eq!(rows[0].telemetry.job_name, "drone evidence"); +} + +#[tokio::test] +async fn terminal_unusable_or_deleted_jobs_are_not_retried() { + for deleted in [false, true] { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let id = lease_for(&h, &token, "borg-ci-light-pr", 82).await; + let lease = h.queen.store().lease(&id).await.unwrap().unwrap(); + h.state().jobs = vec![running_feedback_job(&lease, 82)]; + h.tick().await; + h.state().queued_runs.clear(); + if deleted { + h.state().jobs.clear(); + } else { + h.state().jobs[0]["status"] = json!("completed"); + h.state().jobs[0]["conclusion"] = json!("cancelled"); + // No completion timestamp: no duration may be invented. + } + h.tick().await; + h.tick().await; + assert_eq!(h.state().job_requests, 1); + assert!(h.queen.store().pending_jobs().await.unwrap().is_empty()); + assert!( + h.queen + .store() + .telemetry_since(0, 100) + .await + .unwrap() + .is_empty() + ); + } +} + +#[tokio::test] +async fn a_new_queen_recovers_completion_from_file_backed_bindings() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("queen.sqlite"); + let h = Harness::start_file(json!({}), &path).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let id = lease_for(&h, &token, "borg-ci-light-pr", 83).await; + let lease = h.queen.store().lease(&id).await.unwrap().unwrap(); + h.state().jobs = vec![running_feedback_job(&lease, 83)]; + h.tick().await; + let mut job = h.state().jobs[0].clone(); + job["status"] = json!("completed"); + job["conclusion"] = json!("success"); + job["completed_at"] = json!( + OffsetDateTime::now_utc() + .format(&time::format_description::well_known::Rfc3339) + .unwrap() + ); + let restarted = Harness::start_file(json!({}), &path).await; + restarted.state().runners = h.state().runners.clone(); + restarted.state().jobs = vec![job]; + restarted.tick().await; + restarted.tick().await; + assert_eq!(restarted.state().job_requests, 1); + assert_eq!( + restarted + .queen + .store() + .telemetry_since(0, 100) + .await + .unwrap() + .len(), + 1 + ); +} + +#[tokio::test] +async fn a_timed_out_feedback_lookup_is_bounded_and_retried_later() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + let id = lease_for(&h, &token, "borg-ci-light-pr", 84).await; + let lease = h.queen.store().lease(&id).await.unwrap().unwrap(); + h.state().jobs = vec![running_feedback_job(&lease, 84)]; + h.tick().await; + let gate = std::sync::Arc::new(tokio::sync::Notify::new()); + h.state().queued_runs.clear(); + h.state().job_gate = Some(gate.clone()); + tokio::time::timeout(std::time::Duration::from_secs(3), h.tick()) + .await + .unwrap(); + assert_eq!(h.state().job_requests, 1); + assert_eq!(h.queen.store().pending_jobs().await.unwrap().len(), 1); + h.state().job_gate = None; + h.state().jobs[0]["status"] = json!("completed"); + h.state().jobs[0]["conclusion"] = json!("success"); + h.state().jobs[0]["completed_at"] = json!( + OffsetDateTime::now_utc() + .format(&time::format_description::well_known::Rfc3339) + .unwrap() + ); + gate.notify_one(); + h.tick().await; + assert_eq!(h.state().job_requests, 2); + assert_eq!( + h.queen.store().telemetry_since(0, 100).await.unwrap().len(), + 1 + ); +} diff --git a/tests/sccache_slots.rs b/tests/sccache_slots.rs index b62180a..857310c 100644 --- a/tests/sccache_slots.rs +++ b/tests/sccache_slots.rs @@ -62,6 +62,10 @@ fn workspace_crates_hit_the_compile_cache_when_a_slot_is_reused() { env.extend([ ("RUSTC_WRAPPER".to_owned(), "sccache".to_owned()), ("CARGO_INCREMENTAL".to_owned(), "0".to_owned()), + ( + "CARGO_TARGET_DIR".to_owned(), + tmp.path().join("target").to_string_lossy().into_owned(), + ), ( "CARGO_HOME".to_owned(), cargo_home.to_string_lossy().into_owned(), diff --git a/tests/support/mod.rs b/tests/support/mod.rs index d5b2d90..2fe664d 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -41,6 +41,9 @@ pub struct MockState { pub jwt_rejections: u32, pub fail_mint: bool, pub fail_runner_list: bool, + pub job_requests: u32, + pub hide_completed_jobs: bool, + pub job_gate: Option>, /// Queued `pull_request` runs as the janitor lists them. pub pr_runs: Vec, /// Current head sha per pull request number. @@ -172,7 +175,38 @@ async fn jobs( Path(_p): Path<(String, String, u64)>, ) -> Result, StatusCode> { check_token(&h)?; - Ok(Json(json!({"jobs": lock(&s).jobs.clone()}))) + let state = lock(&s); + let jobs: Vec<_> = state + .jobs + .iter() + .filter(|job| !state.hide_completed_jobs || job["status"] != "completed") + .cloned() + .collect(); + Ok(Json(json!({"jobs": jobs}))) +} + +async fn one_job( + State(s): State, + h: HeaderMap, + Path((_o, _r, id)): Path<(String, String, u64)>, +) -> Result, StatusCode> { + check_token(&h)?; + let (job, gate) = { + let mut state = lock(&s); + state.job_requests += 1; + ( + state + .jobs + .iter() + .find(|job| job["id"].as_u64() == Some(id)) + .cloned(), + state.job_gate.clone(), + ) + }; + if let Some(gate) = gate { + gate.notified().await; + } + job.map(Json).ok_or(StatusCode::NOT_FOUND) } async fn one_run( @@ -321,6 +355,23 @@ impl Harness { /// Like `start`, with a fake tailnet identity table (callers all connect /// from 127.0.0.1) behind `tailnet_enroll` when the config enables it. pub async fn start_with(extra: Value, identities: borg::tailnet::StaticIdentities) -> Self { + Self::start_with_store(extra, identities, Store::open_in_memory().await.unwrap()).await + } + + pub async fn start_file(extra: Value, path: &std::path::Path) -> Self { + Self::start_with_store( + extra, + borg::tailnet::StaticIdentities::default(), + Store::open(path).await.unwrap(), + ) + .await + } + + async fn start_with_store( + extra: Value, + identities: borg::tailnet::StaticIdentities, + store: Store, + ) -> Self { let mock: Shared = Arc::default(); let (private_pem, public_pem) = generate_key(); lock(&mock).public_pem = public_pem; @@ -330,6 +381,7 @@ impl Harness { .route("/app/installations/{id}/access_tokens", post(access_token)) .route("/repos/{owner}/{repo}/actions/runs", get(runs)) .route("/repos/{owner}/{repo}/actions/runs/{run}/jobs", get(jobs)) + .route("/repos/{owner}/{repo}/actions/jobs/{id}", get(one_job)) .route("/repos/{owner}/{repo}/actions/runs/{run}", get(one_run)) .route( "/repos/{owner}/{repo}/actions/runs/{run}/cancel", @@ -358,7 +410,6 @@ impl Harness { } let cfg = QueenConfig::from_json(&cfg.to_string()).unwrap(); let github = GithubClient::new(&gh_url, ORG, "12345", private_pem.as_bytes()).unwrap(); - let store = Store::open_in_memory().await.unwrap(); let janitor = cfg.supersede_janitor.clone().map(|settings| { borg::janitor::Janitor::new(&gh_url, borg::secret::Secret::new(JANITOR_TOKEN), settings) .unwrap() From e20de244f34485489abb2176ecd760b3134f5053 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sun, 11 Oct 2026 22:54:25 +0000 Subject: [PATCH 10/27] chore: project borg from Mono 7bba3d56ad22 --- .repository-projection.json | 6 ++-- docs/operations.md | 12 +++++++ src/classes.rs | 72 +++++++++++++++++++++++++++++++++---- 3 files changed, 81 insertions(+), 9 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 72e1b00..54d60ac 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "5af6e0de217fed962c7f8c796a6f2882c8df6661", + "sourceSha": "7bba3d56ad227703bfe109c3446c993a2633d45e", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "34e8c83d71a6e93badf1fee8996cc0e616da340d", + "priorProjectedBase": "b1d620c02e8f102a826b1bcbd09477c5068b0398", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "35a2309bedf9953a929deb366f5a403f12c75b5dededb129a2377bfbd9a6817e", + "contentDigest": "86cdcfd616bff2c5a72c75fbe4cb88b9e63fe3e16a93806082e73b71204399a4", "publicationEligible": true } diff --git a/docs/operations.md b/docs/operations.md index 8afb081..9dfbf8c 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -51,6 +51,18 @@ label_prefix = "example-ci-" heavy = { mem_mib = 20480, vcpu = 6, docker = true } ``` +Existing `runs-on` labels outside the prefix can be answered too with +`aliases`, a map from label to class: + +```toml +[classes.aliases] +"example-deploy-lane" = "medium" +``` + +Queued jobs with an alias label count as main demand for that class, and every +runner minted for the class also carries the alias labels. Alias labels must +not start with the prefix. + ## Config files Queen and drone configs are TOML or JSON with the same keys, defaults and diff --git a/src/classes.rs b/src/classes.rs index a65824f..0f30e36 100644 --- a/src/classes.rs +++ b/src/classes.rs @@ -4,6 +4,7 @@ //! and merge-group jobs, `-pr` for pull requests). The prefix //! and the per-class requirements are configuration with generic defaults. +use std::collections::BTreeMap; use std::fmt; use serde::{Deserialize, Serialize}; @@ -75,12 +76,18 @@ pub struct Classes { pub light: ClassSpec, pub medium: ClassSpec, pub heavy: ClassSpec, + /// Existing runner labels that borg also answers, each mapped to the class + /// that serves it (for example a deploy lane's `runs-on` label). Queued + /// jobs with an alias label count as main demand for the class, and every + /// runner of that class carries the alias labels so it can take them. + pub aliases: BTreeMap, } impl Default for Classes { fn default() -> Self { Self { label_prefix: DEFAULT_LABEL_PREFIX.to_owned(), + aliases: BTreeMap::new(), // 2.5 GiB p90, 2 vCPU. light: ClassSpec { mem_mib: 2560, @@ -110,6 +117,8 @@ pub enum ClassesError { Parse(#[from] serde_json::Error), #[error("label prefix must not be empty")] EmptyPrefix, + #[error("alias label {0:?} is malformed or starts with the class label prefix")] + BadAlias(String), #[error("class {class} must reserve at least 1 MiB and 1 vCPU")] Empty { class: Class }, #[error("class {larger} must not reserve less than {smaller}")] @@ -138,6 +147,11 @@ impl Classes { if self.label_prefix.is_empty() { return Err(ClassesError::EmptyPrefix); } + for label in self.aliases.keys() { + if !valid_extra_label(label) || label.starts_with(self.label_prefix.as_str()) { + return Err(ClassesError::BadAlias(label.clone())); + } + } for class in Class::ALL { let spec = self.spec(class); if spec.mem_mib == 0 || spec.vcpu == 0 { @@ -193,6 +207,9 @@ impl Classes { /// Parse a job's `runs-on` label into its class and scope. pub fn parse_label(&self, label: &str) -> Option<(Class, Scope)> { + if let Some(&class) = self.aliases.get(label) { + return Some((class, Scope::Main)); + } let rest = label.strip_prefix(self.label_prefix.as_str())?; let (name, scope) = match rest.strip_suffix(PR_SUFFIX) { Some(name) => (name, Scope::Pr), @@ -207,13 +224,17 @@ impl Classes { /// Main grants carry only the bare label so GitHub can hand them nothing /// but main jobs. Pull-request grants carry both labels. pub fn runner_labels(&self, class: Class, scope: Scope) -> Vec { - match scope { - Scope::Main => vec![self.job_label(class, Scope::Main)], - Scope::Pr => vec![ - self.job_label(class, Scope::Main), - self.job_label(class, Scope::Pr), - ], + let mut labels = vec![self.job_label(class, Scope::Main)]; + if scope == Scope::Pr { + labels.push(self.job_label(class, Scope::Pr)); } + labels.extend( + self.aliases + .iter() + .filter(|&(_, &c)| c == class) + .map(|(l, _)| l.clone()), + ); + labels } } @@ -381,6 +402,45 @@ mod tests { ); } + #[test] + fn aliases_are_main_demand_and_ride_on_class_runners() { + let c = Classes::from_json( + r#"{"label_prefix":"x-ci-","aliases":{"deploy-lane":"medium","publisher":"heavy"}}"#, + ) + .unwrap(); + assert_eq!( + c.parse_label("deploy-lane"), + Some((Class::Medium, Scope::Main)) + ); + assert_eq!( + c.parse_label("publisher"), + Some((Class::Heavy, Scope::Main)) + ); + assert_eq!( + c.runner_labels(Class::Medium, Scope::Pr), + ["x-ci-medium", "x-ci-medium-pr", "deploy-lane"] + ); + assert_eq!(c.runner_labels(Class::Light, Scope::Main), ["x-ci-light"]); + assert_eq!( + c.labels_scope(&c.runner_labels(Class::Medium, Scope::Pr)), + Some((Class::Medium, Scope::Pr)) + ); + assert!( + c.sanitize_extra_labels(&["deploy-lane".to_owned()]) + .is_empty(), + "an alias cannot be smuggled in as an extra label" + ); + for bad in [ + r#"{"aliases":{"borg-ci-x":"light"}}"#, + r#"{"aliases":{"a b":"light"}}"#, + ] { + assert!( + matches!(Classes::from_json(bad), Err(ClassesError::BadAlias(_))), + "{bad}" + ); + } + } + #[test] fn extra_label_syntax() { assert!(valid_extra_label("on-prem") && valid_extra_label("a.b_c-1")); From 03f7919330b8eb0f54a6a51aa88af8f8ad98ab34 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Mon, 12 Oct 2026 01:00:16 +0000 Subject: [PATCH 11/27] chore: project borg from Mono 459f737d5b9e --- .repository-projection.json | 6 +- docs/coordination.md | 170 ++++++++ docs/operations.md | 305 ++++++++++++++- src/allocator.rs | 375 +++++++++++++++++- src/api.rs | 353 +++++++++++++++-- src/classes.rs | 230 ++++++++++- src/config.rs | 164 +++++++- src/drone/config.rs | 28 ++ src/drone/fsops.rs | 44 ++- src/drone/install.rs | 9 + src/drone/mod.rs | 169 ++++++-- src/drone/runner.rs | 80 +++- src/drone/standby.rs | 184 +++++++++ src/envcontract.rs | 263 ++++++++++++- src/janitor.rs | 295 ++++++++++++-- src/lead.rs | 750 ++++++++++++++++++++++++++++++++++++ src/lib.rs | 3 + src/main.rs | 65 +++- src/preempt.rs | 409 ++++++++++++++++++++ src/probe/mod.rs | 230 +++++++++++ src/queen.rs | 409 ++++++++++++++++++-- src/requeue.rs | 206 ++++++++++ src/run.rs | 4 +- src/store.rs | 446 ++++++++++++++++++++- src/units.rs | 228 ++++++++++- src/units_api.rs | 22 +- tests/contract.rs | 84 ++++ tests/janitor.rs | 129 +++++++ tests/queen.rs | 213 ++++++++++ tests/requeue.rs | 161 ++++++++ tests/run.rs | 1 + tests/support/mod.rs | 114 +++++- tests/units.rs | 220 +++++++++++ 33 files changed, 6147 insertions(+), 222 deletions(-) create mode 100644 docs/coordination.md create mode 100644 src/drone/standby.rs create mode 100644 src/lead.rs create mode 100644 src/preempt.rs create mode 100644 src/requeue.rs create mode 100644 tests/requeue.rs diff --git a/.repository-projection.json b/.repository-projection.json index 54d60ac..8055428 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "7bba3d56ad227703bfe109c3446c993a2633d45e", + "sourceSha": "459f737d5b9e7e476f6c79467dab5dba31868bd4", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "b1d620c02e8f102a826b1bcbd09477c5068b0398", + "priorProjectedBase": "3d4d719a3dabdd4da94fc7343562e008ce8891e0", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "86cdcfd616bff2c5a72c75fbe4cb88b9e63fe3e16a93806082e73b71204399a4", + "contentDigest": "fb334a68e36619eccf3003521e93a792ad2826c573ebfd1e95d3b96d0e951571", "publicationEligible": true } diff --git a/docs/coordination.md b/docs/coordination.md new file mode 100644 index 0000000..bc185bb --- /dev/null +++ b/docs/coordination.md @@ -0,0 +1,170 @@ +# Coordination: least running time inside one fleet + +How borg and a merge queue should divide work so the fleet spends the least +total machine time per landed change, and so the deploy path never waits +behind pull-request work. This document is the design; the per-feature +mechanics are in the [operations guide](operations.md). + +## The problem it solves + +A monorepo with many agents produces far more CI demand than any fleet can +run at once. Three things then go wrong at the same time: + +1. **Priority inversion.** Jobs that gate a deployment (signing a verdict for + a landed commit, publishing an image, promoting it) queue behind + pull-request jobs that will be superseded by the next push. +2. **Waste.** The same tree is tested several times: on the PR head, on the + merge-queue batch, and again on the landed commit. Runs for stale heads + keep running. Duplicate runs for one commit start from different events. +3. **False capacity.** Slots count memory and vCPUs, so several drones on one + saturated host all report room while the host's CPU pressure is near 100%. + +Mergify's merge queue attacks the second problem on the Git side (batches, +two-step CI, speculative checks, bisection). Borg attacks all three on the +runner side. The two meet at the runner label. + +## Principles + +* **Scopes are the priority ladder.** `deploy` before `main` before `pr` + (`Scope` ordering). Every allocator decision, the starvation guard, the + pressure signal and preemption read the same ladder. +* **Trust follows scope.** A drone is `shared` (main and pull-request work) + or `deploy` (deploy and main work). A deploy drone never runs code that a + pull request produced, and never reads a cache such code wrote. +* **Capacity is CPU pressure, not slots.** Admission reads PSI from each + drone and sheds the lowest scope first. +* **Work that cannot matter is cancelled.** A run for a head that is no + longer the PR's head, or a second run of the same workflow for the same + commit, is stopped before it takes a runner. +* **The queue carries the verdict once.** Heavy suites run once per tested + tree (the batch), and a landed commit whose tree equals a passing batch + tree reuses that verdict instead of running again. + +## Borg mechanisms + +| Mechanism | Where | Status | +| --- | --- | --- | +| Scope ladder `deploy > main > pr`; deploy runners carry only `-deploy` | `classes.rs`, `allocator.rs` | shipped | +| Drone trust (`shared` / `deploy`); aliases with a scope | `classes.rs`, drone config | shipped | +| Admission by CPU pressure (`admission` block; PSI soft sheds `pr`, hard sheds all) | `allocator.rs`, `queen.rs` | shipped | +| Preemption: starving demand cancels one running lower-scope lease (`preemption` block) | `preempt.rs`, drone `cancel_leases` | shipped, including requeue of the cancelled job | +| Supersede janitor: stale heads (queued, and running with `cancel_running`), same-commit duplicates (`dedupe_same_sha`) | `janitor.rs` | shipped | +| Class right-sizing from measured peak RSS (`leases.peak_rss_mib`) | queen config | per fleet | +| Deploy fail-open: standby deploy runner held by deploy drones while the queen is unreachable | queen mints standby; drone starts it on queen loss | shipped | +| Deploy-only trust: a drone that serves nothing but deploy-scope leases | `Trust::DeployOnly` | shipped | +| Queen armada: several queens, one fenced leader | `lead.rs` (new) | design below | + +### Right-sizing classes + +A class reserves memory and vCPUs. Reserve what jobs use, not what the +largest job might use: read `peak_rss_mib` per class from the queen's +`leases` table (p90 and max over the last week) and set `mem_mib` just above +the max, `vcpu` at the p50 of what a job keeps busy. Over-reservation shows +up as "free CPU, no free slots" and under-reservation as `failure_class: +host` with OOM in the drone log. Admission by CPU pressure makes the vCPU +reservation a packing hint rather than a hard budget, so it can be set low. + +### Preemption and requeue + +A lease is only preempted for demand that has waited `after_secs` and that +the allocator could not place. The victim is the lowest-scope, smallest, +most recently started lease on a drone that could then hold the starving +class. The drone sends SIGINT so the Actions runner reports the job as +cancelled, and the queen records (repo, run, job). The requeue step re-runs +that job through the jobs API once its run has completed, with the janitor's +token (it needs `preemption` and `supersede_janitor` both set and respects the +janitor's `dry_run`). + +### Deploy fail-open + +With the queen `standby` block set, each `deploy` or `deploy-only` drone holds +`per_drone` standby leases per class: a JIT runner minted with the `-deploy` +labels that the drone does not start while heartbeats succeed. The queen +starts one through `start_standby` when deploy jobs are queued for its class. +After `fail_open_after_secs` without a queen answer (default 120) the drone +starts all it holds, so a deploy job can still land. A standby is single use: +after fail-open consumes it, no further deploy job runs until the queen +returns, reconciles it (`held_standby`, `active_leases`) and mints another. +Standby leases are exempt from the pending-lease TTL and are never reaped as +unowned. This is the last line of defence; the armada below is the first. +See [operations](operations.md#standby-runners-and-deploy-fail-open). + +### Queen armada (design) + +Several queens, exactly one leader at a time. + +The lease logic (a record of holder, epoch and expiry in an object store +with generation-match preconditions, renewed every `renew_secs` and taken +over after expiry with `epoch + 1`) is in `lead.rs`, with the decision rules +in its module docs and the two-queen, takeover, split-brain and store-outage +cases in its tests. The epoch is the fencing token: `lead.rs` provides the +check a drone applies to a grant (`fence_ok`) and the runner name format +`-e-`. A queen that lost the lease stops +minting at once, and the new leader deletes registrations whose name +carries a lower epoch and that no lease it knows owns. + +* **Only the leader allocates and mints.** Double minting means duplicate + runners and double-run jobs, so followers never call the allocator. + Followers answer `/status` from their last snapshot (marked `follower`), + accept heartbeats and forward them to the leader (or answer `307` with the + leader's URL when the drone can reach it). +* **Drones carry a queen list.** `queen_urls: [...]`; a drone tries the next + URL after `heartbeat_failures` consecutive failures and remembers the one + that answered. Fail-open for deploy still applies when none answers. +* **State is rebuilt, not replicated.** Everything the leader needs is + reconstructible: queued jobs and registered runners from GitHub, drone + capacity and held leases from heartbeats (drones list their active lease + ids and slots). The takeover tick runs the existing reconcile step against + GitHub, adopts registrations that heartbeats claim, deletes the rest, and + only then allocates. SQLite stays per queen; what cannot be rebuilt + (telemetry history, rollout state, health decay) is best-effort and is + mirrored to the store on a timer. +* **Tests.** Two leaders with the same store never both mint (the loser's + CAS fails); a deposed leader's grant with a stale epoch is refused by the + drone and by the store; failover completes within one lease TTL and the + new leader adopts every lease the drones still hold. + +The design keeps the single-authority property the allocator relies on and +adds nothing a follower has to agree on. The fence is the store's generation +check, which is why the store must offer one. + +## Merge-queue side + +Borg does not decide what to run; the repository's CI and merge queue do. +The shape that minimizes running time inside the coordination: + +1. **Two-step CI.** A pull request runs only a light lane on its head; the + heavy graph runs on the merge-queue batch (main plus the batch). The + batch requests main-scope runners, so borg already ranks it above PR work. +2. **Verdict reuse by tree.** A push to the default branch whose git tree + equals a passing batch run's tree runs nothing and signs the batch's + verdict for its commit. The squash of the last PR in a batch has that + tree; the earlier squashes do not. With one pending run per concurrency + group, those intermediate pushes replace each other and only the newest + and the one already running cost anything. +3. **Per-job attestation.** Each heavy job signs its own unit on success, so + a run that is replaced still banks what it finished and the next run tests + only what is left. Without this, a starved default-branch run grows every + time it is replaced. +4. **Parallelism sized to the fleet.** `max_parallel_checks` and per-scope + capacity must not exceed the heavy slots the fleet can give main-scope + work; speculative batches above that only time-share the same CPUs and + finish later, while bisection multiplies the waste on failure. +5. **Priority rules.** Red-main fixes and deploy-path fixes enter at the front + of their scope and never interrupt a running check + (`allow_checks_interruption: false`), because an interrupted heavy batch is + pure loss. +6. **Admin merges disable the queue's economics.** A PR merged outside the + queue while a batch contains it re-cuts the batch, cancels its run and + leaves the landed commit with no batch verdict to reuse. Either changes + land through the queue, or the default-branch run is accepted as the one + heavy run per landed tree and is kept small by 3. + +## Measuring + +Before and after each change, from the Actions API over a fixed window: +job-minutes by event (`pull_request`, batch PR, `push`, `workflow_dispatch`), +job-minutes cancelled after start, duplicate same-commit runs and their +minutes, queue wait p50/p90 for deploy-path labels, job-minutes per landed +commit. From the queen: `/status` `queue` ages per scope, +`borg_drone_admission`, `borg_preemptions_total`, `borg_janitor_*_total`. diff --git a/docs/operations.md b/docs/operations.md index 9dfbf8c..9647432 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -21,10 +21,15 @@ Configuration, scheduling, drone administration, and protocol details for - [Warm-floor controller](#warm-floor-controller) - [Host health](#host-health) - [Tailnet enrollment](#tailnet-enrollment) + - [Cloud burst drones](#cloud-burst-drones) - [Starvation guard](#starvation-guard) + - [Preemption](#preemption) + - [Admission by CPU pressure](#admission-by-cpu-pressure) + - [Standby runners and deploy fail-open](#standby-runners-and-deploy-fail-open) - [Drone](#drone) - [Separate job user (Linux)](#separate-job-user-linux) - [Labels, environment, hooks](#labels-environment-hooks) + - [Drone trust](#drone-trust) - [Owner caps](#owner-caps) - [What a lease does](#what-a-lease-does) - [Host caches](#host-caches) @@ -41,6 +46,15 @@ carries both labels. With this workflow labeling convention, default-branch and merge-queue jobs request the bare label; pull-request jobs request the `-pr` label. Labels route jobs; they do not enforce an event restriction. +Deploy-path jobs (prod OIDC, signing, registry push) ask for +`-deploy`. A runner minted for deploy demand carries only that +label, so GitHub can hand it nothing but deploy jobs. Deploy demand is +scheduled before main demand, and main before pull-request demand. + +A deploy drone never runs a pull-request job or a work unit; a shared drone +never runs a deploy job. Which drones are deploy drones is set per drone with +`trust` (see [Drone trust](#drone-trust)). + ## Classes config The `classes` table is partial; omitted fields keep their defaults: @@ -52,16 +66,21 @@ heavy = { mem_mib = 20480, vcpu = 6, docker = true } ``` Existing `runs-on` labels outside the prefix can be answered too with -`aliases`, a map from label to class: +`aliases`, a map from label to a class name (main scope) or to an object +naming the class and scope: ```toml [classes.aliases] -"example-deploy-lane" = "medium" +"example-legacy-lane" = "medium" +"example-promote-lane" = { class = "medium", scope = "deploy" } ``` -Queued jobs with an alias label count as main demand for that class, and every -runner minted for the class also carries the alias labels. Alias labels must -not start with the prefix. +Queued jobs with an alias label count as demand of the alias's scope, so a +`deploy` alias is deploy demand. A runner carries an alias label only when its +grant scope equals the alias's scope: a deploy alias rides on deploy runners +only, and a main alias rides on main and pull-request runners (both carry the +bare label). Alias labels must not start with the prefix and the scope must be +`main` or `deploy`. ## Config files @@ -116,7 +135,20 @@ through canaries and waves (see [Quorum rollout](#quorum-rollout)). `warm_floors` keep idle runners alive per class and scope; the optional `warm_schedule` raises them in time-of-day windows (`{"class": "light", "scope": "pr", "count": 4, "start": "08:00", "end": "20:00", "days": ["mon", "tue"]}`) -at a fixed `utc_offset_minutes`. A window never lowers a floor. +at a fixed `utc_offset_minutes`. A window never lowers a floor. Floors and +windows may name `"scope": "deploy"`; the allocator keeps those runners on +deploy drones. + +Pending runners absorb queued demand of the same label: a pending +pull-request runner (it carries the bare label) may absorb main demand, but +deploy demand is absorbed only by pending deploy runners, and a pending deploy +runner absorbs nothing else. + +### Leader lease + +The fenced leader lease for running several queens (`src/lead.rs`) exists as a +tested module and is not wired into the queen yet. There is no config key, no +store setting and no flag for it; run a single queen until the wiring ships. ### Job telemetry @@ -244,7 +276,7 @@ Targets are IPv4 addresses or `addr/len` networks. See `/status` `wake`. Optional (`env_contract.enabled`, off by default). Before a CI runner lease or a work unit is placed on a drone, the queen checks what the drone last reported -about its build environment. Four checks, each `required`, `advisory` or `off` +about its build environment. Five checks, each `required`, `advisory` or `off` globally and per class (`env_contract.checks`, `env_contract.classes.`): * `compiler_cache`: `RUSTC_WRAPPER` is an allowed wrapper (`allowed_wrappers`, @@ -262,13 +294,30 @@ globally and per class (`env_contract.checks`, `env_contract.classes.`): * `resources`: memory available covers the class's memory (`mem_available_pct`) and CPU pressure is at most `max_cpu_psi_pm`. (Free capacity after active leases is the allocator's own rule.) +* `tools`: every tool in `required_tools` is on the job user's `PATH`. A required check that fails **or is unknown** makes the drone ineligible for the class; nothing else changes about placement. Facts older than `facts_max_age_secs` (180) count as unknown. A drone that has never reported is skipped unless `require_report`, so enabling the contract does not cut off drones that predate it. Defaults when enabled: `toolchain` and `resources` -required, `compiler_cache` and `cache_store` advisory. +required, `compiler_cache`, `cache_store` and `tools` advisory. + +The `tools` check covers what a job expects to find on `PATH`, such as `gh` or +`node`. `env_contract.required_tools` lists the tools every class needs and +`env_contract.classes..required_tools` replaces that list for one class. +Both lists default to empty, so the check passes until the operator names tools; +names must come from the drone's probe list (`git`, `gh`, `node`, `npm`, `pnpm`, +`python3`, `docker`, `cargo`, `rustc`, `bazel`, `go`, `uv`, `mise`, `jq`, +`cosign`). A drone whose facts show a required tool as `missing` fails the check, +and the reason names the missing tools. A drone whose facts carry no `tools` map +(an older binary) passes with a warning, and fails when `require_report` is set. +`/status` shows the missing tools and every probed version under the check's +observed values. + +```json +"env_contract": { "enabled": true, "checks": { "tools": "required" }, "required_tools": ["git", "gh", "node", "docker"] } +``` `/status` shows each drone's `contract` (the checks with state, reason, observed values and warnings, and the admission per class with the reason) and @@ -296,6 +345,7 @@ old as unknown, so probe at least that often): | `toolchains` | For each toolchain `rustup toolchain list` shows, `{channel, rustc_version}`: the channel as `rust-toolchain.toml` would name it (`1.95.0`, `stable`, `nightly-2026-01-01`; strip the host triple) and the `release:` line of `rustc + -vV`. `null` if rustup could not be queried. | | `mem_available_mib` | Linux `MemAvailable` from `/proc/meminfo`; macOS (free + inactive + speculative pages) times the page size, in MiB. | | `cpu_psi_some_avg10_pm` | Linux `/proc/pressure/cpu` line `some`, `avg10` times ten, rounded; absent where PSI does not exist. | +| `tools` | For each of `git`, `gh`, `node`, `npm`, `pnpm`, `python3`, `docker`, `cargo`, `rustc`, `bazel`, `go`, `uv`, `mise`, `jq`, `cosign`: resolve the executable on the job user's `PATH` (as the job user, with the job environment) and record the first line of ` --version` (`go version` for `go`), trimmed, at most 80 characters. `"missing"` when it is not on `PATH`, `"unknown"` when it is but prints nothing. The map is absent when the probe did not run. | A unit's toolchain (`rust_toolchain`) matches an installed `channel` or `rustc_version` exactly. @@ -340,13 +390,31 @@ with `borg queen enroll-client ` (prints a token once). Client API (bearer: client token): -* `POST /v1/units` with `{repo, tree_sha, argv, platform, toolchain, class, timeout_secs?, force?}` +* `POST /v1/units` with `{repo, tree_sha, argv, platform, toolchain, class, timeout_secs?, force?, scope?, supersede_key?, record?}` returns `201` `{outcome: "new", unit_id, ...}`, or `200` with `outcome: "joined"` (an identical unit is queued or running) or `"receipt"` (an identical unit succeeded recently; `receipt` has the proof). The key is (repo, tree SHA, argv, platform, toolchain); only successes are reused, never for another tree, and `force` skips reuse. `platform` is `-` (`linux-x86_64`, `macos-aarch64`). +* `scope` is `deploy`, `main` or `pr` (default `pr`). The unit holds capacity + under that scope and runs only on drones whose trust admits it (see + [Drone trust](#drone-trust)): a `deploy` unit runs only on a deploy drone, a + `pr` unit never does, and a `main` unit runs on either. Scope is part of the + single-flight key, so a `main` or `deploy` unit never joins or reuses the + result of a lower-scope run. Enroll a client with `borg queen enroll-client + --max-scope deploy|main|pr` (default `pr`) to set the highest scope + its token may submit; a submission above that is refused with `403` and a + message naming both scopes. `borg run --scope main` (or `BORG_RUN_SCOPE`) + sets the scope on the client side. +* `supersede_key` (at most 128 characters) makes the newest submission win: + when a new unit is created with a key, every queued unit from the same client + with the same key is cancelled, and every running one is cancelled through + the `cancel_units` heartbeat entry. A cancelled unit ends in state `failed` + with `failure` set to `superseded by `. A unit submitted with + `"record": true` (default `false`) is never superseded. The response lists + the cancelled ids in `superseded`; `/status` `units.superseded` and the + `borg_units_superseded_total` metric count them since the queen started. * `GET /v1/units/{id}` returns the unit (state `queued|leased|running|done|failed`, `exit_code`, `receipt`). * `GET /v1/units/{id}/log?from=N&wait=S` returns log bytes from offset `N`, @@ -385,7 +453,10 @@ Drone protocol (bearer: drone token): Units are placed by the runner allocator (class capacity, owner caps, bids, profiles, quarantine and probation) on drones whose `-` matches, using -capacity left after CI runner demand. +capacity left after CI runner demand. Among queued units the order is the one +runner leases use: `deploy` before `main` before `pr`, older first within a +scope, and `pr` units are the first held back when a drone is under CPU +pressure (see [Admission by CPU pressure](#admission-by-cpu-pressure)). ### Build-box reclaim @@ -465,7 +536,26 @@ The optional supersede janitor cancels queued pull-request workflow runs whose head commit is no longer the pull request's head. Add a `supersede_janitor` block (`token_path` or `BORG_JANITOR_TOKEN` for a token with `actions: write`, `dry_run` defaulting to true, `interval_secs`, `max_cancels_per_pass`) to turn -it on. It never cancels a run that has started. +it on. By default it never cancels a run that has started. + +Two further switches, both off by default and both subject to `dry_run` and +`max_cancels_per_pass`, remove other waste. `cancel_running = true` also +cancels superseded pull-request runs that are already in progress; a run that +finished between the listing and the cancel is counted as `raced`. +`dedupe_same_sha = true` looks at the queued and in-progress runs of each +repository and groups them by workflow, head commit, event and pull request +number. In a group of more than one it keeps the in-progress run (the lowest id +if several are running), or else the queued run with the lowest id, and cancels +the others only if they are still queued; a running duplicate is never +cancelled. `workflow_dispatch` and `merge_group` runs are never grouped, +because publication waits on those exact runs. Each cancellation is logged with +the id of the run that was kept. With `dry_run` the log and `would_cancel` show +what would have been cancelled. When the janitor is enabled, `GET /status` has +a `janitor` object (`passes`, `last`, `total`, each with the counters +`considered`, `superseded`, `cancelled`, `cancelled_running`, `duplicates`, +`dedup_cancelled`, `would_cancel`, `deferred`, `raced`, `errors`) and +`GET /metrics` has `borg_janitor__total`, summed over passes since the +queen started. The GitHub App needs organization self-hosted runners: read and write, plus read access to Actions on the listed repositories. Set @@ -486,6 +576,66 @@ Manage drones with `borg queen drone list` and `borg queen drone remove ` Drone requests carry `Authorization: Bearer `. The queen stores only the SHA-256 of each token. +### Cloud burst drones + +A burst drone is a short-lived cloud VM that joins the fleet to drain a queue +and then deletes itself. Two things decide whether it is cheap and safe: how +its traffic leaves the cloud network, and how it reaches the queen. + +A burst image must carry the same tools as the on-prem runners (`gh`, `node` +and the rest of the `required_tools` list). The `tools` check is how the queen +refuses a drone whose image does not. + +**Egress: a public IP per VM, no NAT gateway.** Give each VM its own +ephemeral external IPv4 address. Traffic to GitHub, crate registries and +container registries then leaves the VM directly. Do not put burst drones +behind a managed NAT gateway (for example GCP Cloud NAT): the gateway bills +per gateway-hour and per gigabyte processed, on top of normal egress, and a +CI fleet pulls tens of gigabytes per host per day. Inbound transfer into the +cloud is free; you pay normal egress only for what the job uploads +(artifacts, caches, images). + +**Reaching the queen: Tailscale, not NAT.** The queen usually lives on a +private network behind a home or office router. The burst drone joins the +tailnet and reaches the queen over it, so the router needs no port forward +and the cloud side needs no VPN gateway. Tailscale does not change how +traffic to the internet is billed; the public IP above is what avoids NAT +charges. Tailscale only carries the drone-to-queen connection. + +**Containment.** Burst drones run pull-request code, so treat them as +untrusted: + +- Join with an ephemeral, pre-approved, tagged auth key (for example + `tag:borg-drone`). In the tailnet policy, allow that tag to reach only the + queen's port. Nothing else on the tailnet should be reachable from it. +- Enroll through [tailnet enrollment](#tailnet-enrollment) with that tag in + `allowed_tags`, so no token is copied onto the VM by hand. +- Create the VM without cloud credentials: no service account and no API + scopes, so a job cannot read the metadata server for a cloud token. +- Put the VM on a dedicated network with no inbound firewall rules. Manage it + through the cloud's identity-aware tunnel or Tailscale SSH. +- Do not serve deploy-trust work (image publishing, promotion, anything that + signs or holds production credentials) from a burst drone that also runs + pull-request jobs. + +**Lifetime.** Use spot or preemptible capacity with a hard maximum run +duration and delete-on-termination, so a forgotten drone costs at most one +run duration. On GCP: + +```sh +gcloud compute instances create burst-drone-1 \ + --machine-type=c2d-highcpu-56 --provisioning-model=SPOT \ + --instance-termination-action=DELETE --max-run-duration=12h \ + --no-service-account --no-scopes \ + --network=borg-burst --subnet=borg-burst \ + --image-family=ubuntu-2404-lts-amd64 --image-project=ubuntu-os-cloud +``` + +The VM keeps its ephemeral external IP by default; do not pass `--no-address` +unless a NAT gateway exists, which this setup avoids. If the cloud reclaims a +spot VM mid-job, GitHub fails that job when its runner stops responding; rerun +it. Queued jobs are unaffected and go to other drones. + ### Starvation guard When demand for a class has waited longer than `starvation_after_secs`, the @@ -494,6 +644,122 @@ that can free up soonest (fewest in-flight leases, ties by name) and grants no smaller lease there until the starving lease fits. Held drones are listed in `/status` under `reserved_for`. +### Preemption + +Off unless the queen config has a `preemption` block: + +```json +"preemption": { "after_secs": 600, "max_per_pass": 1, "min_victim_age_secs": 60 } +``` + +`after_secs` is how long the oldest queued job of a `(class, scope)` must have +waited, `max_per_pass` caps cancellations in one queen pass, and +`min_victim_age_secs` is how long a job must have been running before it can be +cancelled. All three must be greater than 0. + +When demand the allocator could not place has waited `after_secs`, the queen +cancels one running lease whose scope ranks strictly below the starving +demand's scope (a `pr` lease for `main` demand; the order is the `Scope` +order). It only picks a lease whose drone could hold the starving class once +that lease's resources are freed, and only a lease whose in-progress job GitHub +reports on its runner, since the job's start time sets its age. Among the +candidates it takes the lowest-priority scope, then the smallest class, then the +most recently started lease. A lease that is already being cancelled counts as +covering the starving demand, so one starving job cancels one lease, not one per +pass. + +The queen stamps `preempted_at` on the lease, records the job (repo, run id, +job id) in the `preemptions` table, and returns the lease id in `cancel_leases` +on every heartbeat response to its drone until the lease ends. The drone sends +SIGINT to the runner's process group, which makes the Actions runner report the +job as cancelled to GitHub, waits `preempt_grace_secs` (drone config, default +30, must be greater than 0), then kills the group as it does at `max_job_secs`. +The result is reported with `preempted: true`, and the lease ends normally. +Drones that predate `cancel_leases` ignore it, so update drones before turning +preemption on. + +`/status` lists the last 20 preemptions under `preemptions` (lease, drone, +class, victim `scope`, starving `for_scope`, repo, run, job, `at`, +`requeued_at`) and the all-time counts under `preemption_counts`. The metric +`borg_preemptions_total{from_scope,to_scope}` counts them by victim scope and +starving scope. + +A preempted job is re-run automatically when both `preemption` and +`supersede_janitor` are configured. After each janitor pass the queen reads the +runs of preemptions with `requeued_at` unset and, once a run has completed, +calls `POST /repos/{repo}/actions/jobs/{job_id}/rerun` with the janitor's token +and sets `requeued_at`. It does at most `max_cancels_per_pass` reruns per pass, +marks rows older than 24 hours without re-running them, treats a 403, 404 or 409 +answer as final, retries other errors on the next pass, and in the janitor's +`dry_run` only logs `would rerun`. `/status` shows the counters under `requeue` +and `/metrics` as `borg_requeue__total` (`checked`, `rerun`, `waiting`, +`stale`, `terminal`, `errors`). + +### Admission by CPU pressure + +Slots count memory, vCPUs and leases, not CPU load, so several drones sharing +one saturated host can all show free slots. Optional and off unless the queen +config has an `admission` block: + +```json +"admission": { "psi_soft_pm": 300, "psi_hard_pm": 600, "facts_max_age_secs": 180 } +``` + +Each drone reports Linux PSI `cpu some avg10` in permille on every heartbeat +(the higher of `/proc/pressure/cpu` and the cgroup file +`/sys/fs/cgroup/cpu.pressure`). At or above `psi_soft_pm` the drone takes no +new pull-request grant; main work is still placed. At or above `psi_hard_pm` it +takes no new grant of any scope. Leases already granted are not touched. A +reading older than `facts_max_age_secs`, or no reading, puts no restriction on +the drone. The starvation guard obeys the same rule: a starving pull-request +lease is not placed on a soft-blocked drone. Requires `0 < psi_soft_pm <= +psi_hard_pm <= 1000` and a positive age. + +`/status` shows `admission: {state, cpu_psi_pm}` per drone (`state` is `open`, +`pr_blocked` or `closed`). +Metrics: `borg_drone_cpu_psi_pm{drone}` and `borg_drone_admission{drone,state}`. +The build environment contract's `resources` check is separate and still works. + +### Standby runners and deploy fail-open + +Deploy jobs (image promotion) should not depend on one queen host. Optional and +off unless the queen config has a `standby` block: + +```json +"standby": { "classes": ["medium"], "per_drone": 1 } +``` + +`classes` must not be empty and `per_drone` is 1 to 4. For every fresh drone +with trust `deploy` or `deploy-only`, the queen keeps `per_drone` leases of +each listed class in the state `standby`: a JIT runner minted with the deploy +labels only (plus the drone's extra labels). The queen sends the JIT config in +the heartbeat response (`standby`) and resends it until the drone lists the +lease id in `held_standby`. The drone does not start it. A standby holds +capacity like a pending deploy lease, is exempt from `lease_ttl_secs`, and is +rotated only when `lease_max_secs` passes; its registration is never reaped. + +Promotion rule: when deploy jobs are queued for a class that pending deploy +runners do not cover, the queen does not mint another runner while a standby +exists. It lists the standby in `start_standby` and the drone starts it as an +ordinary lease. When the drone reports the lease in `active_leases` the queen +marks it `pending`, then `running` once GitHub shows the runner busy; the next +tick mints a replacement standby. A standby that ended on the queen's side is +listed in `drop_standby` and the drone forgets it. + +Fail-open rule: the drone config key `fail_open_after_secs` (default 120, `0` +disables it; `borg drone install --fail-open-after-secs N`) is how long the +queen may go without answering a heartbeat. After that the drone starts all +the standbys it holds, with no instruction from the queen, and keeps retrying +heartbeats. The drone keeps standbys in memory and, when `state_dir` is set, in +`/standby/.json` (mode 0600), so a restart while the queen +is down still has them; the file is deleted when the standby starts. When the +queen returns, `held_standby` reports what is unstarted and `active_leases` +what runs, and the queen reconciles. + +A standby is a single-use ephemeral runner, so after fail-open consumes it no +further deploy job runs until the queen returns and mints another. Pair it with +`--trust deploy-only` for a dedicated promotion box. + ## Drone ```sh @@ -551,6 +817,19 @@ jobs run as the drone's own account and the drone logs a warning. * `borg drone run --token-file F` and `borg drone install --token-file F` read the token from a file instead of the command line. +### Drone trust + +`trust` in the drone config is `"shared"` (default), `"deploy"` or +`"deploy-only"`. A `deploy-only` drone serves deploy-scope leases and nothing +else: no main, no pull-request work, no work units. The drone +reports it in every heartbeat and `/status` shows it per drone. The queen +places deploy demand only on `deploy` drones, never places pull-request demand +or work units on them, and may place main demand on either. Set it on the +host, for example `"trust": "deploy"` in `drone.json` or +`borg drone install --trust deploy`; nothing is set on the queen. A deploy +drone should be a host that no pull-request job has touched, since the +invariant only covers what the queen schedules. + ### Owner caps `--max-vcpu`, `--max-mem-mib` and `--max-leases` (config keys of the same @@ -566,7 +845,9 @@ budget, so an owner can say "this PC gives at most 8 vCPUs and 24 GiB". environment. Jobs see `BORG_LEASE_ID`, `BORG_DOCKER_PREFIX`, `COMPOSE_PROJECT_NAME` (both `borg-`) and `BORG_CACHE_ROOT`. 2. When the runner exits, or `max_job_secs` passes, the process group is - killed, and the work directory (including `_diag`) is deleted. + killed, and the work directory (including `_diag`) is deleted. If the queen + preempts the lease (see [Preemption](#preemption)), the group gets SIGINT + first and is killed after `preempt_grace_secs` (default 30). 3. Docker containers, networks and volumes are removed only if `docker inspect` shows they carry the label `borg.lease=`, the compose project `borg-`, or a name that is exactly that prefix or starts diff --git a/src/allocator.rs b/src/allocator.rs index 88f8dff..e80ce86 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -8,11 +8,13 @@ //! //! Rules, in priority order: //! -//! 1. Main demand (bare label) is satisfied before pull-request demand, and -//! within a scope larger classes are placed first so they are not starved -//! of contiguous room by small ones. +//! 1. Deploy demand is satisfied before main demand (bare label), and main +//! before pull-request demand. Within a scope larger classes are placed +//! first so they are not starved of contiguous room by small ones. //! 2. A main grant carries only the bare label; a pull-request grant carries -//! both labels (see [`Classes::runner_labels`]). +//! both labels; a deploy grant carries only the `-deploy` label (see +//! [`Classes::runner_labels`]). Deploy demand is placed only on +//! [`Trust::Deploy`] drones and pull-request demand never is. //! 3. Queued demand is first offset by leases that were granted but have not //! yet taken a job ([`LeaseState::Pending`]); only the shortfall is minted. //! 4. A warm floor keeps that many pending leases of a (class, scope) alive @@ -29,7 +31,7 @@ use serde::{Deserialize, Serialize}; use crate::bid::{ ClassPressure, DroneBid, ProfileScore, ScoreWeights, demand_pressure, profile_term, score, }; -use crate::classes::{Class, ClassSpec, Classes, Scope}; +use crate::classes::{Class, ClassSpec, Classes, Scope, Trust}; use crate::probe::Capacity; /// Owner-set limits for one host ("this PC gives at most 8 cores / 24 GiB"). @@ -76,6 +78,68 @@ pub struct DroneState { /// right now; see [`crate::envcontract`]. The allocator never places them. #[serde(default)] pub blocked_classes: Vec, + /// Which scopes the host may run; see [`Trust::admits`]. + #[serde(default)] + pub trust: Trust, + /// Recent CPU pressure (`some avg10`, permille); `None` is not measured + /// or stale. Only read when [`AllocatorInput::admission`] is set. + #[serde(default)] + pub cpu_psi_pm: Option, +} + +/// CPU-pressure thresholds (permille) for admission; see [`AdmissionState`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdmissionThresholds { + pub psi_soft_pm: u16, + pub psi_hard_pm: u16, +} + +/// What a drone may be granted given its CPU pressure. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AdmissionState { + /// Below the soft threshold, unmeasured, or admission is off. + Open, + /// At or above soft: pull-request grants are shed, everything else placed. + PrBlocked, + /// At or above hard: no new grant of any scope. + Closed, +} + +impl AdmissionState { + pub const ALL: [AdmissionState; 3] = [Self::Open, Self::PrBlocked, Self::Closed]; + + pub fn as_str(self) -> &'static str { + match self { + Self::Open => "open", + Self::PrBlocked => "pr_blocked", + Self::Closed => "closed", + } + } + + pub fn of(psi: Option, t: Option) -> Self { + match (psi, t) { + (Some(p), Some(t)) if p >= t.psi_hard_pm => Self::Closed, + (Some(p), Some(t)) if p >= t.psi_soft_pm => Self::PrBlocked, + _ => Self::Open, + } + } + + fn admits(self, scope: Scope) -> bool { + match self { + Self::Open => true, + Self::PrBlocked => scope != Scope::Pr, + Self::Closed => false, + } + } +} + +/// One drone's admission state this pass. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct DroneAdmission { + pub drone_id: String, + pub state: AdmissionState, + pub cpu_psi_pm: Option, } impl DroneState { @@ -94,6 +158,8 @@ impl DroneState { bids: Vec::new(), profiles: Vec::new(), blocked_classes: Vec::new(), + cpu_psi_pm: None, + trust: Trust::Shared, } } } @@ -102,6 +168,10 @@ impl DroneState { #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum LeaseState { + /// A deploy runner held unstarted by its drone. It holds capacity like a + /// pending deploy lease and absorbs deploy demand by promotion: the drone + /// is told to start it instead of the queen minting another runner. + Standby, /// Granted or minted, runner idle: can still absorb queued demand. Pending, /// Runner is executing a job: only consumes capacity. @@ -164,6 +234,8 @@ pub struct AllocatorInput { /// the guard off. pub starvation_after_secs: Option, pub weights: ScoreWeights, + /// CPU-pressure admission; `None` turns it off. + pub admission: Option, } /// An instruction to mint one runner and hand it to one drone. @@ -204,6 +276,18 @@ pub struct Allocation { pub pressure: Vec, /// Slots each class could still be granted before this pass's grants. pub idle: Vec, + /// Every drone's admission state this pass, by id. + pub admission: Vec, + /// Standby deploy runners to start now, per class: queued deploy demand + /// that pending runners do not cover, up to the standbys that exist. + pub promote: Vec, +} + +/// Start `count` standby deploy runners of `class`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct Promotion { + pub class: Class, + pub count: u32, } /// Idle room for one class. @@ -267,6 +351,7 @@ struct Free<'a> { vcpu: u32, leases_left: Option, in_flight: u32, + admission: AdmissionState, /// Grants made to this drone this pass, per key (bounds `available_slots`). granted: BTreeMap<(Class, Scope), u32>, } @@ -302,6 +387,7 @@ impl Free<'_> { fn fits(&self, spec: &ClassSpec) -> bool { self.drone.accepting + && self.admission != AdmissionState::Closed && (!spec.docker || self.drone.docker) && self.mem_mib >= spec.mem_mib && self.vcpu >= spec.vcpu @@ -319,6 +405,7 @@ impl Free<'_> { fn could_ever_fit(&self, spec: &ClassSpec) -> bool { self.drone.accepting + && self.admission != AdmissionState::Closed && (!spec.docker || self.drone.docker) && self.drone.mem_mib >= spec.mem_mib && self.drone.vcpu >= spec.vcpu @@ -364,6 +451,8 @@ fn place( .iter_mut() .filter(|(id, f)| { f.fits(spec) + && f.admission.admits(scope) + && f.drone.trust.admits(scope) && !f.drone.blocked_classes.contains(&class) && reserved .get(**id) @@ -403,6 +492,7 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { vcpu: drone.vcpu, leases_left: drone.max_leases, in_flight: 0, + admission: AdmissionState::of(drone.cpu_psi_pm, input.admission), granted: BTreeMap::new(), }); } @@ -427,11 +517,38 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { .filter(|l| l.state == LeaseState::Pending && free.contains_key(l.drone_id.as_str())) .map(|l| (l.class, l.scope, 1)), ); + let standby = sum_counts( + input + .leases + .iter() + .filter(|l| l.state == LeaseState::Standby && free.contains_key(l.drone_id.as_str())) + .map(|l| (l.class, l.scope, 1)), + ); let get = |m: &Counts, c, s| m.get(&(c, s)).copied().unwrap_or(0); + let mut promote: Vec = Vec::new(); // How many runners each (scope, class) still has to mint. let mut need: BTreeMap<(Scope, Class), u32> = BTreeMap::new(); for class in Class::ALL { + // Deploy demand is absorbed only by pending deploy runners, which in + // turn absorb nothing else. + let deploy_total = + get(&queued, class, Scope::Deploy).max(get(&floors, class, Scope::Deploy)); + let pending_deploy = get(&pending, class, Scope::Deploy); + // Queued deploy jobs that pending runners do not cover start standby + // runners before the queen mints anything. + let standby_deploy = get(&standby, class, Scope::Deploy); + let promoted = get(&queued, class, Scope::Deploy) + .saturating_sub(pending_deploy) + .min(standby_deploy); + if promoted > 0 { + promote.push(Promotion { + class, + count: promoted, + }); + } + let deploy_need = deploy_total.saturating_sub(pending_deploy + standby_deploy); + let pending_main = get(&pending, class, Scope::Main); let pending_pr = get(&pending, class, Scope::Pr); @@ -446,6 +563,7 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { let pr_total = get(&queued, class, Scope::Pr).max(get(&floors, class, Scope::Pr)); let pr_need = pr_total.saturating_sub(pending_pr_left); + need.insert((Scope::Deploy, class), deploy_need); need.insert((Scope::Main, class), main_need); need.insert((Scope::Pr, class), pr_need); } @@ -487,6 +605,15 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { .map(|(&class, &pressure)| ClassPressure { class, pressure }) .collect(), idle: idle_by_class, + admission: free + .iter() + .map(|(id, f)| DroneAdmission { + drone_id: (*id).to_owned(), + state: f.admission, + cpu_psi_pm: f.drone.cpu_psi_pm, + }) + .collect(), + promote, ..Allocation::default() }; let mut placed: BTreeMap<(Scope, Class), u32> = BTreeMap::new(); @@ -519,7 +646,9 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { .iter() .filter(|(id, f)| { !reserved.contains_key(**id) + && f.drone.trust.admits(scope) && f.could_ever_fit(spec) + && f.admission.admits(scope) && !f.drone.blocked_classes.contains(&class) }) .min_by_key(|(id, f)| (f.in_flight, **id)) @@ -587,6 +716,8 @@ mod tests { bids: vec![], profiles: vec![], blocked_classes: vec![], + cpu_psi_pm: None, + trust: Trust::Shared, } } @@ -607,6 +738,7 @@ mod tests { queue_ages: vec![], starvation_after_secs: None, weights: ScoreWeights::default(), + admission: None, } } @@ -1206,6 +1338,8 @@ mod tests { bids: vec![], profiles: vec![], blocked_classes: vec![], + cpu_psi_pm: None, + trust: Trust::Shared, }, DroneState { id: "c".into(), @@ -1217,6 +1351,8 @@ mod tests { bids: vec![], profiles: vec![], blocked_classes: vec![], + cpu_psi_pm: None, + trust: Trust::Shared, }, ]; for lm in 0..3 { @@ -1489,4 +1625,233 @@ mod tests { assert_eq!(only_drone(&allocate(&i)), "b"); assert_eq!(allocate(&i), allocate(&i)); } + + fn deploy_host(id: &str) -> DroneState { + DroneState { + trust: Trust::Deploy, + ..host(id) + } + } + + #[test] + fn deploy_demand_lands_only_on_deploy_drones_with_only_the_deploy_label() { + let a = allocate(&input( + &[(Class::Medium, Scope::Deploy, 2)], + vec![host("shared"), deploy_host("dep")], + )); + assert_eq!(count(&a, Class::Medium, Scope::Deploy), 2); + assert!(a.grants.iter().all(|g| g.drone_id == "dep")); + assert!( + a.grants + .iter() + .all(|g| g.labels == ["borg-ci-medium-deploy"]) + ); + + let a = allocate(&input( + &[(Class::Medium, Scope::Deploy, 1)], + vec![host("shared")], + )); + assert!(a.grants.is_empty()); + assert_eq!(a.unmet.len(), 1); + assert_eq!(a.unmet[0].scope, Scope::Deploy); + } + + #[test] + fn pr_demand_never_lands_on_deploy_drones_but_main_may() { + let a = allocate(&input( + &[(Class::Light, Scope::Pr, 3), (Class::Light, Scope::Main, 1)], + vec![deploy_host("dep")], + )); + assert_eq!(count(&a, Class::Light, Scope::Pr), 0); + assert_eq!(count(&a, Class::Light, Scope::Main), 1); + assert_eq!(a.unmet[0].scope, Scope::Pr); + } + + #[test] + fn deploy_is_placed_before_main_and_pr_under_contention() { + let mut d = deploy_host("dep"); + d.max_leases = Some(1); + let a = allocate(&input( + &[ + (Class::Light, Scope::Pr, 1), + (Class::Light, Scope::Main, 1), + (Class::Light, Scope::Deploy, 1), + ], + vec![d], + )); + assert_eq!(a.grants.len(), 1); + assert_eq!(a.grants[0].scope, Scope::Deploy); + } + + #[test] + fn starved_deploy_demand_reserves_only_a_deploy_drone() { + let mut i = input( + &[(Class::Heavy, Scope::Deploy, 1)], + vec![host("shared"), deploy_host("dep")], + ); + i.starvation_after_secs = Some(10); + i.queue_ages = vec![QueueAge { + class: Class::Heavy, + scope: Scope::Deploy, + oldest_age_secs: 60, + }]; + i.leases = vec![ + lease("dep", Class::Heavy, Scope::Main, LeaseState::Running), + lease("dep", Class::Light, Scope::Main, LeaseState::Running), + ]; + let a = allocate(&i); + assert!(a.grants.is_empty()); + assert_eq!(a.reserved_for.len(), 1); + assert_eq!(a.reserved_for[0].drone_id, "dep"); + } + + #[test] + fn pending_absorption_respects_trust_scopes() { + // A pending PR runner (bare label) absorbs main demand, never deploy. + let mut i = input(&[(Class::Light, Scope::Main, 1)], vec![host("a")]); + i.leases = vec![lease("a", Class::Light, Scope::Pr, LeaseState::Pending)]; + assert!(allocate(&i).grants.is_empty()); + + let mut i = input( + &[(Class::Light, Scope::Deploy, 1)], + vec![host("a"), deploy_host("dep")], + ); + i.leases = vec![lease("a", Class::Light, Scope::Pr, LeaseState::Pending)]; + assert_eq!(count(&allocate(&i), Class::Light, Scope::Deploy), 1); + i.leases = vec![lease( + "dep", + Class::Light, + Scope::Deploy, + LeaseState::Pending, + )]; + assert!(allocate(&i).grants.is_empty()); + + // A pending deploy runner carries only the deploy label: main demand + // still needs its own runner. + let mut i = input(&[(Class::Light, Scope::Main, 1)], vec![deploy_host("dep")]); + i.leases = vec![lease( + "dep", + Class::Light, + Scope::Deploy, + LeaseState::Pending, + )]; + assert_eq!(count(&allocate(&i), Class::Light, Scope::Main), 1); + } + + #[test] + fn warm_floors_may_name_the_deploy_scope() { + let mut i = input(&[], vec![deploy_host("dep")]); + i.warm_floors = vec![WarmFloor { + class: Class::Light, + scope: Scope::Deploy, + count: 2, + }]; + assert_eq!(count(&allocate(&i), Class::Light, Scope::Deploy), 2); + } + + fn deploy_only_host(id: &str) -> DroneState { + DroneState { + trust: Trust::DeployOnly, + ..host(id) + } + } + + #[test] + fn deploy_only_drones_take_deploy_demand_and_nothing_else() { + let a = allocate(&input( + &[ + (Class::Medium, Scope::Main, 1), + (Class::Medium, Scope::Pr, 1), + ], + vec![deploy_only_host("only")], + )); + assert!(a.grants.is_empty(), "{:?}", a.grants); + assert_eq!(a.unmet.len(), 2); + let a = allocate(&input( + &[(Class::Medium, Scope::Deploy, 1)], + vec![deploy_only_host("only")], + )); + assert_eq!(count(&a, Class::Medium, Scope::Deploy), 1); + assert_eq!(a.grants[0].drone_id, "only"); + } + + #[test] + fn a_standby_holds_capacity_and_counts_as_pending_deploy_capacity() { + // No queued deploy demand: the standby keeps the deploy floor met and + // nothing is promoted or minted. + let mut i = input(&[], vec![deploy_only_host("dep")]); + i.warm_floors = vec![WarmFloor { + class: Class::Medium, + scope: Scope::Deploy, + count: 1, + }]; + i.leases = vec![lease( + "dep", + Class::Medium, + Scope::Deploy, + LeaseState::Standby, + )]; + let a = allocate(&i); + assert!(a.grants.is_empty() && a.promote.is_empty()); + // It also takes room: a 1-lease cap is full. + let mut capped = deploy_only_host("dep"); + capped.max_leases = Some(1); + let mut i = input(&[(Class::Light, Scope::Deploy, 1)], vec![capped]); + i.leases = vec![lease( + "dep", + Class::Medium, + Scope::Deploy, + LeaseState::Standby, + )]; + assert!(allocate(&i).grants.is_empty()); + } + + #[test] + fn deploy_demand_with_a_standby_promotes_it_instead_of_minting() { + let mut i = input( + &[(Class::Medium, Scope::Deploy, 1)], + vec![deploy_only_host("dep")], + ); + i.leases = vec![lease( + "dep", + Class::Medium, + Scope::Deploy, + LeaseState::Standby, + )]; + let a = allocate(&i); + assert!(a.grants.is_empty(), "{:?}", a.grants); + assert!(a.unmet.is_empty()); + assert_eq!( + a.promote, + vec![Promotion { + class: Class::Medium, + count: 1 + }] + ); + // Two queued jobs, one standby: one promotion, one mint. + let mut i = input( + &[(Class::Medium, Scope::Deploy, 2)], + vec![deploy_only_host("dep")], + ); + i.leases = vec![lease( + "dep", + Class::Medium, + Scope::Deploy, + LeaseState::Standby, + )]; + let a = allocate(&i); + assert_eq!(a.promote.len(), 1); + assert_eq!(count(&a, Class::Medium, Scope::Deploy), 1); + // A pending deploy runner already covers the job: no promotion. + let mut i = input( + &[(Class::Medium, Scope::Deploy, 1)], + vec![deploy_only_host("dep")], + ); + i.leases = vec![ + lease("dep", Class::Medium, Scope::Deploy, LeaseState::Standby), + lease("dep", Class::Medium, Scope::Deploy, LeaseState::Pending), + ]; + let a = allocate(&i); + assert!(a.grants.is_empty() && a.promote.is_empty()); + } } diff --git a/src/api.rs b/src/api.rs index 953d101..673e9c4 100644 --- a/src/api.rs +++ b/src/api.rs @@ -20,10 +20,11 @@ use serde::{Deserialize, Serialize}; use time::OffsetDateTime; use crate::allocator::{ - ActiveLease, HostCaps, LeaseState, Reservation, Unmet, free_room, free_slots, max_slots, + ActiveLease, AdmissionState, HostCaps, LeaseState, Reservation, Unmet, free_room, free_slots, + max_slots, }; use crate::bid::{ClassPressure, DroneBid}; -use crate::classes::{Class, Scope}; +use crate::classes::{Class, Scope, Trust}; use crate::enroll::hash_token; use crate::envcontract::{CheckState, EnvFacts, Need, Verdict, blocked_classes, evaluate}; use crate::health::{DroneHealth, HealthState}; @@ -56,6 +57,13 @@ pub struct HeartbeatRequest { /// Labels the drone asks to have added to its runners. #[serde(default)] pub extra_labels: Vec, + /// `shared` (default), `deploy` or `deploy-only`; see [`Trust`]. + #[serde(default)] + pub trust: Trust, + /// Standby lease ids the drone holds unstarted; acknowledges their JIT + /// configs. Leases the drone started are in `active_leases` instead. + #[serde(default)] + pub held_standby: Vec, /// Outcome of an update this drone just ran. The queen may use it as /// rollout evidence or ignore it. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -98,6 +106,18 @@ fn sanitize_facts(f: &EnvFacts, now: i64) -> EnvFacts { t.rustc_version = clip(&t.rustc_version); } } + // Only probe-list names, with bounded printable versions. + out.tools = out.tools.as_ref().map(|m| { + m.iter() + .filter(|(k, _)| crate::envcontract::PROBE_TOOLS.contains(&k.as_str())) + .map(|(k, v)| { + ( + k.clone(), + v.chars().filter(|c| !c.is_control()).take(80).collect(), + ) + }) + .collect() + }); out } @@ -139,7 +159,7 @@ fn yes() -> bool { true } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct LeaseGrant { pub lease_id: String, pub class: Class, @@ -177,6 +197,20 @@ pub struct HeartbeatResponse { /// Demand pressure per class, for the drone to set its bids against. #[serde(default)] pub pressure: Vec, + /// Leases the queen preempted for higher-priority demand; stop their jobs + /// (resent until the lease ends). + #[serde(default)] + pub cancel_leases: Vec, + /// Standby deploy runners to hold unstarted (resent until the drone lists + /// them in `held_standby`). + #[serde(default)] + pub standby: Vec, + /// Held standby leases to start now: deploy jobs are queued for them. + #[serde(default)] + pub start_standby: Vec, + /// Held standby leases that ended on the queen's side; forget them. + #[serde(default)] + pub drop_standby: Vec, } #[derive(Debug, Serialize, Deserialize)] @@ -188,6 +222,9 @@ pub struct ResultRequest { /// Compact per-job summary; see [`crate::telemetry`] for the wire format. #[serde(default)] pub telemetry: Option, + /// The drone stopped the job because the queen preempted the lease. + #[serde(default)] + pub preempted: bool, } #[derive(Debug)] @@ -198,6 +235,8 @@ pub enum ApiError { Unauthorized, NotFound, Forbidden, + /// 403 with a message the client can show. + ForbiddenWith(String), Conflict, Unprocessable, Internal(String), @@ -224,6 +263,7 @@ impl IntoResponse for ApiError { Self::Unauthorized => StatusCode::UNAUTHORIZED.into_response(), Self::NotFound => StatusCode::NOT_FOUND.into_response(), Self::Forbidden => StatusCode::FORBIDDEN.into_response(), + Self::ForbiddenWith(msg) => (StatusCode::FORBIDDEN, msg).into_response(), Self::Conflict => StatusCode::CONFLICT.into_response(), Self::Unprocessable => StatusCode::UNPROCESSABLE_ENTITY.into_response(), Self::Internal(msg) => { @@ -343,6 +383,7 @@ async fn heartbeat( .set_version(&drone, sanitize_version(req.version.as_deref()), now) .await?; store.set_extra_labels(&drone, &req.extra_labels).await?; + store.set_trust(&drone, req.trust).await?; store.set_bids(&drone, &sanitize_bids(&req.bids)).await?; if let Some(facts) = &req.env_facts { store @@ -353,6 +394,42 @@ async fn heartbeat( store.record_reclaim(&drone, stats).await?; } store.acknowledge(&drone, &req.active_leases).await?; + store + .promote_standby( + &drone, + &req.active_leases, + now, + queen.config().lease_ttl_secs, + ) + .await?; + let standby_rows = store.standby_leases(&drone).await?; + let start = queen.snapshot().start_standby; + let start_standby: Vec = standby_rows + .iter() + .filter(|l| start.contains(&l.id)) + .map(|l| l.id.clone()) + .collect(); + let drop_standby: Vec = req + .held_standby + .iter() + .filter(|id| !standby_rows.iter().any(|l| &l.id == *id)) + .cloned() + .collect(); + let standby: Vec = standby_rows + .into_iter() + .filter(|l| !req.held_standby.contains(&l.id)) + .filter_map(|l| { + Some(LeaseGrant { + lease_id: l.id, + class: l.class, + scope: l.scope, + labels: l.labels, + runner_name: l.runner_name, + jit_config: l.jit_config?, + expires_at: l.expires_at, + }) + }) + .collect(); let (units, cancel_units) = if queen.config().units.enabled { store.ack_units(&drone, &req.active_units, now).await?; store @@ -406,6 +483,15 @@ async fn heartbeat( cancel_units, desired_version: queen.desired_version(&drone), pressure: queen.snapshot().pressure, + cancel_leases: store + .preempted_active() + .await? + .into_iter() + .filter_map(|(id, d)| (d == drone).then_some(id)) + .collect(), + standby, + start_standby, + drop_standby, })) } @@ -448,8 +534,10 @@ async fn lease_result( #[derive(Debug, Serialize)] pub struct ClassStatus { + pub queued_deploy: u32, pub queued_main: u32, pub queued_pr: u32, + pub oldest_age_secs_deploy: u64, pub oldest_age_secs_main: u64, pub oldest_age_secs_pr: u64, /// Leases of this class the fleet could hold if it were empty (theoretical @@ -458,12 +546,39 @@ pub struct ClassStatus { /// Leases of this class that fit right now: capacity minus every active /// lease, on drones that are fresh and take new leases. pub free_slots: u32, + pub demand_deploy: ScopeDemand, pub demand_main: ScopeDemand, pub demand_pr: ScopeDemand, pub leases_pending: u32, pub leases_running: u32, } +impl ClassStatus { + fn queued(&self, scope: Scope) -> u32 { + match scope { + Scope::Deploy => self.queued_deploy, + Scope::Main => self.queued_main, + Scope::Pr => self.queued_pr, + } + } + + fn oldest_age_secs(&self, scope: Scope) -> u64 { + match scope { + Scope::Deploy => self.oldest_age_secs_deploy, + Scope::Main => self.oldest_age_secs_main, + Scope::Pr => self.oldest_age_secs_pr, + } + } + + fn demand(&self, scope: Scope) -> &ScopeDemand { + match scope { + Scope::Deploy => &self.demand_deploy, + Scope::Main => &self.demand_main, + Scope::Pr => &self.demand_pr, + } + } +} + /// Queued demand for one label, split by who can serve it. The split is an /// estimate: GitHub, not borg, picks which runner takes a job, and a queued job /// carries no runner yet, so borg counts the idle online runners it did not mint @@ -487,6 +602,8 @@ pub struct DroneStatus { pub heartbeat_age_secs: Option, pub fresh: bool, pub accepting: bool, + /// `shared` or `deploy`: which scopes the drone may run. + pub trust: Trust, pub hostname: Option, pub kind: Option, pub vcpu: Option, @@ -514,6 +631,16 @@ pub struct DroneStatus { /// The build environment contract as it stands for this drone; absent when /// the contract is off. pub contract: Option, + /// CPU-pressure admission from the last pass; absent for drones the last + /// pass did not consider. + pub admission: Option, +} + +/// What CPU pressure currently lets this drone take. +#[derive(Debug, Serialize)] +pub struct DroneAdmissionView { + pub state: AdmissionState, + pub cpu_psi_pm: Option, } /// One drone's build environment contract. @@ -577,8 +704,17 @@ pub struct Status { /// entries. pub reclaim: ReclaimSummary, pub units: UnitsStatus, + /// Present when the supersede janitor is enabled. + pub janitor: Option, + /// Present when preempted jobs are re-run (preemption and the janitor + /// are both configured). + pub requeue: Option, /// The newest quarantine and health audit entries. pub health_events: Vec, + /// The newest leases cancelled for higher-priority demand. + pub preemptions: Vec, + /// All-time preemption counts by victim scope and starving scope. + pub preemption_counts: Vec<(Scope, Scope, u64)>, } #[derive(Debug, Serialize)] @@ -595,6 +731,9 @@ pub struct UnitsStatus { pub joins: u64, /// Submissions answered with a stored receipt since the queen started. pub receipt_hits: u64, + /// Units cancelled by a newer submission with the same `supersede_key` + /// since the queen started. + pub superseded: u64, pub oldest_queued_age_secs: u64, pub recent: Vec, /// Why queued units are not admitted on drones that otherwise match them @@ -756,6 +895,7 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result heartbeat_age_secs: age, fresh, accepting: d.accepting, + trust: d.trust, hostname: d.capacity.as_ref().map(|c| c.hostname.clone()), kind: d.capacity.as_ref().map(|c| c.kind), vcpu: d.capacity.as_ref().map(|c| c.vcpu), @@ -774,6 +914,14 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result version: d.version.clone(), health: snap.health.get(&d.name).cloned(), reclaim: d.reclaim.clone(), + admission: snap + .admission + .iter() + .find(|a| a.drone_id == d.name) + .map(|a| DroneAdmissionView { + state: a.state, + cpu_psi_pm: a.cpu_psi_pm, + }), contract: cfg.env_contract.enabled.then(|| { let verdict = |class| { let need = @@ -823,12 +971,15 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result classes.insert( class, ClassStatus { + queued_deploy: entry(Scope::Deploy).map_or(0, |q| q.queued), queued_main: entry(Scope::Main).map_or(0, |q| q.queued), queued_pr: entry(Scope::Pr).map_or(0, |q| q.queued), + oldest_age_secs_deploy: entry(Scope::Deploy).map_or(0, |q| q.oldest_age_secs), oldest_age_secs_main: entry(Scope::Main).map_or(0, |q| q.oldest_age_secs), oldest_age_secs_pr: entry(Scope::Pr).map_or(0, |q| q.oldest_age_secs), slots: slots.get(&class).copied().unwrap_or(0), free_slots: free_by_class.get(&class).copied().unwrap_or(0), + demand_deploy: scope_demand(&snap, cfg, class, Scope::Deploy), demand_main: scope_demand(&snap, cfg, class, Scope::Main), demand_pr: scope_demand(&snap, cfg, class, Scope::Pr), leases_pending: count(LeaseStatus::Pending), @@ -875,6 +1026,7 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result counts: store.unit_counts().await?, joins, receipt_hits, + superseded: queen.units_superseded(), oldest_queued_age_secs: queued .first() .map_or(0, |u| u64::try_from(now - u.created_at).unwrap_or(0)), @@ -882,11 +1034,15 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result not_admitted: unit_blocks(cfg, &rows, &queued, now), } }, + janitor: queen.janitor().map(crate::janitor::Janitor::view), + requeue: queen.requeuer().map(crate::requeue::Requeuer::view), reclaim: ReclaimSummary { reclaimed_bytes_logged: queen.store().reclaim_bytes_logged().await?, events: queen.store().reclaim_log(20).await?, }, health_events: queen.store().health_log(20).await?, + preemptions: queen.store().preemptions(20).await?, + preemption_counts: queen.store().preemption_counts().await?, }) } @@ -954,13 +1110,7 @@ pub fn render_metrics(status: &Status) -> String { gauge( "borg_queued_jobs", "Queued jobs per class and scope.", - both(&|s, sc| { - f64::from(if sc == Scope::Main { - s.queued_main - } else { - s.queued_pr - }) - }), + both(&|s, sc| f64::from(s.queued(sc))), ); gauge( "borg_oldest_queued_age_seconds", @@ -970,11 +1120,7 @@ pub fn render_metrics(status: &Status) -> String { clippy::cast_precision_loss, reason = "ages are far below 2^52 seconds" )] - let v = (if sc == Scope::Main { - s.oldest_age_secs_main - } else { - s.oldest_age_secs_pr - }) as f64; + let v = s.oldest_age_secs(sc) as f64; v }), ); @@ -1082,24 +1228,12 @@ pub fn render_metrics(status: &Status) -> String { gauge( "borg_queued_borg_placeable", "Queued jobs borg has to place: queued minus idle non-borg runners with the label (estimate).", - both(&|s, sc| { - f64::from(if sc == Scope::Main { - s.demand_main.borg_placeable - } else { - s.demand_pr.borg_placeable - }) - }), + both(&|s, sc| f64::from(s.demand(sc).borg_placeable)), ); gauge( "borg_queued_claimable_by_static", "Queued jobs idle non-borg runners with the label can take (estimate).", - both(&|s, sc| { - f64::from(if sc == Scope::Main { - s.demand_main.claimable_by_static_estimate - } else { - s.demand_pr.claimable_by_static_estimate - }) - }), + both(&|s, sc| f64::from(s.demand(sc).claimable_by_static_estimate)), ); gauge( "borg_active_leases", @@ -1157,6 +1291,42 @@ pub fn render_metrics(status: &Status) -> String { }) .collect(), ); + gauge( + "borg_drone_cpu_psi_pm", + "Fresh CPU pressure (`some avg10`, permille) reported by the drone; absent when unknown or admission is off.", + status + .drones + .iter() + .filter_map(|d| { + let psi = d.admission.as_ref()?.cpu_psi_pm?; + Some(( + format!("drone=\"{}\"", label_value(&d.name)), + f64::from(psi), + )) + }) + .collect(), + ); + gauge( + "borg_drone_admission", + "1 for the drone's CPU-pressure admission state (open, pr_blocked, closed).", + status + .drones + .iter() + .filter_map(|d| d.admission.as_ref().map(|a| (d, a))) + .flat_map(|(d, a)| { + AdmissionState::ALL.into_iter().map(move |st| { + ( + format!( + "drone=\"{}\",state=\"{}\"", + label_value(&d.name), + st.as_str() + ), + f64::from(u8::from(a.state == st)), + ) + }) + }) + .collect(), + ); gauge( "borg_drone_host_failure_score", "Decayed host-caused failure evidence per drone (1 = one fresh failure).", @@ -1441,12 +1611,111 @@ pub fn render_metrics(status: &Status) -> String { "Unit submissions answered with a stored receipt since the queen started.", vec![(String::new(), status.units.receipt_hits as f64)], ); + gauge( + "borg_units_superseded_total", + "Work units cancelled by a newer submission with the same supersede key since the queen started.", + vec![(String::new(), status.units.superseded as f64)], + ); gauge( "borg_unit_oldest_queued_age_seconds", "Age of the oldest queued work unit.", vec![(String::new(), status.units.oldest_queued_age_secs as f64)], ); } + if let Some(j) = &status.janitor { + let t = &j.total; + for (field, help, v) in [ + ( + "considered", + "Pull-request runs the janitor examined", + t.considered, + ), + ( + "superseded", + "Runs whose head was no longer the pull request's head", + t.superseded, + ), + ( + "cancelled", + "Queued superseded runs the janitor cancelled", + t.cancelled, + ), + ( + "cancelled_running", + "Running superseded runs the janitor cancelled", + t.cancelled_running, + ), + ( + "duplicates", + "Same-workflow same-commit runs beyond the one kept", + t.duplicates, + ), + ( + "dedup_cancelled", + "Queued duplicate runs the janitor cancelled", + t.dedup_cancelled, + ), + ( + "would_cancel", + "Runs a dry-run janitor would have cancelled", + t.would_cancel, + ), + ( + "deferred", + "Runs left for a later pass by the per-pass cap", + t.deferred, + ), + ( + "raced", + "Runs that changed state before the janitor cancelled them", + t.raced, + ), + ( + "errors", + "Janitor per-run or per-pull-request failures", + t.errors, + ), + ] { + gauge( + &format!("borg_janitor_{field}_total"), + &format!("{help}, summed over passes since the queen started."), + vec![(String::new(), f64::from(v))], + ); + } + } + if let Some(r) = &status.requeue { + let t = &r.total; + for (field, help, v) in [ + ("checked", "Preempted jobs the requeuer examined", t.checked), + ("rerun", "Preempted jobs re-run", t.rerun), + ( + "waiting", + "Preempted jobs whose run had not completed", + t.waiting, + ), + ( + "stale", + "Preempted jobs older than 24 hours, skipped", + t.stale, + ), + ( + "terminal", + "Preempted jobs GitHub refused to re-run (403, 404 or 409)", + t.terminal, + ), + ( + "errors", + "Requeue failures retried on a later pass", + t.errors, + ), + ] { + gauge( + &format!("borg_requeue_{field}_total"), + &format!("{help}, summed over passes since the queen started."), + vec![(String::new(), f64::from(v))], + ); + } + } gauge( "borg_reservations", "Drones held back for starving demand.", @@ -1524,6 +1793,18 @@ pub fn render_metrics(status: &Status) -> String { }, )], ); + let _ = writeln!( + out, + "# HELP borg_preemptions_total Leases cancelled for higher-priority demand.\n# TYPE borg_preemptions_total counter" + ); + for (from, to, n) in &status.preemption_counts { + let _ = writeln!( + out, + "borg_preemptions_total{{from_scope=\"{}\",to_scope=\"{}\"}} {n}", + from.as_str(), + to.as_str() + ); + } out } @@ -1535,10 +1816,7 @@ fn label_value(v: &str) -> String { } fn scope_name(s: Scope) -> &'static str { - match s { - Scope::Main => "main", - Scope::Pr => "pr", - } + s.as_str() } async fn metrics(State(queen): State>) -> Result { @@ -1598,8 +1876,13 @@ pub async fn serve( // Cancel-safe for the same reason as the ticker above. tokio::select! { _ = every.tick() => { - let report = janitor.pass(&queen.config().repos).await; - tracing::debug!("janitor pass: {report:?}"); + janitor.pass(&queen.config().repos).await; + if let Some(requeuer) = queen.requeuer() { + let now = OffsetDateTime::now_utc().unix_timestamp(); + if let Err(e) = requeuer.pass(queen.store(), now).await { + tracing::warn!("requeue pass failed: {e}"); + } + } } _ = stop_rx.changed() => break, } diff --git a/src/classes.rs b/src/classes.rs index 0f30e36..5867f5e 100644 --- a/src/classes.rs +++ b/src/classes.rs @@ -1,7 +1,8 @@ //! Runner classes: the resources a job of each size needs. //! //! Jobs ask for a class through its runner label (`` for main -//! and merge-group jobs, `-pr` for pull requests). The prefix +//! and merge-group jobs, `-pr` for pull requests, +//! `-deploy` for deploy-path jobs). The prefix //! and the per-class requirements are configuration with generic defaults. use std::collections::BTreeMap; @@ -15,6 +16,8 @@ pub const DEFAULT_LABEL_PREFIX: &str = "borg-ci-"; pub const MAX_EXTRA_LABELS: usize = 16; /// Suffix marking pull-request demand. const PR_SUFFIX: &str = "-pr"; +/// Suffix marking deploy-path demand. +const DEPLOY_SUFFIX: &str = "-deploy"; /// A runner size class. Ordered smallest to largest. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] @@ -44,17 +47,89 @@ impl fmt::Display for Class { } } -/// Where a queued job comes from. Main jobs (main branch, merge groups) are +/// Where a queued job comes from. Deploy jobs (prod OIDC, signing, registry +/// push) are scheduled before main jobs (main branch, merge groups), which are /// scheduled before pull-request jobs. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum Scope { + Deploy, Main, Pr, } impl Scope { - pub const ALL: [Scope; 2] = [Scope::Main, Scope::Pr]; + /// Every scope in scheduling order. + pub const ALL: [Scope; 3] = [Scope::Deploy, Scope::Main, Scope::Pr]; + + pub fn as_str(self) -> &'static str { + match self { + Scope::Deploy => "deploy", + Scope::Main => "main", + Scope::Pr => "pr", + } + } +} + +/// How much a drone's host is trusted with deploy-path jobs. A deploy drone +/// never runs a pull-request job or a work unit; a shared drone never runs a +/// deploy job; a deploy-only drone runs nothing but deploy jobs. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum Trust { + #[default] + Shared, + Deploy, + DeployOnly, +} + +impl Trust { + pub fn as_str(self) -> &'static str { + match self { + Trust::Shared => "shared", + Trust::Deploy => "deploy", + Trust::DeployOnly => "deploy-only", + } + } + + /// Whether demand of `scope` may be placed on a drone of this trust. + pub fn admits(self, scope: Scope) -> bool { + match scope { + Scope::Deploy => self != Trust::Shared, + Scope::Main => self != Trust::DeployOnly, + Scope::Pr => self == Trust::Shared, + } + } + + /// Whether drones of this trust hold standby deploy runners. + pub fn serves_deploy(self) -> bool { + self != Trust::Shared + } +} + +/// What an alias label stands for: a bare class name (main scope) or an object +/// naming both. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum Alias { + Main(Class), + Scoped(AliasTarget), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AliasTarget { + pub class: Class, + pub scope: Scope, +} + +impl Alias { + pub fn target(self) -> (Class, Scope) { + match self { + Alias::Main(class) => (class, Scope::Main), + Alias::Scoped(t) => (t.class, t.scope), + } + } } /// Resources reserved for one running job of a class. @@ -77,10 +152,13 @@ pub struct Classes { pub medium: ClassSpec, pub heavy: ClassSpec, /// Existing runner labels that borg also answers, each mapped to the class - /// that serves it (for example a deploy lane's `runs-on` label). Queued - /// jobs with an alias label count as main demand for the class, and every - /// runner of that class carries the alias labels so it can take them. - pub aliases: BTreeMap, + /// that serves it (for example a deploy lane's `runs-on` label). A plain + /// class name (`"medium"`) is main demand; `{"class": "medium", "scope": + /// "deploy"}` is deploy demand. Queued jobs with an alias label count as + /// demand of its scope, and a runner carries the alias label only when its + /// grant scope equals the alias's (main aliases also ride on PR runners, + /// which carry the bare label). + pub aliases: BTreeMap, } impl Default for Classes { @@ -117,7 +195,9 @@ pub enum ClassesError { Parse(#[from] serde_json::Error), #[error("label prefix must not be empty")] EmptyPrefix, - #[error("alias label {0:?} is malformed or starts with the class label prefix")] + #[error( + "alias label {0:?} is malformed, starts with the class label prefix, or has scope \"pr\"" + )] BadAlias(String), #[error("class {class} must reserve at least 1 MiB and 1 vCPU")] Empty { class: Class }, @@ -147,8 +227,11 @@ impl Classes { if self.label_prefix.is_empty() { return Err(ClassesError::EmptyPrefix); } - for label in self.aliases.keys() { - if !valid_extra_label(label) || label.starts_with(self.label_prefix.as_str()) { + for (label, alias) in &self.aliases { + if !valid_extra_label(label) + || label.starts_with(self.label_prefix.as_str()) + || alias.target().1 == Scope::Pr + { return Err(ClassesError::BadAlias(label.clone())); } } @@ -202,18 +285,22 @@ impl Classes { match scope { Scope::Main => format!("{}{class}", self.label_prefix), Scope::Pr => format!("{}{class}{PR_SUFFIX}", self.label_prefix), + Scope::Deploy => format!("{}{class}{DEPLOY_SUFFIX}", self.label_prefix), } } /// Parse a job's `runs-on` label into its class and scope. pub fn parse_label(&self, label: &str) -> Option<(Class, Scope)> { - if let Some(&class) = self.aliases.get(label) { - return Some((class, Scope::Main)); + if let Some(alias) = self.aliases.get(label) { + return Some(alias.target()); } let rest = label.strip_prefix(self.label_prefix.as_str())?; - let (name, scope) = match rest.strip_suffix(PR_SUFFIX) { - Some(name) => (name, Scope::Pr), - None => (rest, Scope::Main), + let (name, scope) = if let Some(name) = rest.strip_suffix(DEPLOY_SUFFIX) { + (name, Scope::Deploy) + } else if let Some(name) = rest.strip_suffix(PR_SUFFIX) { + (name, Scope::Pr) + } else { + (rest, Scope::Main) }; let class = Class::ALL.into_iter().find(|c| c.as_str() == name)?; Some((class, scope)) @@ -222,16 +309,25 @@ impl Classes { /// The exact label set a runner is minted with. /// /// Main grants carry only the bare label so GitHub can hand them nothing - /// but main jobs. Pull-request grants carry both labels. + /// but main jobs. Pull-request grants carry both labels. Deploy grants + /// carry only the `-deploy` label, so GitHub can hand them nothing but + /// deploy jobs. Alias labels ride on the runners whose scope they name. pub fn runner_labels(&self, class: Class, scope: Scope) -> Vec { - let mut labels = vec![self.job_label(class, Scope::Main)]; - if scope == Scope::Pr { - labels.push(self.job_label(class, Scope::Pr)); - } + let mut labels = match scope { + Scope::Deploy => vec![self.job_label(class, Scope::Deploy)], + Scope::Main => vec![self.job_label(class, Scope::Main)], + Scope::Pr => vec![ + self.job_label(class, Scope::Main), + self.job_label(class, Scope::Pr), + ], + }; labels.extend( self.aliases .iter() - .filter(|&(_, &c)| c == class) + .filter(|&(_, a)| { + let (c, s) = a.target(); + c == class && (s == scope || (s == Scope::Main && scope == Scope::Pr)) + }) .map(|(l, _)| l.clone()), ); labels @@ -441,6 +537,98 @@ mod tests { } } + #[test] + fn deploy_labels_round_trip_and_deploy_runners_carry_only_their_label() { + let c = Classes::default(); + assert_eq!( + c.job_label(Class::Medium, Scope::Deploy), + "borg-ci-medium-deploy" + ); + assert_eq!( + c.parse_label("borg-ci-heavy-deploy"), + Some((Class::Heavy, Scope::Deploy)) + ); + assert_eq!(c.parse_label("borg-ci-heavy-deploy-pr"), None); + assert_eq!( + c.runner_labels(Class::Medium, Scope::Deploy), + ["borg-ci-medium-deploy"] + ); + assert_eq!( + c.labels_scope(&c.runner_labels(Class::Light, Scope::Deploy)), + Some((Class::Light, Scope::Deploy)) + ); + assert!(Scope::Deploy < Scope::Main && Scope::Main < Scope::Pr); + } + + #[test] + fn trust_admits_scopes() { + assert!(Trust::Deploy.admits(Scope::Deploy) && Trust::Deploy.admits(Scope::Main)); + assert!(!Trust::Deploy.admits(Scope::Pr)); + assert!(Trust::Shared.admits(Scope::Main) && Trust::Shared.admits(Scope::Pr)); + assert!(!Trust::Shared.admits(Scope::Deploy)); + assert!(Trust::DeployOnly.admits(Scope::Deploy)); + assert!(!Trust::DeployOnly.admits(Scope::Main)); + assert!(!Trust::DeployOnly.admits(Scope::Pr)); + } + + #[test] + fn deploy_only_trust_serializes_kebab_case() { + assert_eq!(Trust::DeployOnly.as_str(), "deploy-only"); + let t: Trust = serde_json::from_str("\"deploy-only\"").unwrap(); + assert_eq!(t, Trust::DeployOnly); + assert_eq!(serde_json::to_string(&t).unwrap(), "\"deploy-only\""); + assert!(Trust::Deploy.serves_deploy() && !Trust::Shared.serves_deploy()); + } + + #[test] + fn scoped_aliases_ride_only_on_runners_of_their_scope() { + let c = Classes::from_json( + r#"{"label_prefix":"x-ci-","aliases":{ + "promote":{"class":"medium","scope":"deploy"}, + "legacy":"medium"}}"#, + ) + .unwrap(); + assert_eq!( + c.parse_label("promote"), + Some((Class::Medium, Scope::Deploy)) + ); + assert_eq!(c.parse_label("legacy"), Some((Class::Medium, Scope::Main))); + assert_eq!( + c.runner_labels(Class::Medium, Scope::Deploy), + ["x-ci-medium-deploy", "promote"] + ); + assert_eq!( + c.runner_labels(Class::Medium, Scope::Main), + ["x-ci-medium", "legacy"] + ); + assert_eq!( + c.runner_labels(Class::Medium, Scope::Pr), + ["x-ci-medium", "x-ci-medium-pr", "legacy"] + ); + assert_eq!( + c.runner_labels(Class::Heavy, Scope::Deploy), + ["x-ci-heavy-deploy"] + ); + } + + #[test] + fn alias_labels_must_not_look_like_class_labels() { + for bad in [ + r#"{"aliases":{"borg-ci-medium-deploy":"light"}}"#, + r#"{"aliases":{"borg-ci-medium":{"class":"light","scope":"deploy"}}}"#, + r#"{"aliases":{"lane":{"class":"light","scope":"pr"}}}"#, + ] { + assert!( + matches!(Classes::from_json(bad), Err(ClassesError::BadAlias(_))), + "{bad}" + ); + } + assert!(matches!( + Classes::from_json(r#"{"aliases":{"lane":{"class":"light","scope":"deploy","x":1}}}"#), + Err(ClassesError::Parse(_)) + )); + } + #[test] fn extra_label_syntax() { assert!(valid_extra_label("on-prem") && valid_extra_label("a.b_c-1")); diff --git a/src/config.rs b/src/config.rs index eae2f49..3f32157 100644 --- a/src/config.rs +++ b/src/config.rs @@ -7,11 +7,12 @@ use std::path::{Path, PathBuf}; use serde::Deserialize; -use crate::allocator::{WarmFloor, default_warm_floors}; -use crate::classes::{Classes, ClassesError}; +use crate::allocator::{AdmissionThresholds, WarmFloor, default_warm_floors}; +use crate::classes::{Class, Classes, ClassesError}; use crate::envcontract::ContractSettings; use crate::health::HealthSettings; use crate::homeostasis::HomeostasisSettings; +use crate::preempt::PreemptSettings; use crate::profile::ProfileSettings; use crate::rollout::RolloutSettings; use crate::schedule::WarmSchedule; @@ -78,6 +79,13 @@ pub struct JanitorSettings { /// Hard cap on cancellations (or dry-run matches) per pass. #[serde(default = "default_janitor_max")] pub max_cancels_per_pass: u32, + /// Also cancel superseded runs that are already `in_progress`. + #[serde(default)] + pub cancel_running: bool, + /// Cancel queued runs that duplicate another run of the same workflow, + /// commit, event and pull request. + #[serde(default)] + pub dedupe_same_sha: bool, } fn yes() -> bool { @@ -90,6 +98,61 @@ fn default_janitor_max() -> u32 { 20 } +/// Standby deploy runners; present means enabled. See `docs/operations.md`. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct StandbySettings { + /// Classes each deploy or deploy-only drone keeps standby runners for. + pub classes: Vec, + /// Standby runners per drone and class, 1 to 4. + pub per_drone: u32, +} + +impl StandbySettings { + fn validate(&self) -> Result<(), String> { + if self.classes.is_empty() { + return Err("standby.classes must not be empty".to_owned()); + } + if !(1..=4).contains(&self.per_drone) { + return Err("standby.per_drone must be between 1 and 4".to_owned()); + } + Ok(()) + } +} + +/// Admission by CPU pressure; present means enabled. See +/// [`crate::allocator::AdmissionThresholds`]. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AdmissionSettings { + /// From this `cpu some avg10` (permille) a drone takes no new + /// pull-request grants. + pub psi_soft_pm: u16, + /// From this pressure a drone takes no new grants at all. + pub psi_hard_pm: u16, + /// A drone's reported pressure older than this counts as unknown. + pub facts_max_age_secs: i64, +} + +impl AdmissionSettings { + pub fn thresholds(&self) -> AdmissionThresholds { + AdmissionThresholds { + psi_soft_pm: self.psi_soft_pm, + psi_hard_pm: self.psi_hard_pm, + } + } + + fn validate(&self) -> Result<(), String> { + if self.psi_soft_pm == 0 || self.psi_soft_pm > self.psi_hard_pm || self.psi_hard_pm > 1000 { + return Err("admission needs 0 < psi_soft_pm <= psi_hard_pm <= 1000".to_owned()); + } + if self.facts_max_age_secs <= 0 { + return Err("admission.facts_max_age_secs must be positive".to_owned()); + } + Ok(()) + } +} + #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] pub struct QueenConfig { @@ -161,9 +224,18 @@ pub struct QueenConfig { /// `/status`): drones update themselves, the queen never forces it. #[serde(default)] pub rollout_target_version: Option, + /// Absent (the default) disables standby deploy runners. + #[serde(default)] + pub standby: Option, + /// Absent (the default) disables admission by CPU pressure. + #[serde(default)] + pub admission: Option, /// Absent (the default) disables the supersede janitor. #[serde(default)] pub supersede_janitor: Option, + /// Absent (the default) disables preemption; see [`crate::preempt`]. + #[serde(default)] + pub preemption: Option, } fn default_listen() -> SocketAddr { @@ -275,6 +347,12 @@ impl QueenConfig { self.profiles.validate().map_err(ConfigError::Invalid)?; self.homeostasis.validate().map_err(ConfigError::Invalid)?; self.health.validate().map_err(ConfigError::Invalid)?; + if let Some(a) = &self.admission { + a.validate().map_err(ConfigError::Invalid)?; + } + if let Some(st) = &self.standby { + st.validate().map_err(ConfigError::Invalid)?; + } if let Some(t) = &self.tailnet_enroll { t.validate().map_err(ConfigError::Invalid)?; } @@ -288,6 +366,9 @@ impl QueenConfig { "supersede_janitor interval_secs and max_cancels_per_pass must be positive", ); } + if let Some(p) = &self.preemption { + p.validate().map_err(ConfigError::Invalid)?; + } Ok(()) } } @@ -333,6 +414,7 @@ mod tests { .unwrap(); let j = c.supersede_janitor.unwrap(); assert!(j.dry_run); + assert!(!j.cancel_running && !j.dedupe_same_sha); assert_eq!((j.interval_secs, j.max_cancels_per_pass), (60, 20)); for bad in [ "{\"interval_secs\":0}", @@ -364,6 +446,50 @@ mod tests { ); } + #[test] + fn standby_is_off_by_default_and_validated() { + assert!(QueenConfig::from_json(MIN).unwrap().standby.is_none()); + let with = |a: &str| { + QueenConfig::from_json(&MIN.replace("\"repos\"", &format!("\"standby\":{a},\"repos\""))) + }; + let ok = with(r#"{"classes":["medium"],"per_drone":1}"#).unwrap(); + let st = ok.standby.unwrap(); + assert_eq!((st.classes, st.per_drone), (vec![Class::Medium], 1)); + for bad in [ + r#"{"classes":[],"per_drone":1}"#, + r#"{"classes":["medium"],"per_drone":0}"#, + r#"{"classes":["medium"],"per_drone":5}"#, + r#"{"classes":["huge"],"per_drone":1}"#, + r#"{"classes":["medium"],"per_drone":1,"nope":1}"#, + r#"{"per_drone":1}"#, + ] { + assert!(with(bad).is_err(), "{bad}"); + } + } + + #[test] + fn admission_is_off_by_default_and_validated() { + assert!(QueenConfig::from_json(MIN).unwrap().admission.is_none()); + let with = |a: &str| { + QueenConfig::from_json( + &MIN.replace("\"repos\"", &format!("\"admission\":{a},\"repos\"")), + ) + }; + let ok = with(r#"{"psi_soft_pm":300,"psi_hard_pm":600,"facts_max_age_secs":180}"#).unwrap(); + let t = ok.admission.unwrap().thresholds(); + assert_eq!((t.psi_soft_pm, t.psi_hard_pm), (300, 600)); + for bad in [ + r#"{"psi_soft_pm":700,"psi_hard_pm":600,"facts_max_age_secs":180}"#, + r#"{"psi_soft_pm":0,"psi_hard_pm":600,"facts_max_age_secs":180}"#, + r#"{"psi_soft_pm":300,"psi_hard_pm":1001,"facts_max_age_secs":180}"#, + r#"{"psi_soft_pm":300,"psi_hard_pm":600,"facts_max_age_secs":0}"#, + r#"{"psi_soft_pm":300,"psi_hard_pm":600,"facts_max_age_secs":180,"nope":1}"#, + r#"{"psi_soft_pm":300,"psi_hard_pm":600}"#, + ] { + assert!(with(bad).is_err(), "{bad}"); + } + } + #[test] fn zero_telemetry_retention_is_rejected() { let json = MIN.replace("\"repos\"", "\"telemetry_retention_days\":0,\"repos\""); @@ -392,6 +518,27 @@ mod tests { ); } + #[test] + fn preemption_is_off_by_default_and_validated() { + assert!(QueenConfig::from_json(MIN).unwrap().preemption.is_none()); + let with = |p: &str| { + QueenConfig::from_json( + &MIN.replace("\"repos\"", &format!("\"preemption\":{p},\"repos\"")), + ) + }; + let ok = with(r#"{"after_secs":600,"max_per_pass":1,"min_victim_age_secs":60}"#).unwrap(); + assert_eq!(ok.preemption.unwrap().after_secs, 600); + for bad in [ + r#"{"after_secs":0,"max_per_pass":1,"min_victim_age_secs":60}"#, + r#"{"after_secs":600,"max_per_pass":0,"min_victim_age_secs":60}"#, + r#"{"after_secs":600,"max_per_pass":1,"min_victim_age_secs":0}"#, + r#"{"after_secs":600,"max_per_pass":1}"#, + r#"{"after_secs":600,"max_per_pass":1,"min_victim_age_secs":60,"nope":1}"#, + ] { + assert!(with(bad).is_err(), "{bad}"); + } + } + #[test] fn starvation_can_be_disabled() { let c = QueenConfig::from_json( @@ -422,7 +569,8 @@ mod tests { "starvation_after_secs": null, "health": {"auto_quarantine": false}, "units": {"enabled": true}, - "supersede_janitor": {"dry_run": false} + "admission": {"psi_soft_pm": 300, "psi_hard_pm": 600, "facts_max_age_secs": 180}, + "supersede_janitor": {"dry_run": false, "cancel_running": true, "dedupe_same_sha": true} }"#; const FULL_TOML: &str = r#" @@ -452,8 +600,15 @@ auto_quarantine = false [units] enabled = true +[admission] +psi_soft_pm = 300 +psi_hard_pm = 600 +facts_max_age_secs = 180 + [supersede_janitor] dry_run = false +cancel_running = true +dedupe_same_sha = true "#; #[test] @@ -464,6 +619,9 @@ dry_run = false assert_eq!(t.repos.len(), 2); assert_eq!(t.starvation_after_secs, None); assert!(t.units.enabled); + assert_eq!(t.admission.unwrap().psi_hard_pm, 600); + let j = t.supersede_janitor.unwrap(); + assert!(!j.dry_run && j.cancel_running && j.dedupe_same_sha); } #[test] diff --git a/src/drone/config.rs b/src/drone/config.rs index f599697..ee56b32 100644 --- a/src/drone/config.rs +++ b/src/drone/config.rs @@ -213,9 +213,16 @@ pub struct DroneConfig { pub docker_bin: String, #[serde(default = "default_heartbeat")] pub heartbeat_secs: u64, + /// A drone holding standby deploy runners starts them all once the queen + /// has not answered a heartbeat for this long. 0 disables fail-open. + #[serde(default = "default_fail_open")] + pub fail_open_after_secs: u64, /// A lease running longer than this is killed. #[serde(default = "default_max_job")] pub max_job_secs: u64, + /// A preempted lease's runner gets SIGINT, then is killed after this long. + #[serde(default = "default_preempt_grace")] + pub preempt_grace_secs: u64, /// On SIGTERM, wait this long for in-flight jobs before killing them. #[serde(default = "default_drain")] pub drain_timeout_secs: u64, @@ -247,6 +254,12 @@ pub struct DroneConfig { /// ignored by the queen so they cannot change a runner's scope. #[serde(default)] pub extra_labels: Vec, + /// `shared` (default), `deploy` or `deploy-only`. A deploy drone runs only + /// deploy and main jobs, never pull-request jobs or work units; a + /// deploy-only drone runs only deploy jobs; a shared drone never runs + /// deploy jobs. + #[serde(default)] + pub trust: crate::classes::Trust, /// Environment injected into every lease's runner. #[serde(default)] pub runner_env: BTreeMap, @@ -276,9 +289,15 @@ fn default_docker() -> String { fn default_heartbeat() -> u64 { 10 } +fn default_fail_open() -> u64 { + 120 +} fn default_max_job() -> u64 { 6 * 3600 } +fn default_preempt_grace() -> u64 { + 30 +} fn default_drain() -> u64 { 3600 } @@ -316,6 +335,9 @@ impl DroneConfig { if self.heartbeat_secs == 0 || self.max_job_secs == 0 { return bad("heartbeat_secs and max_job_secs must be positive"); } + if self.preempt_grace_secs == 0 { + return bad("preempt_grace_secs must be positive"); + } if self.max_vcpu == Some(0) || self.max_mem_mib == Some(0) || self.max_leases == Some(0) { return bad("owner caps must be positive when set"); } @@ -507,6 +529,12 @@ mod tests { fn defaults_and_caps() { let c = DroneConfig::from_json(BASE).unwrap(); assert_eq!((c.heartbeat_secs, c.docker_bin.as_str()), (10, "docker")); + assert_eq!(c.fail_open_after_secs, 120); + let off = BASE.replace('}', r#","fail_open_after_secs":0}"#); + assert_eq!( + DroneConfig::from_json(&off).unwrap().fail_open_after_secs, + 0 + ); let with = BASE.replace('}', r#","max_vcpu":8,"max_mem_mib":24576}"#); let caps = DroneConfig::from_json(&with).unwrap().host_caps(); assert_eq!((caps.max_vcpu, caps.max_mem_mib), (Some(8), Some(24576))); diff --git a/src/drone/fsops.rs b/src/drone/fsops.rs index ee9fc55..3c095f4 100644 --- a/src/drone/fsops.rs +++ b/src/drone/fsops.rs @@ -97,6 +97,27 @@ impl ExecFs { Self { exec } } + fn run_dirs_writable(&self, path: &Path) -> io::Result<()> { + self.run( + "find", + &[ + &path.to_string_lossy(), + "-type", + "d", + "-exec", + "chmod", + "u+w", + "{}", + "+", + ], + ) + } + + /// Best effort: a missing path or a failed chmod must not block removal. + fn make_dirs_writable(&self, path: &Path) { + let _ = self.run_dirs_writable(path); + } + fn run(&self, program: &str, args: &[&str]) -> io::Result<()> { let out = self.exec.run(program, args)?; if out.success { @@ -130,6 +151,10 @@ impl FsOps for ExecFs { } fn remove_all(&self, path: &Path) -> io::Result<()> { + // Cloned slot directories keep the runner tree's read-only modes; a + // read-only directory cannot have its entries removed. Directories + // only: files are hard links shared with the pinned runner install. + self.make_dirs_writable(path); self.run("rm", &["-rf", "--", &path.to_string_lossy()]) } @@ -141,7 +166,10 @@ impl FsOps for ExecFs { self.run( "cp", &["-al", "--", &src.to_string_lossy(), &dst.to_string_lossy()], - ) + )?; + // The runner writes `run-helper.sh`, `_diag` and `_work` into the + // slot, so its directories must be writable (the pinned install's are not). + self.run_dirs_writable(dst) } } @@ -176,14 +204,16 @@ mod tests { .unwrap(); let calls = fake.calls(); assert_eq!(calls[0], "mkdir -p -- /c/git"); - assert_eq!(calls[1], "rm -rf -- /w/slot-0"); - assert_eq!(calls[2], "mkdir -p -- /w"); - assert_eq!(calls[3], "cp -al -- /r /w/slot-0"); - assert!(calls[4].starts_with("sh -c "), "{}", calls[4]); + assert_eq!(calls[1], "find /w/slot-0 -type d -exec chmod u+w {} +"); + assert_eq!(calls[2], "rm -rf -- /w/slot-0"); + assert_eq!(calls[3], "mkdir -p -- /w"); + assert_eq!(calls[4], "cp -al -- /r /w/slot-0"); + assert_eq!(calls[5], "find /w/slot-0 -type d -exec chmod u+w {} +"); + assert!(calls[6].starts_with("sh -c "), "{}", calls[6]); assert!( - calls[4].ends_with("sh common --x\n; rm -rf / /h/.bazelrc"), + calls[6].ends_with("sh common --x\n; rm -rf / /h/.bazelrc"), "content is an argument, not script text: {}", - calls[4] + calls[6] ); } diff --git a/src/drone/install.rs b/src/drone/install.rs index 0a124d7..ca53265 100644 --- a/src/drone/install.rs +++ b/src/drone/install.rs @@ -84,6 +84,10 @@ pub struct InstallOptions { /// leaving its files as they are. pub adopt_cache: bool, pub extra_labels: Vec, + pub trust: crate::classes::Trust, + /// Seconds without a queen answer before standby deploy runners start; + /// 0 turns fail-open off. + pub fail_open_after_secs: u64, pub runner_env: BTreeMap, pub job_started_hook: Option, pub job_completed_hook: Option, @@ -791,7 +795,9 @@ impl Installer<'_> { cache_root: self.layout.cache_root.clone(), docker_bin: "docker".to_owned(), heartbeat_secs: 10, + fail_open_after_secs: o.fail_open_after_secs, max_job_secs: 6 * 3600, + preempt_grace_secs: 30, drain_timeout_secs: o.drain_timeout_secs, max_vcpu: o.max_vcpu, max_mem_mib: o.max_mem_mib, @@ -805,6 +811,7 @@ impl Installer<'_> { priv_bin: "setpriv".to_owned(), }), extra_labels: o.extra_labels.clone(), + trust: o.trust, runner_env: o.runner_env.clone(), job_started_hook: o.job_started_hook.clone(), job_completed_hook: o.job_completed_hook.clone(), @@ -1012,6 +1019,8 @@ mod tests { cache_root: None, adopt_cache: false, extra_labels: vec!["on-prem".into()], + trust: crate::classes::Trust::Shared, + fail_open_after_secs: 120, runner_env: BTreeMap::new(), job_started_hook: None, job_completed_hook: None, diff --git a/src/drone/mod.rs b/src/drone/mod.rs index 0b9958a..d09e7cb 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -17,8 +17,9 @@ pub mod pins; pub mod runner; pub mod selfupdate; pub mod slots; +pub mod standby; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::Path; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; @@ -38,6 +39,7 @@ use jobuser::JobExec; use runner::{RunSpec, run_lease}; use selfupdate::{Handled, Plan, SelfUpdate}; use slots::{SlotGuard, Slots, slot_dir, slot_name}; +use standby::{Standbys, should_fail_open}; /// Attempts to deliver a lease result before giving up (the lease then /// expires on the queen's side). The waits grow to a minute, so a queen that @@ -46,9 +48,13 @@ const RESULT_ATTEMPTS: u32 = 20; type Held = Arc>>; +/// Per running lease, the switch the queen's `cancel_leases` flips. +type Preempts = Arc>>>; + /// Removes a lease id from the held set when its task ends, even by panic. struct HeldGuard { held: Held, + preempts: Preempts, id: String, } @@ -57,6 +63,9 @@ impl Drop for HeldGuard { if let Ok(mut held) = self.held.lock() { held.remove(&self.id); } + if let Ok(mut preempts) = self.preempts.lock() { + preempts.remove(&self.id); + } } } @@ -70,6 +79,9 @@ pub struct Drone { /// File operations on job-owned directories. fs_ops: Arc, held: Held, + /// Standby deploy runners held unstarted. + standbys: Standbys, + preempts: Preempts, abort: tokio::sync::watch::Sender, proc_root: Arc, caches: Arc, @@ -78,6 +90,8 @@ pub struct Drone { selfupdate: Option, /// The queen wants another version: stop taking leases until idle. update_wanted: AtomicBool, + /// The last tools probe, reused for a minute. + tools_cache: probe::ToolsCache, } impl Drone { @@ -111,10 +125,12 @@ impl Drone { Arc::clone(&host), )); let selfupdate = cfg.state_dir.clone().map(SelfUpdate::new); + let standbys = Standbys::load(cfg.state_dir.clone()); Self { cfg, selfupdate, update_wanted: AtomicBool::new(false), + tools_cache: probe::ToolsCache::default(), client: Arc::new(client), host, exec, @@ -123,6 +139,8 @@ impl Drone { caches, slots: Slots::new(slot_count), held: Held::default(), + standbys, + preempts: Preempts::default(), abort, proc_root: Arc::from(Path::new("/proc")), } @@ -187,6 +205,21 @@ impl Drone { path: format!("probe: {e}"), status: 0, })?; + let host = Arc::clone(&self.host); + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + let mut env_facts = + tokio::task::spawn_blocking(move || probe::env_facts(host.as_ref(), now)) + .await + .map_err(|e| ClientError::Status { + path: format!("probe: {e}"), + status: 0, + })?; + // The tools probe never fails the heartbeat: no map means not measured. + let (exec, cache) = (Arc::clone(&self.job_exec), self.tools_cache.clone()); + env_facts.tools = tokio::task::spawn_blocking(move || cache.get(exec.as_ref(), now)) + .await + .ok() + .flatten(); let req = HeartbeatRequest { capacity, caps: self.cfg.host_caps(), @@ -195,6 +228,8 @@ impl Drone { cache_stats: self.caches.stats(), version: Some(env!("CARGO_PKG_VERSION").to_owned()), extra_labels: self.cfg.extra_labels.clone(), + trust: self.cfg.trust, + held_standby: self.standbys.ids(), bids: Vec::new(), reclaim: None, update_report: self @@ -202,42 +237,105 @@ impl Drone { .as_ref() .and_then(SelfUpdate::pending_report), active_units: Vec::new(), - env_facts: None, + env_facts: Some(env_facts), }; let resp = self.client.heartbeat(&req).await?; let defer = self.after_heartbeat(&req, resp.desired_version.as_deref()); + if let Ok(preempts) = self.preempts.lock() { + for id in &resp.cancel_leases { + if let Some(tx) = preempts.get(id) { + let _ = tx.send(true); + } + } + } + for id in &resp.drop_standby { + self.standbys.remove(id); + } + for grant in resp.standby { + self.standbys.insert(grant); + } let mut started = 0; for grant in resp.leases { - if !accepting || defer || !self.admit(&grant.lease_id) { - continue; + if accepting && !defer && self.start_lease(grant, tasks) { + started += 1; } - let held = HeldGuard { - held: Arc::clone(&self.held), - id: grant.lease_id.clone(), - }; - let Some(slot) = self.slots.acquire(&grant.lease_id) else { - tracing::warn!( - "no free slot for lease {}; leaving it unstarted", - grant.lease_id - ); + } + if accepting && !defer { + started += self.start_standbys(Some(&resp.start_standby), tasks); + } + Ok(started) + } + + /// Start held standby runners: the ones in `only`, or all of them. A + /// standby that cannot start now (no slot) stays held. + fn start_standbys(&self, only: Option<&[String]>, tasks: &mut JoinSet<()>) -> usize { + let ids = match only { + Some(ids) => ids.to_vec(), + None => self.standbys.ids(), + }; + let mut started = 0; + for id in ids { + let Some(grant) = self.standbys.get(&id) else { continue; }; - let ctx = LeaseContext { - cfg: Arc::clone(&self.cfg), - client: Arc::clone(&self.client), - exec: Arc::clone(&self.exec), - job_exec: Arc::clone(&self.job_exec), - fs_ops: Arc::clone(&self.fs_ops), - abort: self.abort.subscribe(), - proc_root: Arc::clone(&self.proc_root), - caches: Arc::clone(&self.caches), - slot, - _held: held, - }; - tasks.spawn(handle_lease(ctx, grant)); - started += 1; + if self.start_lease(grant, tasks) { + // The lease now lives in `active_leases`; the JIT config is + // single use, so a restart must not find it again. + self.standbys.remove(&id); + started += 1; + } } - Ok(started) + started + } + + /// Start every held standby because the queen has been unreachable. + fn fail_open(&self, tasks: &mut JoinSet<()>) { + if self.standbys.ids().is_empty() { + return; + } + let started = self.start_standbys(None, tasks); + if started > 0 { + tracing::warn!("queen unreachable: started {started} standby deploy runner(s)"); + } + } + + /// Claim a slot for `grant` and run it as a lease. Returns whether it + /// started. + fn start_lease(&self, grant: LeaseGrant, tasks: &mut JoinSet<()>) -> bool { + if !self.admit(&grant.lease_id) { + return false; + } + let held = HeldGuard { + held: Arc::clone(&self.held), + preempts: Arc::clone(&self.preempts), + id: grant.lease_id.clone(), + }; + let (preempt_tx, preempt) = tokio::sync::watch::channel(false); + if let Ok(mut preempts) = self.preempts.lock() { + preempts.insert(grant.lease_id.clone(), preempt_tx); + } + let Some(slot) = self.slots.acquire(&grant.lease_id) else { + tracing::warn!( + "no free slot for lease {}; leaving it unstarted", + grant.lease_id + ); + return false; + }; + let ctx = LeaseContext { + cfg: Arc::clone(&self.cfg), + client: Arc::clone(&self.client), + exec: Arc::clone(&self.exec), + job_exec: Arc::clone(&self.job_exec), + fs_ops: Arc::clone(&self.fs_ops), + abort: self.abort.subscribe(), + preempt, + proc_root: Arc::clone(&self.proc_root), + caches: Arc::clone(&self.caches), + slot, + _held: held, + }; + tasks.spawn(handle_lease(ctx, grant)); + true } /// What a successful heartbeat means for self-update: the queen has the @@ -339,14 +437,22 @@ impl Drone { // arrives, and heartbeats are retried with backoff. let mut failures = 0u32; let mut retry_at = tokio::time::Instant::now(); + let mut last_ok_at = time::OffsetDateTime::now_utc().unix_timestamp(); loop { tokio::select! { _ = every.tick() => { + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + if should_fail_open(last_ok_at, now, self.cfg.fail_open_after_secs) { + self.fail_open(&mut tasks); + } if tokio::time::Instant::now() < retry_at { continue; } match self.heartbeat_once(true, &mut tasks).await { - Ok(_) => failures = 0, + Ok(_) => { + failures = 0; + last_ok_at = time::OffsetDateTime::now_utc().unix_timestamp(); + } Err(e) => { failures = failures.saturating_add(1); let wait = selfupdate::backoff(self.cfg.heartbeat_secs, failures); @@ -393,6 +499,7 @@ struct LeaseContext { job_exec: Arc, fs_ops: Arc, abort: tokio::sync::watch::Receiver, + preempt: tokio::sync::watch::Receiver, proc_root: Arc, caches: Arc, /// Dropped before `_held`, so the slot is free whenever the id is. @@ -419,6 +526,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { cache_root: &ctx.cfg.cache_root, extra_env: &extra_env, max_job: Duration::from_secs(ctx.cfg.max_job_secs), + preempt_grace: Duration::from_secs(ctx.cfg.preempt_grace_secs), job_user: ctx.cfg.job_user.as_ref(), proc_root: &ctx.proc_root, }; @@ -427,6 +535,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { ctx.job_exec.as_ref(), Arc::clone(&ctx.fs_ops), ctx.abort.clone(), + ctx.preempt.clone(), ) .await; @@ -472,6 +581,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { duration_secs: o.duration_secs, peak_rss_mib: o.peak_rss_mib, telemetry: None, + preempted: o.preempted, }, Err(e) => { tracing::warn!("lease {} could not run: {e}", grant.lease_id); @@ -480,6 +590,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { duration_secs: 0, peak_rss_mib: None, telemetry: None, + preempted: false, } } }; diff --git a/src/drone/runner.rs b/src/drone/runner.rs index 5355d4f..f3654fe 100644 --- a/src/drone/runner.rs +++ b/src/drone/runner.rs @@ -37,6 +37,9 @@ pub struct RunSpec<'a> { /// Extra job environment (cache directories, git mirrors). pub extra_env: &'a [(String, String)], pub max_job: Duration, + /// After a preempt signal the runner gets SIGINT, then this long before + /// the process group is killed. + pub preempt_grace: Duration, /// Run the job as this user (see [`crate::drone::jobuser`]); `None` runs it /// as the drone's own user. pub job_user: Option<&'a JobUser>, @@ -51,6 +54,8 @@ pub struct RunOutcome { pub peak_rss_mib: Option, pub timed_out: bool, pub aborted: bool, + /// The queen preempted the lease; the runner was interrupted. + pub preempted: bool, } /// Exit code reported when a job was killed (128 + SIGKILL). @@ -132,6 +137,11 @@ async fn spawn_retrying(command: &mut Command) -> io::Result, mut abort: tokio::sync::watch::Receiver, + mut preempt: tokio::sync::watch::Receiver, ) -> io::Result { if !valid_lease_id(spec.lease_id) { return Err(io::Error::new( @@ -216,6 +227,9 @@ pub async fn run_lease( tokio::pin!(deadline); let mut sample = tokio::time::interval(RSS_SAMPLE_EVERY); let (mut peak_kib, mut timed_out, mut aborted) = (0u64, false, false); + let mut preempted = false; + let grace = tokio::time::sleep(spec.preempt_grace); + tokio::pin!(grace); // Every arm is cancel-safe: `Child::wait`, `Sleep`, `Interval::tick` and // `watch::Receiver::changed` keep no partially consumed state. @@ -227,6 +241,17 @@ pub async fn run_lease( if let Some(pgid) = pgid { kill_group(exec, pgid); } break child.wait().await?; } + changed = preempt.changed(), if !preempted => { + if changed.is_ok() && *preempt.borrow() { + preempted = true; + grace.as_mut().reset(tokio::time::Instant::now() + spec.preempt_grace); + if let Some(pgid) = pgid { interrupt_group(exec, pgid); } + } + } + () = &mut grace, if preempted => { + if let Some(pgid) = pgid { kill_group(exec, pgid); } + break child.wait().await?; + } changed = abort.changed() => { if changed.is_ok() && *abort.borrow() { aborted = true; @@ -255,6 +280,7 @@ pub async fn run_lease( peak_rss_mib: (peak_kib > 0).then_some(peak_kib / 1024), timed_out, aborted, + preempted, }) } @@ -289,13 +315,16 @@ mod tests { extra_env: &[], job_user: None, max_job, + preempt_grace: Duration::from_secs(5), proc_root: Path::new("/proc"), }; + let (_ptx, prx) = tokio::sync::watch::channel(false); let out = run_lease( &spec, &SystemExec, Arc::new(crate::drone::fsops::LocalFs), rx, + prx, ) .await .unwrap(); @@ -405,8 +434,10 @@ mod tests { extra_env: &[], job_user: None, max_job: Duration::from_secs(60), + preempt_grace: Duration::from_secs(5), proc_root: Path::new("/proc"), }; + let (_ptx, prx) = tokio::sync::watch::channel(false); let abort = tokio::spawn(async move { tokio::time::sleep(Duration::from_millis(300)).await; let _ = tx.send(true); @@ -417,6 +448,7 @@ mod tests { &SystemExec, Arc::new(crate::drone::fsops::LocalFs), rx, + prx, ) .await .unwrap(); @@ -424,6 +456,49 @@ mod tests { assert!(out.aborted && !out.timed_out); } + #[tokio::test] + async fn preempt_interrupts_the_runner_then_kills_it_after_the_grace_period() { + let dir = tempfile::tempdir().unwrap(); + // Ignores SIGINT, so only the grace-period kill can end it. + let runner = runner_with_script(dir.path(), "trap '' INT\necho up > up.txt\nsleep 300"); + let work = dir.path().join("work"); + let (_tx, rx) = tokio::sync::watch::channel(false); + let (ptx, prx) = tokio::sync::watch::channel(false); + let secret = Secret::new("x"); + let cache = dir.path().join("cache"); + let spec = RunSpec { + lease_id: "lease3", + jit_config: &secret, + runner_dir: &runner, + lease_dir: &work.join("slot-0"), + cache_root: &cache, + extra_env: &[], + job_user: None, + max_job: Duration::from_secs(60), + preempt_grace: Duration::from_millis(500), + proc_root: Path::new("/proc"), + }; + let stop = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(300)).await; + let _ = ptx.send(true); + ptx + }); + let started = Instant::now(); + let out = run_lease( + &spec, + &SystemExec, + Arc::new(crate::drone::fsops::LocalFs), + rx, + prx, + ) + .await + .unwrap(); + let _ptx = stop.await.unwrap(); + assert!(out.preempted && !out.timed_out && !out.aborted); + assert_eq!(out.exit_code, KILLED_EXIT_CODE); + assert!(started.elapsed() < Duration::from_secs(30)); + } + #[tokio::test] async fn unsafe_lease_ids_never_touch_the_filesystem() { let dir = tempfile::tempdir().unwrap(); @@ -439,14 +514,17 @@ mod tests { extra_env: &[], job_user: None, max_job: Duration::from_secs(1), + preempt_grace: Duration::from_secs(1), proc_root: Path::new("/proc"), }; + let (_ptx, prx) = tokio::sync::watch::channel(false); assert!( run_lease( &spec, &SystemExec, Arc::new(crate::drone::fsops::LocalFs), - rx + rx, + prx ) .await .is_err() diff --git a/src/drone/standby.rs b/src/drone/standby.rs new file mode 100644 index 0000000..8c10ab3 --- /dev/null +++ b/src/drone/standby.rs @@ -0,0 +1,184 @@ +//! Standby deploy runners the drone holds unstarted. +//! +//! The queen mints a standby runner for a deploy drone and sends its JIT +//! config. The drone keeps it in memory and, when `state_dir` is set, in a +//! `0600` file so a restart while the queen is down still has it. The drone +//! starts it when the queen says so, or when the queen has been unreachable +//! for `fail_open_after_secs` ([`should_fail_open`]). + +use std::collections::BTreeMap; +use std::io::Write; +use std::os::unix::fs::OpenOptionsExt; +use std::path::PathBuf; +use std::sync::Mutex; + +use crate::api::LeaseGrant; +use crate::drone::cleanup::valid_lease_id; + +/// Directory under `state_dir` holding one `.json` per standby. +const DIR: &str = "standby"; + +/// Whether the queen has been silent long enough to start the held standby +/// runners. `after_secs == 0` turns fail-open off. +pub fn should_fail_open(last_ok_at: i64, now: i64, after_secs: u64) -> bool { + after_secs > 0 + && now.saturating_sub(last_ok_at) >= i64::try_from(after_secs).unwrap_or(i64::MAX) +} + +#[derive(Default)] +pub struct Standbys { + dir: Option, + held: Mutex>, +} + +impl Standbys { + /// Load the standbys persisted under `state_dir`; unreadable files are + /// skipped (the queen resends what the drone does not report). + pub fn load(state_dir: Option) -> Self { + let dir = state_dir.map(|d| d.join(DIR)); + let mut held = BTreeMap::new(); + if let Some(dir) = &dir + && let Ok(entries) = std::fs::read_dir(dir) + { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().is_none_or(|e| e != "json") { + continue; + } + let grant = std::fs::read(&path) + .ok() + .and_then(|b| serde_json::from_slice::(&b).ok()) + .filter(|g| valid_lease_id(&g.lease_id)); + match grant { + Some(g) => { + held.insert(g.lease_id.clone(), g); + } + None => { + tracing::warn!("discarding unreadable standby file {}", path.display()); + let _ = std::fs::remove_file(&path); + } + } + } + } + Self { + dir, + held: Mutex::new(held), + } + } + + pub fn ids(&self) -> Vec { + self.held + .lock() + .map(|h| h.keys().cloned().collect()) + .unwrap_or_default() + } + + pub fn get(&self, id: &str) -> Option { + self.held.lock().ok()?.get(id).cloned() + } + + /// Hold `grant`. Returns false (holding nothing) for an unsafe id. + pub fn insert(&self, grant: LeaseGrant) -> bool { + if !valid_lease_id(&grant.lease_id) { + return false; + } + if let Err(e) = self.persist(&grant) { + tracing::warn!("persist standby {}: {e}", grant.lease_id); + } + if let Ok(mut h) = self.held.lock() { + h.insert(grant.lease_id.clone(), grant); + } + true + } + + /// Forget a standby (started, or ended on the queen's side) and delete its + /// file. + pub fn remove(&self, id: &str) { + if let Ok(mut h) = self.held.lock() { + h.remove(id); + } + if let Some(dir) = &self.dir + && valid_lease_id(id) + { + let _ = std::fs::remove_file(dir.join(format!("{id}.json"))); + } + } + + fn persist(&self, grant: &LeaseGrant) -> std::io::Result<()> { + let Some(dir) = &self.dir else { + return Ok(()); + }; + std::fs::create_dir_all(dir)?; + let tmp = dir.join(format!("{}.tmp", grant.lease_id)); + let bytes = serde_json::to_vec(grant).map_err(std::io::Error::other)?; + let mut f = std::fs::OpenOptions::new() + .write(true) + .create(true) + .truncate(true) + .mode(0o600) + .open(&tmp)?; + f.write_all(&bytes)?; + f.sync_all()?; + std::fs::rename(&tmp, dir.join(format!("{}.json", grant.lease_id))) + } +} + +#[cfg(test)] +mod tests { + use std::os::unix::fs::PermissionsExt; + + use super::*; + use crate::classes::{Class, Scope}; + use crate::secret::Secret; + + fn grant(id: &str) -> LeaseGrant { + LeaseGrant { + lease_id: id.into(), + class: Class::Medium, + scope: Scope::Deploy, + labels: vec!["self-hosted".into()], + runner_name: format!("borg-host-a-{id}"), + jit_config: Secret::new("jit"), + expires_at: i64::MAX, + } + } + + #[test] + fn fail_open_waits_for_the_configured_silence() { + assert!(!should_fail_open(1000, 1119, 120)); + assert!(should_fail_open(1000, 1120, 120)); + assert!(should_fail_open(1000, 5000, 120)); + // 0 disables fail-open however long the queen is silent. + assert!(!should_fail_open(0, i64::MAX, 0)); + // A clock that went backwards does not trigger it. + assert!(!should_fail_open(2000, 1000, 120)); + } + + #[test] + fn standbys_survive_a_restart_in_a_private_file_until_removed() { + let dir = tempfile::tempdir().unwrap(); + let a = Standbys::load(Some(dir.path().to_owned())); + assert!(a.insert(grant("abc123"))); + assert!(!a.insert(grant("../escape"))); + let file = dir.path().join("standby/abc123.json"); + let mode = std::fs::metadata(&file).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o600); + + let b = Standbys::load(Some(dir.path().to_owned())); + assert_eq!(b.ids(), vec!["abc123".to_owned()]); + assert_eq!(b.get("abc123").unwrap().runner_name, "borg-host-a-abc123"); + + b.remove("abc123"); + assert!(b.ids().is_empty() && !file.exists()); + assert!(Standbys::load(Some(dir.path().to_owned())).ids().is_empty()); + } + + #[test] + fn without_a_state_dir_standbys_live_in_memory_only() { + let s = Standbys::load(None); + assert!(s.insert(grant("m1"))); + assert_eq!(s.ids(), vec!["m1".to_owned()]); + s.remove("m1"); + assert!(s.ids().is_empty()); + } +} diff --git a/src/envcontract.rs b/src/envcontract.rs index 62746c8..2ea4612 100644 --- a/src/envcontract.rs +++ b/src/envcontract.rs @@ -1,7 +1,7 @@ //! The build environment contract: facts a drone reports about its build //! environment, and the queen's pure admission check over them. //! -//! Before work (a CI runner lease or a work unit) is placed on a drone, four +//! Before work (a CI runner lease or a work unit) is placed on a drone, five //! checks are evaluated against what the drone last reported: //! //! 1. **compiler_cache**: `RUSTC_WRAPPER` resolves to an allowed wrapper @@ -15,6 +15,9 @@ //! 4. **resources**: memory available now covers the class's memory, and live //! CPU pressure (PSI) is under a ceiling. (Free capacity after active leases //! is enforced by the allocator itself.) +//! 5. **tools**: every tool the operator lists as required (globally or for the +//! class) resolved on the job user's `PATH` when the drone probed it. The +//! list is empty unless the operator writes one. //! //! The drone is a dumb prober: it reports *facts* ([`EnvFacts`]), every one of //! which may be unknown (`None`), cached about a minute on its side. Thresholds @@ -41,14 +44,25 @@ pub enum CheckKind { Toolchain, CacheStore, Resources, + Tools, } +/// The tools a drone looks for on the job user's `PATH`, in report order. +pub const PROBE_TOOLS: [&str; 15] = [ + "git", "gh", "node", "npm", "pnpm", "python3", "docker", "cargo", "rustc", "bazel", "go", "uv", + "mise", "jq", "cosign", +]; + +/// The value reported for a probed tool that is not on `PATH`. +pub const TOOL_MISSING: &str = "missing"; + impl CheckKind { - pub const ALL: [Self; 4] = [ + pub const ALL: [Self; 5] = [ Self::CompilerCache, Self::Toolchain, Self::CacheStore, Self::Resources, + Self::Tools, ]; pub fn as_str(self) -> &'static str { @@ -57,6 +71,7 @@ impl CheckKind { Self::Toolchain => "toolchain", Self::CacheStore => "cache_store", Self::Resources => "resources", + Self::Tools => "tools", } } } @@ -142,6 +157,10 @@ pub struct EnvFacts { /// where PSI does not exist. #[serde(default)] pub cpu_psi_some_avg10_pm: Option, + /// For each tool in [`PROBE_TOOLS`], the first line of its `--version` + /// output, or `"missing"`; `None` when the drone did not measure it. + #[serde(default)] + pub tools: Option>, } /// Settings of one class (or the defaults for all). @@ -152,6 +171,8 @@ pub struct ClassContract { pub checks: BTreeMap, /// The Rust toolchain every build of this class needs. pub rust_toolchain: Option, + /// Tools this class needs on `PATH`; replaces the global list. + pub required_tools: Option>, } #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] @@ -177,6 +198,9 @@ pub struct ContractSettings { pub max_cpu_psi_pm: u16, /// Memory available must be at least this percent of the class's memory. pub mem_available_pct: u64, + /// Tools (from [`PROBE_TOOLS`]) every class needs on `PATH`. Empty means + /// the `tools` check has nothing to require. + pub required_tools: Vec, pub checks: BTreeMap, pub classes: BTreeMap, } @@ -194,6 +218,7 @@ impl Default for ContractSettings { min_requests_for_hit_rate: 200, max_cpu_psi_pm: 700, mem_available_pct: 100, + required_tools: Vec::new(), // Conservative: only the checks that protect a build from failing // outright block by default; the cache checks report. checks: BTreeMap::from([ @@ -201,6 +226,7 @@ impl Default for ContractSettings { (CheckKind::Toolchain, Mode::Required), (CheckKind::CacheStore, Mode::Advisory), (CheckKind::Resources, Mode::Required), + (CheckKind::Tools, Mode::Advisory), ]), classes: BTreeMap::new(), } @@ -238,9 +264,33 @@ impl ContractSettings { "env_contract rust_toolchain must look like 1.95.0 or nightly-2026-01-01".into(), ); } + let unprobed = |t: &String| !PROBE_TOOLS.contains(&t.as_str()); + if self + .required_tools + .iter() + .chain( + self.classes + .values() + .flat_map(|c| c.required_tools.iter().flatten()), + ) + .any(unprobed) + { + return Err(format!( + "env_contract required_tools must be among: {}", + PROBE_TOOLS.join(", ") + )); + } Ok(()) } + /// The tools a class needs on `PATH`: its own list, else the global one. + pub fn class_required_tools(&self, class: Class) -> &[String] { + self.classes + .get(&class) + .and_then(|c| c.required_tools.as_deref()) + .unwrap_or(&self.required_tools) + } + fn mode(&self, class: Class, kind: CheckKind) -> Mode { self.classes .get(&class) @@ -507,6 +557,51 @@ fn resources(s: &ContractSettings, f: &EnvFacts, need: &Need<'_>) -> Acc { a } +fn tools(s: &ContractSettings, f: &EnvFacts, need: &Need<'_>) -> Acc { + let mut a = Acc::new(); + let wanted = s.class_required_tools(need.class); + a.see("required", wanted.join(",")); + let Some(have) = &f.tools else { + if wanted.is_empty() { + return a; + } + if s.require_report { + a.note(CheckState::Fail, "the drone did not report its tools"); + } else { + a.warnings + .push("the drone did not report its tools".to_owned()); + } + return a; + }; + a.see( + "versions", + have.iter() + .map(|(k, v)| format!("{k}={v}")) + .collect::>() + .join("; "), + ); + a.see( + "missing", + have.iter() + .filter(|(_, v)| *v == TOOL_MISSING) + .map(|(k, _)| k.as_str()) + .collect::>() + .join(","), + ); + let missing: Vec<&str> = wanted + .iter() + .filter(|t| have.get(*t).is_none_or(|v| v == TOOL_MISSING)) + .map(String::as_str) + .collect(); + if !missing.is_empty() { + a.note( + CheckState::Fail, + format!("required tools missing on PATH: {}", missing.join(", ")), + ); + } + a +} + /// Evaluate the contract for `need` against `facts` at `now`. Pure. pub fn evaluate( s: &ContractSettings, @@ -565,6 +660,7 @@ pub fn evaluate( CheckKind::Toolchain => toolchain(f, need), CheckKind::CacheStore => cache_store(s, f), CheckKind::Resources => resources(s, f, need), + CheckKind::Tools => tools(s, f, need), } }; checks.push(acc.finish(kind, mode)); @@ -634,6 +730,25 @@ mod tests { }]), mem_available_mib: Some(20_000), cpu_psi_some_avg10_pm: Some(100), + tools: Some(tool_map(&[ + ("git", "git version 2.45.0"), + ("gh", "gh version 2.50.0"), + ])), + } + } + + fn tool_map(pairs: &[(&str, &str)]) -> BTreeMap { + pairs + .iter() + .map(|(k, v)| ((*k).to_owned(), (*v).to_owned())) + .collect() + } + + fn tools_on(mode: Mode, required: &[&str]) -> ContractSettings { + ContractSettings { + checks: BTreeMap::from([(CheckKind::Tools, mode)]), + required_tools: required.iter().map(|t| (*t).to_owned()).collect(), + ..on() } } @@ -1021,4 +1136,148 @@ mod tests { evaluate(&all_required(), Some(&good()), &n, NOW) ); } + + #[test] + fn a_missing_required_tool_is_not_admitted_and_named() { + let s = tools_on(Mode::Required, &["git", "gh", "node"]); + let f = EnvFacts { + tools: Some(tool_map(&[ + ("git", "git version 2.45.0"), + ("gh", "missing"), + ("node", "missing"), + ("jq", "missing"), + ])), + ..good() + }; + let v = evaluate(&s, Some(&f), &need(Class::Light, None), NOW); + assert!(!v.admit); + let c = check(&v, CheckKind::Tools); + assert_eq!(c.state, CheckState::Fail); + assert!(c.reason.contains("gh, node"), "{}", c.reason); + assert!(!c.reason.contains("git"), "{}", c.reason); + assert_eq!(c.observed["missing"], "gh,jq,node"); + assert!(c.observed["versions"].contains("git=git version 2.45.0")); + // A tool that is not in the map at all counts as missing. + let f = EnvFacts { + tools: Some(tool_map(&[("git", "2")])), + ..good() + }; + let v = evaluate(&s, Some(&f), &need(Class::Light, None), NOW); + assert!(check(&v, CheckKind::Tools).reason.contains("gh, node")); + // Everything present: admitted. + let f = EnvFacts { + tools: Some(tool_map(&[("git", "2"), ("gh", "2"), ("node", "22")])), + ..good() + }; + assert!(evaluate(&s, Some(&f), &need(Class::Light, None), NOW).admit); + } + + #[test] + fn advisory_tools_warn_in_the_verdict_but_admit() { + let s = tools_on(Mode::Advisory, &["gh"]); + let f = EnvFacts { + tools: Some(tool_map(&[("gh", "missing")])), + ..good() + }; + let v = evaluate(&s, Some(&f), &need(Class::Light, None), NOW); + assert!(v.admit, "{v:?}"); + let c = check(&v, CheckKind::Tools); + assert_eq!((c.state, c.mode), (CheckState::Fail, Mode::Advisory)); + assert!(c.reason.contains("gh")); + } + + #[test] + fn the_default_requires_nothing_and_a_class_list_replaces_the_global_one() { + // Enabled with defaults: an empty list passes even with tools missing. + let f = EnvFacts { + tools: Some(tool_map(&[("gh", "missing")])), + ..good() + }; + let v = evaluate(&on(), Some(&f), &need(Class::Light, None), NOW); + assert_eq!(check(&v, CheckKind::Tools).mode, Mode::Advisory); + assert_eq!(check(&v, CheckKind::Tools).state, CheckState::Pass); + + let mut s = tools_on(Mode::Required, &["gh"]); + s.classes.insert( + Class::Heavy, + ClassContract { + required_tools: Some(vec!["docker".into()]), + ..ClassContract::default() + }, + ); + s.classes.insert( + Class::Medium, + ClassContract { + required_tools: Some(Vec::new()), + ..ClassContract::default() + }, + ); + let f = EnvFacts { + tools: Some(tool_map(&[("gh", "missing"), ("docker", "27.0")])), + ..good() + }; + let admit = |c| evaluate(&s, Some(&f), &need(c, None), NOW).admit; + assert!(!admit(Class::Light)); // global gh + assert!(admit(Class::Heavy)); // docker only + assert!(admit(Class::Medium)); // explicitly nothing + let off = ContractSettings { + classes: BTreeMap::from([( + Class::Light, + ClassContract { + checks: BTreeMap::from([(CheckKind::Tools, Mode::Off)]), + ..ClassContract::default() + }, + )]), + ..s.clone() + }; + let v = evaluate(&off, Some(&f), &need(Class::Light, None), NOW); + assert!(v.admit && v.checks.iter().all(|c| c.kind != CheckKind::Tools)); + } + + #[test] + fn an_absent_tools_map_fails_only_with_require_report() { + let f = EnvFacts { + tools: None, + ..good() + }; + let s = tools_on(Mode::Required, &["gh"]); + let v = evaluate(&s, Some(&f), &need(Class::Light, None), NOW); + assert!(v.admit, "{v:?}"); + let c = check(&v, CheckKind::Tools); + assert_eq!(c.state, CheckState::Pass); + assert_eq!(c.warnings.len(), 1); + + let strict = ContractSettings { + require_report: true, + ..s.clone() + }; + let v = evaluate(&strict, Some(&f), &need(Class::Light, None), NOW); + assert!(!v.admit); + assert!( + check(&v, CheckKind::Tools) + .reason + .contains("did not report its tools") + ); + // Nothing required: an absent map is fine even with require_report. + let none = ContractSettings { + required_tools: Vec::new(), + ..strict + }; + assert!(evaluate(&none, Some(&f), &need(Class::Light, None), NOW).admit); + } + + #[test] + fn required_tools_must_be_tools_the_drone_probes() { + assert!(tools_on(Mode::Required, &["git", "jq"]).validate().is_ok()); + assert!(tools_on(Mode::Required, &["ffmpeg"]).validate().is_err()); + let mut s = on(); + s.classes.insert( + Class::Light, + ClassContract { + required_tools: Some(vec!["nope".into()]), + ..ClassContract::default() + }, + ); + assert!(s.validate().is_err()); + } } diff --git a/src/janitor.rs b/src/janitor.rs index 46df700..41c7041 100644 --- a/src/janitor.rs +++ b/src/janitor.rs @@ -2,15 +2,20 @@ //! commit is no longer the pull request's head. //! //! A push to a pull request branch makes every run of the older commit -//! pointless, but GitHub keeps them queued and they occupy runner demand. The +//! pointless, but GitHub keeps them queued and they occupy runner demand. Two +//! further opt-ins remove other waste: `cancel_running` also cancels +//! superseded runs that already started, and `dedupe_same_sha` cancels queued +//! duplicates of a run of the same workflow, commit, event and pull request. +//! Duplicate groups never include `workflow_dispatch` or `merge_group` runs, +//! and a running duplicate is never cancelled. The //! janitor is opt-in (`supersede_janitor` in the queen config) and uses its own //! token, because cancelling runs needs `actions: write` and the queen's App //! deliberately does not have it. //! //! Safety rules, all enforced here: -//! * only runs that are `queued` are candidates, and the run is fetched again -//! immediately before cancelling, so a run that started in between is left -//! alone; +//! * only runs that are `queued` are candidates (`in_progress` ones too when +//! `cancel_running` is set), and the run is fetched again immediately before +//! cancelling, so a run whose status changed in between is left alone; //! * only `pull_request` runs that name exactly one pull request are touched; //! runs from forks (no pull request attached) and merge groups never are; //! * the pull request's current head comes from the pulls API, not from the @@ -20,12 +25,13 @@ //! * a cancel that finds the run already finished (409) is a no-op; //! * every cancellation (or, in dry-run, every would-be cancellation) is logged. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Mutex; use std::time::Duration; use reqwest::{Method, RequestBuilder}; -use serde::Deserialize; use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; use crate::config::JanitorSettings; use crate::github::GithubError; @@ -35,17 +41,24 @@ use crate::secret::Secret; const MAX_PAGES: u32 = 5; /// What one pass did. -#[derive(Debug, Default, Clone, PartialEq, Eq)] +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)] pub struct PassReport { - /// Queued pull-request runs examined. + /// Pull-request runs examined (queued, plus running with `cancel_running`). pub considered: u32, /// Runs whose head is no longer the pull request's head. pub superseded: u32, - /// Runs cancelled. + /// Queued superseded runs cancelled. pub cancelled: u32, - /// Runs that dry-run would have cancelled. + /// Running superseded runs cancelled (`cancel_running`). + pub cancelled_running: u32, + /// Runs beyond the first in a same-workflow, same-commit group + /// (`dedupe_same_sha`), queued or running. + pub duplicates: u32, + /// Queued duplicates cancelled. + pub dedup_cancelled: u32, + /// Runs that dry-run would have cancelled, of any kind. pub would_cancel: u32, - /// Superseded runs left for a later pass by the per-pass cap. + /// Runs left for a later pass by the per-pass cap. pub deferred: u32, /// Runs that changed state before the cancel and were left alone. pub raced: u32, @@ -53,9 +66,35 @@ pub struct PassReport { pub errors: u32, } +impl PassReport { + fn add(&mut self, o: &PassReport) { + self.considered += o.considered; + self.superseded += o.superseded; + self.cancelled += o.cancelled; + self.cancelled_running += o.cancelled_running; + self.duplicates += o.duplicates; + self.dedup_cancelled += o.dedup_cancelled; + self.would_cancel += o.would_cancel; + self.deferred += o.deferred; + self.raced += o.raced; + self.errors += o.errors; + } +} + +/// The janitor as `/status` and `/metrics` show it. +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)] +pub struct JanitorView { + pub passes: u64, + pub last: PassReport, + /// Sums over every pass since the queen started. + pub total: PassReport, +} + #[derive(Debug, Deserialize)] struct Run { id: u64, + #[serde(default)] + workflow_id: u64, status: String, event: String, head_sha: String, @@ -83,6 +122,7 @@ pub struct Janitor { api_base: String, token: Secret, settings: JanitorSettings, + view: Mutex, } impl Janitor { @@ -100,9 +140,14 @@ impl Janitor { api_base: api_base.trim_end_matches('/').to_owned(), token, settings, + view: Mutex::default(), }) } + pub fn view(&self) -> JanitorView { + self.view.lock().expect("janitor view lock").clone() + } + pub fn settings(&self) -> &JanitorSettings { &self.settings } @@ -138,16 +183,22 @@ impl Janitor { } } - async fn queued_pr_runs(&self, repo: &str) -> Result, GithubError> { + async fn list_runs( + &self, + repo: &str, + status: &str, + event: Option<&str>, + ) -> Result, GithubError> { #[derive(Deserialize)] struct Runs { workflow_runs: Vec, } + let event = event.map_or_else(String::new, |e| format!("&event={e}")); let mut out = Vec::new(); for page in 1..=MAX_PAGES { let runs: Runs = self .get(&format!( - "/repos/{repo}/actions/runs?status=queued&event=pull_request&per_page=100&page={page}" + "/repos/{repo}/actions/runs?status={status}{event}&per_page=100&page={page}" )) .await?; let n = runs.workflow_runs.len(); @@ -170,6 +221,24 @@ impl Janitor { report.errors += 1; } } + tracing::debug!( + "janitor pass: considered {} superseded {} cancelled {} cancelled_running {} duplicates {} dedup_cancelled {} would_cancel {} deferred {} raced {} errors {}", + report.considered, + report.superseded, + report.cancelled, + report.cancelled_running, + report.duplicates, + report.dedup_cancelled, + report.would_cancel, + report.deferred, + report.raced, + report.errors + ); + let mut view = self.view.lock().expect("janitor view lock"); + view.passes += 1; + view.total.add(&report); + view.last = report.clone(); + drop(view); report } @@ -179,10 +248,22 @@ impl Janitor { budget: &mut u32, report: &mut PassReport, ) -> Result<(), GithubError> { - let runs = self.queued_pr_runs(repo).await?; + let st = &self.settings; + let event = (!st.dedupe_same_sha).then_some("pull_request"); + let mut runs = self.list_runs(repo, "queued", event).await?; + if st.cancel_running || st.dedupe_same_sha { + runs.extend(self.list_runs(repo, "in_progress", event).await?); + // A run that started between the two listings appears in both. + runs.sort_by_key(|r| r.id); + runs.dedup_by_key(|r| r.id); + } + let mut stale = BTreeSet::new(); let mut heads: BTreeMap> = BTreeMap::new(); for run in &runs { - if run.status != "queued" || run.event != "pull_request" { + let running = run.status == "in_progress"; + if run.event != "pull_request" + || !(run.status == "queued" || running && st.cancel_running) + { continue; } // A fork's run carries no pull request; with more than one the @@ -215,57 +296,63 @@ impl Janitor { continue; } report.superseded += 1; + stale.insert(run.id); if *budget == 0 { report.deferred += 1; continue; } - match self.supersede(repo, run, pr.number, ¤t).await { - Ok(Outcome::Cancelled) => { - *budget -= 1; - report.cancelled += 1; - } - Ok(Outcome::DryRun) => { - *budget -= 1; - report.would_cancel += 1; + let why = format!( + "for #{}: head {} superseded by {current}", + pr.number, run.head_sha + ); + let kind = if running { Kind::Running } else { Kind::Stale }; + let outcome = self.cancel_if_unchanged(repo, run, &why).await; + tally(report, budget, repo, run.id, kind, outcome); + } + if st.dedupe_same_sha { + for (run, kept) in surplus(&runs, &stale) { + report.duplicates += 1; + if run.status != "queued" { + continue; } - Ok(Outcome::Raced) => report.raced += 1, - Err(e) => { - tracing::warn!("janitor: {repo} run {}: {e}", run.id); - report.errors += 1; + if *budget == 0 { + report.deferred += 1; + continue; } + let why = format!( + "duplicate of run {kept} (workflow {}, head {})", + run.workflow_id, run.head_sha + ); + let outcome = self.cancel_if_unchanged(repo, run, &why).await; + tally(report, budget, repo, run.id, Kind::Duplicate, outcome); } } Ok(()) } - async fn supersede( + /// Cancel `run` unless a fresh read shows its status or head changed. + async fn cancel_if_unchanged( &self, repo: &str, run: &Run, - pr: u64, - current: &str, + why: &str, ) -> Result { - // Re-read: the run may have started since the list was fetched. let fresh: Run = self .get(&format!("/repos/{repo}/actions/runs/{}", run.id)) .await?; - if fresh.status != "queued" || fresh.head_sha != run.head_sha { + if fresh.status != run.status || fresh.head_sha != run.head_sha { return Ok(Outcome::Raced); } + let state = &run.status; if self.settings.dry_run { tracing::info!( - "janitor (dry-run): would cancel {repo} run {} for #{pr}: head {} superseded by {current}", - run.id, - run.head_sha + "janitor (dry-run): would cancel {state} {repo} run {} {why}", + run.id ); return Ok(Outcome::DryRun); } if self.cancel(repo, run.id).await? { - tracing::info!( - "janitor: cancelled {repo} run {} for #{pr}: head {} superseded by {current}", - run.id, - run.head_sha - ); + tracing::info!("janitor: cancelled {state} {repo} run {} {why}", run.id); Ok(Outcome::Cancelled) } else { Ok(Outcome::Raced) @@ -273,13 +360,79 @@ impl Janitor { } } +#[derive(Clone, Copy)] +enum Kind { + Stale, + Running, + Duplicate, +} + +fn tally( + report: &mut PassReport, + budget: &mut u32, + repo: &str, + run: u64, + kind: Kind, + outcome: Result, +) { + match outcome { + Ok(Outcome::Cancelled) => { + *budget -= 1; + match kind { + Kind::Stale => report.cancelled += 1, + Kind::Running => report.cancelled_running += 1, + Kind::Duplicate => report.dedup_cancelled += 1, + } + } + Ok(Outcome::DryRun) => { + *budget -= 1; + report.would_cancel += 1; + } + Ok(Outcome::Raced) => report.raced += 1, + Err(e) => { + tracing::warn!("janitor: {repo} run {run}: {e}"); + report.errors += 1; + } + } +} + +/// Runs to drop from each group of queued or running runs sharing workflow, +/// head commit, event and pull requests, each with the id of the run kept: the +/// lowest-id running run, else the lowest-id queued run. `workflow_dispatch` +/// and `merge_group` runs are never grouped, and `skip` ids are left out. +fn surplus<'a>(runs: &'a [Run], skip: &BTreeSet) -> Vec<(&'a Run, u64)> { + type Key<'a> = (u64, &'a str, &'a str, Vec); + let mut groups: BTreeMap, Vec<&'a Run>> = BTreeMap::new(); + for run in runs { + if skip.contains(&run.id) + || matches!(run.event.as_str(), "workflow_dispatch" | "merge_group") + || !matches!(run.status.as_str(), "queued" | "in_progress") + { + continue; + } + let mut prs: Vec = run.pull_requests.iter().map(|p| p.number).collect(); + prs.sort_unstable(); + groups + .entry((run.workflow_id, &run.head_sha, &run.event, prs)) + .or_default() + .push(run); + } + let mut out = Vec::new(); + for mut group in groups.into_values() { + group.sort_by_key(|r| (r.status != "in_progress", r.id)); + let kept = group[0].id; + out.extend(group.into_iter().skip(1).map(|r| (r, kept))); + } + out +} + enum Outcome { Cancelled, DryRun, Raced, } -async fn status_error( +pub(crate) async fn status_error( method: &'static str, path: &str, status: u16, @@ -298,3 +451,63 @@ async fn status_error( .collect(), } } + +#[cfg(test)] +mod tests { + use super::*; + + fn run(id: u64, wf: u64, status: &str, event: &str, sha: &str, prs: &[u64]) -> Run { + Run { + id, + workflow_id: wf, + status: status.into(), + event: event.into(), + head_sha: sha.into(), + pull_requests: prs.iter().map(|&number| PullRef { number }).collect(), + } + } + + fn picked(runs: &[Run]) -> Vec<(u64, u64)> { + let mut v: Vec<_> = surplus(runs, &BTreeSet::new()) + .into_iter() + .map(|(r, k)| (r.id, k)) + .collect(); + v.sort_unstable(); + v + } + + #[test] + fn keeps_lowest_queued_run() { + let runs = [ + run(7, 1, "queued", "pull_request", "a", &[3]), + run(5, 1, "queued", "pull_request", "a", &[3]), + run(9, 1, "queued", "pull_request", "a", &[3]), + ]; + assert_eq!(picked(&runs), [(7, 5), (9, 5)]); + } + + #[test] + fn a_running_run_is_kept_over_lower_queued_ones() { + let runs = [ + run(1, 1, "queued", "push", "a", &[]), + run(8, 1, "in_progress", "push", "a", &[]), + run(6, 1, "in_progress", "push", "a", &[]), + ]; + assert_eq!(picked(&runs), [(1, 6), (8, 6)]); + } + + #[test] + fn different_keys_and_exempt_events_do_not_group() { + let runs = [ + run(1, 1, "queued", "pull_request", "a", &[3]), + run(2, 2, "queued", "pull_request", "a", &[3]), + run(3, 1, "queued", "pull_request", "b", &[3]), + run(4, 1, "queued", "pull_request", "a", &[4]), + run(5, 1, "queued", "workflow_dispatch", "a", &[]), + run(6, 1, "queued", "workflow_dispatch", "a", &[]), + run(7, 1, "queued", "merge_group", "a", &[]), + run(8, 1, "queued", "merge_group", "a", &[]), + ]; + assert!(picked(&runs).is_empty()); + } +} diff --git a/src/lead.rs b/src/lead.rs new file mode 100644 index 0000000..72956a5 --- /dev/null +++ b/src/lead.rs @@ -0,0 +1,750 @@ +//! Fenced leader lease for the queen armada. +//! +//! Several queens run at once; exactly one may allocate and mint JIT runners. +//! The leader holds a lease object, `{holder, epoch, expires_at}`, in a store +//! every queen can reach. The store is the only authority: a write succeeds +//! only if the object's generation still matches what the writer last read +//! (compare-and-swap), so two queens that race for the same expired lease +//! cannot both win. The trait [`LeaseStore`] mirrors an object store's +//! generation-match precondition (GCS `ifGenerationMatch`, S3 conditional +//! writes) without binding to either; this module carries no HTTP client. +//! +//! [`Generation`] is an opaque string. GCS generations are integers, but S3 +//! compatible stores condition on an ETag, so a string carries both and the +//! module never does arithmetic on it. +//! +//! Decision rules ([`decide`]), evaluated against the lease as just read: +//! +//! * No lease, or `now >= expires_at` (expiry is inclusive): **acquire** with +//! `epoch = previous epoch + 1`, guarded by the generation just read (or +//! `Absent` when there is no object). +//! * We hold it and `now - renewed_at >= renew_secs`: **renew**, same epoch. +//! * We hold it and no renewal is due: **hold**; nothing is written. +//! * Someone else holds an unexpired lease: **follow**; a follower never writes. +//! * Stale reads: a record whose epoch is lower than one this queen has already +//! seen is never followed or acted on. The epoch only moves forward, so such +//! a record is an out-of-date read; [`decide`] returns `Hold` and +//! [`Leader::step`] keeps its previous role (which still lapses when our own +//! lease expires). +//! +//! [`Leader::step`] applies a decision. It never retries in a loop: on +//! `PreconditionFailed` it re-reads once and reports the role that results; on +//! `Unavailable` it keeps the previous role only while our own lease is +//! unexpired by the local clock (`now < expires_at`), and otherwise steps down. +//! +//! The epoch is the fencing token ([`fence_ok`], [`runner_name`]). Slice 2 +//! wires this into the queen; nothing calls it yet. + +use serde::{Deserialize, Serialize}; + +use crate::wake::BoxFuture; + +/// The lease object as stored (JSON on the wire). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseRecord { + pub holder: String, + pub epoch: u64, + /// Unix seconds. The lease is expired when `now >= expires_at`. + pub expires_at: i64, +} + +/// Opaque version of the lease object, as issued by the store. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Generation(pub String); + +/// What a write requires of the object it replaces. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Precondition { + /// Create only if no object exists. + Absent, + /// Replace only if the object still has this generation. + Generation(Generation), +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum LeaseStoreError { + #[error("lease precondition failed")] + PreconditionFailed, + #[error("lease store unavailable: {0}")] + Unavailable(String), +} + +/// Conditional-write storage for the single lease object. +pub trait LeaseStore: Send + Sync { + fn read(&self) -> BoxFuture<'_, Result, LeaseStoreError>>; + + fn write<'a>( + &'a self, + record: &'a LeaseRecord, + expected: Precondition, + ) -> BoxFuture<'a, Result>; +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LeadSettings { + pub ttl_secs: u64, + pub renew_secs: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum LeadError { + #[error("lease ttl_secs and renew_secs must be greater than zero")] + Zero, + #[error("lease renew_secs ({renew}) must be less than ttl_secs ({ttl})")] + RenewNotBeforeTtl { renew: u64, ttl: u64 }, + #[error("lease holder name must not be empty")] + EmptyName, +} + +impl LeadSettings { + pub fn validate(&self) -> Result<(), LeadError> { + if self.ttl_secs == 0 || self.renew_secs == 0 { + return Err(LeadError::Zero); + } + if self.renew_secs >= self.ttl_secs { + return Err(LeadError::RenewNotBeforeTtl { + renew: self.renew_secs, + ttl: self.ttl_secs, + }); + } + Ok(()) + } + + fn ttl(&self) -> i64 { + i64::try_from(self.ttl_secs).unwrap_or(i64::MAX) + } + + fn renew(&self) -> i64 { + i64::try_from(self.renew_secs).unwrap_or(i64::MAX) + } +} + +/// What this queen believes it holds. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Held { + pub epoch: u64, + pub generation: Generation, + /// Time of the write that created or last renewed the lease. + pub renewed_at: i64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Decision { + Acquire { + record: LeaseRecord, + precondition: Precondition, + }, + Renew { + record: LeaseRecord, + precondition: Precondition, + }, + Follow { + leader: String, + epoch: u64, + }, + Hold, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Role { + Leader { epoch: u64 }, + Follower { leader: Option, epoch: u64 }, +} + +impl Role { + pub fn is_leader(&self) -> bool { + matches!(self, Role::Leader { .. }) + } +} + +/// Pure lease decision; see the module docs for the rules. +pub fn decide( + me: &str, + now: i64, + settings: &LeadSettings, + current: Option<&(LeaseRecord, Generation)>, + held: Option<&Held>, +) -> Decision { + let expires = now.saturating_add(settings.ttl()); + let Some((rec, generation)) = current else { + return Decision::Acquire { + record: LeaseRecord { + holder: me.to_string(), + epoch: held.map_or(0, |h| h.epoch).saturating_add(1), + expires_at: expires, + }, + precondition: Precondition::Absent, + }; + }; + if held.is_some_and(|h| rec.epoch < h.epoch) { + return Decision::Hold; + } + if now >= rec.expires_at { + return Decision::Acquire { + record: LeaseRecord { + holder: me.to_string(), + epoch: rec.epoch.saturating_add(1), + expires_at: expires, + }, + precondition: Precondition::Generation(generation.clone()), + }; + } + if rec.holder != me { + return Decision::Follow { + leader: rec.holder.clone(), + epoch: rec.epoch, + }; + } + // Our name holds an unexpired lease. Without `held` (restart) adopt it. + let renewed_at = held.map_or(rec.expires_at.saturating_sub(settings.ttl()), |h| { + h.renewed_at + }); + if now.saturating_sub(renewed_at) >= settings.renew() { + Decision::Renew { + record: LeaseRecord { + holder: me.to_string(), + epoch: rec.epoch, + expires_at: expires, + }, + precondition: Precondition::Generation(generation.clone()), + } + } else { + Decision::Hold + } +} + +/// A grant whose epoch is lower than the highest seen is rejected. +pub fn fence_ok(seen_epoch: u64, grant_epoch: u64) -> bool { + grant_epoch >= seen_epoch +} + +/// `-e-`. +pub fn runner_name(prefix: &str, drone: &str, epoch: u64, lease_id: &str) -> String { + format!("{prefix}{drone}-e{epoch}-{lease_id}") +} + +/// The epoch in a runner name. Lease ids contain no `-`, so the epoch is the +/// `e` segment just before the last one. The epoch-less format +/// `-` returns `None` (slice 2 treats it as epoch 0). +/// A legacy name whose drone part ends in `-e` is indistinguishable +/// from an epoch segment. +pub fn parse_runner_epoch(name: &str) -> Option { + let (head, lease) = name.rsplit_once('-')?; + if lease.is_empty() { + return None; + } + let (rest, seg) = head.rsplit_once('-')?; + if rest.is_empty() { + return None; + } + let digits = seg.strip_prefix('e')?; + if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + digits.parse().ok() +} + +/// One queen's view of the lease. Drive it with [`Leader::step`] on a timer +/// well inside `renew_secs`. +#[derive(Debug)] +pub struct Leader { + me: String, + settings: LeadSettings, + held: Option, + seen_epoch: u64, + role: Role, +} + +impl Leader { + pub fn new(me: impl Into, settings: LeadSettings) -> Result { + let me = me.into(); + if me.is_empty() { + return Err(LeadError::EmptyName); + } + settings.validate()?; + Ok(Self { + me, + settings, + held: None, + seen_epoch: 0, + role: Role::Follower { + leader: None, + epoch: 0, + }, + }) + } + + pub fn role(&self) -> &Role { + &self.role + } + + /// Highest epoch this queen has seen; the floor for [`fence_ok`]. + pub fn seen_epoch(&self) -> u64 { + self.seen_epoch + } + + fn held_expires_at(&self) -> Option { + self.held + .as_ref() + .map(|h| h.renewed_at.saturating_add(self.settings.ttl())) + } + + fn follow(&mut self, leader: Option, epoch: u64) -> Role { + self.held = None; + self.seen_epoch = self.seen_epoch.max(epoch); + self.role = Role::Follower { + leader, + epoch: self.seen_epoch, + }; + self.role.clone() + } + + fn lead(&mut self, epoch: u64, generation: Generation, renewed_at: i64) -> Role { + self.seen_epoch = self.seen_epoch.max(epoch); + self.held = Some(Held { + epoch, + generation, + renewed_at, + }); + self.role = Role::Leader { epoch }; + self.role.clone() + } + + /// Keep the previous role only while our own lease is unexpired. + fn keep_or_step_down(&mut self, now: i64) -> Role { + match (&self.role, self.held_expires_at()) { + (Role::Leader { .. }, Some(exp)) if now < exp => self.role.clone(), + (Role::Leader { .. }, _) => self.follow(None, 0), + _ => self.role.clone(), + } + } + + /// Adopt whatever the store says without writing. + fn adopt(&mut self, now: i64, current: Option<(LeaseRecord, Generation)>) -> Role { + match current { + Some((rec, generation)) if rec.epoch >= self.seen_epoch && now < rec.expires_at => { + if rec.holder == self.me { + let renewed_at = rec.expires_at.saturating_sub(self.settings.ttl()); + self.lead(rec.epoch, generation, renewed_at) + } else { + self.follow(Some(rec.holder), rec.epoch) + } + } + Some((rec, _)) if rec.epoch >= self.seen_epoch => self.follow(None, rec.epoch), + Some(_) => self.keep_or_step_down(now), + None => self.follow(None, 0), + } + } + + pub async fn step(&mut self, store: &dyn LeaseStore, now: i64) -> Role { + let Ok(current) = store.read().await else { + return self.keep_or_step_down(now); + }; + if current + .as_ref() + .is_some_and(|(r, _)| r.epoch < self.seen_epoch) + { + return self.keep_or_step_down(now); + } + match decide( + &self.me, + now, + &self.settings, + current.as_ref(), + self.held.as_ref(), + ) { + Decision::Hold => self.keep_or_step_down(now), + Decision::Follow { leader, epoch } => self.follow(Some(leader), epoch), + Decision::Acquire { + mut record, + precondition, + } => { + record.epoch = record.epoch.max(self.seen_epoch.saturating_add(1)); + self.write(store, now, record, precondition).await + } + Decision::Renew { + record, + precondition, + } => self.write(store, now, record, precondition).await, + } + } + + async fn write( + &mut self, + store: &dyn LeaseStore, + now: i64, + record: LeaseRecord, + precondition: Precondition, + ) -> Role { + match store.write(&record, precondition).await { + Ok(generation) => self.lead(record.epoch, generation, now), + Err(LeaseStoreError::PreconditionFailed) => match store.read().await { + Ok(current) => self.adopt(now, current), + Err(_) => self.keep_or_step_down(now), + }, + Err(LeaseStoreError::Unavailable(_)) => self.keep_or_step_down(now), + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use super::*; + + #[derive(Default)] + struct State { + object: Option<(LeaseRecord, u64)>, + counter: u64, + writes: Vec<(LeaseRecord, Precondition)>, + fail_next_write: Option, + /// Another queen's write that lands just before the next write is checked. + interleave: Option, + down: bool, + } + + #[derive(Default)] + struct FakeStore(Mutex); + + impl FakeStore { + fn fail_next_write_with(&self, e: LeaseStoreError) { + self.0.lock().unwrap().fail_next_write = Some(e); + } + fn interleave_before_next_write(&self, r: LeaseRecord) { + self.0.lock().unwrap().interleave = Some(r); + } + fn set_down(&self, down: bool) { + self.0.lock().unwrap().down = down; + } + fn writes(&self) -> Vec<(LeaseRecord, Precondition)> { + self.0.lock().unwrap().writes.clone() + } + fn record(&self) -> Option { + self.0.lock().unwrap().object.as_ref().map(|o| o.0.clone()) + } + } + + impl LeaseStore for FakeStore { + fn read( + &self, + ) -> BoxFuture<'_, Result, LeaseStoreError>> { + Box::pin(async move { + let s = self.0.lock().unwrap(); + if s.down { + return Err(LeaseStoreError::Unavailable("down".into())); + } + Ok(s.object + .as_ref() + .map(|(r, g)| (r.clone(), Generation(g.to_string())))) + }) + } + + fn write<'a>( + &'a self, + record: &'a LeaseRecord, + expected: Precondition, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut s = self.0.lock().unwrap(); + s.writes.push((record.clone(), expected.clone())); + if s.down { + return Err(LeaseStoreError::Unavailable("down".into())); + } + if let Some(e) = s.fail_next_write.take() { + return Err(e); + } + if let Some(other) = s.interleave.take() { + s.counter += 1; + let g = s.counter; + s.object = Some((other, g)); + } + let ok = match (&s.object, &expected) { + (None, Precondition::Absent) => true, + (Some((_, g)), Precondition::Generation(want)) => g.to_string() == want.0, + _ => false, + }; + if !ok { + return Err(LeaseStoreError::PreconditionFailed); + } + s.counter += 1; + let g = s.counter; + s.object = Some((record.clone(), g)); + Ok(Generation(g.to_string())) + }) + } + } + + fn settings() -> LeadSettings { + LeadSettings { + ttl_secs: 30, + renew_secs: 10, + } + } + + fn queen(name: &str) -> Leader { + Leader::new(name, settings()).unwrap() + } + + fn generation(s: &str) -> Generation { + Generation(s.to_string()) + } + + fn lease(holder: &str, epoch: u64, expires_at: i64) -> LeaseRecord { + LeaseRecord { + holder: holder.into(), + epoch, + expires_at, + } + } + + fn follower(leader: &str, epoch: u64) -> Role { + Role::Follower { + leader: Some(leader.into()), + epoch, + } + } + + #[tokio::test] + async fn two_queens_one_leader() { + let store = FakeStore::default(); + let mut a = queen("a"); + let mut b = queen("b"); + assert_eq!(a.step(&store, 100).await, Role::Leader { epoch: 1 }); + assert_eq!(b.step(&store, 100).await, follower("a", 1)); + assert_eq!(store.writes().len(), 1); + } + + #[tokio::test] + async fn racing_acquire_loser_follows_winner() { + let store = FakeStore::default(); + let mut a = queen("a"); + // b's create lands between a's read and a's write. + store.interleave_before_next_write(lease("b", 1, 130)); + assert_eq!(a.step(&store, 100).await, follower("b", 1)); + assert_eq!(store.record().unwrap().holder, "b"); + } + + #[tokio::test] + async fn leader_renews_before_ttl_and_keeps_epoch() { + let store = FakeStore::default(); + let mut a = queen("a"); + let mut b = queen("b"); + a.step(&store, 100).await; + b.step(&store, 100).await; + let mut t = 100; + for _ in 0..6 { + t += 5; + assert_eq!(a.step(&store, t).await, Role::Leader { epoch: 1 }); + assert_eq!(b.step(&store, t).await, follower("a", 1)); + } + let rec = store.record().unwrap(); + assert_eq!((rec.holder.as_str(), rec.epoch), ("a", 1)); + assert!(rec.expires_at > t); + // Only a wrote: acquire at 100, renewals at 110, 120, 130. + assert!(store.writes().iter().all(|(r, _)| r.holder == "a")); + assert_eq!(store.writes().len(), 4); + } + + #[tokio::test] + async fn takeover_after_expiry_and_no_split_brain() { + let store = FakeStore::default(); + let mut a = queen("a"); + let mut b = queen("b"); + a.step(&store, 100).await; // expires at 130 + b.step(&store, 100).await; + // a is paused. b follows until expiry, which is inclusive. + assert!(!b.step(&store, 129).await.is_leader()); + assert_eq!(b.step(&store, 130).await, Role::Leader { epoch: 2 }); + // a wakes up and sees the new holder. + assert_eq!(a.step(&store, 131).await, follower("b", 2)); + for now in 131..300 { + let ra = a.step(&store, now).await; + let rb = b.step(&store, now).await; + assert!(!(ra.is_leader() && rb.is_leader()), "split brain at {now}"); + } + } + + #[tokio::test] + async fn old_leader_renew_loses_cas_and_follows() { + let store = FakeStore::default(); + let mut a = queen("a"); + a.step(&store, 100).await; // gen 1, expires 130 + // b takes over between a's read and a's renew write. + store.interleave_before_next_write(lease("b", 2, 200)); + assert_eq!(a.step(&store, 110).await, follower("b", 2)); + assert_eq!(a.seen_epoch(), 2); + // Exactly one renew was attempted, no retry loop. + assert_eq!(store.writes().len(), 2); + assert_eq!(store.record().unwrap().holder, "b"); + } + + #[tokio::test] + async fn precondition_failed_with_us_still_holding_keeps_leading() { + let store = FakeStore::default(); + let mut a = queen("a"); + a.step(&store, 100).await; + store.fail_next_write_with(LeaseStoreError::PreconditionFailed); + assert_eq!(a.step(&store, 110).await, Role::Leader { epoch: 1 }); + } + + #[tokio::test] + async fn follower_never_writes_while_lease_unexpired() { + let store = FakeStore::default(); + let mut a = queen("a"); + let mut b = queen("b"); + a.step(&store, 100).await; + let before = store.writes().len(); + for t in 100..130 { + b.step(&store, t).await; + } + assert_eq!(store.writes().len(), before); + } + + #[test] + fn stale_epoch_is_fenced() { + assert!(!fence_ok(5, 4)); + assert!(fence_ok(5, 5)); + assert!(fence_ok(5, 6)); + let held = Held { + epoch: 5, + generation: generation("9"), + renewed_at: 100, + }; + let stale = (lease("other", 4, 10_000), generation("3")); + let d = decide("me", 101, &settings(), Some(&stale), Some(&held)); + assert_eq!(d, Decision::Hold); + } + + #[tokio::test] + async fn stale_read_is_not_followed_by_step() { + let store = FakeStore::default(); + let mut a = queen("a"); + let mut b = queen("b"); + a.step(&store, 100).await; + b.step(&store, 140).await; // epoch 2 + assert_eq!(b.seen_epoch(), 2); + // The store serves an older epoch (stale replica): b keeps its role. + store.0.lock().unwrap().object = Some((lease("a", 1, 10_000), 77)); + assert_eq!(b.step(&store, 141).await, Role::Leader { epoch: 2 }); + // A follower that has seen epoch 2 does not follow epoch 1 either. + let mut c = queen("c"); + store.0.lock().unwrap().object = Some((lease("b", 2, 10_000), 78)); + assert_eq!(c.step(&store, 142).await, follower("b", 2)); + store.0.lock().unwrap().object = Some((lease("a", 1, 10_000), 79)); + assert_eq!(c.step(&store, 143).await, follower("b", 2)); + } + + #[tokio::test] + async fn store_unavailable_leader_keeps_until_own_expiry() { + let store = FakeStore::default(); + let mut a = queen("a"); + a.step(&store, 100).await; // expires 130 + store.set_down(true); + assert_eq!(a.step(&store, 115).await, Role::Leader { epoch: 1 }); + assert_eq!(a.step(&store, 129).await, Role::Leader { epoch: 1 }); + assert!(!a.step(&store, 130).await.is_leader()); + assert_eq!(a.role().clone(), a.step(&store, 131).await); + assert!(matches!(a.role(), Role::Follower { leader: None, .. })); + } + + #[tokio::test] + async fn failed_renew_write_does_not_extend_local_expiry() { + let store = FakeStore::default(); + let mut a = queen("a"); + a.step(&store, 100).await; + store.fail_next_write_with(LeaseStoreError::Unavailable("x".into())); + assert_eq!(a.step(&store, 110).await, Role::Leader { epoch: 1 }); + store.set_down(true); + assert!(!a.step(&store, 130).await.is_leader()); + } + + #[tokio::test] + async fn store_unavailable_follower_stays_follower() { + let store = FakeStore::default(); + let mut a = queen("a"); + let mut b = queen("b"); + a.step(&store, 100).await; + b.step(&store, 100).await; + store.set_down(true); + assert_eq!(b.step(&store, 200).await, follower("a", 1)); + assert_eq!(store.writes().len(), 1); + } + + #[test] + fn decide_rules() { + let s = settings(); + let rec = |h: &str, e: u64, x: i64| (lease(h, e, x), generation("g")); + let held = Held { + epoch: 3, + generation: generation("g"), + renewed_at: 100, + }; + assert!(matches!( + decide("me", 100, &s, None, None), + Decision::Acquire { record, precondition: Precondition::Absent } if record.epoch == 1 + )); + assert!(matches!( + decide("me", 130, &s, Some(&rec("x", 7, 130)), None), + Decision::Acquire { record, precondition: Precondition::Generation(_) } + if record.epoch == 8 && record.expires_at == 160 + )); + assert_eq!( + decide("me", 105, &s, Some(&rec("me", 3, 130)), Some(&held)), + Decision::Hold + ); + assert!(matches!( + decide("me", 110, &s, Some(&rec("me", 3, 130)), Some(&held)), + Decision::Renew { record, .. } if record.epoch == 3 && record.expires_at == 140 + )); + assert_eq!( + decide("me", 105, &s, Some(&rec("x", 3, 130)), Some(&held)), + Decision::Follow { + leader: "x".into(), + epoch: 3 + } + ); + } + + #[test] + fn lease_record_json_round_trip() { + let r = lease("a", 3, 42); + let json = serde_json::to_string(&r).unwrap(); + assert_eq!(json, r#"{"holder":"a","epoch":3,"expires_at":42}"#); + assert_eq!(serde_json::from_str::(&json).unwrap(), r); + } + + #[test] + fn runner_names_round_trip() { + let n = runner_name("borg-", "host-a", 12, "abc123"); + assert_eq!(n, "borg-host-a-e12-abc123"); + assert_eq!(parse_runner_epoch(&n), Some(12)); + assert_eq!(parse_runner_epoch(&runner_name("", "d", 0, "ff")), Some(0)); + assert_eq!(parse_runner_epoch("borg-host-a-abc123"), None); + assert_eq!(parse_runner_epoch("nodash"), None); + assert_eq!(parse_runner_epoch("borg-e-abc"), None); + } + + #[test] + fn settings_validation() { + let check = |t, r| { + LeadSettings { + ttl_secs: t, + renew_secs: r, + } + .validate() + }; + assert!(check(30, 10).is_ok()); + assert!(check(30, 30).is_err()); + assert!(check(30, 31).is_err()); + assert_eq!(check(0, 0), Err(LeadError::Zero)); + assert_eq!(check(30, 0), Err(LeadError::Zero)); + assert_eq!(check(0, 1), Err(LeadError::Zero)); + let bad = LeadSettings { + ttl_secs: 5, + renew_secs: 5, + }; + assert!(Leader::new("a", bad).is_err()); + assert!(Leader::new("", settings()).is_err()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 26cdffb..7adeda9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,10 +22,13 @@ pub mod github; pub mod health; pub mod homeostasis; pub mod janitor; +pub mod lead; +pub mod preempt; pub mod probe; pub mod profile; pub mod queen; pub mod reclaim; +pub mod requeue; pub mod rollout; pub mod run; pub mod schedule; diff --git a/src/main.rs b/src/main.rs index db42149..0ef3754 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,6 +15,7 @@ use borg::github::GithubClient; use borg::janitor::Janitor; use borg::probe::{self, system::SystemHost}; use borg::queen::Queen; +use borg::requeue::Requeuer; use borg::store::Store; use borg::tailnet::{TailnetEnroll, TailscaledWhois}; use borg::update; @@ -102,6 +103,24 @@ struct UpdateArgs { json: bool, } +/// A unit scope on the command line. +#[derive(Clone, Copy, clap::ValueEnum)] +enum ScopeArg { + Deploy, + Main, + Pr, +} + +impl From for borg::classes::Scope { + fn from(s: ScopeArg) -> Self { + match s { + ScopeArg::Deploy => Self::Deploy, + ScopeArg::Main => Self::Main, + ScopeArg::Pr => Self::Pr, + } + } +} + #[derive(Args)] struct RunArgs { /// Queen base URL. Without it the command simply runs here. @@ -136,6 +155,9 @@ struct RunArgs { /// Run even if an identical command already passed for this tree. #[arg(long)] force: bool, + /// Priority scope of the unit; the client token's maximum applies. + #[arg(long, env = "BORG_RUN_SCOPE", value_enum, default_value = "pr")] + scope: ScopeArg, /// Seconds to wait for a drone before running here. #[arg(long, env = "BORG_RUN_WAIT_SECS", default_value_t = 60)] wait_secs: u64, @@ -194,6 +216,9 @@ enum QueenCommand { config: ConfigArg, /// Client name: lowercase letters, digits and dashes. name: String, + /// Highest unit scope this client may submit. + #[arg(long, value_enum, default_value = "pr")] + max_scope: ScopeArg, }, /// Stop granting a drone new leases (running jobs finish). Quarantine { @@ -320,6 +345,15 @@ struct InstallArgs { /// Add this label to every runner the drone registers; repeatable. #[arg(long = "extra-label")] extra_labels: Vec, + /// `shared` (default), `deploy` or `deploy-only`: a deploy drone runs only + /// deploy and main jobs, a deploy-only drone only deploy jobs, a shared + /// drone never runs deploy jobs. + #[arg(long, value_parser = parse_trust, default_value = "shared")] + trust: borg::classes::Trust, + /// Start standby deploy runners after this many seconds without a queen + /// answer; 0 turns fail-open off. + #[arg(long, default_value_t = 120)] + fail_open_after_secs: u64, /// KEY=VALUE injected into every lease's environment; repeatable. #[arg(long = "runner-env", value_parser = parse_env_pair)] runner_env: Vec<(String, String)>, @@ -351,6 +385,11 @@ struct InstallArgs { dry_run: bool, } +fn parse_trust(s: &str) -> Result { + serde_json::from_value(serde_json::Value::from(s)) + .map_err(|_| "expected shared, deploy or deploy-only".to_owned()) +} + fn parse_env_pair(s: &str) -> Result<(String, String), String> { s.split_once('=') .filter(|(k, _)| !k.is_empty()) @@ -417,6 +456,7 @@ async fn run_queen(path: &Path) -> Result<()> { .await .with_context(|| format!("bind {}", cfg.listen))?; tracing::info!("queen listening on {}", cfg.listen); + let mut requeuer = None; let janitor = match &cfg.supersede_janitor { Some(settings) => { let token = janitor_token(settings)?; @@ -428,6 +468,13 @@ async fn run_queen(path: &Path) -> Result<()> { "cancelling" } ); + if cfg.preemption.is_some() { + requeuer = Some(Requeuer::new( + &cfg.github.api_base, + token.clone(), + settings, + )?); + } Some(Janitor::new(&cfg.github.api_base, token, settings.clone())?) } None => None, @@ -449,6 +496,9 @@ async fn run_queen(path: &Path) -> Result<()> { if let Some(j) = janitor { queen = queen.with_janitor(j); } + if let Some(r) = requeuer { + queen = queen.with_requeuer(r); + } let queen = Arc::new(queen); api::serve(queen, listener, wait_for_shutdown()).await?; Ok(()) @@ -501,7 +551,7 @@ async fn report(a: ReportArgs) -> Result<()> { Ok(()) } -async fn enroll_client(path: &Path, name: &str) -> Result<()> { +async fn enroll_client(path: &Path, name: &str, max_scope: borg::classes::Scope) -> Result<()> { if !enroll::valid_drone_name(name) { bail!( "client name must be 1-40 lowercase letters, digits or dashes, not starting with a dash" @@ -512,7 +562,7 @@ async fn enroll_client(path: &Path, name: &str) -> Result<()> { let token = enroll::new_token(); let now = time::OffsetDateTime::now_utc().unix_timestamp(); store - .enroll_client(name, &enroll::hash_token(&token), now) + .enroll_client_scoped(name, &enroll::hash_token(&token), max_scope, now) .await?; eprintln!("Enrolled client {name}. The token is shown once; the queen stores only its hash."); println!("{token}"); @@ -620,6 +670,8 @@ fn install_drone(a: InstallArgs) -> Result<()> { cache_root: a.cache_root, adopt_cache: a.adopt_cache, extra_labels: a.extra_labels, + trust: a.trust, + fail_open_after_secs: a.fail_open_after_secs, runner_env: a.runner_env.into_iter().collect(), job_started_hook: a.job_started_hook, job_completed_hook: a.job_completed_hook, @@ -701,6 +753,7 @@ async fn run_command(a: RunArgs) -> u8 { toolchain: a.toolchain, timeout_secs: a.timeout, force: a.force, + scope: borg::classes::Scope::from(a.scope).as_str().to_owned(), wait_secs: a.wait_secs, publish: a.publish, argv: a.argv, @@ -788,9 +841,11 @@ async fn dispatch(command: Command) -> Result<()> { reason, name, } => set_quarantine(&config.config, &name, Some(&reason)).await, - QueenCommand::EnrollClient { config, name } => { - enroll_client(&config.config, &name).await - } + QueenCommand::EnrollClient { + config, + name, + max_scope, + } => enroll_client(&config.config, &name, max_scope.into()).await, QueenCommand::Rollout { command } => match command { RolloutCommand::Halt { config, reason } => { rollout_admin(&config.config, Some(&reason)).await diff --git a/src/preempt.rs b/src/preempt.rs new file mode 100644 index 0000000..ddb3628 --- /dev/null +++ b/src/preempt.rs @@ -0,0 +1,409 @@ +//! Preemption: a pure choice of running leases to cancel for starving demand. +//! +//! Demand of a `(class, scope)` that the allocator could not place and that has +//! waited at least `after_secs` may displace a running lease of a strictly +//! lower-priority scope, provided freeing that lease leaves its drone room for +//! the starving class. Priority is the [`Scope`] ordering: smaller outranks. + +use std::collections::BTreeSet; + +use serde::Deserialize; + +use crate::allocator::{ + ActiveLease, DroneState, LeaseState, QueueAge, Unmet, free_room, free_slots, +}; +use crate::classes::{Class, Classes, Scope}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PreemptSettings { + /// Starving demand must have waited this long. + pub after_secs: u64, + /// Most leases cancelled in one queen pass. + pub max_per_pass: u32, + /// A lease younger than this is never cancelled. + pub min_victim_age_secs: u64, +} + +impl PreemptSettings { + pub fn validate(&self) -> Result<(), String> { + if self.after_secs == 0 || self.max_per_pass == 0 || self.min_victim_age_secs == 0 { + return Err( + "preemption after_secs, max_per_pass and min_victim_age_secs must be positive" + .to_owned(), + ); + } + Ok(()) + } +} + +/// A running lease whose job is known, so it can be cancelled. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Running { + pub lease_id: String, + pub drone: String, + pub class: Class, + pub scope: Scope, + pub started_at: i64, + /// Already told to stop; its drone is not yet free. + pub preempted: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Victim { + pub lease_id: String, + pub drone: String, + pub class: Class, + pub scope: Scope, + /// The starving demand this cancellation is for. + pub for_class: Class, + pub for_scope: Scope, +} + +pub struct PreemptInput<'a> { + pub classes: &'a Classes, + pub settings: &'a PreemptSettings, + pub drones: &'a [DroneState], + /// Every active lease that is not in `running` (pending, unit-placed, or + /// running with no known job); they hold capacity and are never victims. + pub others: &'a [ActiveLease], + pub running: &'a [Running], + /// Demand the allocator could not place this pass. + pub unmet: &'a [Unmet], + pub queue_ages: &'a [QueueAge], + pub now: i64, +} + +/// Capacity held by each lease: `Some(i)` marks `running[i]`. +type Held = Vec<(Option, ActiveLease)>; + +/// Whether `running[i]`'s drone could hold a `class` lease once it is gone. +fn frees_room_for(input: &PreemptInput<'_>, held: &Held, i: usize, class: Class) -> bool { + let Some(drone) = input.drones.iter().find(|d| d.id == input.running[i].drone) else { + return false; + }; + let rest: Vec = held + .iter() + .filter(|(c, _)| *c != Some(i)) + .map(|(_, l)| l.clone()) + .collect(); + let room = free_room(drone, &rest, input.classes); + !drone.blocked_classes.contains(&class) + && free_slots(drone, &room, input.classes.spec(class)) > 0 +} + +/// Up to `max_per_pass` leases to cancel. Deterministic. +pub fn select(input: &PreemptInput<'_>) -> Vec { + let s = input.settings; + let min_age = i64::try_from(s.min_victim_age_secs).unwrap_or(i64::MAX); + let mut held: Held = input + .others + .iter() + .map(|l| (None, l.clone())) + .chain(input.running.iter().enumerate().map(|(i, r)| { + ( + Some(i), + ActiveLease { + drone_id: r.drone.clone(), + class: r.class, + scope: r.scope, + state: LeaseState::Running, + }, + ) + })) + .collect(); + + let mut starving: Vec<(Scope, Class)> = input + .unmet + .iter() + .filter(|u| u.count > 0) + .filter(|u| { + input.queue_ages.iter().any(|a| { + a.class == u.class && a.scope == u.scope && a.oldest_age_secs >= s.after_secs + }) + }) + .map(|u| (u.scope, u.class)) + .collect(); + starving.sort_by_key(|&(scope, class)| (scope, std::cmp::Reverse(class))); + starving.dedup(); + + let mut out = Vec::new(); + let mut picked = BTreeSet::new(); + for (scope, class) in starving { + if out.len() >= s.max_per_pass as usize { + break; + } + // A cancellation already under way that frees room for this demand + // covers it; cancelling more would take another job for nothing. + let covered = input + .running + .iter() + .enumerate() + .any(|(i, r)| r.preempted && frees_room_for(input, &held, i, class)); + if covered { + continue; + } + let best = input + .running + .iter() + .enumerate() + .filter(|(i, r)| { + !r.preempted + && !picked.contains(i) + && r.scope > scope + && input.now - r.started_at >= min_age + && frees_room_for(input, &held, *i, class) + }) + .min_by(|(_, a), (_, b)| { + b.scope + .cmp(&a.scope) + .then(a.class.cmp(&b.class)) + .then(b.started_at.cmp(&a.started_at)) + .then(a.lease_id.cmp(&b.lease_id)) + }); + let Some((i, r)) = best else { + continue; + }; + picked.insert(i); + // The freed room now belongs to the starving demand. + if let Some((owner, lease)) = held.iter_mut().find(|(c, _)| *c == Some(i)) { + *owner = None; + lease.class = class; + lease.scope = scope; + } + out.push(Victim { + lease_id: r.lease_id.clone(), + drone: r.drone.clone(), + class: r.class, + scope: r.scope, + for_class: class, + for_scope: scope, + }); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + const NOW: i64 = 10_000; + + fn settings() -> PreemptSettings { + PreemptSettings { + after_secs: 600, + max_per_pass: 1, + min_victim_age_secs: 60, + } + } + + fn drone(id: &str, mem_mib: u64, vcpu: u32) -> DroneState { + DroneState { + id: id.to_owned(), + vcpu, + mem_mib, + docker: true, + accepting: true, + max_leases: None, + bids: Vec::new(), + profiles: Vec::new(), + blocked_classes: Vec::new(), + trust: crate::classes::Trust::Shared, + cpu_psi_pm: None, + } + } + + /// A drone sized for `n` leases of `class`. + fn sized_for(id: &str, class: Class, n: u32) -> DroneState { + let classes = Classes::default(); + let spec = classes.spec(class); + drone(id, spec.mem_mib * u64::from(n), spec.vcpu * n) + } + + fn run(id: &str, drone: &str, class: Class, scope: Scope, started_at: i64) -> Running { + Running { + lease_id: id.to_owned(), + drone: drone.to_owned(), + class, + scope, + started_at, + preempted: false, + } + } + + type Demand = (Vec, Vec); + + fn starving(keys: &[(Class, Scope)], age: u64) -> Demand { + ( + keys.iter() + .map(|&(class, scope)| Unmet { + class, + scope, + count: 1, + }) + .collect(), + keys.iter() + .map(|&(class, scope)| QueueAge { + class, + scope, + oldest_age_secs: age, + }) + .collect(), + ) + } + + fn pick( + settings: &PreemptSettings, + drones: &[DroneState], + running: &[Running], + demand: &Demand, + ) -> Vec { + select(&PreemptInput { + classes: &Classes::default(), + settings, + drones, + others: &[], + running, + unmet: &demand.0, + queue_ages: &demand.1, + now: NOW, + }) + .into_iter() + .map(|v| v.lease_id) + .collect() + } + + const MEDIUM_MAIN: (Class, Scope) = (Class::Medium, Scope::Main); + + #[test] + fn nothing_starves_before_the_threshold_or_without_unmet_demand() { + let drones = [sized_for("host-a", Class::Medium, 1)]; + let running = [run("a", "host-a", Class::Medium, Scope::Pr, 0)]; + assert!( + pick( + &settings(), + &drones, + &running, + &starving(&[MEDIUM_MAIN], 599) + ) + .is_empty() + ); + assert!(pick(&settings(), &drones, &running, &(vec![], vec![])).is_empty()); + let unplaced_but_not_queued = (starving(&[MEDIUM_MAIN], 9000).0, vec![]); + assert!(pick(&settings(), &drones, &running, &unplaced_but_not_queued).is_empty()); + } + + #[test] + fn starving_main_displaces_a_running_pr_lease() { + let drones = [sized_for("host-a", Class::Medium, 1)]; + let running = [run("a", "host-a", Class::Medium, Scope::Pr, 0)]; + assert_eq!( + pick( + &settings(), + &drones, + &running, + &starving(&[MEDIUM_MAIN], 600) + ), + ["a"] + ); + } + + #[test] + fn equal_or_higher_priority_leases_are_never_victims() { + let drones = [sized_for("host-a", Class::Medium, 1)]; + let main = [run("a", "host-a", Class::Medium, Scope::Main, 0)]; + let demand = starving(&[MEDIUM_MAIN], 9000); + assert!(pick(&settings(), &drones, &main, &demand).is_empty()); + let demand = starving(&[(Class::Medium, Scope::Pr)], 9000); + assert!(pick(&settings(), &drones, &main, &demand).is_empty()); + } + + #[test] + fn no_victim_when_freeing_it_would_not_fit_the_class() { + // The drone fits one heavy lease; a medium Main lease stays beside the + // light Pr one, so removing the light one cannot make room. + let drones = [sized_for("host-a", Class::Heavy, 1)]; + let running = [ + run("a", "host-a", Class::Light, Scope::Pr, 0), + run("b", "host-a", Class::Medium, Scope::Main, 0), + ]; + let demand = starving(&[(Class::Heavy, Scope::Main)], 9000); + assert!(pick(&settings(), &drones, &running, &demand).is_empty()); + } + + #[test] + fn prefers_lowest_priority_scope_then_smallest_class_then_newest() { + let drones = [sized_for("host-a", Class::Heavy, 4)]; + let demand = starving(&[MEDIUM_MAIN], 9000); + let running = [ + run("main-light", "host-a", Class::Light, Scope::Main, 900), + run("old-pr", "host-a", Class::Light, Scope::Pr, 100), + run("new-pr", "host-a", Class::Light, Scope::Pr, 500), + run("medium-pr", "host-a", Class::Medium, Scope::Pr, 900), + ]; + assert_eq!(pick(&settings(), &drones, &running, &demand), ["new-pr"]); + } + + #[test] + fn respects_the_minimum_victim_age() { + let drones = [sized_for("host-a", Class::Medium, 1)]; + let demand = starving(&[MEDIUM_MAIN], 9000); + let young = [run("a", "host-a", Class::Medium, Scope::Pr, NOW - 59)]; + assert!(pick(&settings(), &drones, &young, &demand).is_empty()); + let old = [run("a", "host-a", Class::Medium, Scope::Pr, NOW - 60)]; + assert_eq!(pick(&settings(), &drones, &old, &demand), ["a"]); + } + + #[test] + fn respects_max_per_pass_and_never_reuses_freed_room() { + let drones = [ + sized_for("host-a", Class::Medium, 1), + sized_for("host-b", Class::Medium, 1), + ]; + let running = [ + run("a", "host-a", Class::Medium, Scope::Pr, 0), + run("b", "host-b", Class::Medium, Scope::Pr, 0), + ]; + let two = PreemptSettings { + max_per_pass: 2, + ..settings() + }; + // One starving key claims one drone, whatever the cap. + let one_key = starving(&[MEDIUM_MAIN], 9000); + assert_eq!(pick(&two, &drones, &running, &one_key), ["a"]); + let two_keys = starving(&[MEDIUM_MAIN, (Class::Light, Scope::Main)], 9000); + assert_eq!(pick(&two, &drones, &running, &two_keys), ["a", "b"]); + assert_eq!(pick(&settings(), &drones, &running, &two_keys), ["a"]); + } + + #[test] + fn a_cancellation_under_way_covers_the_demand() { + let drones = [sized_for("host-a", Class::Medium, 2)]; + let mut stopping = run("a", "host-a", Class::Medium, Scope::Pr, 0); + stopping.preempted = true; + let running = [stopping, run("b", "host-a", Class::Medium, Scope::Pr, 0)]; + let demand = starving(&[MEDIUM_MAIN], 9000); + assert!(pick(&settings(), &drones, &running, &demand).is_empty()); + } + + #[test] + fn settings_reject_zero_values() { + assert!(settings().validate().is_ok()); + for bad in [ + PreemptSettings { + after_secs: 0, + ..settings() + }, + PreemptSettings { + max_per_pass: 0, + ..settings() + }, + PreemptSettings { + min_victim_age_secs: 0, + ..settings() + }, + ] { + assert!(bad.validate().is_err()); + } + } +} diff --git a/src/probe/mod.rs b/src/probe/mod.rs index b617cc2..9975e43 100644 --- a/src/probe/mod.rs +++ b/src/probe/mod.rs @@ -5,10 +5,15 @@ pub mod system; +use std::collections::BTreeMap; use std::path::Path; +use std::sync::{Arc, Mutex}; use serde::{Deserialize, Serialize}; +use crate::drone::exec::Exec; +use crate::envcontract::{EnvFacts, PROBE_TOOLS, TOOL_MISSING}; + /// What kind of machine the drone runs on. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] @@ -71,6 +76,120 @@ pub fn probe(host: &dyn Host) -> Capacity { } } +/// Build-environment facts that need no tooling: memory available and CPU +/// pressure. Sent on every heartbeat; the admission rule reads the pressure. +pub fn env_facts(host: &dyn Host, now: i64) -> EnvFacts { + EnvFacts { + observed_at: now, + mem_available_mib: mem_available_mib(host), + cpu_psi_some_avg10_pm: cpu_psi_pm(host), + ..EnvFacts::default() + } +} + +/// How long a tools probe is reused before the tools are looked up again. +const TOOLS_TTL_SECS: i64 = 60; +/// Longest version string reported for a tool. +const TOOL_VERSION_MAX: usize = 80; + +/// The first line of a tool's `--version` output, trimmed and clipped to +/// [`TOOL_VERSION_MAX`] characters; `"unknown"` when there is no output. +pub fn tool_version_line(output: &str) -> String { + let line = output.lines().map(str::trim).find(|l| !l.is_empty()); + match line { + Some(l) => l.chars().take(TOOL_VERSION_MAX).collect(), + None => "unknown".to_owned(), + } +} + +/// For each tool in [`PROBE_TOOLS`], its version line, or `"missing"` when the +/// shell cannot find it on `PATH`. `exec` runs as the job user, so this is the +/// `PATH` a job gets. A tool that exists but prints no version is `"unknown"`. +pub fn tool_versions(exec: &dyn Exec) -> BTreeMap { + PROBE_TOOLS + .iter() + .map(|&tool| { + let found = exec + .run( + "sh", + &["-c", "command -v \"$1\" >/dev/null 2>&1", "sh", tool], + ) + .is_ok_and(|o| o.success); + let value = if found { + // `go` has no `--version` flag. + let args: &[&str] = if tool == "go" { + &["version"] + } else { + &["--version"] + }; + match exec.run(tool, args) { + Ok(o) if o.success => tool_version_line(&o.stdout), + _ => "unknown".to_owned(), + } + } else { + TOOL_MISSING.to_owned() + }; + (tool.to_owned(), value) + }) + .collect() +} + +/// The last tools probe, reused for [`TOOLS_TTL_SECS`]. +#[derive(Clone, Default)] +pub struct ToolsCache(Arc>>); + +/// When the tools were probed, and what was found. +type ToolsProbe = (i64, BTreeMap); + +impl ToolsCache { + /// The cached tools if they are under a minute old, else a fresh probe. + pub fn get(&self, exec: &dyn Exec, now: i64) -> Option> { + let mut slot = self.0.lock().ok()?; + if let Some((at, tools)) = slot.as_ref() + && (0..TOOLS_TTL_SECS).contains(&(now - at)) + { + return Some(tools.clone()); + } + let tools = tool_versions(exec); + *slot = Some((now, tools.clone())); + Some(tools) + } +} + +fn mem_available_mib(host: &dyn Host) -> Option { + let meminfo = host.read_to_string(Path::new("/proc/meminfo"))?; + let kib = meminfo + .lines() + .find_map(|l| l.strip_prefix("MemAvailable:"))? + .split_whitespace() + .next()? + .parse::() + .ok()?; + Some(kib / 1024) +} + +/// The higher of the host-wide and the cgroup `cpu some avg10`, in permille. +/// Inside an LXC container `/proc/pressure/cpu` is the host's or the +/// container's depending on lxcfs, so both files are read. +fn cpu_psi_pm(host: &dyn Host) -> Option { + ["/proc/pressure/cpu", "/sys/fs/cgroup/cpu.pressure"] + .iter() + .filter_map(|path| parse_psi_some_avg10_pm(&host.read_to_string(Path::new(path))?)) + .max() +} + +/// `some avg10=1.23 ...` (percent) as permille, clamped to 1000. +fn parse_psi_some_avg10_pm(text: &str) -> Option { + let line = text.lines().find(|l| l.starts_with("some "))?; + let pct = line + .split_whitespace() + .find_map(|f| f.strip_prefix("avg10="))? + .parse::() + .ok()?; + pct.is_finite() + .then(|| (pct * 10.0).round().clamp(0.0, 1000.0) as u16) +} + fn hostname(host: &dyn Host) -> String { host.read_to_string(Path::new("/proc/sys/kernel/hostname")) .or_else(|| host.run("hostname", &[])) @@ -246,6 +365,117 @@ mod tests { } } + struct FakeExec { + calls: Mutex>, + } + + impl Exec for FakeExec { + fn run(&self, program: &str, args: &[&str]) -> std::io::Result { + self.calls + .lock() + .unwrap() + .push(format!("{program} {}", args.join(" "))); + let ok = |stdout: &str| { + Ok(crate::drone::exec::Output { + success: true, + stdout: stdout.to_owned(), + }) + }; + let no = || { + Ok(crate::drone::exec::Output { + success: false, + stdout: String::new(), + }) + }; + match (program, args) { + ("sh", [_, _, _, tool]) => { + if ["git", "gh", "go", "jq"].contains(tool) { + ok("") + } else { + no() + } + } + ("git", ["--version"]) => ok("git version 2.45.0\nbuilt from source"), + ("gh", ["--version"]) => { + ok("\n gh version 2.50.0 (2026-01-01) \nhttps://example.test") + } + ("go", ["version"]) => ok("go version go1.25.0 linux/amd64"), + _ => no(), + } + } + } + + #[test] + fn version_lines_are_the_first_non_empty_line_trimmed_and_clipped() { + assert_eq!( + tool_version_line("git version 2.45.0\nmore"), + "git version 2.45.0" + ); + assert_eq!(tool_version_line("\n v22.1.0 \n"), "v22.1.0"); + assert_eq!(tool_version_line(" \n"), "unknown"); + assert_eq!(tool_version_line(&"x".repeat(200)).len(), 80); + } + + #[test] + fn tools_are_found_on_path_then_asked_for_their_version() { + let exec = FakeExec { + calls: Mutex::default(), + }; + let t = tool_versions(&exec); + assert_eq!(t.len(), PROBE_TOOLS.len()); + assert_eq!(t["git"], "git version 2.45.0"); + assert_eq!(t["gh"], "gh version 2.50.0 (2026-01-01)"); + assert_eq!(t["go"], "go version go1.25.0 linux/amd64"); + // On PATH but silent: present, version unknown. + assert_eq!(t["jq"], "unknown"); + assert_eq!(t["node"], "missing"); + assert_eq!(t["cosign"], "missing"); + // A missing tool is never executed. + assert!( + !exec + .calls + .lock() + .unwrap() + .iter() + .any(|c| c.starts_with("node ")) + ); + } + + #[test] + fn the_tools_probe_is_cached_for_a_minute() { + let exec = FakeExec { + calls: Mutex::default(), + }; + let cache = ToolsCache::default(); + let first = cache.get(&exec, 1000).unwrap(); + let n = exec.calls.lock().unwrap().len(); + assert_eq!(cache.get(&exec, 1059).unwrap(), first); + assert_eq!(exec.calls.lock().unwrap().len(), n); + cache.get(&exec, 1060).unwrap(); + assert!(exec.calls.lock().unwrap().len() > n); + } + + #[test] + fn psi_is_the_higher_of_host_and_cgroup_pressure() { + let mut h = linux(); + h.files.insert( + "/proc/pressure/cpu", + "some avg10=12.34 avg60=1.00 avg300=1.00 total=1\nfull avg10=0.00\n", + ); + assert_eq!(cpu_psi_pm(&h), Some(123)); + h.files.insert( + "/sys/fs/cgroup/cpu.pressure", + "some avg10=45.60 avg60=1.00 avg300=1.00 total=1\n", + ); + assert_eq!(cpu_psi_pm(&h), Some(456)); + assert_eq!(cpu_psi_pm(&linux()), None); + assert_eq!(parse_psi_some_avg10_pm("full avg10=5.00\n"), None); + assert_eq!( + parse_psi_some_avg10_pm("some avg10=100.00 avg60=0\n"), + Some(1000) + ); + } + const MEMINFO_28G: &str = "MemTotal: 29360128 kB\nMemFree: 1 kB\n"; #[test] diff --git a/src/queen.rs b/src/queen.rs index 2ddf1dc..976ead4 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -21,26 +21,29 @@ use rand::RngCore; use rand::rngs::OsRng; use serde::Serialize; use time::OffsetDateTime; +use time::format_description::well_known::Rfc3339; use crate::allocator::{ - ActiveLease, AllocatorInput, ClassIdle, Demand, DroneState, Grant, LeaseState, QueueAge, - Reservation, Unmet, WarmFloor, allocate, max_slots, + ActiveLease, AllocatorInput, ClassIdle, Demand, DroneAdmission, DroneState, Grant, LeaseState, + Promotion, QueueAge, Reservation, Unmet, WarmFloor, allocate, free_room, free_slots, max_slots, }; use crate::bid::{ClassPressure, ScoreWeights}; use crate::classes::{Class, Scope}; -use crate::config::QueenConfig; +use crate::config::{QueenConfig, StandbySettings}; use crate::envcontract::blocked_classes; use crate::github::{GithubClient, GithubError, QueuedJob, Runner, WorkflowJob}; use crate::health::{DroneHealth, DroneInfo, Outcome, OutcomeKind, evaluate}; use crate::homeostasis::{Controller, KeyInput, KeyReport}; use crate::janitor::Janitor; +use crate::preempt::{PreemptInput, Running, select as select_victims}; use crate::probe::Capacity; use crate::profile::{ClassProfile, build_profiles}; +use crate::requeue::Requeuer; use crate::rollout::{ Check as RolloutCheck, DroneView, Observed as RolloutObserved, RolloutState, advance, desired_for, }; -use crate::store::{LeaseRow, LeaseStatus, Store, StoreError}; +use crate::store::{LeaseRow, LeaseStatus, PreemptionRow, Store, StoreError}; use crate::tailnet::TailnetEnroll; use crate::telemetry::{Conclusion, FailureClass, MAX_REPORT_ROWS, TelemetryRow}; use crate::units::{Unit, UnitState}; @@ -72,6 +75,8 @@ pub struct Snapshot { pub health_systemic: bool, /// Slots each class could still be granted, from the last pass. pub idle: Vec, + /// CPU-pressure admission per fresh drone from the last pass. + pub admission: Vec, /// Warm-floor controller state per key (empty when off). pub homeostasis: Vec, /// Learned per-drone, per-class profiles (empty when off). @@ -81,6 +86,9 @@ pub struct Snapshot { pub wake: WakeView, /// Non-borg online runners by label, from the last pass. pub static_runners: BTreeMap, + /// Standby leases the drones are told to start now (deploy demand is queued + /// for them); read by the heartbeat. + pub start_standby: BTreeSet, /// `unmet` minus what idle non-borg runners with the label can take. pub unmet_placeable: Vec, } @@ -157,12 +165,14 @@ pub struct Queen { github: GithubClient, snapshot: RwLock, janitor: Option, + requeuer: Option, tailnet: Option, homeostasis: std::sync::Mutex, profile_cache: std::sync::Mutex)>>, wake: std::sync::Mutex, unit_joins: AtomicU64, unit_receipts: AtomicU64, + units_superseded: AtomicU64, wake_sender: Arc, tick_lock: tokio::sync::Mutex<()>, } @@ -235,6 +245,93 @@ fn new_lease_id() -> String { hex::encode(b) } +fn as_active(l: &LeaseRow) -> ActiveLease { + ActiveLease { + drone_id: l.drone.clone(), + class: l.class, + scope: l.scope, + state: match l.state { + LeaseStatus::Running => LeaseState::Running, + LeaseStatus::Standby => LeaseState::Standby, + _ => LeaseState::Pending, + }, + } +} + +/// Where a preemptible lease's job lives: repo, run id, job id. +type JobRef<'a> = (&'a str, u64, u64); + +/// Split the active leases into those whose in-progress job is known (the +/// job's `runner_name` is the lease's runner) and everything else. +fn split_running<'a>( + active: &[LeaseRow], + unit_leases: &[ActiveLease], + observed: &[(&'a str, WorkflowJob)], + preempted: &HashSet<&str>, +) -> (Vec, Vec, HashMap>) { + let (mut running, mut others, mut jobs) = (Vec::new(), unit_leases.to_vec(), HashMap::new()); + for l in active { + let found = observed.iter().find_map(|(repo, j)| { + let started = OffsetDateTime::parse(j.started_at.as_deref()?, &Rfc3339).ok()?; + (j.status == "in_progress" && j.runner_name.as_deref() == Some(&l.runner_name)) + .then_some((*repo, j.run_id, j.id, started.unix_timestamp())) + }); + match found { + Some((repo, run_id, job_id, started_at)) => { + jobs.insert(l.id.clone(), (repo, run_id, job_id)); + running.push(Running { + lease_id: l.id.clone(), + drone: l.drone.clone(), + class: l.class, + scope: l.scope, + started_at, + preempted: preempted.contains(l.id.as_str()), + }); + } + None => others.push(as_active(l)), + } + } + (running, others, jobs) +} + +/// What standby minting needs to know about the fleet. +struct StandbyView<'a> { + active: &'a [LeaseRow], + drones: &'a [DroneState], + capacities: &'a HashMap, + extra_labels: &'a HashMap>, + serves_deploy: &'a BTreeSet, +} + +/// The standby leases to start now: per promoted class, the oldest standby +/// leases on fresh drones. +fn promotions( + promote: &[Promotion], + active: &[LeaseRow], + fresh: &HashMap, +) -> BTreeSet { + let mut out = BTreeSet::new(); + for p in promote { + let mut standby: Vec<&LeaseRow> = active + .iter() + .filter(|l| { + l.state == LeaseStatus::Standby + && l.scope == Scope::Deploy + && l.class == p.class + && fresh.contains_key(&l.drone) + }) + .collect(); + standby.sort_by(|a, b| (a.created_at, &a.id).cmp(&(b.created_at, &b.id))); + out.extend( + standby + .into_iter() + .take(usize::try_from(p.count).unwrap_or(usize::MAX)) + .map(|l| l.id.clone()), + ); + } + out +} + fn age_secs(now: OffsetDateTime, then: OffsetDateTime) -> u64 { u64::try_from((now - then).whole_seconds()).unwrap_or(0) } @@ -247,17 +344,27 @@ impl Queen { github, snapshot: RwLock::new(Snapshot::default()), janitor: None, + requeuer: None, tailnet: None, homeostasis: std::sync::Mutex::new(Controller::new()), profile_cache: std::sync::Mutex::new(None), wake: std::sync::Mutex::new(WakePlanner::new()), unit_joins: AtomicU64::new(0), unit_receipts: AtomicU64::new(0), + units_superseded: AtomicU64::new(0), wake_sender: Arc::new(UdpWakeSender), tick_lock: tokio::sync::Mutex::new(()), } } + /// Attach the preemption requeuer; `api::serve` runs it after each + /// janitor pass. + #[must_use] + pub fn with_requeuer(mut self, requeuer: Requeuer) -> Self { + self.requeuer = Some(requeuer); + self + } + /// Attach the opt-in supersede janitor; `api::serve` runs it. #[must_use] pub fn with_janitor(mut self, janitor: Janitor) -> Self { @@ -287,6 +394,17 @@ impl Queen { self.unit_receipts.fetch_add(1, Ordering::Relaxed); } + pub fn count_units_superseded(&self, n: usize) { + self.units_superseded + .fetch_add(u64::try_from(n).unwrap_or(u64::MAX), Ordering::Relaxed); + } + + /// Units cancelled by a newer submission with the same supersede key since + /// start. + pub fn units_superseded(&self) -> u64 { + self.units_superseded.load(Ordering::Relaxed) + } + /// Submissions that joined a running unit / reused a receipt since start. pub fn unit_reuse_counts(&self) -> (u64, u64) { ( @@ -303,6 +421,10 @@ impl Queen { self.janitor.as_ref() } + pub fn requeuer(&self) -> Option<&Requeuer> { + self.requeuer.as_ref() + } + pub fn config(&self) -> &QueenConfig { &self.cfg } @@ -372,6 +494,7 @@ impl Queen { let mut fresh_names: BTreeSet = BTreeSet::new(); let mut capacities: HashMap = HashMap::new(); let mut extra_labels: HashMap> = HashMap::new(); + let mut serves_deploy: BTreeSet = BTreeSet::new(); let mut drones = Vec::new(); for d in &drone_rows { let fresh = d @@ -389,6 +512,13 @@ impl Queen { None, now_ts, ); + if let Some(a) = &self.cfg.admission { + state.cpu_psi_pm = d + .env_facts + .as_ref() + .filter(|f| now_ts - f.observed_at <= a.facts_max_age_secs) + .and_then(|f| f.cpu_psi_some_avg10_pm); + } state.profiles = profiles .iter() .filter(|p| p.drone == d.name) @@ -396,6 +526,9 @@ impl Queen { .collect(); drones.push(state); capacities.insert(d.name.clone(), cap.clone()); + if d.trust.serves_deploy() { + serves_deploy.insert(d.name.clone()); + } extra_labels.insert( d.name.clone(), self.cfg.classes.sanitize_extra_labels(&d.extra_labels), @@ -408,17 +541,8 @@ impl Queen { let unit_leases = self.unit_leases().await; let leases = active .iter() - .map(|l| ActiveLease { - drone_id: l.drone.clone(), - class: l.class, - scope: l.scope, - state: if l.state == LeaseStatus::Running { - LeaseState::Running - } else { - LeaseState::Pending - }, - }) - .chain(unit_leases) + .map(as_active) + .chain(unit_leases.iter().cloned()) .collect::>(); let schedule_floors = self .cfg @@ -427,6 +551,7 @@ impl Queen { let (warm_floors, homeostasis) = self.regulate(now_ts, &schedule_floors, &queue); let placement_drones = drones.clone(); let placement_leases = leases.clone(); + let preempt_ages = queue_ages.clone(); let allocation = allocate(&AllocatorInput { classes: self.cfg.classes.clone(), demand, @@ -436,6 +561,7 @@ impl Queen { queue_ages, starvation_after_secs: self.cfg.starvation_after_secs, weights: ScoreWeights::default(), + admission: self.cfg.admission.as_ref().map(|a| a.thresholds()), }); if let Err(e) = self @@ -443,7 +569,7 @@ impl Queen { now_ts, &placement_drones, &capacities, - placement_leases, + placement_leases.clone(), &allocation.grants, ) .await @@ -502,6 +628,83 @@ impl Queen { } } + let start_standby = promotions(&allocation.promote, &active, &capacities); + if mint_error.is_none() + && let Some(st) = &self.cfg.standby + { + let mut held = placement_leases.clone(); + held.extend(allocation.grants.iter().map(|g| ActiveLease { + drone_id: g.drone_id.clone(), + class: g.class, + scope: g.scope, + state: LeaseState::Pending, + })); + mint_error = self + .top_up_standby( + st, + &StandbyView { + active: &active, + drones: &placement_drones, + capacities: &capacities, + extra_labels: &extra_labels, + serves_deploy: &serves_deploy, + }, + held, + &prefix, + now_ts, + ) + .await?; + } + + if let Some(p) = &self.cfg.preemption { + let preempted = self.store.preempted_active().await?; + let (running, others, jobs) = split_running( + &active, + &unit_leases, + &observed, + &preempted.iter().map(|(id, _)| id.as_str()).collect(), + ); + let victims = select_victims(&PreemptInput { + classes: &self.cfg.classes, + settings: p, + drones: &placement_drones, + others: &others, + running: &running, + unmet: &unmet_placeable, + queue_ages: &preempt_ages, + now: now_ts, + }); + for v in victims { + let Some((repo, run_id, job_id)) = jobs.get(&v.lease_id) else { + continue; + }; + let row = PreemptionRow { + lease_id: v.lease_id, + drone: v.drone, + class: v.class, + scope: v.scope, + for_scope: v.for_scope, + repo: (*repo).to_owned(), + run_id: *run_id, + job_id: *job_id, + at: now_ts, + requeued_at: None, + }; + if self.store.record_preemption(&row, now_ts).await? { + tracing::warn!( + "preempting {} lease {} on {} (job {} of {}) for starving {} {} demand", + row.scope.as_str(), + row.lease_id, + row.drone, + row.job_id, + row.repo, + v.for_class, + row.for_scope.as_str(), + ); + } + } + } + self.update_snapshot(|s| { s.last_tick_at = Some(now_ts); s.last_error = mint_error.as_ref().map(|e| format!("mint failed: {e}")); @@ -511,10 +714,12 @@ impl Queen { s.warm_floors = warm_floors; s.pressure = allocation.pressure.clone(); s.idle = allocation.idle.clone(); + s.admission = allocation.admission.clone(); s.homeostasis = homeostasis; s.profiles = profiles; s.wake = wake; s.static_runners = statics; + s.start_standby = start_standby; s.unmet_placeable = unmet_placeable; }); if let Err(e) = self @@ -543,6 +748,95 @@ impl Queen { Ok(()) } + /// Mint standby deploy runners until every drone that serves deploy work + /// holds `per_drone` of each configured class. Returns a mint error to + /// surface in the snapshot; store failures abort the tick like grants do. + async fn top_up_standby( + &self, + st: &StandbySettings, + view: &StandbyView<'_>, + mut held: Vec, + prefix: &str, + now_ts: i64, + ) -> Result, QueenError> { + for drone in view.drones { + let (Some(cap), true) = ( + view.capacities.get(&drone.id), + view.serves_deploy.contains(&drone.id), + ) else { + continue; + }; + for &class in &st.classes { + if drone.blocked_classes.contains(&class) { + continue; + } + let have = view + .active + .iter() + .filter(|l| { + l.state == LeaseStatus::Standby && l.drone == drone.id && l.class == class + }) + .count(); + let spec = self.cfg.classes.spec(class); + for _ in have..usize::try_from(st.per_drone).unwrap_or(usize::MAX) { + let room = free_room(drone, &held, &self.cfg.classes); + if free_slots(drone, &room, spec) == 0 { + break; + } + let lease_id = new_lease_id(); + let runner_name = format!("{prefix}{}-{lease_id}", drone.id); + let extra = view + .extra_labels + .get(&drone.id) + .map_or(&[][..], Vec::as_slice); + let labels = jit_labels( + &cap.os, + &cap.arch, + extra, + &self.cfg.classes.runner_labels(class, Scope::Deploy), + ); + let minted = match self + .github + .mint_jit_runner( + &runner_name, + self.cfg.github.runner_group_id, + &labels, + &self.cfg.work_folder, + ) + .await + { + Ok(m) => m, + Err(e) => return Ok(Some(e)), + }; + let row = LeaseRow { + id: lease_id, + drone: drone.id.clone(), + class, + scope: Scope::Deploy, + labels, + runner_name, + runner_id: minted.runner_id, + jit_config: Some(minted.jit_config), + state: LeaseStatus::Standby, + created_at: now_ts, + expires_at: i64::MAX, + }; + if let Err(e) = self.store.insert_lease(&row).await { + let _ = self.github.delete_runner(minted.runner_id).await; + return Err(e.into()); + } + held.push(ActiveLease { + drone_id: drone.id.clone(), + class, + scope: Scope::Deploy, + state: LeaseState::Standby, + }); + } + } + } + Ok(None) + } + /// The version `drone` should run: from the gated rollout when it is on, /// otherwise the configured target (report-only mode), otherwise no /// instruction. @@ -639,7 +933,7 @@ impl Queen { Some(ActiveLease { drone_id: u.drone?, class: u.class, - scope: Scope::Pr, + scope: u.scope, state: LeaseState::Running, }) }) @@ -761,17 +1055,17 @@ impl Queen { if group.is_empty() { continue; } - let mut demand: BTreeMap = BTreeMap::new(); + let mut demand: BTreeMap<(Class, Scope), u32> = BTreeMap::new(); for u in &units { - *demand.entry(u.class).or_default() += 1; + *demand.entry((u.class, u.scope)).or_default() += 1; } let allocation = allocate(&AllocatorInput { classes: self.cfg.classes.clone(), demand: demand .into_iter() - .map(|(class, count)| Demand { + .map(|((class, scope), count)| Demand { class, - scope: Scope::Pr, + scope, count, }) .collect(), @@ -781,18 +1075,25 @@ impl Queen { queue_ages: Vec::new(), starvation_after_secs: None, weights: ScoreWeights::default(), + admission: self.cfg.admission.as_ref().map(|a| a.thresholds()), }); - // Grants of a class are interchangeable: hand them to the units - // of that class in submission order. + // Grants of a (class, scope) are interchangeable: hand them out + // deploy first, then main, then pull request, each in submission + // order. let mut grants: Vec<&Grant> = allocation.grants.iter().collect(); - for u in units { - if let Some(i) = grants.iter().position(|g| g.class == u.class) { + let mut ordered = units; + ordered.sort_by_key(|u| u.scope); + for u in ordered { + if let Some(i) = grants + .iter() + .position(|g| g.class == u.class && g.scope == u.scope) + { let g = grants.remove(i); if self.store.lease_unit(&u.id, &g.drone_id, now).await? { base.push(ActiveLease { drone_id: g.drone_id.clone(), class: g.class, - scope: Scope::Pr, + scope: g.scope, state: LeaseState::Running, }); tracing::info!("unit {} placed on {}", u.id, g.drone_id); @@ -1073,15 +1374,18 @@ impl Queen { for lease in self.store.active_leases().await? { let runner = ours.get(lease.runner_name.as_str()); let overdue = now - lease.created_at > self.cfg.lease_max_secs; + let standby = lease.state == LeaseStatus::Standby; + // A standby waits for its drone indefinitely: only `lease_max_secs` + // rotates it. let never_came_online = - now > lease.expires_at && runner.is_none_or(|r| r.status != "online"); + !standby && now > lease.expires_at && runner.is_none_or(|r| r.status != "online"); match runner { _ if overdue => { self.end(&lease, runner.copied(), LeaseStatus::Expired, now) .await? } Some(r) if r.busy => { - if lease.state == LeaseStatus::Pending { + if matches!(lease.state, LeaseStatus::Pending | LeaseStatus::Standby) { self.store .set_state(&lease.id, LeaseStatus::Running, now) .await?; @@ -1089,7 +1393,9 @@ impl Queen { still_active.push(lease); } // Ephemeral runners deregister after their one job. - None if lease.state == LeaseStatus::Running => { + // A standby whose registration is gone was started and has + // finished (or GitHub dropped it): the next tick mints another. + None if matches!(lease.state, LeaseStatus::Running | LeaseStatus::Standby) => { self.store .set_state(&lease.id, LeaseStatus::Done, now) .await?; @@ -1229,4 +1535,49 @@ mod tests { "extra labels sit between the platform and class labels" ); } + + fn lease_row(id: &str, runner: &str) -> LeaseRow { + LeaseRow { + id: id.to_owned(), + drone: "host-a".to_owned(), + class: Class::Medium, + scope: Scope::Pr, + labels: Vec::new(), + runner_name: runner.to_owned(), + runner_id: 1, + jit_config: None, + state: LeaseStatus::Running, + created_at: 0, + expires_at: 0, + } + } + + fn job(id: u64, status: &str, runner: &str) -> WorkflowJob { + serde_json::from_value(serde_json::json!({ + "id": id, "run_id": 77, "status": status, "created_at": "2026-01-01T00:00:00Z", + "started_at": "2026-01-01T00:10:00Z", "runner_name": runner, + })) + .unwrap() + } + + #[test] + fn only_leases_with_an_in_progress_job_on_their_runner_are_preemptible() { + let active = [ + lease_row("l1", "borg-host-a-l1"), + lease_row("l2", "borg-host-a-l2"), + lease_row("l3", "borg-host-a-l3"), + ]; + let observed = [ + ("example-org/app", job(5, "in_progress", "borg-host-a-l1")), + ("example-org/app", job(6, "queued", "borg-host-a-l2")), + ]; + let preempted = HashSet::from(["l1"]); + let (running, others, jobs) = split_running(&active, &[], &observed, &preempted); + assert_eq!(running.len(), 1); + assert_eq!(running[0].lease_id, "l1"); + assert_eq!(running[0].started_at, 1_767_226_200); + assert!(running[0].preempted); + assert_eq!(others.len(), 2); + assert_eq!(jobs["l1"], ("example-org/app", 77, 5)); + } } diff --git a/src/requeue.rs b/src/requeue.rs new file mode 100644 index 0000000..0183a59 --- /dev/null +++ b/src/requeue.rs @@ -0,0 +1,206 @@ +//! Re-runs jobs the queen preempted. +//! +//! Preemption cancels a running job so a higher-priority one can have its +//! slot. The queen records `(repo, run, job)` in the `preemptions` table with +//! `requeued_at` unset. This pass re-runs each recorded job through the jobs +//! rerun API once its run has completed (GitHub refuses to re-run a job of a +//! run that is still in progress). It uses the supersede janitor's token, +//! which has `actions: write`, and the janitor's `dry_run` and +//! `max_cancels_per_pass` settings. +//! +//! * a run that is not `completed` is left for a later pass; +//! * a rerun answered 403, 404 or 409 (job gone, retention expired, or already +//! re-run) is terminal: the row is marked and a warning is logged; +//! * any other error is logged and retried on the next pass; +//! * a row older than 24 hours is marked without a call, so a backlog never +//! re-runs ancient jobs; +//! * in dry-run nothing is called and nothing is marked. + +use std::sync::Mutex; +use std::time::Duration; + +use reqwest::{Method, RequestBuilder}; +use serde::{Deserialize, Serialize}; + +use crate::config::JanitorSettings; +use crate::github::GithubError; +use crate::janitor::status_error; +use crate::secret::Secret; +use crate::store::{Store, StoreError}; + +/// Rows older than this are marked stale instead of re-run. +const MAX_AGE_SECS: i64 = 24 * 3600; +/// Most unrequeued rows read per pass. +const BATCH: i64 = 200; + +/// What one pass did. +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)] +pub struct PassReport { + /// Preemption rows examined. + pub checked: u32, + /// Jobs re-run. + pub rerun: u32, + /// Rows whose run has not completed yet. + pub waiting: u32, + /// Rows older than 24 hours, marked without a call. + pub stale: u32, + /// Rows whose rerun answered 403, 404 or 409. + pub terminal: u32, + /// Failures that will be retried next pass. + pub errors: u32, +} + +impl PassReport { + fn add(&mut self, o: &PassReport) { + self.checked += o.checked; + self.rerun += o.rerun; + self.waiting += o.waiting; + self.stale += o.stale; + self.terminal += o.terminal; + self.errors += o.errors; + } +} + +/// The requeuer as `/status` and `/metrics` show it. +#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)] +pub struct RequeueView { + pub passes: u64, + pub last: PassReport, + /// Sums over every pass since the queen started. + pub total: PassReport, +} + +#[derive(Debug, Deserialize)] +struct Run { + status: String, +} + +pub struct Requeuer { + http: reqwest::Client, + api_base: String, + token: Secret, + dry_run: bool, + max_per_pass: u32, + view: Mutex, +} + +impl Requeuer { + /// Built from the janitor's token and settings. + pub fn new( + api_base: &str, + token: Secret, + settings: &JanitorSettings, + ) -> Result { + let http = reqwest::Client::builder() + .user_agent("borg-queen-requeue") + .timeout(Duration::from_secs(30)) + .build()?; + Ok(Self { + http, + api_base: api_base.trim_end_matches('/').to_owned(), + token, + dry_run: settings.dry_run, + max_per_pass: settings.max_cancels_per_pass, + view: Mutex::default(), + }) + } + + pub fn view(&self) -> RequeueView { + self.view.lock().expect("requeue view lock").clone() + } + + fn request(&self, method: Method, path: &str) -> RequestBuilder { + self.http + .request(method, format!("{}{path}", self.api_base)) + .header("Accept", "application/vnd.github+json") + .header("X-GitHub-Api-Version", "2022-11-28") + .bearer_auth(self.token.expose()) + } + + async fn run_status(&self, repo: &str, run: u64) -> Result { + let path = format!("/repos/{repo}/actions/runs/{run}"); + let resp = self.request(Method::GET, &path).send().await?; + let status = resp.status(); + if !status.is_success() { + return Err(status_error("GET", &path, status.as_u16(), resp).await); + } + let bytes = resp.bytes().await?; + serde_json::from_slice::(&bytes) + .map(|r| r.status) + .map_err(|e| GithubError::Decode(format!("{path}: {e}"))) + } + + /// Re-run a job. `Ok(Some(code))` is a terminal refusal (403, 404, 409). + async fn rerun(&self, repo: &str, job: u64) -> Result, GithubError> { + let path = format!("/repos/{repo}/actions/jobs/{job}/rerun"); + let resp = self.request(Method::POST, &path).send().await?; + match resp.status().as_u16() { + 200..=299 => Ok(None), + code @ (403 | 404 | 409) => Ok(Some(code)), + code => Err(status_error("POST", &path, code, resp).await), + } + } + + /// One pass over the preemptions not yet re-run. + pub async fn pass(&self, store: &Store, now: i64) -> Result { + let mut report = PassReport::default(); + let mut budget = self.max_per_pass; + for row in store.unrequeued_preemptions(BATCH).await? { + report.checked += 1; + let (repo, run, job) = (row.repo.as_str(), row.run_id, row.job_id); + if now - row.at > MAX_AGE_SECS { + tracing::info!("requeue: {repo} job {job} (run {run}) is stale, not re-running"); + store.mark_requeued(&row.lease_id, now).await?; + report.stale += 1; + continue; + } + if budget == 0 { + break; + } + if self.dry_run { + tracing::info!( + "requeue (dry-run): would rerun {repo} job {job} once run {run} completes" + ); + budget -= 1; + continue; + } + match self.run_status(repo, run).await { + Ok(status) if status == "completed" => {} + Ok(_) => { + report.waiting += 1; + continue; + } + Err(e) => { + tracing::warn!("requeue: {repo} run {run}: {e}"); + report.errors += 1; + continue; + } + } + budget -= 1; + match self.rerun(repo, job).await { + Ok(None) => { + tracing::info!("requeue: re-ran {repo} job {job} (run {run})"); + store.mark_requeued(&row.lease_id, now).await?; + report.rerun += 1; + } + Ok(Some(code)) => { + tracing::warn!( + "requeue: {repo} job {job} (run {run}) rerun answered {code}, giving up" + ); + store.mark_requeued(&row.lease_id, now).await?; + report.terminal += 1; + } + Err(e) => { + tracing::warn!("requeue: {repo} job {job} (run {run}): {e}"); + report.errors += 1; + } + } + } + let mut view = self.view.lock().expect("requeue view lock"); + view.passes += 1; + view.total.add(&report); + view.last = report.clone(); + drop(view); + Ok(report) + } +} diff --git a/src/run.rs b/src/run.rs index 4f4dde7..0b5e6e8 100644 --- a/src/run.rs +++ b/src/run.rs @@ -55,6 +55,8 @@ pub struct RunOptions { pub toolchain: String, pub timeout_secs: Option, pub force: bool, + /// Unit priority scope: `deploy`, `main` or `pr`. + pub scope: String, /// How long to wait for a drone to pick the unit up before running locally. pub wait_secs: u64, pub publish: bool, @@ -347,7 +349,7 @@ async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result let body = serde_json::json!({ "repo": repo, "tree_sha": tree, "argv": opts.argv, "platform": opts.platform, "toolchain": opts.toolchain, "class": opts.class, - "timeout_secs": opts.timeout_secs, "force": opts.force, "bundle": needs_bundle, + "timeout_secs": opts.timeout_secs, "force": opts.force, "bundle": needs_bundle, "scope": opts.scope, "rust_toolchain": tree_rust_toolchain(&opts.dir, &tree), }); let resp = q diff --git a/src/store.rs b/src/store.rs index 3c0b897..30f54bb 100644 --- a/src/store.rs +++ b/src/store.rs @@ -9,13 +9,13 @@ use std::str::FromStr; use serde::Serialize; use serde::de::DeserializeOwned; -use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; -use sqlx::{Row, SqlitePool}; +use sqlx::sqlite::{SqliteConnectOptions, SqliteConnection, SqliteJournalMode, SqlitePoolOptions}; +use sqlx::{ConnectOptions, Connection, Row, SqlitePool}; use crate::allocator::{DroneState, HostCaps}; use crate::api::CacheStat; use crate::bid::DroneBid; -use crate::classes::{Class, Scope}; +use crate::classes::{Class, Scope, Trust}; use crate::envcontract::EnvFacts; use crate::health::{Action, Hold}; use crate::probe::Capacity; @@ -45,7 +45,8 @@ CREATE TABLE IF NOT EXISTS drones ( health_cleared_at INTEGER, version_since INTEGER, reclaim TEXT, - env_facts TEXT + env_facts TEXT, + trust TEXT NOT NULL DEFAULT 'shared' ); CREATE TABLE IF NOT EXISTS reclaim_log ( drone TEXT NOT NULL, @@ -165,6 +166,18 @@ CREATE TABLE IF NOT EXISTS job_telemetry ( finished_at INTEGER NOT NULL ); CREATE INDEX IF NOT EXISTS job_telemetry_finished ON job_telemetry(finished_at); +CREATE TABLE IF NOT EXISTS preemptions ( + lease_id TEXT PRIMARY KEY, + drone TEXT NOT NULL, + class TEXT NOT NULL, + scope TEXT NOT NULL, + for_scope TEXT NOT NULL, + repo TEXT NOT NULL, + run_id INTEGER NOT NULL, + job_id INTEGER NOT NULL, + at INTEGER NOT NULL, + requeued_at INTEGER +); "; #[derive(Debug, thiserror::Error)] @@ -180,6 +193,9 @@ pub enum StoreError { #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] #[serde(rename_all = "lowercase")] pub enum LeaseStatus { + /// A deploy runner the drone holds unstarted: started when deploy demand + /// is queued or the queen is unreachable. Never expires by lease TTL. + Standby, /// Runner minted, no job yet. Pending, /// Runner is executing a job. @@ -191,6 +207,7 @@ pub enum LeaseStatus { impl LeaseStatus { pub fn as_str(self) -> &'static str { match self { + Self::Standby => "standby", Self::Pending => "pending", Self::Running => "running", Self::Done => "done", @@ -200,6 +217,7 @@ impl LeaseStatus { fn parse(s: &str) -> Result { Ok(match s { + "standby" => Self::Standby, "pending" => Self::Pending, "running" => Self::Running, "done" => Self::Done, @@ -209,7 +227,7 @@ impl LeaseStatus { } pub fn is_active(self) -> bool { - matches!(self, Self::Pending | Self::Running) + matches!(self, Self::Standby | Self::Pending | Self::Running) } } @@ -227,6 +245,8 @@ pub struct DroneRow { /// The borg version the drone last reported, if it reports one. pub version: Option, pub extra_labels: Vec, + /// The trust the drone last reported. + pub trust: Trust, /// The drone's last reported bids (may have expired; see `expires_at`). pub bids: Vec, /// `operator` or `auto` while quarantined. @@ -284,6 +304,7 @@ impl DroneRow { pub fn allocator_state(&self, cap: &Capacity, now: i64) -> DroneState { let mut s = DroneState::from_capacity(self.name.clone(), cap, &self.caps); s.accepting = self.schedulable(); + s.trust = self.trust; if self.on_probation() { s.max_leases = Some(s.max_leases.map_or(1, |m| m.min(1))); } @@ -332,6 +353,24 @@ pub struct LeaseRow { pub expires_at: i64, } +/// A lease cancelled for higher-priority demand, and the job it was running. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct PreemptionRow { + pub lease_id: String, + pub drone: String, + pub class: Class, + /// The cancelled lease's scope. + pub scope: Scope, + /// The starving demand's scope. + pub for_scope: Scope, + pub repo: String, + pub run_id: u64, + pub job_id: u64, + pub at: i64, + /// When the job was re-run; unset until then. + pub requeued_at: Option, +} + #[derive(Debug, Clone)] pub struct LeaseResult { pub exit_code: i32, @@ -369,13 +408,21 @@ impl Store { .filename(path) .create_if_missing(true) .journal_mode(SqliteJournalMode::Wal); + // Migrate on a dedicated connection first so every pooled connection + // prepares its statements against the final schema. + let mut conn = opts.connect().await?; + Self::migrate(&mut conn).await?; + conn.close().await?; Self::connect(opts, 4).await } pub async fn open_in_memory() -> Result { // One connection: every pooled in-memory connection is its own database. let opts = SqliteConnectOptions::from_str("sqlite::memory:")?; - Self::connect(opts, 1).await + let store = Self::connect(opts, 1).await?; + // The pool has one connection, so migrating through it is safe. + Self::migrate(&mut *store.pool.acquire().await?).await?; + Ok(store) } async fn connect(opts: SqliteConnectOptions, max: u32) -> Result { @@ -383,7 +430,11 @@ impl Store { .max_connections(max) .connect_with(opts) .await?; - sqlx::raw_sql(SCHEMA).execute(&pool).await?; + Ok(Self { pool }) + } + + async fn migrate(conn: &mut SqliteConnection) -> Result<(), StoreError> { + sqlx::raw_sql(SCHEMA).execute(&mut *conn).await?; // Databases created before these columns existed lack them. for ddl in [ "ALTER TABLE drones ADD COLUMN cache_stats TEXT", @@ -400,15 +451,21 @@ impl Store { "ALTER TABLE drones ADD COLUMN version_since INTEGER", "ALTER TABLE drones ADD COLUMN reclaim TEXT", "ALTER TABLE drones ADD COLUMN env_facts TEXT", + "ALTER TABLE leases ADD COLUMN preempted_at INTEGER", + "ALTER TABLE drones ADD COLUMN trust TEXT NOT NULL DEFAULT 'shared'", + "ALTER TABLE units ADD COLUMN scope TEXT NOT NULL DEFAULT 'pr'", + "ALTER TABLE units ADD COLUMN supersede_key TEXT", + "ALTER TABLE units ADD COLUMN record INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE clients ADD COLUMN max_scope TEXT NOT NULL DEFAULT 'pr'", ] { // sqlx-runtime-query: one-off additive migration, no data involved - if let Err(e) = sqlx::query(ddl).execute(&pool).await + if let Err(e) = sqlx::query(ddl).execute(&mut *conn).await && !e.to_string().contains("duplicate column") { return Err(e.into()); } } - Ok(Self { pool }) + Ok(()) } pub async fn enroll(&self, name: &str, token_hash: &str, now: i64) -> Result<(), StoreError> { @@ -473,6 +530,17 @@ impl Store { Ok(()) } + /// The trust the drone reports for itself. + pub async fn set_trust(&self, name: &str, trust: Trust) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query("UPDATE drones SET trust = ? WHERE name = ?") + .bind(trust.as_str()) + .bind(name) + .execute(&self.pool) + .await?; + Ok(()) + } + /// Delete a drone (its token hash with it) and its lease records. Returns /// the number of unfinished leases dropped, or `None` if there is no such /// drone. With the lease rows gone the queen no longer treats those @@ -511,7 +579,7 @@ impl Store { pub async fn drones(&self) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let rows = sqlx::query( - "SELECT name, last_heartbeat, capacity, caps, accepting, cache_stats, quarantined, quarantine_reason, version, extra_labels, bids, quarantine_source, quarantined_at, probation_since, quarantine_strikes, health_cleared_at, version_since, reclaim, env_facts FROM drones ORDER BY name", + "SELECT name, last_heartbeat, capacity, caps, accepting, cache_stats, quarantined, quarantine_reason, version, extra_labels, bids, quarantine_source, quarantined_at, probation_since, quarantine_strikes, health_cleared_at, version_since, reclaim, env_facts, trust FROM drones ORDER BY name", ) .fetch_all(&self.pool) .await?; @@ -544,6 +612,7 @@ impl Store { .map(from_text) .transpose()? .unwrap_or_default(), + trust: enum_from(&r.get::("trust"))?, quarantine_source: r.get("quarantine_source"), quarantined_at: r.get("quarantined_at"), probation_since: r.get("probation_since"), @@ -929,7 +998,7 @@ impl Store { pub async fn active_leases(&self) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let rows = sqlx::query( - "SELECT * FROM leases WHERE state IN ('pending', 'running') ORDER BY created_at, id", + "SELECT id, drone, class, scope, labels, runner_name, runner_id, jit_config, state, created_at, expires_at FROM leases WHERE state IN ('standby', 'pending', 'running') ORDER BY created_at, id", ) .fetch_all(&self.pool) .await?; @@ -938,7 +1007,7 @@ impl Store { pub async fn lease(&self, id: &str) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database - let row = sqlx::query("SELECT * FROM leases WHERE id = ?") + let row = sqlx::query("SELECT id, drone, class, scope, labels, runner_name, runner_id, jit_config, state, created_at, expires_at FROM leases WHERE id = ?") .bind(id) .fetch_optional(&self.pool) .await?; @@ -951,7 +1020,7 @@ impl Store { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database sqlx::query( "UPDATE leases SET state = ?, finished_at = COALESCE(?, finished_at), jit_config = CASE WHEN ? THEN NULL ELSE jit_config END - WHERE id = ? AND state IN ('pending', 'running')", + WHERE id = ? AND state IN ('standby', 'pending', 'running')", ) .bind(to.as_str()) .bind(finished) @@ -984,7 +1053,7 @@ impl Store { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database sqlx::query( "UPDATE leases SET state = 'done', finished_at = ?, jit_config = NULL, exit_code = ?, duration_secs = ?, peak_rss_mib = ? - WHERE id = ? AND state IN ('pending', 'running', 'expired')", + WHERE id = ? AND state IN ('standby', 'pending', 'running', 'expired')", ) .bind(now) .bind(i64::from(result.exit_code)) @@ -996,6 +1065,122 @@ impl Store { Ok(true) } + /// Mark an active lease preempted and record the job it was running, once. + /// Returns false if it was already marked or is no longer active. + pub async fn record_preemption(&self, p: &PreemptionRow, now: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let marked = sqlx::query( + "UPDATE leases SET preempted_at = ? WHERE id = ? AND preempted_at IS NULL AND state IN ('pending', 'running')", + ) + .bind(now) + .bind(&p.lease_id) + .execute(&self.pool) + .await? + .rows_affected(); + if marked == 0 { + return Ok(false); + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT OR IGNORE INTO preemptions (lease_id, drone, class, scope, for_scope, repo, run_id, job_id, at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&p.lease_id) + .bind(&p.drone) + .bind(enum_text(&p.class)?) + .bind(enum_text(&p.scope)?) + .bind(enum_text(&p.for_scope)?) + .bind(&p.repo) + .bind(i64::try_from(p.run_id).unwrap_or(i64::MAX)) + .bind(i64::try_from(p.job_id).unwrap_or(i64::MAX)) + .bind(now) + .execute(&self.pool) + .await?; + Ok(true) + } + + /// `(lease id, drone)` of preempted leases that have not ended: the drone + /// is told to cancel each on every heartbeat until the lease finishes. + pub async fn preempted_active(&self) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT id, drone FROM leases WHERE preempted_at IS NOT NULL AND state IN ('pending', 'running') ORDER BY id", + ) + .fetch_all(&self.pool) + .await?; + Ok(rows.iter().map(|r| (r.get("id"), r.get("drone"))).collect()) + } + + /// The newest `limit` preemptions, newest first. + pub async fn preemptions(&self, limit: i64) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT lease_id, drone, class, scope, for_scope, repo, run_id, job_id, at, requeued_at FROM preemptions ORDER BY at DESC, lease_id LIMIT ?") + .bind(limit) + .fetch_all(&self.pool) + .await?; + rows.iter().map(Self::preemption_from_row).collect() + } + + /// Up to `limit` preemptions whose job has not been re-run, oldest first. + pub async fn unrequeued_preemptions( + &self, + limit: i64, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query("SELECT lease_id, drone, class, scope, for_scope, repo, run_id, job_id, at, requeued_at FROM preemptions WHERE requeued_at IS NULL ORDER BY at, lease_id LIMIT ?") + .bind(limit) + .fetch_all(&self.pool) + .await?; + rows.iter().map(Self::preemption_from_row).collect() + } + + /// Record that a preempted job needs no further requeue attempts. + pub async fn mark_requeued(&self, lease_id: &str, now: i64) -> Result<(), StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "UPDATE preemptions SET requeued_at = ? WHERE lease_id = ? AND requeued_at IS NULL", + ) + .bind(now) + .bind(lease_id) + .execute(&self.pool) + .await?; + Ok(()) + } + + fn preemption_from_row(r: &sqlx::sqlite::SqliteRow) -> Result { + Ok(PreemptionRow { + lease_id: r.get("lease_id"), + drone: r.get("drone"), + class: enum_from(&r.get::("class"))?, + scope: enum_from(&r.get::("scope"))?, + for_scope: enum_from(&r.get::("for_scope"))?, + repo: r.get("repo"), + run_id: u64::try_from(r.get::("run_id")).unwrap_or(0), + job_id: u64::try_from(r.get::("job_id")).unwrap_or(0), + at: r.get("at"), + requeued_at: r.get("requeued_at"), + }) + } + + /// Preemption counts by `(victim scope, starving scope)`, for metrics. + pub async fn preemption_counts(&self) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT scope, for_scope, COUNT(*) AS n FROM preemptions GROUP BY scope, for_scope ORDER BY scope, for_scope", + ) + .fetch_all(&self.pool) + .await?; + rows.iter() + .map(|r| { + Ok(( + enum_from(&r.get::("scope"))?, + enum_from(&r.get::("for_scope"))?, + u64::try_from(r.get::("n")).unwrap_or(0), + )) + }) + .collect() + } + /// Drop the JIT config of leases `drone` says it holds. pub async fn acknowledge(&self, drone: &str, lease_ids: &[String]) -> Result<(), StoreError> { for id in lease_ids { @@ -1013,7 +1198,7 @@ impl Store { pub async fn undelivered(&self, drone: &str) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let rows = sqlx::query( - "SELECT * FROM leases WHERE drone = ? AND state = 'pending' AND jit_config IS NOT NULL ORDER BY created_at, id", + "SELECT id, drone, class, scope, labels, runner_name, runner_id, jit_config, state, created_at, expires_at FROM leases WHERE drone = ? AND state = 'pending' AND jit_config IS NOT NULL ORDER BY created_at, id", ) .bind(drone) .fetch_all(&self.pool) @@ -1021,6 +1206,45 @@ impl Store { rows.iter().map(Self::lease_from_row).collect() } + /// Unstarted standby leases of `drone`, oldest first. + pub async fn standby_leases(&self, drone: &str) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query( + "SELECT id, drone, class, scope, labels, runner_name, runner_id, jit_config, state, created_at, expires_at FROM leases WHERE drone = ? AND state = 'standby' ORDER BY created_at, id", + ) + .bind(drone) + .fetch_all(&self.pool) + .await?; + rows.iter().map(Self::lease_from_row).collect() + } + + /// Turn the standby leases `ids` of `drone` into pending ones: the drone + /// started them. The usual TTL applies from `now`, and the JIT config is + /// dropped. Returns how many changed. + pub async fn promote_standby( + &self, + drone: &str, + ids: &[String], + now: i64, + ttl_secs: i64, + ) -> Result { + let mut changed = 0; + for id in ids { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + changed += sqlx::query( + "UPDATE leases SET state = 'pending', expires_at = ?, jit_config = NULL + WHERE id = ? AND drone = ? AND state = 'standby'", + ) + .bind(now.saturating_add(ttl_secs)) + .bind(id) + .bind(drone) + .execute(&self.pool) + .await? + .rows_affected(); + } + Ok(changed) + } + /// Lease counts by (class, state), for metrics. pub async fn lease_counts(&self) -> Result, StoreError> { let rows = @@ -1066,6 +1290,74 @@ mod tests { s } + /// A database created before the newest columns existed must open, migrate + /// and read back without sqlx tripping over a stale column-name map. + #[tokio::test] + async fn opens_a_pre_migration_database_and_migrates_idempotently() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("borg.db"); + let opts = SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal); + let mut conn = opts.connect().await.unwrap(); + sqlx::raw_sql( + " +CREATE TABLE drones ( + name TEXT PRIMARY KEY, + token_hash TEXT NOT NULL UNIQUE, + enrolled_at INTEGER NOT NULL, + last_heartbeat INTEGER, + capacity TEXT, + caps TEXT, + accepting INTEGER NOT NULL DEFAULT 1 +); +CREATE TABLE leases ( + id TEXT PRIMARY KEY, + drone TEXT NOT NULL REFERENCES drones(name), + class TEXT NOT NULL, + scope TEXT NOT NULL, + labels TEXT NOT NULL, + runner_name TEXT NOT NULL, + runner_id INTEGER NOT NULL, + jit_config TEXT, + state TEXT NOT NULL, + created_at INTEGER NOT NULL, + expires_at INTEGER NOT NULL, + finished_at INTEGER, + exit_code INTEGER, + duration_secs INTEGER, + peak_rss_mib INTEGER +); +INSERT INTO drones (name, token_hash, enrolled_at) VALUES ('d1', 'hash-d1', 1); +INSERT INTO leases (id, drone, class, scope, labels, runner_name, runner_id, jit_config, state, created_at, expires_at) +VALUES ('l1', 'd1', 'medium', 'pr', '[\"a\"]', 'borg-d1-l1', 7, 'jit', 'pending', 100, 400); +", + ) + .execute(&mut conn) + .await + .unwrap(); + conn.close().await.unwrap(); + + for _ in 0..2 { + let s = Store::open(&path).await.unwrap(); + // Concurrent reads warm several pooled connections. + let (a, b, c, d) = tokio::join!( + s.active_leases(), + s.active_leases(), + s.active_leases(), + s.drones() + ); + assert_eq!(a.unwrap().len(), 1); + assert_eq!(b.unwrap().len(), 1); + assert_eq!(c.unwrap().len(), 1); + let drones = d.unwrap(); + assert_eq!(drones.len(), 1); + assert_eq!(drones[0].name, "d1"); + assert_eq!(s.lease("l1").await.unwrap().unwrap().runner_id, 7); + } + } + #[tokio::test] async fn enrollment_is_unique_and_lookup_is_by_hash() { let s = store_with("d1").await; @@ -1117,6 +1409,126 @@ mod tests { assert!(s.active_leases().await.unwrap().is_empty()); } + #[tokio::test] + async fn preemption_is_recorded_once_and_listed_until_the_lease_ends() { + let s = store_with("host-a").await; + s.insert_lease(&lease("l1", "host-a")).await.unwrap(); + let row = PreemptionRow { + lease_id: "l1".into(), + drone: "host-a".into(), + class: Class::Medium, + scope: Scope::Pr, + for_scope: Scope::Main, + repo: "example-org/app".into(), + run_id: 5, + job_id: 6, + at: 0, + requeued_at: None, + }; + assert!(s.record_preemption(&row, 500).await.unwrap()); + assert!(!s.record_preemption(&row, 501).await.unwrap()); + assert_eq!( + s.preempted_active().await.unwrap(), + [("l1".to_owned(), "host-a".to_owned())] + ); + assert_eq!( + s.preemptions(20).await.unwrap(), + [PreemptionRow { at: 500, ..row }] + ); + assert_eq!( + s.preemption_counts().await.unwrap(), + [(Scope::Pr, Scope::Main, 1)] + ); + s.set_state("l1", LeaseStatus::Done, 600).await.unwrap(); + assert!(s.preempted_active().await.unwrap().is_empty()); + } + + #[tokio::test] + async fn unrequeued_preemptions_list_oldest_first_until_marked() { + let s = store_with("host-a").await; + for (id, at) in [("l1", 300), ("l2", 100), ("l3", 200)] { + s.insert_lease(&lease(id, "host-a")).await.unwrap(); + let row = PreemptionRow { + lease_id: id.into(), + drone: "host-a".into(), + class: Class::Medium, + scope: Scope::Pr, + for_scope: Scope::Main, + repo: "example-org/app".into(), + run_id: 5, + job_id: 6, + at: 0, + requeued_at: None, + }; + assert!(s.record_preemption(&row, at).await.unwrap()); + } + let ids = |rows: Vec| -> Vec { + rows.into_iter().map(|r| r.lease_id).collect() + }; + assert_eq!( + ids(s.unrequeued_preemptions(10).await.unwrap()), + ["l2", "l3", "l1"] + ); + assert_eq!( + ids(s.unrequeued_preemptions(2).await.unwrap()), + ["l2", "l3"] + ); + s.mark_requeued("l2", 900).await.unwrap(); + s.mark_requeued("l2", 950).await.unwrap(); + assert_eq!( + ids(s.unrequeued_preemptions(10).await.unwrap()), + ["l3", "l1"] + ); + let all = s.preemptions(10).await.unwrap(); + let l2 = all.iter().find(|r| r.lease_id == "l2").unwrap(); + assert_eq!(l2.requeued_at, Some(900), "the first mark sticks"); + } + + #[tokio::test] + async fn standby_leases_round_trip_stay_active_and_promote() { + let s = store_with("d1").await; + let mut l = lease("sb1", "d1"); + l.state = LeaseStatus::Standby; + l.expires_at = i64::MAX; + s.insert_lease(&l).await.unwrap(); + s.insert_lease(&lease("p1", "d1")).await.unwrap(); + let got = s.lease("sb1").await.unwrap().unwrap(); + assert_eq!(got.state, LeaseStatus::Standby); + assert_eq!(got.expires_at, i64::MAX); + assert_eq!(s.active_leases().await.unwrap().len(), 2); + let standby = s.standby_leases("d1").await.unwrap(); + assert_eq!(standby.len(), 1); + assert!(standby[0].jit_config.is_some()); + // A standby is not delivered as an ordinary pending lease. + let undelivered = s.undelivered("d1").await.unwrap(); + assert!(undelivered.iter().all(|u| u.id != "sb1")); + assert_eq!( + s.promote_standby("d1", &["sb1".into()], 1000, 60) + .await + .unwrap(), + 1 + ); + let got = s.lease("sb1").await.unwrap().unwrap(); + assert_eq!((got.state, got.expires_at), (LeaseStatus::Pending, 1060)); + assert!(got.jit_config.is_none()); + assert_eq!( + s.promote_standby("d1", &["sb1".into()], 1000, 60) + .await + .unwrap(), + 0 + ); + // A standby can end without ever starting. + let mut l2 = lease("sb2", "d1"); + l2.state = LeaseStatus::Standby; + s.insert_lease(&l2).await.unwrap(); + s.set_state("sb2", LeaseStatus::Done, 1100).await.unwrap(); + assert_eq!( + s.lease("sb2").await.unwrap().unwrap().state, + LeaseStatus::Done + ); + assert!(s.standby_leases("d1").await.unwrap().is_empty()); + } + #[tokio::test] async fn set_state_never_revives_a_finished_lease() { let s = store_with("d1").await; @@ -1257,7 +1669,11 @@ mod tests { s.set_extra_labels("d1", &["on-prem".to_owned()]) .await .unwrap(); + assert_eq!(s.drones().await.unwrap()[0].trust, Trust::Shared); + s.set_trust("d1", Trust::Deploy).await.unwrap(); let d = &s.drones().await.unwrap()[0]; + assert_eq!(d.trust, Trust::Deploy); + assert_eq!(d.allocator_state(&cap, 0).trust, Trust::Deploy); assert_eq!((d.last_heartbeat, d.accepting), (Some(123), false)); assert_eq!(d.cache_stats, [stat]); assert_eq!(d.extra_labels, ["on-prem"]); diff --git a/src/units.rs b/src/units.rs index cf734a7..986cb78 100644 --- a/src/units.rs +++ b/src/units.rs @@ -32,7 +32,7 @@ use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use sqlx::Row; -use crate::classes::Class; +use crate::classes::{Class, Scope}; use crate::store::{Store, StoreError}; #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] @@ -169,6 +169,25 @@ pub struct SubmitUnit { /// (`PUT /v1/units/{id}/bundle`); the unit is not placed until it arrives. #[serde(default)] pub bundle: bool, + /// Priority scope: `deploy`, `main` or `pr` (default). Decides the + /// capacity the unit competes for and the drones it may run on. + #[serde(default = "default_scope")] + pub scope: Scope, + /// Newest wins: a new submission with this key from the same client + /// cancels that client's older queued or running units with the same key, + /// unless they were submitted with `record`. + #[serde(default)] + pub supersede_key: Option, + /// The result is wanted as a record: this unit is never superseded. + #[serde(default)] + pub record: bool, +} + +/// Longest `supersede_key`. +pub const MAX_SUPERSEDE_KEY: usize = 128; + +fn default_scope() -> Scope { + Scope::Pr } fn default_class() -> Class { @@ -215,6 +234,13 @@ impl SubmitUnit { if self.toolchain.len() > 200 || self.toolchain.chars().any(char::is_control) { return Err("toolchain must be at most 200 printable characters".into()); } + if let Some(k) = &self.supersede_key + && (k.is_empty() || k.len() > MAX_SUPERSEDE_KEY || k.chars().any(char::is_control)) + { + return Err(format!( + "supersede_key must be 1 to {MAX_SUPERSEDE_KEY} printable characters" + )); + } Ok(()) } @@ -234,6 +260,11 @@ impl SubmitUnit { for a in &self.argv { field(a); } + // Work of a higher scope never shares a unit or a receipt with a lower + // one. Pull-request units add nothing, so existing keys stay valid. + if self.scope != Scope::Pr { + field(self.scope.as_str()); + } hex::encode(h.finalize()) } } @@ -268,6 +299,9 @@ pub struct Unit { pub awaiting_bundle: bool, /// Size of the uploaded bundle, 0 if none. pub bundle_bytes: u64, + pub scope: Scope, + pub supersede_key: Option, + pub record: bool, } /// The proof a unit succeeded for exactly this tree. @@ -416,6 +450,10 @@ fn unit_from_row(r: &sqlx::sqlite::SqliteRow) -> Result { log_bytes: u64::try_from(r.get::("log_bytes")).unwrap_or(0), awaiting_bundle: r.get::("awaiting_bundle") != 0, bundle_bytes: u64::try_from(r.get::("bundle_bytes")).unwrap_or(0), + scope: serde_json::from_str(&format!("\"{}\"", r.get::("scope"))) + .map_err(|e| StoreError::Corrupt(e.to_string()))?, + supersede_key: r.get("supersede_key"), + record: r.get::("record") != 0, }) } @@ -427,14 +465,29 @@ impl Store { name: &str, token_hash: &str, now: i64, + ) -> Result<(), StoreError> { + self.enroll_client_scoped(name, token_hash, Scope::Pr, now) + .await + } + + /// Enroll a client that may submit units up to `max_scope`. + pub async fn enroll_client_scoped( + &self, + name: &str, + token_hash: &str, + max_scope: Scope, + now: i64, ) -> Result<(), StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database - let r = sqlx::query("INSERT INTO clients (name, token_hash, created_at) VALUES (?, ?, ?)") - .bind(name) - .bind(token_hash) - .bind(now) - .execute(&self.pool) - .await; + let r = sqlx::query( + "INSERT INTO clients (name, token_hash, created_at, max_scope) VALUES (?, ?, ?, ?)", + ) + .bind(name) + .bind(token_hash) + .bind(now) + .bind(max_scope.as_str()) + .execute(&self.pool) + .await; match r { Ok(_) => Ok(()), Err(sqlx::Error::Database(e)) if e.is_unique_violation() => { @@ -456,6 +509,69 @@ impl Store { Ok(row.map(|r| r.get("name"))) } + /// The highest scope a client may submit; `pr` for an unknown client. + pub async fn client_max_scope(&self, name: &str) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query("SELECT max_scope FROM clients WHERE name = ?") + .bind(name) + .fetch_optional(&self.pool) + .await?; + let Some(row) = row else { + return Ok(Scope::Pr); + }; + serde_json::from_str(&format!("\"{}\"", row.get::("max_scope"))) + .map_err(|e| StoreError::Corrupt(e.to_string())) + } + + /// Cancel `client`'s queued and running non-record units that carry + /// `key`, except `except`. Queued ones end at once; running ones are + /// flagged for their drone's next heartbeat. Returns the ids. + pub async fn supersede_units( + &self, + client: &str, + key: &str, + except: &str, + now: i64, + ) -> Result, StoreError> { + let sql = format!( + "SELECT id FROM units WHERE client = ? AND supersede_key = ? AND record = 0 AND id != ? AND state IN {ACTIVE} AND cancel_requested = 0 ORDER BY created_at, id" + ); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query(&sql) + .bind(client) + .bind(key) + .bind(except) + .fetch_all(&self.pool) + .await?; + let why = format!("superseded by {except}"); + let mut out = Vec::new(); + for r in rows { + let id: String = r.get("id"); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let ended = sqlx::query("UPDATE units SET state = 'failed', failure = ?, finished_at = ?, cancel_requested = 1 WHERE id = ? AND state IN ('queued', 'leased')") + .bind(&why) + .bind(now) + .bind(&id) + .execute(&self.pool) + .await?; + if ended.rows_affected() > 0 { + self.delete_unit_bundle(&id).await?; + out.push(id); + continue; + } + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let flagged = sqlx::query("UPDATE units SET cancel_requested = 1, failure = ? WHERE id = ? AND state = 'running'") + .bind(&why) + .bind(&id) + .execute(&self.pool) + .await?; + if flagged.rows_affected() > 0 { + out.push(id); + } + } + Ok(out) + } + pub async fn unit(&self, id: &str) -> Result, StoreError> { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let row = sqlx::query("SELECT * FROM units WHERE id = ?") @@ -502,7 +618,7 @@ impl Store { return Ok(existing); } // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database - let inserted = sqlx::query("INSERT INTO units (id, key, repo, tree_sha, argv, platform, toolchain, rust_toolchain, class, timeout_secs, client, state, attempt, created_at, awaiting_bundle) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?, ?)") + let inserted = sqlx::query("INSERT INTO units (id, key, repo, tree_sha, argv, platform, toolchain, rust_toolchain, class, timeout_secs, client, state, attempt, created_at, awaiting_bundle, scope, supersede_key, record) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?, ?, ?, ?, ?)") .bind(id) .bind(&key) .bind(&s.repo) @@ -516,6 +632,9 @@ impl Store { .bind(client) .bind(now) .bind(i64::from(s.bundle)) + .bind(s.scope.as_str()) + .bind(s.supersede_key.as_deref()) + .bind(i64::from(s.record)) .execute(&self.pool) .await; match inserted { @@ -942,6 +1061,9 @@ mod tests { force: false, rust_toolchain: None, bundle: false, + scope: Scope::Pr, + supersede_key: None, + record: false, } } @@ -982,6 +1104,32 @@ mod tests { ] { assert_ne!(changed.key(), k, "{changed:?}"); } + // Work of a higher scope is its own unit; pull-request keys are unchanged. + for scope in [Scope::Main, Scope::Deploy] { + assert_ne!(SubmitUnit { scope, ..submit() }.key(), k); + } + assert_ne!( + SubmitUnit { + scope: Scope::Main, + ..submit() + } + .key(), + SubmitUnit { + scope: Scope::Deploy, + ..submit() + } + .key() + ); + // Superseding and record do not change identity either. + assert_eq!( + SubmitUnit { + supersede_key: Some("pr-1".into()), + record: true, + ..submit() + } + .key(), + k + ); // Case of the object id is not part of the identity; class and force are not either. let upper = SubmitUnit { tree_sha: "A".repeat(40), @@ -1235,4 +1383,68 @@ mod tests { ); assert!(s.client_for_token_hash("nope").await.unwrap().is_none()); } + + #[tokio::test] + async fn scope_supersede_key_record_and_client_maximum_are_stored() { + let s = store().await; + s.enroll_client("plain", "h1", 1).await.unwrap(); + s.enroll_client_scoped("gate", "h2", Scope::Deploy, 1) + .await + .unwrap(); + assert_eq!(s.client_max_scope("plain").await.unwrap(), Scope::Pr); + assert_eq!(s.client_max_scope("gate").await.unwrap(), Scope::Deploy); + assert_eq!(s.client_max_scope("unknown").await.unwrap(), Scope::Pr); + + let mut a = submit(); + a.scope = Scope::Main; + a.supersede_key = Some("pr-1".into()); + s.submit_unit(&a, "gate", "u1", 60, 100, 10).await.unwrap(); + let u = s.unit("u1").await.unwrap().unwrap(); + assert_eq!( + (u.scope, u.supersede_key.as_deref(), u.record), + (Scope::Main, Some("pr-1"), false) + ); + let plain = SubmitUnit { + tree_sha: "b".repeat(40), + ..submit() + }; + s.submit_unit(&plain, "gate", "u2", 60, 100, 11) + .await + .unwrap(); + let u = s.unit("u2").await.unwrap().unwrap(); + assert_eq!( + (u.scope, u.supersede_key, u.record), + (Scope::Pr, None, false) + ); + + // Records and other clients' units with the same key survive. + let keep = SubmitUnit { + tree_sha: "c".repeat(40), + supersede_key: Some("pr-1".into()), + record: true, + ..submit() + }; + s.submit_unit(&keep, "gate", "u3", 60, 100, 12) + .await + .unwrap(); + let theirs = SubmitUnit { + tree_sha: "d".repeat(40), + supersede_key: Some("pr-1".into()), + ..submit() + }; + s.submit_unit(&theirs, "plain", "u4", 60, 100, 13) + .await + .unwrap(); + let hit = s.supersede_units("gate", "pr-1", "u9", 20).await.unwrap(); + assert_eq!(hit, ["u1"]); + let u = s.unit("u1").await.unwrap().unwrap(); + assert_eq!(u.state, UnitState::Failed); + assert_eq!(u.failure.as_deref(), Some("superseded by u9")); + assert!( + s.supersede_units("gate", "pr-1", "u9", 21) + .await + .unwrap() + .is_empty() + ); + } } diff --git a/src/units_api.rs b/src/units_api.rs index bf9ce60..f64f614 100644 --- a/src/units_api.rs +++ b/src/units_api.rs @@ -72,6 +72,10 @@ pub struct SubmitResponse { pub unit_id: String, pub state: String, pub receipt: Option, + /// Older units of the same client and `supersede_key` this submission + /// cancelled. + #[serde(default)] + pub superseded: Vec, } async fn submit( @@ -82,6 +86,14 @@ async fn submit( enabled(&queen)?; let client = authenticate_client(&queen, &headers).await?; req.validate().map_err(ApiError::BadRequest)?; + let max_scope = queen.store().client_max_scope(&client).await?; + if req.scope < max_scope { + return Err(ApiError::ForbiddenWith(format!( + "this client may submit units up to scope {} but asked for {}", + max_scope.as_str(), + req.scope.as_str() + ))); + } let cfg = &queen.config().units; let timeout = req .timeout_secs @@ -103,8 +115,15 @@ async fn submit( let submitted = store .submit_unit(&req, &client, &id, timeout, cfg.receipt_ttl_secs, t) .await?; + let mut superseded = Vec::new(); let (outcome, unit, status) = match submitted { - Submitted::New(u) => ("new", u, StatusCode::CREATED), + Submitted::New(u) => { + if let Some(key) = &req.supersede_key { + superseded = store.supersede_units(&client, key, &u.id, t).await?; + queen.count_units_superseded(superseded.len()); + } + ("new", u, StatusCode::CREATED) + } Submitted::Joined(u) => { queen.count_unit_join(); ("joined", u, StatusCode::OK) @@ -127,6 +146,7 @@ async fn submit( unit_id: unit.id.clone(), state: unit.state.as_str().to_owned(), receipt: unit.receipt(), + superseded, }), )) } diff --git a/tests/contract.rs b/tests/contract.rs index fe18dca..716a476 100644 --- a/tests/contract.rs +++ b/tests/contract.rs @@ -285,3 +285,87 @@ async fn the_contract_is_off_by_default_and_require_report_is_opt_in() { .is_empty() ); } + +#[tokio::test] +async fn a_drone_missing_a_required_tool_gets_no_lease_and_status_names_it() { + let h = Harness::start(cfg(json!({ + "env_contract": { + "enabled": true, + "checks": {"tools": "required"}, + "required_tools": ["git", "gh", "node"], + "classes": {"heavy": {"required_tools": ["git"]}} + } + }))) + .await; + // "a-bare" would win every name tie; it came up without gh and node. + let bare = h.enroll("a-bare").await; + let full = h.enroll("b-full").await; + let mut f = good("1.95.0"); + f["tools"] = + json!({"git": "git version 2.45.0", "gh": "missing", "node": "missing", "jq": "jq-1.7"}); + beat(&h, &bare, Some(f)).await; + let mut f = good("1.95.0"); + f["tools"] = json!({"git": "git version 2.45.0", "gh": "gh version 2.50.0", "node": "v22.1.0"}); + beat(&h, &full, Some(f)).await; + + h.queue_job(1, "borg-ci-light-pr", 5); + h.tick().await; + assert!( + beat(&h, &bare, None).await["leases"] + .as_array() + .unwrap() + .is_empty() + ); + assert_eq!( + beat(&h, &full, None).await["leases"] + .as_array() + .unwrap() + .len(), + 1 + ); + + let s = h.status().await; + let c = &drone(&s, "a-bare")["contract"]; + assert_eq!(c["admission"]["light"]["admit"], false); + assert!( + c["admission"]["light"]["why"] + .as_str() + .unwrap() + .contains("required tools missing on PATH: gh, node") + ); + // The heavy class only asks for git. + assert_eq!(c["admission"]["heavy"]["admit"], true); + let find_tools = |c: &Value| { + c["checks"] + .as_array() + .unwrap() + .iter() + .find(|x| x["kind"] == "tools") + .unwrap() + .clone() + }; + let t = find_tools(c); + assert_eq!( + (t["state"].as_str(), t["mode"].as_str()), + (Some("fail"), Some("required")) + ); + assert_eq!(t["observed"]["missing"], "gh,node"); + assert!( + t["observed"]["versions"] + .as_str() + .unwrap() + .contains("git=git version 2.45.0") + ); + assert_eq!( + find_tools(&drone(&s, "b-full")["contract"])["state"], + "pass" + ); + + let m = h.metrics().await; + for line in [ + "borg_env_contract_check{drone=\"a-bare\",check=\"tools\",state=\"fail\"} 1", + "borg_env_contract_check{drone=\"b-full\",check=\"tools\",state=\"pass\"} 1", + ] { + assert!(m.contains(line), "missing {line}\n{m}"); + } +} diff --git a/tests/janitor.rs b/tests/janitor.rs index d5dc208..886b007 100644 --- a/tests/janitor.rs +++ b/tests/janitor.rs @@ -3,6 +3,14 @@ mod support; use serde_json::json; use support::Harness; +fn with(flags: serde_json::Value) -> serde_json::Value { + let mut j = json!({"dry_run": false}); + j.as_object_mut() + .unwrap() + .extend(flags.as_object().unwrap().clone()); + json!({"supersede_janitor": j}) +} + fn enabled(dry_run: bool, max: u32) -> serde_json::Value { json!({"supersede_janitor": {"dry_run": dry_run, "max_cancels_per_pass": max}}) } @@ -114,3 +122,124 @@ async fn the_janitor_uses_its_own_token() { h.queue_pr_run(1, 10, "old", "new"); assert_eq!(h.janitor_pass().await.cancelled, 1); } + +#[tokio::test] +async fn running_stale_runs_are_cancelled_only_with_cancel_running() { + let off = Harness::start(enabled(false, 20)).await; + off.state().pulls.insert(10, "new".into()); + off.add_run(1, "in_progress", "pull_request", "old", Some(10)); + let r = off.janitor_pass().await; + assert_eq!((r.considered, r.cancelled_running), (0, 0), "{r:?}"); + assert!(off.state().cancelled.is_empty()); + + let on = Harness::start(with(json!({"cancel_running": true}))).await; + on.state().pulls.insert(10, "new".into()); + on.add_run(1, "in_progress", "pull_request", "old", Some(10)); + on.add_run(2, "in_progress", "pull_request", "new", Some(10)); + on.add_run(3, "in_progress", "pull_request", "old", Some(10)); + on.state().fresh_status.insert(3, "completed".into()); + let r = on.janitor_pass().await; + assert_eq!( + (r.superseded, r.cancelled, r.cancelled_running, r.raced), + (2, 0, 1, 1), + "{r:?}" + ); + assert_eq!(on.state().cancelled, [1]); +} + +#[tokio::test] +async fn queued_duplicates_keep_the_lowest_id() { + let h = Harness::start(with(json!({"dedupe_same_sha": true}))).await; + h.state().pulls.insert(10, "sha".into()); + for id in [9, 4, 6] { + h.add_run(id, "queued", "pull_request", "sha", Some(10)); + } + let r = h.janitor_pass().await; + assert_eq!( + (r.duplicates, r.dedup_cancelled, r.cancelled), + (2, 2, 0), + "{r:?}" + ); + let mut cancelled = h.state().cancelled.clone(); + cancelled.sort_unstable(); + assert_eq!(cancelled, [6, 9]); + let r = h.janitor_pass().await; + assert_eq!(r.duplicates, 0, "{r:?}"); +} + +#[tokio::test] +async fn a_running_duplicate_is_kept_and_never_cancelled() { + let h = Harness::start(with(json!({"dedupe_same_sha": true}))).await; + h.state().pulls.insert(10, "sha".into()); + h.add_run(5, "queued", "pull_request", "sha", Some(10)); + h.add_run(7, "in_progress", "pull_request", "sha", Some(10)); + h.add_run(8, "in_progress", "pull_request", "sha", Some(10)); + h.add_run(9, "queued", "pull_request", "sha", Some(10)); + let r = h.janitor_pass().await; + assert_eq!((r.duplicates, r.dedup_cancelled), (3, 2), "{r:?}"); + let mut cancelled = h.state().cancelled.clone(); + cancelled.sort_unstable(); + assert_eq!(cancelled, [5, 9]); +} + +#[tokio::test] +async fn dispatch_and_merge_group_duplicates_are_untouched() { + let h = Harness::start(with(json!({"dedupe_same_sha": true}))).await; + for id in 1..=3 { + h.add_run(id, "queued", "workflow_dispatch", "sha", None); + } + for id in 4..=5 { + h.add_run(id, "queued", "merge_group", "sha", None); + } + let r = h.janitor_pass().await; + assert_eq!((r.duplicates, r.dedup_cancelled), (0, 0), "{r:?}"); + assert!(h.state().cancelled.is_empty()); +} + +#[tokio::test] +async fn dry_run_counts_both_kinds_and_calls_no_cancel() { + let h = Harness::start(json!({"supersede_janitor": + {"cancel_running": true, "dedupe_same_sha": true}})) + .await; + h.state().pulls.insert(10, "new".into()); + h.add_run(1, "in_progress", "pull_request", "old", Some(10)); + h.add_run(2, "queued", "pull_request", "new", Some(10)); + h.add_run(3, "queued", "pull_request", "new", Some(10)); + let r = h.janitor_pass().await; + assert_eq!( + ( + r.would_cancel, + r.cancelled, + r.cancelled_running, + r.dedup_cancelled + ), + (2, 0, 0, 0), + "{r:?}" + ); + assert!(h.state().cancelled.is_empty()); +} + +#[tokio::test] +async fn totals_reach_status_and_metrics() { + let h = Harness::start(enabled(false, 20)).await; + h.queue_pr_run(1, 10, "old", "new"); + h.janitor_pass().await; + h.janitor_pass().await; + let status = h.status().await; + assert_eq!(status["janitor"]["passes"], 2); + assert_eq!(status["janitor"]["total"]["cancelled"], 1); + assert_eq!(status["janitor"]["last"]["cancelled"], 0); + let metrics = h + .http + .get(format!("{}/metrics", h.url)) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + assert!( + metrics.contains("borg_janitor_cancelled_total 1"), + "{metrics}" + ); +} diff --git a/tests/queen.rs b/tests/queen.rs index dac71ad..9f4a8f6 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -1028,3 +1028,216 @@ async fn a_timed_out_feedback_lookup_is_bounded_and_retried_later() { 1 ); } + +#[tokio::test] +async fn starving_main_demand_preempts_one_running_pr_lease_and_tells_its_drone() { + let h = Harness::start(json!({ + "preemption": {"after_secs": 600, "max_per_pass": 1, "min_victim_age_secs": 60} + })) + .await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-medium-pr", 1); + h.queue_job(2, "borg-ci-medium-pr", 1); + h.tick().await; + h.set_runner("borg-d1-", "online", true); + let started = (OffsetDateTime::now_utc() - time::Duration::minutes(10)) + .format(&time::format_description::well_known::Rfc3339) + .unwrap(); + let running: Vec = h + .state() + .mints + .iter() + .enumerate() + .map(|(i, m)| { + json!({"id": 100 + i, "run_id": 9, "status": "in_progress", "labels": [], + "created_at": started, "started_at": started, "runner_name": m["name"]}) + }) + .collect(); + h.state().jobs = running; + h.queue_job(3, "borg-ci-medium", 700); + + h.tick().await; + h.tick().await; + + let rows = h.queen.store().preemptions(20).await.unwrap(); + assert_eq!(rows.len(), 1, "one cancellation covers the demand"); + assert_eq!((rows[0].run_id, rows[0].job_id / 100), (1, 1)); + let body = h.heartbeat_with(&token, json!({})).await; + assert_eq!(body["cancel_leases"], json!([rows[0].lease_id])); + let status = h.status().await; + assert_eq!( + status["preemptions"][0]["lease_id"], + json!(rows[0].lease_id) + ); + assert!( + h.metrics() + .await + .contains("borg_preemptions_total{from_scope=\"pr\",to_scope=\"main\"} 1") + ); +} + +#[tokio::test] +async fn nothing_is_preempted_unless_configured() { + let h = Harness::start(json!({})).await; + let token = h.enroll("d1").await; + h.heartbeat(&token, &[]).await; + h.queue_job(1, "borg-ci-medium", 5000); + h.tick().await; + assert!(h.queen.store().preemptions(20).await.unwrap().is_empty()); + let body = h.heartbeat_with(&token, json!({})).await; + assert_eq!(body["cancel_leases"], json!([])); +} + +fn standby_cfg() -> Value { + json!({"standby": {"classes": ["medium"], "per_drone": 1}}) +} + +async fn standby_beat(h: &Harness, token: &str, held: &[&str], active: &[&str]) -> Value { + h.heartbeat_with( + token, + json!({"trust": "deploy-only", "held_standby": held, "active_leases": active}), + ) + .await +} + +#[tokio::test] +async fn standby_runners_are_minted_for_deploy_drones_only_and_never_expire_by_ttl() { + let mut cfg = standby_cfg(); + cfg["lease_ttl_secs"] = json!(300); + let h = Harness::start(cfg).await; + let dep = h.enroll("dep").await; + let shared = h.enroll("shr").await; + standby_beat(&h, &dep, &[], &[]).await; + h.heartbeat(&shared, &[]).await; + + h.tick().await; + + let mints = h.state().mints.clone(); + assert_eq!(mints.len(), 1, "one standby for the deploy-only drone"); + assert!(mints[0]["name"].as_str().unwrap().starts_with("borg-dep-")); + assert_eq!( + labels(&mints[0]), + ["self-hosted", "Linux", "X64", "borg-ci-medium-deploy"], + "deploy labels only" + ); + let body = standby_beat(&h, &dep, &[], &[]).await; + assert_eq!(body["leases"], json!([]), "not an ordinary lease"); + assert_eq!(body["start_standby"], json!([])); + let grants = body["standby"].as_array().unwrap(); + assert_eq!(grants.len(), 1); + assert_eq!(grants[0]["scope"], "deploy"); + let name = mints[0]["name"].as_str().unwrap(); + assert_eq!(grants[0]["jit_config"], format!("jit-{name}")); + let id = grants[0]["lease_id"].as_str().unwrap().to_owned(); + let shared_body = h.heartbeat_with(&shared, json!({})).await; + assert_eq!(shared_body["standby"], json!([])); + + // Acknowledged by listing it as held: not resent. + let body = standby_beat(&h, &dep, &[&id], &[]).await; + assert_eq!(body["standby"], json!([])); + + // Far beyond the pending-lease TTL the standby is still there, its + // registration is not deleted and no replacement is minted. + let runner_id = h.state().runners[0]["id"].as_u64().unwrap(); + let later = OffsetDateTime::now_utc() + time::Duration::seconds(301); + h.queen.tick(later).await.unwrap(); + assert!(!h.state().deleted.contains(&runner_id)); + assert_eq!(h.state().mints.len(), 1); + let lease = h.queen.store().lease(&id).await.unwrap().unwrap(); + assert_eq!(lease.state, borg::store::LeaseStatus::Standby); + // The reaper treats the offline standby registration as owned. + assert_eq!(h.state().runners.len(), 1); +} + +#[tokio::test] +async fn a_lost_standby_config_is_resent_and_a_forgotten_one_is_dropped() { + let h = Harness::start(standby_cfg()).await; + let dep = h.enroll("dep").await; + standby_beat(&h, &dep, &[], &[]).await; + h.tick().await; + let id = standby_beat(&h, &dep, &[], &[]).await["standby"][0]["lease_id"] + .as_str() + .unwrap() + .to_owned(); + // The drone lost its copy: the queen sends it again. + let body = standby_beat(&h, &dep, &[], &[]).await; + assert_eq!(body["standby"][0]["lease_id"], json!(id)); + // A held id the queen has no standby for is dropped. + let body = standby_beat(&h, &dep, &[&id, "gone1"], &[]).await; + assert_eq!(body["drop_standby"], json!(["gone1"])); +} + +#[tokio::test] +async fn queued_deploy_demand_starts_the_standby_instead_of_minting_and_a_replacement_follows() { + let h = Harness::start(standby_cfg()).await; + let dep = h.enroll("dep").await; + standby_beat(&h, &dep, &[], &[]).await; + h.tick().await; + let id = standby_beat(&h, &dep, &[], &[]).await["standby"][0]["lease_id"] + .as_str() + .unwrap() + .to_owned(); + + h.queue_job(1, "borg-ci-medium-deploy", 5); + h.tick().await; + assert_eq!(h.state().mints.len(), 1, "the standby absorbs the demand"); + let body = standby_beat(&h, &dep, &[&id], &[]).await; + assert_eq!(body["start_standby"], json!([id])); + assert_eq!(body["leases"], json!([])); + + // The drone started it: it reports the id as active and no longer held. + standby_beat(&h, &dep, &[], &[&id]).await; + let lease = h.queen.store().lease(&id).await.unwrap().unwrap(); + assert_eq!(lease.state, borg::store::LeaseStatus::Pending); + assert!(lease.jit_config.is_none()); + + // The runner takes the job; the queen marks the lease running and mints + // a new standby. + h.set_runner("borg-dep-", "online", true); + h.state().jobs.clear(); + h.tick().await; + let lease = h.queen.store().lease(&id).await.unwrap().unwrap(); + assert_eq!(lease.state, borg::store::LeaseStatus::Running); + assert_eq!(h.state().mints.len(), 2, "replacement standby minted"); + let body = standby_beat(&h, &dep, &[], &[&id]).await; + assert_eq!(body["start_standby"], json!([])); + assert_eq!(body["standby"].as_array().unwrap().len(), 1); + + // The job finishes and the ephemeral runner deregisters. + h.state().runners.retain(|r| r["busy"] != json!(true)); + h.tick().await; + let lease = h.queen.store().lease(&id).await.unwrap().unwrap(); + assert_eq!(lease.state, borg::store::LeaseStatus::Done); + assert_eq!(h.state().mints.len(), 2); +} + +#[tokio::test] +async fn a_standby_started_during_an_outage_is_reconciled_when_the_queen_returns() { + let h = Harness::start(standby_cfg()).await; + let dep = h.enroll("dep").await; + standby_beat(&h, &dep, &[], &[]).await; + h.tick().await; + let id = standby_beat(&h, &dep, &[], &[]).await["standby"][0]["lease_id"] + .as_str() + .unwrap() + .to_owned(); + // The drone fail-opened on its own and the job ran while the queen could + // not hear it; the runner is gone from GitHub before the next heartbeat. + h.set_runner("borg-dep-", "online", true); + h.tick().await; + assert_eq!( + h.queen.store().lease(&id).await.unwrap().unwrap().state, + borg::store::LeaseStatus::Running + ); + h.state().runners.retain(|r| r["busy"] != json!(true)); + h.tick().await; + assert_eq!( + h.queen.store().lease(&id).await.unwrap().unwrap().state, + borg::store::LeaseStatus::Done + ); + // The drone, back online, still lists it as held; the queen does not + // revive it and tells the drone to forget it. + let body = standby_beat(&h, &dep, &[&id], &[]).await; + assert_eq!(body["drop_standby"], json!([id])); +} diff --git a/tests/requeue.rs b/tests/requeue.rs new file mode 100644 index 0000000..513b09e --- /dev/null +++ b/tests/requeue.rs @@ -0,0 +1,161 @@ +mod support; + +use serde_json::json; +use support::Harness; +use time::OffsetDateTime; + +fn config(dry_run: bool, max: u32) -> serde_json::Value { + json!({ + "preemption": {"after_secs": 600, "max_per_pass": 1, "min_victim_age_secs": 60}, + "supersede_janitor": {"dry_run": dry_run, "max_cancels_per_pass": max} + }) +} + +fn now() -> i64 { + OffsetDateTime::now_utc().unix_timestamp() +} + +async fn harness(dry_run: bool, max: u32) -> Harness { + let h = Harness::start(config(dry_run, max)).await; + h.enroll("d1").await; + h +} + +async fn requeued_at(h: &Harness, lease: &str) -> Option { + h.queen + .store() + .preemptions(50) + .await + .unwrap() + .into_iter() + .find(|r| r.lease_id == lease) + .unwrap() + .requeued_at +} + +#[tokio::test] +async fn a_completed_runs_job_is_rerun_and_marked() { + let h = harness(false, 20).await; + h.add_run(7, "completed", "pull_request", "sha", None); + h.preempted("l1", 7, 70, now()).await; + + let r = h.requeue_pass().await; + assert_eq!((r.checked, r.rerun, r.errors), (1, 1, 0), "{r:?}"); + assert_eq!(h.state().reruns, [70]); + assert!(requeued_at(&h, "l1").await.is_some()); + + let r = h.requeue_pass().await; + assert_eq!((r.checked, r.rerun), (0, 0), "{r:?}"); + assert_eq!(h.state().reruns.len(), 1); + + let status = h.status().await; + assert_eq!(status["requeue"]["total"]["rerun"], json!(1)); + assert!(h.metrics().await.contains("borg_requeue_rerun_total 1")); +} + +#[tokio::test] +async fn an_in_progress_run_is_left_waiting() { + let h = harness(false, 20).await; + h.add_run(7, "in_progress", "pull_request", "sha", None); + h.preempted("l1", 7, 70, now()).await; + + let r = h.requeue_pass().await; + assert_eq!((r.checked, r.waiting, r.rerun), (1, 1, 0), "{r:?}"); + assert!(h.state().reruns.is_empty()); + assert!(requeued_at(&h, "l1").await.is_none()); + + h.state().fresh_status.insert(7, "completed".into()); + let r = h.requeue_pass().await; + assert_eq!(r.rerun, 1, "{r:?}"); + assert_eq!(h.state().reruns, [70]); +} + +#[tokio::test] +async fn a_404_on_rerun_is_terminal() { + let h = harness(false, 20).await; + h.add_run(7, "completed", "pull_request", "sha", None); + h.preempted("l1", 7, 70, now()).await; + h.state().rerun_status.insert(70, 404); + + let r = h.requeue_pass().await; + assert_eq!((r.terminal, r.rerun, r.errors), (1, 0, 0), "{r:?}"); + assert!(requeued_at(&h, "l1").await.is_some()); + let r = h.requeue_pass().await; + assert_eq!(r.checked, 0, "{r:?}"); +} + +#[tokio::test] +async fn other_rerun_errors_are_retried_next_pass() { + let h = harness(false, 20).await; + h.add_run(7, "completed", "pull_request", "sha", None); + h.preempted("l1", 7, 70, now()).await; + h.state().rerun_status.insert(70, 500); + + let r = h.requeue_pass().await; + assert_eq!((r.errors, r.rerun, r.terminal), (1, 0, 0), "{r:?}"); + assert!(requeued_at(&h, "l1").await.is_none()); + + h.state().rerun_status.clear(); + let r = h.requeue_pass().await; + assert_eq!(r.rerun, 1, "{r:?}"); +} + +#[tokio::test] +async fn dry_run_makes_no_calls_and_marks_nothing() { + let h = harness(true, 20).await; + h.add_run(7, "completed", "pull_request", "sha", None); + h.preempted("l1", 7, 70, now()).await; + + let r = h.requeue_pass().await; + assert_eq!((r.checked, r.rerun), (1, 0), "{r:?}"); + assert_eq!(h.state().run_reads, 0); + assert!(h.state().reruns.is_empty()); + assert!(requeued_at(&h, "l1").await.is_none()); +} + +#[tokio::test] +async fn the_per_pass_cap_limits_reruns() { + let h = harness(false, 2).await; + for i in 1..=5u64 { + h.add_run(i, "completed", "pull_request", "sha", None); + h.preempted( + &format!("l{i}"), + i, + i * 10, + now() - 100 + i64::try_from(i).unwrap(), + ) + .await; + } + let r = h.requeue_pass().await; + assert_eq!(r.rerun, 2, "{r:?}"); + assert_eq!(h.state().reruns, [10, 20], "oldest first"); + let r = h.requeue_pass().await; + assert_eq!(r.rerun, 2, "{r:?}"); + let r = h.requeue_pass().await; + assert_eq!(r.rerun, 1, "{r:?}"); + assert_eq!(h.state().reruns.len(), 5); +} + +#[tokio::test] +async fn a_row_older_than_24_hours_is_marked_stale_without_a_call() { + let h = harness(false, 20).await; + h.add_run(7, "completed", "pull_request", "sha", None); + h.preempted("l1", 7, 70, now() - 25 * 3600).await; + + let r = h.requeue_pass().await; + assert_eq!((r.checked, r.stale, r.rerun), (1, 1, 0), "{r:?}"); + assert_eq!(h.state().run_reads, 0); + assert!(h.state().reruns.is_empty()); + assert!(requeued_at(&h, "l1").await.is_some()); +} + +#[tokio::test] +async fn the_requeuer_needs_both_preemption_and_the_janitor() { + let h = Harness::start(json!({"supersede_janitor": {}})).await; + assert!(h.queen.requeuer().is_none()); + let h = Harness::start(json!({ + "preemption": {"after_secs": 600, "max_per_pass": 1, "min_victim_age_secs": 60} + })) + .await; + assert!(h.queen.requeuer().is_none()); +} diff --git a/tests/run.rs b/tests/run.rs index abf5325..07160c7 100644 --- a/tests/run.rs +++ b/tests/run.rs @@ -190,6 +190,7 @@ fn opts(r: &Repo, queen: Option, marker: &str) -> RunOptions { toolchain: "tc".into(), timeout_secs: None, force: false, + scope: "pr".into(), wait_secs: 1, publish: false, argv: vec!["sh".into(), "-c".into(), format!("touch {marker}; exit 7")], diff --git a/tests/support/mod.rs b/tests/support/mod.rs index 2fe664d..cfe274b 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -54,6 +54,12 @@ pub struct MockState { pub cancelled: Vec, /// Runs whose cancel answers 409. pub uncancellable: Vec, + /// Jobs the requeuer re-ran. + pub reruns: Vec, + /// Status code a job's rerun answers instead of 201. + pub rerun_status: HashMap, + /// Reads of a single run (janitor or requeuer). + pub run_reads: u32, next_runner_id: u64, } @@ -152,9 +158,18 @@ async fn runs( h: HeaderMap, Query(q): Query>, ) -> Result, StatusCode> { - if q.get("event").map(String::as_str) == Some("pull_request") { - check_janitor(&h)?; - return Ok(Json(json!({"workflow_runs": lock(&s).pr_runs.clone()}))); + if bearer(&h) == Some(JANITOR_TOKEN) { + let runs: Vec = lock(&s) + .pr_runs + .iter() + .filter(|r| { + ["status", "event"] + .iter() + .all(|k| q.get(*k).is_none_or(|v| r[*k].as_str() == Some(v))) + }) + .cloned() + .collect(); + return Ok(Json(json!({"workflow_runs": runs}))); } check_token(&h)?; let ids: Vec = if q.get("status").map(String::as_str) == Some("queued") { @@ -215,7 +230,8 @@ async fn one_run( Path((_o, _r, id)): Path<(String, String, u64)>, ) -> Result, StatusCode> { check_janitor(&h)?; - let st = lock(&s); + let mut st = lock(&s); + st.run_reads += 1; let mut run = st .pr_runs .iter() @@ -245,6 +261,22 @@ async fn cancel_run( StatusCode::ACCEPTED } +async fn rerun_job( + State(s): State, + h: HeaderMap, + Path((_o, _r, id)): Path<(String, String, u64)>, +) -> StatusCode { + if check_janitor(&h).is_err() { + return StatusCode::UNAUTHORIZED; + } + let mut st = lock(&s); + if let Some(code) = st.rerun_status.get(&id) { + return StatusCode::from_u16(*code).unwrap(); + } + st.reruns.push(id); + StatusCode::CREATED +} + async fn pull( State(s): State, h: HeaderMap, @@ -387,6 +419,10 @@ impl Harness { "/repos/{owner}/{repo}/actions/runs/{run}/cancel", post(cancel_run), ) + .route( + "/repos/{owner}/{repo}/actions/jobs/{id}/rerun", + post(rerun_job), + ) .route("/repos/{owner}/{repo}/pulls/{n}", get(pull)) .route("/orgs/{org}/actions/runners", get(list_runners)) .route( @@ -414,11 +450,26 @@ impl Harness { borg::janitor::Janitor::new(&gh_url, borg::secret::Secret::new(JANITOR_TOKEN), settings) .unwrap() }); + let requeuer = cfg + .supersede_janitor + .as_ref() + .filter(|_| cfg.preemption.is_some()) + .map(|settings| { + borg::requeue::Requeuer::new( + &gh_url, + borg::secret::Secret::new(JANITOR_TOKEN), + settings, + ) + .unwrap() + }); let wakes = Arc::new(RecordingWake::default()); let mut queen = Queen::new(cfg, store, github).with_wake_sender(wakes.clone()); if let Some(j) = janitor { queen = queen.with_janitor(j); } + if let Some(r) = requeuer { + queen = queen.with_requeuer(r); + } if let Some(settings) = queen.config().tailnet_enroll.clone() { queen = queen.with_tailnet_enroll(borg::tailnet::TailnetEnroll { verifier: Arc::new(identities), @@ -538,11 +589,66 @@ impl Harness { })); } + /// Add a run of any status or event as the janitor lists it. + pub fn add_run(&self, id: u64, status: &str, event: &str, sha: &str, pr: Option) { + let prs: Vec = pr.iter().map(|n| json!({"number": n})).collect(); + self.state().pr_runs.push(json!({ + "id": id, "workflow_id": 1, "status": status, "event": event, + "head_sha": sha, "pull_requests": prs + })); + } + pub async fn janitor_pass(&self) -> borg::janitor::PassReport { let repos = self.queen.config().repos.clone(); self.queen.janitor().unwrap().pass(&repos).await } + /// Record a preemption of `job` in `run` at unix time `at`. + pub async fn preempted(&self, lease: &str, run: u64, job: u64, at: i64) { + use borg::classes::{Class, Scope}; + use borg::store::{LeaseRow, LeaseStatus, PreemptionRow}; + let store = self.queen.store(); + store + .insert_lease(&LeaseRow { + id: lease.into(), + drone: "d1".into(), + class: Class::Medium, + scope: Scope::Pr, + labels: vec![], + runner_name: format!("borg-d1-{lease}"), + runner_id: 1, + jit_config: None, + state: LeaseStatus::Running, + created_at: at, + expires_at: at + 600, + }) + .await + .unwrap(); + let row = PreemptionRow { + lease_id: lease.into(), + drone: "d1".into(), + class: Class::Medium, + scope: Scope::Pr, + for_scope: Scope::Main, + repo: format!("{ORG}/app"), + run_id: run, + job_id: job, + at, + requeued_at: None, + }; + assert!(store.record_preemption(&row, at).await.unwrap()); + } + + pub async fn requeue_pass(&self) -> borg::requeue::PassReport { + let now = OffsetDateTime::now_utc().unix_timestamp(); + self.queen + .requeuer() + .unwrap() + .pass(self.queen.store(), now) + .await + .unwrap() + } + pub async fn tick(&self) { self.queen.tick(OffsetDateTime::now_utc()).await.unwrap(); } diff --git a/tests/units.rs b/tests/units.rs index 0dc87e2..e2b4c91 100644 --- a/tests/units.rs +++ b/tests/units.rs @@ -597,3 +597,223 @@ async fn a_bundle_that_never_arrives_fails_the_unit_and_bundles_go_with_their_un assert!(r.h.queen.store().unit(&c).await.unwrap().is_none()); assert!(r.h.queen.store().unit_bundle(&c).await.unwrap().is_none()); } + +fn tree(c: char) -> String { + c.to_string().repeat(40) +} + +async fn state_of(r: &Rig, id: &str) -> String { + r.h.queen + .store() + .unit(id) + .await + .unwrap() + .unwrap() + .state + .as_str() + .to_owned() +} + +async fn deploy_drone(r: &Rig, name: &str) -> String { + let t = r.h.enroll(name).await; + r.h.heartbeat_with(&t, json!({"trust": "deploy"})).await; + t +} + +async fn scoped_client(r: &Rig, name: &str, max: borg::classes::Scope) -> String { + let tok = enroll::new_token(); + r.h.queen + .store() + .enroll_client_scoped(name, &enroll::hash_token(&tok), max, 0) + .await + .unwrap(); + tok +} + +#[tokio::test] +async fn scope_round_trips_and_defaults_to_pr() { + let r = Rig::new(json!({})).await; + r.drone("d1").await; + let (_, a) = r.submit(&r.unit(json!({}))).await; + let u = r + .get(&format!("/v1/units/{}", a["unit_id"].as_str().unwrap())) + .await; + assert_eq!(u.json::().await.unwrap()["scope"], "pr"); + + let tok = scoped_client(&r, "gate", borg::classes::Scope::Main).await; + let body = r.unit(json!({"scope": "main", "tree_sha": tree('1')})); + let (code, sub) = r.submit_as(&tok, &body).await; + assert_eq!(code, 201, "{sub}"); + let u = r + .get(&format!("/v1/units/{}", sub["unit_id"].as_str().unwrap())) + .await; + assert_eq!(u.json::().await.unwrap()["scope"], "main"); + assert_eq!( + r.h.queen.store().client_max_scope("gate").await.unwrap(), + borg::classes::Scope::Main + ); +} + +#[tokio::test] +async fn a_unit_above_the_clients_maximum_scope_is_refused() { + let r = Rig::new(json!({})).await; + let hook = scoped_client(&r, "hook", borg::classes::Scope::Pr).await; + let resp = + r.h.http + .post(format!("{}/v1/units", r.h.url)) + .bearer_auth(&hook) + .json(&r.unit(json!({"scope": "deploy"}))) + .send() + .await + .unwrap(); + assert_eq!(resp.status().as_u16(), 403); + let msg = resp.text().await.unwrap(); + assert!( + msg.contains("up to scope pr") && msg.contains("asked for deploy"), + "{msg}" + ); + let (code, _) = r.submit_as(&hook, &r.unit(json!({"scope": "main"}))).await; + assert_eq!(code, 403); + let (code, _) = r.submit_as(&hook, &r.unit(json!({"scope": "pr"}))).await; + assert_eq!(code, 201); +} + +#[tokio::test] +async fn one_free_slot_goes_to_deploy_then_main_then_pr() { + let r = Rig::new(json!({})).await; + let tok = scoped_client(&r, "gate", borg::classes::Scope::Deploy).await; + let d = deploy_drone(&r, "d1").await; + // Oldest first: pr, then main, then deploy. A heavy unit fills the drone. + let mut ids = Vec::new(); + for (scope, c) in [("pr", '1'), ("main", '2'), ("deploy", '3')] { + let body = r.unit(json!({"scope": scope, "class": "heavy", "tree_sha": tree(c)})); + let (code, sub) = r.submit_as(&tok, &body).await; + assert_eq!(code, 201, "{sub}"); + ids.push(sub["unit_id"].as_str().unwrap().to_owned()); + } + for (n, want) in [2usize, 1, 0].into_iter().enumerate() { + r.tick(10 + i64::try_from(n).unwrap()).await; + let grants = r.beat(&d, &[]).await["units"].clone(); + assert_eq!(grants.as_array().unwrap().len(), 1, "round {n}: {grants}"); + assert_eq!(grants[0]["unit_id"], ids[want].as_str(), "round {n}"); + r.beat(&d, &[&ids[want]]).await; + assert_eq!(r.post_result(&d, &ids[want], 0).await, 204); + } +} + +#[tokio::test] +async fn a_deploy_unit_is_not_placed_on_a_shared_drone() { + let r = Rig::new(json!({})).await; + let tok = scoped_client(&r, "gate", borg::classes::Scope::Deploy).await; + let shared = r.drone("shared").await; + let (_, dep) = r.submit_as(&tok, &r.unit(json!({"scope": "deploy"}))).await; + let dep = dep["unit_id"].as_str().unwrap().to_owned(); + let (_, main) = r + .submit_as( + &tok, + &r.unit(json!({"scope": "main", "tree_sha": tree('4')})), + ) + .await; + let main = main["unit_id"].as_str().unwrap().to_owned(); + r.tick(5).await; + assert_eq!(state_of(&r, &dep).await, "queued"); + // A shared drone does run main work, so only the deploy unit waits. + let grants = r.beat(&shared, &[]).await["units"].clone(); + assert_eq!(grants.as_array().unwrap().len(), 1); + assert_eq!(grants[0]["unit_id"], main.as_str()); + + // A deploy-trust drone takes it. + let d = deploy_drone(&r, "dep").await; + r.tick(6).await; + assert_eq!(state_of(&r, &dep).await, "leased"); + assert_eq!(r.beat(&d, &[]).await["units"][0]["unit_id"], dep.as_str()); + + // A pr unit never runs on a deploy drone: with only one, it waits. + let r2 = Rig::new(json!({})).await; + let d2 = deploy_drone(&r2, "dep").await; + let (_, pr) = r2.submit(&r2.unit(json!({}))).await; + r2.tick(5).await; + assert_eq!( + state_of(&r2, pr["unit_id"].as_str().unwrap()).await, + "queued" + ); + assert!( + r2.beat(&d2, &[]).await["units"] + .as_array() + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn the_newest_unit_with_a_supersede_key_wins() { + let r = Rig::new(json!({})).await; + let d = r.drone("d1").await; + let sub = |c: char, extra: Value| { + let mut b = r.unit(json!({"tree_sha": tree(c)})); + for (k, v) in extra.as_object().cloned().unwrap_or_default() { + b[k] = v; + } + b + }; + let id = |v: &Value| v["unit_id"].as_str().unwrap().to_owned(); + + // A record unit, then running (placed and acknowledged), queued, and + // other keys. Submitting a record unit with a key supersedes older + // non-record units too, so it goes first. + let (_, record) = r + .submit(&sub('3', json!({"supersede_key": "pr-7", "record": true}))) + .await; + let record = id(&record); + let (_, running) = r.submit(&sub('1', json!({"supersede_key": "pr-7"}))).await; + let running = id(&running); + r.tick(1).await; + r.beat(&d, &[]).await; + r.beat(&d, &[&running]).await; + assert_eq!(state_of(&r, &running).await, "running"); + let (_, other) = r.submit(&sub('4', json!({"supersede_key": "pr-8"}))).await; + let other = id(&other); + let (_, plain) = r.submit(&sub('5', json!({}))).await; + let plain = id(&plain); + + let (code, newest) = r.submit(&sub('6', json!({"supersede_key": "pr-7"}))).await; + assert_eq!(code, 201, "{newest}"); + assert_eq!(newest["superseded"], json!([running])); + let rec = r.h.queen.store().unit(&record).await.unwrap().unwrap(); + assert!(!rec.cancel_requested && rec.state.as_str() != "failed"); + assert_eq!(state_of(&r, &other).await, "queued"); + assert_eq!(state_of(&r, &plain).await, "queued"); + + // The running one is told to stop on the drone's next heartbeat. + let hb = r.beat(&d, &[&running]).await; + assert_eq!(hb["cancel_units"], json!([running])); + + // A queued unit is cancelled at once by the next submission with the key. + let (_, third) = r.submit(&sub('2', json!({"supersede_key": "pr-7"}))).await; + assert_eq!(third["superseded"], json!([id(&newest)])); + let q = r.h.queen.store().unit(&id(&newest)).await.unwrap().unwrap(); + assert_eq!(q.state.as_str(), "failed"); + assert_eq!(q.failure.unwrap(), format!("superseded by {}", id(&third))); + + assert_eq!(r.h.status().await["units"]["superseded"], 2); + assert!( + r.h.metrics() + .await + .contains("borg_units_superseded_total 2") + ); + + // Another client's units with the same key are left alone. + let tok = scoped_client(&r, "other", borg::classes::Scope::Pr).await; + let (_, theirs) = r + .submit_as(&tok, &sub('7', json!({"supersede_key": "pr-9"}))) + .await; + let (_, mine) = r.submit(&sub('8', json!({"supersede_key": "pr-9"}))).await; + assert!(mine["superseded"].as_array().unwrap().is_empty()); + assert_eq!(state_of(&r, &id(&theirs)).await, "queued"); + + // The key is bounded. + let (code, _) = r + .submit(&sub('9', json!({"supersede_key": "k".repeat(129)}))) + .await; + assert_eq!(code, 400); +} From 6fb07e4c9aaec2e652028bb9c56ed10e21b7ef19 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Mon, 12 Oct 2026 03:32:55 +0000 Subject: [PATCH 12/27] chore: project borg from Mono 2e833192cc63 --- .repository-projection.json | 6 +- Cargo.toml | 20 +++ README.md | 7 ++ docs/operations.md | 45 ++++++- src/api.rs | 26 ++++ src/config.rs | 107 +++++++++++++++++ src/drone/caches/git.rs | 234 ++++++++++++++++++++++++++++++++---- src/drone/caches/mod.rs | 10 ++ src/drone/config.rs | 8 +- src/drone/mod.rs | 3 + src/github.rs | 88 +++++++++++--- src/janitor.rs | 32 +++-- src/lib.rs | 1 + src/main.rs | 65 +++++++--- src/mirror_auth.rs | 102 ++++++++++++++++ src/queen.rs | 14 +++ src/requeue.rs | 30 +++-- tests/capacity.rs | 2 +- tests/contract.rs | 2 +- tests/drone.rs | 2 +- tests/health.rs | 2 +- tests/host_suite.rs | 14 +++ tests/janitor.rs | 2 +- tests/job_user.rs | 2 +- tests/queen.rs | 84 +++++++++++++ tests/requeue.rs | 2 +- tests/rollout.rs | 2 +- tests/sccache_slots.rs | 94 +++++++++++++-- tests/suite.rs | 12 ++ tests/support/mod.rs | 47 ++++++-- tests/tailnet.rs | 2 +- tests/units.rs | 2 +- tests/wake.rs | 2 +- 33 files changed, 956 insertions(+), 115 deletions(-) create mode 100644 src/mirror_auth.rs create mode 100644 tests/host_suite.rs create mode 100644 tests/suite.rs diff --git a/.repository-projection.json b/.repository-projection.json index 8055428..cd48cd7 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "459f737d5b9e7e476f6c79467dab5dba31868bd4", + "sourceSha": "2e833192cc635eb2dcb5f4cb02562547636a982a", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "3d4d719a3dabdd4da94fc7343562e008ce8891e0", + "priorProjectedBase": "c51a232c724d6915b1fa0b001a635c49519b01ae", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "fb334a68e36619eccf3003521e93a792ad2826c573ebfd1e95d3b96d0e951571", + "contentDigest": "58e3b42b548a0e1b22cbee864fb0e787e0a8970a40e8c6d23a43758175bc8928", "publicationEligible": true } diff --git a/Cargo.toml b/Cargo.toml index 79672c5..d2f6702 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ edition = "2024" license = "BUSL-1.1" rust-version = "1.95" publish = false +autotests = false repository = "https://github.com/dx-corp/borg" description = "Self-organizing CI fleet agent: drones probe their host and run ephemeral runners that a queen leases to them." @@ -17,6 +18,25 @@ name = "borg" path = "src/main.rs" test = false +# Shared fixtures compile once per suite. The cache test spawns itself by +# its exact test name and keeps its own target. +[[test]] +name = "suite" +path = "tests/suite.rs" + +[[test]] +name = "host_suite" +path = "tests/host_suite.rs" + +# Keep the largest test module independently compilable for focused edits. +[[test]] +name = "queen" +path = "tests/queen.rs" + +[[test]] +name = "sccache_slots" +path = "tests/sccache_slots.rs" + [dependencies] anyhow = "1.0" axum = "0.8" diff --git a/README.md b/README.md index 2a2946d..2a3c0a5 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,13 @@ cargo test --locked cargo clippy --locked --all-targets -- -D warnings ``` +The integration tests share fixtures across two suites. Run service-state tests +with `cargo test --locked --test suite units::`, or host tests with +`cargo test --locked --test host_suite drone::`. The largest module, `queen`, +keeps its own target for focused edits: `cargo test --locked --test queen`. +The compiler-cache test has its own target because it invokes itself by exact test name: +`cargo test --locked --test sccache_slots`. + | Source | Responsibility | | --- | --- | | `src/probe/` | Host capacity and tool detection. | diff --git a/docs/operations.md b/docs/operations.md index 9647432..b208562 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -534,7 +534,9 @@ front), and its service user needs access to the tailscaled socket The optional supersede janitor cancels queued pull-request workflow runs whose head commit is no longer the pull request's head. Add a `supersede_janitor` -block (`token_path` or `BORG_JANITOR_TOKEN` for a token with `actions: write`, +block (optionally `token_path` or `BORG_JANITOR_TOKEN` for a static token with +`actions: write`; without either the janitor uses the queen's own App installation +token, minted and refreshed on demand, so the App needs Actions: read and write; `dry_run` defaulting to true, `interval_secs`, `max_cancels_per_pass`) to turn it on. By default it never cancels a run that has started. @@ -925,10 +927,43 @@ Notes on each kind: ``` Mirrors fetch from `remote_base` (default `https://github.com`) plus - `/.git`. The drone holds no GitHub credentials: for a private - repository give the service user a credential helper, or set `remote_base` to - an ssh base such as `git@host:` with a read-only deploy key. A failed refresh - keeps the previous mirror and never blocks a job. + `/.git`. The drone holds no long-lived GitHub credentials and no + App key. For a private repository either let the queen hand it a short-lived + read-only token (see Mirror credentials below), give the service user a + credential helper, or set `remote_base` to an ssh base such as `git@host:` + with a read-only deploy key. A failed refresh keeps the previous mirror and + never blocks a job. + + **Mirror credentials.** When deploy keys are disabled, a second GitHub App + (installed on the organization with repository permission Contents: read) + can refresh private mirrors. Its key stays on the queen. Add to `queen.toml`: + + ```toml + [mirror_auth] + app_id = "12345" + private_key_path = "/etc/borg/mirror-app.pem" + repositories = ["app"] + ``` + + `repositories` are bare names in `github.org` (no slashes). Set + `BORG_MIRROR_APP_PRIVATE_KEY` to the PEM text instead of `private_key_path` + to pass the key in the environment. Without a `mirror_auth` block nothing + changes. The queen mints an installation token limited to `contents: read` on + those repositories, which GitHub expires after one hour, keeps it, and mints a + new one when under ten minutes of life remain. Every heartbeat response then + carries `mirror_credential` (`token`, `expires_at`, and the `owner/name` + repositories it covers), so resending is the refresh. The drone keeps it in + memory only, replaces it on each heartbeat and drops it at `expires_at`. For a + repository in the credential, and only when `caches.git.remote_base` is an + `https://` base, the drone fetches with `git -c http.extraheader=...`: the + token is never written to the mirror's config, the stored remote URL stays + credential-free, and an ssh base ignores the credential. While a fetch runs + the header is visible in that drone's process list. The queen speaks plain + HTTP, so run the drone-to-queen channel over TLS or a trusted LAN. `GET + /status` shows `mirror_auth` as `{"configured": true, "expires_at": }` and never the token; the queen never logs it. The first + credential after a drone starts refreshes its mirrors at once. Drones need a + borg build that understands `mirror_credential`; older drones ignore it. #### Slots, path remapping and compile-cache keys diff --git a/src/api.rs b/src/api.rs index 673e9c4..a8ce5a6 100644 --- a/src/api.rs +++ b/src/api.rs @@ -211,6 +211,11 @@ pub struct HeartbeatResponse { /// Held standby leases that ended on the queen's side; forget them. #[serde(default)] pub drop_standby: Vec, + /// A short-lived read-only token for the git mirrors, present on every + /// heartbeat while the queen has `mirror_auth` configured. Keep it in + /// memory only and never log it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mirror_credential: Option, } #[derive(Debug, Serialize, Deserialize)] @@ -477,6 +482,16 @@ async fn heartbeat( }) }) .collect(); + let mirror_credential = match queen.mirror_auth() { + Some(mirror) => match mirror.credential().await { + Ok(credential) => Some(credential), + Err(e) => { + tracing::warn!("mint mirror token: {e}"); + None + } + }, + None => None, + }; Ok(Json(HeartbeatResponse { leases, units, @@ -492,6 +507,7 @@ async fn heartbeat( standby, start_standby, drop_standby, + mirror_credential, })) } @@ -706,6 +722,9 @@ pub struct Status { pub units: UnitsStatus, /// Present when the supersede janitor is enabled. pub janitor: Option, + /// Whether drones are handed git mirror tokens, and when the current one + /// expires. Never the token. + pub mirror_auth: crate::mirror_auth::MirrorAuthView, /// Present when preempted jobs are re-run (preemption and the janitor /// are both configured). pub requeue: Option, @@ -1035,6 +1054,13 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result } }, janitor: queen.janitor().map(crate::janitor::Janitor::view), + mirror_auth: queen.mirror_auth().map_or( + crate::mirror_auth::MirrorAuthView { + configured: false, + expires_at: None, + }, + crate::mirror_auth::MirrorAuth::view, + ), requeue: queen.requeuer().map(crate::requeue::Requeuer::view), reclaim: ReclaimSummary { reclaimed_bytes_logged: queen.store().reclaim_bytes_logged().await?, diff --git a/src/config.rs b/src/config.rs index 3f32157..1256379 100644 --- a/src/config.rs +++ b/src/config.rs @@ -236,6 +236,62 @@ pub struct QueenConfig { /// Absent (the default) disables preemption; see [`crate::preempt`]. #[serde(default)] pub preemption: Option, + /// Absent (the default) disables git mirror credentials; see + /// [`crate::mirror_auth`]. + #[serde(default)] + pub mirror_auth: Option, +} + +/// A second GitHub App, held by the queen only, whose short-lived read-only +/// installation tokens let drones refresh private git mirrors. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MirrorAuthSettings { + /// The mirror App's id (or client id). + pub app_id: String, + /// PEM private key file. If absent, `BORG_MIRROR_APP_PRIVATE_KEY` must + /// hold the PEM text. + #[serde(default)] + pub private_key_path: Option, + /// Bare repository names in `github.org` the token may read. + pub repositories: Vec, +} + +impl MirrorAuthSettings { + fn validate(&self) -> Result<(), ConfigError> { + let bad = |msg: &str| Err(ConfigError::Invalid(msg.to_owned())); + if self.app_id.is_empty() { + return bad("mirror_auth.app_id is required"); + } + if self.repositories.is_empty() { + return bad("mirror_auth.repositories must name at least one repository"); + } + if self.repositories.iter().any(|r| { + r.is_empty() + || r.contains('/') + || r.chars().any(|c| c.is_whitespace() || c.is_control()) + }) { + return bad("every mirror_auth.repositories entry must be a bare repository name"); + } + Ok(()) + } + + /// The App private key: `private_key_path`, else `env_pem` (the value of + /// `BORG_MIRROR_APP_PRIVATE_KEY`). + pub fn private_key(&self, env_pem: Option) -> Result, ConfigError> { + match (&self.private_key_path, env_pem) { + (Some(path), _) => std::fs::read(path).map_err(|e| { + ConfigError::Invalid(format!( + "read mirror_auth.private_key_path {}: {e}", + path.display() + )) + }), + (None, Some(pem)) if !pem.trim().is_empty() => Ok(pem.into_bytes()), + (None, _) => Err(ConfigError::Invalid( + "set mirror_auth.private_key_path or BORG_MIRROR_APP_PRIVATE_KEY".to_owned(), + )), + } + } } fn default_listen() -> SocketAddr { @@ -313,6 +369,9 @@ impl QueenConfig { fn validate(&self) -> Result<(), ConfigError> { self.classes.validate()?; let bad = |msg: &str| Err(ConfigError::Invalid(msg.to_owned())); + if let Some(m) = &self.mirror_auth { + m.validate()?; + } if self.github.org.is_empty() || self.github.app_id.is_empty() { return bad("github.org and github.app_id are required"); } @@ -389,6 +448,54 @@ mod tests { assert_eq!(c.listen.ip().to_string(), "127.0.0.1"); } + fn with_mirror(block: &str) -> String { + MIN.replace("\"repos\"", &format!("\"mirror_auth\":{block},\"repos\"")) + } + + #[test] + fn mirror_auth_is_off_unless_present_and_validated() { + assert!(QueenConfig::from_json(MIN).unwrap().mirror_auth.is_none()); + let ok = with_mirror( + r#"{"app_id":"12345","private_key_path":"/etc/borg/mirror-app.pem","repositories":["app"]}"#, + ); + let c = QueenConfig::from_json(&ok).unwrap(); + assert_eq!(c.mirror_auth.unwrap().repositories, ["app"]); + for bad in [ + r#"{"app_id":"12345","repositories":[]}"#, + r#"{"app_id":"12345","repositories":["example-org/app"]}"#, + r#"{"app_id":"12345","repositories":[""]}"#, + r#"{"app_id":"12345","repositories":["a b"]}"#, + r#"{"app_id":"","repositories":["app"]}"#, + r#"{"app_id":"12345","repositories":["app"],"nope":1}"#, + r#"{"repositories":["app"]}"#, + ] { + assert!(QueenConfig::from_json(&with_mirror(bad)).is_err(), "{bad}"); + } + } + + #[test] + fn mirror_key_comes_from_the_file_or_the_environment_value() { + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("mirror-app.pem"); + std::fs::write(&file, "PEM-FROM-FILE").unwrap(); + let settings = |path: Option| MirrorAuthSettings { + app_id: "12345".to_owned(), + private_key_path: path, + repositories: vec!["app".to_owned()], + }; + let from_file = settings(Some(file)).private_key(Some("ignored".to_owned())); + assert_eq!(from_file.unwrap(), b"PEM-FROM-FILE"); + let from_env = settings(None).private_key(Some("PEM-FROM-ENV".to_owned())); + assert_eq!(from_env.unwrap(), b"PEM-FROM-ENV"); + assert!(settings(None).private_key(None).is_err()); + assert!(settings(None).private_key(Some(" \n".to_owned())).is_err()); + assert!( + settings(Some(dir.path().join("missing.pem"))) + .private_key(None) + .is_err() + ); + } + #[test] fn rejects_bad_repos_unknown_fields_and_empty_org() { for bad in [ diff --git a/src/drone/caches/git.rs b/src/drone/caches/git.rs index e668329..c602fa7 100644 --- a/src/drone/caches/git.rs +++ b/src/drone/caches/git.rs @@ -13,6 +13,7 @@ use tokio::sync::Mutex; use crate::drone::config::GitCaches; use crate::drone::exec::Exec; use crate::drone::fsops::FsOps; +use crate::mirror_auth::MirrorCredential; /// A mirror refreshed this recently is not refreshed again for a new lease. const MIN_REFRESH_GAP: Duration = Duration::from_secs(15); @@ -52,14 +53,63 @@ pub fn remote_url(base: &str, repo: &str) -> String { } } +const EXTRAHEADER: &str = "http.extraheader="; + +/// `args` for a log line, with the credential header hidden. +fn redacted(args: &[&str]) -> String { + args.iter() + .map(|a| { + if a.starts_with(EXTRAHEADER) { + "http.extraheader=" + } else { + a + } + }) + .collect::>() + .join(" ") +} + +fn base64(input: &[u8]) -> String { + const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + let mut out = String::with_capacity(input.len().div_ceil(3) * 4); + for chunk in input.chunks(3) { + let n = chunk + .iter() + .enumerate() + .fold(0u32, |n, (i, b)| n | (u32::from(*b) << (16 - 8 * i))); + for i in 0..4 { + if i <= chunk.len() { + out.push(char::from(ALPHABET[((n >> (18 - 6 * i)) & 63) as usize])); + } else { + out.push('='); + } + } + } + out +} + +/// The `-c` value that authenticates a fetch from GitHub over https as an +/// App installation. Passed per command, so the token is never written to the +/// mirror's config or the stored remote URL. +fn auth_header_config(token: &str) -> String { + let basic = base64(format!("x-access-token:{token}").as_bytes()); + format!("{EXTRAHEADER}AUTHORIZATION: basic {basic}") +} + fn run_git(exec: &dyn Exec, git: &str, args: &[&str]) -> Result<(), String> { match exec.run(git, args) { Ok(out) if out.success => Ok(()), - Ok(_) => Err(format!("`git {}` failed", args.join(" "))), + Ok(_) => Err(format!("`git {}` failed", redacted(args))), Err(e) => Err(format!("{git}: {e}")), } } +fn unix_now() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(0)) +} + /// Create the mirror if needed, point it at `url`, and fetch every branch and /// tag into it. Idempotent, so an interrupted first clone heals on the next run. fn refresh_mirror( @@ -68,6 +118,7 @@ fn refresh_mirror( git: &str, path: &Path, url: &str, + auth: Option<&str>, ) -> Result<(), String> { if let Some(parent) = path.parent() { fs_ops @@ -93,24 +144,19 @@ fn refresh_mirror( // Objects only accumulate: a job borrowing them via alternates must never // see them pruned underneath it. run_git(exec, git, &["-C", &p, "config", "gc.auto", "0"])?; - run_git( - exec, - git, - &[ - "-C", - &p, - "-c", - "http.lowSpeedLimit=1000", - "-c", - "http.lowSpeedTime=60", - "fetch", - "--prune", - "--tags", - "--force", - "--quiet", - "origin", - ], - ) + let mut args = vec![ + "-C", + &p, + "-c", + "http.lowSpeedLimit=1000", + "-c", + "http.lowSpeedTime=60", + ]; + if let Some(header) = auth { + args.extend(["-c", header]); + } + args.extend(["fetch", "--prune", "--tags", "--force", "--quiet", "origin"]); + run_git(exec, git, &args) } struct Mirror { @@ -125,6 +171,10 @@ pub struct GitMirrors { exec: Arc, fs_ops: Arc, mirrors: Vec, + /// Fetch credentials apply only to an `https://` remote base. + https: bool, + /// The queen's latest token; memory only, replaced on each heartbeat. + credential: std::sync::Mutex>, /// Leases that were handed a mirror / ran without one. served: AtomicU64, missed: AtomicU64, @@ -152,6 +202,8 @@ impl GitMirrors { exec, fs_ops, mirrors, + https: cfg.remote_base.starts_with("https://"), + credential: std::sync::Mutex::new(None), served: AtomicU64::new(0), missed: AtomicU64::new(0), } @@ -161,20 +213,58 @@ impl GitMirrors { self.mirrors.is_empty() } + /// Keep the queen's latest credential. Returns true when mirrors were + /// configured and no usable credential was held before, so the caller can + /// refresh them now instead of at the next timer tick. + pub fn set_credential(&self, credential: MirrorCredential) -> bool { + let now = unix_now(); + let Ok(mut held) = self.credential.lock() else { + return false; + }; + let was_usable = held.as_ref().is_some_and(|c| c.expires_at > now); + *held = Some(credential); + !was_usable && !self.mirrors.is_empty() && self.https + } + + /// The `-c` auth setting for `repo`, when a held credential covers it and + /// has not expired; an expired one is dropped. + fn auth_for(&self, repo: &str) -> Option { + if !self.https { + return None; + } + let mut held = self.credential.lock().ok()?; + if held.as_ref().is_some_and(|c| c.expires_at <= unix_now()) { + *held = None; + } + let c = held.as_ref()?; + c.repositories + .iter() + .any(|r| r == repo) + .then(|| auth_header_config(c.token.expose())) + } + async fn refresh(&self, mirror: &Mirror, force: bool) { let mut last = mirror.refreshed.lock().await; if !force && last.is_some_and(|t| t.elapsed() < MIN_REFRESH_GAP) { return; } - let (exec, fs_ops, git, path, url) = ( + let (exec, fs_ops, git, path, url, auth) = ( Arc::clone(&self.exec), Arc::clone(&self.fs_ops), self.git.clone(), mirror.path.clone(), mirror.url.clone(), + self.auth_for(&mirror.repo), ); let job = tokio::task::spawn_blocking(move || { - refresh_mirror(exec.as_ref(), fs_ops.as_ref(), &git, &path, &url) + refresh_mirror( + exec.as_ref(), + fs_ops.as_ref(), + &git, + &path, + &url, + auth.as_deref(), + ) }); match tokio::time::timeout(REFRESH_TIMEOUT, job).await { Ok(Ok(Ok(()))) => *last = Some(Instant::now()), @@ -306,6 +396,108 @@ mod tests { ); } + fn credential(expires_at: i64, repos: &[&str]) -> MirrorCredential { + MirrorCredential { + token: crate::secret::Secret::new("ghs_mirror_token"), + expires_at, + repositories: repos.iter().map(|r| (*r).to_owned()).collect(), + } + } + + #[test] + fn base64_matches_the_standard_alphabet_and_padding() { + assert_eq!(base64(b""), ""); + assert_eq!(base64(b"f"), "Zg=="); + assert_eq!(base64(b"fo"), "Zm8="); + assert_eq!(base64(b"foo"), "Zm9v"); + assert_eq!(base64(b"foob"), "Zm9vYg=="); + assert_eq!(base64(b"x-access-token:abc"), "eC1hY2Nlc3MtdG9rZW46YWJj"); + } + + #[tokio::test] + async fn the_credential_rides_the_fetch_command_and_never_the_stored_url() { + let root = tempfile::tempdir().unwrap(); + let x = Arc::new(FakeExec::default()); + let m = GitMirrors::new( + &cfg(&["example-org/app", "example-org/other"]), + root.path(), + x.clone(), + Arc::new(LocalFs), + ); + let far = unix_now() + 3600; + assert!( + m.set_credential(credential(far, &["example-org/app"])), + "the first usable credential asks for a refresh" + ); + assert!( + !m.set_credential(credential(far, &["example-org/app"])), + "a replacement does not" + ); + m.refresh_all().await; + + let header = format!( + "http.extraheader=AUTHORIZATION: basic {}", + base64(b"x-access-token:ghs_mirror_token") + ); + let calls = x.calls(); + let fetch = |repo: &str| { + calls + .iter() + .find(|c| c.contains("fetch --prune") && c.contains(&format!("{repo}.git"))) + .cloned() + .unwrap() + }; + assert!(fetch("example-org/app").contains(&format!("-c {header} fetch"))); + assert!( + !fetch("example-org/other").contains("extraheader"), + "repositories the credential does not name get no header" + ); + for call in &calls { + if call.contains("config remote.origin.url") { + assert!( + call.ends_with("https://git.example.test/example-org/app.git") + || call.ends_with("https://git.example.test/example-org/other.git"), + "{call}" + ); + } + if !call.contains("fetch --prune") { + assert!(!call.contains("ghs_mirror_token"), "{call}"); + } + } + } + + #[tokio::test] + async fn an_expired_credential_is_dropped_and_ssh_bases_ignore_it() { + let root = tempfile::tempdir().unwrap(); + let x = Arc::new(FakeExec::default()); + let m = GitMirrors::new( + &cfg(&["example-org/app"]), + root.path(), + x.clone(), + Arc::new(LocalFs), + ); + m.set_credential(credential(unix_now() - 1, &["example-org/app"])); + m.refresh_all().await; + assert!(x.calls().iter().all(|c| !c.contains("extraheader"))); + assert!(m.credential.lock().unwrap().is_none()); + + let ssh = GitCaches { + remote_base: "git@git.example.test:".to_owned(), + ..cfg(&["example-org/app"]) + }; + let m = GitMirrors::new(&ssh, root.path(), x.clone(), Arc::new(LocalFs)); + assert!(!m.set_credential(credential(unix_now() + 3600, &["example-org/app"]))); + m.refresh_all().await; + assert!(x.calls().iter().all(|c| !c.contains("extraheader"))); + } + + #[test] + fn failure_messages_hide_the_header() { + let header = auth_header_config("ghs_mirror_token"); + let shown = redacted(&["-c", &header, "fetch"]); + assert_eq!(shown, "-c http.extraheader= fetch"); + } + #[tokio::test] async fn a_failed_refresh_is_retried_by_the_next_lease() { let root = tempfile::tempdir().unwrap(); diff --git a/src/drone/caches/mod.rs b/src/drone/caches/mod.rs index f025b52..e25573c 100644 --- a/src/drone/caches/mod.rs +++ b/src/drone/caches/mod.rs @@ -25,6 +25,7 @@ use crate::api::CacheStat; use crate::drone::config::DroneConfig; use crate::drone::exec::Exec; use crate::drone::fsops::FsOps; +use crate::mirror_auth::MirrorCredential; use crate::probe::{self, Host}; use git::GitMirrors; use kinds::Kind; @@ -363,6 +364,15 @@ impl Caches { .collect() } + /// Keep the queen's latest git mirror credential (memory only). The first + /// usable one refreshes the mirrors now, not at the next timer tick. + pub fn set_mirror_credential(self: &Arc, credential: MirrorCredential) { + if self.git.set_credential(credential) { + let caches = Arc::clone(self); + tokio::spawn(async move { caches.git.refresh_stale().await }); + } + } + /// Refresh git mirrors and run maintenance on a timer until aborted. pub fn spawn_background(self: &Arc) -> tokio::task::JoinHandle<()> { let caches = Arc::clone(self); diff --git a/src/drone/config.rs b/src/drone/config.rs index ee56b32..7b080a5 100644 --- a/src/drone/config.rs +++ b/src/drone/config.rs @@ -133,9 +133,11 @@ pub struct GitCaches { /// Repositories (`owner/name`). Jobs get each mirror's path in /// `BORG_GIT_MIRROR__`. pub mirrors: Vec, - /// Where mirrors fetch from; `//.git`. Credentials are - /// the host's business (a credential helper, or an ssh base such as - /// `git@host:` with a read-only deploy key for the service user). + /// Where mirrors fetch from; `//.git`. For an `https://` + /// base the queen's short-lived `mirror_credential` authenticates fetches + /// of the repositories it names; otherwise credentials are the host's + /// business (a credential helper, or an ssh base such as `git@host:` with + /// a read-only deploy key for the service user). pub remote_base: String, pub bin: String, pub interval_secs: u64, diff --git a/src/drone/mod.rs b/src/drone/mod.rs index d09e7cb..fd345f5 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -241,6 +241,9 @@ impl Drone { }; let resp = self.client.heartbeat(&req).await?; let defer = self.after_heartbeat(&req, resp.desired_version.as_deref()); + if let Some(credential) = resp.mirror_credential { + self.caches.set_mirror_credential(credential); + } if let Ok(preempts) = self.preempts.lock() { for id in &resp.cancel_leases { if let Some(tx) = preempts.get(id) { diff --git a/src/github.rs b/src/github.rs index 70056fd..72ca6f2 100644 --- a/src/github.rs +++ b/src/github.rs @@ -196,6 +196,23 @@ fn unix_now() -> i64 { .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(0)) } +/// Where a GitHub API caller gets its bearer token: a fixed secret, or the +/// queen's own App installation token, minted and refreshed on demand. +#[derive(Clone)] +pub enum TokenSource { + Static(Secret), + App(std::sync::Arc), +} + +impl TokenSource { + pub async fn bearer(&self) -> Result { + match self { + TokenSource::Static(s) => Ok(s.clone()), + TokenSource::App(c) => c.installation_token().await, + } + } +} + impl GithubClient { pub fn new( api_base: &str, @@ -262,18 +279,66 @@ impl GithubClient { serde_json::from_slice(&bytes).map_err(|e| GithubError::Decode(format!("{path}: {e}"))) } - /// A valid installation token, refreshed when near expiry. - async fn installation_token(&self) -> Result { + /// The App's installation id in the configured org, looked up once. + async fn installation_id(&self, auth: &mut AuthState, jwt: &str) -> Result { #[derive(Deserialize)] struct Installation { id: u64, } + if let Some(id) = auth.installation_id { + return Ok(id); + } + let path = format!("/orgs/{}/installation", self.org); + let rb = self.request(Method::GET, &path).bearer_auth(jwt); + let found: Installation = self.send(Method::GET, &path, rb).await?; + auth.installation_id = Some(found.id); + Ok(found.id) + } + + /// POST `/app/installations/{id}/access_tokens`; a `body` down-scopes the token. + async fn mint( + &self, + auth: &mut AuthState, + body: Option<&serde_json::Value>, + ) -> Result<(Secret, i64), GithubError> { #[derive(Deserialize)] struct Issued { token: String, expires_at: String, } + let jwt = self.app_jwt()?; + let installation_id = self.installation_id(auth, &jwt).await?; + let path = format!("/app/installations/{installation_id}/access_tokens"); + let mut rb = self.request(Method::POST, &path).bearer_auth(&jwt); + if let Some(body) = body { + rb = rb.json(body); + } + let issued: Issued = self.send(Method::POST, &path, rb).await?; + let expires_at = OffsetDateTime::parse(&issued.expires_at, &Rfc3339) + .map_err(|e| GithubError::Decode(format!("expires_at: {e}")))? + .unix_timestamp(); + Ok((Secret::new(issued.token), expires_at)) + } + /// Mint a fresh installation token limited to `permissions` (for example + /// `{"contents": "read"}`) on the named `repositories` (bare names in the + /// org), returning it with its expiry (unix seconds). Not cached here: the + /// caller decides how long to keep it. + pub async fn installation_token_scoped( + &self, + permissions: &serde_json::Value, + repositories: &[String], + ) -> Result<(Secret, i64), GithubError> { + let body = serde_json::json!({ + "permissions": permissions, + "repositories": repositories, + }); + let mut auth = self.auth.lock().await; + self.mint(&mut auth, Some(&body)).await + } + + /// A valid installation token, refreshed when near expiry. + pub(crate) async fn installation_token(&self) -> Result { let mut auth = self.auth.lock().await; let now = unix_now(); if let Some(t) = &auth.token @@ -281,24 +346,7 @@ impl GithubClient { { return Ok(t.value.clone()); } - let jwt = self.app_jwt()?; - let installation_id = match auth.installation_id { - Some(id) => id, - None => { - let path = format!("/orgs/{}/installation", self.org); - let rb = self.request(Method::GET, &path).bearer_auth(&jwt); - let found: Installation = self.send(Method::GET, &path, rb).await?; - auth.installation_id = Some(found.id); - found.id - } - }; - let path = format!("/app/installations/{installation_id}/access_tokens"); - let rb = self.request(Method::POST, &path).bearer_auth(&jwt); - let issued: Issued = self.send(Method::POST, &path, rb).await?; - let expires_at = OffsetDateTime::parse(&issued.expires_at, &Rfc3339) - .map_err(|e| GithubError::Decode(format!("expires_at: {e}")))? - .unix_timestamp(); - let value = Secret::new(issued.token); + let (value, expires_at) = self.mint(&mut auth, None).await?; auth.token = Some(CachedToken { value: value.clone(), expires_at, diff --git a/src/janitor.rs b/src/janitor.rs index 41c7041..ea39384 100644 --- a/src/janitor.rs +++ b/src/janitor.rs @@ -8,9 +8,10 @@ //! duplicates of a run of the same workflow, commit, event and pull request. //! Duplicate groups never include `workflow_dispatch` or `merge_group` runs, //! and a running duplicate is never cancelled. The -//! janitor is opt-in (`supersede_janitor` in the queen config) and uses its own -//! token, because cancelling runs needs `actions: write` and the queen's App -//! deliberately does not have it. +//! janitor is opt-in (`supersede_janitor` in the queen config). Cancelling +//! runs needs `actions: write`: the janitor uses the queen's App installation +//! token (minted and refreshed on demand) unless the config names a static +//! token. //! //! Safety rules, all enforced here: //! * only runs that are `queued` are candidates (`in_progress` ones too when @@ -34,7 +35,7 @@ use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; use crate::config::JanitorSettings; -use crate::github::GithubError; +use crate::github::{GithubError, TokenSource}; use crate::secret::Secret; /// Upper bound on pages of queued runs read per repository (100 per page). @@ -120,16 +121,25 @@ struct PullHead { pub struct Janitor { http: reqwest::Client, api_base: String, - token: Secret, + token: TokenSource, settings: JanitorSettings, view: Mutex, } impl Janitor { + /// A janitor with a fixed token. pub fn new( api_base: &str, token: Secret, settings: JanitorSettings, + ) -> Result { + Self::with_token_source(api_base, TokenSource::Static(token), settings) + } + + pub fn with_token_source( + api_base: &str, + token: TokenSource, + settings: JanitorSettings, ) -> Result { let http = reqwest::Client::builder() .user_agent("borg-queen-janitor") @@ -152,16 +162,18 @@ impl Janitor { &self.settings } - fn request(&self, method: Method, path: &str) -> RequestBuilder { - self.http + async fn request(&self, method: Method, path: &str) -> Result { + let token = self.token.bearer().await?; + Ok(self + .http .request(method, format!("{}{path}", self.api_base)) .header("Accept", "application/vnd.github+json") .header("X-GitHub-Api-Version", "2022-11-28") - .bearer_auth(self.token.expose()) + .bearer_auth(token.expose())) } async fn get(&self, path: &str) -> Result { - let resp = self.request(Method::GET, path).send().await?; + let resp = self.request(Method::GET, path).await?.send().await?; let status = resp.status(); if !status.is_success() { return Err(status_error("GET", path, status.as_u16(), resp).await); @@ -174,7 +186,7 @@ impl Janitor { /// cancelled (already finished), which is fine. async fn cancel(&self, repo: &str, run: u64) -> Result { let path = format!("/repos/{repo}/actions/runs/{run}/cancel"); - let resp = self.request(Method::POST, &path).send().await?; + let resp = self.request(Method::POST, &path).await?.send().await?; let status = resp.status(); match status.as_u16() { 200..=299 => Ok(true), diff --git a/src/lib.rs b/src/lib.rs index 7adeda9..8be6a35 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,6 +23,7 @@ pub mod health; pub mod homeostasis; pub mod janitor; pub mod lead; +pub mod mirror_auth; pub mod preempt; pub mod probe; pub mod profile; diff --git a/src/main.rs b/src/main.rs index 0ef3754..dc9be4f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,8 +11,9 @@ use borg::drone::exec::SystemExec; use borg::drone::install::{self, InstallOptions, Platform}; use borg::drone::{Drone, pins}; use borg::enroll; -use borg::github::GithubClient; +use borg::github::{GithubClient, TokenSource}; use borg::janitor::Janitor; +use borg::mirror_auth::MirrorAuth; use borg::probe::{self, system::SystemHost}; use borg::queen::Queen; use borg::requeue::Requeuer; @@ -408,18 +409,31 @@ fn private_key(cfg: &QueenConfig) -> Result> { } } -fn janitor_token(settings: &JanitorSettings) -> Result { +/// The janitor's token: `token_path`, else `BORG_JANITOR_TOKEN`, else the +/// queen's own App installation token (the App then needs Actions: write). +fn janitor_token(settings: &JanitorSettings, cfg: &QueenConfig) -> Result { let raw = match &settings.token_path { - Some(path) => std::fs::read_to_string(path) - .with_context(|| format!("read janitor token {}", path.display()))?, - None => std::env::var("BORG_JANITOR_TOKEN") - .context("set supersede_janitor.token_path or BORG_JANITOR_TOKEN")?, + Some(path) => Some( + std::fs::read_to_string(path) + .with_context(|| format!("read janitor token {}", path.display()))?, + ), + None => std::env::var("BORG_JANITOR_TOKEN").ok(), }; - let token = raw.trim(); - if token.is_empty() { - bail!("the janitor token is empty"); + match raw { + Some(raw) => { + let token = raw.trim(); + if token.is_empty() { + bail!("the janitor token is empty"); + } + Ok(TokenSource::Static(borg::secret::Secret::new(token))) + } + None => Ok(TokenSource::App(Arc::new(GithubClient::new( + &cfg.github.api_base, + &cfg.github.org, + &cfg.github.app_id, + &private_key(cfg)?, + )?))), } - Ok(borg::secret::Secret::new(token)) } async fn open_store(cfg: &QueenConfig) -> Result { @@ -459,23 +473,31 @@ async fn run_queen(path: &Path) -> Result<()> { let mut requeuer = None; let janitor = match &cfg.supersede_janitor { Some(settings) => { - let token = janitor_token(settings)?; + let token = janitor_token(settings, &cfg)?; tracing::info!( - "supersede janitor enabled ({})", + "supersede janitor enabled ({}, {} token)", if settings.dry_run { "dry-run" } else { "cancelling" + }, + match token { + TokenSource::Static(_) => "static", + TokenSource::App(_) => "App installation", } ); if cfg.preemption.is_some() { - requeuer = Some(Requeuer::new( + requeuer = Some(Requeuer::with_token_source( &cfg.github.api_base, token.clone(), settings, )?); } - Some(Janitor::new(&cfg.github.api_base, token, settings.clone())?) + Some(Janitor::with_token_source( + &cfg.github.api_base, + token, + settings.clone(), + )?) } None => None, }; @@ -499,6 +521,21 @@ async fn run_queen(path: &Path) -> Result<()> { if let Some(r) = requeuer { queen = queen.with_requeuer(r); } + if let Some(settings) = &queen.config().mirror_auth { + let pem = settings.private_key(std::env::var("BORG_MIRROR_APP_PRIVATE_KEY").ok())?; + let client = GithubClient::new( + &queen.config().github.api_base, + &queen.config().github.org, + &settings.app_id, + &pem, + )?; + let mirror = MirrorAuth::new(client, &queen.config().github.org, &settings.repositories); + tracing::info!( + "git mirror credentials enabled for {} repositories", + settings.repositories.len() + ); + queen = queen.with_mirror_auth(mirror); + } let queen = Arc::new(queen); api::serve(queen, listener, wait_for_shutdown()).await?; Ok(()) diff --git a/src/mirror_auth.rs b/src/mirror_auth.rs new file mode 100644 index 0000000..f40a443 --- /dev/null +++ b/src/mirror_auth.rs @@ -0,0 +1,102 @@ +//! Short-lived read-only GitHub tokens for the drones' git mirrors. +//! +//! A second GitHub App key lives on the queen only. The queen mints +//! installation tokens down-scoped to `contents: read` on the mirrored +//! repositories and hands them to drones in the heartbeat response; the App +//! key never reaches a drone. The token is never logged and `/status` shows +//! only whether mirror auth is configured and when the current token expires. + +use std::time::{SystemTime, UNIX_EPOCH}; + +use serde::{Deserialize, Serialize}; + +use crate::github::{GithubClient, GithubError}; +use crate::secret::Secret; + +/// Mint a new token when the cached one has less than this much life left. +const REMINT_BELOW_SECS: i64 = 10 * 60; + +/// What a drone needs to fetch private mirrors with. `repositories` are +/// `owner/name`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct MirrorCredential { + pub token: Secret, + pub expires_at: i64, + pub repositories: Vec, +} + +/// `/status`: never the token. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub struct MirrorAuthView { + pub configured: bool, + pub expires_at: Option, +} + +pub struct MirrorAuth { + client: GithubClient, + org: String, + /// Bare names, as GitHub's `repositories` parameter wants them. + repositories: Vec, + cached: tokio::sync::Mutex>, + expires_at: std::sync::atomic::AtomicI64, +} + +fn unix_now() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(0)) +} + +impl MirrorAuth { + pub fn new(client: GithubClient, org: &str, repositories: &[String]) -> Self { + Self { + client, + org: org.to_owned(), + repositories: repositories.to_vec(), + cached: tokio::sync::Mutex::new(None), + expires_at: std::sync::atomic::AtomicI64::new(0), + } + } + + /// The current credential, minting a new token when under ten minutes of + /// life remain. + pub async fn credential(&self) -> Result { + let mut cached = self.cached.lock().await; + let reusable = cached + .as_ref() + .filter(|(_, expires_at)| expires_at - REMINT_BELOW_SECS > unix_now()); + let (token, expires_at) = match reusable { + Some((token, expires_at)) => (token.clone(), *expires_at), + None => { + let minted = self + .client + .installation_token_scoped( + &serde_json::json!({"contents": "read"}), + &self.repositories, + ) + .await?; + *cached = Some(minted.clone()); + self.expires_at + .store(minted.1, std::sync::atomic::Ordering::Relaxed); + minted + } + }; + Ok(MirrorCredential { + token, + expires_at, + repositories: self + .repositories + .iter() + .map(|r| format!("{}/{r}", self.org)) + .collect(), + }) + } + + pub fn view(&self) -> MirrorAuthView { + let expires_at = self.expires_at.load(std::sync::atomic::Ordering::Relaxed); + MirrorAuthView { + configured: true, + expires_at: (expires_at > 0).then_some(expires_at), + } + } +} diff --git a/src/queen.rs b/src/queen.rs index 976ead4..cbe5b8a 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -35,6 +35,7 @@ use crate::github::{GithubClient, GithubError, QueuedJob, Runner, WorkflowJob}; use crate::health::{DroneHealth, DroneInfo, Outcome, OutcomeKind, evaluate}; use crate::homeostasis::{Controller, KeyInput, KeyReport}; use crate::janitor::Janitor; +use crate::mirror_auth::MirrorAuth; use crate::preempt::{PreemptInput, Running, select as select_victims}; use crate::probe::Capacity; use crate::profile::{ClassProfile, build_profiles}; @@ -167,6 +168,7 @@ pub struct Queen { janitor: Option, requeuer: Option, tailnet: Option, + mirror_auth: Option, homeostasis: std::sync::Mutex, profile_cache: std::sync::Mutex)>>, wake: std::sync::Mutex, @@ -346,6 +348,7 @@ impl Queen { janitor: None, requeuer: None, tailnet: None, + mirror_auth: None, homeostasis: std::sync::Mutex::new(Controller::new()), profile_cache: std::sync::Mutex::new(None), wake: std::sync::Mutex::new(WakePlanner::new()), @@ -379,6 +382,17 @@ impl Queen { self } + /// Hand drones short-lived read-only tokens for their git mirrors. + #[must_use] + pub fn with_mirror_auth(mut self, mirror_auth: MirrorAuth) -> Self { + self.mirror_auth = Some(mirror_auth); + self + } + + pub fn mirror_auth(&self) -> Option<&MirrorAuth> { + self.mirror_auth.as_ref() + } + /// Replace how magic packets are sent (tests use a recording fake). #[must_use] pub fn with_wake_sender(mut self, sender: Arc) -> Self { diff --git a/src/requeue.rs b/src/requeue.rs index 0183a59..dc26ae2 100644 --- a/src/requeue.rs +++ b/src/requeue.rs @@ -4,9 +4,9 @@ //! slot. The queen records `(repo, run, job)` in the `preemptions` table with //! `requeued_at` unset. This pass re-runs each recorded job through the jobs //! rerun API once its run has completed (GitHub refuses to re-run a job of a -//! run that is still in progress). It uses the supersede janitor's token, -//! which has `actions: write`, and the janitor's `dry_run` and -//! `max_cancels_per_pass` settings. +//! run that is still in progress). It uses the supersede janitor's token +//! source (the queen's App token by default), which has `actions: write`, and +//! the janitor's `dry_run` and `max_cancels_per_pass` settings. //! //! * a run that is not `completed` is left for a later pass; //! * a rerun answered 403, 404 or 409 (job gone, retention expired, or already @@ -23,7 +23,7 @@ use reqwest::{Method, RequestBuilder}; use serde::{Deserialize, Serialize}; use crate::config::JanitorSettings; -use crate::github::GithubError; +use crate::github::{GithubError, TokenSource}; use crate::janitor::status_error; use crate::secret::Secret; use crate::store::{Store, StoreError}; @@ -78,7 +78,7 @@ struct Run { pub struct Requeuer { http: reqwest::Client, api_base: String, - token: Secret, + token: TokenSource, dry_run: bool, max_per_pass: u32, view: Mutex, @@ -90,6 +90,14 @@ impl Requeuer { api_base: &str, token: Secret, settings: &JanitorSettings, + ) -> Result { + Self::with_token_source(api_base, TokenSource::Static(token), settings) + } + + pub fn with_token_source( + api_base: &str, + token: TokenSource, + settings: &JanitorSettings, ) -> Result { let http = reqwest::Client::builder() .user_agent("borg-queen-requeue") @@ -109,17 +117,19 @@ impl Requeuer { self.view.lock().expect("requeue view lock").clone() } - fn request(&self, method: Method, path: &str) -> RequestBuilder { - self.http + async fn request(&self, method: Method, path: &str) -> Result { + let token = self.token.bearer().await?; + Ok(self + .http .request(method, format!("{}{path}", self.api_base)) .header("Accept", "application/vnd.github+json") .header("X-GitHub-Api-Version", "2022-11-28") - .bearer_auth(self.token.expose()) + .bearer_auth(token.expose())) } async fn run_status(&self, repo: &str, run: u64) -> Result { let path = format!("/repos/{repo}/actions/runs/{run}"); - let resp = self.request(Method::GET, &path).send().await?; + let resp = self.request(Method::GET, &path).await?.send().await?; let status = resp.status(); if !status.is_success() { return Err(status_error("GET", &path, status.as_u16(), resp).await); @@ -133,7 +143,7 @@ impl Requeuer { /// Re-run a job. `Ok(Some(code))` is a terminal refusal (403, 404, 409). async fn rerun(&self, repo: &str, job: u64) -> Result, GithubError> { let path = format!("/repos/{repo}/actions/jobs/{job}/rerun"); - let resp = self.request(Method::POST, &path).send().await?; + let resp = self.request(Method::POST, &path).await?.send().await?; match resp.status().as_u16() { 200..=299 => Ok(None), code @ (403 | 404 | 409) => Ok(Some(code)), diff --git a/tests/capacity.rs b/tests/capacity.rs index 3b20130..e871659 100644 --- a/tests/capacity.rs +++ b/tests/capacity.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use serde_json::{Value, json}; use support::Harness; diff --git a/tests/contract.rs b/tests/contract.rs index 716a476..5705909 100644 --- a/tests/contract.rs +++ b/tests/contract.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use borg::enroll; use serde_json::{Value, json}; diff --git a/tests/drone.rs b/tests/drone.rs index 992de54..e5b44cf 100644 --- a/tests/drone.rs +++ b/tests/drone.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use std::fs; use std::io; diff --git a/tests/health.rs b/tests/health.rs index 60ec8ba..87041b5 100644 --- a/tests/health.rs +++ b/tests/health.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use borg::classes::{Class, Scope}; use borg::store::{LeaseRow, LeaseStatus}; diff --git a/tests/host_suite.rs b/tests/host_suite.rs new file mode 100644 index 0000000..fb447a6 --- /dev/null +++ b/tests/host_suite.rs @@ -0,0 +1,14 @@ +//! Host execution, updates, and protocol tests share one fixture module. +//! Filter a module with `cargo test --test host_suite drone::` in the standalone +//! crate, or add `-p borg` from the Mono Rust workspace. + +mod support; + +mod admin_cli; +mod contract; +mod drone; +mod job_user; +mod rollout; +mod run; +mod tailnet; +mod update; diff --git a/tests/janitor.rs b/tests/janitor.rs index 886b007..2554b81 100644 --- a/tests/janitor.rs +++ b/tests/janitor.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use serde_json::json; use support::Harness; diff --git a/tests/job_user.rs b/tests/job_user.rs index 2b8b722..19cd41a 100644 --- a/tests/job_user.rs +++ b/tests/job_user.rs @@ -1,6 +1,6 @@ //! Real privilege separation. Needs root (to change user) and util-linux //! `setpriv`; skipped otherwise. Run it with, for example, -//! `sudo -E cargo test --test job_user -- --nocapture`. +//! `sudo -E cargo test --test host_suite job_user:: -- --nocapture`. use std::fs; use std::os::unix::fs::PermissionsExt; diff --git a/tests/queen.rs b/tests/queen.rs index 9f4a8f6..f2fe066 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -1241,3 +1241,87 @@ async fn a_standby_started_during_an_outage_is_reconciled_when_the_queen_returns let body = standby_beat(&h, &dep, &[&id], &[]).await; assert_eq!(body["drop_standby"], json!([id])); } + +fn mirror_cfg() -> Value { + json!({"mirror_auth": { + "app_id": "12345", + "private_key_path": "/etc/borg/mirror-app.pem", + "repositories": ["app"] + }}) +} + +#[tokio::test] +async fn the_minter_scopes_the_token_to_contents_read_on_the_named_repositories() { + let h = Harness::start(json!({})).await; + let client = + borg::github::GithubClient::new(&h.gh_url, support::ORG, "12345", h.private_pem.as_bytes()) + .unwrap(); + let (token, expires_at) = client + .installation_token_scoped(&json!({"contents": "read"}), &["app".to_owned()]) + .await + .unwrap(); + assert_eq!(token.expose(), support::MIRROR_TOKEN); + assert!(expires_at > OffsetDateTime::now_utc().unix_timestamp()); + let bodies = h.state().scoped_requests.clone(); + assert_eq!(bodies.len(), 1); + assert_eq!(bodies[0]["permissions"]["contents"], "read"); + assert_eq!(bodies[0]["repositories"], json!(["app"])); + assert_eq!( + bodies[0]["permissions"].as_object().unwrap().len(), + 1, + "no other permission is requested" + ); +} + +#[tokio::test] +async fn heartbeats_carry_the_mirror_credential_only_when_configured() { + let off = Harness::start(json!({})).await; + let token = off.enroll("d1").await; + let body = off.heartbeat_with(&token, json!({})).await; + assert!(body.get("mirror_credential").is_none(), "{body}"); + assert_eq!( + off.status().await["mirror_auth"], + json!({"configured": false, "expires_at": null}) + ); + + let on = Harness::start(mirror_cfg()).await; + let token = on.enroll("d1").await; + let body = on.heartbeat_with(&token, json!({})).await; + let cred = &body["mirror_credential"]; + assert_eq!(cred["token"], support::MIRROR_TOKEN); + assert_eq!(cred["repositories"], json!(["example-org/app"])); + let expires_at = cred["expires_at"].as_i64().unwrap(); + assert!(expires_at > OffsetDateTime::now_utc().unix_timestamp()); + + let status = on.status().await; + assert_eq!( + status["mirror_auth"], + json!({"configured": true, "expires_at": expires_at}) + ); + assert!( + !status.to_string().contains(support::MIRROR_TOKEN), + "/status never shows the token" + ); +} + +#[tokio::test] +async fn the_mirror_token_is_cached_until_ten_minutes_of_life_remain() { + let h = Harness::start(mirror_cfg()).await; + let token = h.enroll("d1").await; + for _ in 0..3 { + h.heartbeat_with(&token, json!({})).await; + } + assert_eq!( + h.state().scoped_requests.len(), + 1, + "heartbeats reuse the cached token" + ); + + // A token that lives five minutes is always inside the re-mint window. + let short = Harness::start(mirror_cfg()).await; + short.state().token_ttl_secs = Some(300); + let token = short.enroll("d1").await; + short.heartbeat_with(&token, json!({})).await; + short.heartbeat_with(&token, json!({})).await; + assert_eq!(short.state().scoped_requests.len(), 2); +} diff --git a/tests/requeue.rs b/tests/requeue.rs index 513b09e..c5d4a1b 100644 --- a/tests/requeue.rs +++ b/tests/requeue.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use serde_json::json; use support::Harness; diff --git a/tests/rollout.rs b/tests/rollout.rs index 1f1fe44..4f45d10 100644 --- a/tests/rollout.rs +++ b/tests/rollout.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use borg::classes::{Class, Scope}; use borg::store::{LeaseRow, LeaseStatus}; diff --git a/tests/sccache_slots.rs b/tests/sccache_slots.rs index 857310c..ce322b5 100644 --- a/tests/sccache_slots.rs +++ b/tests/sccache_slots.rs @@ -31,6 +31,39 @@ fn run_with(env: &[(String, String)], program: &str, args: &[&str]) -> bool { SystemExec.run("env", &argv).is_ok_and(|o| o.success) } +struct TestServer { + env: Vec<(String, String)>, +} + +impl TestServer { + fn start(root: &Path) -> Self { + let env = sccache::server_env(&root.join("cache").to_string_lossy(), "1G", free_port()); + // sccache gives an inherited Unix socket precedence over the TCP port. + // Keep start, compiler, stats and cleanup commands on our own endpoint. + #[cfg(unix)] + let env = { + let mut env = env; + env.push(( + "SCCACHE_SERVER_UDS".to_owned(), + root.join("socket").to_string_lossy().into_owned(), + )); + env + }; + let server = Self { env }; + assert!( + sccache::start_server(&SystemExec, "sccache", &server.env), + "sccache server starts" + ); + server + } +} + +impl Drop for TestServer { + fn drop(&mut self) { + let _ = run_with(&self.env, "sccache", &["--stop-server"]); + } +} + fn write_crate(dir: &Path) { fs::create_dir_all(dir.join("src")).unwrap(); fs::write( @@ -54,11 +87,10 @@ fn workspace_crates_hit_the_compile_cache_when_a_slot_is_reused() { let tmp = tempfile::tempdir().unwrap(); let work = tmp.path().join("work"); let cargo_home = tmp.path().join("cargo"); - let port = free_port(); - let dir = tmp.path().join("sccache").to_string_lossy().into_owned(); let flags = rustflags::remap_flags(&work, 2, &cargo_home).join(" "); - let mut env = sccache::server_env(&dir, "1G", port); + let mut server = TestServer::start(tmp.path()); + let env = &mut server.env; env.extend([ ("RUSTC_WRAPPER".to_owned(), "sccache".to_owned()), ("CARGO_INCREMENTAL".to_owned(), "0".to_owned()), @@ -72,17 +104,12 @@ fn workspace_crates_hit_the_compile_cache_when_a_slot_is_reused() { ), ("RUSTFLAGS".to_owned(), flags), ]); - assert!( - sccache::start_server(&SystemExec, "sccache", &env), - "sccache server starts" - ); - let build = |slot: usize| { let crate_dir = work.join(format!("slot-{slot}")).join("demo"); write_crate(&crate_dir); let manifest = crate_dir.join("Cargo.toml"); assert!(run_with( - &env, + env, "cargo", &[ "build", @@ -92,7 +119,7 @@ fn workspace_crates_hit_the_compile_cache_when_a_slot_is_reused() { manifest.to_str().unwrap() ] )); - sccache::show_stats(&SystemExec, "sccache", &env).expect("sccache answers") + sccache::show_stats(&SystemExec, "sccache", env).expect("sccache answers") }; let first = build(0); @@ -122,6 +149,51 @@ fn workspace_crates_hit_the_compile_cache_when_a_slot_is_reused() { // working directory and every CARGO_* variable (CARGO_MANIFEST_DIR), which // differ between slots whatever rustc flags are passed. The numbers above // record what happens. +} - let _ = run_with(&env, "sccache", &["--stop-server"]); +#[cfg(unix)] +#[test] +fn slot_fixture_preserves_an_inherited_sccache_server() { + if !(available("cargo") && available("rustc") && available("sccache")) { + println!("skipping: cargo, rustc and sccache are required"); + return; + } + let tmp = tempfile::tempdir().unwrap(); + let outer = TestServer::start(tmp.path()); + let source = tmp.path().join("sentinel.rs"); + fs::write(&source, "pub fn sentinel() -> u8 { 7 }\n").unwrap(); + assert!(run_with( + &outer.env, + "sccache", + &[ + "rustc", + "--crate-name", + "sentinel", + "--crate-type", + "lib", + "--emit=dep-info,link", + source.to_str().unwrap(), + "--out-dir", + tmp.path().to_str().unwrap(), + ], + )); + let before = sccache::show_stats(&SystemExec, "sccache", &outer.env).unwrap(); + assert!(before.1 > 0, "the outer server has a sentinel cache miss"); + + // Use a child process so ambient environment is exercised without changing + // this test process's environment while other tests may be running. + assert!(run_with( + &outer.env, + std::env::current_exe().unwrap().to_str().unwrap(), + &[ + "--exact", + "workspace_crates_hit_the_compile_cache_when_a_slot_is_reused", + "--nocapture", + ], + )); + assert_eq!( + sccache::show_stats(&SystemExec, "sccache", &outer.env), + Some(before), + "the outer server survives fixture cleanup with unchanged counters" + ); } diff --git a/tests/suite.rs b/tests/suite.rs new file mode 100644 index 0000000..6137ee8 --- /dev/null +++ b/tests/suite.rs @@ -0,0 +1,12 @@ +//! Scheduling and service-state tests share one executable and fixture module. +//! Filter a module with `cargo test --test suite units::` in the standalone +//! crate, or add `-p borg` from the Mono Rust workspace. + +mod support; + +mod capacity; +mod health; +mod janitor; +mod requeue; +mod units; +mod wake; diff --git a/tests/support/mod.rs b/tests/support/mod.rs index cfe274b..a7d0880 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -1,5 +1,5 @@ //! A fake GitHub REST API and a running queen, for integration tests. -//! Shared by several test binaries, each of which uses only part of it. +//! Compiled once per integration suite; individual modules use parts of it. #![allow(dead_code)] use std::collections::HashMap; @@ -26,6 +26,7 @@ use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; pub const INSTALLATION_TOKEN: &str = "ghs_installation_for_tests"; +pub const MIRROR_TOKEN: &str = "ghs_mirror_for_tests"; pub const JANITOR_TOKEN: &str = "ghp_janitor_for_tests"; pub const ORG: &str = "example-org"; @@ -38,6 +39,10 @@ pub struct MockState { pub mints: Vec, pub deleted: Vec, pub token_requests: u32, + /// Bodies of down-scoped token requests (the mirror App's). + pub scoped_requests: Vec, + /// Lifetime of issued tokens in seconds (default one hour). + pub token_ttl_secs: Option, pub jwt_rejections: u32, pub fail_mint: bool, pub fail_runner_list: bool, @@ -136,15 +141,29 @@ async fn installation(State(s): State, h: HeaderMap) -> Result, h: HeaderMap) -> Result, StatusCode> { +async fn access_token( + State(s): State, + h: HeaderMap, + body: axum::body::Bytes, +) -> Result, StatusCode> { check_jwt(&s, &h)?; - lock(&s).token_requests += 1; - let expires = (OffsetDateTime::now_utc() + time::Duration::hours(1)) + let scoped: Option = serde_json::from_slice(&body).ok(); + let mut st = lock(&s); + let ttl = st.token_ttl_secs.unwrap_or(3600); + let expires = (OffsetDateTime::now_utc() + time::Duration::seconds(ttl)) .format(&Rfc3339) .unwrap(); - Ok(Json( - json!({"token": INSTALLATION_TOKEN, "expires_at": expires}), - )) + let token = match scoped { + Some(body) => { + st.scoped_requests.push(body); + MIRROR_TOKEN + } + None => { + st.token_requests += 1; + INSTALLATION_TOKEN + } + }; + Ok(Json(json!({"token": token, "expires_at": expires}))) } fn check_janitor(headers: &HeaderMap) -> Result<(), StatusCode> { @@ -360,6 +379,9 @@ pub struct Harness { pub queen: Arc, pub url: String, pub http: reqwest::Client, + /// The fake GitHub API base and the App key it accepts. + pub gh_url: String, + pub private_pem: String, } pub fn capacity() -> Capacity { @@ -476,6 +498,15 @@ impl Harness { settings, }); } + if let Some(settings) = queen.config().mirror_auth.clone() { + let client = + GithubClient::new(&gh_url, ORG, &settings.app_id, private_pem.as_bytes()).unwrap(); + queen = queen.with_mirror_auth(borg::mirror_auth::MirrorAuth::new( + client, + ORG, + &settings.repositories, + )); + } let queen = Arc::new(queen); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -494,6 +525,8 @@ impl Harness { queen, url, http: reqwest::Client::new(), + gh_url, + private_pem, } } diff --git a/tests/tailnet.rs b/tests/tailnet.rs index 6942e05..350cebf 100644 --- a/tests/tailnet.rs +++ b/tests/tailnet.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use std::net::{IpAddr, Ipv4Addr}; diff --git a/tests/units.rs b/tests/units.rs index e2b4c91..e0d61b0 100644 --- a/tests/units.rs +++ b/tests/units.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use borg::enroll; use serde_json::{Value, json}; diff --git a/tests/wake.rs b/tests/wake.rs index e638eb9..434ccba 100644 --- a/tests/wake.rs +++ b/tests/wake.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use serde_json::{Value, json}; use support::Harness; From 810c48d8e419b3964fd5a11a8de333fbe6f4b149 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Mon, 12 Oct 2026 04:19:57 +0000 Subject: [PATCH 13/27] chore: project borg from Mono 6a106187736f --- .repository-projection.json | 6 +- README.md | 5 +- docs/capability-admission.md | 69 +++ docs/operations.md | 7 + docs/placement-diagnostics.md | 33 ++ docs/remote-execution.md | 39 ++ docs/work-unit-fairness.md | 29 ++ src/allocator.rs | 347 +++++++++++++-- src/api.rs | 44 ++ src/drone/capabilities.rs | 707 +++++++++++++++++++++++++++++++ src/drone/config.rs | 6 + src/drone/exec.rs | 83 +++- src/drone/install.rs | 1 + src/drone/jobuser.rs | 44 ++ src/drone/mod.rs | 17 +- src/envcontract.rs | 367 +++++++++++++++- src/explain.rs | 163 ++++++++ src/lib.rs | 1 + src/main.rs | 61 ++- src/probe/mod.rs | 3 +- src/queen.rs | 612 ++++++++++++++++++++++----- src/run.rs | 351 +++++++++++++--- src/store.rs | 113 +++++ src/units.rs | 35 +- src/units_api.rs | 5 +- tests/contract.rs | 60 +++ tests/drone.rs | 194 +++++++++ tests/explain.rs | 547 ++++++++++++++++++++++++ tests/explain_cli.rs | 224 ++++++++++ tests/placement_batch.rs | 238 +++++++++++ tests/run.rs | 476 ++++++++++++++++++++- tests/sccache_slots.rs | 55 ++- tests/units.rs | 766 ++++++++++++++++++++++++++++++++++ 33 files changed, 5473 insertions(+), 235 deletions(-) create mode 100644 docs/capability-admission.md create mode 100644 docs/placement-diagnostics.md create mode 100644 docs/remote-execution.md create mode 100644 docs/work-unit-fairness.md create mode 100644 src/drone/capabilities.rs create mode 100644 src/explain.rs create mode 100644 tests/explain.rs create mode 100644 tests/explain_cli.rs create mode 100644 tests/placement_batch.rs diff --git a/.repository-projection.json b/.repository-projection.json index cd48cd7..45d0b70 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "2e833192cc635eb2dcb5f4cb02562547636a982a", + "sourceSha": "6a106187736f994ddbc328e1f8bfe3d4786de24d", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "c51a232c724d6915b1fa0b001a635c49519b01ae", + "priorProjectedBase": "9ea7dac52a1c013677042fb1d313452ade1946c3", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "58e3b42b548a0e1b22cbee864fb0e787e0a8970a40e8c6d23a43758175bc8928", + "contentDigest": "17ec6815470882c73d267d67b062fe816ee7a97a411e18b70e598a8ff73bc250", "publicationEligible": true } diff --git a/README.md b/README.md index 2a3c0a5..258ce21 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,10 @@ and job telemetry at `/report`. | --- | --- | | [Configuration](docs/operations.md#config-files) | TOML and JSON, validation, defaults, and redacted secrets. | | [Scheduling](docs/operations.md#drone-bids-and-demand-pressure) | Bids, capacity, starvation protection, host health, and warm floors. | -| [Work units](docs/operations.md#borg-run) | Run commands on the fleet with `borg run`, including local fallback. | +| [Work units](docs/remote-execution.md) | Remote-only execution, local fallback, and JSON execution results. | +| [Capability admission](docs/capability-admission.md) | Job-user smoke tests and per-class requirements. | +| [Work-unit fairness](docs/work-unit-fairness.md) | Capacity shares, client rotation, and starvation protection. | +| [Placement diagnostics](docs/placement-diagnostics.md) | Explain recorded placement choices with `borg explain`. | | [Host caches](docs/operations.md#host-caches) | Cache limits, collection, git mirrors, slots, and compile-cache behavior. | | [Drone updates](docs/operations.md#updating-a-drone) | Canary rollout, update manifests, health checks, and rollback. | | [Operations reference](docs/operations.md) | APIs, telemetry fields, enrollment, and optional controllers. | diff --git a/docs/capability-admission.md b/docs/capability-admission.md new file mode 100644 index 0000000..6c20464 --- /dev/null +++ b/docs/capability-admission.md @@ -0,0 +1,69 @@ +# Execution capability admission + +The build environment contract can require execution proofs in addition to +reported tools and capacity. Probes run through the same job-user executor used +by jobs. Service-user `docker info` remains inventory and does not prove that a +job can run a container. + +The named capabilities are sandbox namespace execution, Docker execution, work +filesystem writability, and cache filesystem writability. Each reported proof +keeps its own observation timestamp. Failed, absent, or stale required proofs +make the affected class ineligible. Classes with no capability requirements +retain their existing admission behavior. + +Configure required capabilities in the queen's existing `env_contract` block, +globally or per class. Class lists replace the global list. The `capabilities` +check supports `required`, `advisory`, and `off`, like other contract checks. +Use `/status` to inspect the reason before changing workload routing. + +```toml +[env_contract] +enabled = true +required_capabilities = ["work_write", "cache_write"] + +[env_contract.classes.heavy] +required_capabilities = ["sandbox", "docker", "work_write", "cache_write"] +``` + +Retain the other checks and thresholds in the existing queen configuration. + +In the drone configuration, select an image already loaded on that host: + +```toml +[capability_probe] +timeout_bin = "timeout" +timeout_secs = 5 +bwrap_bin = "/usr/bin/bwrap" +docker_image = "example/smoke:local" +``` + +Drone `capability_probe` settings choose the timeout binary, timeout duration, +bubblewrap binary, and an optional Docker smoke image. The default command +timeout is five seconds. On Unix, Borg starts each probe in a private process +group. A trusted shell supervisor kills that group when timeout sends TERM; +this does not depend on the timeout utility's final KILL reaching descendants. +The utility's two-second KILL escalation remains a fallback. macOS installations +using GNU coreutils can configure `gtimeout`. Executors without private process +group support report unknown without starting a smoke command. + +The Docker image must already be present and contain `/bin/true`. Borg uses +`--pull=never`, disables container networking, and removes its uniquely named +smoke container after execution. Without a configured image, Docker execution +is reported as unknown. No probe installs packages or changes AppArmor, +namespace, or host permission settings. + +Work and cache probes create, write, read and remove a private temporary child +of the configured root. Smoke commands and cleanup are bounded, discard output, +and reuse observations for a minute. Refreshes run in the drone's existing +background task set; heartbeats send the cached proofs without waiting for a +smoke test. The first heartbeats may report no proof, which keeps required +capabilities ineligible until the refresh completes. Drain starts no refreshes +and joins those already running, including their cleanup. + +If a refresh cannot start, a cached proof retains its original observation +time. A completed failing smoke test records a fresh failure. Cleanup is best +effort when the Docker daemon is unresponsive. + +Deploy the new drone binary first, verify capability observations, then add +requirements for the classes whose workloads need them. Older drones cannot +satisfy a required proof they do not report. diff --git a/docs/operations.md b/docs/operations.md index b208562..4aae546 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -274,6 +274,9 @@ Targets are IPv4 addresses or `addr/len` networks. See `/status` `wake`. ### Build environment contract +[Capability admission](capability-admission.md) adds bounded job-user execution +proofs for sandbox, Docker, work-root and cache-root access. + Optional (`env_contract.enabled`, off by default). Before a CI runner lease or a work unit is placed on a drone, the queen checks what the drone last reported about its build environment. Five checks, each `required`, `advisory` or `off` @@ -384,6 +387,10 @@ the command locally, so wrappers can call it unconditionally. ### Work units (gate and hook runs) +See [remote execution](remote-execution.md) for remote-only runs and JSON +results, [work-unit fairness](work-unit-fairness.md) for capacity shares, and +[placement diagnostics](placement-diagnostics.md) for `borg explain`. + Optional (`units.enabled`, off by default). Clients such as `borg run`, the pre-commit hook and `mq` submit whole commands; drones run them. Enroll a client with `borg queen enroll-client ` (prints a token once). diff --git a/docs/placement-diagnostics.md b/docs/placement-diagnostics.md new file mode 100644 index 0000000..e7b8936 --- /dev/null +++ b/docs/placement-diagnostics.md @@ -0,0 +1,33 @@ +# Explaining placement + +`borg explain` reads decisions captured by the allocator at placement time. +These records retain the considered hosts, remaining owner-capped resources, +admission rejections, reservations, bid and profile scores, and selected host. +A current `/status` snapshot cannot reconstruct that earlier capacity state. + +```sh +borg explain borg-ci-medium-pr --queen https://queen.example.com +borg explain unit-id --queen https://queen.example.com --json +borg explain unit-id --config /etc/borg/queen.toml +``` + +Remote unit queries use the existing client token from `BORG_CLIENT_TOKEN`, +`BORG_CLIENT_TOKEN_FILE`, or the default client-token file. They return only +that client's diagnostic records. Runner-label queries expose runner decisions. +Local `--config` access reads the queen database as an operator. + +A selected host is an allocator decision. The record's outcome separately +reports whether the queen successfully minted a runner or persisted the unit +lease. A failed owning operation must not appear as a completed placement. + +The existing queen SQLite database retains at most 1,000 decision records and +returns the newest 20 for a query. Each allocator call captures at most 128 +decisions and 128 candidates per decision, retaining the selected host when +candidates are truncated. Each stored decision is limited to 128 KiB. +Truncated candidate lists carry the original count. No command +arguments, environment values or tokens enter placement history. Queries with +no retained decision report that state explicitly. + +History is diagnostic: insertion failure is logged without changing scheduling. +Old rows are pruned during insertion, and database startup creates the table +for an existing installation. A queen database backup preserves this history. diff --git a/docs/remote-execution.md b/docs/remote-execution.md new file mode 100644 index 0000000..6e7ab85 --- /dev/null +++ b/docs/remote-execution.md @@ -0,0 +1,39 @@ +# Fleet-only command execution + +Use `borg run --remote-only` when the command must run on an approved fleet host. +If Borg cannot complete remote execution, it returns exit code 125 and does not +start the command locally. A command that runs remotely and exits nonzero keeps +its command exit code. Without this option, the existing local fallback remains. + +```sh +borg run --remote-only --platform linux-x86_64 --wait-secs 120 \ + --timeout 1200 --result-file ./build-result.json -- cargo test -p example +``` + +`BORG_RUN_REMOTE_ONLY=true` also enables the mode. `--wait-secs` bounds pickup +waiting. After reading unit metadata, the client bounds polling by the pickup +wait plus the queen-admitted command timeout if the queen stops making progress. +Connection and cancellation requests have their own bounded timeouts. An unpublished tree needs `--publish` +so a drone can fetch its bundle. Without a configured queen or client token, +remote-only execution fails. + +The optional `--result-file` writes JSON separately from streamed command output. +Its fields include `schema_version`, `requested_tree_sha`, `argv`, `execution` +(`remote`, `local`, or `none`), `provenance` (`new`, `joined`, `reused`, or +`fallback`), `unit_id`, `host`, `exit_code`, `infrastructure_error`, and `receipt`. +A null tree or host means it could not be established. The requested tree does +not describe the mutable checkout used by local fallback. + +Borg validates a reused receipt against the submitted repository, tree, command, +platform and toolchain. The receipt digest checks consistency. Configure +authenticated TLS for the queen endpoint. A reused successful receipt names the +original execution host. + +An invocation that joins another unit does not cancel that unit when its own +wait ends. Cancellation of a newly created unit is best effort; a client timeout +does not establish that a disconnected drone has stopped executing it. + +An unwritable result destination fails before execution. A later write failure +can return 125 after the command ran. If terminal metadata cannot be verified, +Borg preserves the observed remote exit code and records the metadata error; +host and receipt remain unavailable. Provenance may be null before submission. diff --git a/docs/work-unit-fairness.md b/docs/work-unit-fairness.md new file mode 100644 index 0000000..9fe1afe --- /dev/null +++ b/docs/work-unit-fairness.md @@ -0,0 +1,29 @@ +# Work-unit capacity and fairness + +When work units are enabled, `units.capacity_share_percent` reserves a share +of eligible capacity for queued units before main and pull-request runner +demand consumes the remainder. The default is 25 percent. Deploy runner demand +keeps priority. The share admits whole units, so one unit may cross the budget +boundary. CPU, memory and lease-slot budgets round up. Existing active units count against +the share; idle capacity +can be borrowed when the other demand has no eligible work. + +```toml +[units] +enabled = true +capacity_share_percent = 25 +``` + +Set the share to zero to retain strict CI-first scheduling. Placement within the +scheduling budget still depends on owner caps, available CPU and memory, lease +limits, trust, quarantine, platform, toolchain, environment contracts and bids. A host with no eligible capacity +cannot make progress on a queued unit. + +Client selection considers units already in flight and rotates across clients. +Within a priority scope, starving requests take precedence over fresh requests; +larger starving classes get the first chance to reserve contiguous room. +The rotation cursor is transient and resets when the queen restarts. Running +work is not preempted to satisfy the share. +The existing queen starvation threshold ages queued unit demand so larger +classes can reserve contiguous room. Unit priority scopes and client maximum +scopes retain their existing authorization checks. diff --git a/src/allocator.rs b/src/allocator.rs index e80ce86..964bf7e 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -281,6 +281,76 @@ pub struct Allocation { /// Standby deploy runners to start now, per class: queued deploy demand /// that pending runners do not cover, up to the standbys that exist. pub promote: Vec, + /// Decisions captured by the same evaluation that selected each grant. + #[serde(default)] + pub decisions: Vec, + /// Extra decisions omitted to bound diagnostic memory; placement continues. + #[serde(default)] + pub omitted_decisions: u32, +} + +/// Diagnostic bounds do not limit fleet size or the number of grants. +pub const MAX_DECISIONS: usize = 128; +pub const MAX_CANDIDATES: usize = 128; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Rejection { + NotAccepting, + AdmissionClosed, + PrPressure, + Trust, + Environment, + Docker, + Memory, + Vcpu, + LeaseCap, + Reserved, + NoBid, + BidExhausted, + BidThreshold, +} + +/// Facts at one placement attempt, before its winner consumes capacity. The +/// budget is already net of the owner's caps and OS reserve; room additionally +/// subtracts active leases and earlier grants. No command or credentials enter +/// this diagnostic type. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Candidate { + pub drone_id: String, + pub budget: Room, + pub room: Room, + pub accepting: bool, + pub trust: Trust, + pub admission: AdmissionState, + pub reserved_for: Option<(Class, Scope)>, + pub rejection_reasons: Vec, + /// The queen may attach required failing checks from the same environment + /// evaluation that blocked this class. Raw observed values are excluded. + #[serde(default)] + pub environment_checks: BTreeMap, + #[serde(default)] + pub capability_checks: BTreeMap, + pub bid: Option, + pub bid_score: i64, + pub profile_score: i64, + pub score: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PlacementDecision { + pub class: Class, + pub scope: Scope, + /// Index in this allocation's grant order. Failed attempts do not advance it. + pub grant_order: usize, + pub requested: ClassSpec, + pub pressure: u16, + pub candidates: Vec, + pub total_candidates: usize, + /// A host held for this demand after an unsuccessful starving attempt. + pub reservation: Option, + /// Retained even when the winning host is beyond the candidate limit. + pub winner: Option, } /// Start `count` standby deploy runners of `class`. @@ -356,33 +426,145 @@ struct Free<'a> { granted: BTreeMap<(Class, Scope), u32>, } +/// One evaluation shared by placement diagnostics and reservation admission. +struct BidEvaluation { + bid: Option, + bid_score: i64, + profile_score: i64, + rejection_reasons: Vec, + score: Option, +} + impl Free<'_> { - /// The score this drone offers for `(class, scope)` at `pressure`, or - /// `None` if it declines. - fn offer( + fn evaluate_bid( &self, class: Class, scope: Scope, pressure: u16, weights: &ScoreWeights, - ) -> Option { - let learned = self + ) -> BidEvaluation { + let profile_score = self .drone .profiles .iter() .find(|p| p.class == class) .map_or(0, |p| profile_term(p, weights)); - if self.drone.bids.is_empty() { - return Some(learned); - } let bid = self .drone .bids .iter() - .find(|b| b.class == class && b.scope == scope)?; - let used = self.granted.get(&(class, scope)).copied().unwrap_or(0); - (used < bid.available_slots && pressure >= bid.min_pressure.min(1000)) - .then(|| score(bid, weights).saturating_add(learned)) + .find(|b| b.class == class && b.scope == scope) + .copied(); + let mut rejection_reasons = Vec::new(); + if !self.drone.bids.is_empty() { + match bid { + None => rejection_reasons.push(Rejection::NoBid), + Some(b) => { + let used = self.granted.get(&(class, scope)).copied().unwrap_or(0); + if used >= b.available_slots { + rejection_reasons.push(Rejection::BidExhausted); + } + if pressure < b.min_pressure.min(1000) { + rejection_reasons.push(Rejection::BidThreshold); + } + } + } + } + let bid_score = bid.as_ref().map_or(0, |b| score(b, weights)); + BidEvaluation { + bid, + bid_score, + profile_score, + score: rejection_reasons + .is_empty() + .then_some(bid_score.saturating_add(profile_score)), + rejection_reasons, + } + } + + /// The score this drone offers for `(class, scope)` at `pressure`, or + /// `None` if it declines. + fn offer( + &self, + class: Class, + scope: Scope, + pressure: u16, + weights: &ScoreWeights, + ) -> Option { + self.evaluate_bid(class, scope, pressure, weights).score + } + + /// Eligibility and scores are evaluated once, for both selection and its + /// explanation. Keeping this next to `offer` prevents a second scheduler. + fn candidate( + &self, + input: &AllocatorInput, + class: Class, + scope: Scope, + pressure: u16, + reserved_for: Option<(Class, Scope)>, + ) -> Candidate { + let spec = input.classes.spec(class); + let mut reasons = Vec::new(); + if !self.drone.accepting { + reasons.push(Rejection::NotAccepting); + } + match self.admission { + AdmissionState::Closed => reasons.push(Rejection::AdmissionClosed), + AdmissionState::PrBlocked if scope == Scope::Pr => reasons.push(Rejection::PrPressure), + _ => {} + } + if !self.drone.trust.admits(scope) { + reasons.push(Rejection::Trust); + } + if self.drone.blocked_classes.contains(&class) { + reasons.push(Rejection::Environment); + } + if spec.docker && !self.drone.docker { + reasons.push(Rejection::Docker); + } + if self.mem_mib < spec.mem_mib { + reasons.push(Rejection::Memory); + } + if self.vcpu < spec.vcpu { + reasons.push(Rejection::Vcpu); + } + if self.leases_left == Some(0) { + reasons.push(Rejection::LeaseCap); + } + if reserved_for.is_some_and(|key| key != (class, scope)) { + reasons.push(Rejection::Reserved); + } + let evaluated = self.evaluate_bid(class, scope, pressure, &input.weights); + reasons.extend(evaluated.rejection_reasons); + Candidate { + drone_id: self.drone.id.clone(), + budget: Room { + mem_mib: self.drone.mem_mib, + vcpu: self.drone.vcpu, + leases_left: self.drone.max_leases, + }, + room: Room { + mem_mib: self.mem_mib, + vcpu: self.vcpu, + leases_left: self.leases_left, + }, + accepting: self.drone.accepting, + trust: self.drone.trust, + admission: self.admission, + reserved_for, + score: if reasons.is_empty() { + evaluated.score + } else { + None + }, + rejection_reasons: reasons, + environment_checks: BTreeMap::new(), + capability_checks: BTreeMap::new(), + bid: evaluated.bid, + bid_score: evaluated.bid_score, + profile_score: evaluated.profile_score, + } } fn fits(&self, spec: &ClassSpec) -> bool { @@ -439,7 +621,7 @@ fn place( reserved: &BTreeMap, (scope, class): (Scope, Class), wanted: u32, - out: &mut Vec, + out: &mut Allocation, ) -> u32 { let level = pressure.get(&class).copied().unwrap_or(0); let spec = input.classes.spec(class); @@ -447,28 +629,51 @@ fn place( for _ in 0..wanted { // Highest bid score; then tightest fit, which keeps hosts packed; then // id, which keeps the output stable. - let best = free - .iter_mut() - .filter(|(id, f)| { - f.fits(spec) - && f.admission.admits(scope) - && f.drone.trust.admits(scope) - && !f.drone.blocked_classes.contains(&class) - && reserved - .get(**id) - .is_none_or(|held| *held == (class, scope)) - }) - .filter_map(|(id, f)| { - let s = f.offer(class, scope, level, &input.weights)?; - Some((s, id, f)) - }) - .min_by_key(|(s, id, f)| (std::cmp::Reverse(*s), f.mem_mib - spec.mem_mib, **id)) - .map(|(_, _, f)| f); - let Some(f) = best else { break }; + let mut candidates = Vec::new(); + let mut winning_candidate = None; + let mut best: Option<(i64, u64, &str)> = None; + for (id, f) in free.iter() { + let candidate = f.candidate(input, class, scope, level, reserved.get(*id).copied()); + if let Some(score) = candidate.score { + let rank = (std::cmp::Reverse(score), f.mem_mib - spec.mem_mib, *id); + if best.is_none_or(|(s, mem, name)| rank < (std::cmp::Reverse(s), mem, name)) { + best = Some((score, f.mem_mib - spec.mem_mib, *id)); + winning_candidate = Some(candidate.clone()); + } + } + if candidates.len() < MAX_CANDIDATES { + candidates.push(candidate); + } + } + let winner = best.map(|(_, _, id)| id); + if let Some(candidate) = winning_candidate + && !candidates.iter().any(|c| c.drone_id == candidate.drone_id) + { + candidates.pop(); + candidates.push(candidate); + } + if out.decisions.len() < MAX_DECISIONS { + out.decisions.push(PlacementDecision { + class, + scope, + grant_order: out.grants.len(), + requested: *spec, + pressure: level, + candidates, + total_candidates: free.len(), + reservation: None, + winner: winner.map(str::to_owned), + }); + } else { + out.omitted_decisions = out.omitted_decisions.saturating_add(1); + } + let Some(f) = winner.and_then(|id| free.get_mut(id)) else { + break; + }; f.take(spec); *f.granted.entry((class, scope)).or_default() += 1; placed += 1; - out.push(Grant { + out.grants.push(Grant { drone_id: f.drone.id.clone(), class, scope, @@ -482,6 +687,41 @@ fn place( /// Decide which runners to mint and where. Pure and deterministic. pub fn allocate(input: &AllocatorInput) -> Allocation { + allocate_reserved(input, &[]) +} + +/// Allocate only one scope while retaining the full queue's bid pressure. +/// Coverage by pending/shared runners and standby deploy runners is calculated +/// before restricting placement, just as in a complete allocation pass. +pub fn allocate_scope(input: &AllocatorInput, scope: Scope) -> Allocation { + allocate_inner(input, &[], None, Some(scope)) +} + +/// Apply reservations inherited from an earlier allocation pass. These are +/// enforced at the same eligibility boundary as new starvation reservations, +/// so their diagnostic reason is preserved instead of masquerading as drain. +pub fn allocate_reserved(input: &AllocatorInput, inherited: &[Reservation]) -> Allocation { + allocate_inner(input, inherited, None, None) +} + +/// Select only `input.demand`, but evaluate bid thresholds against the actual +/// uncovered queue. Per-client fairness selects one unit at a time; that must +/// not make a saturated queue look like one isolated request. Callers pass only +/// demand not already covered by a pending runner or leased work unit. +pub fn allocate_reserved_with_backlog( + input: &AllocatorInput, + inherited: &[Reservation], + uncovered_backlog: &[Demand], +) -> Allocation { + allocate_inner(input, inherited, Some(uncovered_backlog), None) +} + +fn allocate_inner( + input: &AllocatorInput, + inherited: &[Reservation], + backlog: Option<&[Demand]>, + scope: Option, +) -> Allocation { // Remaining room per drone after leases already in flight. Duplicate // drone ids collapse to the first entry. let mut free: BTreeMap<&str, Free<'_>> = BTreeMap::new(); @@ -568,7 +808,9 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { need.insert((Scope::Pr, class), pr_need); } - let initial_need = need.clone(); + let mut initial_need = need.clone(); + let pressure_backlog = + backlog.map(|demand| sum_counts(demand.iter().map(|d| (d.class, d.scope, d.count)))); let aged = |class, scope| { input.starvation_after_secs.is_some_and(|after| { input @@ -584,8 +826,13 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { let mut shortfall = 0u32; let mut starving = false; for scope in Scope::ALL { - let queued_here = get(&queued, class, scope); - shortfall += queued_here.min(initial_need.get(&(scope, class)).copied().unwrap_or(0)); + let queued_here = get(pressure_backlog.as_ref().unwrap_or(&queued), class, scope); + let uncovered = if pressure_backlog.is_some() { + queued_here + } else { + queued_here.min(initial_need.get(&(scope, class)).copied().unwrap_or(0)) + }; + shortfall = shortfall.saturating_add(uncovered); starving |= queued_here > 0 && aged(class, scope); } let idle = free @@ -599,6 +846,13 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { idle_slots: idle, }); } + if let Some(scope) = scope { + need.retain(|&(s, _), _| s == scope); + initial_need.retain(|&(s, _), _| s == scope); + if scope != Scope::Deploy { + promote.clear(); + } + } let mut out = Allocation { pressure: pressure .iter() @@ -617,7 +871,10 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { ..Allocation::default() }; let mut placed: BTreeMap<(Scope, Class), u32> = BTreeMap::new(); - let mut reserved: BTreeMap = BTreeMap::new(); + let mut reserved: BTreeMap = inherited + .iter() + .map(|r| (r.drone_id.clone(), (r.class, r.scope))) + .collect(); // Starvation guard. Demand that has waited too long is placed first; if it // still cannot fit, hold back the drone that will free up soonest so the @@ -636,7 +893,7 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { &reserved, (scope, class), wanted, - &mut out.grants, + &mut out, ); placed.insert((scope, class), n); need.insert((scope, class), 0); @@ -650,10 +907,24 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { && f.could_ever_fit(spec) && f.admission.admits(scope) && !f.drone.blocked_classes.contains(&class) + && f.offer( + class, + scope, + pressure.get(&class).copied().unwrap_or(0), + &input.weights, + ) + .is_some() }) .min_by_key(|(id, f)| (f.in_flight, **id)) .map(|(id, _)| (*id).to_owned()); if let Some(id) = target { + if let Some(decision) = out.decisions.last_mut() + && decision.class == class + && decision.scope == scope + && decision.winner.is_none() + { + decision.reservation = Some(id.clone()); + } reserved.insert(id.clone(), (class, scope)); out.reserved_for.push(Reservation { drone_id: id, @@ -676,7 +947,7 @@ pub fn allocate(input: &AllocatorInput) -> Allocation { &reserved, (scope, class), wanted, - &mut out.grants, + &mut out, ); *placed.entry((scope, class)).or_default() += n; } diff --git a/src/api.rs b/src/api.rs index a8ce5a6..e83a33b 100644 --- a/src/api.rs +++ b/src/api.rs @@ -97,6 +97,9 @@ fn sanitize_facts(f: &EnvFacts, now: i64) -> EnvFacts { }; let mut out = f.clone(); out.observed_at = out.observed_at.min(now); + for fact in out.capabilities.values_mut() { + fact.observed_at = fact.observed_at.min(now); + } out.rustc_wrapper = out.rustc_wrapper.as_ref().map(clip); out.sccache_dir = out.sccache_dir.as_ref().map(clip); if let Some(list) = &mut out.toolchains { @@ -284,6 +287,7 @@ pub fn router(queen: Arc) -> Router { .route("/status", get(status)) .route("/metrics", get(metrics)) .route("/report", get(report)) + .route("/explain", get(explain)) .route("/v1/enroll", post(tailnet_enroll)) .route("/v1/heartbeat", post(heartbeat)) .route("/v1/leases/{id}/result", post(lease_result)) @@ -291,6 +295,46 @@ pub fn router(queen: Arc) -> Router { .with_state(queen) } +#[derive(Debug, Deserialize)] +struct ExplainQuery { + key: String, +} + +/// Public label queries contain runner decisions only. Unit diagnostics require +/// the existing client token and a row owned by that client; another client's +/// unit is indistinguishable from an unknown key. +async fn explain( + State(queen): State>, + headers: HeaderMap, + Query(query): Query, +) -> Result, ApiError> { + use crate::explain::{Explanation, PlacementKind}; + if query.key.is_empty() || query.key.len() > 256 { + return Err(ApiError::BadRequest( + "explain key must be 1-256 bytes".into(), + )); + } + let (kind, key, client) = + if let Some((class, scope)) = queen.config().classes.parse_label(&query.key) { + ( + PlacementKind::Runner, + queen.config().classes.job_label(class, scope), + None, + ) + } else { + if !queen.config().units.enabled { + return Err(ApiError::NotFound); + } + let owner = crate::units_api::authenticate_client(&queen, &headers).await?; + (PlacementKind::Unit, query.key.clone(), Some(owner)) + }; + let history = queen + .store() + .placement_history(kind, &key, client.as_deref()) + .await?; + Ok(Json(Explanation::new(query.key, history))) +} + /// The drone name behind the request's bearer token. pub(crate) async fn authenticate(queen: &Queen, headers: &HeaderMap) -> Result { let presented = headers diff --git a/src/drone/capabilities.rs b/src/drone/capabilities.rs new file mode 100644 index 0000000..be451ca --- /dev/null +++ b/src/drone/capabilities.rs @@ -0,0 +1,707 @@ +//! Bounded job-user execution proofs, distinct from service-user inventory. +//! No probe installs tools, pulls an image, changes host policy or accesses the +//! network. Every child command (including cleanup) has a hard timeout and +//! discards output before it reaches the daemon's process boundary. + +use std::collections::BTreeMap; +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use serde::{Deserialize, Serialize}; +use tokio::task::JoinSet; + +use crate::drone::config::DroneConfig; +use crate::drone::exec::Exec; +use crate::envcontract::{Capability, CapabilityFact, CheckState}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct ProbeSettings { + /// GNU-compatible timeout, or `gtimeout` where coreutils uses that name. + pub timeout_bin: String, + pub timeout_secs: u64, + pub bwrap_bin: String, + /// An already-loaded image containing `/bin/true`. Never pulled by Borg. + /// Absent: Docker execution is unknown, even if `docker info` works. + pub docker_image: Option, +} + +impl Default for ProbeSettings { + fn default() -> Self { + Self { + timeout_bin: "timeout".into(), + timeout_secs: 5, + bwrap_bin: "bwrap".into(), + docker_image: None, + } + } +} + +impl ProbeSettings { + pub fn validate(&self) -> Result<(), String> { + let binary_ok = |s: &str| { + !s.is_empty() + && s.len() <= 256 + && !s.starts_with('-') + && !s.chars().any(|c| c.is_control() || c.is_whitespace()) + }; + if !binary_ok(&self.timeout_bin) || !binary_ok(&self.bwrap_bin) { + return Err( + "capability_probe binaries must be nonempty executable names or paths".into(), + ); + } + if !(1..=10).contains(&self.timeout_secs) { + return Err("capability_probe.timeout_secs must be between 1 and 10".into()); + } + if self.docker_image.as_ref().is_some_and(|s| { + s.is_empty() + || s.len() > 256 + || s.starts_with('-') + || !s.bytes().all(|b| { + b.is_ascii_alphanumeric() + || matches!(b, b'/' | b'.' | b'_' | b'-' | b':' | b'@') + }) + }) { + return Err("capability_probe.docker_image must be an image reference".into()); + } + Ok(()) + } +} + +/// Keep a trusted supervisor alive while the smoke runs. Its TERM trap kills +/// the isolated group immediately: some timeout implementations kill only the +/// direct child on their final KILL escalation. The grouped Exec boundary is +/// mandatory so signaling group zero cannot affect the drone's ambient group. +/// Redirection happens before any smoke output can reach SystemExec. +const QUIET_EXEC: &str = r#"exec >/dev/null 2>&1 +trap 'kill -KILL 0' TERM HUP INT +"$@" & +child=$! +wait "$child" +status=$? +trap - TERM HUP INT +exit "$status""#; +const TTL_SECS: i64 = 60; + +fn bounded(exec: &dyn Exec, settings: &ProbeSettings, program: &str, args: &[&str]) -> CheckState { + let seconds = format!("{}s", settings.timeout_secs); + let mut argv = vec![ + "--signal=TERM", + "--kill-after=2s", + &seconds, + "sh", + "-c", + QUIET_EXEC, + "sh", + program, + ]; + argv.extend_from_slice(args); + match exec.run_grouped(&settings.timeout_bin, &argv) { + Ok(out) if out.success => CheckState::Pass, + Ok(_) => CheckState::Fail, + Err(_) => CheckState::Unknown, + } +} + +const WRITE_PROBE: &str = r#"mkdir -m 700 -- "$1" || exit 1 +trap 'rm -rf -- "$1"' EXIT HUP INT TERM +printf borg >"$1/proof" && test "$(cat "$1/proof")" = borg"#; + +fn write_probe(exec: &dyn Exec, settings: &ProbeSettings, root: &Path, nonce: &str) -> CheckState { + let path = root.join(format!(".borg-capability-{nonce}")); + let Some(path) = path.to_str() else { + return CheckState::Unknown; + }; + let state = bounded(exec, settings, "sh", &["-c", WRITE_PROBE, "sh", path]); + // Also clean after a timeout/failed shell. Exact generated child only; + // neither work nor cache root itself can be passed to removal. + let cleanup = bounded(exec, settings, "rm", &["-rf", "--", path]); + if state == CheckState::Pass && cleanup != CheckState::Pass { + CheckState::Fail + } else { + state + } +} + +fn docker_probe(exec: &dyn Exec, cfg: &DroneConfig, nonce: &str) -> CheckState { + let Some(image) = cfg.capability_probe.docker_image.as_deref() else { + return CheckState::Unknown; + }; + let name = format!("borg-capability-{nonce}"); + let state = bounded( + exec, + &cfg.capability_probe, + &cfg.docker_bin, + &[ + "run", + "--pull=never", + "--network=none", + "--read-only", + "--cap-drop=ALL", + "--security-opt=no-new-privileges", + "--pids-limit=16", + "--memory=64m", + "--cpus=0.1", + "--rm", + "--name", + &name, + "--entrypoint=/bin/true", + image, + ], + ); + // `--rm` cleans normal completion. An interrupted client can leave the + // daemon running its container: always attempt bounded exact-name removal. + let _ = bounded( + exec, + &cfg.capability_probe, + &cfg.docker_bin, + &["rm", "-f", "-v", &name], + ); + state +} + +fn probe(exec: &dyn Exec, cfg: &DroneConfig, now: i64) -> BTreeMap { + let nonce = format!("{:016x}", rand::random::()); + let sandbox = bounded( + exec, + &cfg.capability_probe, + &cfg.capability_probe.bwrap_bin, + &[ + "--unshare-user", + "--unshare-pid", + "--unshare-net", + "--unshare-ipc", + "--unshare-uts", + "--die-with-parent", + "--new-session", + "--ro-bind", + "/", + "/", + "--proc", + "/proc", + "--dev", + "/dev", + "--", + "/bin/true", + ], + ); + [ + (Capability::Sandbox, sandbox), + (Capability::Docker, docker_probe(exec, cfg, &nonce)), + ( + Capability::WorkWrite, + write_probe(exec, &cfg.capability_probe, &cfg.work_root, &nonce), + ), + ( + Capability::CacheWrite, + write_probe(exec, &cfg.capability_probe, &cfg.cache_root, &nonce), + ), + ] + .into_iter() + .map(|(kind, state)| { + ( + kind, + CapabilityFact { + observed_at: now, + state, + }, + ) + }) + .collect() +} + +#[derive(Clone, Default)] +pub struct CapabilityCache(Arc>); + +#[derive(Default)] +struct Cached { + attempted_at: Option, + refreshing: bool, + facts: BTreeMap, +} + +/// Owns the single-flight reservation until the blocking probe and cleanup +/// finish, including when its task panics. No lock is held during commands. +struct Refresh(CapabilityCache); + +impl Refresh { + fn complete(self, mut facts: BTreeMap) { + if let Ok(mut cache) = self.0.0.lock() { + for (kind, fact) in &mut facts { + if fact.state == CheckState::Unknown + && let Some(old) = cache.facts.get(kind) + { + *fact = old.clone(); + } + } + cache.facts = facts; + } + } +} + +impl Drop for Refresh { + fn drop(&mut self) { + if let Ok(mut cache) = self.0.0.lock() { + cache.refreshing = false; + } + } +} + +impl CapabilityCache { + fn begin(&self, now: i64) -> Option { + let mut cache = self.0.lock().ok()?; + if cache.refreshing + || cache + .attempted_at + .is_some_and(|at| (0..TTL_SECS).contains(&now.saturating_sub(at))) + { + return None; + } + cache.refreshing = true; + cache.attempted_at = Some(now); + Some(Refresh(self.clone())) + } + + /// Snapshot only: a slow smoke test never delays a heartbeat. Empty or + /// aged proofs still fail required admission until a refresh completes. + pub fn snapshot(&self) -> BTreeMap { + self.0 + .lock() + .map(|cache| cache.facts.clone()) + .unwrap_or_default() + } + + /// Schedule at most one refresh, once per minute, under the drone's + /// existing task/drain lifecycle. The caller stops scheduling on drain. + pub fn refresh_if_due( + &self, + exec: Arc, + cfg: Arc, + now: i64, + tasks: &mut JoinSet<()>, + ) { + let Some(refresh) = self.begin(now) else { + return; + }; + tasks.spawn_blocking(move || { + refresh.complete(probe(exec.as_ref(), &cfg, now)); + }); + } + + /// Synchronous test driver for the same reservation/completion path. + #[cfg(test)] + fn get( + &self, + exec: &dyn Exec, + cfg: &DroneConfig, + now: i64, + ) -> BTreeMap { + if let Some(refresh) = self.begin(now) { + refresh.complete(probe(exec, cfg, now)); + } + self.snapshot() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::drone::exec::Output; + use crate::drone::jobuser::{JobExec, JobUser}; + use std::io; + + #[derive(Default)] + struct FakeExec { + calls: Mutex)>>, + fail: Mutex>, + unavailable: Mutex, + } + + impl Exec for FakeExec { + fn run_grouped(&self, program: &str, args: &[&str]) -> io::Result { + self.run(program, args) + } + + fn run(&self, program: &str, args: &[&str]) -> io::Result { + self.calls + .lock() + .unwrap() + .push((program.into(), args.iter().map(|a| (*a).into()).collect())); + if *self.unavailable.lock().unwrap() { + return Err(io::Error::new( + io::ErrorKind::NotFound, + "timeout unavailable", + )); + } + let success = !self + .fail + .lock() + .unwrap() + .as_ref() + .is_some_and(|p| args.contains(&p.as_str())); + Ok(Output { + success, + stdout: "ignored and never reported".into(), + }) + } + } + + fn config() -> DroneConfig { + DroneConfig::from_json(r#"{"queen_url":"https://queen.example.test","runner_dir":"/r","work_root":"/w","cache_root":"/c","capability_probe":{"docker_image":"busybox@sha256:abcd"}}"#).unwrap() + } + + #[test] + fn an_executor_without_group_isolation_reports_unknown_before_starting_a_smoke() { + struct Ungrouped; + impl Exec for Ungrouped { + fn run(&self, _program: &str, _args: &[&str]) -> io::Result { + panic!("unsafe supervisor must not start in an ambient group"); + } + } + assert_eq!( + bounded(&Ungrouped, &ProbeSettings::default(), "sh", &["-c", "true"]), + CheckState::Unknown + ); + } + + #[test] + fn proofs_use_timeouts_quiet_output_isolation_and_exact_cleanup() { + let fake = FakeExec::default(); + let facts = CapabilityCache::default().get(&fake, &config(), 100); + assert!( + facts + .values() + .all(|f| f.state == CheckState::Pass && f.observed_at == 100) + ); + let calls = fake.calls.lock().unwrap(); + assert_eq!(calls.len(), 7); + for (program, args) in calls.iter() { + assert_eq!(program, "timeout"); + assert_eq!( + &args[..7], + [ + "--signal=TERM", + "--kill-after=2s", + "5s", + "sh", + "-c", + QUIET_EXEC, + "sh" + ] + ); + } + let docker = calls + .iter() + .find(|(_, a)| a.iter().any(|x| x == "run")) + .unwrap(); + for argument in [ + "--pull=never", + "--network=none", + "--read-only", + "--cap-drop=ALL", + "--entrypoint=/bin/true", + ] { + assert!(docker.1.iter().any(|a| a == argument)); + } + let name = &docker.1[docker.1.iter().position(|a| a == "--name").unwrap() + 1]; + assert!(calls.iter().any(|(_, a)| a.ends_with(&[ + "docker".into(), + "rm".into(), + "-f".into(), + "-v".into(), + name.clone() + ]))); + assert!( + calls + .iter() + .any(|(_, a)| a.contains(&"--unshare-user".into()) + && a.contains(&"--unshare-net".into())) + ); + for root in ["/w/.borg-capability-", "/c/.borg-capability-"] { + let path = calls + .iter() + .flat_map(|(_, a)| a) + .find(|a| a.starts_with(root)) + .unwrap(); + assert!(calls.iter().any(|(_, a)| a.ends_with(&[ + "rm".into(), + "-rf".into(), + "--".into(), + path.clone() + ]))); + } + } + + #[test] + fn failed_or_timed_out_commands_are_not_positive_proof_and_cleanup_still_runs() { + let fake = FakeExec::default(); + *fake.fail.lock().unwrap() = Some("run".into()); + let facts = CapabilityCache::default().get(&fake, &config(), 100); + assert_eq!(facts[&Capability::Docker].state, CheckState::Fail); + assert!( + fake.calls + .lock() + .unwrap() + .iter() + .any(|(_, a)| a.contains(&"-f".into())) + ); + // The process boundary represents timeout's nonzero exit as success=false. + *fake.fail.lock().unwrap() = Some("bwrap".into()); + let facts = CapabilityCache::default().get(&fake, &config(), 100); + assert_eq!(facts[&Capability::Sandbox].state, CheckState::Fail); + } + + #[test] + fn filesystem_write_or_cleanup_failure_is_not_positive_proof() { + let fake = FakeExec::default(); + *fake.fail.lock().unwrap() = Some(WRITE_PROBE.into()); + let facts = CapabilityCache::default().get(&fake, &config(), 100); + for capability in [Capability::WorkWrite, Capability::CacheWrite] { + assert_eq!(facts[&capability].state, CheckState::Fail); + } + assert_eq!( + fake.calls + .lock() + .unwrap() + .iter() + .filter(|(_, a)| a.contains(&"-rf".into())) + .count(), + 2 + ); + *fake.fail.lock().unwrap() = Some("-rf".into()); + let facts = CapabilityCache::default().get(&fake, &config(), 100); + assert_eq!(facts[&Capability::WorkWrite].state, CheckState::Fail); + assert_eq!(facts[&Capability::CacheWrite].state, CheckState::Fail); + } + + #[test] + fn missing_docker_image_is_unknown_without_a_docker_invocation() { + let fake = FakeExec::default(); + let mut cfg = config(); + cfg.capability_probe.docker_image = None; + let facts = CapabilityCache::default().get(&fake, &cfg, 100); + assert_eq!(facts[&Capability::Docker].state, CheckState::Unknown); + assert!( + !fake + .calls + .lock() + .unwrap() + .iter() + .any(|(_, a)| a.contains(&"docker".into())) + ); + } + + #[test] + fn cached_proofs_keep_their_time_even_when_refresh_cannot_start() { + let fake = FakeExec::default(); + let cache = CapabilityCache::default(); + let cfg = config(); + let first = cache.get(&fake, &cfg, 100); + assert_eq!(cache.get(&fake, &cfg, 159), first); + assert_eq!(fake.calls.lock().unwrap().len(), 7); + *fake.unavailable.lock().unwrap() = true; + assert_eq!(cache.get(&fake, &cfg, 160), first); + let calls = fake.calls.lock().unwrap().len(); + assert_eq!(cache.get(&fake, &cfg, 161), first); + assert_eq!(fake.calls.lock().unwrap().len(), calls); + } + + #[test] + fn every_smoke_command_crosses_the_existing_job_user_boundary() { + let inner = Arc::new(FakeExec::default()); + let exec = JobExec::new( + inner.clone(), + JobUser { + name: "borg-job".into(), + home: "/home/borg-job".into(), + priv_bin: "setpriv".into(), + }, + ); + CapabilityCache::default().get(&exec, &config(), 100); + let calls = inner.calls.lock().unwrap(); + for (program, args) in calls.iter() { + assert_eq!(program, "setpriv"); + assert!(args.starts_with(&[ + "--reuid".into(), + "borg-job".into(), + "--regid".into(), + "borg-job".into() + ])); + assert!(args.contains(&"--bounding-set=-all".into())); + assert!(args.contains(&"timeout".into())); + } + } + + #[test] + fn a_refresh_keeps_single_flight_past_the_ttl_and_does_not_lock_snapshots() { + let cache = CapabilityCache::default(); + let refresh = cache.begin(100).unwrap(); + assert!(cache.snapshot().is_empty()); + assert!( + cache.begin(200).is_none(), + "slow refresh stays single-flight after the TTL" + ); + let facts = probe(&FakeExec::default(), &config(), 100); + refresh.complete(facts.clone()); + assert_eq!(cache.snapshot(), facts); + assert!( + cache.begin(200).is_some(), + "completion releases the reservation" + ); + } + + #[test] + fn a_panicking_refresh_releases_its_single_flight_reservation() { + let cache = CapabilityCache::default(); + let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _refresh = cache.begin(100).unwrap(); + panic!("fake probe panicked outside the cache lock"); + })); + assert!(panic.is_err()); + assert!(cache.snapshot().is_empty()); + assert!(cache.begin(200).is_some()); + } + + /// Acceptance for the Linux timeout path; a missing timeout binary + /// is a fixture failure rather than silently skipped subprocess evidence. + #[cfg(target_os = "linux")] + fn native_timeout_settings() -> ProbeSettings { + use crate::drone::exec::SystemExec; + + let settings = ProbeSettings { + timeout_bin: "/usr/bin/timeout".into(), + timeout_secs: 1, + ..ProbeSettings::default() + }; + let version = SystemExec + .run(&settings.timeout_bin, &["--version"]) + .expect("Linux capability acceptance requires configured /usr/bin/timeout"); + assert!(version.success, "timeout --version failed: {version:?}"); + settings + } + + #[test] + #[cfg(target_os = "linux")] + fn native_timeout_kills_term_ignoring_descendants_before_their_delayed_write() { + use crate::drone::exec::SystemExec; + use std::os::unix::fs::PermissionsExt; + use std::time::{Duration, Instant}; + + let settings = native_timeout_settings(); + let dir = tempfile::tempdir().unwrap(); + let script = dir.path().join("term-ignoring-smoke.sh"); + let started = dir.path().join("descendant-started"); + let leaked = dir.path().join("descendant-survived"); + // Even a broken timeout leaves only a finite fixture: its descendant + // exits after six seconds and the parent then exits. All writes stay + // under this exact temporary directory, without global process kills. + std::fs::write( + &script, + r#"#!/bin/sh +trap '' TERM +( + trap '' TERM + printf started > "$2" + sleep 6 + printf leaked > "$1" +) & +wait "$!" +"#, + ) + .unwrap(); + std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o700)).unwrap(); + let began = Instant::now(); + let state = bounded( + &SystemExec, + &settings, + script.to_str().unwrap(), + &[leaked.to_str().unwrap(), started.to_str().unwrap()], + ); + let elapsed = began.elapsed(); + // Wait past the descendant's original marker deadline before any + // assertion can unwind and delete its temporary paths. This catches + // killing just the shell while leaving its descendant alive. + std::thread::sleep(Duration::from_secs(7)); + assert_eq!( + state, + CheckState::Fail, + "a killed smoke must not be positive proof" + ); + assert!( + elapsed < Duration::from_secs(10), + "one-second timeout plus two-second kill fallback must terminate the finite fixture: {elapsed:?}" + ); + assert!( + started.exists(), + "fixture must have started the real descendant" + ); + assert!( + !leaked.exists(), + "timed-out descendant survived to its delayed write" + ); + } + + #[test] + #[cfg(target_os = "linux")] + fn native_filesystem_smokes_write_read_and_remove_only_their_temporary_children() { + use crate::drone::exec::SystemExec; + + let settings = native_timeout_settings(); + let dir = tempfile::tempdir().unwrap(); + for name in ["work", "cache"] { + let root = dir.path().join(name); + std::fs::create_dir(&root).unwrap(); + let keep = root.join("existing-data"); + std::fs::write(&keep, "preserve").unwrap(); + assert_eq!( + write_probe(&SystemExec, &settings, &root, "native-fixture"), + CheckState::Pass + ); + assert_eq!(std::fs::read_to_string(&keep).unwrap(), "preserve"); + assert!(!root.join(".borg-capability-native-fixture").exists()); + assert_eq!(std::fs::read_dir(&root).unwrap().count(), 1); + } + let absent = dir.path().join("absent-root"); + assert_eq!( + write_probe(&SystemExec, &settings, &absent, "native-fixture"), + CheckState::Fail + ); + assert!( + !absent.exists(), + "a smoke must not create its configured root" + ); + } + + #[test] + fn settings_reject_unbounded_or_malformed_commands() { + assert!(ProbeSettings::default().validate().is_ok()); + for bad in [ + ProbeSettings { + timeout_secs: 0, + ..ProbeSettings::default() + }, + ProbeSettings { + timeout_secs: 11, + ..ProbeSettings::default() + }, + ProbeSettings { + timeout_bin: "".into(), + ..ProbeSettings::default() + }, + ProbeSettings { + bwrap_bin: "bwrap --help".into(), + ..ProbeSettings::default() + }, + ProbeSettings { + docker_image: Some("--privileged".into()), + ..ProbeSettings::default() + }, + ProbeSettings { + docker_image: Some("busybox latest".into()), + ..ProbeSettings::default() + }, + ] { + assert!(bad.validate().is_err()); + } + } +} diff --git a/src/drone/config.rs b/src/drone/config.rs index 7b080a5..b79a766 100644 --- a/src/drone/config.rs +++ b/src/drone/config.rs @@ -213,6 +213,9 @@ pub struct DroneConfig { pub cache_root: PathBuf, #[serde(default = "default_docker")] pub docker_bin: String, + /// Bounded execution smoke tests run through the job-user boundary. + #[serde(default)] + pub capability_probe: crate::drone::capabilities::ProbeSettings, #[serde(default = "default_heartbeat")] pub heartbeat_secs: u64, /// A drone holding standby deploy runners starts them all once the queen @@ -349,6 +352,9 @@ impl DroneConfig { if let Some(user) = &self.job_user { user.validate().map_err(DroneConfigError::Invalid)?; } + self.capability_probe + .validate() + .map_err(DroneConfigError::Invalid)?; self.validate_job_settings()?; self.reclaim.validate().map_err(DroneConfigError::Invalid)?; self.validate_caches() diff --git a/src/drone/exec.rs b/src/drone/exec.rs index d718c34..80ebca5 100644 --- a/src/drone/exec.rs +++ b/src/drone/exec.rs @@ -16,6 +16,15 @@ pub trait Exec: Send + Sync { /// Run `program` with `args` to completion. `Err` means it could not be /// started; a non-zero exit is `Ok` with `success == false`. fn run(&self, program: &str, args: &[&str]) -> io::Result; + + /// Run in a private process group. Group-signaling supervisors must use + /// this boundary; executors that cannot guarantee isolation fail closed. + fn run_grouped(&self, _program: &str, _args: &[&str]) -> io::Result { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "executor cannot guarantee a private process group", + )) + } } /// Runs commands on the local machine. Programs and arguments are chosen by @@ -23,20 +32,84 @@ pub trait Exec: Send + Sync { /// from the queen. pub struct SystemExec; -impl Exec for SystemExec { - fn run(&self, program: &str, args: &[&str]) -> io::Result { +impl SystemExec { + fn output(&self, program: &str, args: &[&str], grouped: bool) -> io::Result { + #[cfg(not(unix))] + if grouped { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "private process groups require Unix", + )); + } #[allow( clippy::disallowed_methods, reason = "drone host commands: fixed programs and arguments built by borg, never queen-supplied" )] - let out = Command::new(program) + let mut command = Command::new(program); + command .args(args) .stdin(Stdio::null()) - .stderr(Stdio::null()) - .output()?; + .stderr(Stdio::null()); + #[cfg(unix)] + if grouped { + use std::os::unix::process::CommandExt; + command.process_group(0); + } + let out = command.output()?; Ok(Output { success: out.status.success(), stdout: String::from_utf8_lossy(&out.stdout).trim().to_owned(), }) } } + +impl Exec for SystemExec { + fn run(&self, program: &str, args: &[&str]) -> io::Result { + self.output(program, args, false) + } + + fn run_grouped(&self, program: &str, args: &[&str]) -> io::Result { + self.output(program, args, true) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn an_executor_must_explicitly_opt_in_to_grouped_execution() { + struct Ungrouped; + impl Exec for Ungrouped { + fn run(&self, _program: &str, _args: &[&str]) -> io::Result { + panic!("unsupported grouped execution must not run in an ambient group"); + } + } + assert_eq!( + Ungrouped + .run_grouped("sh", &["-c", "true"]) + .unwrap_err() + .kind(), + io::ErrorKind::Unsupported + ); + } + + #[test] + #[cfg(unix)] + fn system_grouped_execution_creates_a_group_led_by_its_child() { + let output = SystemExec + .run_grouped("sh", &["-c", "printf '%s ' \"$$\"; ps -o pgid= -p \"$$\""]) + .unwrap(); + assert!(output.success); + let ids: Vec = output + .stdout + .split_whitespace() + .map(|id| id.parse().unwrap()) + .collect(); + assert_eq!(ids.len(), 2); + assert_eq!( + ids[0], ids[1], + "child PID must equal its private process group ID" + ); + } +} diff --git a/src/drone/install.rs b/src/drone/install.rs index ca53265..f4280b5 100644 --- a/src/drone/install.rs +++ b/src/drone/install.rs @@ -794,6 +794,7 @@ impl Installer<'_> { work_root: self.layout.work(), cache_root: self.layout.cache_root.clone(), docker_bin: "docker".to_owned(), + capability_probe: Default::default(), heartbeat_secs: 10, fail_open_after_secs: o.fail_open_after_secs, max_job_secs: 6 * 3600, diff --git a/src/drone/jobuser.rs b/src/drone/jobuser.rs index b848276..ab4eff3 100644 --- a/src/drone/jobuser.rs +++ b/src/drone/jobuser.rs @@ -95,6 +95,14 @@ impl Exec for JobExec { let argv: Vec<&str> = argv.iter().map(String::as_str).collect(); self.inner.run(&self.user.priv_bin, &argv) } + + fn run_grouped(&self, program: &str, args: &[&str]) -> io::Result { + let mut argv = self.user.setpriv_args(); + argv.push(program.to_owned()); + argv.extend(args.iter().map(|a| (*a).to_owned())); + let argv: Vec<&str> = argv.iter().map(String::as_str).collect(); + self.inner.run_grouped(&self.user.priv_bin, &argv) + } } #[cfg(test)] @@ -123,6 +131,42 @@ mod tests { ); } + #[test] + fn grouped_commands_preserve_both_isolation_and_the_job_user_boundary() { + struct GroupOnly; + impl Exec for GroupOnly { + fn run(&self, _program: &str, _args: &[&str]) -> io::Result { + panic!("JobExec must forward the grouped boundary"); + } + fn run_grouped(&self, program: &str, args: &[&str]) -> io::Result { + assert_eq!(program, "setpriv"); + assert_eq!( + args, + [ + "--reuid", + "borg-job", + "--regid", + "borg-job", + "--init-groups", + "--inh-caps=-all", + "--ambient-caps=-all", + "--bounding-set=-all", + "--", + "timeout", + "1s", + "sh" + ] + ); + Ok(Output { + success: true, + stdout: String::new(), + }) + } + } + let exec = JobExec::new(Arc::new(GroupOnly), user()); + assert!(exec.run_grouped("timeout", &["1s", "sh"]).unwrap().success); + } + #[test] fn validation_rejects_bad_names_and_relative_homes() { assert!(user().validate().is_ok()); diff --git a/src/drone/mod.rs b/src/drone/mod.rs index fd345f5..7f0537a 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -6,6 +6,7 @@ //! stops accepting leases and lets in-flight jobs finish, up to a timeout. pub mod caches; +pub mod capabilities; pub mod cleanup; pub mod client; pub mod config; @@ -92,6 +93,7 @@ pub struct Drone { update_wanted: AtomicBool, /// The last tools probe, reused for a minute. tools_cache: probe::ToolsCache, + capability_cache: capabilities::CapabilityCache, } impl Drone { @@ -131,6 +133,7 @@ impl Drone { selfupdate, update_wanted: AtomicBool::new(false), tools_cache: probe::ToolsCache::default(), + capability_cache: capabilities::CapabilityCache::default(), client: Arc::new(client), host, exec, @@ -220,6 +223,18 @@ impl Drone { .await .ok() .flatten(); + // Slow execution proofs must not make an otherwise healthy drone + // miss its heartbeat and cause the queen to requeue active work. + // Refreshes share the existing task set, and drain starts none. + if accepting { + self.capability_cache.refresh_if_due( + Arc::clone(&self.job_exec), + Arc::clone(&self.cfg), + now, + tasks, + ); + } + env_facts.capabilities = self.capability_cache.snapshot(); let req = HeartbeatRequest { capacity, caps: self.cfg.host_caps(), @@ -470,7 +485,7 @@ impl Drone { } tracing::info!( - "draining: no new leases; waiting up to {}s for {} job(s)", + "draining: no new leases; waiting up to {}s for {} task(s)", self.cfg.drain_timeout_secs, tasks.len() ); diff --git a/src/envcontract.rs b/src/envcontract.rs index 2ea4612..650c6a3 100644 --- a/src/envcontract.rs +++ b/src/envcontract.rs @@ -1,7 +1,7 @@ //! The build environment contract: facts a drone reports about its build //! environment, and the queen's pure admission check over them. //! -//! Before work (a CI runner lease or a work unit) is placed on a drone, five +//! Before work (a CI runner lease or a work unit) is placed on a drone, six //! checks are evaluated against what the drone last reported: //! //! 1. **compiler_cache**: `RUSTC_WRAPPER` resolves to an allowed wrapper @@ -18,6 +18,9 @@ //! 5. **tools**: every tool the operator lists as required (globally or for the //! class) resolved on the job user's `PATH` when the drone probed it. The //! list is empty unless the operator writes one. +//! 6. **capabilities**: actual sandbox, Docker and filesystem smoke tests run +//! as the job user. Named requirements fail closed even on older drones; +//! the list is empty unless the operator writes one. //! //! The drone is a dumb prober: it reports *facts* ([`EnvFacts`]), every one of //! which may be unknown (`None`), cached about a minute on its side. Thresholds @@ -29,7 +32,8 @@ //! //! Nothing is enforced unless `env_contract.enabled`; a drone that has never //! reported is skipped unless `require_report` is set, so enabling the contract -//! does not cut off drones that predate it. +//! does not cut off drones that predate it. Explicit required capability +//! names are the exception: a missing job-user execution proof blocks them. use std::collections::BTreeMap; @@ -45,6 +49,7 @@ pub enum CheckKind { CacheStore, Resources, Tools, + Capabilities, } /// The tools a drone looks for on the job user's `PATH`, in report order. @@ -57,12 +62,13 @@ pub const PROBE_TOOLS: [&str; 15] = [ pub const TOOL_MISSING: &str = "missing"; impl CheckKind { - pub const ALL: [Self; 5] = [ + pub const ALL: [Self; 6] = [ Self::CompilerCache, Self::Toolchain, Self::CacheStore, Self::Resources, Self::Tools, + Self::Capabilities, ]; pub fn as_str(self) -> &'static str { @@ -72,6 +78,7 @@ impl CheckKind { Self::CacheStore => "cache_store", Self::Resources => "resources", Self::Tools => "tools", + Self::Capabilities => "capabilities", } } } @@ -106,6 +113,42 @@ impl CheckState { pub const ALL: [Self; 3] = [Self::Pass, Self::Fail, Self::Unknown]; } +/// Execution capabilities proven by a smoke test as the configured job user. +/// Inventory (for example `docker info` as the service user) is not proof. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Capability { + Sandbox, + Docker, + WorkWrite, + CacheWrite, +} + +impl Capability { + pub const ALL: [Self; 4] = [ + Self::Sandbox, + Self::Docker, + Self::WorkWrite, + Self::CacheWrite, + ]; + + pub fn as_str(self) -> &'static str { + match self { + Self::Sandbox => "sandbox", + Self::Docker => "docker", + Self::WorkWrite => "work_write", + Self::CacheWrite => "cache_write", + } + } +} + +/// A capability's own observation time survives cached heartbeat reports. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct CapabilityFact { + pub observed_at: i64, + pub state: CheckState, +} + /// An installed Rust toolchain, as `rustc -vV` reports it for the channel. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Toolchain { @@ -161,6 +204,9 @@ pub struct EnvFacts { /// output, or `"missing"`; `None` when the drone did not measure it. #[serde(default)] pub tools: Option>, + /// Actual execution proofs as the job user, each with its probe timestamp. + #[serde(default)] + pub capabilities: BTreeMap, } /// Settings of one class (or the defaults for all). @@ -173,6 +219,8 @@ pub struct ClassContract { pub rust_toolchain: Option, /// Tools this class needs on `PATH`; replaces the global list. pub required_tools: Option>, + /// Execution proofs this class needs; replaces the global list. + pub required_capabilities: Option>, } #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] @@ -201,6 +249,8 @@ pub struct ContractSettings { /// Tools (from [`PROBE_TOOLS`]) every class needs on `PATH`. Empty means /// the `tools` check has nothing to require. pub required_tools: Vec, + /// Execution proofs every class needs. Empty preserves older drones. + pub required_capabilities: Vec, pub checks: BTreeMap, pub classes: BTreeMap, } @@ -219,6 +269,7 @@ impl Default for ContractSettings { max_cpu_psi_pm: 700, mem_available_pct: 100, required_tools: Vec::new(), + required_capabilities: Vec::new(), // Conservative: only the checks that protect a build from failing // outright block by default; the cache checks report. checks: BTreeMap::from([ @@ -227,6 +278,7 @@ impl Default for ContractSettings { (CheckKind::CacheStore, Mode::Advisory), (CheckKind::Resources, Mode::Required), (CheckKind::Tools, Mode::Advisory), + (CheckKind::Capabilities, Mode::Required), ]), classes: BTreeMap::new(), } @@ -291,13 +343,25 @@ impl ContractSettings { .unwrap_or(&self.required_tools) } + /// Execution proofs a class needs: its own list, else the global one. + pub fn class_required_capabilities(&self, class: Class) -> &[Capability] { + self.classes + .get(&class) + .and_then(|c| c.required_capabilities.as_deref()) + .unwrap_or(&self.required_capabilities) + } + fn mode(&self, class: Class, kind: CheckKind) -> Mode { self.classes .get(&class) .and_then(|c| c.checks.get(&kind)) .or_else(|| self.checks.get(&kind)) .copied() - .unwrap_or(Mode::Off) + .unwrap_or(if kind == CheckKind::Capabilities { + Mode::Required + } else { + Mode::Off + }) } /// The toolchain a class needs by configuration. @@ -602,6 +666,87 @@ fn tools(s: &ContractSettings, f: &EnvFacts, need: &Need<'_>) -> Acc { a } +fn capability_fresh(s: &ContractSettings, fact: &CapabilityFact, now: i64) -> bool { + fact.observed_at <= now && now.saturating_sub(fact.observed_at) <= s.facts_max_age_secs +} + +fn capability_state(s: &ContractSettings, fact: Option<&CapabilityFact>, now: i64) -> CheckState { + match fact { + Some(fact) if capability_fresh(s, fact, now) => fact.state, + _ => CheckState::Unknown, + } +} + +/// Effective states of the named proofs this class needs, using the same +/// per-proof freshness rules as admission. Check mode is deliberately left +/// to the caller: diagnostics may also show advisory requirements. +pub fn capability_states( + s: &ContractSettings, + facts: Option<&EnvFacts>, + class: Class, + now: i64, +) -> BTreeMap { + s.class_required_capabilities(class) + .iter() + .map(|&kind| { + let fact = facts.and_then(|f| f.capabilities.get(&kind)); + (kind, capability_state(s, fact, now)) + }) + .collect() +} + +fn capabilities(s: &ContractSettings, f: &EnvFacts, need: &Need<'_>, now: i64) -> Acc { + let mut a = Acc::new(); + let required = s.class_required_capabilities(need.class); + a.see( + "required", + required + .iter() + .map(|c| c.as_str()) + .collect::>() + .join(","), + ); + for &capability in required { + let name = capability.as_str(); + match f.capabilities.get(&capability) { + None => a.note( + CheckState::Unknown, + format!("{name}: job-user proof was not reported"), + ), + Some(fact) => { + let age = now.saturating_sub(fact.observed_at); + a.see(&format!("{name}_age_secs"), age); + a.see(name, fact.state.as_str()); + if !capability_fresh(s, fact, now) { + a.note( + CheckState::Unknown, + format!("{name}: proof is stale or future-dated ({age}s old)"), + ); + } else { + a.note( + capability_state(s, Some(fact), now), + format!("{name}: job-user smoke test {}", fact.state.as_str()), + ); + } + } + } + } + if s.mode(need.class, CheckKind::Capabilities) == Mode::Advisory { + a.warnings.extend(a.reasons.iter().cloned()); + } + // Unrequired failures remain diagnostic and never exclude a class. + for (capability, fact) in &f.capabilities { + if !required.contains(capability) && fact.state != CheckState::Pass { + a.warnings.push(format!( + "{}: unrequired job-user smoke test {}", + capability.as_str(), + fact.state.as_str() + )); + } + } + a +} + /// Evaluate the contract for `need` against `facts` at `now`. Pure. pub fn evaluate( s: &ContractSettings, @@ -624,19 +769,32 @@ pub fn evaluate( .filter(|(_, m)| *m != Mode::Off) .collect(); let Some(f) = facts else { - let blocked = s.require_report && modes.iter().any(|(_, m)| *m == Mode::Required); + let blocked = modes.iter().any(|&(kind, mode)| { + mode == Mode::Required + && if kind == CheckKind::Capabilities { + !s.class_required_capabilities(need.class).is_empty() + } else { + s.require_report + } + }); return Verdict { reported: false, age_secs: None, checks: modes .iter() - .map(|&(kind, mode)| CheckOutcome { - kind, - mode, - state: CheckState::Unknown, - reason: "the drone has not reported its environment".into(), - observed: BTreeMap::new(), - warnings: Vec::new(), + .map(|&(kind, mode)| { + if kind == CheckKind::Capabilities { + capabilities(s, &EnvFacts::default(), need, now).finish(kind, mode) + } else { + CheckOutcome { + kind, + mode, + state: CheckState::Unknown, + reason: "the drone has not reported its environment".into(), + observed: BTreeMap::new(), + warnings: Vec::new(), + } + } }) .collect(), admit: !blocked, @@ -647,7 +805,11 @@ pub fn evaluate( let stale = age > s.facts_max_age_secs; let mut checks = Vec::new(); for (kind, mode) in modes { - let acc = if stale { + let acc = if kind == CheckKind::Capabilities { + // Each execution proof has its own clock. An empty requirement + // must not become blocking just because unrelated facts aged. + capabilities(s, f, need, now) + } else if stale { let mut a = Acc::new(); a.note( CheckState::Unknown, @@ -661,6 +823,7 @@ pub fn evaluate( CheckKind::CacheStore => cache_store(s, f), CheckKind::Resources => resources(s, f, need), CheckKind::Tools => tools(s, f, need), + CheckKind::Capabilities => capabilities(s, f, need, now), } }; checks.push(acc.finish(kind, mode)); @@ -734,6 +897,7 @@ mod tests { ("git", "git version 2.45.0"), ("gh", "gh version 2.50.0"), ])), + capabilities: BTreeMap::new(), } } @@ -1062,7 +1226,13 @@ mod tests { }; let v = evaluate(&s, Some(&old), &need(Class::Light, None), NOW); assert!(!v.admit); - assert!(v.checks.iter().all(|c| c.state == CheckState::Unknown)); + assert!( + v.checks + .iter() + .filter(|c| c.kind != CheckKind::Capabilities) + .all(|c| c.state == CheckState::Unknown) + ); + assert_eq!(check(&v, CheckKind::Capabilities).state, CheckState::Pass); assert!(v.checks[0].reason.contains("1000s old")); // A drone that never reported: admitted by default, blocked on request. let v = evaluate(&s, None, &need(Class::Light, None), NOW); @@ -1280,4 +1450,173 @@ mod tests { ); assert!(s.validate().is_err()); } + + fn capability_settings(mode: Mode) -> ContractSettings { + ContractSettings { + checks: BTreeMap::from([(CheckKind::Capabilities, mode)]), + required_capabilities: vec![Capability::Sandbox], + ..on() + } + } + + fn capability_facts(state: CheckState, observed_at: i64) -> EnvFacts { + EnvFacts { + capabilities: BTreeMap::from([( + Capability::Sandbox, + CapabilityFact { observed_at, state }, + )]), + ..good() + } + } + + #[test] + fn required_job_user_capabilities_fail_closed_including_old_drones() { + let settings = capability_settings(Mode::Required); + let need = need(Class::Light, None); + assert!( + evaluate( + &settings, + Some(&capability_facts(CheckState::Pass, NOW)), + &need, + NOW + ) + .admit + ); + for state in [CheckState::Fail, CheckState::Unknown] { + let v = evaluate(&settings, Some(&capability_facts(state, NOW)), &need, NOW); + assert!(!v.admit); + assert_eq!(check(&v, CheckKind::Capabilities).state, state); + assert!(v.why.unwrap().contains("sandbox")); + } + assert!(!evaluate(&settings, Some(&good()), &need, NOW).admit); + // require_report=false cannot skip an explicitly required named proof. + assert!(!evaluate(&settings, None, &need, NOW).admit); + let old: EnvFacts = serde_json::from_str(r#"{"observed_at":10000}"#).unwrap(); + assert!(old.capabilities.is_empty()); + assert!(!evaluate(&settings, Some(&old), &need, NOW).admit); + } + + #[test] + fn a_fresh_heartbeat_cannot_refresh_an_old_or_future_capability_proof() { + let settings = capability_settings(Mode::Required); + for at in [NOW - settings.facts_max_age_secs - 1, NOW + 1] { + let v = evaluate( + &settings, + Some(&capability_facts(CheckState::Pass, at)), + &need(Class::Light, None), + NOW, + ); + assert_eq!(v.age_secs, Some(10)); + assert!(!v.admit); + assert_eq!( + check(&v, CheckKind::Capabilities).state, + CheckState::Unknown + ); + } + } + + #[test] + fn capability_requirements_and_advisory_modes_are_class_specific() { + let mut settings = capability_settings(Mode::Required); + settings.required_capabilities.clear(); + settings.classes.insert( + Class::Heavy, + ClassContract { + required_capabilities: Some(vec![Capability::Sandbox]), + ..ClassContract::default() + }, + ); + let facts = capability_facts(CheckState::Fail, NOW); + assert!(evaluate(&settings, Some(&facts), &need(Class::Light, None), NOW).admit); + assert!(!evaluate(&settings, Some(&facts), &need(Class::Heavy, None), NOW).admit); + let v = evaluate(&settings, Some(&facts), &need(Class::Light, None), NOW); + assert!(check(&v, CheckKind::Capabilities).warnings[0].contains("sandbox")); + settings + .classes + .get_mut(&Class::Heavy) + .unwrap() + .checks + .insert(CheckKind::Capabilities, Mode::Advisory); + let v = evaluate(&settings, Some(&facts), &need(Class::Heavy, None), NOW); + assert!(v.admit); + assert_eq!(check(&v, CheckKind::Capabilities).state, CheckState::Fail); + assert!(check(&v, CheckKind::Capabilities).warnings[0].contains("sandbox")); + assert!(evaluate(&settings, None, &need(Class::Light, None), NOW).admit); + } + + #[test] + fn capability_names_are_closed_and_the_class_list_replaces_global_requirements() { + let json = r#"{"enabled":true,"required_capabilities":["sandbox"],"classes":{"heavy":{"required_capabilities":["docker"]},"medium":{"required_capabilities":[]}}}"#; + let settings: ContractSettings = serde_json::from_str(json).unwrap(); + assert!(settings.validate().is_ok()); + assert_eq!( + settings.class_required_capabilities(Class::Light), + [Capability::Sandbox] + ); + assert_eq!( + settings.class_required_capabilities(Class::Heavy), + [Capability::Docker] + ); + assert!( + settings + .class_required_capabilities(Class::Medium) + .is_empty() + ); + assert!( + serde_json::from_str::(&json.replace("sandbox", "arbitrary_command")) + .is_err() + ); + // No requirements is backward compatible, even with the new check required. + assert!(evaluate(&on(), Some(&good()), &need(Class::Light, None), NOW).admit); + let empty = ContractSettings { + checks: BTreeMap::from([(CheckKind::Capabilities, Mode::Required)]), + required_capabilities: Vec::new(), + require_report: true, + ..on() + }; + assert!(evaluate(&empty, None, &need(Class::Light, None), NOW).admit); + let stale = EnvFacts { + observed_at: NOW - 3600, + ..good() + }; + assert!(evaluate(&empty, Some(&stale), &need(Class::Light, None), NOW).admit); + } + + #[test] + fn capability_diagnostics_use_admissions_effective_proof_states() { + let settings = capability_settings(Mode::Required); + for (state, at, expected) in [ + (CheckState::Pass, NOW, CheckState::Pass), + (CheckState::Fail, NOW, CheckState::Fail), + (CheckState::Unknown, NOW, CheckState::Unknown), + (CheckState::Pass, NOW - 3600, CheckState::Unknown), + (CheckState::Fail, NOW + 1, CheckState::Unknown), + ] { + let facts = capability_facts(state, at); + let states = capability_states(&settings, Some(&facts), Class::Light, NOW); + let verdict = evaluate(&settings, Some(&facts), &need(Class::Light, None), NOW); + assert_eq!(states[&Capability::Sandbox], expected); + assert_eq!( + states[&Capability::Sandbox], + check(&verdict, CheckKind::Capabilities).state + ); + } + assert_eq!( + capability_states(&settings, None, Class::Light, NOW)[&Capability::Sandbox], + CheckState::Unknown + ); + assert_eq!( + capability_states(&settings, Some(&good()), Class::Light, NOW)[&Capability::Sandbox], + CheckState::Unknown + ); + let mut settings = settings; + settings.classes.insert( + Class::Heavy, + ClassContract { + required_capabilities: Some(Vec::new()), + ..ClassContract::default() + }, + ); + assert!(capability_states(&settings, None, Class::Heavy, NOW).is_empty()); + } } diff --git a/src/explain.rs b/src/explain.rs new file mode 100644 index 0000000..0261c1b --- /dev/null +++ b/src/explain.rs @@ -0,0 +1,163 @@ +//! Placement history captured by the allocator, rather than reconstructed from +//! today's fleet. Capacity, bids and reservations change after each grant, so +//! `/status` cannot explain an earlier choice. The existing queen SQLite store +//! owns this bounded diagnostic history; it contains no argv, env or tokens. + +use std::fmt::Write as _; + +use serde::{Deserialize, Serialize}; + +use crate::allocator::PlacementDecision; +use crate::secret::Secret; + +/// Both the row count and per-row size are bounded. Old decisions are deleted +/// transactionally on insert; dropping this table loses diagnostics only. +pub const MAX_HISTORY: i64 = 1000; +pub const MAX_DECISION_BYTES: usize = 128 * 1024; +pub const MAX_QUERY_ROWS: i64 = 20; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PlacementKind { + Runner, + Unit, +} + +impl PlacementKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Runner => "runner", + Self::Unit => "unit", + } + } +} + +/// Allocator selection is not proof that minting or the persisted unit lease +/// succeeded. Record the result only after the owning operation finishes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum PlacementOutcome { + Granted, + Unplaced, + Failed, +} + +pub struct PlacementRecord<'a> { + pub at: i64, + pub kind: PlacementKind, + pub key: &'a str, + pub client: Option<&'a str>, + pub outcome: PlacementOutcome, + pub decision: &'a PlacementDecision, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct HistoryEntry { + pub id: i64, + pub at: i64, + pub kind: PlacementKind, + pub key: String, + pub outcome: PlacementOutcome, + pub decision: PlacementDecision, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Explanation { + pub key: String, + /// Empty explicitly means no retained decision, not successful placement. + pub history: Vec, + pub retention_limit: i64, +} + +impl Explanation { + pub fn new(key: impl Into, history: Vec) -> Self { + Self { + key: key.into(), + history, + retention_limit: MAX_HISTORY, + } + } +} + +pub fn render_text(report: &Explanation) -> String { + let mut out = String::new(); + if report.history.is_empty() { + let _ = writeln!(out, "{}: no retained placement decision", report.key); + return out; + } + for row in &report.history { + let d = &row.decision; + let _ = writeln!( + out, + "{} {} at {}: {:?}, selected {} (grant order {})", + row.kind.as_str(), + row.key, + row.at, + row.outcome, + d.winner.as_deref().unwrap_or("none"), + d.grant_order + ); + let _ = writeln!( + out, + " {} {} needs {} vCPU / {} MiB; pressure {}", + d.class, + d.scope.as_str(), + d.requested.vcpu, + d.requested.mem_mib, + d.pressure + ); + for host in &d.candidates { + let _ = writeln!( + out, + " {}: {} / {} MiB, {} / {} vCPU free; lease slots {:?}; score {:?} (bid {}, profile {}); {:?}; reservation {:?}", + host.drone_id, + host.room.mem_mib, + host.budget.mem_mib, + host.room.vcpu, + host.budget.vcpu, + host.room.leases_left, + host.score, + host.bid_score, + host.profile_score, + host.rejection_reasons, + host.reserved_for + ); + if !host.environment_checks.is_empty() { + let _ = writeln!(out, " environment checks: {:?}", host.environment_checks); + } + if !host.capability_checks.is_empty() { + let _ = writeln!(out, " capability checks: {:?}", host.capability_checks); + } + } + if d.candidates.len() < d.total_candidates { + let _ = writeln!( + out, + " showing {} of {} considered hosts", + d.candidates.len(), + d.total_candidates + ); + } + if let Some(host) = &d.reservation { + let _ = writeln!(out, " reserved {host} for this demand to free enough room"); + } + } + out +} + +/// Fetch the same typed history the CLI reads locally. Unit keys require the +/// existing client bearer token; runner labels expose runner metadata only. +pub async fn fetch( + queen: &str, + key: &str, + token: Option<&Secret>, +) -> Result { + let url = format!("{}/explain", queen.trim_end_matches('/')); + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(15)) + .build()?; + let mut request = client.get(url).query(&[("key", key)]); + if let Some(token) = token { + request = request.bearer_auth(token.expose()); + } + request.send().await?.error_for_status()?.json().await +} diff --git a/src/lib.rs b/src/lib.rs index 8be6a35..f54d334 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,6 +18,7 @@ pub mod configfile; pub mod drone; pub mod enroll; pub mod envcontract; +pub mod explain; pub mod github; pub mod health; pub mod homeostasis; diff --git a/src/main.rs b/src/main.rs index dc9be4f..3f1d643 100644 --- a/src/main.rs +++ b/src/main.rs @@ -49,6 +49,8 @@ enum Command { /// Summarize stored job telemetry: time by job and step, queue wait per /// class, cancellation share and host-caused failures per drone. Report(ReportArgs), + /// Explain a retained placement decision for a runner label or work unit. + Explain(ExplainArgs), /// Run a command on the fleet when capacity is free, otherwise here. /// Everything after `--` is the command. Run(RunArgs), @@ -159,7 +161,13 @@ struct RunArgs { /// Priority scope of the unit; the client token's maximum applies. #[arg(long, env = "BORG_RUN_SCOPE", value_enum, default_value = "pr")] scope: ScopeArg, - /// Seconds to wait for a drone before running here. + /// Require fleet execution; infrastructure failures exit 125 without running here. + #[arg(long, env = "BORG_RUN_REMOTE_ONLY")] + remote_only: bool, + /// Write a JSON execution result to this file, separately from command stdout. + #[arg(long)] + result_file: Option, + /// Seconds to wait for a drone before fallback (or failure with --remote-only). #[arg(long, env = "BORG_RUN_WAIT_SECS", default_value_t = 60)] wait_secs: u64, /// Upload a git bundle of the tree (since the merge base with origin/main) @@ -191,6 +199,24 @@ struct ReportArgs { json: bool, } +#[derive(Args)] +struct ExplainArgs { + /// Work-unit id or configured runner label. + key: String, + /// Fetch from a running queen. + #[arg(long, conflicts_with = "config", required_unless_present = "config")] + queen: Option, + /// Read the queen's SQLite file directly as a local operator. + #[arg(long)] + config: Option, + /// Client token file for a remote unit query (same defaults as `borg run`). + #[arg(long, env = "BORG_CLIENT_TOKEN_FILE")] + token_file: Option, + /// Print the typed JSON history instead of text. + #[arg(long)] + json: bool, +} + #[derive(Args)] struct ConfigArg { /// Path to the queen's config file (TOML or JSON). @@ -588,6 +614,36 @@ async fn report(a: ReportArgs) -> Result<()> { Ok(()) } +async fn explain(a: ExplainArgs) -> Result<()> { + use borg::explain::{Explanation, PlacementKind}; + let report = match (&a.queen, &a.config) { + (Some(url), _) => { + let token = client_token(a.token_file.as_deref()).map(borg::secret::Secret::new); + borg::explain::fetch(url, &a.key, token.as_ref()) + .await + .context("fetch placement history")? + } + (None, Some(path)) => { + let cfg = QueenConfig::load(path)?; + let store = open_store(&cfg).await?; + let (kind, key) = match cfg.classes.parse_label(&a.key) { + Some((class, scope)) => { + (PlacementKind::Runner, cfg.classes.job_label(class, scope)) + } + None => (PlacementKind::Unit, a.key.clone()), + }; + Explanation::new(a.key, store.placement_history(kind, &key, None).await?) + } + (None, None) => bail!("give --queen or --config"), + }; + if a.json { + println!("{}", serde_json::to_string_pretty(&report)?); + } else { + print!("{}", borg::explain::render_text(&report)); + } + Ok(()) +} + async fn enroll_client(path: &Path, name: &str, max_scope: borg::classes::Scope) -> Result<()> { if !enroll::valid_drone_name(name) { bail!( @@ -793,6 +849,8 @@ async fn run_command(a: RunArgs) -> u8 { scope: borg::classes::Scope::from(a.scope).as_str().to_owned(), wait_secs: a.wait_secs, publish: a.publish, + remote_only: a.remote_only, + result_file: a.result_file, argv: a.argv, }; let code = borg::run::run(&opts, &mut std::io::stdout()).await; @@ -811,6 +869,7 @@ async fn main() -> Result { async fn dispatch(command: Command) -> Result<()> { match command { + Command::Explain(a) => explain(a).await, Command::Run(_) => Ok(()), Command::Update(a) => { run_update(update::cli::Args { diff --git a/src/probe/mod.rs b/src/probe/mod.rs index 9975e43..33d2bad 100644 --- a/src/probe/mod.rs +++ b/src/probe/mod.rs @@ -50,7 +50,8 @@ pub struct Capacity { pub mem_mib: u64, /// Free space on the filesystem holding the work directory. pub free_disk_mib: Option, - /// `docker info` reached a daemon. + /// `docker info` reached a daemon as the service user. This is inventory; + /// job-user execution proof lives in `EnvFacts::capabilities`. pub docker_usable: bool, pub mise_version: Option, pub git_version: Option, diff --git a/src/queen.rs b/src/queen.rs index cbe5b8a..7cfee02 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -12,7 +12,7 @@ //! on a partial view of the queue would over- or under-grant. use std::borrow::Cow; -use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque}; use std::sync::Arc; use std::sync::RwLock; use std::sync::atomic::{AtomicU64, Ordering}; @@ -24,13 +24,16 @@ use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; use crate::allocator::{ - ActiveLease, AllocatorInput, ClassIdle, Demand, DroneAdmission, DroneState, Grant, LeaseState, - Promotion, QueueAge, Reservation, Unmet, WarmFloor, allocate, free_room, free_slots, max_slots, + ActiveLease, AdmissionState, AllocatorInput, ClassIdle, Demand, DroneAdmission, DroneState, + LeaseState, PlacementDecision, Promotion, QueueAge, Reservation, Unmet, WarmFloor, + allocate_reserved, allocate_reserved_with_backlog, allocate_scope, free_room, free_slots, + max_slots, }; use crate::bid::{ClassPressure, ScoreWeights}; use crate::classes::{Class, Scope}; use crate::config::{QueenConfig, StandbySettings}; -use crate::envcontract::blocked_classes; +use crate::envcontract::{CheckKind, CheckState, Mode, Need, blocked_classes, capability_states}; +use crate::explain::{MAX_HISTORY, PlacementKind, PlacementOutcome, PlacementRecord}; use crate::github::{GithubClient, GithubError, QueuedJob, Runner, WorkflowJob}; use crate::health::{DroneHealth, DroneInfo, Outcome, OutcomeKind, evaluate}; use crate::homeostasis::{Controller, KeyInput, KeyReport}; @@ -44,7 +47,7 @@ use crate::rollout::{ Check as RolloutCheck, DroneView, Observed as RolloutObserved, RolloutState, advance, desired_for, }; -use crate::store::{LeaseRow, LeaseStatus, PreemptionRow, Store, StoreError}; +use crate::store::{DroneRow, LeaseRow, LeaseStatus, PreemptionRow, Store, StoreError}; use crate::tailnet::TailnetEnroll; use crate::telemetry::{Conclusion, FailureClass, MAX_REPORT_ROWS, TelemetryRow}; use crate::units::{Unit, UnitState}; @@ -160,6 +163,36 @@ pub struct QueueEntry { pub oldest_age_secs: u64, } +#[derive(Default)] +struct UnitPlacement { + leases: Vec, + reserved: Vec, +} + +type UnitGroupKey<'a> = (&'a str, Option<&'a str>); + +struct UnitPlacementPolicy<'a> { + share: Option, + reserved: &'a [Reservation], + spent_bids: &'a [ActiveLease], +} + +/// Heartbeat bids describe additional slots, so grants made earlier in the +/// same reconciliation pass spend them even when a different queue allocates. +fn spend_bids(drones: &mut [DroneState], granted: &[ActiveLease]) { + for d in drones { + for bid in &mut d.bids { + let used = granted + .iter() + .filter(|g| g.drone_id == d.id && g.class == bid.class && g.scope == bid.scope) + .count(); + bid.available_slots = bid + .available_slots + .saturating_sub(u32::try_from(used).unwrap_or(u32::MAX)); + } + } +} + pub struct Queen { pub(crate) cfg: QueenConfig, pub(crate) store: Store, @@ -175,6 +208,8 @@ pub struct Queen { unit_joins: AtomicU64, unit_receipts: AtomicU64, units_superseded: AtomicU64, + /// Transient round-robin cursor; authoritative work stays in the store. + unit_cursor: std::sync::Mutex>, wake_sender: Arc, tick_lock: tokio::sync::Mutex<()>, } @@ -355,6 +390,7 @@ impl Queen { unit_joins: AtomicU64::new(0), unit_receipts: AtomicU64::new(0), units_superseded: AtomicU64::new(0), + unit_cursor: std::sync::Mutex::new(BTreeMap::new()), wake_sender: Arc::new(UdpWakeSender), tick_lock: tokio::sync::Mutex::new(()), } @@ -564,9 +600,9 @@ impl Queen { .effective_floors(&self.cfg.warm_floors, now); let (warm_floors, homeostasis) = self.regulate(now_ts, &schedule_floors, &queue); let placement_drones = drones.clone(); - let placement_leases = leases.clone(); + let mut placement_leases = leases.clone(); let preempt_ages = queue_ages.clone(); - let allocation = allocate(&AllocatorInput { + let mut runner_input = AllocatorInput { classes: self.cfg.classes.clone(), demand, drones, @@ -576,26 +612,136 @@ impl Queen { starvation_after_secs: self.cfg.starvation_after_secs, weights: ScoreWeights::default(), admission: self.cfg.admission.as_ref().map(|a| a.thresholds()), - }); - + }; + let mut unit_reserved = Vec::new(); + let mut deploy_grants = Vec::new(); + let mut deploy_decisions = Vec::new(); + let mut spent_bids = Vec::new(); + if self.cfg.units.enabled && self.cfg.units.capacity_share_percent > 0 { + // Deploy CI retains first claim. Hold these grants while admitting + // the unit share, then let the ordinary allocator fill the rest. + // Pressure still comes from the full runner queue after the usual + // pending/standby coverage, including shared PR runners for Main. + let deploy = allocate_scope(&runner_input, Scope::Deploy); + deploy_grants = deploy.grants; + deploy_decisions = deploy.decisions; + let mut held = placement_leases.clone(); + spent_bids.extend(deploy_grants.iter().map(|g| ActiveLease { + drone_id: g.drone_id.clone(), + class: g.class, + scope: g.scope, + state: LeaseState::Pending, + })); + held.extend(spent_bids.iter().cloned()); + let units = self + .place_units( + now_ts, + &placement_drones, + &capacities, + held, + UnitPlacementPolicy { + share: Some(self.cfg.units.capacity_share_percent), + reserved: &deploy.reserved_for, + spent_bids: &spent_bids, + }, + ) + .await?; + unit_reserved = units.reserved; + spent_bids.extend(units.leases); + // Re-read persisted placements; units hold capacity without + // offsetting runner demand. + placement_leases = active + .iter() + .map(as_active) + .chain(self.unit_leases().await) + .collect(); + runner_input.leases = placement_leases.clone(); + runner_input + .leases + .extend(deploy_grants.iter().map(|g| ActiveLease { + drone_id: g.drone_id.clone(), + class: g.class, + scope: g.scope, + state: LeaseState::Pending, + })); + // A starving unit needs contiguous room as existing work drains. + // Do not refill its eligible reserved host with new CI this tick. + spend_bids(&mut runner_input.drones, &spent_bids); + } + let mut allocation = allocate_reserved(&runner_input, &unit_reserved); + deploy_grants.append(&mut allocation.grants); + allocation.grants = deploy_grants; + deploy_decisions.append(&mut allocation.decisions); + allocation.decisions = deploy_decisions; + allocation.reserved_for.extend(unit_reserved); + let mut held = placement_leases.clone(); + let runner_spent: Vec = allocation + .grants + .iter() + .map(|g| ActiveLease { + drone_id: g.drone_id.clone(), + class: g.class, + scope: g.scope, + state: LeaseState::Pending, + }) + .collect(); + held.extend(runner_spent.iter().cloned()); + // Deploy grants already occur in spent_bids; add only later CI grants. + spent_bids.extend( + runner_spent + .into_iter() + .filter(|g| self.cfg.units.capacity_share_percent == 0 || g.scope != Scope::Deploy), + ); if let Err(e) = self .place_units( now_ts, &placement_drones, &capacities, - placement_leases.clone(), - &allocation.grants, + held, + UnitPlacementPolicy { + share: None, + reserved: &allocation.reserved_for, + spent_bids: &spent_bids, + }, ) .await { tracing::warn!("unit placement failed: {e}"); } + let unit_leases = self.unit_leases().await; + placement_leases = active + .iter() + .map(as_active) + .chain(unit_leases.iter().cloned()) + .collect(); // Demand that idle non-borg runners will take is not borg's to wake // hosts for. let unmet_placeable = placeable_unmet(&allocation.unmet, &statics, &self.cfg.classes); let wake = self.wake_pass(now_ts, &unmet_placeable, &fresh_names).await; let mut mint_error = None; - for grant in &allocation.grants { + let mut lease_error = None; + let mut persisted = BTreeSet::new(); + // Deploy decisions came from the preliminary pass; align their order + // with the combined grants actually minted below. + let mut assigned = BTreeSet::new(); + for decision in allocation + .decisions + .iter_mut() + .filter(|d| d.winner.is_some()) + { + if let Some((order, _)) = allocation.grants.iter().enumerate().find(|(i, g)| { + !assigned.contains(i) + && Some(g.drone_id.as_str()) == decision.winner.as_deref() + && g.class == decision.class + && g.scope == decision.scope + }) { + decision.grant_order = order; + assigned.insert(order); + } + } + let rows_by_id = drone_rows.iter().map(|d| (d.name.as_str(), d)).collect(); + self.enrich_decisions(&mut allocation.decisions, &rows_by_id, None, now_ts); + for (order, grant) in allocation.grants.iter().enumerate() { let Some(cap) = capacities.get(&grant.drone_id) else { continue; }; @@ -638,8 +784,32 @@ impl Queen { if let Err(e) = self.store.insert_lease(&row).await { // Do not leave an unowned registration behind. let _ = self.github.delete_runner(minted.runner_id).await; - return Err(e.into()); + lease_error = Some(e); + break; } + persisted.insert(order); + } + for decision in &allocation.decisions { + let key = self.cfg.classes.job_label(decision.class, decision.scope); + let outcome = if decision.winner.is_none() { + PlacementOutcome::Unplaced + } else if persisted.contains(&decision.grant_order) { + PlacementOutcome::Granted + } else { + PlacementOutcome::Failed + }; + self.record_placement(PlacementRecord { + at: now_ts, + kind: PlacementKind::Runner, + key: &key, + client: None, + outcome, + decision, + }) + .await; + } + if let Some(e) = lease_error { + return Err(e.into()); } let start_standby = promotions(&allocation.promote, &active, &capacities); @@ -1003,18 +1173,88 @@ impl Queen { Ok(()) } - /// Place queued units with the runner allocator, on drones of each unit's - /// platform, in the capacity left after CI runner grants. + /// Diagnostic writes never change admission or lease persistence. + async fn record_placement(&self, record: PlacementRecord<'_>) { + if let Err(e) = self.store.record_placement(record).await { + tracing::warn!("placement history write failed: {e}"); + } + } + + async fn record_unit_placements( + &self, + now: i64, + history: &VecDeque<(String, String, PlacementOutcome, PlacementDecision)>, + ) { + let records: Vec<_> = history + .iter() + .map(|(key, client, outcome, decision)| PlacementRecord { + at: now, + kind: PlacementKind::Unit, + key, + client: Some(client), + outcome: *outcome, + decision, + }) + .collect(); + if let Err(e) = self.store.record_placements(&records).await { + tracing::warn!("placement history write failed: {e}"); + } + } + + /// Copy only typed effective blockers, never raw environment values or + /// free-form reasons, from the same snapshot used for admission. + fn enrich_decisions( + &self, + decisions: &mut [PlacementDecision], + rows: &HashMap<&str, &DroneRow>, + toolchain: Option<&str>, + now: i64, + ) { + for decision in decisions { + let need = Need::for_class( + &self.cfg.env_contract, + decision.class, + self.cfg.classes.spec(decision.class), + toolchain, + ); + for candidate in &mut decision.candidates { + let facts = rows + .get(candidate.drone_id.as_str()) + .and_then(|d| d.env_facts.as_ref()); + let verdict = + crate::envcontract::evaluate(&self.cfg.env_contract, facts, &need, now); + if verdict.admit { + continue; + } + candidate.environment_checks = verdict + .checks + .iter() + .filter(|c| c.mode == Mode::Required && c.state != CheckState::Pass) + .map(|c| (c.kind, c.state)) + .collect(); + if candidate + .environment_checks + .contains_key(&CheckKind::Capabilities) + { + candidate.capability_checks = + capability_states(&self.cfg.env_contract, facts, decision.class, now); + } + } + } + } + + /// Place units fairly across clients, platforms and toolchains. A share + /// pass precedes non-deploy CI; a borrowing pass consumes only leftover room. async fn place_units( &self, now: i64, drones: &[DroneState], capacities: &HashMap, - leases: Vec, - runner_grants: &[Grant], - ) -> Result<(), QueenError> { + mut base: Vec, + policy: UnitPlacementPolicy<'_>, + ) -> Result { if !self.cfg.units.enabled { - return Ok(()); + return Ok(UnitPlacement::default()); } let queued: Vec = self .store @@ -1024,98 +1264,274 @@ impl Queen { .filter(|u| !u.awaiting_bundle) .collect(); if queued.is_empty() { - return Ok(()); + return Ok(UnitPlacement::default()); } - let mut base = leases; - base.extend(runner_grants.iter().map(|g| ActiveLease { - drone_id: g.drone_id.clone(), - class: g.class, - scope: g.scope, - state: LeaseState::Pending, - })); - // Units of one platform and Rust toolchain share a candidate set: the - // drones of that platform whose build environment contract admits the - // class for that toolchain. - let mut by_platform: BTreeMap<(&str, Option<&str>), Vec<&Unit>> = BTreeMap::new(); + let inflight = self + .store + .units_in(&[UnitState::Leased, UnitState::Running]) + .await?; + let rows = self.store.drones().await?; + let rows_by_id: HashMap<&str, &DroneRow> = + rows.iter().map(|d| (d.name.as_str(), d)).collect(); + let admission = self.cfg.admission.as_ref().map(|a| a.thresholds()); + // Platform and toolchain determine the contract snapshot, independently + // of unit class/scope. Share it across queued units while keeping blocked + // hosts visible to diagnostics and scope eligibility per request. + let mut groups: BTreeMap, Vec> = BTreeMap::new(); + let mut backlogs: BTreeMap, BTreeMap<(Class, Scope), u32>> = + BTreeMap::new(); for u in &queued { - by_platform - .entry((u.platform.as_str(), u.rust_toolchain.as_deref())) + let key = (u.platform.as_str(), u.rust_toolchain.as_deref()); + groups.entry(key).or_insert_with(|| { + let mut group: Vec = drones + .iter() + .filter(|d| { + capacities + .get(&d.id) + .is_some_and(|c| platform_of(c) == u.platform) + }) + .map(|d| { + let mut d = d.clone(); + d.blocked_classes = blocked_classes( + &self.cfg.env_contract, + rows_by_id + .get(d.id.as_str()) + .and_then(|r| r.env_facts.as_ref()), + &self.cfg.classes, + u.rust_toolchain.as_deref(), + now, + ); + d + }) + .collect(); + spend_bids(&mut group, policy.spent_bids); + group + }); + let count = backlogs + .entry(key) .or_default() - .push(u); + .entry((u.class, u.scope)) + .or_default(); + *count = count.saturating_add(1); } - let rows = self.store.drones().await?; - for ((platform, toolchain), units) in by_platform { - let group: Vec = drones + let eligible: BTreeSet<&str> = queued + .iter() + .flat_map(|u| { + groups[&(u.platform.as_str(), u.rust_toolchain.as_deref())] + .iter() + .filter(|d| { + let state = AdmissionState::of(d.cpu_psi_pm, admission); + d.accepting + && d.trust.admits(u.scope) + && state != AdmissionState::Closed + && (state != AdmissionState::PrBlocked || u.scope != Scope::Pr) + && !d.blocked_classes.contains(&u.class) + && max_slots(d, self.cfg.classes.spec(u.class)) > 0 + && (d.bids.is_empty() + || d.bids.iter().any(|b| { + b.class == u.class + && b.scope == u.scope + && b.available_slots > 0 + })) + }) + .map(|d| d.id.as_str()) + }) + .collect(); + let (vcpu, mem, slots) = drones + .iter() + .filter(|d| eligible.contains(d.id.as_str())) + .fold((0u64, 0u64, 0u64), |(c, m, slots), d| { + let capacity = Class::ALL + .into_iter() + .map(|class| max_slots(d, self.cfg.classes.spec(class))) + .max() + .unwrap_or(0); + ( + c.saturating_add(u64::from(d.vcpu)), + m.saturating_add(d.mem_mib), + slots.saturating_add(u64::from(capacity)), + ) + }); + let (mut used_cpu, mut used_mem, mut used_slots) = inflight + .iter() + .filter(|u| u.drone.as_deref().is_some_and(|d| eligible.contains(d))) + .fold((0u64, 0u64, 0u64), |(c, m, slots), u| { + let spec = self.cfg.classes.spec(u.class); + ( + c.saturating_add(u64::from(spec.vcpu)), + m.saturating_add(spec.mem_mib), + slots.saturating_add(1), + ) + }); + let budget = policy.share.map(|percent| { + let portion = |n: u64| n.saturating_mul(u64::from(percent)).div_ceil(100); + (portion(vcpu), portion(mem), portion(slots)) + }); + let mut loads: BTreeMap<&str, usize> = BTreeMap::new(); + for u in &inflight { + *loads.entry(&u.client).or_default() += 1; + } + let mut cursors = self + .unit_cursor + .lock() + .unwrap_or_else(|e| e.into_inner()) + .clone(); + let mut remaining: Vec = (0..queued.len()).collect(); + let mut reservations = policy.reserved.to_vec(); + let mut placed = UnitPlacement::default(); + let mut history = VecDeque::new(); + while !remaining.is_empty() { + // The share is a soft boundary: allow one whole unit to cross it, + // so even a one-slot fleet and heavy units can make progress. + if budget + .is_some_and(|(c, m, slots)| used_cpu >= c || used_mem >= m || used_slots >= slots) + { + break; + } + let next = remaining .iter() - .filter(|d| { - capacities - .get(&d.id) - .is_some_and(|c| platform_of(c) == platform) + .enumerate() + .min_by_key(|(_, i)| { + let u = &queued[**i]; + let age = u64::try_from(now.saturating_sub(u.created_at)).unwrap_or(0); + let starving = self + .cfg + .starvation_after_secs + .is_some_and(|after| age >= after); + let after_cursor = cursors.get(&u.scope).is_none_or(|last| &u.client > last); + ( + u.scope, + std::cmp::Reverse(starving), + starving.then_some(std::cmp::Reverse(u.class)), + loads.get(u.client.as_str()).copied().unwrap_or(0), + std::cmp::Reverse(after_cursor), + &u.client, + u.created_at, + &u.id, + ) }) - .map(|d| { - let mut d = d.clone(); - d.blocked_classes = blocked_classes( - &self.cfg.env_contract, - rows.iter() - .find(|r| r.name == d.id) - .and_then(|r| r.env_facts.as_ref()), - &self.cfg.classes, - toolchain, - now, - ); - d + .map(|(i, _)| i) + .expect("remaining is nonempty"); + let i = remaining.remove(next); + let u = &queued[i]; + let key = (u.platform.as_str(), u.rust_toolchain.as_deref()); + let mut group = groups[&key].clone(); + spend_bids(&mut group, &placed.leases); + let backlog: Vec = backlogs[&key] + .iter() + .filter(|(_, count)| **count > 0) + .map(|(&(class, scope), &count)| Demand { + class, + scope, + count, }) .collect(); - if group.is_empty() { - continue; + let mut allocation = allocate_reserved_with_backlog( + &AllocatorInput { + classes: self.cfg.classes.clone(), + demand: vec![Demand { + class: u.class, + scope: u.scope, + count: 1, + }], + drones: group, + // Pending runner leases must consume room without offsetting + // unit demand: they execute a different queue's work. + leases: base + .iter() + .cloned() + .map(|mut l| { + l.state = LeaseState::Running; + l + }) + .collect(), + warm_floors: Vec::new(), + queue_ages: vec![QueueAge { + class: u.class, + scope: u.scope, + oldest_age_secs: u64::try_from(now.saturating_sub(u.created_at)) + .unwrap_or(0), + }], + starvation_after_secs: self.cfg.starvation_after_secs, + weights: ScoreWeights::default(), + admission, + }, + &reservations, + &backlog, + ); + for r in allocation.reserved_for { + if !reservations.iter().any(|old| old.drone_id == r.drone_id) { + reservations.push(r.clone()); + placed.reserved.push(r); + } } - let mut demand: BTreeMap<(Class, Scope), u32> = BTreeMap::new(); - for u in &units { - *demand.entry((u.class, u.scope)).or_default() += 1; + self.enrich_decisions( + &mut allocation.decisions, + &rows_by_id, + u.rust_toolchain.as_deref(), + now, + ); + let leased = match allocation.grants.first() { + Some(g) => self.store.lease_unit(&u.id, &g.drone_id, now).await, + None => Ok(false), + }; + let outcome = if allocation.grants.is_empty() { + PlacementOutcome::Unplaced + } else if matches!(leased, Ok(true)) { + PlacementOutcome::Granted + } else { + PlacementOutcome::Failed + }; + for decision in allocation.decisions { + // Only the newest retained records can survive store pruning. + if history.len() == MAX_HISTORY as usize { + history.pop_front(); + } + let outcome = if decision.winner.is_none() { + PlacementOutcome::Unplaced + } else { + outcome + }; + history.push_back((u.id.clone(), u.client.clone(), outcome, decision)); } - let allocation = allocate(&AllocatorInput { - classes: self.cfg.classes.clone(), - demand: demand - .into_iter() - .map(|((class, scope), count)| Demand { - class, - scope, - count, - }) - .collect(), - drones: group, - leases: base.clone(), - warm_floors: Vec::new(), - queue_ages: Vec::new(), - starvation_after_secs: None, - weights: ScoreWeights::default(), - admission: self.cfg.admission.as_ref().map(|a| a.thresholds()), - }); - // Grants of a (class, scope) are interchangeable: hand them out - // deploy first, then main, then pull request, each in submission - // order. - let mut grants: Vec<&Grant> = allocation.grants.iter().collect(); - let mut ordered = units; - ordered.sort_by_key(|u| u.scope); - for u in ordered { - if let Some(i) = grants - .iter() - .position(|g| g.class == u.class && g.scope == u.scope) - { - let g = grants.remove(i); - if self.store.lease_unit(&u.id, &g.drone_id, now).await? { - base.push(ActiveLease { - drone_id: g.drone_id.clone(), - class: g.class, - scope: g.scope, - state: LeaseState::Running, - }); - tracing::info!("unit {} placed on {}", u.id, g.drone_id); - } + let leased = match leased { + Ok(leased) => leased, + Err(e) => { + self.record_unit_placements(now, &history).await; + return Err(e.into()); } + }; + if leased && let Some(g) = allocation.grants.first() { + // Only persisted leases cover backlog; failed and attempted + // unplaced work continues to contribute to demand pressure. + let count = backlogs + .get_mut(&key) + .and_then(|counts| counts.get_mut(&(u.class, u.scope))) + .expect("queued unit has a backlog count"); + *count = count.saturating_sub(1); + let lease = ActiveLease { + drone_id: g.drone_id.clone(), + class: g.class, + scope: g.scope, + state: LeaseState::Running, + }; + base.push(lease.clone()); + placed.leases.push(lease); + let spec = self.cfg.classes.spec(u.class); + used_cpu = used_cpu.saturating_add(u64::from(spec.vcpu)); + used_mem = used_mem.saturating_add(spec.mem_mib); + used_slots = used_slots.saturating_add(1); + *loads.entry(&u.client).or_default() += 1; + cursors.insert(u.scope, u.client.clone()); + self.unit_cursor + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(u.scope, u.client.clone()); + tracing::info!("unit {} placed on {}", u.id, g.drone_id); } } - Ok(()) + self.record_unit_placements(now, &history).await; + Ok(placed) } /// Send wake-on-LAN packets for demand that has stayed unmet. diff --git a/src/run.rs b/src/run.rs index 0b5e6e8..b77a9de 100644 --- a/src/run.rs +++ b/src/run.rs @@ -2,10 +2,12 @@ //! //! The client resolves the tree it is asked to run against, submits a work unit //! to the queen (`POST /v1/units`, see [`crate::units`]), streams the unit's log -//! to this process's stdout and exits with the unit's exit code. It never makes -//! things worse than running the command directly: when the queen is not -//! configured, unreachable, refuses the unit, or no drone picks it up within -//! `wait_secs`, the command runs locally with this process's own stdio. +//! to this process's stdout and exits with the unit's exit code. By default, +//! when the queen is not configured, unreachable, refuses the unit, or no drone +//! picks it up within `wait_secs`, the command runs locally with this process's +//! own stdio. `--remote-only` disables every local fallback and returns +//! infrastructure exit code 125 instead. `--result-file` writes JSON execution +//! provenance separately from command output. //! //! * The tree is `git write-tree` of the index (`--index`, for pre-commit //! hooks) or `HEAD^{tree}` (default), or an explicit SHA. @@ -13,21 +15,24 @@ //! streamed); one that already succeeded is not run at all: the receipt is //! reported and the exit code is 0. //! * A unit that fails for infrastructure reasons (cancelled, timed out, out of -//! attempts, or the queen lost track of it) is run locally rather than -//! reported as the command's failure. A command that ran and exited nonzero +//! attempts, or the queen lost track of it) uses local fallback by default +//! or returns 125 in remote-only mode. A command that ran and exited nonzero //! on a drone is a real failure: its code is returned as is. //! * The drone must be able to get the tree. A tree from a pushed commit is in //! its git mirror. Otherwise `--publish` uploads a `git bundle` of just the //! missing objects (everything since the merge base with `origin/main`) to the //! queen, which hands it to the drone; nothing is pushed to any remote. Without -//! `--publish`, an unpublished tree runs locally. +//! `--publish`, an unpublished tree uses local fallback by default or returns +//! 125 in remote-only mode. use std::io::Write; use std::path::{Path, PathBuf}; use std::process::Command as StdCommand; use std::time::{Duration, Instant}; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; + +use crate::units::{Receipt, UnitState}; use crate::secret::Secret; @@ -60,6 +65,10 @@ pub struct RunOptions { /// How long to wait for a drone to pick the unit up before running locally. pub wait_secs: u64, pub publish: bool, + /// Refuse local fallback and return 125 when infrastructure cannot finish the unit. + pub remote_only: bool, + /// Write a JSON execution result separately from command output. + pub result_file: Option, pub argv: Vec, } @@ -264,7 +273,7 @@ fn signal_of(_: &std::process::ExitStatus) -> i32 { struct Submitted { outcome: String, unit_id: String, - receipt: Option, + receipt: Option, } struct Remote<'a> { @@ -277,6 +286,25 @@ impl Remote<'_> { fn url(&self, path: &str) -> String { format!("{}{path}", self.base.trim_end_matches('/')) } + + async fn cancel_created(&self, id: &str, created: bool) { + if created { + let _ = self + .http + .post(self.url(&format!("/v1/units/{id}/cancel"))) + .bearer_auth(self.token.expose()) + .timeout(Duration::from_secs(2)) + .send() + .await; + } + } +} + +async fn wait_deadline(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending::<()>().await, + } } async fn upload_bundle(q: &Remote<'_>, dir: &Path, tree: &str, unit: &str) -> Result<(), String> { @@ -307,28 +335,170 @@ enum Fallback { Local(String), } -/// Run `opts.argv`: remotely when possible, otherwise here. Returns the exit -/// code to exit with. Remote log bytes go to `out`. +/// Machine-readable execution result, separate from streamed command output. +/// `requested_tree_sha` identifies the requested snapshot, not the mutable +/// working directory used by local fallback. A null host means unavailable. +#[derive(Debug, Serialize)] +pub struct RunResult { + pub schema_version: u32, + pub requested_tree_sha: Option, + pub argv: Vec, + /// `remote`, `local`, or `none` (no execution was observed). + pub execution: String, + /// `new`, `joined`, `reused`, or `fallback`. + pub provenance: Option, + pub unit_id: Option, + pub host: Option, + pub exit_code: i32, + pub infrastructure_error: Option, + pub receipt: Option, +} + +#[derive(Deserialize)] +struct UnitStatus { + id: String, + repo: String, + tree_sha: String, + argv: Vec, + platform: String, + toolchain: String, + state: String, + exit_code: Option, + drone: Option, + timeout_secs: u64, + receipt: Option, +} + +async fn unit_status(q: &Remote<'_>, id: &str) -> Result { + q.http + .get(q.url(&format!("/v1/units/{id}"))) + .bearer_auth(q.token.expose()) + .timeout(Duration::from_secs(10)) + .send() + .await + .map_err(|e| format!("cannot read unit metadata ({e})"))? + .error_for_status() + .map_err(|e| format!("cannot read unit metadata ({e})"))? + .json() + .await + .map_err(|e| format!("unreadable unit metadata ({e})")) +} + +fn record_status( + result: &mut RunResult, + unit: UnitStatus, + exit: Option, + opts: &RunOptions, + repo: &str, +) -> Result { + if result.unit_id.as_deref() != Some(unit.id.as_str()) + || result.requested_tree_sha.as_deref() != Some(unit.tree_sha.as_str()) + || result.argv != unit.argv + || unit.repo != repo + || unit.platform != opts.platform + || unit.toolchain != opts.toolchain + || exit.is_some_and(|code| { + unit.state != UnitState::Done.as_str() || unit.exit_code != Some(code) + }) + { + return Err("unit metadata does not match the submitted command".into()); + } + if let Some(receipt) = &unit.receipt + && (receipt.unit_id != unit.id + || receipt.tree_sha != unit.tree_sha + || receipt.argv != unit.argv + || receipt.repo != unit.repo + || receipt.platform != unit.platform + || receipt.toolchain != unit.toolchain + || unit.drone.as_deref() != Some(receipt.drone.as_str()) + || unit.state != UnitState::Done.as_str() + || unit.exit_code != Some(0) + || receipt.exit_code != 0 + || !receipt.verifies()) + { + return Err("unit receipt does not match the completed command".into()); + } + if unit.state == UnitState::Running.as_str() || unit.state == UnitState::Done.as_str() { + result.execution = "remote".into(); + } + result.host = unit.drone; + result.receipt = unit.receipt; + Ok(unit.timeout_secs) +} + +/// Run remotely when possible. `remote_only` prevents every local fallback. +/// Returns 125 for infrastructure errors and result-file failures. pub async fn run(opts: &RunOptions, out: &mut dyn Write) -> i32 { - match remote(opts, out).await { + // Open before execution so an invalid result destination cannot quietly + // execute work without the requested result. + let file = match opts + .result_file + .as_ref() + .map(std::fs::File::create) + .transpose() + { + Ok(file) => file, + Err(e) => { + eprintln!("borg run: cannot create result file: {e}"); + return 125; + } + }; + let mut result = RunResult { + schema_version: 1, + requested_tree_sha: None, + argv: opts.argv.clone(), + execution: "none".into(), + provenance: None, + unit_id: None, + host: None, + exit_code: 125, + infrastructure_error: None, + receipt: None, + }; + result.exit_code = match remote(opts, out, &mut result).await { Ok(code) => code, Err(Fallback::Local(why)) => { - eprintln!("borg run: {why}; running locally"); - run_local(&opts.argv, &opts.dir) + result.infrastructure_error = Some(why.clone()); + if opts.remote_only { + eprintln!("borg run: {why}; remote-only execution failed"); + 125 + } else { + eprintln!("borg run: {why}; running locally"); + result.execution = "local".into(); + result.provenance = Some("fallback".into()); + result.host = None; + result.receipt = None; + run_local(&opts.argv, &opts.dir) + } } + }; + if let Some(mut file) = file + && let Err(e) = serde_json::to_writer_pretty(&mut file, &result) + .map_err(std::io::Error::other) + .and_then(|()| file.write_all(b"\n")) + .and_then(|()| file.flush()) + { + eprintln!("borg run: cannot write result file: {e}"); + return 125; } + result.exit_code } #[allow( clippy::too_many_lines, reason = "one linear submit, wait, stream, finish flow" )] -async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result { +async fn remote( + opts: &RunOptions, + out: &mut dyn Write, + result: &mut RunResult, +) -> Result { let local = |why: &str| Fallback::Local(why.to_owned()); let (Some(base), Some(token)) = (&opts.queen, &opts.token) else { return Err(local("no queen configured")); }; let tree = resolve_tree(&opts.dir, &opts.tree).map_err(Fallback::Local)?; + result.requested_tree_sha = Some(tree.clone()); let repo = match &opts.repo { Some(r) => r.clone(), None => git(&opts.dir, &["remote", "get-url", "origin"]) @@ -371,29 +541,46 @@ async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result .json() .await .map_err(|e| Fallback::Local(format!("unreadable queen answer ({e})")))?; + result.unit_id = Some(sub.unit_id.clone()); if sub.outcome == "receipt" { - let drone = sub + let receipt = sub .receipt - .as_ref() - .and_then(|r| r["drone"].as_str()) - .unwrap_or("a drone"); + .ok_or_else(|| local("queen returned no reusable receipt"))?; + if receipt.unit_id != sub.unit_id + || receipt.tree_sha != tree + || receipt.argv != opts.argv + || receipt.repo != repo + || receipt.platform != opts.platform + || receipt.toolchain != opts.toolchain + || receipt.exit_code != 0 + || !receipt.verifies() + { + return Err(local( + "queen returned a receipt that does not match the command", + )); + } eprintln!( - "borg run: tree {} already passed on {drone} (unit {}); not running again", + "borg run: tree {} already passed on {} (unit {}); not running again", &tree[..tree.len().min(12)], + receipt.drone, sub.unit_id ); + result.execution = "remote".into(); + result.provenance = Some("reused".into()); + result.host = Some(receipt.drone.clone()); + result.receipt = Some(receipt); return Ok(0); } - let created = sub.outcome == "new"; + let created = match sub.outcome.as_str() { + "new" => true, + "joined" => false, + _ => return Err(local("queen returned an unknown submission outcome")), + }; + result.provenance = Some(if created { "new" } else { "joined" }.into()); if created && needs_bundle { let sent = upload_bundle(&q, &opts.dir, &tree, &sub.unit_id).await; if let Err(why) = sent { - let _ = q - .http - .post(q.url(&format!("/v1/units/{}/cancel", sub.unit_id))) - .bearer_auth(q.token.expose()) - .send() - .await; + q.cancel_created(&sub.unit_id, created).await; return Err(Fallback::Local(why)); } } @@ -403,9 +590,27 @@ async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result if created { "queued" } else { "joined" } ); + let timeout_secs = if opts.remote_only { + // Use the admitted timeout, including the queen's configured default + // and cap. The client also enforces it if the queen stops making progress. + let metadata = unit_status(&q, &sub.unit_id) + .await + .and_then(|unit| record_status(result, unit, None, opts, &repo)); + match metadata { + Ok(timeout) => Some(timeout), + Err(why) => { + q.cancel_created(&sub.unit_id, created).await; + return Err(Fallback::Local(why)); + } + } + } else { + None + }; let started = Instant::now(); + let queued_deadline = tokio::time::Instant::now() + Duration::from_secs(opts.wait_secs); + let deadline = timeout_secs.map(|t| queued_deadline + Duration::from_secs(t)); let mut offset: u64 = 0; - let mut running = false; + let mut running = result.execution == "remote"; let mut unreachable_since: Option = None; loop { let poll = q @@ -418,11 +623,16 @@ async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result .timeout(Duration::from_secs(25)); let resp = tokio::select! { r = poll.send() => r, + _ = wait_deadline(deadline.map(|d| if running { d } else { d.min(queued_deadline) })) => { + q.cancel_created(&sub.unit_id, created).await; + return Err(local(if running { + "remote execution exceeded the command deadline" + } else { + "no drone took the unit within the admission deadline" + })); + }, _ = tokio::signal::ctrl_c() => { - if created { - let _ = q.http.post(q.url(&format!("/v1/units/{}/cancel", sub.unit_id))) - .bearer_auth(q.token.expose()).send().await; - } + q.cancel_created(&sub.unit_id, created).await; eprintln!("borg run: interrupted"); return Ok(130); } @@ -434,9 +644,16 @@ async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result } other => { let since = *unreachable_since.get_or_insert_with(Instant::now); - if since.elapsed() > Duration::from_secs(60) { + let loss_limit = if opts.remote_only { + opts.wait_secs.min(60) + } else { + 60 + }; + if since.elapsed() >= Duration::from_secs(loss_limit) { + q.cancel_created(&sub.unit_id, created && opts.remote_only) + .await; return Err(Fallback::Local(match other { - Ok(r) => format!("queen answered {} for 60s", r.status()), + Ok(r) => format!("queen answered {} while reading logs", r.status()), Err(e) => format!("lost the queen ({e})"), })); } @@ -454,25 +671,62 @@ async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result let next = hdr("x-borg-next").and_then(|v| v.parse::().ok()); let exit = hdr("x-borg-exit").and_then(|v| v.parse::().ok()); let failure = hdr("x-borg-failure"); - let bytes = resp - .bytes() - .await - .map_err(|e| Fallback::Local(format!("log read failed ({e})")))?; - if !bytes.is_empty() { - let _ = out.write_all(&bytes); - let _ = out.flush(); + if state == "running" || state == "done" { + running = true; + result.execution = "remote".into(); + } + let bytes = tokio::select! { + bytes = resp.bytes() => match bytes { + Ok(bytes) => bytes, + Err(e) => { + q.cancel_created(&sub.unit_id, created && opts.remote_only).await; + return Err(Fallback::Local(format!("log read failed ({e})"))); + } + }, + _ = wait_deadline(deadline.map(|d| if running { d } else { d.min(queued_deadline) })) => { + q.cancel_created(&sub.unit_id, created).await; + return Err(local("remote log exceeded the execution deadline")); + }, + }; + if !bytes.is_empty() + && let Err(e) = out.write_all(&bytes).and_then(|()| out.flush()) + { + let why = format!("cannot write remote log ({e})"); + if opts.remote_only { + q.cancel_created(&sub.unit_id, created).await; + return Err(Fallback::Local(why)); + } + // A broken output consumer must never rerun a command that already + // executed remotely. Preserve the default behavior: drain the log + // and return the remote exit code even if displaying output fails. + result.infrastructure_error = Some(why); } if let Some(n) = next { offset = n; } - if state == "running" || state == "done" { - running = true; - } match state.as_str() { "done" => { if let Some(code) = exit { + if opts.result_file.is_some() { + match unit_status(&q, &sub.unit_id) + .await + .and_then(|unit| record_status(result, unit, Some(code), opts, &repo)) + { + Ok(_) => {} + Err(why) => { + // The command finished. Metadata loss must never + // trigger a second execution through local fallback. + result.host = None; + result.receipt = None; + result.infrastructure_error = Some(why); + } + } + } return Ok(code); } + // The queen labels every chunk of a completed unit "done", + // but sends the exit code only after its log is fully drained. + // Keep polling from `next` until that final chunk arrives. } "failed" => { let why = failure.unwrap_or_else(|| "failed".into()); @@ -483,14 +737,7 @@ async fn remote(opts: &RunOptions, out: &mut dyn Write) -> Result "queued" | "leased" if !running && started.elapsed() >= Duration::from_secs(opts.wait_secs) => { - if created { - let _ = q - .http - .post(q.url(&format!("/v1/units/{}/cancel", sub.unit_id))) - .bearer_auth(q.token.expose()) - .send() - .await; - } + q.cancel_created(&sub.unit_id, created).await; return Err(Fallback::Local(format!( "no drone took it within {}s", opts.wait_secs diff --git a/src/store.rs b/src/store.rs index 30f54bb..d24198f 100644 --- a/src/store.rs +++ b/src/store.rs @@ -61,6 +61,18 @@ CREATE TABLE IF NOT EXISTS reclaim_log ( PRIMARY KEY (drone, seq) ); CREATE INDEX IF NOT EXISTS reclaim_log_at ON reclaim_log(at); +-- Diagnostic metadata only, owned and pruned with the queen's existing store. +CREATE TABLE IF NOT EXISTS placement_history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + at INTEGER NOT NULL, + kind TEXT NOT NULL CHECK (kind IN ('runner', 'unit')), + key TEXT NOT NULL, + client TEXT, + outcome TEXT NOT NULL, + decision TEXT NOT NULL, + CHECK ((kind = 'runner' AND client IS NULL) OR (kind = 'unit' AND client IS NOT NULL)) +); +CREATE INDEX IF NOT EXISTS placement_history_key ON placement_history(kind, key, id); CREATE TABLE IF NOT EXISTS clients ( name TEXT PRIMARY KEY, token_hash TEXT NOT NULL UNIQUE, @@ -403,6 +415,107 @@ fn enum_from(s: &str) -> Result { } impl Store { + /// Historical placement cannot be derived from current heartbeats. Keep it + /// in this existing store, bounded by rows and serialized bytes, pruning in + /// the same transaction as insertion. Unit ownership is checked against its + /// authoritative row before persisting the diagnostic owner. + pub async fn record_placement( + &self, + record: crate::explain::PlacementRecord<'_>, + ) -> Result<(), StoreError> { + self.record_placements(&[record]).await + } + + /// Persist one bounded placement pass atomically, pruning once after its + /// inserts. Invalid bounds or a mismatched unit owner reject the whole batch. + pub async fn record_placements( + &self, + records: &[crate::explain::PlacementRecord<'_>], + ) -> Result<(), StoreError> { + use crate::explain::{MAX_DECISION_BYTES, MAX_HISTORY, PlacementKind}; + if records.len() > MAX_HISTORY as usize { + return Err(StoreError::Corrupt( + "placement batch exceeds diagnostic bounds".into(), + )); + } + if records.is_empty() { + return Ok(()); + } + let mut encoded = Vec::with_capacity(records.len()); + for record in records { + if record.key.is_empty() || record.key.len() > 256 { + return Err(StoreError::Corrupt( + "placement key exceeds diagnostic bounds".into(), + )); + } + let decision = to_text(record.decision)?; + if decision.len() > MAX_DECISION_BYTES + || record.decision.candidates.len() > crate::allocator::MAX_CANDIDATES + { + return Err(StoreError::Corrupt( + "placement exceeds diagnostic bounds".into(), + )); + } + encoded.push(decision); + } + let mut tx = self.pool.begin().await?; + for (record, decision) in records.iter().zip(encoded) { + match record.kind { + PlacementKind::Runner if record.client.is_some() => { + return Err(StoreError::Corrupt("runner placement has a client".into())); + } + PlacementKind::Unit => { + let owner: Option = + // sqlx-runtime-query: standalone SQLite diagnostic owner check + sqlx::query_scalar("SELECT client FROM units WHERE id = ?") + .bind(record.key) + .fetch_optional(&mut *tx) + .await?; + if record.client.is_none() || owner.as_deref() != record.client { + return Err(StoreError::Corrupt("unit placement owner mismatch".into())); + } + } + PlacementKind::Runner => {} + } + // sqlx-runtime-query: standalone-buildable bounded SQLite diagnostics + sqlx::query("INSERT INTO placement_history (at, kind, key, client, outcome, decision) VALUES (?, ?, ?, ?, ?, ?)") + .bind(record.at).bind(record.kind.as_str()).bind(record.key).bind(record.client) + .bind(enum_text(&record.outcome)?).bind(decision).execute(&mut *tx).await?; + } + // sqlx-runtime-query: all batch inserts and retention form one transaction + sqlx::query("DELETE FROM placement_history WHERE id NOT IN (SELECT id FROM placement_history ORDER BY id DESC LIMIT ?)") + .bind(MAX_HISTORY).execute(&mut *tx).await?; + tx.commit().await?; + Ok(()) + } + + /// `client` restricts unit history to its persisted owner. `None` for a unit + /// is deliberately queen-local administration; HTTP callers must supply + /// their authenticated client. Runner reads always exclude unit records. + pub async fn placement_history( + &self, + kind: crate::explain::PlacementKind, + key: &str, + client: Option<&str>, + ) -> Result, StoreError> { + // sqlx-runtime-query: typed bounded query; no command payloads selected + let rows = sqlx::query("SELECT id, at, kind, key, outcome, decision FROM placement_history WHERE kind = ? AND key = ? AND (? IS NULL OR client = ?) ORDER BY id DESC LIMIT ?") + .bind(kind.as_str()).bind(key).bind(client).bind(client) + .bind(crate::explain::MAX_QUERY_ROWS).fetch_all(&self.pool).await?; + rows.into_iter() + .map(|row| { + Ok(crate::explain::HistoryEntry { + id: row.try_get("id")?, + at: row.try_get("at")?, + kind: enum_from(&row.try_get::("kind")?)?, + key: row.try_get("key")?, + outcome: enum_from(&row.try_get::("outcome")?)?, + decision: from_text(&row.try_get::("decision")?)?, + }) + }) + .collect() + } + pub async fn open(path: &Path) -> Result { let opts = SqliteConnectOptions::new() .filename(path) diff --git a/src/units.rs b/src/units.rs index 986cb78..7199109 100644 --- a/src/units.rs +++ b/src/units.rs @@ -12,8 +12,9 @@ //! never stand in for a different tree. `force` skips reuse. //! * **Placement.** Queued units are placed by the same allocator that places //! runner leases (class capacity, owner caps, bids, learned profiles, -//! quarantine and probation), on drones of the unit's platform, after CI -//! runner demand has been served. +//! quarantine and probation), on drones of the unit's platform. A configurable +//! capacity share serves queued units before non-deploy CI demand; either +//! queue may borrow unused capacity. //! * **Execution** is the drone's: it is told about a unit in its heartbeat //! response (`units`), acknowledges it with `active_units`, appends log bytes //! with `POST /v1/units/{id}/log?offset=N`, and ends it with @@ -40,6 +41,12 @@ use crate::store::{Store, StoreError}; pub struct UnitSettings { /// Off unless the queen opts in; the endpoints answer 404 otherwise. pub enabled: bool, + /// Guaranteed share of eligible fleet CPU, memory and lease slots for work units, + /// before non-deploy CI demand, in percent. In-flight units count toward + /// the share; an indivisible unit may cross the budget once. Unused share + /// is available to CI, and units may borrow room left by CI. Zero keeps + /// strict CI-first scheduling. Does not preempt work already running. + pub capacity_share_percent: u8, /// A successful result is reused for this long. pub receipt_ttl_secs: i64, /// A placed unit the drone has not acknowledged within this is re-queued. @@ -64,6 +71,7 @@ impl Default for UnitSettings { fn default() -> Self { Self { enabled: false, + capacity_share_percent: 25, receipt_ttl_secs: 7 * 86_400, lease_ttl_secs: 120, default_timeout_secs: 3600, @@ -80,6 +88,9 @@ impl Default for UnitSettings { impl UnitSettings { pub fn validate(&self) -> Result<(), String> { + if self.capacity_share_percent > 100 { + return Err("units.capacity_share_percent must be between 0 and 100".into()); + } if self.receipt_ttl_secs <= 0 || self.lease_ttl_secs <= 0 || self.keep_secs <= 0 @@ -1067,6 +1078,26 @@ mod tests { } } + #[test] + fn capacity_share_defaults_and_bounds() { + let defaults: UnitSettings = serde_json::from_str("{}").unwrap(); + assert_eq!(defaults.capacity_share_percent, 25); + for value in [0, 100] { + let settings: UnitSettings = + serde_json::from_value(serde_json::json!({"capacity_share_percent": value})) + .unwrap(); + assert!(settings.validate().is_ok()); + } + let invalid: UnitSettings = + serde_json::from_str(r#"{"capacity_share_percent":101}"#).unwrap(); + assert!( + invalid + .validate() + .unwrap_err() + .contains("capacity_share_percent") + ); + } + #[test] fn the_key_covers_every_field_and_cannot_collide_by_concatenation() { let base = submit(); diff --git a/src/units_api.rs b/src/units_api.rs index f64f614..e2b6684 100644 --- a/src/units_api.rs +++ b/src/units_api.rs @@ -52,7 +52,10 @@ fn enabled(queen: &Queen) -> Result<(), ApiError> { } } -async fn authenticate_client(queen: &Queen, headers: &HeaderMap) -> Result { +pub(crate) async fn authenticate_client( + queen: &Queen, + headers: &HeaderMap, +) -> Result { let presented = headers .get("authorization") .and_then(|v| v.to_str().ok()) diff --git a/tests/contract.rs b/tests/contract.rs index 5705909..098bf37 100644 --- a/tests/contract.rs +++ b/tests/contract.rs @@ -369,3 +369,63 @@ async fn a_drone_missing_a_required_tool_gets_no_lease_and_status_names_it() { assert!(m.contains(line), "missing {line}\n{m}"); } } + +#[tokio::test] +async fn class_specific_job_user_proofs_gate_ci_and_survive_heartbeat_storage() { + let h = Harness::start(cfg(json!({ + "env_contract": { + "enabled": true, + "checks": {"capabilities": "required"}, + "classes": {"heavy": {"required_capabilities": ["sandbox", "docker"]}} + } + }))) + .await; + let old = h.enroll("a-inventory-only").await; + let ok = h.enroll("b-proven").await; + let stale = h.enroll("c-stale-proof").await; + let mut fresh = good("1.95.0"); + fresh["capabilities"] = json!({ + "sandbox": {"observed_at": now(), "state": "pass"}, + "docker": {"observed_at": now(), "state": "pass"} + }); + beat(&h, &old, Some(good("1.95.0"))).await; + beat(&h, &ok, Some(fresh.clone())).await; + fresh["capabilities"]["sandbox"]["observed_at"] = json!(now() - 3600); + beat(&h, &stale, Some(fresh)).await; + h.queue_job(1, "borg-ci-heavy-pr", 5); + h.tick().await; + for token in [&old, &stale] { + assert!( + beat(&h, token, None).await["leases"] + .as_array() + .unwrap() + .is_empty() + ); + } + assert_eq!( + beat(&h, &ok, None).await["leases"] + .as_array() + .unwrap() + .len(), + 1 + ); + let status = h.status().await; + for name in ["a-inventory-only", "c-stale-proof"] { + let admission = &drone(&status, name)["contract"]["admission"]; + assert_eq!(admission["heavy"]["admit"], false); + assert_eq!(admission["light"]["admit"], true); + assert!( + admission["heavy"]["why"] + .as_str() + .unwrap() + .contains("sandbox") + ); + } + let stored = h.queen.store().drones().await.unwrap(); + let stale = stored.iter().find(|d| d.name == "c-stale-proof").unwrap(); + assert!( + stale.env_facts.as_ref().unwrap().capabilities[&borg::envcontract::Capability::Sandbox] + .observed_at + < now() - 3500 + ); +} diff --git a/tests/drone.rs b/tests/drone.rs index e5b44cf..75da712 100644 --- a/tests/drone.rs +++ b/tests/drone.rs @@ -46,6 +46,10 @@ impl Host for FixedHost { struct RecordingExec(Mutex>); impl Exec for RecordingExec { + fn run_grouped(&self, program: &str, args: &[&str]) -> io::Result { + self.run(program, args) + } + fn run(&self, program: &str, args: &[&str]) -> io::Result { self.0 .lock() @@ -792,3 +796,193 @@ async fn an_idle_drone_asks_for_the_version_the_queen_wants_and_takes_no_lease() "no new lease while an update is pending" ); } + +#[tokio::test] +async fn capability_heartbeat_proves_job_user_execution_and_reuses_probe_timestamps() { + let rig = Rig::new().await; + let drone = rig.drone(json!({ + "job_user": {"name": "borg-job", "home": "/var/lib/borg-job"}, + "capability_probe": {"docker_image": "busybox:local"} + })); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + drain(&mut tasks).await; + // The next heartbeat publishes the asynchronously completed proof. + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + let first = rig.h.queen.store().drones().await.unwrap(); + let facts = &first[0].env_facts.as_ref().unwrap().capabilities; + for name in borg::envcontract::Capability::ALL { + assert_eq!(facts[&name].state, borg::envcontract::CheckState::Pass); + assert!(facts[&name].observed_at > 0); + } + let calls = rig.exec.0.lock().unwrap().clone(); + let probes: Vec<_> = calls + .iter() + .filter(|c| c.contains("--kill-after=2s")) + .collect(); + assert_eq!(probes.len(), 7); + assert!(probes.iter().all( + |c| c.starts_with("setpriv --reuid borg-job --regid borg-job") + && c.contains("--bounding-set=-all -- timeout") + )); + drone.heartbeat_once(true, &mut tasks).await.unwrap(); + let second = rig.h.queen.store().drones().await.unwrap(); + assert_eq!(facts, &second[0].env_facts.as_ref().unwrap().capabilities); + assert_eq!( + rig.exec + .0 + .lock() + .unwrap() + .iter() + .filter(|c| c.contains("--kill-after=2s")) + .count(), + 7 + ); +} + +/// Deliberately blocks only capability commands; inventory/tool probes stay +/// fast. Releasing the gate lets the full smoke and cleanup sequence finish. +#[derive(Default)] +struct GatedCapabilityExec { + released: (Mutex, std::sync::Condvar), + started: tokio::sync::Notify, + commands: std::sync::atomic::AtomicUsize, +} + +impl GatedCapabilityExec { + fn release(&self) { + *self.released.0.lock().unwrap() = true; + self.released.1.notify_all(); + } +} + +impl Exec for GatedCapabilityExec { + fn run_grouped(&self, program: &str, args: &[&str]) -> io::Result { + self.run(program, args) + } + + fn run(&self, _program: &str, args: &[&str]) -> io::Result { + if args.contains(&"--kill-after=2s") { + self.commands + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + self.started.notify_one(); + let mut released = self.released.0.lock().unwrap(); + while !*released { + released = self.released.1.wait(released).unwrap(); + } + } + Ok(Output { + success: true, + stdout: String::new(), + }) + } +} + +/// A failing assertion must also release the blocking thread, so a regression +/// fails promptly instead of hanging the Tokio runtime's shutdown. +struct ReleaseCapabilityGate(Arc); +impl Drop for ReleaseCapabilityGate { + fn drop(&mut self) { + self.0.release(); + } +} + +#[tokio::test] +async fn blocked_capability_refresh_does_not_block_heartbeats_or_duplicate_probes() { + let rig = Rig::with_queen(json!({ + "env_contract": { + "enabled": true, + "checks": {"capabilities": "required"}, + "required_capabilities": ["sandbox"] + } + })) + .await; + let exec = Arc::new(GatedCapabilityExec::default()); + let release = ReleaseCapabilityGate(Arc::clone(&exec)); + let drone = rig.drone_with( + json!({ + "capability_probe": {"timeout_secs": 10, "docker_image": "busybox:local"} + }), + exec.clone(), + ); + let mut tasks = JoinSet::new(); + tokio::time::timeout( + Duration::from_secs(30), + drone.heartbeat_once(true, &mut tasks), + ) + .await + .expect("heartbeat must not wait for smoke completion") + .unwrap(); + tokio::time::timeout(Duration::from_secs(30), exec.started.notified()) + .await + .expect("smoke should start in background"); + assert_eq!(exec.commands.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!( + tasks.len(), + 1, + "unfinished refresh belongs to the drain task set" + ); + let stored = rig.h.queen.store().drones().await.unwrap(); + assert!( + stored[0].last_heartbeat.is_some(), + "heartbeat reached the queen while the smoke was blocked" + ); + assert!( + stored[0] + .env_facts + .as_ref() + .unwrap() + .capabilities + .is_empty() + ); + assert_eq!( + rig.h.status().await["drones"][0]["contract"]["admission"]["light"]["admit"], + false + ); + tokio::time::timeout( + Duration::from_secs(30), + drone.heartbeat_once(true, &mut tasks), + ) + .await + .expect("later heartbeat must also remain independent of the refresh") + .unwrap(); + assert_eq!(exec.commands.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(tasks.len(), 1, "refresh is single-flight"); + tokio::time::timeout( + Duration::from_secs(30), + drone.heartbeat_once(false, &mut tasks), + ) + .await + .expect("draining heartbeat must not wait for the refresh") + .unwrap(); + assert_eq!(tasks.len(), 1); + drop(release); + tokio::time::timeout(Duration::from_secs(30), drain(&mut tasks)) + .await + .expect("drain joins refresh and its cleanup"); + assert_eq!(exec.commands.load(std::sync::atomic::Ordering::SeqCst), 7); + drone.heartbeat_once(false, &mut tasks).await.unwrap(); + let stored = rig.h.queen.store().drones().await.unwrap(); + assert_eq!( + stored[0].env_facts.as_ref().unwrap().capabilities[&borg::envcontract::Capability::Sandbox] + .state, + borg::envcontract::CheckState::Pass + ); +} + +#[tokio::test] +async fn draining_drones_do_not_start_capability_refreshes() { + let rig = Rig::new().await; + let drone = rig.drone(json!({})); + let mut tasks = JoinSet::new(); + drone.heartbeat_once(false, &mut tasks).await.unwrap(); + assert!(tasks.is_empty()); + assert!( + !rig.exec + .0 + .lock() + .unwrap() + .iter() + .any(|c| c.contains("--kill-after=2s")) + ); +} diff --git a/tests/explain.rs b/tests/explain.rs new file mode 100644 index 0000000..5779f74 --- /dev/null +++ b/tests/explain.rs @@ -0,0 +1,547 @@ +mod support; + +use borg::allocator::{ + ActiveLease, Allocation, AllocatorInput, Demand, DroneState, HostCaps, LeaseState, QueueAge, + Rejection, allocate, +}; +use borg::bid::ScoreWeights; +use borg::classes::{Class, Classes, Scope}; +use borg::explain::{Explanation, MAX_HISTORY, PlacementKind, PlacementOutcome, PlacementRecord}; +use borg::store::Store; +use serde_json::json; +use support::Harness; + +fn host(id: &str) -> DroneState { + DroneState::from_capacity(id, &support::capacity(), &HostCaps::default()) +} + +fn input(drones: Vec, count: u32) -> AllocatorInput { + AllocatorInput { + classes: Classes::default(), + demand: vec![Demand { + class: Class::Light, + scope: Scope::Main, + count, + }], + drones, + leases: vec![], + warm_floors: vec![], + queue_ages: vec![], + starvation_after_secs: None, + weights: ScoreWeights::default(), + admission: None, + } +} + +fn one() -> Allocation { + allocate(&input(vec![host("a")], 1)) +} + +#[test] +fn traces_actual_winner_and_order_with_bid_and_profile_terms() { + let mut a = host("a"); + a.bids = serde_json::from_value(json!([{"class":"light","scope":"main","available_slots":1,"cache_affinity":1000,"expires_at":100}])).unwrap(); + a.profiles = serde_json::from_value(json!([{"class":"light","speed_permille":1500,"speed_confidence_permille":1000,"reliability_permille":1000}])).unwrap(); + let result = allocate(&input(vec![a, host("b")], 2)); + assert_eq!( + result + .grants + .iter() + .map(|g| g.drone_id.as_str()) + .collect::>(), + vec!["a", "b"] + ); + assert_eq!( + result + .decisions + .iter() + .map(|d| d.winner.as_deref()) + .collect::>(), + vec![Some("a"), Some("b")] + ); + assert_eq!(result.decisions[1].grant_order, 1); + let first = &result.decisions[0].candidates[0]; + assert_eq!( + (first.bid_score, first.profile_score, first.score), + (3000, 500, Some(3500)) + ); + let second = &result.decisions[1].candidates[0]; + assert!(second.rejection_reasons.contains(&Rejection::BidExhausted)); + assert_eq!( + second.room.mem_mib, + first.room.mem_mib - result.grants[0].mem_mib + ); +} + +#[test] +fn traces_owner_net_capacity_blocked_environment_and_nonaccepting_hosts() { + let mut cap = support::capacity(); + cap.vcpu = 32; + cap.mem_mib = 64 * 1024; + let limited = DroneState::from_capacity( + "limited", + &cap, + &HostCaps { + max_vcpu: Some(1), + max_mem_mib: Some(1000), + max_leases: Some(0), + reserve_mem_mib: 2048, + }, + ); + let mut blocked = host("blocked"); + blocked.blocked_classes.push(Class::Light); + let mut quarantined = host("quarantined"); + quarantined.accepting = false; + let result = allocate(&input(vec![limited, blocked, quarantined], 1)); + assert!(result.grants.is_empty()); + let decision = &result.decisions[0]; + assert_eq!(decision.winner, None); + let find = |id| { + decision + .candidates + .iter() + .find(|c| c.drone_id == id) + .unwrap() + }; + assert!( + find("blocked") + .rejection_reasons + .contains(&Rejection::Environment) + ); + assert!( + find("quarantined") + .rejection_reasons + .contains(&Rejection::NotAccepting) + ); + assert_eq!( + (find("limited").budget.vcpu, find("limited").budget.mem_mib), + (1, 1000) + ); + assert_eq!( + find("limited").rejection_reasons, + vec![Rejection::Memory, Rejection::Vcpu, Rejection::LeaseCap] + ); +} + +#[test] +fn records_starvation_reservation_and_why_smaller_work_cannot_take_it() { + let mut drone = host("a"); + drone.mem_mib = 28 * 1024; + drone.vcpu = 16; + let mut i = input(vec![drone], 1); + i.demand.push(Demand { + class: Class::Heavy, + scope: Scope::Main, + count: 1, + }); + i.leases.push(ActiveLease { + drone_id: "a".into(), + class: Class::Medium, + scope: Scope::Main, + state: LeaseState::Running, + }); + i.queue_ages.push(QueueAge { + class: Class::Heavy, + scope: Scope::Main, + oldest_age_secs: 60, + }); + i.starvation_after_secs = Some(30); + let result = allocate(&i); + assert_eq!(result.decisions[0].reservation.as_deref(), Some("a")); + let light = result + .decisions + .iter() + .find(|d| d.class == Class::Light) + .unwrap(); + assert!( + light.candidates[0] + .rejection_reasons + .contains(&Rejection::Reserved) + ); + assert_eq!( + light.candidates[0].reserved_for, + Some((Class::Heavy, Scope::Main)) + ); +} + +#[test] +fn inherited_reservations_are_enforced_and_explained_at_the_selection_boundary() { + let i = input(vec![host("a")], 1); + let reserved = [borg::allocator::Reservation { + drone_id: "a".into(), + class: Class::Heavy, + scope: Scope::Main, + }]; + let result = borg::allocator::allocate_reserved(&i, &reserved); + assert!(result.grants.is_empty()); + assert_eq!( + result.decisions[0].candidates[0].rejection_reasons, + vec![Rejection::Reserved] + ); +} + +#[test] +fn fair_single_unit_selection_uses_the_real_backlog_for_bid_pressure() { + let mut drone = host("a"); + drone.bids = serde_json::from_value(json!([{"class":"light","scope":"main","available_slots":10,"min_pressure":500,"expires_at":100}])).unwrap(); + let i = input(vec![drone], 1); + assert!(allocate(&i).grants.is_empty()); + let backlog = [Demand { + class: Class::Light, + scope: Scope::Main, + count: 20, + }]; + let result = borg::allocator::allocate_reserved_with_backlog(&i, &[], &backlog); + assert_eq!( + result.grants.len(), + 1, + "backlog controls pressure, not the number selected by fairness" + ); + assert_eq!(result.decisions[0].pressure, 1000); + assert_eq!(result.decisions[0].winner.as_deref(), Some("a")); +} + +#[test] +fn starvation_never_reserves_a_host_that_declines_the_demand() { + let mut drone = host("a"); + drone.bids = serde_json::from_value( + json!([{"class":"heavy","scope":"main","available_slots":0,"expires_at":100}]), + ) + .unwrap(); + let mut i = input(vec![drone], 0); + i.demand.push(Demand { + class: Class::Heavy, + scope: Scope::Main, + count: 1, + }); + i.leases.push(ActiveLease { + drone_id: "a".into(), + class: Class::Light, + scope: Scope::Main, + state: LeaseState::Running, + }); + i.queue_ages.push(QueueAge { + class: Class::Heavy, + scope: Scope::Main, + oldest_age_secs: 60, + }); + i.starvation_after_secs = Some(30); + let result = allocate(&i); + assert!(result.reserved_for.is_empty()); + assert_eq!(result.decisions[0].reservation, None); + assert!( + result.decisions[0].candidates[0] + .rejection_reasons + .contains(&Rejection::BidExhausted) + ); +} + +#[test] +fn diagnostic_bounds_never_change_the_actual_winner_or_grant_count() { + let mut drones = (0..borg::allocator::MAX_CANDIDATES + 1) + .map(|i| host(&format!("host-{i:04}"))) + .collect::>(); + drones.last_mut().unwrap().profiles = serde_json::from_value(json!([{"class":"light","speed_permille":1500,"speed_confidence_permille":1000,"reliability_permille":1000}])).unwrap(); + let result = allocate(&input(drones, 1)); + assert_eq!( + result.decisions[0].total_candidates, + borg::allocator::MAX_CANDIDATES + 1 + ); + assert_eq!( + result.decisions[0].candidates.len(), + borg::allocator::MAX_CANDIDATES + ); + assert_eq!(result.decisions[0].winner.as_deref(), Some("host-0128")); + assert_eq!(result.grants[0].drone_id, "host-0128"); + assert!( + result.decisions[0] + .candidates + .iter() + .any(|c| c.drone_id == "host-0128") + ); + let mut drone = host("huge"); + drone.vcpu = 1000; + drone.mem_mib = 1_000_000; + let result = allocate(&input(vec![drone], 130)); + assert_eq!(result.grants.len(), 130); + assert_eq!(result.decisions.len(), borg::allocator::MAX_DECISIONS); + assert_eq!(result.omitted_decisions, 2); +} + +#[tokio::test] +async fn history_is_bounded_and_survives_reopening_the_owned_store() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("queen.db"); + let store = Store::open(&path).await.unwrap(); + let result = one(); + for n in 0..MAX_HISTORY + 5 { + let key = format!("label-{n}"); + store + .record_placement(PlacementRecord { + at: n, + kind: PlacementKind::Runner, + key: &key, + client: None, + outcome: PlacementOutcome::Granted, + decision: &result.decisions[0], + }) + .await + .unwrap(); + } + drop(store); + let reopened = Store::open(&path).await.unwrap(); + assert!( + reopened + .placement_history(PlacementKind::Runner, "label-4", None) + .await + .unwrap() + .is_empty() + ); + assert_eq!( + reopened + .placement_history(PlacementKind::Runner, "label-5", None) + .await + .unwrap() + .len(), + 1 + ); + let latest = reopened + .placement_history(PlacementKind::Runner, "label-1004", None) + .await + .unwrap(); + assert_eq!(latest[0].decision, result.decisions[0]); + assert_eq!(latest[0].outcome, PlacementOutcome::Granted); +} + +#[tokio::test] +async fn history_rejects_oversized_and_wrong_owner_records_without_pruning() { + let store = Store::open_in_memory().await.unwrap(); + let mut decision = one().decisions.remove(0); + decision.winner = Some("x".repeat(borg::explain::MAX_DECISION_BYTES)); + assert!( + store + .record_placement(PlacementRecord { + at: 0, + kind: PlacementKind::Runner, + key: "label", + client: None, + outcome: PlacementOutcome::Granted, + decision: &decision + }) + .await + .is_err() + ); + assert!( + store + .record_placement(PlacementRecord { + at: 0, + kind: PlacementKind::Unit, + key: "missing-unit", + client: Some("wrong"), + outcome: PlacementOutcome::Granted, + decision: &one().decisions[0] + }) + .await + .is_err() + ); + assert!( + store + .placement_history(PlacementKind::Runner, "label", None) + .await + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn api_requires_unit_owner_and_public_label_queries_never_include_units() { + let h = Harness::start(json!({"units":{"enabled":true}})).await; + for (name, token) in [("alice", "alice-token"), ("bob", "bob-token")] { + h.queen + .store() + .enroll_client(name, &borg::enroll::hash_token(token), 0) + .await + .unwrap(); + } + let req = h.http.post(format!("{}/v1/units", h.url)).bearer_auth("alice-token") + .json(&json!({"repo":"example-org/app","tree_sha":"0123456789abcdef0123456789abcdef01234567","argv":["private-command"],"platform":"linux-x86_64","class":"light"})).send().await.unwrap(); + assert_eq!(req.status(), 201); + let submitted: serde_json::Value = req.json().await.unwrap(); + let id = submitted["unit_id"].as_str().unwrap(); + let result = one(); + let store = h.queen.store(); + store + .record_placement(PlacementRecord { + at: 1, + kind: PlacementKind::Unit, + key: id, + client: Some("alice"), + outcome: PlacementOutcome::Granted, + decision: &result.decisions[0], + }) + .await + .unwrap(); + assert!( + store + .record_placement(PlacementRecord { + at: 1, + kind: PlacementKind::Unit, + key: id, + client: Some("bob"), + outcome: PlacementOutcome::Granted, + decision: &result.decisions[0] + }) + .await + .is_err() + ); + store + .record_placement(PlacementRecord { + at: 1, + kind: PlacementKind::Runner, + key: "borg-ci-light", + client: None, + outcome: PlacementOutcome::Granted, + decision: &result.decisions[0], + }) + .await + .unwrap(); + let url = format!("{}/explain", h.url); + assert_eq!( + h.http + .get(&url) + .query(&[("key", id)]) + .send() + .await + .unwrap() + .status(), + 401 + ); + let bob: Explanation = h + .http + .get(&url) + .query(&[("key", id)]) + .bearer_auth("bob-token") + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert!(bob.history.is_empty()); + let alice: Explanation = h + .http + .get(&url) + .query(&[("key", id)]) + .bearer_auth("alice-token") + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(alice.history.len(), 1); + assert_eq!(alice.history[0].key, id); + let public: Explanation = h + .http + .get(&url) + .query(&[("key", "borg-ci-light")]) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(public.history.len(), 1); + assert_eq!(public.history[0].kind, PlacementKind::Runner); + assert!(!serde_json::to_string(&public).unwrap().contains(id)); + assert!( + !serde_json::to_string(&alice) + .unwrap() + .contains("private-command") + ); + let missing: Explanation = h + .http + .get(&url) + .query(&[("key", "unknown-unit")]) + .bearer_auth("alice-token") + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert!(missing.history.is_empty()); + assert!(borg::explain::render_text(&missing).contains("no retained placement decision")); +} + +#[tokio::test] +async fn disabled_units_hide_retained_unit_history_but_keep_runner_explanations() { + let h = Harness::start(json!({"units":{"enabled":false}})).await; + let store = h.queen.store(); + store + .enroll_client("alice", &borg::enroll::hash_token("alice-token"), 0) + .await + .unwrap(); + // Persist directly to model history retained after an operator disables units. + let unit: borg::units::SubmitUnit = serde_json::from_value(json!({ + "repo":"example-org/app", + "tree_sha":"0123456789abcdef0123456789abcdef01234567", + "argv":["private-command"], + "platform":"linux-x86_64", + "class":"light" + })) + .unwrap(); + unit.validate().unwrap(); + store + .submit_unit(&unit, "alice", "retained-unit", 3600, 3600, 0) + .await + .unwrap(); + let allocation = one(); + for (kind, key, client) in [ + (PlacementKind::Unit, "retained-unit", Some("alice")), + (PlacementKind::Runner, "borg-ci-light", None), + ] { + store + .record_placement(PlacementRecord { + at: 1, + kind, + key, + client, + outcome: PlacementOutcome::Granted, + decision: &allocation.decisions[0], + }) + .await + .unwrap(); + } + assert_eq!( + store + .placement_history(PlacementKind::Unit, "retained-unit", Some("alice")) + .await + .unwrap() + .len(), + 1 + ); + for authenticated in [false, true] { + let request = h + .http + .get(format!("{}/explain", h.url)) + .query(&[("key", "retained-unit")]); + let request = if authenticated { + request.bearer_auth("alice-token") + } else { + request + }; + assert_eq!(request.send().await.unwrap().status(), 404); + } + let response = h + .http + .get(format!("{}/explain", h.url)) + .query(&[("key", "borg-ci-light")]) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 200); + let runner: Explanation = response.json().await.unwrap(); + assert_eq!(runner.history.len(), 1); + assert_eq!(runner.history[0].kind, PlacementKind::Runner); +} diff --git a/tests/explain_cli.rs b/tests/explain_cli.rs new file mode 100644 index 0000000..3f178a8 --- /dev/null +++ b/tests/explain_cli.rs @@ -0,0 +1,224 @@ +//! Exercise the public CLI's local and authenticated HTTP history reads. +mod support; + +use std::path::Path; +use std::process::Output; + +use borg::allocator::{AllocatorInput, Demand, DroneState, HostCaps, allocate}; +use borg::bid::ScoreWeights; +use borg::classes::{Class, Classes, Scope}; +use borg::explain::{Explanation, PlacementKind, PlacementOutcome, PlacementRecord}; +use borg::store::Store; +use serde_json::json; +use support::Harness; + +fn borg_bin() -> String { + option_env!("CARGO_BIN_EXE_borg") + .map(str::to_owned) + .unwrap_or_else(|| { + std::env::current_exe() + .unwrap() + .parent() + .and_then(Path::parent) + .unwrap() + .join("borg") + .to_string_lossy() + .into_owned() + }) +} + +async fn run(args: &[&str], home: &Path) -> Output { + #[allow( + clippy::disallowed_methods, + reason = "CLI boundary test: execute the compiled borg binary with isolated credential sources" + )] + tokio::process::Command::new(borg_bin()) + .args(args) + .env("HOME", home) + .env_remove("BORG_CLIENT_TOKEN") + .env_remove("BORG_CLIENT_TOKEN_FILE") + .output() + .await + .unwrap() +} + +async fn record(store: &Store, kind: PlacementKind, key: &str, client: Option<&str>) { + let input = AllocatorInput { + classes: Classes::default(), + demand: vec![Demand { + class: Class::Light, + scope: Scope::Main, + count: 1, + }], + drones: vec![DroneState::from_capacity( + "host-a", + &support::capacity(), + &HostCaps::default(), + )], + leases: vec![], + warm_floors: vec![], + queue_ages: vec![], + starvation_after_secs: None, + weights: ScoreWeights::default(), + admission: None, + }; + store + .record_placement(PlacementRecord { + at: 123, + kind, + key, + client, + outcome: PlacementOutcome::Granted, + decision: &allocate(&input).decisions[0], + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn local_config_resolves_aliases_and_renders_real_typed_history() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("queen.db"); + let cfg = dir.path().join("queen.json"); + std::fs::write( + &cfg, + json!({"github":{"org":"example-org","app_id":"1"},"repos":["example-org/app"], + "db_path":path,"classes":{"aliases":{"build-light":"light"}}}) + .to_string(), + ) + .unwrap(); + let store = Store::open(&path).await.unwrap(); + record(&store, PlacementKind::Runner, "borg-ci-light", None).await; + let output = run( + &[ + "explain", + "build-light", + "--config", + cfg.to_str().unwrap(), + "--json", + ], + dir.path(), + ) + .await; + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let report: Explanation = serde_json::from_slice(&output.stdout).unwrap(); + assert_eq!(report.key, "build-light"); + assert_eq!(report.history.len(), 1); + assert_eq!(report.history[0].key, "borg-ci-light"); + assert_eq!(report.history[0].decision.winner.as_deref(), Some("host-a")); + let text = run( + &["explain", "build-light", "--config", cfg.to_str().unwrap()], + dir.path(), + ) + .await; + assert!(text.status.success()); + let text = String::from_utf8(text.stdout).unwrap(); + assert!(text.contains("selected host-a"), "{text}"); + assert!(text.contains("bid 0, profile 0"), "{text}"); + let unknown = run( + &["explain", "unknown-unit", "--config", cfg.to_str().unwrap()], + dir.path(), + ) + .await; + assert!(unknown.status.success()); + assert!( + String::from_utf8(unknown.stdout) + .unwrap() + .contains("no retained placement decision") + ); +} + +#[tokio::test] +async fn remote_unit_cli_uses_token_file_and_does_not_disclose_another_clients_history() { + let h = Harness::start(json!({"units":{"enabled":true}})).await; + for (name, token) in [ + ("alice", "alice-private-token"), + ("bob", "bob-private-token"), + ] { + h.queen + .store() + .enroll_client(name, &borg::enroll::hash_token(token), 0) + .await + .unwrap(); + } + let response = h.http.post(format!("{}/v1/units", h.url)).bearer_auth("alice-private-token") + .json(&json!({"repo":"example-org/app","tree_sha":"0123456789abcdef0123456789abcdef01234567", "argv":["do-private-work"],"platform":"linux-x86_64","class":"light"})).send().await.unwrap(); + assert_eq!(response.status(), 201); + let submitted: serde_json::Value = response.json().await.unwrap(); + let id = submitted["unit_id"].as_str().unwrap(); + record(h.queen.store(), PlacementKind::Unit, id, Some("alice")).await; + record( + h.queen.store(), + PlacementKind::Runner, + "borg-ci-light", + None, + ) + .await; + let dir = tempfile::tempdir().unwrap(); + let token = dir.path().join("client-token"); + std::fs::write(&token, "alice-private-token\n").unwrap(); + let args = [ + "explain", + id, + "--queen", + &h.url, + "--token-file", + token.to_str().unwrap(), + "--json", + ]; + let alice = run(&args, dir.path()).await; + assert!( + alice.status.success(), + "{}", + String::from_utf8_lossy(&alice.stderr) + ); + let report: Explanation = serde_json::from_slice(&alice.stdout).unwrap(); + assert_eq!(report.history.len(), 1); + assert_eq!(report.history[0].key, id); + assert!(!String::from_utf8_lossy(&alice.stdout).contains("do-private-work")); + assert!(!String::from_utf8_lossy(&alice.stderr).contains("alice-private-token")); + std::fs::write(&token, "bob-private-token\n").unwrap(); + let bob = run(&args, dir.path()).await; + assert!(bob.status.success()); + let report: Explanation = serde_json::from_slice(&bob.stdout).unwrap(); + assert!(report.history.is_empty()); + let anonymous = run(&["explain", id, "--queen", &h.url, "--json"], dir.path()).await; + assert!(!anonymous.status.success()); + assert!(anonymous.stdout.is_empty()); + let public = run( + &["explain", "borg-ci-light", "--queen", &h.url, "--json"], + dir.path(), + ) + .await; + assert!(public.status.success()); + let report: Explanation = serde_json::from_slice(&public.stdout).unwrap(); + assert_eq!(report.history.len(), 1); + assert_eq!(report.history[0].kind, PlacementKind::Runner); +} + +#[tokio::test] +async fn cli_rejects_missing_or_conflicting_sources() { + let dir = tempfile::tempdir().unwrap(); + let missing = run(&["explain", "borg-ci-light"], dir.path()).await; + assert!(!missing.status.success()); + assert!(missing.stdout.is_empty()); + let conflict = run( + &[ + "explain", + "borg-ci-light", + "--queen", + "http://127.0.0.1:1", + "--config", + "queen.json", + ], + dir.path(), + ) + .await; + assert!(!conflict.status.success()); + assert!(conflict.stdout.is_empty()); + assert!(String::from_utf8_lossy(&conflict.stderr).contains("cannot be used with")); +} diff --git a/tests/placement_batch.rs b/tests/placement_batch.rs new file mode 100644 index 0000000..d8797a1 --- /dev/null +++ b/tests/placement_batch.rs @@ -0,0 +1,238 @@ +mod support; + +use borg::allocator::{AllocatorInput, Demand, DroneState, HostCaps, PlacementDecision, allocate}; +use borg::bid::ScoreWeights; +use borg::classes::{Class, Classes, Scope}; +use borg::explain::{ + MAX_DECISION_BYTES, MAX_HISTORY, PlacementKind, PlacementOutcome, PlacementRecord, +}; +use borg::store::Store; +use serde_json::json; + +fn decision() -> PlacementDecision { + allocate(&AllocatorInput { + classes: Classes::default(), + demand: vec![Demand { + class: Class::Light, + scope: Scope::Main, + count: 1, + }], + drones: vec![DroneState::from_capacity( + "host", + &support::capacity(), + &HostCaps::default(), + )], + leases: vec![], + warm_floors: vec![], + queue_ages: vec![], + starvation_after_secs: None, + weights: ScoreWeights::default(), + admission: None, + }) + .decisions + .remove(0) +} + +fn runner<'a>(key: &'a str, at: i64, decision: &'a PlacementDecision) -> PlacementRecord<'a> { + PlacementRecord { + at, + kind: PlacementKind::Runner, + key, + client: None, + outcome: PlacementOutcome::Granted, + decision, + } +} + +#[tokio::test] +async fn batch_rejects_wrong_authoritative_unit_owner_without_partial_inserts() { + let store = Store::open_in_memory().await.unwrap(); + let decision = decision(); + let unit = serde_json::from_value(json!({ + "repo":"example-org/app", "tree_sha":"0123456789abcdef0123456789abcdef01234567", + "argv":["cargo","check"], "platform":"linux-x86_64", "toolchain":"rust-1.95", + "class":"light" + })) + .unwrap(); + store + .submit_unit(&unit, "alice", "unit", 3600, 3600, 0) + .await + .unwrap(); + store + .record_placement(runner("sentinel", 0, &decision)) + .await + .unwrap(); + let bad = PlacementRecord { + at: 1, + kind: PlacementKind::Unit, + key: "unit", + client: Some("bob"), + outcome: PlacementOutcome::Unplaced, + decision: &decision, + }; + assert!( + store + .record_placements(&[runner("partial", 1, &decision), bad]) + .await + .is_err() + ); + assert!( + store + .placement_history(PlacementKind::Runner, "partial", None) + .await + .unwrap() + .is_empty() + ); + assert!( + store + .placement_history(PlacementKind::Unit, "unit", Some("alice")) + .await + .unwrap() + .is_empty() + ); + assert_eq!( + store + .placement_history(PlacementKind::Runner, "sentinel", None) + .await + .unwrap() + .len(), + 1 + ); + store + .record_placements(&[ + runner("accepted", 2, &decision), + PlacementRecord { + at: 2, + kind: PlacementKind::Unit, + key: "unit", + client: Some("alice"), + outcome: PlacementOutcome::Unplaced, + decision: &decision, + }, + ]) + .await + .unwrap(); + let unit_history = store + .placement_history(PlacementKind::Unit, "unit", Some("alice")) + .await + .unwrap(); + assert_eq!(unit_history.len(), 1); + assert_eq!(unit_history[0].outcome, PlacementOutcome::Unplaced); + assert_eq!( + store + .placement_history(PlacementKind::Runner, "accepted", None) + .await + .unwrap() + .len(), + 1 + ); +} + +#[tokio::test] +async fn invalid_batch_bounds_preserve_existing_history_and_insert_nothing() { + let store = Store::open_in_memory().await.unwrap(); + let decision = decision(); + store + .record_placement(runner("sentinel", 0, &decision)) + .await + .unwrap(); + let mut oversized = decision.clone(); + oversized.winner = Some("x".repeat(MAX_DECISION_BYTES)); + assert!( + store + .record_placements(&[ + runner("partial", 1, &decision), + runner("oversized", 1, &oversized), + ]) + .await + .is_err() + ); + let mut too_many_candidates = decision.clone(); + too_many_candidates.candidates = + vec![decision.candidates[0].clone(); borg::allocator::MAX_CANDIDATES + 1]; + assert!( + store + .record_placements(&[ + runner("partial", 1, &decision), + runner("candidates", 1, &too_many_candidates), + ]) + .await + .is_err() + ); + for bad_key in [String::new(), "k".repeat(257)] { + assert!( + store + .record_placements(&[ + runner("partial", 1, &decision), + runner(&bad_key, 1, &decision), + ]) + .await + .is_err() + ); + } + let too_many_rows: Vec<_> = (0..=MAX_HISTORY) + .map(|at| runner("partial", at, &decision)) + .collect(); + assert!(store.record_placements(&too_many_rows).await.is_err()); + assert!( + store + .placement_history(PlacementKind::Runner, "partial", None) + .await + .unwrap() + .is_empty() + ); + assert_eq!( + store + .placement_history(PlacementKind::Runner, "sentinel", None) + .await + .unwrap() + .len(), + 1 + ); +} + +#[tokio::test] +async fn batch_retains_newest_records_at_the_global_cap_across_reopening() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("queen.db"); + let store = Store::open(&path).await.unwrap(); + let decision = decision(); + store + .record_placement(runner("old", 0, &decision)) + .await + .unwrap(); + let records: Vec<_> = (1..=MAX_HISTORY) + .map(|at| runner("batch", at, &decision)) + .collect(); + store.record_placements(&records).await.unwrap(); + drop(store); + let store = Store::open(&path).await.unwrap(); + assert!( + store + .placement_history(PlacementKind::Runner, "old", None) + .await + .unwrap() + .is_empty() + ); + let latest = store + .placement_history(PlacementKind::Runner, "batch", None) + .await + .unwrap(); + assert_eq!(latest.len(), borg::explain::MAX_QUERY_ROWS as usize); + assert_eq!(latest[0].at, MAX_HISTORY); + assert_eq!(latest[0].decision, decision); + assert!(latest.windows(2).all(|pair| pair[0].at == pair[1].at + 1)); + // Check total retention through the same SQLite file, independently of the + // public read limit, rather than adding test-only store accessors. + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .connect_with(sqlx::sqlite::SqliteConnectOptions::new().filename(&path)) + .await + .unwrap(); + let (count, oldest): (i64, i64) = + // sqlx-runtime-query: integration assertion on bounded SQLite diagnostics + sqlx::query_as("SELECT COUNT(*), MIN(at) FROM placement_history") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!((count, oldest), (MAX_HISTORY, 1)); +} diff --git a/tests/run.rs b/tests/run.rs index 07160c7..6c4fe5d 100644 --- a/tests/run.rs +++ b/tests/run.rs @@ -18,9 +18,19 @@ use serde_json::{Value, json}; #[derive(Clone, Copy, PartialEq)] enum Mode { Receipt, + BadReusedDigest, Run, Failed, Queued, + Joined, + LostLog, + JoinedLostLog, + JoinedRun, + MultiChunk, + RunningForever, + Unauthorized, + BadTerminalReceipt, + BadTerminalDigest, Refuse, /// Runs like `Run` but rejects the bundle upload. RefuseBundle, @@ -31,19 +41,44 @@ struct Fake { submitted: Mutex>, cancelled: Mutex, bundles: Mutex>>, + log_reads: Mutex, } type Shared = Arc; +fn receipt_for(command: &Value) -> borg::units::Receipt { + let mut receipt = borg::units::Receipt { + unit_id: "u-1".into(), + repo: command["repo"].as_str().unwrap().into(), + tree_sha: command["tree_sha"].as_str().unwrap().into(), + argv: serde_json::from_value(command["argv"].clone()).unwrap(), + platform: command["platform"].as_str().unwrap().into(), + toolchain: command["toolchain"].as_str().unwrap().into(), + drone: "d9".into(), + exit_code: 0, + finished_at: 123, + duration_secs: 2, + digest: String::new(), + }; + receipt.digest = receipt.compute_digest(); + receipt +} + async fn submit(State(f): State, Json(b): Json) -> Response { - f.submitted.lock().unwrap().push(b); + f.submitted.lock().unwrap().push(b.clone()); match f.mode { + Mode::Joined | Mode::JoinedLostLog | Mode::JoinedRun => { + Json(json!({"outcome": "joined", "unit_id": "u-1", "state": "queued"})).into_response() + } Mode::Refuse => StatusCode::INTERNAL_SERVER_ERROR.into_response(), - Mode::Receipt => Json( - json!({"outcome": "receipt", "unit_id": "u-1", "state": "done", - "receipt": {"drone": "d9", "exit_code": 0}}), - ) - .into_response(), + Mode::Unauthorized => StatusCode::UNAUTHORIZED.into_response(), + Mode::Receipt | Mode::BadReusedDigest => { + let mut receipt = receipt_for(&b); + if f.mode == Mode::BadReusedDigest { + receipt.digest = "tampered".into(); + } + Json(json!({"outcome": "receipt", "unit_id": "u-1", "state": "done", "receipt": receipt})).into_response() + } _ => ( StatusCode::CREATED, Json(json!({"outcome": "new", "unit_id": "u-1", "state": "queued"})), @@ -52,9 +87,80 @@ async fn submit(State(f): State, Json(b): Json) -> Response { } } +async fn show(State(f): State, AxPath(_id): AxPath) -> Json { + let mut unit = f.submitted.lock().unwrap().last().unwrap().clone(); + unit["id"] = json!("u-1"); + unit["drone"] = json!("d9"); + unit["timeout_secs"] = json!(if f.mode == Mode::RunningForever { + 1 + } else { + 30 + }); + unit["state"] = json!(match f.mode { + Mode::Run + | Mode::RefuseBundle + | Mode::BadTerminalReceipt + | Mode::BadTerminalDigest + | Mode::JoinedRun + | Mode::MultiChunk => "done", + Mode::RunningForever | Mode::LostLog | Mode::JoinedLostLog => "running", + Mode::Failed => "failed", + _ => "queued", + }); + unit["exit_code"] = json!(if matches!( + f.mode, + Mode::BadTerminalReceipt | Mode::BadTerminalDigest + ) { + 0 + } else { + 3 + }); + if f.mode == Mode::BadTerminalReceipt { + unit["receipt"] = json!({"unit_id": "u-1", "repo": "other/repo", + "tree_sha": unit["tree_sha"], "argv": unit["argv"], + "platform": unit["platform"], "toolchain": unit["toolchain"], + "drone": "d9", "exit_code": 0, "finished_at": 123, + "duration_secs": 1, "digest": "tampered"}); + } + if f.mode == Mode::BadTerminalDigest { + let mut receipt = receipt_for(&unit); + receipt.digest = "tampered".into(); + unit["receipt"] = serde_json::to_value(receipt).unwrap(); + } + Json(unit) +} + async fn log(State(f): State, AxPath(_id): AxPath) -> Response { + if matches!(f.mode, Mode::LostLog | Mode::JoinedLostLog) { + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + if f.mode == Mode::MultiChunk { + let mut reads = f.log_reads.lock().unwrap(); + *reads += 1; + let last = *reads > 1; + let mut response = if last { + "second chunk\n" + } else { + "first chunk\n" + } + .into_response(); + let headers = response.headers_mut(); + headers.insert("x-borg-state", "done".parse().unwrap()); + headers.insert( + "x-borg-next", + if last { "25" } else { "12" }.parse().unwrap(), + ); + if last { + headers.insert("x-borg-exit", "3".parse().unwrap()); + } + return response; + } let mut r = match f.mode { - Mode::Run | Mode::RefuseBundle => "hello from the fleet\n".into_response(), + Mode::Run + | Mode::RefuseBundle + | Mode::BadTerminalReceipt + | Mode::BadTerminalDigest + | Mode::JoinedRun => "hello from the fleet\n".into_response(), _ => { tokio::time::sleep(std::time::Duration::from_millis(150)).await; "".into_response() @@ -62,14 +168,35 @@ async fn log(State(f): State, AxPath(_id): AxPath) -> Response { }; let h = r.headers_mut(); let (state, next) = match f.mode { - Mode::Run | Mode::RefuseBundle => ("done", "21"), + Mode::Run + | Mode::RefuseBundle + | Mode::BadTerminalReceipt + | Mode::BadTerminalDigest + | Mode::JoinedRun => ("done", "21"), Mode::Failed => ("failed", "0"), + Mode::RunningForever => ("running", "0"), _ => ("queued", "0"), }; h.insert("x-borg-state", state.parse().unwrap()); h.insert("x-borg-next", next.parse().unwrap()); - if matches!(f.mode, Mode::Run | Mode::RefuseBundle) { - h.insert("x-borg-exit", "3".parse().unwrap()); + if matches!( + f.mode, + Mode::Run + | Mode::RefuseBundle + | Mode::BadTerminalReceipt + | Mode::BadTerminalDigest + | Mode::JoinedRun + ) { + h.insert( + "x-borg-exit", + if matches!(f.mode, Mode::BadTerminalReceipt | Mode::BadTerminalDigest) { + "0" + } else { + "3" + } + .parse() + .unwrap(), + ); } if f.mode == Mode::Failed { h.insert("x-borg-failure", "timed out".parse().unwrap()); @@ -100,9 +227,11 @@ async fn fake_queen(mode: Mode) -> (String, Shared) { submitted: Mutex::new(Vec::new()), cancelled: Mutex::new(0), bundles: Mutex::new(Vec::new()), + log_reads: Mutex::new(0), }); let app = Router::new() .route("/v1/units", post(submit)) + .route("/v1/units/{id}", get(show)) .route("/v1/units/{id}/log", get(log)) .route("/v1/units/{id}/bundle", axum::routing::put(bundle)) .route("/v1/units/{id}/cancel", post(cancel)) @@ -193,6 +322,8 @@ fn opts(r: &Repo, queen: Option, marker: &str) -> RunOptions { scope: "pr".into(), wait_secs: 1, publish: false, + remote_only: false, + result_file: None, argv: vec!["sh".into(), "-c".into(), format!("touch {marker}; exit 7")], } } @@ -448,3 +579,328 @@ async fn the_submission_carries_the_trees_toolchain() { "1.95.0" ); } + +#[tokio::test] +async fn remote_only_never_runs_the_local_command_on_infrastructure_failures() { + for mode in [ + Mode::Refuse, + Mode::Unauthorized, + Mode::Failed, + Mode::Queued, + Mode::LostLog, + ] { + let r = repo(); + let (url, _) = fake_queen(mode).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = true; + o.wait_secs = 0; + o.result_file = Some(r.dir.join("result.json")); + assert_eq!(run(&o, &mut Vec::new()).await, 125); + assert!(!Path::new(&marker(&r)).exists()); + let result: Value = + serde_json::from_slice(&std::fs::read(o.result_file.unwrap()).unwrap()).unwrap(); + assert_eq!(result["exit_code"], 125); + assert!(result["infrastructure_error"].is_string()); + } +} + +#[tokio::test] +async fn remote_only_requires_configuration_a_resolvable_published_tree_and_a_reachable_queen() { + let r = repo(); + let mut o = opts(&r, None, &marker(&r)); + o.remote_only = true; + assert_eq!(run(&o, &mut Vec::new()).await, 125); + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let dead = format!("http://{}", listener.local_addr().unwrap()); + drop(listener); + o.queen = Some(dead); + o.token = Some(borg::secret::Secret::new("token")); + assert_eq!(run(&o, &mut Vec::new()).await, 125); + o.tree = TreeSource::Explicit("missing-tree".into()); + assert_eq!(run(&o, &mut Vec::new()).await, 125); + std::fs::write(r.dir.join("a.txt"), "staged\n").unwrap(); + git(&r.dir, &["add", "a.txt"]); + o.tree = TreeSource::Index; + assert_eq!(run(&o, &mut Vec::new()).await, 125); + let (url, fake) = fake_queen(Mode::RefuseBundle).await; + o.queen = Some(url); + o.publish = true; + assert_eq!(run(&o, &mut Vec::new()).await, 125); + assert_eq!(*fake.cancelled.lock().unwrap(), 1); + assert!(!Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn leaving_a_joined_unit_does_not_cancel_another_subscribers_work() { + for remote_only in [false, true] { + let r = repo(); + let (url, fake) = fake_queen(Mode::Joined).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = remote_only; + o.wait_secs = 0; + assert_eq!( + run(&o, &mut Vec::new()).await, + if remote_only { 125 } else { 7 } + ); + assert_eq!(*fake.cancelled.lock().unwrap(), 0); + assert_eq!(Path::new(&marker(&r)).exists(), !remote_only); + } +} + +#[tokio::test] +async fn result_files_identify_remote_execution_and_reused_receipts_without_mixing_stdout() { + for mode in [Mode::Run, Mode::Receipt] { + let r = repo(); + let (url, _) = fake_queen(mode).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = true; + o.result_file = Some(r.dir.join("result.json")); + let mut out = Vec::new(); + assert_eq!( + run(&o, &mut out).await, + if mode == Mode::Receipt { 0 } else { 3 } + ); + let result: Value = + serde_json::from_slice(&std::fs::read(o.result_file.unwrap()).unwrap()).unwrap(); + assert_eq!(result["schema_version"], 1); + assert_eq!(result["execution"], "remote"); + assert_eq!(result["host"], "d9"); + assert_eq!(result["unit_id"], "u-1"); + assert_eq!( + result["requested_tree_sha"], + resolve_tree(&r.dir, &TreeSource::Head).unwrap() + ); + assert_eq!(result["argv"], json!(o.argv)); + assert_eq!( + result["provenance"], + if mode == Mode::Receipt { + "reused" + } else { + "new" + } + ); + if mode == Mode::Receipt { + assert_eq!(out, b""); + assert_eq!(result["receipt"]["drone"], "d9"); + } else { + assert_eq!(out, b"hello from the fleet\n"); + } + assert!(!Path::new(&marker(&r)).exists()); + } +} + +#[tokio::test] +async fn a_local_result_does_not_claim_execution_against_a_git_tree() { + let r = repo(); + let mut o = opts(&r, None, &marker(&r)); + o.result_file = Some(r.dir.join("result.json")); + assert_eq!(run(&o, &mut Vec::new()).await, 7); + let result: Value = + serde_json::from_slice(&std::fs::read(o.result_file.unwrap()).unwrap()).unwrap(); + assert_eq!(result["execution"], "local"); + assert_eq!(result["provenance"], "fallback"); + assert!(result["host"].is_null()); + assert!(result["receipt"].is_null()); +} + +#[test] +fn the_cli_remote_only_flag_prevents_local_execution_and_emits_a_result() { + let r = repo(); + let result_path = r.dir.join("result.json"); + #[allow( + clippy::disallowed_methods, + reason = "test invokes the Borg CLI with a local sentinel" + )] + let output = Command::new(env!("CARGO_BIN_EXE_borg")) + .current_dir(&r.dir) + .env_remove("BORG_QUEEN") + .env_remove("BORG_RUN_REMOTE_ONLY") + .args(["run", "--remote-only", "--result-file"]) + .arg(&result_path) + .args(["--", "sh", "-c", "touch ran-locally"]) + .output() + .unwrap(); + assert_eq!(output.status.code(), Some(125)); + assert!(output.stdout.is_empty()); + assert!(!Path::new(&marker(&r)).exists()); + let result: Value = serde_json::from_slice(&std::fs::read(result_path).unwrap()).unwrap(); + assert_eq!(result["execution"], "none"); + assert_eq!(result["exit_code"], 125); +} + +#[tokio::test] +async fn remote_only_enforces_the_admitted_command_deadline() { + let r = repo(); + let (url, fake) = fake_queen(Mode::RunningForever).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = true; + let started = std::time::Instant::now(); + assert_eq!(run(&o, &mut Vec::new()).await, 125); + assert!(started.elapsed() < std::time::Duration::from_secs(5)); + assert_eq!(*fake.cancelled.lock().unwrap(), 1); + assert!(!Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn broken_output_consumers_never_rerun_remote_commands_locally() { + struct BrokenOutput { + fail_flush: bool, + } + impl std::io::Write for BrokenOutput { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if self.fail_flush { + Ok(bytes.len()) + } else { + Err(std::io::Error::other("closed log consumer")) + } + } + fn flush(&mut self) -> std::io::Result<()> { + if self.fail_flush { + Err(std::io::Error::other("cannot flush log consumer")) + } else { + Ok(()) + } + } + } + for remote_only in [false, true] { + for fail_flush in [false, true] { + for mode in [Mode::Run, Mode::JoinedRun, Mode::MultiChunk] { + let r = repo(); + let (url, fake) = fake_queen(mode).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = remote_only; + o.result_file = Some(r.dir.join("result.json")); + assert_eq!( + run(&o, &mut BrokenOutput { fail_flush }).await, + if remote_only { 125 } else { 3 } + ); + assert_eq!( + *fake.cancelled.lock().unwrap(), + u32::from(remote_only && mode != Mode::JoinedRun) + ); + assert!(!Path::new(&marker(&r)).exists()); + if mode == Mode::MultiChunk && !remote_only { + assert_eq!(*fake.log_reads.lock().unwrap(), 2); + } + let result: Value = + serde_json::from_slice(&std::fs::read(o.result_file.unwrap()).unwrap()) + .unwrap(); + assert_eq!(result["execution"], "remote"); + assert!( + result["infrastructure_error"] + .as_str() + .unwrap() + .contains("remote log") + ); + } + } + } +} + +#[tokio::test] +async fn an_unwritable_result_destination_prevents_execution() { + let r = repo(); + let mut o = opts(&r, None, &marker(&r)); + o.result_file = Some(r.dir.join("missing-parent/result.json")); + assert_eq!(run(&o, &mut Vec::new()).await, 125); + assert!(!Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn terminal_metadata_never_publishes_a_tampered_receipt_or_reexecutes_the_command() { + let r = repo(); + let (url, _) = fake_queen(Mode::BadTerminalReceipt).await; + let mut o = opts(&r, Some(url), &marker(&r)); + // Default execution can finish; optional metadata validation must not cause + // a second local execution after the remote command has completed. + o.result_file = Some(r.dir.join("result.json")); + assert_eq!(run(&o, &mut Vec::new()).await, 0); + let result: Value = + serde_json::from_slice(&std::fs::read(o.result_file.unwrap()).unwrap()).unwrap(); + assert!(result["receipt"].is_null()); + assert!(result["host"].is_null()); + assert!( + result["infrastructure_error"] + .as_str() + .unwrap() + .contains("receipt") + ); + assert!(!Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn completed_units_drain_all_log_chunks_before_returning_the_exit_code() { + for remote_only in [false, true] { + let r = repo(); + let (url, fake) = fake_queen(Mode::MultiChunk).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = remote_only; + let mut out = Vec::new(); + assert_eq!(run(&o, &mut out).await, 3); + assert_eq!(out, b"first chunk\nsecond chunk\n"); + assert_eq!(*fake.log_reads.lock().unwrap(), 2); + assert_eq!(*fake.cancelled.lock().unwrap(), 0); + assert!(!Path::new(&marker(&r)).exists()); + } +} + +#[tokio::test] +async fn losing_remote_logs_cancels_only_a_newly_created_unit() { + for mode in [Mode::LostLog, Mode::JoinedLostLog] { + let r = repo(); + let (url, fake) = fake_queen(mode).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = true; + o.wait_secs = 0; + assert_eq!(run(&o, &mut Vec::new()).await, 125); + assert_eq!( + *fake.cancelled.lock().unwrap(), + u32::from(mode == Mode::LostLog) + ); + assert!(!Path::new(&marker(&r)).exists()); + } +} + +#[tokio::test] +async fn a_matching_terminal_receipt_with_only_its_digest_corrupted_is_rejected() { + let r = repo(); + let (url, _) = fake_queen(Mode::BadTerminalDigest).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.result_file = Some(r.dir.join("result.json")); + assert_eq!(run(&o, &mut Vec::new()).await, 0); + let result: Value = + serde_json::from_slice(&std::fs::read(o.result_file.unwrap()).unwrap()).unwrap(); + assert_eq!(result["execution"], "remote"); + assert_eq!(result["exit_code"], 0); + assert!(result["receipt"].is_null()); + assert!(result["host"].is_null()); + assert!( + result["infrastructure_error"] + .as_str() + .unwrap() + .contains("receipt") + ); + assert!(!Path::new(&marker(&r)).exists()); +} + +#[tokio::test] +async fn a_matching_reused_receipt_with_only_its_digest_corrupted_is_rejected() { + let r = repo(); + let (url, _) = fake_queen(Mode::BadReusedDigest).await; + let mut o = opts(&r, Some(url), &marker(&r)); + o.remote_only = true; + o.result_file = Some(r.dir.join("result.json")); + assert_eq!(run(&o, &mut Vec::new()).await, 125); + let result: Value = + serde_json::from_slice(&std::fs::read(o.result_file.unwrap()).unwrap()).unwrap(); + assert_eq!(result["exit_code"], 125); + assert!(result["receipt"].is_null()); + assert!(result["host"].is_null()); + assert!( + result["infrastructure_error"] + .as_str() + .unwrap() + .contains("receipt") + ); + assert!(!Path::new(&marker(&r)).exists()); +} diff --git a/tests/sccache_slots.rs b/tests/sccache_slots.rs index ce322b5..48b34e2 100644 --- a/tests/sccache_slots.rs +++ b/tests/sccache_slots.rs @@ -6,7 +6,26 @@ use std::fs; use std::path::Path; use borg::drone::caches::{rustflags, sccache}; -use borg::drone::exec::{Exec, SystemExec}; +use borg::drone::exec::{Exec, Output, SystemExec}; + +/// Keep each fixture's local cache independent of runner backend settings. +struct FixtureExec; + +impl Exec for FixtureExec { + fn run(&self, program: &str, args: &[&str]) -> std::io::Result { + let mut argv = Vec::new(); + for name in std::env::vars_os() + .filter_map(|(name, _)| name.into_string().ok()) + .filter(|name| name.starts_with("SCCACHE_")) + { + argv.extend(["-u".to_owned(), name]); + } + argv.push(program.to_owned()); + argv.extend(args.iter().map(|arg| (*arg).to_owned())); + let argv: Vec<&str> = argv.iter().map(String::as_str).collect(); + SystemExec.run("env", &argv) + } +} fn available(program: &str) -> bool { SystemExec @@ -28,7 +47,7 @@ fn run_with(env: &[(String, String)], program: &str, args: &[&str]) -> bool { argv.push(program.to_owned()); argv.extend(args.iter().map(|a| (*a).to_owned())); let argv: Vec<&str> = argv.iter().map(String::as_str).collect(); - SystemExec.run("env", &argv).is_ok_and(|o| o.success) + FixtureExec.run("env", &argv).is_ok_and(|o| o.success) } struct TestServer { @@ -37,7 +56,19 @@ struct TestServer { impl TestServer { fn start(root: &Path) -> Self { - let env = sccache::server_env(&root.join("cache").to_string_lossy(), "1G", free_port()); + let mut env = sccache::server_env(&root.join("cache").to_string_lossy(), "1G", free_port()); + let config = root.join("config"); + fs::write(&config, "").unwrap(); + env.extend([ + ( + "SCCACHE_CONF".to_owned(), + config.to_string_lossy().into_owned(), + ), + ( + "SCCACHE_CACHED_CONF".to_owned(), + root.join("cached-config").to_string_lossy().into_owned(), + ), + ]); // sccache gives an inherited Unix socket precedence over the TCP port. // Keep start, compiler, stats and cleanup commands on our own endpoint. #[cfg(unix)] @@ -51,7 +82,7 @@ impl TestServer { }; let server = Self { env }; assert!( - sccache::start_server(&SystemExec, "sccache", &server.env), + sccache::start_server(&FixtureExec, "sccache", &server.env), "sccache server starts" ); server @@ -119,7 +150,7 @@ fn workspace_crates_hit_the_compile_cache_when_a_slot_is_reused() { manifest.to_str().unwrap() ] )); - sccache::show_stats(&SystemExec, "sccache", env).expect("sccache answers") + sccache::show_stats(&FixtureExec, "sccache", env).expect("sccache answers") }; let first = build(0); @@ -177,13 +208,21 @@ fn slot_fixture_preserves_an_inherited_sccache_server() { tmp.path().to_str().unwrap(), ], )); - let before = sccache::show_stats(&SystemExec, "sccache", &outer.env).unwrap(); + let before = sccache::show_stats(&FixtureExec, "sccache", &outer.env).unwrap(); assert!(before.1 > 0, "the outer server has a sentinel cache miss"); // Use a child process so ambient environment is exercised without changing // this test process's environment while other tests may be running. + // CI can retain backend selectors after removing Actions credentials. + // The child must still use its own cold local cache and private socket. + let mut child_env = outer.env.clone(); + child_env.extend([ + ("SCCACHE_GCS_BUCKET".to_owned(), String::new()), + ("SCCACHE_GHA_ENABLED".to_owned(), "true".to_owned()), + ("SCCACHE_GHA_VERSION".to_owned(), "inherited".to_owned()), + ]); assert!(run_with( - &outer.env, + &child_env, std::env::current_exe().unwrap().to_str().unwrap(), &[ "--exact", @@ -192,7 +231,7 @@ fn slot_fixture_preserves_an_inherited_sccache_server() { ], )); assert_eq!( - sccache::show_stats(&SystemExec, "sccache", &outer.env), + sccache::show_stats(&FixtureExec, "sccache", &outer.env), Some(before), "the outer server survives fixture cleanup with unchanged counters" ); diff --git a/tests/units.rs b/tests/units.rs index e0d61b0..e5a3ba0 100644 --- a/tests/units.rs +++ b/tests/units.rs @@ -817,3 +817,769 @@ async fn the_newest_unit_with_a_supersede_key_wins() { .await; assert_eq!(code, 400); } + +async fn retire_ci(r: &Rig, count: usize, now: i64) { + let leases = r.h.queen.store().active_leases().await.unwrap(); + for lease in leases.into_iter().take(count) { + r.h.queen + .store() + .set_state(&lease.id, borg::store::LeaseStatus::Done, now) + .await + .unwrap(); + } +} + +fn sustained_ci(r: &Rig) { + for id in 1..=20 { + r.h.queue_job(id, "borg-ci-light-pr", 0); + } +} + +#[tokio::test] +async fn unit_share_survives_sustained_ci_and_rotates_clients_across_ticks() { + let r = Rig::new(json!({})).await; + let d = r.drone("d1").await; + let other = scoped_client(&r, "other", borg::classes::Scope::Pr).await; + sustained_ci(&r); + for n in 0..4 { + for (token, name) in [(&r.client, "agent"), (&other, "other")] { + let (_, sub) = r + .submit_as( + token, + &r.unit(json!({"argv": ["job", n.to_string(), name]})), + ) + .await; + assert_eq!(sub["outcome"], "new"); + } + } + let mut served = Vec::new(); + for round in 0..4 { + r.tick(10 + round * 2).await; + let units = + r.h.queen + .store() + .units_in(&[ + borg::units::UnitState::Leased, + borg::units::UnitState::Running, + ]) + .await + .unwrap(); + // 25% of 16 vCPU = two light units; CI consumes the remaining six. + assert_eq!(units.len(), 2); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 6); + assert_ne!( + units[0].client, units[1].client, + "both clients receive their share" + ); + if round == 0 { + // Held units count toward the next tick's budget. + r.tick(11).await; + assert_eq!( + r.h.queen + .store() + .units_in(&[ + borg::units::UnitState::Leased, + borg::units::UnitState::Running + ]) + .await + .unwrap() + .len(), + 2 + ); + } + for unit in units { + served.push(unit.client.clone()); + r.beat(&d, &[&unit.id]).await; + assert_eq!(r.post_result(&d, &unit.id, 0).await, 204); + } + retire_ci(&r, usize::MAX, r.base.unix_timestamp() + 11 + round * 2).await; + } + assert_eq!(served.iter().filter(|c| c.as_str() == "agent").count(), 4); + assert_eq!(served.iter().filter(|c| c.as_str() == "other").count(), 4); +} + +#[tokio::test] +async fn zero_share_keeps_ci_first_and_idle_ci_borrows_the_unit_share() { + for percent in [0, 25] { + let r = + Rig::new(json!({"units": {"enabled": true, "capacity_share_percent": percent}})).await; + r.drone("d1").await; + sustained_ci(&r); + // No unit demand reserves no capacity, even with the feature enabled. + r.tick(5).await; + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 8); + let (_, sub) = r.submit(&r.unit(json!({}))).await; + retire_ci(&r, usize::MAX, r.base.unix_timestamp() + 6).await; + r.tick(10).await; + assert_eq!( + state_of(&r, sub["unit_id"].as_str().unwrap()).await, + if percent == 0 { "queued" } else { "leased" } + ); + assert_eq!( + r.h.queen.store().active_leases().await.unwrap().len(), + if percent == 0 { 8 } else { 7 } + ); + } +} + +#[tokio::test] +async fn pending_ci_runners_do_not_absorb_work_unit_demand() { + let r = Rig::new(json!({"units": {"enabled": true, "capacity_share_percent": 0}})).await; + r.drone("d1").await; + r.h.queue_job(1, "borg-ci-light-pr", 0); + r.tick(5).await; + let (_, sub) = r.submit(&r.unit(json!({}))).await; + r.tick(10).await; + assert_eq!( + state_of(&r, sub["unit_id"].as_str().unwrap()).await, + "leased" + ); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn aging_reserves_contiguous_room_and_a_heavy_unit_can_cross_the_share() { + let r = Rig::new(json!({"starvation_after_secs": 30})).await; + r.drone("d1").await; + sustained_ci(&r); + r.tick(1).await; + let (_, sub) = r.submit(&r.unit(json!({"class": "heavy"}))).await; + let id = sub["unit_id"].as_str().unwrap(); + r.tick(5).await; + assert_eq!(state_of(&r, id).await, "queued"); + retire_ci(&r, 1, r.base.unix_timestamp() + 35).await; + r.tick(40).await; + assert_eq!(state_of(&r, id).await, "queued"); + assert_eq!( + r.h.queen.store().active_leases().await.unwrap().len(), + 7, + "an aged heavy unit prevents CI refilling its reserved host" + ); + assert!( + r.h.status().await["reserved_for"] + .as_array() + .unwrap() + .iter() + .any(|v| v["class"] == "heavy") + ); + retire_ci(&r, usize::MAX, r.base.unix_timestamp() + 41).await; + r.tick(45).await; + assert_eq!( + state_of(&r, id).await, + "leased", + "24 GiB unit exceeds the 7 GiB share but remains indivisible" + ); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn active_units_reduce_client_priority_even_across_toolchains() { + let r = Rig::new(json!({"units": {"enabled": true, "capacity_share_percent": 50}})).await; + let d = r.drone("d1").await; + let other = scoped_client(&r, "other", borg::classes::Scope::Pr).await; + sustained_ci(&r); + let mut occupied = Vec::new(); + for n in 0..4 { + let (_, sub) = r + .submit(&r.unit(json!({"argv": ["job", n.to_string()], "rust_toolchain": "1.95.0"}))) + .await; + if n < 2 { + let id = sub["unit_id"].as_str().unwrap(); + assert!( + r.h.queen + .store() + .lease_unit(id, "d1", r.base.unix_timestamp()) + .await + .unwrap() + ); + occupied.push(id.to_owned()); + } + r.submit_as( + &other, + &r.unit(json!({"argv": ["other", n.to_string()], "rust_toolchain": "1.96.0"})), + ) + .await; + } + r.beat(&d, &occupied.iter().map(String::as_str).collect::>()) + .await; + r.tick(5).await; + let units = + r.h.queen + .store() + .units_in(&[ + borg::units::UnitState::Leased, + borg::units::UnitState::Running, + ]) + .await + .unwrap(); + assert_eq!(units.len(), 4); + assert_eq!(units.iter().filter(|u| u.client == "agent").count(), 2); + assert_eq!( + units.iter().filter(|u| u.client == "other").count(), + 2, + "the less-loaded client wins despite its lexically later toolchain" + ); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 4); +} + +#[tokio::test] +async fn unit_share_never_overrides_owner_caps_or_unavailable_bids() { + let r = Rig::new(json!({"starvation_after_secs": 1})).await; + let d = r.h.enroll("d1").await; + r.h.heartbeat_with(&d, json!({"caps": {"max_vcpu": 4, "max_leases": 2}, + "bids": [{"class":"light", "scope":"pr", "available_slots": 1, "expires_at": r.base.unix_timestamp() + 1000}]})).await; + sustained_ci(&r); + let (_, heavy) = r.submit(&r.unit(json!({"class": "heavy"}))).await; + for n in 0..3 { + r.submit(&r.unit(json!({"argv": ["light", n.to_string()]}))) + .await; + } + r.tick(5).await; + assert_eq!( + state_of(&r, heavy["unit_id"].as_str().unwrap()).await, + "queued" + ); + let units = + r.h.queen + .store() + .units_in(&[borg::units::UnitState::Leased]) + .await + .unwrap(); + assert_eq!( + units.len(), + 1, + "a heartbeat bid slot is shared by both allocation passes" + ); + assert!(r.h.queen.store().active_leases().await.unwrap().is_empty()); + assert!( + r.h.status().await["reserved_for"] + .as_array() + .unwrap() + .is_empty(), + "an impossible heavy unit cannot reserve this capped host" + ); +} + +#[tokio::test] +async fn deploy_ci_retains_priority_over_the_unit_share() { + let r = Rig::new(json!({})).await; + deploy_drone(&r, "deploy").await; + let client = scoped_client(&r, "gate", borg::classes::Scope::Main).await; + let (_, sub) = r + .submit_as(&client, &r.unit(json!({"scope":"main", "class":"heavy"}))) + .await; + r.h.queue_job(1, "borg-ci-heavy-deploy", 0); + r.tick(5).await; + assert_eq!( + state_of(&r, sub["unit_id"].as_str().unwrap()).await, + "queued" + ); + let leases = r.h.queen.store().active_leases().await.unwrap(); + assert_eq!(leases.len(), 1); + assert_eq!(leases[0].scope, borg::classes::Scope::Deploy); +} + +#[tokio::test] +async fn one_slot_share_rotates_clients_without_toolchain_lexical_priority() { + let r = Rig::new(json!({})).await; + let d = r.drone("d1").await; + let other = scoped_client(&r, "other", borg::classes::Scope::Pr).await; + sustained_ci(&r); + for n in 0..2 { + for (token, toolchain) in [(&r.client, "1.96.0"), (&other, "1.95.0")] { + r.submit_as(token, &r.unit(json!({"class":"heavy", "argv":["job",n.to_string(),toolchain], "rust_toolchain":toolchain}))).await; + } + } + let mut clients = Vec::new(); + for round in 0..4 { + r.tick(5 + round).await; + let units = + r.h.queen + .store() + .units_in(&[borg::units::UnitState::Leased]) + .await + .unwrap(); + assert_eq!(units.len(), 1); + clients.push(units[0].client.clone()); + r.beat(&d, &[&units[0].id]).await; + assert_eq!(r.post_result(&d, &units[0].id, 0).await, 204); + retire_ci(&r, usize::MAX, r.base.unix_timestamp() + 5 + round).await; + } + assert_ne!(clients[0], clients[1]); + assert_eq!(clients[0], clients[2]); + assert_eq!(clients[1], clients[3]); +} + +#[tokio::test] +async fn aging_does_not_reserve_a_host_missing_the_units_toolchain() { + let r = Rig::new(json!({"starvation_after_secs":1, + "env_contract":{"enabled":true,"require_report":true, + "checks":{"compiler_cache":"off","toolchain":"required","cache_store":"off","resources":"off","tools":"off"}}})).await; + let d = r.h.enroll("d1").await; + r.h.heartbeat_with( + &d, + json!({"env_facts":{"observed_at":r.base.unix_timestamp(), + "toolchains":[{"channel":"1.95.0","rustc_version":"1.95.0"}]}}), + ) + .await; + sustained_ci(&r); + // Age only the unit, so an empty reservation set specifically proves + // that its missing toolchain cannot hold capacity from young CI demand. + let id = submit_at( + &r, + "agent", + json!({"class":"heavy","rust_toolchain":"1.96.0"}), + r.base.unix_timestamp() - 10, + ) + .await; + r.tick(0).await; + assert_eq!(state_of(&r, &id).await, "queued"); + assert!( + r.h.status().await["reserved_for"] + .as_array() + .unwrap() + .is_empty() + ); + assert_eq!( + r.h.queen.store().active_leases().await.unwrap().len(), + 8, + "ineligible work holds no CI capacity" + ); + // Once CI itself ages, its residual demand may legitimately reserve the + // host. That reservation must be light CI, never the blocked heavy unit. + r.tick(5).await; + assert_eq!(state_of(&r, &id).await, "queued"); + assert_eq!( + r.h.status().await["reserved_for"], + json!([{"drone_id":"d1","class":"light","scope":"pr"}]) + ); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 8); + let report = explanation(&r, &id).await; + assert!( + report + .history + .iter() + .all(|row| row.decision.reservation.is_none()) + ); +} + +#[tokio::test] +async fn per_unit_fairness_preserves_backlog_bid_pressure() { + let r = Rig::new(json!({})).await; + let d = r.h.enroll("d1").await; + r.h.heartbeat_with( + &d, + json!({"bids":[{"class":"light","scope":"pr", + "available_slots":8,"min_pressure":500,"expires_at":r.base.unix_timestamp()+1000}]}), + ) + .await; + for n in 0..8 { + r.submit(&r.unit(json!({"argv":["job",n.to_string()]}))) + .await; + } + sustained_ci(&r); + r.tick(5).await; + assert_eq!( + r.h.queen + .store() + .units_in(&[borg::units::UnitState::Leased]) + .await + .unwrap() + .len(), + 2, + "eight queued units create pressure500; selecting one fairly must retain the backlog pressure" + ); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 6); +} + +#[tokio::test] +async fn unit_share_accounts_for_owner_lease_slots_across_platforms() { + let r = Rig::new(json!({})).await; + let linux = r.h.enroll("linux").await; + r.h.heartbeat_with(&linux, json!({"caps":{"max_leases":1}})) + .await; + let mac = r.h.enroll("mac").await; + let mut cap = serde_json::to_value(support::capacity()).unwrap(); + cap["os"] = json!("macos"); + cap["arch"] = json!("aarch64"); + r.h.heartbeat_with(&mac, json!({"capacity":cap,"caps":{"max_leases":1}})) + .await; + r.submit(&r.unit(json!({"platform":"linux-x86_64"}))).await; + r.submit(&r.unit(json!({"platform":"macos-aarch64"}))).await; + sustained_ci(&r); + r.tick(5).await; + assert_eq!( + r.h.queen + .store() + .units_in(&[borg::units::UnitState::Leased]) + .await + .unwrap() + .len(), + 1, + "25 percent of two owner-capped slots rounds to one, even with spare CPU and memory" + ); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 1); +} + +async fn explanation(r: &Rig, key: &str) -> borg::explain::Explanation { + r.h.http + .get(format!("{}/explain", r.h.url)) + .query(&[("key", key)]) + .bearer_auth(&r.client) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap() +} + +#[tokio::test] +async fn queen_automatically_retains_unit_and_runner_decisions_after_persistence() { + use borg::explain::{PlacementKind, PlacementOutcome}; + let r = Rig::new(json!({})).await; + r.drone("d1").await; + let (_, sub) = r.submit(&r.unit(json!({}))).await; + let (_, wrong_platform) = r.submit(&r.unit(json!({"platform":"macos-aarch64"}))).await; + r.h.queue_job(1, "borg-ci-light-pr", 0); + r.tick(5).await; + let success = explanation(&r, sub["unit_id"].as_str().unwrap()).await; + assert!(!success.history.is_empty()); + assert_eq!(success.history[0].kind, PlacementKind::Unit); + assert_eq!(success.history[0].outcome, PlacementOutcome::Granted); + assert_eq!(success.history[0].decision.winner.as_deref(), Some("d1")); + let unplaced = explanation(&r, wrong_platform["unit_id"].as_str().unwrap()).await; + assert!(!unplaced.history.is_empty()); + assert_eq!(unplaced.history[0].outcome, PlacementOutcome::Unplaced); + assert!(unplaced.history[0].decision.winner.is_none()); + assert!(unplaced.history[0].decision.candidates.is_empty()); + let runner = explanation(&r, "borg-ci-light-pr").await; + assert!(!runner.history.is_empty()); + assert_eq!(runner.history[0].kind, PlacementKind::Runner); + assert_eq!(runner.history[0].outcome, PlacementOutcome::Granted); +} + +#[tokio::test] +async fn a_selected_runner_is_failed_in_history_when_minting_fails() { + let r = Rig::new(json!({})).await; + r.drone("d1").await; + r.h.queue_job(1, "borg-ci-light-pr", 0); + r.h.state().fail_mint = true; + r.tick(5).await; + assert!(r.h.queen.store().active_leases().await.unwrap().is_empty()); + let report = explanation(&r, "borg-ci-light-pr").await; + assert_eq!(report.history.len(), 1); + assert_eq!( + report.history[0].outcome, + borg::explain::PlacementOutcome::Failed + ); + assert_eq!(report.history[0].decision.winner.as_deref(), Some("d1")); +} + +#[tokio::test] +async fn blocked_unit_history_contains_effective_typed_checks_without_raw_environment() { + use borg::envcontract::{CheckKind, CheckState}; + let r = Rig::new(json!({"env_contract":{"enabled":true,"require_report":true, + "checks":{"compiler_cache":"off","toolchain":"required","cache_store":"off","resources":"off","tools":"off"}}})).await; + let d = r.h.enroll("d1").await; + r.h.heartbeat_with(&d,json!({"env_facts":{"observed_at":r.base.unix_timestamp(), + "rustc_wrapper":"private-path-should-not-appear", "toolchains":[{"channel":"1.95.0","rustc_version":"private-version-should-not-appear"}]}})).await; + let (_, sub) = r.submit(&r.unit(json!({"rust_toolchain":"1.96.0"}))).await; + r.tick(5).await; + let report = explanation(&r, sub["unit_id"].as_str().unwrap()).await; + assert_eq!( + report.history[0].outcome, + borg::explain::PlacementOutcome::Unplaced + ); + let candidate = &report.history[0].decision.candidates[0]; + assert_eq!( + candidate.environment_checks.get(&CheckKind::Toolchain), + Some(&CheckState::Fail) + ); + assert!( + candidate + .rejection_reasons + .contains(&borg::allocator::Rejection::Environment) + ); + let encoded = serde_json::to_string(&report).unwrap(); + assert!(!encoded.contains("private-path-should-not-appear")); + assert!(!encoded.contains("private-version-should-not-appear")); +} + +#[tokio::test] +async fn unit_history_uses_effective_capability_requirements_and_freshness() { + use borg::envcontract::{Capability, CheckState}; + let r = Rig::new(json!({"env_contract":{"enabled":true,"facts_max_age_secs":2, + "checks":{"compiler_cache":"off","toolchain":"off","cache_store":"off","resources":"off","tools":"off"}, + "required_capabilities":["sandbox"],"classes":{"light":{"required_capabilities":["docker"]}}}})).await; + let d = r.h.enroll("d1").await; + r.h.heartbeat_with( + &d, + json!({"env_facts":{"observed_at":r.base.unix_timestamp(),"capabilities":{ + "sandbox":{"observed_at":r.base.unix_timestamp(),"state":"fail"}, + "docker":{"observed_at":r.base.unix_timestamp(),"state":"pass"}}}}), + ) + .await; + let (_, sub) = r.submit(&r.unit(json!({}))).await; + r.tick(5).await; + let report = explanation(&r, sub["unit_id"].as_str().unwrap()).await; + assert_eq!( + report.history[0].outcome, + borg::explain::PlacementOutcome::Unplaced + ); + let candidate = &report.history[0].decision.candidates[0]; + assert_eq!(candidate.capability_checks.len(), 1); + assert_eq!( + candidate.capability_checks.get(&Capability::Docker), + Some(&CheckState::Unknown) + ); + assert!( + !candidate + .capability_checks + .contains_key(&Capability::Sandbox), + "a class replacement excludes the global requirement from diagnostics too" + ); +} + +#[tokio::test] +async fn leased_units_leave_the_backlog_before_the_next_bid_threshold_check() { + let r = Rig::new(json!({})).await; + let d = r.h.enroll("d1").await; + r.h.heartbeat_with( + &d, + json!({"bids":[{"class":"light","scope":"pr", + "available_slots":8,"min_pressure":250,"expires_at":r.base.unix_timestamp()+1000}]}), + ) + .await; + for n in 0..4 { + r.submit(&r.unit(json!({"argv":["job",n.to_string()]}))) + .await; + } + sustained_ci(&r); + r.tick(5).await; + assert_eq!( + r.h.queen + .store() + .units_in(&[borg::units::UnitState::Leased]) + .await + .unwrap() + .len(), + 1, + "initial pressure4*500/8=250 admits one; remaining3*500/7=214 is below the bid threshold" + ); + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 7); + let queued = + r.h.queen + .store() + .units_in(&[borg::units::UnitState::Queued]) + .await + .unwrap(); + assert_eq!(queued.len(), 3); + let report = explanation(&r, &queued[0].id).await; + assert!(report.history.iter().any(|row| row.decision.pressure == 214 + && row.outcome == borg::explain::PlacementOutcome::Unplaced)); +} + +#[tokio::test] +async fn young_deploy_ci_keeps_priority_when_main_ci_has_starved() { + let r = Rig::new(json!({"starvation_after_secs":30})).await; + deploy_drone(&r, "deploy").await; + let client = scoped_client(&r, "gate", borg::classes::Scope::Main).await; + let (_, unit) = r + .submit_as(&client, &r.unit(json!({"scope":"main","class":"heavy"}))) + .await; + r.h.queue_job(1, "borg-ci-heavy", 300); + r.h.queue_job(2, "borg-ci-heavy-deploy", 0); + r.tick(5).await; + assert_eq!( + state_of(&r, unit["unit_id"].as_str().unwrap()).await, + "queued" + ); + let leases = r.h.queen.store().active_leases().await.unwrap(); + assert_eq!(leases.len(), 1); + assert_eq!( + leases[0].scope, + borg::classes::Scope::Deploy, + "an aged main preallocation must not hide deploy demand before the unit share" + ); +} + +async fn submit_at(r: &Rig, client: &str, extra: Value, at: i64) -> String { + let unit: borg::units::SubmitUnit = serde_json::from_value(r.unit(extra)).unwrap(); + unit.validate().unwrap(); + let id = borg::queen::new_unit_id(); + r.h.queen + .store() + .submit_unit( + &unit, + client, + &id, + 3600, + r.h.queen.config().units.receipt_ttl_secs, + at, + ) + .await + .unwrap(); + id +} + +#[tokio::test] +async fn aged_heavy_unit_reserves_before_fresh_work_from_a_less_loaded_client() { + // Cover both a fresh light unit and an already-aged smaller class. + for light_created_offset in [119, 0] { + let r = Rig::new(json!({"starvation_after_secs":30})).await; + let d = r.drone("d1").await; + scoped_client(&r, "other", borg::classes::Scope::Pr).await; + let active = submit_at( + &r, + "agent", + json!({"argv":["active"]}), + r.base.unix_timestamp(), + ) + .await; + r.h.queen + .store() + .lease_unit(&active, "d1", r.base.unix_timestamp()) + .await + .unwrap(); + r.beat(&d, &[&active]).await; + sustained_ci(&r); + r.tick(1).await; + assert_eq!(r.h.queen.store().active_leases().await.unwrap().len(), 7); + let heavy = submit_at( + &r, + "agent", + json!({"class":"heavy","argv":["aged-heavy"]}), + r.base.unix_timestamp(), + ) + .await; + let fresh = submit_at( + &r, + "other", + json!({"argv":["fresh-light"]}), + r.base.unix_timestamp() + light_created_offset, + ) + .await; + retire_ci(&r, 1, r.base.unix_timestamp() + 119).await; + r.tick(120).await; + assert_eq!(state_of(&r, &heavy).await, "queued"); + assert_eq!( + state_of(&r, &fresh).await, + "queued", + "fresh light work cannot consume the last share slot before the aged heavy reservation" + ); + assert_eq!( + r.h.queen.store().active_leases().await.unwrap().len(), + 6, + "continuous CI cannot refill the host reserved for heavy work" + ); + assert!( + r.h.status().await["reserved_for"] + .as_array() + .unwrap() + .iter() + .any(|v| v["class"] == "heavy") + ); + retire_ci(&r, usize::MAX, r.base.unix_timestamp() + 121).await; + r.tick(122).await; + assert_eq!(state_of(&r, &heavy).await, "leased"); + assert_eq!(state_of(&r, &active).await, "running"); + assert_eq!(state_of(&r, &fresh).await, "queued"); + } +} + +#[tokio::test] +async fn deploy_prepass_uses_full_uncovered_runner_pressure_before_unit_share() { + use borg::classes::Scope; + use borg::explain::{PlacementKind, PlacementOutcome}; + let r = Rig::new(json!({})).await; + let drone = r.h.enroll("deploy").await; + r.h.heartbeat_with(&drone, json!({"trust":"deploy","bids":[ + {"class":"heavy","scope":"deploy","available_slots":1,"min_pressure":750,"expires_at":r.base.unix_timestamp()+1000}, + {"class":"heavy","scope":"main","available_slots":1,"expires_at":r.base.unix_timestamp()+1000} + ]})).await; + let client = scoped_client(&r, "gate", Scope::Main).await; + let (_, unit) = r + .submit_as(&client, &r.unit(json!({"scope":"main","class":"heavy"}))) + .await; + r.h.queue_job(1, "borg-ci-heavy", 0); + r.h.queue_job(2, "borg-ci-heavy-deploy", 0); + r.tick(5).await; + let leases = r.h.queen.store().active_leases().await.unwrap(); + assert_eq!(leases.len(), 1); + assert_eq!(leases[0].scope, Scope::Deploy); + assert_eq!( + state_of(&r, unit["unit_id"].as_str().unwrap()).await, + "queued", + "full runner pressure admits Deploy before a Main unit can consume its only heavy slot" + ); + let history = + r.h.queen + .store() + .placement_history(PlacementKind::Runner, "borg-ci-heavy-deploy", None) + .await + .unwrap(); + assert_eq!(history[0].outcome, PlacementOutcome::Granted); + assert_eq!(history[0].decision.pressure, 1000); + assert_eq!(history[0].decision.winner.as_deref(), Some("deploy")); +} + +#[tokio::test] +async fn pending_pr_runner_covering_main_does_not_inflate_deploy_pressure() { + use borg::allocator::Rejection; + use borg::classes::Scope; + use borg::explain::{PlacementKind, PlacementOutcome}; + let r = Rig::new(json!({})).await; + r.drone("shared").await; + r.h.queue_job(1, "borg-ci-heavy-pr", 0); + r.tick(0).await; + let pending = r.h.queen.store().active_leases().await.unwrap(); + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].scope, Scope::Pr); + assert_eq!(pending[0].state, borg::store::LeaseStatus::Pending); + + let drone = r.h.enroll("deploy").await; + r.h.heartbeat_with(&drone, json!({"trust":"deploy","bids":[ + {"class":"heavy","scope":"deploy","available_slots":1,"min_pressure":750,"expires_at":r.base.unix_timestamp()+1000}, + {"class":"heavy","scope":"main","available_slots":1,"expires_at":r.base.unix_timestamp()+1000} + ]})).await; + let client = scoped_client(&r, "gate", Scope::Main).await; + let (_, unit) = r + .submit_as(&client, &r.unit(json!({"scope":"main","class":"heavy"}))) + .await; + r.h.state().jobs.clear(); + r.h.queue_job(2, "borg-ci-heavy", 0); + r.h.queue_job(3, "borg-ci-heavy-deploy", 0); + r.tick(5).await; + let leases = r.h.queen.store().active_leases().await.unwrap(); + assert_eq!( + leases.len(), + 1, + "the existing pending PR runner covers Main CI" + ); + assert_eq!(leases[0].id, pending[0].id); + assert_eq!( + state_of(&r, unit["unit_id"].as_str().unwrap()).await, + "leased" + ); + let history = + r.h.queen + .store() + .placement_history(PlacementKind::Runner, "borg-ci-heavy-deploy", None) + .await + .unwrap(); + assert!( + history + .iter() + .all(|row| row.outcome == PlacementOutcome::Unplaced) + ); + assert!(history.iter().any(|row| row.decision.pressure == 500 + && row.decision.candidates.iter().any(|candidate| { + candidate.drone_id == "deploy" + && candidate + .rejection_reasons + .contains(&Rejection::BidThreshold) + }))); +} From 985a0a05a82b4052cbadc3effade42dfb1f5c0ef Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Mon, 12 Oct 2026 05:54:58 +0000 Subject: [PATCH 14/27] chore: project borg from Mono e52156fbb589 --- .repository-projection.json | 6 +++--- tests/run.rs | 23 ++++++++++++++--------- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 45d0b70..a5debda 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "6a106187736f994ddbc328e1f8bfe3d4786de24d", + "sourceSha": "e52156fbb589b7ed0e4bff6b1294a5abcc776efb", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "9ea7dac52a1c013677042fb1d313452ade1946c3", + "priorProjectedBase": "384e00bcad612da87223088490d005633c4cfe9f", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "17ec6815470882c73d267d67b062fe816ee7a97a411e18b70e598a8ff73bc250", + "contentDigest": "e603d60c1f5b99d7558d2da165b57d7e5001497a894a005115fc89cf58918874", "publicationEligible": true } diff --git a/tests/run.rs b/tests/run.rs index 6c4fe5d..fe3256e 100644 --- a/tests/run.rs +++ b/tests/run.rs @@ -711,15 +711,20 @@ fn the_cli_remote_only_flag_prevents_local_execution_and_emits_a_result() { clippy::disallowed_methods, reason = "test invokes the Borg CLI with a local sentinel" )] - let output = Command::new(env!("CARGO_BIN_EXE_borg")) - .current_dir(&r.dir) - .env_remove("BORG_QUEEN") - .env_remove("BORG_RUN_REMOTE_ONLY") - .args(["run", "--remote-only", "--result-file"]) - .arg(&result_path) - .args(["--", "sh", "-c", "touch ran-locally"]) - .output() - .unwrap(); + let output = Command::new( + option_env!("CARGO_BIN_EXE_borg") + .map(PathBuf::from) + .or_else(|| std::env::var_os("CARGO_BIN_EXE_borg").map(PathBuf::from)) + .expect("Cargo binary path available at runtime"), + ) + .current_dir(&r.dir) + .env_remove("BORG_QUEEN") + .env_remove("BORG_RUN_REMOTE_ONLY") + .args(["run", "--remote-only", "--result-file"]) + .arg(&result_path) + .args(["--", "sh", "-c", "touch ran-locally"]) + .output() + .unwrap(); assert_eq!(output.status.code(), Some(125)); assert!(output.stdout.is_empty()); assert!(!Path::new(&marker(&r)).exists()); From eb79fe9e794c5a5cea1239bad39ccde7d9fb8f20 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Mon, 12 Oct 2026 07:15:57 +0000 Subject: [PATCH 15/27] chore: project borg from Mono 50bb0b080bad --- .repository-projection.json | 6 +++--- Cargo.toml | 9 --------- tests/explain.rs | 2 +- tests/explain_cli.rs | 2 +- tests/host_suite.rs | 14 -------------- tests/placement_batch.rs | 2 +- tests/queen.rs | 2 +- tests/suite.rs | 21 ++++++++++++++++++--- 8 files changed, 25 insertions(+), 33 deletions(-) delete mode 100644 tests/host_suite.rs diff --git a/.repository-projection.json b/.repository-projection.json index a5debda..eda9b34 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "e52156fbb589b7ed0e4bff6b1294a5abcc776efb", + "sourceSha": "50bb0b080bad2817661adde6e362cc44426ec130", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "384e00bcad612da87223088490d005633c4cfe9f", + "priorProjectedBase": "c62234764b2520e62f3336d0503ba290b9efb78b", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", - "contentDigest": "e603d60c1f5b99d7558d2da165b57d7e5001497a894a005115fc89cf58918874", + "contentDigest": "e60726dff18dd8020a9127ef26a1c1c9f004f30f939cb8973f82da5d39a8ee2d", "publicationEligible": true } diff --git a/Cargo.toml b/Cargo.toml index d2f6702..335c8e3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,15 +24,6 @@ test = false name = "suite" path = "tests/suite.rs" -[[test]] -name = "host_suite" -path = "tests/host_suite.rs" - -# Keep the largest test module independently compilable for focused edits. -[[test]] -name = "queen" -path = "tests/queen.rs" - [[test]] name = "sccache_slots" path = "tests/sccache_slots.rs" diff --git a/tests/explain.rs b/tests/explain.rs index 5779f74..7b7ad9c 100644 --- a/tests/explain.rs +++ b/tests/explain.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use borg::allocator::{ ActiveLease, Allocation, AllocatorInput, Demand, DroneState, HostCaps, LeaseState, QueueAge, diff --git a/tests/explain_cli.rs b/tests/explain_cli.rs index 3f178a8..298ea85 100644 --- a/tests/explain_cli.rs +++ b/tests/explain_cli.rs @@ -1,5 +1,5 @@ //! Exercise the public CLI's local and authenticated HTTP history reads. -mod support; +use crate::support; use std::path::Path; use std::process::Output; diff --git a/tests/host_suite.rs b/tests/host_suite.rs deleted file mode 100644 index fb447a6..0000000 --- a/tests/host_suite.rs +++ /dev/null @@ -1,14 +0,0 @@ -//! Host execution, updates, and protocol tests share one fixture module. -//! Filter a module with `cargo test --test host_suite drone::` in the standalone -//! crate, or add `-p borg` from the Mono Rust workspace. - -mod support; - -mod admin_cli; -mod contract; -mod drone; -mod job_user; -mod rollout; -mod run; -mod tailnet; -mod update; diff --git a/tests/placement_batch.rs b/tests/placement_batch.rs index d8797a1..a97a82e 100644 --- a/tests/placement_batch.rs +++ b/tests/placement_batch.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use borg::allocator::{AllocatorInput, Demand, DroneState, HostCaps, PlacementDecision, allocate}; use borg::bid::ScoreWeights; diff --git a/tests/queen.rs b/tests/queen.rs index f2fe066..a3dfe63 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -1,4 +1,4 @@ -mod support; +use crate::support; use serde_json::{Value, json}; use support::Harness; diff --git a/tests/suite.rs b/tests/suite.rs index 6137ee8..e6bfa58 100644 --- a/tests/suite.rs +++ b/tests/suite.rs @@ -1,12 +1,27 @@ -//! Scheduling and service-state tests share one executable and fixture module. -//! Filter a module with `cargo test --test suite units::` in the standalone -//! crate, or add `-p borg` from the Mono Rust workspace. +//! Single integration-test binary for `borg`: scheduling, service-state, host +//! execution, update, and protocol tests share one executable and one fixture +//! module. Filter a module with `cargo test --test suite units::` in the +//! standalone crate, or add `-p borg` from the Mono Rust workspace. +//! `sccache_slots` stays a separate binary because it drives real `cargo` and +//! `sccache` processes. mod support; +mod admin_cli; mod capacity; +mod contract; +mod drone; +mod explain; +mod explain_cli; mod health; mod janitor; +mod job_user; +mod placement_batch; +mod queen; mod requeue; +mod rollout; +mod run; +mod tailnet; mod units; +mod update; mod wake; From a74b3af4b7e8b52bfba4228b033379115a38dda5 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Mon, 12 Oct 2026 08:00:52 +0000 Subject: [PATCH 16/27] chore: project borg from Mono fcb4796cdc8e --- .repository-projection.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index eda9b34..2961d59 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "50bb0b080bad2817661adde6e362cc44426ec130", + "sourceSha": "fcb4796cdc8ec682819173017c5c8d2c50f6b730", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "c62234764b2520e62f3336d0503ba290b9efb78b", + "priorProjectedBase": "1badb7a8070f7d5e4b6a61dd7d529f62e68aaf7a", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", - "toolDigest": "535e3cb77f0b59898e63a5f097fb3308564bafbc", + "toolDigest": "4d061573580b6f1387676e7c0edeed284e29c97b", "contentDigest": "e60726dff18dd8020a9127ef26a1c1c9f004f30f939cb8973f82da5d39a8ee2d", "publicationEligible": true } From 256f424eb9fe63e475e2f3b9372dd7b574ca046c Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Mon, 12 Oct 2026 09:44:33 +0000 Subject: [PATCH 17/27] chore: project borg from Mono 004e0bd914d0 --- .repository-projection.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 2961d59..d748a47 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,9 +3,9 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "fcb4796cdc8ec682819173017c5c8d2c50f6b730", + "sourceSha": "004e0bd914d0218e1049e3a48ee4c996b966d48d", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "1badb7a8070f7d5e4b6a61dd7d529f62e68aaf7a", + "priorProjectedBase": "50f13d9f1c8ea6caf75d17fa0f49c27fbac478cd", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "4d061573580b6f1387676e7c0edeed284e29c97b", "contentDigest": "e60726dff18dd8020a9127ef26a1c1c9f004f30f939cb8973f82da5d39a8ee2d", From 1af517a532f30d927be85bab17735fc862ce8151 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Fri, 16 Oct 2026 02:55:49 +0000 Subject: [PATCH 18/27] chore: project borg from Mono 0fd8f678d750 --- .repository-projection.json | 8 +- Cargo.lock | 304 ++++++++++++++++++++++++++++++++---- Cargo.toml | 3 +- README.md | 8 +- docs/operations.md | 71 +++++++++ src/api.rs | 41 ++++- src/drone/caches/git.rs | 103 +++++++++++- src/drone/client.rs | 4 +- src/drone/mod.rs | 144 +++++++++++++++-- src/envcontract.rs | 67 +++++++- src/main.rs | 23 ++- src/probe/mod.rs | 66 +++++++- src/queen.rs | 12 ++ tests/contract.rs | 86 ++++++++++ tests/run.rs | 30 ++-- 15 files changed, 879 insertions(+), 91 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index d748a47..d14eb4c 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "004e0bd914d0218e1049e3a48ee4c996b966d48d", + "sourceSha": "0fd8f678d7506046a6e241de983a94ddfdc4f019", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "50f13d9f1c8ea6caf75d17fa0f49c27fbac478cd", + "priorProjectedBase": "76bef692ff5ff193ccd54ed79cfe36408172b130", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", - "toolDigest": "4d061573580b6f1387676e7c0edeed284e29c97b", - "contentDigest": "e60726dff18dd8020a9127ef26a1c1c9f004f30f939cb8973f82da5d39a8ee2d", + "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", + "contentDigest": "12df8b19ca3ba327fecc607d78f8f681be2e2b5c294e1524d9f827b785c825b7", "publicationEligible": true } diff --git a/Cargo.lock b/Cargo.lock index 65b6f37..aa0ce85 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -204,6 +204,7 @@ dependencies = [ "jsonwebtoken", "rand 0.8.6", "reqwest", + "rustls", "serde", "serde_json", "sha2", @@ -253,6 +254,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "clap" version = "4.6.1" @@ -308,6 +320,32 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -317,6 +355,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc" version = "3.4.0" @@ -576,11 +623,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", - "js-sys", "libc", "r-efi 5.3.0", "wasip2", - "wasm-bindgen", ] [[package]] @@ -590,8 +635,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", + "wasm-bindgen", ] [[package]] @@ -711,7 +759,6 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.8", ] [[package]] @@ -868,6 +915,55 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror", + "walkdir", + "windows-link", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn", +] + [[package]] name = "jobserver" version = "0.1.34" @@ -1055,6 +1151,12 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + [[package]] name = "parking" version = "2.2.1" @@ -1167,14 +1269,16 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.15" +version = "0.11.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" +checksum = "a9746dbde176634f4f2f1faf2404e30a31b2bc1e9cafb5329c95d8177a18c9fc" dependencies = [ + "aws-lc-rs", "bytes", - "getrandom 0.3.4", + "getrandom 0.4.3", "lru-slab", - "rand 0.9.4", + "rand 0.10.1", + "rand_pcg", "ring", "rustc-hash", "rustls", @@ -1228,18 +1332,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" dependencies = [ "libc", - "rand_chacha 0.3.1", + "rand_chacha", "rand_core 0.6.4", ] [[package]] name = "rand" -version = "0.9.4" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "d2e8e8bcc7961af1fdac401278c6a831614941f6164ee3bf4ce61b7edb162207" dependencies = [ - "rand_chacha 0.9.0", - "rand_core 0.9.5", + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", ] [[package]] @@ -1252,16 +1357,6 @@ dependencies = [ "rand_core 0.6.4", ] -[[package]] -name = "rand_chacha" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" -dependencies = [ - "ppv-lite86", - "rand_core 0.9.5", -] - [[package]] name = "rand_core" version = "0.6.4" @@ -1273,11 +1368,17 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.9.5" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" dependencies = [ - "getrandom 0.3.4", + "rand_core 0.10.1", ] [[package]] @@ -1308,9 +1409,9 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "reqwest" -version = "0.12.28" +version = "0.13.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ "base64", "bytes", @@ -1328,6 +1429,7 @@ dependencies = [ "quinn", "rustls", "rustls-pki-types", + "rustls-platform-verifier", "serde", "serde_json", "serde_urlencoded", @@ -1341,7 +1443,6 @@ dependencies = [ "wasm-bindgen", "wasm-bindgen-futures", "web-sys", - "webpki-roots 1.0.8", ] [[package]] @@ -1364,6 +1465,15 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustix" version = "1.1.4" @@ -1379,10 +1489,11 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.45" +version = "0.23.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +checksum = "48e13bd8c0e9365c43cfa5c9e8f9ad49d3c8444926c9aac819e0e4dc503c8fdf" dependencies = [ + "aws-lc-rs", "once_cell", "ring", "rustls-pki-types", @@ -1391,6 +1502,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + [[package]] name = "rustls-pki-types" version = "1.15.0" @@ -1401,12 +1524,40 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-platform-verifier" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" +dependencies = [ + "core-foundation", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + [[package]] name = "rustls-webpki" version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted 0.9.0", @@ -1424,12 +1575,59 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + [[package]] name = "serde" version = "1.0.228" @@ -1512,7 +1710,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -1550,6 +1748,22 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "simd_cesu8" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f90157bb87cddf702797c5dadfa0be7d266cdf49e22da2fcaa32eff75b2c33" +dependencies = [ + "rustc_version", + "simdutf8", +] + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "simple_asn1" version = "0.6.4" @@ -2107,6 +2321,16 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -2206,6 +2430,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d46a5a140e6f7afeccd8eae97eff335163939eac8b929834875168b29b3d267" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webpki-roots" version = "0.26.11" @@ -2224,6 +2457,15 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "windows-link" version = "0.2.1" diff --git a/Cargo.toml b/Cargo.toml index 335c8e3..2da85db 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,7 +35,8 @@ clap = { version = "4.5", features = ["derive", "env"] } hex = "0.4" jsonwebtoken = { version = "10.3", default-features = false, features = ["aws_lc_rs", "use_pem"] } rand = "0.8" -reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "json"] } +rustls = { version = "0.23", default-features = false, features = ["aws_lc_rs"] } +reqwest = { version = "0.13", default-features = false, features = ["rustls", "json", "query"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/README.md b/README.md index 258ce21..062bd9a 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,7 @@ borg queen enroll --config queen.toml host-a > host-a.token borg queen run --config queen.toml ``` -Enrollment prints the token once, to stdout; diagnostics go to stderr. Keep +Enrollment prints the token once, to `stdout`; diagnostics go to `stderr`. Keep `host-a.token` private. The queen stores its hash. The queen serves plain HTTP. Put its loopback listener behind an HTTPS @@ -72,8 +72,8 @@ sudo "$(command -v borg)" drone install \ ``` On Linux, the installer creates a daemon account and a separate job account, -installs the pinned GitHub runner, and writes a systemd unit. WSL needs -systemd enabled. On macOS, add `--run-as ` to select the +installs the pinned GitHub runner, and writes a `systemd` unit. WSL needs +`systemd` enabled. On macOS, add `--run-as ` to select the LaunchDaemon's account. See the [drone guide](docs/operations.md#drone) for platform setup and isolation details. @@ -160,4 +160,4 @@ upstream so the next projection preserves them. ## License -[Business Source License 1.1](LICENSE) (BUSL-1.1). +[Business Source License 1.1](../../../LICENSE) (BUSL-1.1). diff --git a/docs/operations.md b/docs/operations.md index 4aae546..cac28a3 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -22,6 +22,7 @@ Configuration, scheduling, drone administration, and protocol details for - [Host health](#host-health) - [Tailnet enrollment](#tailnet-enrollment) - [Cloud burst drones](#cloud-burst-drones) + - [Deploy-trust drone](#deploy-trust-drone) - [Starvation guard](#starvation-guard) - [Preemption](#preemption) - [Admission by CPU pressure](#admission-by-cpu-pressure) @@ -645,6 +646,76 @@ unless a NAT gateway exists, which this setup avoids. If the cloud reclaims a spot VM mid-job, GitHub fails that job when its runner stops responding; rerun it. Queued jobs are unaffected and go to other drones. +### Deploy-trust drone + +Deploy-scope jobs (labels ending `-deploy`) run only on a drone whose trust is +`deploy` or `deploy-only`. Without one they queue indefinitely +(`queued_deploy` in `/status`). A deploy-only drone is a separate, on-demand +VM that never takes pull-request or main work. Do not reuse a burst drone. + +Create the VM. Standard provisioning, no max run duration, no service account, +the same dedicated network, and an ephemeral external IP: + +```sh +gcloud compute networks subnets create borg-burst-usc1 --network=borg-burst \ + --region=us-central1 --range=192.0.2.0/24 # replace with your private subnet range +gcloud compute instances create borg-deploy-1 --zone=us-central1-c \ + --machine-type=n2d-standard-8 --provisioning-model=STANDARD \ + --no-service-account --no-scopes \ + --network=borg-burst --subnet=borg-burst-usc1 \ + --image-family=ubuntu-2204-lts --image-project=ubuntu-os-cloud \ + --boot-disk-size=200GB --boot-disk-type=pd-balanced +``` + +Use Ubuntu 22.04: iac `scripts/install-owned-runner-toolchain.sh` refuses +24.04. Zones can be out of capacity (`stockout`); try another zone in the +region. Reach the VM with `gcloud compute ssh --tunnel-through-iap`. + +Install the toolchain and arm64 emulation: + +```sh +sudo useradd -m -s /bin/bash runner # the script checks as this user +sudo apt-get install -y shellcheck qemu-user-static binfmt-support +sudo bash install-owned-runner-toolchain.sh # docker, gh, node 22, rust, kani +sudo docker run --rm --platform linux/arm64 alpine uname -m # aarch64 +``` + +The script ends with `sudo: a password is required` from its +`sudo -n true` check on the `runner` user. Every tool check before it has +passed; the failure is expected and the user stays without sudo. + +Join the tailnet with an ephemeral, pre-approved auth key tagged +`tag:borg-deploy` (its ACL reaches the queen on 443 only; mint it with the +Tailscale OAuth client in 1Password). Do not use `tag:borg-drone`, which is for +drones that run pull requests: + +```sh +sudo tailscale up --auth-key=file:/tmp/ts.key --hostname=borg-deploy-1 \ + --advertise-tags=tag:borg-deploy +shred -u /tmp/ts.key +``` + +Leave MagicDNS on: the queen is addressed by name. The queen's +`tailscale serve` publishes its HTTPS URL on port 443, proxying the LAN +listener. Use that URL in place of `https://queen.example.test` below; the drone needs no +`--allow-insecure` and no `--ca-cert`. + +Enroll and install. Use the queen's own binary on the VM; it is a glibc 2.35 +build and runs on 22.04 and 24.04: + +```sh +# on the Proxmox node hosting the queen container (CT 102) +pct exec 102 -- /usr/local/bin/borg queen enroll --config /etc/borg/queen.json borg-deploy-1 > drone.token +# copy the binary and token to the VM, then +sudo ./borg drone install --queen https://queen.example.test \ + --token-file ./drone.token --trust deploy-only --max-vcpu 8 --max-mem-mib 24576 +shred -u drone.token +``` + +Check `/status`: `borg-deploy-1` is `fresh`, `accepting`, `trust: +deploy-only`, and `queued_deploy` falls as leases start. Neither the queen +config nor the queen binary changes. + ### Starvation guard When demand for a class has waited longer than `starvation_after_secs`, the diff --git a/src/api.rs b/src/api.rs index e83a33b..f9db125 100644 --- a/src/api.rs +++ b/src/api.rs @@ -185,6 +185,9 @@ pub struct UpdateReport { #[derive(Debug, Serialize, Deserialize)] pub struct HeartbeatResponse { + /// Existing Queen disk floor; older Queens omit it and drones use 10 GiB. + #[serde(default = "crate::envcontract::default_min_disk_free_bytes")] + pub min_free_disk_bytes: u64, pub leases: Vec, /// Work units to run (resent until acknowledged in `active_units`). #[serde(default)] @@ -418,12 +421,24 @@ async fn heartbeat( let drone = authenticate(&queen, &headers).await?; let now = OffsetDateTime::now_utc().unix_timestamp(); let store = queen.store(); + let disk = crate::envcontract::disk_admission( + queen.config().env_contract.min_cache_dir_free_bytes, + req.capacity.free_disk_mib, + req.env_facts + .as_ref() + .filter(|facts| { + (0..=queen.config().env_contract.facts_max_age_secs) + .contains(&(now - facts.observed_at)) + }) + .and_then(|facts| facts.cache_dir_free_bytes), + ); + let accepting = req.accepting && !disk.blocked(); store .record_heartbeat( &drone, &req.capacity, &req.caps, - req.accepting, + accepting, &req.cache_stats, now, ) @@ -455,7 +470,7 @@ async fn heartbeat( let start = queen.snapshot().start_standby; let start_standby: Vec = standby_rows .iter() - .filter(|l| start.contains(&l.id)) + .filter(|l| accepting && start.contains(&l.id)) .map(|l| l.id.clone()) .collect(); let drop_standby: Vec = req @@ -494,6 +509,7 @@ async fn heartbeat( .unacked_units(&drone) .await? .into_iter() + .filter(|_| accepting) .map(|u| UnitGrant { unit_id: u.id, repo: u.repo, @@ -514,6 +530,7 @@ async fn heartbeat( .undelivered(&drone) .await? .into_iter() + .filter(|_| accepting) .filter_map(|l| { Some(LeaseGrant { lease_id: l.id, @@ -537,6 +554,7 @@ async fn heartbeat( None => None, }; Ok(Json(HeartbeatResponse { + min_free_disk_bytes: queen.config().env_contract.min_cache_dir_free_bytes, leases, units, cancel_units, @@ -662,6 +680,8 @@ pub struct DroneStatus { pub heartbeat_age_secs: Option, pub fresh: bool, pub accepting: bool, + /// Enforced independently of the optional environment contract. + pub disk_admission: crate::envcontract::DiskAdmission, /// `shared` or `deploy`: which scopes the drone may run. pub trust: Trust, pub hostname: Option, @@ -919,11 +939,25 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result for d in &rows { let age = d.last_heartbeat.map(|t| now - t); let fresh = age.is_some_and(|a| a <= cfg.drone_stale_secs); + let disk = crate::envcontract::disk_admission( + cfg.env_contract.min_cache_dir_free_bytes, + d.capacity + .as_ref() + .filter(|_| fresh) + .and_then(|cap| cap.free_disk_mib), + d.env_facts + .as_ref() + .filter(|facts| { + (0..=cfg.env_contract.facts_max_age_secs).contains(&(now - facts.observed_at)) + }) + .and_then(|facts| facts.cache_dir_free_bytes), + ); let mut budget = None; let mut free_slots_here = BTreeMap::new(); let mut theoretical_here = BTreeMap::new(); if let (true, Some(cap)) = (fresh, &d.capacity) { let mut state = d.allocator_state(cap, now); + state.accepting &= !disk.blocked(); state.blocked_classes = blocked_classes( &cfg.env_contract, d.env_facts.as_ref(), @@ -957,7 +991,8 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result name: d.name.clone(), heartbeat_age_secs: age, fresh, - accepting: d.accepting, + accepting: d.accepting && !disk.blocked(), + disk_admission: disk, trust: d.trust, hostname: d.capacity.as_ref().map(|c| c.hostname.clone()), kind: d.capacity.as_ref().map(|c| c.kind), diff --git a/src/drone/caches/git.rs b/src/drone/caches/git.rs index c602fa7..18082fe 100644 --- a/src/drone/caches/git.rs +++ b/src/drone/caches/git.rs @@ -286,12 +286,26 @@ impl GitMirrors { } } - /// The `BORG_GIT_MIRROR_*` variables for mirrors that exist, counting the - /// lease as served or not. + /// Advertise only mirrors with a real branch ref. `git init --bare` + /// creates HEAD before the first fetch; if that fetch fails, handing the + /// empty repository to a job makes its "local" checkout fetch everything. + /// `show-ref` reads loose and packed refs without fetching blob contents. + fn usable(&self, mirror: &Mirror) -> bool { + if !mirror.path.join("HEAD").is_file() { + return false; + } + let path = mirror.path.to_string_lossy(); + self.exec + .run(&self.git, &["-C", &path, "show-ref", "--heads", "--quiet"]) + .is_ok_and(|output| output.success) + } + + /// The `BORG_GIT_MIRROR_*` variables for usable mirrors only, counting + /// a lease as a miss if initialization or the first fetch never finished. pub fn env(&self) -> Vec<(String, String)> { let mut env = Vec::new(); for m in &self.mirrors { - if m.path.join("HEAD").exists() { + if self.usable(m) { self.served.fetch_add(1, Ordering::Relaxed); env.push((mirror_var(&m.repo), m.path.to_string_lossy().into_owned())); } else { @@ -523,6 +537,89 @@ mod tests { ); } + #[test] + fn empty_bare_mirror_is_not_advertised_after_failed_fetch() { + // Real Git fixture: a failed first fetch can leave a valid HEAD file + // but no branches or objects. The runner must use its normal fallback. + let root = tempfile::tempdir().unwrap(); + let path = mirror_path(root.path(), "example-org/app"); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + let init = crate::drone::exec::SystemExec + .run( + "git", + &["init", "--bare", "--quiet", path.to_str().unwrap()], + ) + .unwrap(); + assert!(init.success); + assert!(path.join("HEAD").exists()); + let m = GitMirrors::new( + &cfg(&["example-org/app"]), + root.path(), + Arc::new(crate::drone::exec::SystemExec), + Arc::new(LocalFs), + ); + assert!(m.env().is_empty(), "a bare HEAD is not a populated mirror"); + assert_eq!(m.hit_counts(), (0, 1)); + } + + #[test] + fn populated_bare_mirror_is_advertised() { + let root = tempfile::tempdir().unwrap(); + let path = mirror_path(root.path(), "example-org/app"); + let seed = root.path().join("seed"); + let exec = crate::drone::exec::SystemExec; + for args in [ + vec!["init", "--quiet", "-b", "main", seed.to_str().unwrap()], + vec![ + "-C", + seed.to_str().unwrap(), + "-c", + "user.email=ci@example.test", + "-c", + "user.name=CI", + "commit", + "--allow-empty", + "-m", + "init", + "--quiet", + ], + ] { + assert!(exec.run("git", &args).unwrap().success); + } + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + assert!( + exec.run( + "git", + &["init", "--bare", "--quiet", path.to_str().unwrap()] + ) + .unwrap() + .success + ); + assert!( + exec.run( + "git", + &[ + "-C", + path.to_str().unwrap(), + "fetch", + "--quiet", + seed.to_str().unwrap(), + "main:main" + ] + ) + .unwrap() + .success + ); + let m = GitMirrors::new( + &cfg(&["example-org/app"]), + root.path(), + Arc::new(crate::drone::exec::SystemExec), + Arc::new(LocalFs), + ); + assert_eq!(m.env().len(), 1); + assert_eq!(m.hit_counts(), (1, 0)); + } + #[test] fn env_lists_only_existing_mirrors_and_counts_hits_and_misses() { let root = tempfile::tempdir().unwrap(); diff --git a/src/drone/client.rs b/src/drone/client.rs index 111fa76..7409cb0 100644 --- a/src/drone/client.rs +++ b/src/drone/client.rs @@ -40,9 +40,7 @@ impl QueenClient { source, })?; let cert = reqwest::Certificate::from_pem(&pem).map_err(ClientError::CaParse)?; - builder = builder - .tls_built_in_root_certs(false) - .add_root_certificate(cert); + builder = builder.tls_certs_only([cert]); } Ok(Self { http: builder.build()?, diff --git a/src/drone/mod.rs b/src/drone/mod.rs index 7f0537a..897f581 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -22,7 +22,7 @@ pub mod standby; use std::collections::{HashMap, HashSet}; use std::path::Path; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -93,6 +93,8 @@ pub struct Drone { update_wanted: AtomicBool, /// The last tools probe, reused for a minute. tools_cache: probe::ToolsCache, + /// Last Queen floor, retained while the Queen is unreachable. + disk_floor: AtomicU64, capability_cache: capabilities::CapabilityCache, } @@ -133,6 +135,7 @@ impl Drone { selfupdate, update_wanted: AtomicBool::new(false), tools_cache: probe::ToolsCache::default(), + disk_floor: AtomicU64::new(crate::envcontract::DEFAULT_MIN_DISK_FREE_BYTES), capability_cache: capabilities::CapabilityCache::default(), client: Arc::new(client), host, @@ -202,12 +205,18 @@ impl Drone { // another version it stops accepting and updates when idle. let accepting = accepting && !self.update_wanted.load(Ordering::Relaxed); let host = Arc::clone(&self.host); - let capacity = tokio::task::spawn_blocking(move || probe::probe(host.as_ref())) - .await - .map_err(|e| ClientError::Status { - path: format!("probe: {e}"), - status: 0, - })?; + let cfg = Arc::clone(&self.cfg); + let (capacity, cache_free) = tokio::task::spawn_blocking(move || { + let capacity = probe::probe_at(host.as_ref(), &cfg.work_root); + let cache_free = probe::free_disk_mib_at(host.as_ref(), &cfg.cache_root) + .map(|mib| mib.saturating_mul(1 << 20)); + (capacity, cache_free) + }) + .await + .map_err(|e| ClientError::Status { + path: format!("probe: {e}"), + status: 0, + })?; let host = Arc::clone(&self.host); let now = time::OffsetDateTime::now_utc().unix_timestamp(); let mut env_facts = @@ -217,6 +226,14 @@ impl Drone { path: format!("probe: {e}"), status: 0, })?; + env_facts.cache_dir_free_bytes = cache_free; + let accepting = accepting + && !crate::envcontract::disk_admission( + self.disk_floor.load(Ordering::Relaxed), + capacity.free_disk_mib, + cache_free, + ) + .blocked(); // The tools probe never fails the heartbeat: no map means not measured. let (exec, cache) = (Arc::clone(&self.job_exec), self.tools_cache.clone()); env_facts.tools = tokio::task::spawn_blocking(move || cache.get(exec.as_ref(), now)) @@ -255,6 +272,15 @@ impl Drone { env_facts: Some(env_facts), }; let resp = self.client.heartbeat(&req).await?; + self.disk_floor + .store(resp.min_free_disk_bytes, Ordering::Relaxed); + let accepting = accepting + && !crate::envcontract::disk_admission( + resp.min_free_disk_bytes, + req.capacity.free_disk_mib, + cache_free, + ) + .blocked(); let defer = self.after_heartbeat(&req, resp.desired_version.as_deref()); if let Some(credential) = resp.mirror_credential { self.caches.set_mirror_credential(credential); @@ -307,10 +333,29 @@ impl Drone { } /// Start every held standby because the queen has been unreachable. - fn fail_open(&self, tasks: &mut JoinSet<()>) { + async fn fail_open(&self, tasks: &mut JoinSet<()>) { if self.standbys.ids().is_empty() { return; } + let host = Arc::clone(&self.host); + let cfg = Arc::clone(&self.cfg); + let Ok((work_free, cache_free)) = tokio::task::spawn_blocking(move || { + probe::job_disk_space(host.as_ref(), &cfg.work_root, &cfg.cache_root) + }) + .await + else { + return; + }; + if crate::envcontract::disk_admission( + self.disk_floor.load(Ordering::Relaxed), + work_free, + cache_free, + ) + .blocked() + { + tracing::warn!("disk below admission floor: leaving standby runners unstarted"); + return; + } let started = self.start_standbys(None, tasks); if started > 0 { tracing::warn!("queen unreachable: started {started} standby deploy runner(s)"); @@ -461,7 +506,7 @@ impl Drone { _ = every.tick() => { let now = time::OffsetDateTime::now_utc().unix_timestamp(); if should_fail_open(last_ok_at, now, self.cfg.fail_open_after_secs) { - self.fail_open(&mut tasks); + self.fail_open(&mut tasks).await; } if tokio::time::Instant::now() < retry_at { continue; @@ -625,3 +670,84 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { } } } + +#[cfg(test)] +mod disk_tests { + use super::*; + use crate::classes::{Class, Scope as JobScope}; + use crate::drone::exec::Output; + use crate::secret::Secret; + + struct DiskHost; + impl Host for DiskHost { + fn read_to_string(&self, _: &Path) -> Option { + None + } + fn exists(&self, _: &Path) -> bool { + false + } + fn run(&self, program: &str, args: &[&str]) -> Option { + if program != "df" { + return None; + } + let available_kib = if args.last() == Some(&"/work") { + 20 * 1024 * 1024 + } else { + 100 * 1024 * 1024 + }; + Some(format!( + "fs blocks used avail cap mount\ndisk 999999999 0 {available_kib} 0% /" + )) + } + fn os(&self) -> &'static str { + "linux" + } + fn arch(&self) -> &'static str { + "x86_64" + } + fn available_parallelism(&self) -> Option { + Some(1) + } + } + struct NoExec; + impl Exec for NoExec { + fn run(&self, _: &str, _: &[&str]) -> std::io::Result { + panic!("disk-blocked standby must not run commands") + } + } + + #[tokio::test] + async fn fail_open_retains_standby_and_active_lease_below_learned_queen_floor() { + let cfg = DroneConfig::from_json(r#"{"queen_url":"http://127.0.0.1:1","allow_insecure":true,"runner_dir":"/runner","work_root":"/work","cache_root":"/cache"}"#).unwrap(); + let client = QueenClient::new(&cfg, Secret::new("test".to_owned())).unwrap(); + let drone = Drone::new(cfg, client, Arc::new(DiskHost), Arc::new(NoExec)); + drone.disk_floor.store(32 << 30, Ordering::Relaxed); + drone.held.lock().unwrap().insert("active".into()); + drone.standbys.insert(LeaseGrant { + lease_id: "standby".into(), + class: Class::Light, + scope: JobScope::Deploy, + labels: vec![], + runner_name: "standby".into(), + jit_config: Secret::new("jit".to_owned()), + expires_at: i64::MAX, + }); + let mut tasks = JoinSet::new(); + drone.fail_open(&mut tasks).await; + assert!(tasks.is_empty()); + assert_eq!(drone.standbys.ids(), vec!["standby"]); + assert_eq!(drone.held_ids(), vec!["active"]); + assert!(drone.slot_assignments().is_empty()); + assert!(!*drone.abort.subscribe().borrow()); + } + + #[test] + fn older_queen_response_keeps_default_floor_and_new_response_carries_configured_floor() { + let mut response = serde_json::json!({"leases": []}); + let old: crate::api::HeartbeatResponse = serde_json::from_value(response.clone()).unwrap(); + assert_eq!(old.min_free_disk_bytes, 10 << 30); + response["min_free_disk_bytes"] = serde_json::json!(32u64 << 30); + let new: crate::api::HeartbeatResponse = serde_json::from_value(response).unwrap(); + assert_eq!(new.min_free_disk_bytes, 32 << 30); + } +} diff --git a/src/envcontract.rs b/src/envcontract.rs index 650c6a3..a258c5c 100644 --- a/src/envcontract.rs +++ b/src/envcontract.rs @@ -209,6 +209,53 @@ pub struct EnvFacts { pub capabilities: BTreeMap, } +/// Admission safety, independent of optional cache efficiency checks. +#[derive(Debug, Clone, Serialize)] +pub struct DiskAdmission { + pub state: CheckState, + pub minimum_free_bytes: u64, + pub work_free_bytes: Option, + pub cache_free_bytes: Option, +} + +impl DiskAdmission { + pub fn blocked(&self) -> bool { + self.state == CheckState::Fail + } +} + +pub const DEFAULT_MIN_DISK_FREE_BYTES: u64 = 10 << 30; +pub fn default_min_disk_free_bytes() -> u64 { + DEFAULT_MIN_DISK_FREE_BYTES +} + +/// A floor guards admission; it does not reserve space for active jobs. +/// Any known low filesystem blocks. Unknown measurements retain compatibility. +pub fn disk_admission( + minimum_free_bytes: u64, + work_free_mib: Option, + cache_free_bytes: Option, +) -> DiskAdmission { + let work_free_bytes = work_free_mib.map(|mib| mib.saturating_mul(1 << 20)); + let low = [work_free_bytes, cache_free_bytes] + .into_iter() + .flatten() + .any(|free| free < minimum_free_bytes); + let state = if low { + CheckState::Fail + } else if work_free_bytes.is_none() || cache_free_bytes.is_none() { + CheckState::Unknown + } else { + CheckState::Pass + }; + DiskAdmission { + state, + minimum_free_bytes, + work_free_bytes, + cache_free_bytes, + } +} + /// Settings of one class (or the defaults for all). #[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] #[serde(default, deny_unknown_fields)] @@ -262,7 +309,7 @@ impl Default for ContractSettings { require_report: false, facts_max_age_secs: 180, allowed_wrappers: vec!["sccache".into()], - min_cache_dir_free_bytes: 10 << 30, + min_cache_dir_free_bytes: DEFAULT_MIN_DISK_FREE_BYTES, max_cache_error_pm: 50, hit_rate_warn_pm: 10, min_requests_for_hit_rate: 200, @@ -866,6 +913,24 @@ pub fn blocked_classes( mod tests { use super::*; + #[test] + fn disk_admission_blocks_either_known_low_mount_and_keeps_unknown_explicit() { + let floor = 32 << 30; + assert!(disk_admission(floor, Some(8 * 1024), Some(100 << 30)).blocked()); + assert!(disk_admission(floor, Some(100 * 1024), Some(8 << 30)).blocked()); + assert!(disk_admission(floor, None, Some(0)).blocked()); + assert_eq!( + disk_admission(floor, Some(32 * 1024), Some(floor)).state, + CheckState::Pass + ); + assert_eq!(disk_admission(floor, None, None).state, CheckState::Unknown); + assert_eq!( + disk_admission(floor, Some(100 * 1024), None).state, + CheckState::Unknown + ); + assert!(!disk_admission(floor, None, None).blocked()); + } + const NOW: i64 = 10_000; fn on() -> ContractSettings { diff --git a/src/main.rs b/src/main.rs index 3f1d643..3c592c7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -500,18 +500,16 @@ async fn run_queen(path: &Path) -> Result<()> { let janitor = match &cfg.supersede_janitor { Some(settings) => { let token = janitor_token(settings, &cfg)?; - tracing::info!( - "supersede janitor enabled ({}, {} token)", - if settings.dry_run { - "dry-run" - } else { - "cancelling" - }, - match token { - TokenSource::Static(_) => "static", - TokenSource::App(_) => "App installation", - } - ); + let mode = if settings.dry_run { + "dry-run" + } else { + "cancelling" + }; + let credential_kind = match token { + TokenSource::Static(_) => "static", + TokenSource::App(_) => "App installation", + }; + tracing::info!("supersede janitor enabled ({mode}, {credential_kind} credential)"); if cfg.preemption.is_some() { requeuer = Some(Requeuer::with_token_source( &cfg.github.api_base, @@ -859,6 +857,7 @@ async fn run_command(a: RunArgs) -> u8 { #[tokio::main] async fn main() -> Result { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let cli = Cli::parse(); if let Command::Run(a) = cli.command { return Ok(std::process::ExitCode::from(run_command(a).await)); diff --git a/src/probe/mod.rs b/src/probe/mod.rs index 33d2bad..b9d5cd4 100644 --- a/src/probe/mod.rs +++ b/src/probe/mod.rs @@ -61,6 +61,11 @@ pub struct Capacity { const DISK_PROBE_PATH: &str = "/"; pub fn probe(host: &dyn Host) -> Capacity { + probe_at(host, Path::new(DISK_PROBE_PATH)) +} + +/// Capacity of the filesystem actually used for leases. +pub fn probe_at(host: &dyn Host, work_root: &Path) -> Capacity { Capacity { hostname: hostname(host), os: host.os().to_owned(), @@ -68,7 +73,7 @@ pub fn probe(host: &dyn Host) -> Capacity { kind: host_kind(host), vcpu: host.available_parallelism().unwrap_or(1).max(1), mem_mib: mem_mib(host).unwrap_or(0), - free_disk_mib: free_disk_mib(host), + free_disk_mib: free_disk_mib_at(host, work_root), docker_usable: host .run("docker", &["info", "--format", "{{.ServerVersion}}"]) .is_some(), @@ -305,8 +310,16 @@ fn cgroup_mem_limit_mib(host: &dyn Host) -> Option { .map(|bytes| bytes / (1024 * 1024)) } -fn free_disk_mib(host: &dyn Host) -> Option { - free_disk_mib_at(host, Path::new(DISK_PROBE_PATH)) +/// Fresh available blocks on the configured job filesystems (which may differ). +pub fn job_disk_space( + host: &dyn Host, + work_root: &Path, + cache_root: &Path, +) -> (Option, Option) { + ( + free_disk_mib_at(host, work_root), + free_disk_mib_at(host, cache_root).map(|mib| mib.saturating_mul(1 << 20)), + ) } /// Free space, in MiB, on the filesystem holding `path`. @@ -344,8 +357,16 @@ mod tests { fn exists(&self, path: &Path) -> bool { path.to_str().is_some_and(|p| self.files.contains_key(p)) } - fn run(&self, program: &str, _args: &[&str]) -> Option { - self.commands.get(program).map(|s| (*s).to_owned()) + fn run(&self, program: &str, args: &[&str]) -> Option { + let key = if program == "df" { + args.last() + .copied() + .filter(|path| self.commands.contains_key(path)) + .unwrap_or(program) + } else { + program + }; + self.commands.get(key).map(|s| (*s).to_owned()) } fn os(&self) -> &'static str { self.os @@ -358,6 +379,39 @@ mod tests { } } + #[test] + fn job_disk_probes_use_each_configured_filesystem_and_preserve_unknown() { + let host = FakeHost { + commands: HashMap::from([ + ( + "/", + "fs blocks used available cap mount\nroot 999999999 0 999999999 0% /", + ), + ( + "/work", + "fs blocks used available cap mount\nwork 99999999 0 8388608 0% /work", + ), + ( + "/cache", + "fs blocks used available cap mount\ncache 99999999 0 52428800 0% /cache", + ), + ]), + ..linux() + }; + assert_eq!( + probe_at(&host, Path::new("/work")).free_disk_mib, + Some(8192) + ); + assert_eq!( + job_disk_space(&host, Path::new("/work"), Path::new("/cache")), + (Some(8192), Some(50 << 30)) + ); + assert_eq!( + job_disk_space(&host, Path::new("/missing"), Path::new("/cache")), + (None, Some(50 << 30)) + ); + } + fn linux() -> FakeHost { FakeHost { os: "linux", @@ -577,7 +631,7 @@ mod tests { "df", "Filesystem 1024-blocks Used Available Capacity Mounted on\n/dev/sda1 100000000 1000000 52428800 2% /", ); - assert_eq!(free_disk_mib(&h), Some(51200)); + assert_eq!(free_disk_mib_at(&h, Path::new("/")), Some(51200)); } #[test] diff --git a/src/queen.rs b/src/queen.rs index 7cfee02..c3d1741 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -555,6 +555,18 @@ impl Queen { } if let (true, Some(cap)) = (fresh, &d.capacity) { let mut state = d.allocator_state(cap, now_ts); + state.accepting &= !crate::envcontract::disk_admission( + self.cfg.env_contract.min_cache_dir_free_bytes, + cap.free_disk_mib, + d.env_facts + .as_ref() + .filter(|facts| { + (0..=self.cfg.env_contract.facts_max_age_secs) + .contains(&(now_ts - facts.observed_at)) + }) + .and_then(|facts| facts.cache_dir_free_bytes), + ) + .blocked(); state.blocked_classes = blocked_classes( &self.cfg.env_contract, d.env_facts.as_ref(), diff --git a/tests/contract.rs b/tests/contract.rs index 098bf37..42b4e1b 100644 --- a/tests/contract.rs +++ b/tests/contract.rs @@ -429,3 +429,89 @@ async fn class_specific_job_user_proofs_gate_ci_and_survive_heartbeat_storage() < now() - 3500 ); } + +#[tokio::test] +async fn disk_floor_blocks_new_delivery_and_placement_with_contract_off_but_keeps_active_work() { + let h = Harness::start( + json!({"env_contract": {"enabled": false, "min_cache_dir_free_bytes": 32u64 << 30}}), + ) + .await; + let token = h.enroll("disk").await; + h.heartbeat_with(&token, json!({})).await; + h.queue_job(101, "borg-ci-light-pr", 1); + h.tick().await; + let initial = h.heartbeat_with(&token, json!({})).await; + assert_eq!(initial["min_free_disk_bytes"], 32u64 << 30); + let active = initial["leases"][0]["lease_id"].as_str().unwrap(); + let mut low = serde_json::to_value(support::capacity()).unwrap(); + low["free_disk_mib"] = json!(8192); + let held = h + .heartbeat_with(&token, json!({"capacity": low, "active_leases": [active]})) + .await; + assert!(held["leases"].as_array().unwrap().is_empty()); + assert!(held["cancel_leases"].as_array().unwrap().is_empty()); + h.queue_job(102, "borg-ci-light-pr", 2); + h.tick().await; + assert_eq!(h.state().mints.len(), 1, "no new lease under the floor"); + let status = h.status().await; + let d = drone(&status, "disk"); + assert_eq!(d["accepting"], false); + assert_eq!(d["active_leases"], 1); + assert_eq!(d["disk_admission"]["state"], "fail"); + assert!( + d["contract"].is_null(), + "safety does not enable the optional contract" + ); + h.heartbeat_with(&token, json!({"active_leases": [active]})) + .await; + h.tick().await; + assert_eq!(h.state().mints.len(), 2, "recovery restores admission"); +} + +#[tokio::test] +async fn low_cache_blocks_pending_delivery_without_cache_stats_and_recovers_at_exact_floor() { + let h = Harness::start(json!({"env_contract": {"checks": {"cache_store": "advisory"}}})).await; + let token = h.enroll("disk").await; + h.heartbeat_with(&token, json!({})).await; + h.queue_job(201, "borg-ci-light-pr", 1); + h.tick().await; + let low = h + .heartbeat_with( + &token, + json!({"env_facts": {"observed_at": now(), "cache_dir_free_bytes": 1}}), + ) + .await; + assert!( + low["leases"].as_array().unwrap().is_empty(), + "pending JIT remains unstarted" + ); + let leases = h.queen.store().active_leases().await.unwrap(); + assert_eq!(leases.len(), 1); + assert_eq!(leases[0].state, borg::store::LeaseStatus::Pending); + let recovered = h + .heartbeat_with( + &token, + json!({"env_facts": {"observed_at": now(), "cache_dir_free_bytes": 10u64 << 30}}), + ) + .await; + assert_eq!(recovered["leases"].as_array().unwrap().len(), 1); +} + +#[tokio::test] +async fn stale_cache_space_is_unknown_and_does_not_override_current_work_capacity() { + let h = Harness::start(json!({"env_contract": {"facts_max_age_secs": 2}})).await; + let token = h.enroll("disk").await; + h.heartbeat_with( + &token, + json!({"env_facts": {"observed_at": now() - 10, "cache_dir_free_bytes": 0}}), + ) + .await; + h.queue_job(301, "borg-ci-light-pr", 1); + h.tick().await; + assert_eq!(h.state().mints.len(), 1); + let status = h.status().await; + let d = drone(&status, "disk"); + assert_eq!(d["accepting"], true); + assert_eq!(d["disk_admission"]["state"], "unknown"); + assert!(d["disk_admission"]["cache_free_bytes"].is_null()); +} diff --git a/tests/run.rs b/tests/run.rs index fe3256e..1967463 100644 --- a/tests/run.rs +++ b/tests/run.rs @@ -242,6 +242,13 @@ async fn fake_queen(mode: Mode) -> (String, Shared) { (url, f) } +fn borg_bin() -> PathBuf { + option_env!("CARGO_BIN_EXE_borg") + .map(PathBuf::from) + .or_else(|| std::env::var_os("CARGO_BIN_EXE_borg").map(PathBuf::from)) + .expect("Cargo binary path available at runtime") +} + fn git(dir: &Path, args: &[&str]) -> String { #[allow( clippy::disallowed_methods, @@ -711,20 +718,15 @@ fn the_cli_remote_only_flag_prevents_local_execution_and_emits_a_result() { clippy::disallowed_methods, reason = "test invokes the Borg CLI with a local sentinel" )] - let output = Command::new( - option_env!("CARGO_BIN_EXE_borg") - .map(PathBuf::from) - .or_else(|| std::env::var_os("CARGO_BIN_EXE_borg").map(PathBuf::from)) - .expect("Cargo binary path available at runtime"), - ) - .current_dir(&r.dir) - .env_remove("BORG_QUEEN") - .env_remove("BORG_RUN_REMOTE_ONLY") - .args(["run", "--remote-only", "--result-file"]) - .arg(&result_path) - .args(["--", "sh", "-c", "touch ran-locally"]) - .output() - .unwrap(); + let output = Command::new(borg_bin()) + .current_dir(&r.dir) + .env_remove("BORG_QUEEN") + .env_remove("BORG_RUN_REMOTE_ONLY") + .args(["run", "--remote-only", "--result-file"]) + .arg(&result_path) + .args(["--", "sh", "-c", "touch ran-locally"]) + .output() + .unwrap(); assert_eq!(output.status.code(), Some(125)); assert!(output.stdout.is_empty()); assert!(!Path::new(&marker(&r)).exists()); From 273c4f6510f19f5c544c98bbe59f8314e050f3c3 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Fri, 16 Oct 2026 07:37:41 +0000 Subject: [PATCH 19/27] chore: project borg from Mono 4578fdbb3924 --- .repository-projection.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index d14eb4c..04b3387 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,9 +3,9 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "0fd8f678d7506046a6e241de983a94ddfdc4f019", + "sourceSha": "4578fdbb3924d7849d48c45a7ed814c4c2e2e947", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "76bef692ff5ff193ccd54ed79cfe36408172b130", + "priorProjectedBase": "b67719b3fdf78215423602f53c04e5e82e89864e", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", "contentDigest": "12df8b19ca3ba327fecc607d78f8f681be2e2b5c294e1524d9f827b785c825b7", From a3a4c6d2a02fa9f9d2cd942812747d42a644b44c Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Fri, 16 Oct 2026 21:13:30 +0000 Subject: [PATCH 20/27] chore: project borg from Mono 91f10217da2a --- .repository-projection.json | 6 +- docs/operations.md | 58 +++++++- src/api.rs | 19 ++- src/config.rs | 10 ++ src/janitor.rs | 260 ++++++++++++++++++++++++++++++------ src/probe/mod.rs | 4 +- tests/janitor.rs | 104 ++++++++++++++- 7 files changed, 403 insertions(+), 58 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 04b3387..f84ced8 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "4578fdbb3924d7849d48c45a7ed814c4c2e2e947", + "sourceSha": "91f10217da2a2a8dceb3aedbc480aab1eba1bc95", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "b67719b3fdf78215423602f53c04e5e82e89864e", + "priorProjectedBase": "bc61279726a1ea1773112633c4769aa7128ccefa", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", - "contentDigest": "12df8b19ca3ba327fecc607d78f8f681be2e2b5c294e1524d9f827b785c825b7", + "contentDigest": "021a908917cb1b06a2a4e35505a9faafce219868aa67d9a800c37ccb0de15c25", "publicationEligible": true } diff --git a/docs/operations.md b/docs/operations.md index cac28a3..45ff441 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -560,12 +560,58 @@ the others only if they are still queued; a running duplicate is never cancelled. `workflow_dispatch` and `merge_group` runs are never grouped, because publication waits on those exact runs. Each cancellation is logged with the id of the run that was kept. With `dry_run` the log and `would_cancel` show -what would have been cancelled. When the janitor is enabled, `GET /status` has -a `janitor` object (`passes`, `last`, `total`, each with the counters -`considered`, `superseded`, `cancelled`, `cancelled_running`, `duplicates`, -`dedup_cancelled`, `would_cancel`, `deferred`, `raced`, `errors`) and -`GET /metrics` has `borg_janitor__total`, summed over passes since the -queen started. +what would have been cancelled. + +The pull request's current head comes from the run's own `pull_requests` +entry, which GitHub fills at read time with the open pull requests whose head +branch matches. The janitor reads the pulls API only when that entry carries no +head, so the App does not need pull-request read access. + +Two more lists, both empty by default, cover the runs a closed pull request or +a newer push leaves behind: + +* `closed_pr_workflows`: workflow file paths (`.github/workflows/ci.yml`) or + names. A queued `pull_request` run of one of them (a running one too with + `cancel_running`) whose head repository is the repository itself and which + names no open pull request belongs to a pull request that was closed or + merged, and is cancelled. A fresh read must still show no pull request. Do + not list a workflow that runs on the `closed` event (teardown, cleanup): its + runs look the same. +* `superseded_push_branches`: branch names, for example `["main"]`. A queued + `push` run on one of them is cancelled when a newer queued or running run of + the same workflow on the same branch is on another commit. Running push runs + are never cancelled, so the branch always finishes its newest started run. + +When the janitor is enabled, `GET /status` has a `janitor` object: + +* `since`: unix time the queen started; `total` counts from here. +* `last_pass_at`: unix time of the last pass. +* `passes`: passes since `since`. +* `last`: what the last pass alone found. A run that stays queued is found + again by every pass, so `last` is the current state. +* `total`: `last` summed over every pass. Because a run that stays queued or + running is counted again on each pass, `duplicates`, `raced`, + `uncancellable` and `errors` grow with time even when nothing changes; read + `last` to judge the present and the rate of `total` to judge a trend. + +| Counter | Meaning | +| --- | --- | +| `considered` | `pull_request` runs examined: queued ones, plus running ones with `cancel_running`, that name exactly one pull request or (with `closed_pr_workflows`) none. | +| `superseded` | Of those, runs whose head is no longer the pull request's head. | +| `closed_pr` | Runs of `closed_pr_workflows` whose pull request is closed or merged. | +| `superseded_push` | Queued push runs on `superseded_push_branches` with a newer run of the same workflow. | +| `cancelled` | Queued runs cancelled for any of the three reasons above. | +| `cancelled_running` | Running runs cancelled as superseded or closed-PR (`cancel_running`). | +| `duplicates` | Runs beyond the one kept in a same-workflow, same-commit group (`dedupe_same_sha`), queued or running. Running duplicates are never cancelled, so they are counted again on every pass while they run. | +| `dedup_cancelled` | Queued duplicates cancelled. | +| `would_cancel` | Runs `dry_run` would have cancelled, of any kind. | +| `deferred` | Runs left for a later pass by `max_cancels_per_pass`. | +| `raced` | Runs whose status, head or pull request changed between the listing and the re-read just before the cancel; left alone. | +| `uncancellable` | Runs whose cancel GitHub answered with 409. Runs stuck in `queued` for days (for example `workflow_run` runs GitHub never started) answer 409 on every pass and are counted each time. | +| `errors` | Failed GitHub calls: a run listing for a repository, a pulls API read, a run re-read or a cancel. Each is logged as a `janitor:` warning in the queen journal. | + +`GET /metrics` has `borg_janitor__total` for each counter, summed over +passes since the queen started. The GitHub App needs organization self-hosted runners: read and write, plus read access to Actions on the listed repositories. Set diff --git a/src/api.rs b/src/api.rs index f9db125..ed4a506 100644 --- a/src/api.rs +++ b/src/api.rs @@ -1740,14 +1740,24 @@ pub fn render_metrics(status: &Status) -> String { "Runs whose head was no longer the pull request's head", t.superseded, ), + ( + "closed_pr", + "Runs of listed workflows whose pull request was closed or merged", + t.closed_pr, + ), + ( + "superseded_push", + "Queued push runs with a newer run of the same workflow on the branch", + t.superseded_push, + ), ( "cancelled", - "Queued superseded runs the janitor cancelled", + "Queued superseded, closed-PR or superseded-push runs the janitor cancelled", t.cancelled, ), ( "cancelled_running", - "Running superseded runs the janitor cancelled", + "Running superseded or closed-PR runs the janitor cancelled", t.cancelled_running, ), ( @@ -1775,6 +1785,11 @@ pub fn render_metrics(status: &Status) -> String { "Runs that changed state before the janitor cancelled them", t.raced, ), + ( + "uncancellable", + "Runs whose cancel GitHub answered with 409", + t.uncancellable, + ), ( "errors", "Janitor per-run or per-pull-request failures", diff --git a/src/config.rs b/src/config.rs index 1256379..8eaefc6 100644 --- a/src/config.rs +++ b/src/config.rs @@ -86,6 +86,16 @@ pub struct JanitorSettings { /// commit, event and pull request. #[serde(default)] pub dedupe_same_sha: bool, + /// Workflows, by file path (`.github/workflows/ci.yml`) or name, whose + /// pull-request runs are cancelled once the pull request is closed or + /// merged. Leave out workflows that run on the `closed` event. + #[serde(default)] + pub closed_pr_workflows: Vec, + /// Branches whose queued `push` runs are cancelled when a newer run of the + /// same workflow on the branch is queued or running. Running push runs + /// are never cancelled. + #[serde(default)] + pub superseded_push_branches: Vec, } fn yes() -> bool { diff --git a/src/janitor.rs b/src/janitor.rs index ea39384..691908c 100644 --- a/src/janitor.rs +++ b/src/janitor.rs @@ -1,5 +1,8 @@ //! The supersede janitor: cancels queued pull-request workflow runs whose head -//! commit is no longer the pull request's head. +//! commit is no longer the pull request's head, runs of listed workflows for +//! pull requests that were closed or merged (`closed_pr_workflows`), and queued +//! push runs on listed branches that a newer run of the same workflow +//! supersedes (`superseded_push_branches`). //! //! A push to a pull request branch makes every run of the older commit //! pointless, but GitHub keeps them queued and they occupy runner demand. Two @@ -19,11 +22,17 @@ //! cancelling, so a run whose status changed in between is left alone; //! * only `pull_request` runs that name exactly one pull request are touched; //! runs from forks (no pull request attached) and merge groups never are; -//! * the pull request's current head comes from the pulls API, not from the -//! run's own snapshot; +//! * the pull request's current head comes from the run's `pull_requests` +//! entry, which GitHub fills at read time with the open pull requests whose +//! head matches, and from the pulls API only when that entry has no head; +//! * a `pull_request` run with no open pull request is treated as belonging +//! to a closed one only when its head repository is the repository itself +//! (forks never) and its workflow is listed in `closed_pr_workflows`, and a +//! fresh read must still show no pull request; //! * at most `max_cancels_per_pass` cancellations per pass, one pass per //! `interval_secs`; -//! * a cancel that finds the run already finished (409) is a no-op; +//! * a cancel that GitHub refuses with 409 is a no-op, counted as +//! `uncancellable`; //! * every cancellation (or, in dry-run, every would-be cancellation) is logged. use std::collections::{BTreeMap, BTreeSet}; @@ -48,9 +57,14 @@ pub struct PassReport { pub considered: u32, /// Runs whose head is no longer the pull request's head. pub superseded: u32, - /// Queued superseded runs cancelled. + /// Runs of `closed_pr_workflows` whose pull request is closed or merged. + pub closed_pr: u32, + /// Queued push runs on `superseded_push_branches` with a newer run of the + /// same workflow on the branch. + pub superseded_push: u32, + /// Queued runs cancelled as superseded, closed-PR or superseded-push. pub cancelled: u32, - /// Running superseded runs cancelled (`cancel_running`). + /// Running runs cancelled as superseded or closed-PR (`cancel_running`). pub cancelled_running: u32, /// Runs beyond the first in a same-workflow, same-commit group /// (`dedupe_same_sha`), queued or running. @@ -61,8 +75,12 @@ pub struct PassReport { pub would_cancel: u32, /// Runs left for a later pass by the per-pass cap. pub deferred: u32, - /// Runs that changed state before the cancel and were left alone. + /// Runs that changed state between the listing and the cancel and were + /// left alone. pub raced: u32, + /// Runs whose cancel GitHub answered with 409 (already finished, or stuck + /// in a state GitHub will not cancel). + pub uncancellable: u32, /// Per-run or per-pull-request failures (logged, not fatal). pub errors: u32, } @@ -71,6 +89,8 @@ impl PassReport { fn add(&mut self, o: &PassReport) { self.considered += o.considered; self.superseded += o.superseded; + self.closed_pr += o.closed_pr; + self.superseded_push += o.superseded_push; self.cancelled += o.cancelled; self.cancelled_running += o.cancelled_running; self.duplicates += o.duplicates; @@ -78,6 +98,7 @@ impl PassReport { self.would_cancel += o.would_cancel; self.deferred += o.deferred; self.raced += o.raced; + self.uncancellable += o.uncancellable; self.errors += o.errors; } } @@ -86,8 +107,15 @@ impl PassReport { #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize)] pub struct JanitorView { pub passes: u64, + /// Unix time the queen started; `total` counts from here. + pub since: i64, + /// Unix time of the last pass; `None` before the first. + pub last_pass_at: Option, + /// What the last pass alone found. Runs still queued are found again by + /// every pass, so this is the current state. pub last: PassReport, - /// Sums over every pass since the queen started. + /// Sums over every pass since `since`. A run that stays queued is counted + /// once per pass, so these grow with time even when nothing changes. pub total: PassReport, } @@ -96,16 +124,31 @@ struct Run { id: u64, #[serde(default)] workflow_id: u64, + #[serde(default)] + name: String, + #[serde(default)] + path: String, status: String, event: String, head_sha: String, #[serde(default)] + head_branch: Option, + #[serde(default)] + head_repository: Option, + #[serde(default)] pull_requests: Vec, } +#[derive(Debug, Deserialize)] +struct RepoRef { + full_name: String, +} + #[derive(Debug, Deserialize)] struct PullRef { number: u64, + #[serde(default)] + head: Option, } #[derive(Debug, Deserialize)] @@ -150,7 +193,10 @@ impl Janitor { api_base: api_base.trim_end_matches('/').to_owned(), token, settings, - view: Mutex::default(), + view: Mutex::new(JanitorView { + since: now(), + ..JanitorView::default() + }), }) } @@ -234,9 +280,11 @@ impl Janitor { } } tracing::debug!( - "janitor pass: considered {} superseded {} cancelled {} cancelled_running {} duplicates {} dedup_cancelled {} would_cancel {} deferred {} raced {} errors {}", + "janitor pass: considered {} superseded {} closed_pr {} superseded_push {} cancelled {} cancelled_running {} duplicates {} dedup_cancelled {} would_cancel {} deferred {} raced {} uncancellable {} errors {}", report.considered, report.superseded, + report.closed_pr, + report.superseded_push, report.cancelled, report.cancelled_running, report.duplicates, @@ -244,10 +292,12 @@ impl Janitor { report.would_cancel, report.deferred, report.raced, + report.uncancellable, report.errors ); let mut view = self.view.lock().expect("janitor view lock"); view.passes += 1; + view.last_pass_at = Some(now()); view.total.add(&report); view.last = report.clone(); drop(view); @@ -261,7 +311,10 @@ impl Janitor { report: &mut PassReport, ) -> Result<(), GithubError> { let st = &self.settings; - let event = (!st.dedupe_same_sha).then_some("pull_request"); + let all_events = st.dedupe_same_sha + || !st.closed_pr_workflows.is_empty() + || !st.superseded_push_branches.is_empty(); + let event = (!all_events).then_some("pull_request"); let mut runs = self.list_runs(repo, "queued", event).await?; if st.cancel_running || st.dedupe_same_sha { runs.extend(self.list_runs(repo, "in_progress", event).await?); @@ -278,48 +331,57 @@ impl Janitor { { continue; } - // A fork's run carries no pull request; with more than one the - // answer is ambiguous. Leave both alone. - let [pr] = run.pull_requests.as_slice() else { - continue; - }; - report.considered += 1; - let head = if let Some(cached) = heads.get(&pr.number) { - cached.clone() - } else { - let fetched = match self - .get::(&format!("/repos/{repo}/pulls/{}", pr.number)) - .await - { - Ok(p) => Some(p.head.sha), - Err(e) => { - tracing::warn!("janitor: {repo}#{}: {e}", pr.number); - report.errors += 1; - None + let why = match run.pull_requests.as_slice() { + // No open pull request matches this same-repository run: its + // pull request was closed or merged. + [] if self.closed_pr_workflow(repo, run) => { + report.considered += 1; + report.closed_pr += 1; + format!( + "for branch {}: no open pull request", + run.head_branch.as_deref().unwrap_or("?") + ) + } + [pr] => { + report.considered += 1; + let Some(current) = self.pr_head(repo, pr, &mut heads, report).await else { + continue; + }; + if run.head_sha == current { + continue; } - }; - heads.insert(pr.number, fetched.clone()); - fetched - }; - let Some(current) = head else { - continue; + report.superseded += 1; + format!( + "for #{}: head {} superseded by {current}", + pr.number, run.head_sha + ) + } + // A fork's run carries no pull request; with more than one + // the answer is ambiguous. Leave both alone. + _ => continue, }; - if run.head_sha == current { + stale.insert(run.id); + if *budget == 0 { + report.deferred += 1; continue; } - report.superseded += 1; + let kind = if running { Kind::Running } else { Kind::Stale }; + let outcome = self.cancel_if_unchanged(repo, run, &why).await; + tally(report, budget, repo, run.id, kind, outcome); + } + for (run, newer) in superseded_pushes(&runs, &st.superseded_push_branches) { + report.superseded_push += 1; stale.insert(run.id); if *budget == 0 { report.deferred += 1; continue; } let why = format!( - "for #{}: head {} superseded by {current}", - pr.number, run.head_sha + "on {}: superseded by run {newer}", + run.head_branch.as_deref().unwrap_or("?") ); - let kind = if running { Kind::Running } else { Kind::Stale }; let outcome = self.cancel_if_unchanged(repo, run, &why).await; - tally(report, budget, repo, run.id, kind, outcome); + tally(report, budget, repo, run.id, Kind::Stale, outcome); } if st.dedupe_same_sha { for (run, kept) in surplus(&runs, &stale) { @@ -342,6 +404,48 @@ impl Janitor { Ok(()) } + /// Whether `run` is a same-repository run of a `closed_pr_workflows` entry. + fn closed_pr_workflow(&self, repo: &str, run: &Run) -> bool { + run.head_repository + .as_ref() + .is_some_and(|r| r.full_name.eq_ignore_ascii_case(repo)) + && self + .settings + .closed_pr_workflows + .iter() + .any(|w| *w == run.path || *w == run.name) + } + + /// The pull request's current head: the run's own entry when GitHub filled + /// it, else the pulls API (cached per pass; `None` after a logged error). + async fn pr_head( + &self, + repo: &str, + pr: &PullRef, + heads: &mut BTreeMap>, + report: &mut PassReport, + ) -> Option { + if let Some(head) = &pr.head { + return Some(head.sha.clone()); + } + if let Some(cached) = heads.get(&pr.number) { + return cached.clone(); + } + let fetched = match self + .get::(&format!("/repos/{repo}/pulls/{}", pr.number)) + .await + { + Ok(p) => Some(p.head.sha), + Err(e) => { + tracing::warn!("janitor: {repo}#{}: {e}", pr.number); + report.errors += 1; + None + } + }; + heads.insert(pr.number, fetched.clone()); + fetched + } + /// Cancel `run` unless a fresh read shows its status or head changed. async fn cancel_if_unchanged( &self, @@ -352,7 +456,10 @@ impl Janitor { let fresh: Run = self .get(&format!("/repos/{repo}/actions/runs/{}", run.id)) .await?; - if fresh.status != run.status || fresh.head_sha != run.head_sha { + if fresh.status != run.status + || fresh.head_sha != run.head_sha + || fresh.pull_requests.len() != run.pull_requests.len() + { return Ok(Outcome::Raced); } let state = &run.status; @@ -367,7 +474,8 @@ impl Janitor { tracing::info!("janitor: cancelled {state} {repo} run {} {why}", run.id); Ok(Outcome::Cancelled) } else { - Ok(Outcome::Raced) + tracing::debug!("janitor: {repo} run {} answered 409 to cancel", run.id); + Ok(Outcome::Uncancellable) } } } @@ -401,6 +509,7 @@ fn tally( report.would_cancel += 1; } Ok(Outcome::Raced) => report.raced += 1, + Ok(Outcome::Uncancellable) => report.uncancellable += 1, Err(e) => { tracing::warn!("janitor: {repo} run {run}: {e}"); report.errors += 1; @@ -442,6 +551,41 @@ enum Outcome { Cancelled, DryRun, Raced, + Uncancellable, +} + +fn now() -> i64 { + time::OffsetDateTime::now_utc().unix_timestamp() +} + +/// Queued `push` runs on `branches`, each with the id of the newest queued or +/// running run of the same workflow on the same branch, when that newer run +/// is on another commit. +fn superseded_pushes<'a>(runs: &'a [Run], branches: &[String]) -> Vec<(&'a Run, u64)> { + let mut newest: BTreeMap<(u64, &str), &Run> = BTreeMap::new(); + let on_branch = |r: &&'a Run| { + r.event == "push" + && matches!(r.status.as_str(), "queued" | "in_progress") + && r.head_branch.as_ref().is_some_and(|b| branches.contains(b)) + }; + for run in runs.iter().filter(on_branch) { + let key = ( + run.workflow_id, + run.head_branch.as_deref().unwrap_or_default(), + ); + let e = newest.entry(key).or_insert(run); + if run.id > e.id { + *e = run; + } + } + runs.iter() + .filter(on_branch) + .filter(|r| r.status == "queued") + .filter_map(|r| { + let n = newest[&(r.workflow_id, r.head_branch.as_deref().unwrap_or_default())]; + (n.id > r.id && n.head_sha != r.head_sha).then_some((r, n.id)) + }) + .collect() } pub(crate) async fn status_error( @@ -472,10 +616,17 @@ mod tests { Run { id, workflow_id: wf, + name: String::new(), + path: String::new(), status: status.into(), event: event.into(), head_sha: sha.into(), - pull_requests: prs.iter().map(|&number| PullRef { number }).collect(), + head_branch: Some("main".into()), + head_repository: None, + pull_requests: prs + .iter() + .map(|&number| PullRef { number, head: None }) + .collect(), } } @@ -488,6 +639,27 @@ mod tests { v } + #[test] + fn only_older_queued_pushes_on_listed_branches_are_superseded() { + let runs = [ + run(1, 1, "queued", "push", "a", &[]), + run(2, 1, "in_progress", "push", "b", &[]), + run(3, 1, "queued", "push", "c", &[]), + run(4, 2, "queued", "push", "c", &[]), + run(5, 2, "queued", "push", "c", &[]), + run(6, 3, "queued", "pull_request", "a", &[9]), + run(7, 3, "queued", "pull_request", "b", &[9]), + ]; + let mut got: Vec<_> = superseded_pushes(&runs, &["main".into()]) + .into_iter() + .map(|(r, n)| (r.id, n)) + .collect(); + got.sort_unstable(); + // Run 2 is running (never cancelled); run 4 shares run 5's commit. + assert_eq!(got, [(1, 3)]); + assert!(superseded_pushes(&runs, &["release".into()]).is_empty()); + } + #[test] fn keeps_lowest_queued_run() { let runs = [ diff --git a/src/probe/mod.rs b/src/probe/mod.rs index b9d5cd4..76b9904 100644 --- a/src/probe/mod.rs +++ b/src/probe/mod.rs @@ -99,7 +99,7 @@ const TOOLS_TTL_SECS: i64 = 60; const TOOL_VERSION_MAX: usize = 80; /// The first line of a tool's `--version` output, trimmed and clipped to -/// [`TOOL_VERSION_MAX`] characters; `"unknown"` when there is no output. +/// `TOOL_VERSION_MAX` characters; `"unknown"` when there is no output. pub fn tool_version_line(output: &str) -> String { let line = output.lines().map(str::trim).find(|l| !l.is_empty()); match line { @@ -140,7 +140,7 @@ pub fn tool_versions(exec: &dyn Exec) -> BTreeMap { .collect() } -/// The last tools probe, reused for [`TOOLS_TTL_SECS`]. +/// The last tools probe, reused for `TOOLS_TTL_SECS`. #[derive(Clone, Default)] pub struct ToolsCache(Arc>>); diff --git a/tests/janitor.rs b/tests/janitor.rs index 2554b81..9eaf82c 100644 --- a/tests/janitor.rs +++ b/tests/janitor.rs @@ -67,7 +67,107 @@ async fn an_already_finished_run_answering_409_is_not_an_error() { h.queue_pr_run(1, 10, "old", "new"); h.state().uncancellable.push(1); let r = h.janitor_pass().await; - assert_eq!((r.raced, r.cancelled, r.errors), (1, 0, 0), "{r:?}"); + assert_eq!( + (r.uncancellable, r.raced, r.cancelled, r.errors), + (1, 0, 0, 0), + "{r:?}" + ); +} + +#[tokio::test] +async fn the_run_s_own_pull_request_head_needs_no_pulls_api() { + // The queen's App may lack pull-request read access: with the head in the + // run's entry the pulls API is never asked. + let h = Harness::start(enabled(false, 20)).await; + h.state().pr_runs.push(json!({ + "id": 1, "status": "queued", "event": "pull_request", "head_sha": "old", + "pull_requests": [{"number": 10, "head": {"sha": "new"}}] + })); + h.state().pr_runs.push(json!({ + "id": 2, "status": "queued", "event": "pull_request", "head_sha": "new", + "pull_requests": [{"number": 10, "head": {"sha": "new"}}] + })); + let r = h.janitor_pass().await; + assert_eq!((r.superseded, r.cancelled, r.errors), (1, 1, 0), "{r:?}"); + assert_eq!(h.state().cancelled, [1]); +} + +#[tokio::test] +async fn runs_of_closed_pull_requests_are_cancelled_for_listed_workflows_only() { + let h = Harness::start(with(json!({ + "cancel_running": true, + "closed_pr_workflows": [".github/workflows/pr-gates.yml", "platform-ci"] + }))) + .await; + let closed = |id: u64, status: &str, path: &str, name: &str, repo: &str| { + json!({"id": id, "status": status, "event": "pull_request", "head_sha": "a", + "head_branch": "feat/x", "path": path, "name": name, + "head_repository": {"full_name": repo}, "pull_requests": []}) + }; + { + let mut st = h.state(); + st.pr_runs.push(closed( + 1, + "queued", + ".github/workflows/pr-gates.yml", + "pr-gates", + "example-org/app", + )); + st.pr_runs.push(closed( + 2, + "in_progress", + ".github/workflows/platform-ci.yml", + "platform-ci", + "example-org/app", + )); + // Not listed: a teardown that runs on the closed event. + st.pr_runs.push(closed( + 3, + "queued", + ".github/workflows/teardown.yml", + "teardown", + "example-org/app", + )); + // A fork's run never names a pull request either. + st.pr_runs.push(closed( + 4, + "queued", + ".github/workflows/pr-gates.yml", + "pr-gates", + "fork/r", + )); + } + let r = h.janitor_pass().await; + assert_eq!( + (r.closed_pr, r.cancelled, r.cancelled_running), + (2, 1, 1), + "{r:?}" + ); + let mut cancelled = h.state().cancelled.clone(); + cancelled.sort_unstable(); + assert_eq!(cancelled, [1, 2]); +} + +#[tokio::test] +async fn older_queued_main_pushes_are_cancelled_and_running_ones_kept() { + let h = Harness::start(with(json!({"superseded_push_branches": ["main"]}))).await; + let push = |id: u64, status: &str, sha: &str, branch: &str| { + json!({"id": id, "workflow_id": 1, "status": status, "event": "push", + "head_sha": sha, "head_branch": branch, "pull_requests": []}) + }; + { + let mut st = h.state(); + st.pr_runs.push(push(1, "queued", "a", "main")); + st.pr_runs.push(push(2, "queued", "b", "main")); + st.pr_runs.push(push(3, "queued", "c", "main")); + st.pr_runs.push(push(4, "queued", "a", "feat")); + st.pr_runs.push(push(5, "queued", "b", "feat")); + } + let r = h.janitor_pass().await; + assert_eq!((r.superseded_push, r.cancelled), (2, 2), "{r:?}"); + let mut cancelled = h.state().cancelled.clone(); + cancelled.sort_unstable(); + assert_eq!(cancelled, [1, 2]); } #[tokio::test] @@ -227,6 +327,8 @@ async fn totals_reach_status_and_metrics() { h.janitor_pass().await; let status = h.status().await; assert_eq!(status["janitor"]["passes"], 2); + assert!(status["janitor"]["since"].as_i64().unwrap() > 0); + assert!(status["janitor"]["last_pass_at"].as_i64().unwrap() > 0); assert_eq!(status["janitor"]["total"]["cancelled"], 1); assert_eq!(status["janitor"]["last"]["cancelled"], 0); let metrics = h From 3594dd62546faa3dde2de019677bccada71b78da Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Fri, 16 Oct 2026 21:32:57 +0000 Subject: [PATCH 21/27] chore: project borg from Mono 296fa1359ec1 --- .repository-projection.json | 6 +- docs/operations.md | 65 ++++++++++++++- src/api.rs | 161 +++++++++++++++++++++++++++++++++++- src/config.rs | 70 ++++++++++++++++ src/drone/caches/gc.rs | 46 ++++++++++- src/drone/caches/mod.rs | 39 +++++++-- src/drone/install.rs | 30 ++++++- src/queen.rs | 147 ++++++++++++++++++++++++++++++++ src/store.rs | 14 ++++ tests/capacity.rs | 115 ++++++++++++++++++++++++++ tests/units.rs | 10 +++ 11 files changed, 685 insertions(+), 18 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index f84ced8..245c3f4 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "91f10217da2a2a8dceb3aedbc480aab1eba1bc95", + "sourceSha": "296fa1359ec1145ea6d30631ec2ee02e38d9aa2e", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "bc61279726a1ea1773112633c4769aa7128ccefa", + "priorProjectedBase": "1c7af61c966a6ab9a7b2d566b3c1c2a65c12de4c", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", - "contentDigest": "021a908917cb1b06a2a4e35505a9faafce219868aa67d9a800c37ccb0de15c25", + "contentDigest": "d1f9aa12cfbdad0dbefe3fd520a413d63e09fb36cc7ac977092a62698df564e0", "publicationEligible": true } diff --git a/docs/operations.md b/docs/operations.md index 45ff441..2fbca43 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -392,7 +392,7 @@ See [remote execution](remote-execution.md) for remote-only runs and JSON results, [work-unit fairness](work-unit-fairness.md) for capacity shares, and [placement diagnostics](placement-diagnostics.md) for `borg explain`. -Optional (`units.enabled`, off by default). Clients such as `borg run`, the +Optional (`units.enabled`, off by default). While off, `/status` `units` shows only `enabled: false`, with empty counts and no recent units. Clients such as `borg run`, the pre-commit hook and `mq` submit whole commands; drones run them. Enroll a client with `borg queen enroll-client ` (prints a token once). @@ -528,6 +528,51 @@ Tune it with a `health` block in the queen config (`auto_quarantine`, `half_life_secs`, `window_secs`, `corroboration_runs`, `cooldown_secs`, `proof_jobs`, and the bounds); `/status` shows each drone's `health`. +### Drone availability and expiry + +A drone whose last heartbeat is older than `drone_stale_secs` (60) is +`stale`. It holds no schedulable capacity: it is left out of `slots`, +`free_slots` and placement, whatever its `health` says. `health` and +`profiles` describe past failures and speed only. Each drone in `/status` +carries `availability` (`alive` or `stale`), `last_seen_age_secs` and, for a +stale drone that will be removed, `expires_in_secs`. `stale_drones` lists the +stale names. When a drone goes stale, a `drone_stale` health event is logged. + +A drone that stays silent past its expiry is removed with its token and lease +records, and a `drone_expired` health event is logged. Configure the expiry in +the queen config: + +```json +"drone_expiry": {"after_secs": 604800, "burst_prefixes": ["gcp-burst-"], "burst_after_secs": 1800} +``` + +`after_secs` (7 days by default) applies to every drone. `burst_after_secs` +(30 minutes) applies to drones whose name starts with one of +`burst_prefixes`; cloud burst drones delete themselves and never come back +under the same token. 0 turns either expiry off. A drone named in +`wake.hosts`, or one that has never sent a heartbeat, is never removed. A +removed drone that comes back has to be enrolled again. + +`/status` `capacity` shows, per class, `schedulable_slots` (what the alive, +accepting drones could hold if empty) and `free_slots` next to `unmet_main`, +`unmet_pr` and `unmet_deploy`. The unmet counts are the last pass's +`unmet_placeable`. When main demand of a class in `capacity_alarm.classes` +stays unmet for `capacity_alarm.after_secs` while that class has at most +`capacity_alarm.max_slots` schedulable slots, the queen logs a +`capacity_starved` event (drone `fleet`; the reason names the stale drones). +When the condition clears, it logs `capacity_restored`. `capacity..alarm` +and the metric `borg_capacity_starved{class}` show whether the alarm is +raised, and `borg_unmet_jobs{class,scope}` exports the unmet counts. The +alarm is on by default for heavy and medium, with `max_slots` 2 and +`after_secs` 600: + +```json +"capacity_alarm": {"enabled": true, "classes": ["heavy", "medium"], "max_slots": 2, "after_secs": 600} +``` + +All of these events go to the same audit trail as quarantines: `health_log`, +shown as `health_events` in `/status`, and a `WARN` line in the queen journal. + ### Tailnet enrollment Optional and off by default. With a `tailnet_enroll` block @@ -1128,7 +1173,23 @@ sccache server per host (its configuration then applies to all jobs, and a job whose sccache version differs from the host binary's restarts it); the drone reads its totals with `sccache --show-stats --stats-format json` and reports hits and misses for `sccache` in heartbeats, `/status` and `borg_cache_hits`. -Without it those fields stay empty. +Without it, `hits` and `misses` are null and each cache's `counters` says +`unknown`. mono's `setup-sccache` action starts one server per job on its own +socket, so the real counts for those builds are in each job's log, in the +`sccache --show-stats` step. `counters` is `reported` when the numbers are +real, and `not_tracked` for kinds that have no hit counters. + +`size_status` says whether `size_bytes` covers the whole cache. `measured` +means every directory was listed. `partial` means some directories could not +be listed by the drone's user, so the size is a lower bound and the drone logs +a warning. With a separate job user, sccache creates its directories `0750` +and owned by the job user. `borg drone install` therefore adds the drone's +user to the job user's group. On a drone installed before that change, run +`usermod -aG borg-job borg` and restart the drone between jobs (a self-update +restarts it too). Until then, `/status` shows `partial` and a size far below +the real one. On a GCP burst drone on Oct 16, `du` as root measured 4.3G in +`sccache` while the drone reported 852K. `unknown` means the drone is older +than these fields. Every heartbeat carries each cache's size, cap, evicted bytes and last collection time (and, for git, how many leases were given a mirror and how many diff --git a/src/api.rs b/src/api.rs index ed4a506..35c170d 100644 --- a/src/api.rs +++ b/src/api.rs @@ -150,9 +150,45 @@ pub struct CacheStat { pub max_bytes: Option, pub evicted_bytes: u64, pub last_gc_at: Option, - /// Git mirrors only: leases that were handed a mirror / ran without one. + /// Git: leases that were handed a mirror / ran without one. sccache: + /// the host server's totals (`caches.sccache.server`). See `counters`. pub hits: Option, pub misses: Option, + /// Whether `size_bytes` covers the whole cache. + #[serde(default)] + pub size_status: SizeStatus, + /// Whether `hits` and `misses` are real counts. + #[serde(default)] + pub counters: Counters, +} + +/// How complete a cache's `size_bytes` is. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SizeStatus { + /// Every directory was listed. + Measured, + /// Some directories could not be listed by the drone's user (for + /// example the job user's 0750 sccache directories): a lower bound. + Partial, + /// The drone did not say (a drone older than this field). + #[default] + Unknown, +} + +/// What a cache's `hits` and `misses` mean. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Counters { + /// Real counts. + Reported, + /// Counts exist but the drone cannot read them: sccache without a host + /// server keeps its counters in each job's own server, which exits with + /// the job. Read the job's `sccache --show-stats` output instead. + #[default] + Unknown, + /// The cache kind has no hit counters. + NotTracked, } /// A running unit the drone stops listing is considered lost after this. @@ -679,6 +715,15 @@ pub struct DroneStatus { pub name: String, pub heartbeat_age_secs: Option, pub fresh: bool, + /// `alive` (heartbeat within `drone_stale_secs`) or `stale`. A stale + /// drone holds no schedulable capacity whatever its `health` says; + /// `health` is the historical failure evidence only. + pub availability: Availability, + /// Seconds since the last heartbeat; `None` if it never sent one. + pub last_seen_age_secs: Option, + /// Seconds until a stale drone is removed (`drone_expiry`); `None` while + /// alive or when it is never removed. + pub expires_in_secs: Option, pub accepting: bool, /// Enforced independently of the optional environment contract. pub disk_admission: crate::envcontract::DiskAdmission, @@ -716,6 +761,31 @@ pub struct DroneStatus { pub admission: Option, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum Availability { + Alive, + Stale, +} + +/// One class's schedulable capacity next to its unmet demand. +#[derive(Debug, Serialize)] +pub struct CapacityStatus { + /// Slots the alive, accepting drones could hold if empty. + pub schedulable_slots: u32, + /// Slots free right now on alive, accepting drones. + pub free_slots: u32, + /// Jobs borg could not place at the last pass (`unmet_placeable`). + pub unmet_main: u32, + pub unmet_pr: u32, + pub unmet_deploy: u32, + /// Since when main demand has been unmet with schedulable slots at or + /// under `capacity_alarm.max_slots`. + pub starved_since: Option, + /// The `capacity_starved` health event is raised and not yet restored. + pub alarm: bool, +} + /// What CPU pressure currently lets this drone take. #[derive(Debug, Serialize)] pub struct DroneAdmissionView { @@ -765,6 +835,10 @@ pub struct Status { pub classes: BTreeMap, pub queue: Vec, pub drones: Vec, + /// Per class: schedulable slots on alive drones next to unmet demand. + pub capacity: BTreeMap, + /// Names of drones whose heartbeat is older than `drone_stale_secs`. + pub stale_drones: Vec, /// Demand the allocator could not place. Counts jobs that non-borg runners /// may serve; see `unmet_placeable`. pub unmet: Vec, @@ -991,6 +1065,20 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result name: d.name.clone(), heartbeat_age_secs: age, fresh, + availability: if fresh { + Availability::Alive + } else { + Availability::Stale + }, + last_seen_age_secs: age, + expires_in_secs: match (fresh, age, cfg.drone_expiry.expiry_for(&d.name)) { + (false, Some(age), Some(after)) + if !cfg.wake.hosts.iter().any(|h| h.name == d.name) => + { + Some((after.max(cfg.drone_stale_secs) - age).max(0)) + } + _ => None, + }, accepting: d.accepting && !disk.blocked(), disk_admission: disk, trust: d.trust, @@ -1050,6 +1138,35 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result } let slots = queen.class_slots(&states); + let unmet_of = |class: Class, scope: Scope| -> u32 { + snap.unmet_placeable + .iter() + .filter(|u| u.class == class && u.scope == scope) + .map(|u| u.count) + .sum() + }; + let capacity: BTreeMap = Class::ALL + .into_iter() + .map(|class| { + ( + class, + CapacityStatus { + schedulable_slots: slots.get(&class).copied().unwrap_or(0), + free_slots: free_by_class.get(&class).copied().unwrap_or(0), + unmet_main: unmet_of(class, Scope::Main), + unmet_pr: unmet_of(class, Scope::Pr), + unmet_deploy: unmet_of(class, Scope::Deploy), + starved_since: snap.starved_since.get(&class).copied(), + alarm: snap.starved.contains(&class), + }, + ) + }) + .collect(); + let stale_drones: Vec = drones + .iter() + .filter(|d| d.availability == Availability::Stale) + .map(|d| d.name.clone()) + .collect(); let mut classes = BTreeMap::new(); for class in Class::ALL { let entry = |scope| { @@ -1105,6 +1222,8 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result classes, queue: snap.queue, drones, + capacity, + stale_drones, unmet: snap.unmet, unmet_placeable: snap.unmet_placeable, reserved_for: snap.reserved_for, @@ -1115,7 +1234,20 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result homeostasis: snap.homeostasis, profiles: snap.profiles, wake: snap.wake, - units: { + // Disabled units show only `enabled: false`: old records are not + // live state, and no unit is placed or served. + units: if !cfg.units.enabled { + UnitsStatus { + enabled: false, + counts: BTreeMap::new(), + joins: 0, + receipt_hits: 0, + superseded: 0, + oldest_queued_age_secs: 0, + recent: Vec::new(), + not_admitted: Vec::new(), + } + } else { let (joins, receipt_hits) = queen.unit_reuse_counts(); let store = queen.store(); let queued = store.units_in(&[crate::units::UnitState::Queued]).await?; @@ -1299,6 +1431,31 @@ pub fn render_metrics(status: &Status) -> String { .map(|(c, s)| (format!("class=\"{c}\""), f64::from(s.free_slots))) .collect(), ); + gauge( + "borg_unmet_jobs", + "Jobs borg could not place at the last pass, less what idle non-borg runners can take.", + status + .capacity + .iter() + .flat_map(|(c, v)| { + [ + ("main", v.unmet_main), + ("pr", v.unmet_pr), + ("deploy", v.unmet_deploy), + ] + .map(|(scope, n)| (format!("class=\"{c}\",scope=\"{scope}\""), f64::from(n))) + }) + .collect(), + ); + gauge( + "borg_capacity_starved", + "1 while the capacity_starved health event is raised for a class.", + status + .capacity + .iter() + .map(|(c, v)| (format!("class=\"{c}\""), f64::from(u8::from(v.alarm)))) + .collect(), + ); gauge( "borg_drone_free_mem_mib", "Memory budget left on a drone after its active leases.", diff --git a/src/config.rs b/src/config.rs index 8eaefc6..5033ab2 100644 --- a/src/config.rs +++ b/src/config.rs @@ -108,6 +108,70 @@ fn default_janitor_max() -> u32 { 20 } +/// A drone silent past `drone_stale_secs` is `stale`: it holds no +/// schedulable capacity. One silent past its expiry is removed (token and +/// lease records with it) and the removal is logged as a health event. +/// Drones named in `wake.hosts` are never removed. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields, default)] +pub struct DroneExpirySettings { + /// Remove a drone silent this long; 0 never removes. + pub after_secs: i64, + /// Name prefixes of short-lived cloud drones (for example `gcp-burst-`). + pub burst_prefixes: Vec, + /// Remove a drone matching `burst_prefixes` silent this long; 0 never. + pub burst_after_secs: i64, +} + +impl Default for DroneExpirySettings { + fn default() -> Self { + Self { + after_secs: 7 * 86_400, + burst_prefixes: Vec::new(), + burst_after_secs: 1800, + } + } +} + +impl DroneExpirySettings { + /// The silence after which `name` is removed, if ever. + pub fn expiry_for(&self, name: &str) -> Option { + let secs = if self + .burst_prefixes + .iter() + .any(|p| name.starts_with(p.as_str())) + { + self.burst_after_secs + } else { + self.after_secs + }; + (secs > 0).then_some(secs) + } +} + +/// Raise a `capacity_starved` health event when main demand of a class stays +/// unmet while the fleet's schedulable slots for it are at most `max_slots` +/// for `after_secs`, and `capacity_restored` when that ends. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields, default)] +pub struct CapacityAlarmSettings { + pub enabled: bool, + pub classes: Vec, + pub max_slots: u32, + pub after_secs: i64, +} + +impl Default for CapacityAlarmSettings { + fn default() -> Self { + Self { + enabled: true, + classes: vec![Class::Heavy, Class::Medium], + max_slots: 2, + after_secs: 600, + } + } +} + /// Standby deploy runners; present means enabled. See `docs/operations.md`. #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] @@ -187,6 +251,12 @@ pub struct QueenConfig { /// A drone silent for longer than this gets no new leases. #[serde(default = "default_drone_stale")] pub drone_stale_secs: i64, + /// When silent drones are removed; see [`DroneExpirySettings`]. + #[serde(default)] + pub drone_expiry: DroneExpirySettings, + /// The fleet-capacity alarm; see [`CapacityAlarmSettings`]. + #[serde(default)] + pub capacity_alarm: CapacityAlarmSettings, /// Demand older than this triggers the starvation guard. `false` (or JSON /// `null`) disables it; TOML has no null. #[serde(default = "default_starvation", deserialize_with = "secs_or_off")] diff --git a/src/drone/caches/gc.rs b/src/drone/caches/gc.rs index 70f9aff..c4d066a 100644 --- a/src/drone/caches/gc.rs +++ b/src/drone/caches/gc.rs @@ -35,8 +35,25 @@ pub fn parse_size(text: &str) -> Option { /// Total size and the newest file modification/access time under `path`. /// Symlinks are counted as themselves, never followed. pub fn measure(path: &Path) -> (u64, SystemTime) { + let m = measure_tree(path); + (m.size, m.newest) +} + +/// What [`measure_tree`] found. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Measured { + pub size: u64, + pub newest: SystemTime, + /// Directories that could not be listed (for example owned by the job + /// user and closed to the drone's user): `size` is then a lower bound. + pub unreadable_dirs: u64, +} + +/// [`measure`], also counting directories that could not be listed. +pub fn measure_tree(path: &Path) -> Measured { let mut size = 0; let mut newest = SystemTime::UNIX_EPOCH; + let mut unreadable_dirs = 0; let mut stack = vec![path.to_owned()]; while let Some(dir) = stack.pop() { let Ok(meta) = fs::symlink_metadata(&dir) else { @@ -46,8 +63,11 @@ pub fn measure(path: &Path) -> (u64, SystemTime) { if meta.is_dir() { // Directory times only record when entries were added or listed; // so recency comes from the files. - if let Ok(entries) = fs::read_dir(&dir) { - stack.extend(entries.flatten().map(|e| e.path())); + match fs::read_dir(&dir) { + Ok(entries) => stack.extend(entries.flatten().map(|e| e.path())), + // Removed between the listing and now: nothing to count. + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(_) => unreadable_dirs += 1, } } else { for t in [meta.modified(), meta.accessed()].into_iter().flatten() { @@ -55,7 +75,11 @@ pub fn measure(path: &Path) -> (u64, SystemTime) { } } } - (size, newest) + Measured { + size, + newest, + unreadable_dirs, + } } /// One evictable entry. @@ -162,6 +186,22 @@ mod tests { } } + #[cfg(unix)] + #[test] + fn an_unlistable_directory_makes_the_size_a_lower_bound() { + use std::os::unix::fs::PermissionsExt; + let dir = tempfile::tempdir().unwrap(); + file(&dir.path().join("open/a"), 10, 5); + file(&dir.path().join("closed/b"), 10, 5); + let closed = dir.path().join("closed"); + fs::set_permissions(&closed, fs::Permissions::from_mode(0o000)).unwrap(); + let m = measure_tree(dir.path()); + let listable = fs::read_dir(&closed).is_ok(); // root reads anything + fs::set_permissions(&closed, fs::Permissions::from_mode(0o755)).unwrap(); + assert_eq!(m.unreadable_dirs, u64::from(!listable)); + assert_eq!(measure_tree(dir.path()).unreadable_dirs, 0); + } + #[test] fn measure_sums_a_tree_and_finds_the_newest_file() { let dir = tempfile::tempdir().unwrap(); diff --git a/src/drone/caches/mod.rs b/src/drone/caches/mod.rs index e25573c..d7f15b1 100644 --- a/src/drone/caches/mod.rs +++ b/src/drone/caches/mod.rs @@ -21,7 +21,7 @@ use std::time::Duration; use tokio::sync::{OnceCell, OwnedRwLockReadGuard, RwLock}; -use crate::api::CacheStat; +use crate::api::{CacheStat, Counters, SizeStatus}; use crate::drone::config::DroneConfig; use crate::drone::exec::Exec; use crate::drone::fsops::FsOps; @@ -59,6 +59,8 @@ struct Track { /// sccache only, from the host server. hits: Option, misses: Option, + /// The last measurement missed directories the drone could not list. + partial: bool, } pub struct Caches { @@ -296,9 +298,10 @@ impl Caches { let root = self.cfg.cache_root.clone(); let fs_ops = Arc::clone(&self.fs_ops); let work = tokio::task::spawn_blocking(move || { - let mut out: Vec<(&'static str, u64, u64)> = Vec::new(); + let mut out: Vec<(&'static str, u64, u64, bool)> = Vec::new(); for (kind, cap) in plan { - let (mut size, _) = gc::measure(&kind.dir(&root)); + let measured = gc::measure_tree(&kind.dir(&root)); + let mut size = measured.size; let mut freed = 0; if let Some(cap) = cap && size > cap @@ -308,19 +311,25 @@ impl Caches { freed = gc::evict_lru(fs_ops.as_ref(), units, pool, cap).freed_bytes; size = size.saturating_sub(freed); } - out.push((kind.name(), size, freed)); + out.push((kind.name(), size, freed, measured.unreadable_dirs > 0)); } - let (git_size, _) = gc::measure(&root.join("git")); - out.push(("git", git_size, 0)); + let git = gc::measure_tree(&root.join("git")); + out.push(("git", git.size, 0, git.unreadable_dirs > 0)); out }) .await; let Ok(results) = work else { return }; let now = unix_now(); if let Ok(mut tracks) = self.tracks.lock() { - for (name, size, freed) in results { + for (name, size, freed, partial) in results { + if partial { + tracing::warn!( + "cache {name}: some directories are not readable by the drone's user; size {size} is a lower bound" + ); + } let t = tracks.entry(name).or_default(); t.size_bytes = size; + t.partial = partial; if freed > 0 { t.evicted_bytes += freed; t.last_gc_at = Some(now); @@ -351,6 +360,11 @@ impl Caches { } else { (t.hits, t.misses) }; + let counters = match (name, hits) { + (_, Some(_)) => Counters::Reported, + ("sccache", None) => Counters::Unknown, + _ => Counters::NotTracked, + }; CacheStat { name: name.to_owned(), size_bytes: t.size_bytes, @@ -359,6 +373,12 @@ impl Caches { last_gc_at: t.last_gc_at, hits, misses, + size_status: if t.partial { + SizeStatus::Partial + } else { + SizeStatus::Measured + }, + counters, } }) .collect() @@ -592,6 +612,7 @@ mod tests { c.maintain().await; let git = c.stats().into_iter().find(|s| s.name == "git").unwrap(); assert_eq!((git.hits, git.misses), (Some(1), Some(0))); + assert_eq!(git.counters, Counters::Reported); } #[tokio::test] @@ -638,6 +659,10 @@ mod tests { Some(100 << 30), "sccache reports its own cap" ); + // No host server: the counts live in each job's own server. + assert_eq!((sccache.hits, sccache.counters), (None, Counters::Unknown)); + assert_eq!(sccache.size_status, SizeStatus::Measured); + assert_eq!(cargo.counters, Counters::NotTracked); } #[tokio::test] diff --git a/src/drone/install.rs b/src/drone/install.rs index f4280b5..d2fc355 100644 --- a/src/drone/install.rs +++ b/src/drone/install.rs @@ -580,6 +580,30 @@ impl Installer<'_> { && job != daemon { self.ensure_user(&job, &self.layout.job_home.clone())?; + // The job user's caches (sccache makes its directories 0750) must + // be readable by the drone to measure them; without the group the + // drone reports a fraction of their size. + if self.opts.platform == Platform::Linux { + self.ensure_in_group(&daemon, &job)?; + } + } + Ok(()) + } + + fn ensure_in_group(&mut self, user: &str, group: &str) -> Result<(), InstallError> { + let member = self + .exec + .run("id", &["-nG", user]) + .is_ok_and(|o| o.success && o.stdout.split_whitespace().any(|g| g == group)); + if member { + self.note("unchanged", format_args!("user {user} in group {group}")); + } else if self.opts.dry_run { + self.changed = true; + self.note("would add", format_args!("user {user} to group {group}")); + } else { + self.changed = true; + self.must("usermod", &["-aG", group, user])?; + self.note("added", format_args!("user {user} to group {group}")); } Ok(()) } @@ -1043,7 +1067,7 @@ mod tests { fn exec() -> FakeExec { let x = FakeExec::default(); x.on("getent group docker", "docker:x:999:"); - x.on("id -nG borg", "borg docker"); + x.on("id -nG borg", "borg docker borg-job"); x.on("id -nG borg-job", "borg-job docker"); x } @@ -1374,6 +1398,10 @@ mod tests { "the job user is never given the drone's files: {calls:?}" ); assert!(calls.contains(&"usermod -aG docker borg-job".to_owned())); + assert!( + calls.contains(&"usermod -aG borg-job borg".to_owned()), + "the drone reads the job user's caches: {calls:?}" + ); let unit = fs::read_to_string(e.root.join("etc/systemd/system/borg-drone.service")).unwrap(); assert!( diff --git a/src/queen.rs b/src/queen.rs index c3d1741..b983335 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -95,6 +95,22 @@ pub struct Snapshot { pub start_standby: BTreeSet, /// `unmet` minus what idle non-borg runners with the label can take. pub unmet_placeable: Vec, + /// Per class, slots the fresh accepting drones could hold if empty. + pub schedulable_slots: BTreeMap, + /// Per class, since when main demand has been unmet with schedulable + /// slots at or under `capacity_alarm.max_slots`. + pub starved_since: BTreeMap, + /// Classes whose `capacity_starved` event is raised and not yet restored. + pub starved: BTreeSet, +} + +/// What the fleet watch remembers between passes. +#[derive(Debug, Default)] +struct FleetWatch { + /// Fresh drones at the previous pass; `None` before the first. + prev_fresh: Option>, + starved_since: BTreeMap, + raised: BTreeSet, } /// Runners in the organization that borg did not mint, for one label. @@ -212,6 +228,7 @@ pub struct Queen { unit_cursor: std::sync::Mutex>, wake_sender: Arc, tick_lock: tokio::sync::Mutex<()>, + watch: std::sync::Mutex, } /// The full label set a JIT runner is registered with: the labels GitHub @@ -393,6 +410,7 @@ impl Queen { unit_cursor: std::sync::Mutex::new(BTreeMap::new()), wake_sender: Arc::new(UdpWakeSender), tick_lock: tokio::sync::Mutex::new(()), + watch: std::sync::Mutex::new(FleetWatch::default()), } } @@ -540,6 +558,9 @@ impl Queen { tracing::warn!("rollout pass failed: {e}"); } let profiles = self.profiles(now_ts).await; + if let Err(e) = self.expire_drones(now_ts).await { + tracing::warn!("drone expiry failed: {e}"); + } let drone_rows = self.store.drones().await?; let mut fresh_names: BTreeSet = BTreeSet::new(); let mut capacities: HashMap = HashMap::new(); @@ -597,6 +618,7 @@ impl Queen { ); } } + let schedulable_slots = self.class_slots(&drones); if let Err(e) = self.advance_units(now_ts, &fresh_names).await { tracing::warn!("unit pass failed: {e}"); } @@ -730,6 +752,15 @@ impl Queen { // hosts for. let unmet_placeable = placeable_unmet(&allocation.unmet, &statics, &self.cfg.classes); let wake = self.wake_pass(now_ts, &unmet_placeable, &fresh_names).await; + let (starved_since, starved) = self + .watch_fleet( + now_ts, + &drone_rows, + &fresh_names, + &schedulable_slots, + &unmet_placeable, + ) + .await; let mut mint_error = None; let mut lease_error = None; let mut persisted = BTreeSet::new(); @@ -917,6 +948,9 @@ impl Queen { s.static_runners = statics; s.start_standby = start_standby; s.unmet_placeable = unmet_placeable; + s.schedulable_slots = schedulable_slots; + s.starved_since = starved_since; + s.starved = starved; }); if let Err(e) = self .store @@ -1931,6 +1965,119 @@ impl Queen { } /// Per-class totals for status and metrics: slots the fleet could hold. + /// Remove drones silent past their `drone_expiry` window (never one named + /// in `wake.hosts`, never one that has not heartbeated yet) and log each + /// removal as a `drone_expired` health event. + async fn expire_drones(&self, now: i64) -> Result<(), StoreError> { + for d in self.store.drones().await? { + let Some(last) = d.last_heartbeat else { + continue; + }; + let Some(after) = self.cfg.drone_expiry.expiry_for(&d.name) else { + continue; + }; + let silent = now - last; + if silent <= after.max(self.cfg.drone_stale_secs) + || self.cfg.wake.hosts.iter().any(|h| h.name == d.name) + { + continue; + } + if let Some(dropped) = self.store.remove_drone(&d.name).await? { + let reason = format!( + "no heartbeat for {silent}s (expiry {after}s); {dropped} unfinished leases dropped" + ); + tracing::warn!("expired drone {}: {reason}", d.name); + self.store + .record_health_event(now, &d.name, "drone_expired", &reason) + .await?; + } + } + Ok(()) + } + + /// Log drones that went stale since the last pass, and raise or clear the + /// capacity alarm per class. Returns the alarm state for the snapshot. + async fn watch_fleet( + &self, + now: i64, + rows: &[DroneRow], + fresh: &BTreeSet, + slots: &BTreeMap, + unmet: &[Unmet], + ) -> (BTreeMap, BTreeSet) { + let alarm = &self.cfg.capacity_alarm; + let stale: Vec<&str> = rows + .iter() + .filter(|d| !fresh.contains(&d.name)) + .map(|d| d.name.as_str()) + .collect(); + let mut events: Vec<(String, &'static str, String)> = Vec::new(); + let state = { + let mut w = self.watch.lock().expect("fleet watch lock"); + if let Some(prev) = &w.prev_fresh { + for d in rows + .iter() + .filter(|d| prev.contains(&d.name) && !fresh.contains(&d.name)) + { + let age = d.last_heartbeat.map_or(-1, |t| now - t); + events.push(( + d.name.clone(), + "drone_stale", + format!("no heartbeat for {age}s"), + )); + } + } + w.prev_fresh = Some(fresh.clone()); + for class in Class::ALL { + let main: u32 = unmet + .iter() + .filter(|u| u.class == class && u.scope == Scope::Main) + .map(|u| u.count) + .sum(); + let have = slots.get(&class).copied().unwrap_or(0); + if alarm.enabled + && alarm.classes.contains(&class) + && main > 0 + && have <= alarm.max_slots + { + let since = *w.starved_since.entry(class).or_insert(now); + if now - since >= alarm.after_secs && w.raised.insert(class) { + events.push(( + "fleet".to_owned(), + "capacity_starved", + format!( + "{class}: {main} main jobs unmet for {}s with {have} schedulable slots; stale drones: [{}]", + now - since, + stale.join(", ") + ), + )); + } + } else { + w.starved_since.remove(&class); + if w.raised.remove(&class) { + events.push(( + "fleet".to_owned(), + "capacity_restored", + format!("{class}: {have} schedulable slots, {main} main jobs unmet"), + )); + } + } + } + (w.starved_since.clone(), w.raised.clone()) + }; + for (drone, event, reason) in events { + tracing::warn!("{event} {drone}: {reason}"); + if let Err(e) = self + .store + .record_health_event(now, &drone, event, &reason) + .await + { + tracing::warn!("record {event} failed: {e}"); + } + } + state + } + pub(crate) fn class_slots(&self, drones: &[DroneState]) -> BTreeMap { Class::ALL .into_iter() diff --git a/src/store.rs b/src/store.rs index d24198f..fff1d44 100644 --- a/src/store.rs +++ b/src/store.rs @@ -785,6 +785,18 @@ impl Store { Ok(true) } + /// Append a fleet or drone event (stale, expired, capacity) to the + /// health audit trail that `/status` shows as `health_events`. + pub async fn record_health_event( + &self, + at: i64, + drone: &str, + event: &str, + reason: &str, + ) -> Result<(), StoreError> { + self.log_health(at, drone, event, Some(reason), None).await + } + async fn log_health( &self, at: i64, @@ -1768,6 +1780,8 @@ VALUES ('l1', 'd1', 'medium', 'pr', '[\"a\"]', 'borg-d1-l1', 7, 'jit', 'pending' last_gc_at: Some(100), hits: None, misses: None, + size_status: crate::api::SizeStatus::Measured, + counters: crate::api::Counters::NotTracked, }; s.record_heartbeat( "d1", diff --git a/tests/capacity.rs b/tests/capacity.rs index e871659..bb60075 100644 --- a/tests/capacity.rs +++ b/tests/capacity.rs @@ -127,3 +127,118 @@ async fn wake_on_lan_ignores_demand_that_non_borg_runners_will_take() { } assert_eq!(h.wakes.0.lock().unwrap().len(), 1); } + +fn events(s: &Value) -> Vec<(String, String)> { + s["health_events"] + .as_array() + .unwrap() + .iter() + .map(|e| { + ( + e["drone"].as_str().unwrap().to_owned(), + e["event"].as_str().unwrap().to_owned(), + ) + }) + .collect() +} + +async fn status_at(h: &Harness, at: OffsetDateTime) -> Value { + let s = borg::api::build_status(&h.queen, at.unix_timestamp()) + .await + .unwrap(); + serde_json::to_value(s).unwrap() +} + +#[tokio::test] +async fn silent_drones_are_stale_then_expired_burst_drones_first() { + let h = Harness::start(json!({ + "drone_expiry": {"after_secs": 3600, "burst_prefixes": ["gcp-burst-"], "burst_after_secs": 600}, + "capacity_alarm": {"enabled": false} + })) + .await; + let burst = h.enroll("gcp-burst-1").await; + let onprem = h.enroll("pve-1").await; + h.heartbeat(&burst, &[]).await; + h.heartbeat(&onprem, &[]).await; + let t0 = OffsetDateTime::now_utc(); + h.queen.tick(t0).await.unwrap(); + + let s = status_at(&h, t0).await; + assert_eq!(drone(&s, "gcp-burst-1")["availability"], "alive"); + assert_eq!(s["capacity"]["heavy"]["schedulable_slots"], 2); + assert_eq!(s["stale_drones"], json!([])); + + // Silent past drone_stale_secs: stale, no schedulable capacity, health + // untouched, and one drone_stale event each. + let t1 = t0 + time::Duration::seconds(120); + h.queen.tick(t1).await.unwrap(); + let s = status_at(&h, t1).await; + let d = drone(&s, "gcp-burst-1"); + assert_eq!(d["availability"], "stale"); + assert!(d["last_seen_age_secs"].as_i64().unwrap() >= 120); + assert!(d["expires_in_secs"].as_i64().unwrap() <= 480); + assert_eq!(s["capacity"]["heavy"]["schedulable_slots"], 0); + assert_eq!(s["stale_drones"], json!(["gcp-burst-1", "pve-1"])); + let ev = events(&s); + assert!( + ev.contains(&("gcp-burst-1".into(), "drone_stale".into())), + "{ev:?}" + ); + assert!( + ev.contains(&("pve-1".into(), "drone_stale".into())), + "{ev:?}" + ); + + // Past the burst expiry only the burst drone is removed. + let t2 = t0 + time::Duration::seconds(700); + h.queen.tick(t2).await.unwrap(); + let s = status_at(&h, t2).await; + let names: Vec<&str> = s["drones"] + .as_array() + .unwrap() + .iter() + .map(|d| d["name"].as_str().unwrap()) + .collect(); + assert_eq!(names, ["pve-1"]); + assert!(events(&s).contains(&("gcp-burst-1".into(), "drone_expired".into()))); + assert!(h.heartbeat(&burst, &[]).await.status().is_client_error()); + + let t3 = t0 + time::Duration::seconds(3700); + h.queen.tick(t3).await.unwrap(); + assert_eq!(status_at(&h, t3).await["drones"], json!([])); +} + +#[tokio::test] +async fn unmet_main_demand_without_capacity_raises_and_clears_the_alarm() { + let h = Harness::start(json!({ + "capacity_alarm": {"after_secs": 60, "max_slots": 0} + })) + .await; + h.queue_job(1, "borg-ci-heavy", 5); + let t0 = OffsetDateTime::now_utc(); + h.queen.tick(t0).await.unwrap(); + let s = status_at(&h, t0).await; + assert_eq!(s["capacity"]["heavy"]["unmet_main"], 1); + assert_eq!(s["capacity"]["heavy"]["schedulable_slots"], 0); + assert_eq!(s["capacity"]["heavy"]["alarm"], false); + assert!(s["capacity"]["heavy"]["starved_since"].is_i64()); + + let t1 = t0 + time::Duration::seconds(61); + h.queen.tick(t1).await.unwrap(); + let s = status_at(&h, t1).await; + assert_eq!(s["capacity"]["heavy"]["alarm"], true); + assert!(events(&s).contains(&("fleet".into(), "capacity_starved".into()))); + assert!( + h.metrics() + .await + .contains("borg_capacity_starved{class=\"heavy\"} 1") + ); + + // A drone that can take the job arrives: the alarm clears. + let t = h.enroll("d1").await; + h.heartbeat(&t, &[]).await; + h.tick().await; + let s = h.status().await; + assert_eq!(s["capacity"]["heavy"]["alarm"], false); + assert!(events(&s).contains(&("fleet".into(), "capacity_restored".into()))); +} diff --git a/tests/units.rs b/tests/units.rs index e5a3ba0..932898e 100644 --- a/tests/units.rs +++ b/tests/units.rs @@ -1583,3 +1583,13 @@ async fn pending_pr_runner_covering_main_does_not_inflate_deploy_pressure() { .contains(&Rejection::BidThreshold) }))); } + +#[tokio::test] +async fn disabled_units_show_only_that_they_are_off() { + let h = support::Harness::start(json!({})).await; + let s = h.status().await; + assert_eq!(s["units"]["enabled"], false); + assert_eq!(s["units"]["counts"], json!({})); + assert_eq!(s["units"]["recent"], json!([])); + assert!(!h.metrics().await.contains("borg_units{")); +} From 1824b8da7c58ab90042efec2d93c78a2f12b25be Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 17 Oct 2026 01:27:05 +0000 Subject: [PATCH 22/27] chore: project borg from Mono 54af7346321a --- .repository-projection.json | 6 +- docs/operations.md | 121 ++++- src/allocator.rs | 158 +++++- src/api.rs | 91 +++- src/config.rs | 22 + src/drone/mod.rs | 77 ++- src/envcontract.rs | 10 + src/hostdiag.rs | 978 ++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/preempt.rs | 1 + src/probe/mod.rs | 28 +- src/queen.rs | 30 +- src/store.rs | 21 + tests/drone.rs | 13 + tests/queen.rs | 162 ++++++ 15 files changed, 1676 insertions(+), 43 deletions(-) create mode 100644 src/hostdiag.rs diff --git a/.repository-projection.json b/.repository-projection.json index 245c3f4..36e2b12 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "296fa1359ec1145ea6d30631ec2ee02e38d9aa2e", + "sourceSha": "54af7346321ae5f3a1aeb1ae703bd34e80c204ec", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "1c7af61c966a6ab9a7b2d566b3c1c2a65c12de4c", + "priorProjectedBase": "ec3cb38cabf0371e8e514cab74a1d0b494c02395", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", - "contentDigest": "d1f9aa12cfbdad0dbefe3fd520a413d63e09fb36cc7ac977092a62698df564e0", + "contentDigest": "cd2134012915159bfbe895ad883939b40f10ccea2675fd6869366f7c13de4979", "publicationEligible": true } diff --git a/docs/operations.md b/docs/operations.md index 2fbca43..c3e72c6 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -10,6 +10,7 @@ Configuration, scheduling, drone administration, and protocol details for - [Config files](#config-files) - [Queen](#queen) - [Job telemetry](#job-telemetry) + - [Per-job host diagnostics](#per-job-host-diagnostics) - [Drone bids and demand pressure](#drone-bids-and-demand-pressure) - [Wake-on-LAN](#wake-on-lan) - [Build environment contract](#build-environment-contract) @@ -236,6 +237,66 @@ their meaning. Keep it consistent across drones: `duration_secs`. A retried result may be sent again; the queen keeps the latest. +### Per-job host diagnostics + +A job that fails on an overloaded drone looks like a code failure unless +someone compares runners by hand. The drone samples its host just before the +runner starts and again when it exits, and sends both samples with the lease +result (`host: {start, end}` next to `exit_code`). Each sample has: + +| Field | Source | +| --- | --- | +| `at`, `cpus` | Unix seconds; CPUs the drone may use. | +| `load` | `/proc/loadavg` 1, 5 and 15 minute averages. Inside an LXC, lxcfs computes this and it over-reports, so use the pressure and busy fields instead. | +| `cpu_psi`, `memory_psi`, `io_psi` | `/proc/pressure/*`, `some`/`full` `avg10`/`avg60` in percent. | +| `cgroup_cpu_psi`, `cgroup_memory_psi` | `/sys/fs/cgroup/*.pressure`: the cgroup-namespace root, which is the whole container in an LXC. | +| `host_cpu_usage_usec`, `host_cpu_busy_cores` | `usage_usec` of `/sys/fs/cgroup/cpu.stat`, and the cores busy since the drone's previous reading (heartbeats feed it, so at most one heartbeat old). | +| `drone_cpu_usage_usec`, `drone_cpu_throttled_usec` | `cpu.stat` of the drone's own cgroup (`/proc/self/cgroup`). Jobs run inside it. | +| `mem_available_mib`, `work_free_mib`, `cache_free_mib` | `MemAvailable`; `df` of `work_root` and `cache_root`. | +| `active_leases` | Leases the drone held at that moment, this one included. | + +Every field is optional, and a field the drone cannot read is left out. A drone +that predates this sends no `host`, and the row still records which drone ran +the job. + +The queen stores one row per lease in `lease_hosts`. It adds the GitHub run id, +job id and job name when its queue poll sees the job on a runner it minted +(the same match as [GitHub job feedback](#github-job-feedback)). Rows are kept +for 7 days and at most 50,000 rows, oldest first. Two read endpoints are +unauthenticated like `/status`: + +- `GET /jobs/{job_id}` returns the lease that ran the job, or 404. +- `GET /runs/{run_id}` returns `{run_id, jobs: [...]}`, oldest lease first. + +```json +{ + "lease_id": "a1b2c3", "drone": "shared-box", "class": "heavy", "scope": "main", + "repo": "example-org/app", "run_id": 123456789, "job_id": 987654321, "job_name": "ui-e2e (shard 3)", + "leased_at": 1760650000, "finished_at": 1760651800, "duration_secs": 1790, "exit_code": 1, + "start": { "at": 1760650010, "cpus": 16, "load": [77.1, 60.2, 50.3], + "cgroup_cpu_psi": { "some_avg10": 54.9, "some_avg60": 52.4 }, + "host_cpu_busy_cores": 15.2, "active_leases": 2, "…": "…" }, + "end": { "…": "…" }, + "during": { "secs": 1790, "host_cpu_busy_cores": 15.5, "drone_cpu_busy_cores": 3.0, + "drone_cpu_throttled_secs": 0.0 }, + "slow_host": true, + "slow_host_reasons": ["cpu pressure 54.9% at start (threshold 30%)"] +} +``` + +`slow_host` is true when any of these held: CPU pressure (`some avg10`, the +higher of the host and cgroup files) at start at or above the admission soft +threshold (`psi_soft_pm / 10`, or 30% without an `admission` block); memory +pressure at start at or above 10%; load average per CPU at start at or above +1.5, used only when the host has no CPU pressure reading; 90% of the CPUs busy +at start or on average during the job; less than 10 GiB free on the work or +cache volume at start. `during` is computed from the two samples. + +The job id is known only if a queue poll saw the job while it ran or a direct +lookup found it afterwards (see [GitHub job feedback](#github-job-feedback)). +A job that starts and finishes between polls has a row without `job_id`, and +`/jobs/{job_id}` returns 404 for it. + ### Drone bids and demand pressure A drone may add `bids` to its heartbeat: at most one per `(class, scope)` it is @@ -349,6 +410,8 @@ old as unknown, so probe at least that often): | `toolchains` | For each toolchain `rustup toolchain list` shows, `{channel, rustc_version}`: the channel as `rust-toolchain.toml` would name it (`1.95.0`, `stable`, `nightly-2026-01-01`; strip the host triple) and the `release:` line of `rustc + -vV`. `null` if rustup could not be queried. | | `mem_available_mib` | Linux `MemAvailable` from `/proc/meminfo`; macOS (free + inactive + speculative pages) times the page size, in MiB. | | `cpu_psi_some_avg10_pm` | Linux `/proc/pressure/cpu` line `some`, `avg10` times ten, rounded; absent where PSI does not exist. | +| `load_per_cpu_pm` | `/proc/loadavg` 1-minute average divided by the usable CPUs, times 1000. | +| `cpu_busy_pm` | Share of the CPUs busy since the previous heartbeat, from `usage_usec` in `/sys/fs/cgroup/cpu.stat`, in permille. | | `tools` | For each of `git`, `gh`, `node`, `npm`, `pnpm`, `python3`, `docker`, `cargo`, `rustc`, `bazel`, `go`, `uv`, `mise`, `jq`, `cosign`: resolve the executable on the job user's `PATH` (as the job user, with the job environment) and record the first line of ` --version` (`go version` for `go`), trimmed, at most 80 characters. `"missing"` when it is not on `PATH`, `"unknown"` when it is but prints nothing. The map is absent when the probe did not run. | A unit's toolchain (`rust_toolchain`) matches an installed `channel` or @@ -873,23 +936,55 @@ one saturated host can all show free slots. Optional and off unless the queen config has an `admission` block: ```json -"admission": { "psi_soft_pm": 300, "psi_hard_pm": 600, "facts_max_age_secs": 180 } +"admission": { "psi_soft_pm": 300, "psi_light_only_pm": 500, "psi_hard_pm": 600, "facts_max_age_secs": 180 } ``` Each drone reports Linux PSI `cpu some avg10` in permille on every heartbeat (the higher of `/proc/pressure/cpu` and the cgroup file -`/sys/fs/cgroup/cpu.pressure`). At or above `psi_soft_pm` the drone takes no -new pull-request grant; main work is still placed. At or above `psi_hard_pm` it -takes no new grant of any scope. Leases already granted are not touched. A -reading older than `facts_max_age_secs`, or no reading, puts no restriction on -the drone. The starvation guard obeys the same rule: a starving pull-request -lease is not placed on a soft-blocked drone. Requires `0 < psi_soft_pm <= -psi_hard_pm <= 1000` and a positive age. - -`/status` shows `admission: {state, cpu_psi_pm}` per drone (`state` is `open`, -`pr_blocked` or `closed`). -Metrics: `borg_drone_cpu_psi_pm{drone}` and `borg_drone_admission{drone,state}`. -The build environment contract's `resources` check is separate and still works. +`/sys/fs/cgroup/cpu.pressure`). The queen compares it with three thresholds: + +| Pressure | State | New grants | +| --- | --- | --- | +| below `psi_soft_pm` | `open` | all | +| at or above `psi_soft_pm` | `pr_blocked` | no pull-request work | +| at or above `psi_light_only_pm` | `light_only` | no pull-request work and no main medium or heavy work; main light and deploy work still placed | +| at or above `psi_hard_pm` | `closed` | none | + +`psi_light_only_pm` is optional and must lie between the other two. Without +it a drone goes from `pr_blocked` straight to `closed`. Leases already granted +are not touched. A reading older than `facts_max_age_secs`, or no reading at +all, puts no restriction on the drone. The starvation guard obeys the same +rule: a starving pull-request lease is not placed on a soft-blocked drone, and +a starving main heavy lease is not placed on a `light_only` drone. Requires +`0 < psi_soft_pm <= psi_hard_pm <= 1000` and a positive age. + +**Hosts without PSI.** A kernel without PSI (or booted with `psi=0`) reports +no pressure. Such a drone would never be blocked, however loaded. The drone +therefore also reports `load_per_cpu_pm`, the 1-minute load average divided by +its usable CPUs in permille, and `cpu_busy_pm`, the share of CPUs busy since the +previous heartbeat from the root cgroup's `cpu.stat`. When PSI is missing the +queen estimates pressure as `(load - cpus) / load`, the share of runnable tasks +that wait for a CPU (load 93 on 16 CPUs gives 828 permille). The estimate is 0 +when the load is at or below one task per CPU, or when `cpu_busy_pm` is under +900. In that case the CPUs have idle time, so the run queue is I/O wait or an +lxcfs-computed load average, not CPU contention. PSI always wins when present. + +**A null `cpu_psi_pm` usually means an old drone binary.** Drones built before +CPU-pressure admission send no pressure at all, so `/status` shows +`cpu_psi_pm: null` and the drone is never blocked. Before trusting admission, +check that every drone runs a current binary (see +[Updating a drone](#updating-a-drone)). Inside an LXC, `/proc/loadavg` is +lxcfs's estimate and over-reports badly, while the container's +`/sys/fs/cgroup/cpu.pressure` is kernel-computed. An LXC drone with a current +binary therefore admits on real pressure. + +`/status` shows `admission: {state, cpu_psi_pm, cpu_load_pm}` per drone +(`state` is `open`, `pr_blocked`, `light_only` or `closed`; `cpu_load_pm` is the +load estimate, used only when `cpu_psi_pm` is null). Placement diagnostics name +the rejection `pr_pressure`, `light_only_pressure` or `admission_closed`. +Metrics: `borg_drone_cpu_psi_pm{drone}`, `borg_drone_cpu_load_pm{drone}` and +`borg_drone_admission{drone,state}`. The build environment contract's +`resources` check is separate and still works. ### Standby runners and deploy fail-open diff --git a/src/allocator.rs b/src/allocator.rs index 964bf7e..5900680 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -85,12 +85,34 @@ pub struct DroneState { /// or stale. Only read when [`AllocatorInput::admission`] is set. #[serde(default)] pub cpu_psi_pm: Option, + /// CPU pressure estimated from the load average, for a drone without PSI + /// (see [`load_pressure_pm`]). Admission uses it only when `cpu_psi_pm` + /// is `None`. + #[serde(default)] + pub cpu_load_pm: Option, +} + +/// CPU pressure (permille) estimated from the 1-minute load average per CPU +/// (`load_per_cpu_pm`, 1000 = one runnable task per CPU) for hosts without +/// PSI. With `L` runnable tasks on `N` CPUs, `(L - N) / L` of them wait, which +/// approximates PSI `cpu some`. When the root cgroup shows the CPUs under 90% +/// busy (`cpu_busy_pm`), the run queue is not CPU contention (I/O wait, or an +/// lxcfs load average) and the estimate is 0. +pub fn load_pressure_pm(load_per_cpu_pm: Option, cpu_busy_pm: Option) -> Option { + let load = u64::from(load_per_cpu_pm?); + if load <= 1000 || cpu_busy_pm.is_some_and(|b| b < 900) { + return Some(0); + } + u16::try_from((load - 1000) * 1000 / load).ok() } /// CPU-pressure thresholds (permille) for admission; see [`AdmissionState`]. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct AdmissionThresholds { pub psi_soft_pm: u16, + /// From this pressure only light main and deploy grants; at or above + /// `psi_hard_pm` it has no effect. + pub psi_light_only_pm: u16, pub psi_hard_pm: u16, } @@ -102,44 +124,66 @@ pub enum AdmissionState { Open, /// At or above soft: pull-request grants are shed, everything else placed. PrBlocked, + /// At or above `psi_light_only_pm`: pull-request grants are shed, and main + /// work only of the light class; deploy work is still placed. + LightOnly, /// At or above hard: no new grant of any scope. Closed, } impl AdmissionState { - pub const ALL: [AdmissionState; 3] = [Self::Open, Self::PrBlocked, Self::Closed]; + pub const ALL: [AdmissionState; 4] = + [Self::Open, Self::PrBlocked, Self::LightOnly, Self::Closed]; pub fn as_str(self) -> &'static str { match self { Self::Open => "open", Self::PrBlocked => "pr_blocked", + Self::LightOnly => "light_only", Self::Closed => "closed", } } - pub fn of(psi: Option, t: Option) -> Self { - match (psi, t) { + /// The state for a drone's pressure: its PSI reading, or the load + /// estimate when it has no PSI. + pub fn of(pressure: Option, t: Option) -> Self { + match (pressure, t) { (Some(p), Some(t)) if p >= t.psi_hard_pm => Self::Closed, + (Some(p), Some(t)) if p >= t.psi_light_only_pm => Self::LightOnly, (Some(p), Some(t)) if p >= t.psi_soft_pm => Self::PrBlocked, _ => Self::Open, } } - fn admits(self, scope: Scope) -> bool { + pub fn admits(self, class: Class, scope: Scope) -> bool { match self { Self::Open => true, Self::PrBlocked => scope != Scope::Pr, + Self::LightOnly => { + scope == Scope::Deploy || (scope == Scope::Main && class == Class::Light) + } Self::Closed => false, } } } +impl DroneState { + /// The pressure admission uses: PSI when measured, else the load estimate. + pub fn admission_pressure_pm(&self) -> Option { + self.cpu_psi_pm.or(self.cpu_load_pm) + } +} + /// One drone's admission state this pass. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct DroneAdmission { pub drone_id: String, pub state: AdmissionState, pub cpu_psi_pm: Option, + /// The load-average estimate; drives `state` only when `cpu_psi_pm` is + /// `None`. + #[serde(default)] + pub cpu_load_pm: Option, } impl DroneState { @@ -159,6 +203,7 @@ impl DroneState { profiles: Vec::new(), blocked_classes: Vec::new(), cpu_psi_pm: None, + cpu_load_pm: None, trust: Trust::Shared, } } @@ -299,6 +344,8 @@ pub enum Rejection { NotAccepting, AdmissionClosed, PrPressure, + /// CPU pressure at or above `psi_light_only_pm`: no main medium or heavy. + LightOnlyPressure, Trust, Environment, Docker, @@ -511,7 +558,12 @@ impl Free<'_> { } match self.admission { AdmissionState::Closed => reasons.push(Rejection::AdmissionClosed), - AdmissionState::PrBlocked if scope == Scope::Pr => reasons.push(Rejection::PrPressure), + AdmissionState::PrBlocked | AdmissionState::LightOnly if scope == Scope::Pr => { + reasons.push(Rejection::PrPressure); + } + AdmissionState::LightOnly if !self.admission.admits(class, scope) => { + reasons.push(Rejection::LightOnlyPressure); + } _ => {} } if !self.drone.trust.admits(scope) { @@ -732,7 +784,7 @@ fn allocate_inner( vcpu: drone.vcpu, leases_left: drone.max_leases, in_flight: 0, - admission: AdmissionState::of(drone.cpu_psi_pm, input.admission), + admission: AdmissionState::of(drone.admission_pressure_pm(), input.admission), granted: BTreeMap::new(), }); } @@ -865,6 +917,7 @@ fn allocate_inner( drone_id: (*id).to_owned(), state: f.admission, cpu_psi_pm: f.drone.cpu_psi_pm, + cpu_load_pm: f.drone.cpu_load_pm, }) .collect(), promote, @@ -905,7 +958,7 @@ fn allocate_inner( !reserved.contains_key(**id) && f.drone.trust.admits(scope) && f.could_ever_fit(spec) - && f.admission.admits(scope) + && f.admission.admits(class, scope) && !f.drone.blocked_classes.contains(&class) && f.offer( class, @@ -988,6 +1041,7 @@ mod tests { profiles: vec![], blocked_classes: vec![], cpu_psi_pm: None, + cpu_load_pm: None, trust: Trust::Shared, } } @@ -1070,6 +1124,94 @@ mod tests { } } + #[test] + fn load_estimate_approximates_waiting_share_and_ignores_idle_cpus() { + assert_eq!(load_pressure_pm(None, Some(1000)), None); + assert_eq!(load_pressure_pm(Some(800), None), Some(0)); + assert_eq!(load_pressure_pm(Some(1000), Some(1000)), Some(0)); + // Load 93 on 16 CPUs, CPUs saturated: ~83% of runnable tasks wait. + assert_eq!(load_pressure_pm(Some(5838), Some(990)), Some(828)); + assert_eq!(load_pressure_pm(Some(2000), None), Some(500)); + // Same run queue with idle CPU time is not CPU contention. + assert_eq!(load_pressure_pm(Some(5838), Some(550)), Some(0)); + } + + #[test] + fn light_only_sheds_pr_and_main_medium_heavy_but_keeps_light_main_and_deploy() { + let t = AdmissionThresholds { + psi_soft_pm: 300, + psi_light_only_pm: 500, + psi_hard_pm: 600, + }; + let of = |p| AdmissionState::of(Some(p), Some(t)); + assert_eq!(of(299), AdmissionState::Open); + assert_eq!(of(300), AdmissionState::PrBlocked); + assert_eq!(of(540), AdmissionState::LightOnly); + assert_eq!(of(600), AdmissionState::Closed); + let s = AdmissionState::LightOnly; + assert!(s.admits(Class::Light, Scope::Main)); + assert!(s.admits(Class::Heavy, Scope::Deploy)); + assert!(!s.admits(Class::Medium, Scope::Main)); + assert!(!s.admits(Class::Heavy, Scope::Main)); + assert!(!s.admits(Class::Light, Scope::Pr)); + // Unset light-only tier (== hard): PR-blocked straight to closed. + let old = AdmissionThresholds { + psi_light_only_pm: 600, + ..t + }; + assert_eq!( + AdmissionState::of(Some(540), Some(old)), + AdmissionState::PrBlocked + ); + + // A loaded drone without PSI: its load estimate drives admission. + let loaded = DroneState { + cpu_load_pm: Some(828), + ..host("loaded") + }; + let mut i = input( + &[ + (Class::Heavy, Scope::Main, 1), + (Class::Medium, Scope::Main, 1), + (Class::Light, Scope::Main, 1), + ], + vec![loaded.clone()], + ); + i.admission = Some(t); + let a = allocate(&i); + assert_eq!(a.admission[0].state, AdmissionState::Closed); + assert!(a.grants.is_empty()); + + let busy = DroneState { + cpu_psi_pm: Some(540), + ..loaded + }; + i.drones = vec![busy]; + let a = allocate(&i); + assert_eq!( + a.admission[0].state, + AdmissionState::LightOnly, + "PSI wins over the load estimate" + ); + assert_eq!(count(&a, Class::Light, Scope::Main), 1); + assert_eq!(count(&a, Class::Medium, Scope::Main), 0); + assert_eq!(count(&a, Class::Heavy, Scope::Main), 0); + let rejected: Vec<_> = a + .decisions + .iter() + .filter(|d| d.class == Class::Heavy) + .flat_map(|d| { + d.candidates + .iter() + .flat_map(|c| c.rejection_reasons.clone()) + }) + .collect(); + assert!( + rejected.contains(&Rejection::LightOnlyPressure), + "{rejected:?}" + ); + } + #[test] fn free_room_and_slots_subtract_what_the_drone_already_holds() { let classes = Classes::default(); @@ -1610,6 +1752,7 @@ mod tests { profiles: vec![], blocked_classes: vec![], cpu_psi_pm: None, + cpu_load_pm: None, trust: Trust::Shared, }, DroneState { @@ -1623,6 +1766,7 @@ mod tests { profiles: vec![], blocked_classes: vec![], cpu_psi_pm: None, + cpu_load_pm: None, trust: Trust::Shared, }, ]; diff --git a/src/api.rs b/src/api.rs index 35c170d..701d667 100644 --- a/src/api.rs +++ b/src/api.rs @@ -1,7 +1,8 @@ //! HTTP surface of the queen. //! -//! * `GET /status` and `GET /metrics` are read-only and unauthenticated; bind -//! them to a private interface. +//! * `GET /status`, `GET /metrics`, `GET /report`, `GET /jobs/{job_id}` and +//! `GET /runs/{run_id}` are read-only and unauthenticated; bind them to a +//! private interface. //! * `POST /v1/heartbeat` and `POST /v1/leases/{id}/result` are the drone API, //! authenticated with `Authorization: Bearer `. //! @@ -29,6 +30,7 @@ use crate::enroll::hash_token; use crate::envcontract::{CheckState, EnvFacts, Need, Verdict, blocked_classes, evaluate}; use crate::health::{DroneHealth, HealthState}; use crate::homeostasis::KeyReport; +use crate::hostdiag::{JobHostView, LeaseHostReport}; use crate::probe::Capacity; use crate::queen::RolloutView; use crate::queen::{Queen, QueenError, QueueEntry}; @@ -97,6 +99,7 @@ fn sanitize_facts(f: &EnvFacts, now: i64) -> EnvFacts { }; let mut out = f.clone(); out.observed_at = out.observed_at.min(now); + out.cpu_busy_pm = out.cpu_busy_pm.map(|b| b.min(1000)); for fact in out.capabilities.values_mut() { fact.observed_at = fact.observed_at.min(now); } @@ -272,6 +275,10 @@ pub struct ResultRequest { /// The drone stopped the job because the queen preempted the lease. #[serde(default)] pub preempted: bool, + /// Host load when the runner started and when it exited; see + /// [`crate::hostdiag`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub host: Option, } #[derive(Debug)] @@ -327,6 +334,8 @@ pub fn router(queen: Arc) -> Router { .route("/metrics", get(metrics)) .route("/report", get(report)) .route("/explain", get(explain)) + .route("/jobs/{job_id}", get(job_host)) + .route("/runs/{run_id}", get(run_hosts)) .route("/v1/enroll", post(tailnet_enroll)) .route("/v1/heartbeat", post(heartbeat)) .route("/v1/leases/{id}/result", post(lease_result)) @@ -627,6 +636,15 @@ async fn lease_result( .finish_lease(&id, &drone, &result, now) .await? { + if let Some(lease) = queen.store().lease(&id).await? + && let Err(e) = queen + .store() + .record_lease_host(&lease, &result, req.host.as_ref(), now) + .await + { + // Diagnostics never fail a result. + tracing::warn!("lease {id}: host diagnostics not stored: {e}"); + } if let Some(t) = req.telemetry { match t.sanitized() { Some(t) => { @@ -791,6 +809,9 @@ pub struct CapacityStatus { pub struct DroneAdmissionView { pub state: AdmissionState, pub cpu_psi_pm: Option, + /// Pressure estimated from the load average; `state` uses it only when + /// `cpu_psi_pm` is null. + pub cpu_load_pm: Option, } /// One drone's build environment contract. @@ -1107,6 +1128,7 @@ pub async fn build_status(queen: &Queen, now: i64) -> Result .map(|a| DroneAdmissionView { state: a.state, cpu_psi_pm: a.cpu_psi_pm, + cpu_load_pm: a.cpu_load_pm, }), contract: cfg.env_contract.enabled.then(|| { let verdict = |class| { @@ -1311,6 +1333,54 @@ async fn report( ))) } +/// The CPU pressure (percent) above which a job's host counts as slow: the +/// admission soft threshold when admission is on. +fn slow_cpu_psi_pct(queen: &Queen) -> f64 { + queen + .config() + .admission + .as_ref() + .map_or(crate::hostdiag::SLOW_CPU_PSI_PCT, |a| { + f64::from(a.psi_soft_pm) / 10.0 + }) +} + +/// `GET /jobs/{job_id}`: the drone that ran GitHub job `job_id` and its host +/// load at start and end. 404 when the queen never saw that job on a runner it +/// minted, or the row aged out. Unauthenticated like `/status`. +async fn job_host( + State(queen): State>, + Path(job_id): Path, +) -> Result, ApiError> { + queen + .store() + .job_host(job_id, slow_cpu_psi_pct(&queen)) + .await? + .map(Json) + .ok_or(ApiError::NotFound) +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RunHosts { + pub run_id: u64, + pub jobs: Vec, +} + +/// `GET /runs/{run_id}`: every job of a workflow run that ran on a borg +/// runner, oldest lease first. An unknown run has an empty `jobs` list. +async fn run_hosts( + State(queen): State>, + Path(run_id): Path, +) -> Result, ApiError> { + Ok(Json(RunHosts { + run_id, + jobs: queen + .store() + .run_hosts(run_id, slow_cpu_psi_pct(&queen)) + .await?, + })) +} + async fn status(State(queen): State>) -> Result, ApiError> { Ok(Json( build_status(&queen, OffsetDateTime::now_utc().unix_timestamp()).await?, @@ -1568,9 +1638,24 @@ pub fn render_metrics(status: &Status) -> String { }) .collect(), ); + gauge( + "borg_drone_cpu_load_pm", + "CPU pressure (permille) estimated from the drone's load average; admission uses it when PSI is absent.", + status + .drones + .iter() + .filter_map(|d| { + let load = d.admission.as_ref()?.cpu_load_pm?; + Some(( + format!("drone=\"{}\"", label_value(&d.name)), + f64::from(load), + )) + }) + .collect(), + ); gauge( "borg_drone_admission", - "1 for the drone's CPU-pressure admission state (open, pr_blocked, closed).", + "1 for the drone's CPU-pressure admission state (open, pr_blocked, light_only, closed).", status .drones .iter() diff --git a/src/config.rs b/src/config.rs index 5033ab2..dae08b5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -202,6 +202,12 @@ pub struct AdmissionSettings { /// From this `cpu some avg10` (permille) a drone takes no new /// pull-request grants. pub psi_soft_pm: u16, + /// From this pressure a drone takes, besides deploy work, only light main + /// grants: no medium or heavy main work. Optional; between `psi_soft_pm` + /// and `psi_hard_pm`. Unset, the drone goes from pull-request blocked + /// straight to closed. + #[serde(default)] + pub psi_light_only_pm: Option, /// From this pressure a drone takes no new grants at all. pub psi_hard_pm: u16, /// A drone's reported pressure older than this counts as unknown. @@ -212,6 +218,7 @@ impl AdmissionSettings { pub fn thresholds(&self) -> AdmissionThresholds { AdmissionThresholds { psi_soft_pm: self.psi_soft_pm, + psi_light_only_pm: self.psi_light_only_pm.unwrap_or(self.psi_hard_pm), psi_hard_pm: self.psi_hard_pm, } } @@ -220,6 +227,13 @@ impl AdmissionSettings { if self.psi_soft_pm == 0 || self.psi_soft_pm > self.psi_hard_pm || self.psi_hard_pm > 1000 { return Err("admission needs 0 < psi_soft_pm <= psi_hard_pm <= 1000".to_owned()); } + if let Some(l) = self.psi_light_only_pm + && !(self.psi_soft_pm..=self.psi_hard_pm).contains(&l) + { + return Err( + "admission needs psi_soft_pm <= psi_light_only_pm <= psi_hard_pm".to_owned(), + ); + } if self.facts_max_age_secs <= 0 { return Err("admission.facts_max_age_secs must be positive".to_owned()); } @@ -665,6 +679,12 @@ mod tests { let ok = with(r#"{"psi_soft_pm":300,"psi_hard_pm":600,"facts_max_age_secs":180}"#).unwrap(); let t = ok.admission.unwrap().thresholds(); assert_eq!((t.psi_soft_pm, t.psi_hard_pm), (300, 600)); + assert_eq!(t.psi_light_only_pm, 600, "unset means no light-only tier"); + let t = with(r#"{"psi_soft_pm":300,"psi_light_only_pm":500,"psi_hard_pm":600,"facts_max_age_secs":180}"#) + .unwrap() + .admission + .unwrap(); + assert_eq!(t.thresholds().psi_light_only_pm, 500); for bad in [ r#"{"psi_soft_pm":700,"psi_hard_pm":600,"facts_max_age_secs":180}"#, r#"{"psi_soft_pm":0,"psi_hard_pm":600,"facts_max_age_secs":180}"#, @@ -672,6 +692,8 @@ mod tests { r#"{"psi_soft_pm":300,"psi_hard_pm":600,"facts_max_age_secs":0}"#, r#"{"psi_soft_pm":300,"psi_hard_pm":600,"facts_max_age_secs":180,"nope":1}"#, r#"{"psi_soft_pm":300,"psi_hard_pm":600}"#, + r#"{"psi_soft_pm":300,"psi_light_only_pm":200,"psi_hard_pm":600,"facts_max_age_secs":180}"#, + r#"{"psi_soft_pm":300,"psi_light_only_pm":700,"psi_hard_pm":600,"facts_max_age_secs":180}"#, ] { assert!(with(bad).is_err(), "{bad}"); } diff --git a/src/drone/mod.rs b/src/drone/mod.rs index 897f581..20f4c1a 100644 --- a/src/drone/mod.rs +++ b/src/drone/mod.rs @@ -96,6 +96,8 @@ pub struct Drone { /// Last Queen floor, retained while the Queen is unreachable. disk_floor: AtomicU64, capability_cache: capabilities::CapabilityCache, + /// Busy-cores rate of the host, fed by heartbeats and lease samples. + cpu_meter: Arc, } impl Drone { @@ -137,6 +139,7 @@ impl Drone { tools_cache: probe::ToolsCache::default(), disk_floor: AtomicU64::new(crate::envcontract::DEFAULT_MIN_DISK_FREE_BYTES), capability_cache: capabilities::CapabilityCache::default(), + cpu_meter: Arc::default(), client: Arc::new(client), host, exec, @@ -218,14 +221,26 @@ impl Drone { status: 0, })?; let host = Arc::clone(&self.host); + let meter = Arc::clone(&self.cpu_meter); let now = time::OffsetDateTime::now_utc().unix_timestamp(); - let mut env_facts = - tokio::task::spawn_blocking(move || probe::env_facts(host.as_ref(), now)) - .await - .map_err(|e| ClientError::Status { - path: format!("probe: {e}"), - status: 0, - })?; + let mut env_facts = tokio::task::spawn_blocking(move || { + let now_ms = + i64::try_from(time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) + .unwrap_or(i64::MAX); + let busy = crate::hostdiag::observe_cpu(host.as_ref(), &meter, now_ms); + let mut facts = probe::env_facts(host.as_ref(), now); + facts.cpu_busy_pm = busy.zip(host.available_parallelism()).map(|(b, cpus)| { + (b / f64::from(cpus.max(1)) * 1000.0) + .round() + .clamp(0.0, 1000.0) as u16 + }); + facts + }) + .await + .map_err(|e| ClientError::Status { + path: format!("probe: {e}"), + status: 0, + })?; env_facts.cache_dir_free_bytes = cache_free; let accepting = accepting && !crate::envcontract::disk_admission( @@ -386,6 +401,8 @@ impl Drone { }; let ctx = LeaseContext { cfg: Arc::clone(&self.cfg), + host: Arc::clone(&self.host), + cpu_meter: Arc::clone(&self.cpu_meter), client: Arc::clone(&self.client), exec: Arc::clone(&self.exec), job_exec: Arc::clone(&self.job_exec), @@ -395,7 +412,7 @@ impl Drone { proc_root: Arc::clone(&self.proc_root), caches: Arc::clone(&self.caches), slot, - _held: held, + held, }; tasks.spawn(handle_lease(ctx, grant)); true @@ -557,6 +574,8 @@ impl Drone { struct LeaseContext { cfg: Arc, + host: Arc, + cpu_meter: Arc, client: Arc, exec: Arc, job_exec: Arc, @@ -565,9 +584,39 @@ struct LeaseContext { preempt: tokio::sync::watch::Receiver, proc_root: Arc, caches: Arc, - /// Dropped before `_held`, so the slot is free whenever the id is. + /// Dropped before `held`, so the slot is free whenever the id is. slot: SlotGuard, - _held: HeldGuard, + held: HeldGuard, +} + +/// Sample the host for this lease's diagnostics. Never fails the lease. +async fn sample_host(ctx: &LeaseContext) -> Option { + let (host, cfg, meter) = ( + Arc::clone(&ctx.host), + Arc::clone(&ctx.cfg), + Arc::clone(&ctx.cpu_meter), + ); + let active = ctx + .held + .held + .lock() + .map(|h| u32::try_from(h.len()).unwrap_or(u32::MAX)) + .unwrap_or(0); + tokio::task::spawn_blocking(move || { + let now_ms = + i64::try_from(time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) + .unwrap_or(i64::MAX); + crate::hostdiag::sample( + host.as_ref(), + &cfg.work_root, + &cfg.cache_root, + active, + &meter, + now_ms, + ) + }) + .await + .ok() } /// Run one lease, clean up unconditionally, report the result. @@ -593,6 +642,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { job_user: ctx.cfg.job_user.as_ref(), proc_root: &ctx.proc_root, }; + let host_start = sample_host(&ctx).await; let outcome = run_lease( &spec, ctx.job_exec.as_ref(), @@ -601,6 +651,11 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { ctx.preempt.clone(), ) .await; + let host_end = sample_host(&ctx).await; + let host = Some(crate::hostdiag::LeaseHostReport { + start: host_start, + end: host_end, + }); // Runs whether the job passed, failed, timed out, was aborted or never started. let (work_root, id, slot_name, docker, exec, fs_ops) = ( @@ -645,6 +700,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { peak_rss_mib: o.peak_rss_mib, telemetry: None, preempted: o.preempted, + host, }, Err(e) => { tracing::warn!("lease {} could not run: {e}", grant.lease_id); @@ -654,6 +710,7 @@ async fn handle_lease(ctx: LeaseContext, grant: LeaseGrant) { peak_rss_mib: None, telemetry: None, preempted: false, + host, } } }; diff --git a/src/envcontract.rs b/src/envcontract.rs index a258c5c..06442f5 100644 --- a/src/envcontract.rs +++ b/src/envcontract.rs @@ -200,6 +200,14 @@ pub struct EnvFacts { /// where PSI does not exist. #[serde(default)] pub cpu_psi_some_avg10_pm: Option, + /// `/proc/loadavg` 1-minute average per CPU, in permille (1000 is one + /// runnable task per CPU). Admission falls back to it without PSI. + #[serde(default)] + pub load_per_cpu_pm: Option, + /// Share of the host's CPUs busy since the previous heartbeat, from the + /// root cgroup's `cpu.stat`, in permille. + #[serde(default)] + pub cpu_busy_pm: Option, /// For each tool in [`PROBE_TOOLS`], the first line of its `--version` /// output, or `"missing"`; `None` when the drone did not measure it. #[serde(default)] @@ -958,6 +966,8 @@ mod tests { }]), mem_available_mib: Some(20_000), cpu_psi_some_avg10_pm: Some(100), + load_per_cpu_pm: None, + cpu_busy_pm: None, tools: Some(tool_map(&[ ("git", "git version 2.45.0"), ("gh", "gh version 2.50.0"), diff --git a/src/hostdiag.rs b/src/hostdiag.rs new file mode 100644 index 0000000..56b9a64 --- /dev/null +++ b/src/hostdiag.rs @@ -0,0 +1,978 @@ +//! Per-lease host diagnostics: how loaded the drone was when a job started +//! and when it ended. +//! +//! The drone samples its host just before it starts the runner and again when +//! the runner exits, and sends both samples with the lease result +//! (`POST /v1/leases/{id}/result`, field `host`). The queen stores one row per +//! lease, adds the GitHub run id and job id when it sees the job on a runner it +//! minted, and serves the rows on `GET /jobs/{job_id}` and `GET /runs/{run_id}`. +//! +//! A slow job on an overloaded drone looks like a code failure unless someone +//! compares hosts by hand. These rows let a triage tool show which drone ran a +//! job and how loaded that drone was. + +use std::path::Path; +use std::sync::Mutex; + +use serde::{Deserialize, Serialize}; +use sqlx::Row; + +use crate::classes::{Class, Scope}; +use crate::probe::{self, Host}; +use crate::store::{LeaseResult, LeaseRow, Store, StoreError}; + +/// Rows are kept this long after the lease was granted. +pub const KEEP_SECS: i64 = 7 * 86_400; +/// At most this many rows are kept; the oldest go first. +pub const MAX_ROWS: i64 = 50_000; +/// CPU pressure (`some avg10`, percent) at job start that marks a slow host +/// when the queen has no `admission` block. +pub const SLOW_CPU_PSI_PCT: f64 = 30.0; +/// Memory pressure (`some avg10`, percent) at job start that marks a slow host. +pub const SLOW_MEM_PSI_PCT: f64 = 10.0; +/// Load average per CPU that marks a slow host. Only used when the host has no +/// CPU pressure reading, because an LXC's load average is computed by lxcfs and +/// over-reports. +pub const SLOW_LOAD_PER_CPU: f64 = 1.5; +/// Share of the host's CPUs busy that marks a slow host. +pub const SLOW_BUSY_SHARE: f64 = 0.9; +/// Free disk on the work or cache volume that marks a slow host. +pub const SLOW_FREE_DISK_MIB: u64 = 10 * 1024; + +const MAX_NAME_CHARS: usize = 120; + +/// One Linux PSI file (`/proc/pressure/*` or a cgroup `*.pressure`), in percent. +#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)] +pub struct Psi { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub some_avg10: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub some_avg60: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub full_avg10: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub full_avg60: Option, +} + +impl Psi { + /// Parse `some avg10=1.00 avg60=...` / `full ...` lines. `None` when no + /// field could be read. + pub fn parse(text: &str) -> Option { + let field = |kind: &str, key: &str| -> Option { + text.lines() + .find(|l| l.starts_with(kind))? + .split_whitespace() + .find_map(|f| f.strip_prefix(key))? + .parse::() + .ok() + .filter(|v| v.is_finite()) + .map(|v| v.clamp(0.0, 100.0)) + }; + let psi = Self { + some_avg10: field("some ", "avg10="), + some_avg60: field("some ", "avg60="), + full_avg10: field("full ", "avg10="), + full_avg60: field("full ", "avg60="), + }; + (psi != Self::default()).then_some(psi) + } + + fn sanitized(self) -> Self { + let pct = |v: Option| v.filter(|v| v.is_finite()).map(|v| v.clamp(0.0, 100.0)); + Self { + some_avg10: pct(self.some_avg10), + some_avg60: pct(self.some_avg60), + full_avg10: pct(self.full_avg10), + full_avg60: pct(self.full_avg60), + } + } +} + +/// The drone's host at one moment. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct HostSample { + /// Unix seconds. + pub at: i64, + /// CPUs the drone process may use. + #[serde(default)] + pub cpus: Option, + /// `/proc/loadavg` 1, 5 and 15 minute averages. Inside an LXC this is + /// lxcfs's estimate and over-reports; prefer the pressure and busy fields. + #[serde(default)] + pub load: Option<[f64; 3]>, + /// `/proc/pressure/cpu`. + #[serde(default)] + pub cpu_psi: Option, + /// `/proc/pressure/memory`. + #[serde(default)] + pub memory_psi: Option, + /// `/proc/pressure/io`. + #[serde(default)] + pub io_psi: Option, + /// `/sys/fs/cgroup/cpu.pressure`: the cgroup-namespace root, which is the + /// whole container inside an LXC. + #[serde(default)] + pub cgroup_cpu_psi: Option, + /// `/sys/fs/cgroup/memory.pressure`. + #[serde(default)] + pub cgroup_memory_psi: Option, + /// `usage_usec` of `/sys/fs/cgroup/cpu.stat` (whole host or container). + #[serde(default)] + pub host_cpu_usage_usec: Option, + /// Cores busy in that cgroup, averaged since the drone's previous sample + /// (at most one heartbeat ago). + #[serde(default)] + pub host_cpu_busy_cores: Option, + /// `usage_usec` of the drone's own cgroup; jobs run inside it. + #[serde(default)] + pub drone_cpu_usage_usec: Option, + /// `throttled_usec` of the drone's own cgroup. + #[serde(default)] + pub drone_cpu_throttled_usec: Option, + #[serde(default)] + pub mem_available_mib: Option, + /// Free space on the filesystem holding the drone's `work_root`. + #[serde(default)] + pub work_free_mib: Option, + /// Free space on the filesystem holding the drone's `cache_root`. + #[serde(default)] + pub cache_free_mib: Option, + /// Leases the drone held at this moment, this one included. + #[serde(default)] + pub active_leases: u32, +} + +impl HostSample { + /// Clamp what a drone sent: finite non-negative numbers, no future time. + pub fn sanitized(mut self, now: i64) -> Self { + let num = |v: f64| v.is_finite().then(|| v.max(0.0)); + self.at = self.at.min(now); + self.load = self + .load + .and_then(|l| Some([num(l[0])?, num(l[1])?, num(l[2])?])); + for psi in [ + &mut self.cpu_psi, + &mut self.memory_psi, + &mut self.io_psi, + &mut self.cgroup_cpu_psi, + &mut self.cgroup_memory_psi, + ] { + *psi = psi.map(Psi::sanitized); + } + self.host_cpu_busy_cores = self.host_cpu_busy_cores.and_then(num); + self + } + + /// The higher CPU `some avg10` of the host file and the cgroup file. + pub fn cpu_some_avg10(&self) -> Option { + [self.cpu_psi, self.cgroup_cpu_psi] + .into_iter() + .flatten() + .filter_map(|p| p.some_avg10) + .reduce(f64::max) + } + + fn mem_some_avg10(&self) -> Option { + [self.memory_psi, self.cgroup_memory_psi] + .into_iter() + .flatten() + .filter_map(|p| p.some_avg10) + .reduce(f64::max) + } +} + +/// What a drone sends in its lease result. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct LeaseHostReport { + #[serde(default)] + pub start: Option, + #[serde(default)] + pub end: Option, +} + +/// Turns successive `cpu.stat` readings into a busy-cores rate. +#[derive(Debug, Default)] +pub struct CpuMeter(Mutex>); + +#[derive(Debug, Clone, Copy)] +struct MeterState { + at_ms: i64, + usage_usec: u64, + busy_cores: Option, +} + +/// Readings closer together than this keep the previous rate. +const METER_MIN_GAP_MS: i64 = 2_000; + +impl CpuMeter { + /// Record a reading and return the latest rate. A reading that arrives + /// within two seconds of the stored one returns the stored rate, so a lease + /// that starts right after a heartbeat still gets a measured value. + pub fn observe(&self, at_ms: i64, usage_usec: u64) -> Option { + let mut slot = self.0.lock().ok()?; + let next = match *slot { + Some(prev) if at_ms - prev.at_ms < METER_MIN_GAP_MS && at_ms >= prev.at_ms => { + return prev.busy_cores; + } + Some(prev) if at_ms > prev.at_ms && usage_usec >= prev.usage_usec => MeterState { + at_ms, + usage_usec, + busy_cores: Some( + (usage_usec - prev.usage_usec) as f64 / ((at_ms - prev.at_ms) as f64 * 1000.0), + ), + }, + _ => MeterState { + at_ms, + usage_usec, + busy_cores: None, + }, + }; + *slot = Some(next); + next.busy_cores + } +} + +/// `usage_usec` and `throttled_usec` from a cgroup v2 `cpu.stat`. +fn cpu_stat(host: &dyn Host, dir: &Path) -> (Option, Option) { + let Some(text) = host.read_to_string(&dir.join("cpu.stat")) else { + return (None, None); + }; + let field = |key: &str| { + text.lines() + .find_map(|l| l.strip_prefix(key)?.trim().parse::().ok()) + }; + (field("usage_usec "), field("throttled_usec ")) +} + +/// The drone's own cgroup directory under `/sys/fs/cgroup`, from the cgroup v2 +/// line of `/proc/self/cgroup`. +fn own_cgroup(host: &dyn Host) -> Option { + let text = host.read_to_string(Path::new("/proc/self/cgroup"))?; + let rel = text.lines().find_map(|l| l.strip_prefix("0::"))?.trim(); + let rel = rel.trim_start_matches('/'); + if rel.contains("..") { + return None; + } + Some(Path::new("/sys/fs/cgroup").join(rel)) +} + +fn loadavg(host: &dyn Host) -> Option<[f64; 3]> { + let text = host.read_to_string(Path::new("/proc/loadavg"))?; + let mut it = text + .split_whitespace() + .map(|f| f.parse::().ok().filter(|v| v.is_finite())); + Some([it.next()??, it.next()??, it.next()??]) +} + +/// Feed the meter from the root cgroup and return the busy-cores rate. Called +/// on every heartbeat so a lease start has a recent rate. +pub fn observe_cpu(host: &dyn Host, meter: &CpuMeter, at_ms: i64) -> Option { + let (usage, _) = cpu_stat(host, Path::new("/sys/fs/cgroup")); + meter.observe(at_ms, usage?) +} + +/// Sample the host for a lease start or end. Every field degrades to `None` +/// on its own; this never fails. +pub fn sample( + host: &dyn Host, + work_root: &Path, + cache_root: &Path, + active_leases: u32, + meter: &CpuMeter, + at_ms: i64, +) -> HostSample { + let psi = |path: &str| Psi::parse(&host.read_to_string(Path::new(path))?); + let (host_usage, _) = cpu_stat(host, Path::new("/sys/fs/cgroup")); + let (drone_usage, drone_throttled) = own_cgroup(host) + .map(|dir| cpu_stat(host, &dir)) + .unwrap_or_default(); + HostSample { + at: at_ms.div_euclid(1000), + cpus: host.available_parallelism(), + load: loadavg(host), + cpu_psi: psi("/proc/pressure/cpu"), + memory_psi: psi("/proc/pressure/memory"), + io_psi: psi("/proc/pressure/io"), + cgroup_cpu_psi: psi("/sys/fs/cgroup/cpu.pressure"), + cgroup_memory_psi: psi("/sys/fs/cgroup/memory.pressure"), + host_cpu_usage_usec: host_usage, + host_cpu_busy_cores: host_usage.and_then(|u| meter.observe(at_ms, u)), + drone_cpu_usage_usec: drone_usage, + drone_cpu_throttled_usec: drone_throttled, + mem_available_mib: probe::mem_available_mib(host), + work_free_mib: probe::free_disk_mib_at(host, work_root), + cache_free_mib: probe::free_disk_mib_at(host, cache_root), + active_leases, + } +} + +// ---------------------------------------------------------------- queen side + +/// Averages over the job, from the start and end samples. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct During { + pub secs: i64, + /// Cores busy on the whole host (or container). + #[serde(skip_serializing_if = "Option::is_none")] + pub host_cpu_busy_cores: Option, + /// Cores used by the drone's cgroup: this job and any other lease. + #[serde(skip_serializing_if = "Option::is_none")] + pub drone_cpu_busy_cores: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub drone_cpu_throttled_secs: Option, +} + +fn during(start: &HostSample, end: &HostSample) -> Option { + let secs = end.at - start.at; + if secs <= 0 { + return None; + } + let rate = + |a: Option, b: Option| Some(b?.checked_sub(a?)? as f64 / 1e6 / secs as f64); + Some(During { + secs, + host_cpu_busy_cores: rate(start.host_cpu_usage_usec, end.host_cpu_usage_usec), + drone_cpu_busy_cores: rate(start.drone_cpu_usage_usec, end.drone_cpu_usage_usec), + drone_cpu_throttled_secs: end + .drone_cpu_throttled_usec + .zip(start.drone_cpu_throttled_usec) + .and_then(|(b, a)| b.checked_sub(a)) + .map(|d| d as f64 / 1e6), + }) +} + +/// Why a job's host was overloaded, or nothing. `cpu_psi_pct` is the CPU +/// pressure threshold (the queen's admission soft threshold when set). +pub fn slow_host_reasons( + start: Option<&HostSample>, + during: Option<&During>, + cpu_psi_pct: f64, +) -> Vec { + let mut out = Vec::new(); + if let Some(s) = start { + let cpus = s.cpus.filter(|&c| c > 0).map(f64::from); + match s.cpu_some_avg10() { + Some(p) if p >= cpu_psi_pct => { + out.push(format!( + "cpu pressure {p:.1}% at start (threshold {cpu_psi_pct:.0}%)" + )); + } + Some(_) => {} + None => { + if let (Some(l), Some(c)) = (s.load, cpus) + && l[0] / c >= SLOW_LOAD_PER_CPU + { + out.push(format!( + "load {:.1} on {c:.0} cpus at start ({:.1} per cpu, threshold {SLOW_LOAD_PER_CPU})", + l[0], + l[0] / c + )); + } + } + } + if let Some(p) = s.mem_some_avg10() + && p >= SLOW_MEM_PSI_PCT + { + out.push(format!( + "memory pressure {p:.1}% at start (threshold {SLOW_MEM_PSI_PCT:.0}%)" + )); + } + if let (Some(b), Some(c)) = (s.host_cpu_busy_cores, cpus) + && b / c >= SLOW_BUSY_SHARE + { + out.push(format!("{b:.1} of {c:.0} cpus busy at start")); + } + for (name, free) in [("work", s.work_free_mib), ("cache", s.cache_free_mib)] { + if let Some(f) = free + && f < SLOW_FREE_DISK_MIB + { + out.push(format!("{name} volume {f} MiB free at start")); + } + } + if let (Some(d), Some(c)) = (during, cpus) + && let Some(b) = d.host_cpu_busy_cores + && b / c >= SLOW_BUSY_SHARE + { + out.push(format!( + "{b:.1} of {c:.0} cpus busy on average during the job" + )); + } + } + out +} + +/// One lease as `GET /jobs/{id}` and `GET /runs/{id}` show it. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct JobHostView { + pub lease_id: String, + pub drone: String, + pub class: Class, + pub scope: Scope, + pub repo: Option, + pub run_id: Option, + pub job_id: Option, + pub job_name: Option, + /// When the queen granted the lease. + pub leased_at: i64, + /// When the drone reported the result. + pub finished_at: Option, + /// Runner wall time the drone reported. + pub duration_secs: Option, + pub exit_code: Option, + pub start: Option, + pub end: Option, + pub during: Option, + /// The drone was overloaded when the job started or while it ran. + pub slow_host: bool, + pub slow_host_reasons: Vec, +} + +fn to_i64(v: u64) -> i64 { + i64::try_from(v).unwrap_or(i64::MAX) +} + +fn json(v: Option<&T>) -> Result, StoreError> { + v.map(|v| serde_json::to_string(v).map_err(|e| StoreError::Corrupt(e.to_string()))) + .transpose() +} + +fn enum_text(v: &T) -> Result { + match serde_json::to_value(v) { + Ok(serde_json::Value::String(s)) => Ok(s), + _ => Err(StoreError::Corrupt("enum is not a string".to_owned())), + } +} + +fn enum_from Deserialize<'de>>(s: &str) -> Result { + serde_json::from_value(serde_json::Value::String(s.to_owned())) + .map_err(|e| StoreError::Corrupt(e.to_string())) +} + +fn clean(s: &str) -> String { + s.chars() + .filter(|c| !c.is_control()) + .take(MAX_NAME_CHARS) + .collect::() + .trim() + .to_owned() +} + +const SELECT: &str = "SELECT lease_id, drone, class, scope, repo, run_id, job_id, job_name, leased_at, finished_at, exit_code, duration_secs, start_sample, end_sample FROM lease_hosts"; + +impl Store { + /// Record a drone's result and host samples for its lease. + pub async fn record_lease_host( + &self, + lease: &LeaseRow, + result: &LeaseResult, + report: Option<&LeaseHostReport>, + now: i64, + ) -> Result<(), StoreError> { + let start = report + .and_then(|r| r.start.clone()) + .map(|s| s.sanitized(now)); + let end = report.and_then(|r| r.end.clone()).map(|s| s.sanitized(now)); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT INTO lease_hosts (lease_id, drone, class, scope, leased_at, finished_at, exit_code, duration_secs, start_sample, end_sample) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(lease_id) DO UPDATE SET + finished_at = excluded.finished_at, exit_code = excluded.exit_code, + duration_secs = excluded.duration_secs, + start_sample = COALESCE(excluded.start_sample, start_sample), + end_sample = COALESCE(excluded.end_sample, end_sample)", + ) + .bind(&lease.id) + .bind(&lease.drone) + .bind(enum_text(&lease.class)?) + .bind(enum_text(&lease.scope)?) + .bind(lease.created_at) + .bind(now) + .bind(i64::from(result.exit_code)) + .bind(to_i64(result.duration_secs)) + .bind(json(start.as_ref())?) + .bind(json(end.as_ref())?) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// Attach the GitHub job a lease's runner picked up. + pub async fn bind_lease_host_job( + &self, + lease: &LeaseRow, + repo: &str, + run_id: u64, + job_id: u64, + job_name: &str, + ) -> Result<(), StoreError> { + let name = clean(job_name); + let run_id = (run_id != 0).then(|| to_i64(run_id)); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + sqlx::query( + "INSERT INTO lease_hosts (lease_id, drone, class, scope, leased_at, repo, run_id, job_id, job_name) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(lease_id) DO UPDATE SET + repo = excluded.repo, run_id = COALESCE(excluded.run_id, run_id), + job_id = excluded.job_id, job_name = COALESCE(excluded.job_name, job_name) + WHERE job_id IS NOT excluded.job_id OR run_id IS NOT COALESCE(excluded.run_id, run_id) + OR job_name IS NOT COALESCE(excluded.job_name, job_name)", + ) + .bind(&lease.id) + .bind(&lease.drone) + .bind(enum_text(&lease.class)?) + .bind(enum_text(&lease.scope)?) + .bind(lease.created_at) + .bind(repo) + .bind(run_id) + .bind(to_i64(job_id)) + .bind((!name.is_empty()).then_some(name)) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// The lease that ran GitHub job `job_id` (the newest, if several did). + pub async fn job_host( + &self, + job_id: u64, + cpu_psi_pct: f64, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query(&format!( + "{SELECT} WHERE job_id = ? ORDER BY leased_at DESC LIMIT 1" + )) + .bind(to_i64(job_id)) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(|r| view(r, cpu_psi_pct)).transpose() + } + + /// One lease's row, by lease id. + pub async fn lease_host( + &self, + lease_id: &str, + cpu_psi_pct: f64, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let row = sqlx::query(&format!("{SELECT} WHERE lease_id = ?")) + .bind(lease_id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(|r| view(r, cpu_psi_pct)).transpose() + } + + /// Every lease that ran a job of workflow run `run_id`, oldest first. + pub async fn run_hosts( + &self, + run_id: u64, + cpu_psi_pct: f64, + ) -> Result, StoreError> { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let rows = sqlx::query(&format!( + "{SELECT} WHERE run_id = ? ORDER BY leased_at, lease_id LIMIT 1000" + )) + .bind(to_i64(run_id)) + .fetch_all(&self.pool) + .await?; + rows.iter().map(|r| view(r, cpu_psi_pct)).collect() + } + + /// Drop rows granted before `before`, then all but the newest `max_rows`. + pub async fn prune_lease_hosts(&self, before: i64, max_rows: i64) -> Result { + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let old = sqlx::query("DELETE FROM lease_hosts WHERE leased_at < ?") + .bind(before) + .execute(&self.pool) + .await? + .rows_affected(); + // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database + let over = sqlx::query( + "DELETE FROM lease_hosts WHERE lease_id IN + (SELECT lease_id FROM lease_hosts ORDER BY leased_at DESC, lease_id DESC LIMIT -1 OFFSET ?)", + ) + .bind(max_rows) + .execute(&self.pool) + .await? + .rows_affected(); + Ok(old + over) + } +} + +fn view(r: &sqlx::sqlite::SqliteRow, cpu_psi_pct: f64) -> Result { + let sample = |col: &str| -> Result, StoreError> { + r.get::, _>(col) + .map(|s| serde_json::from_str(&s).map_err(|e| StoreError::Corrupt(e.to_string()))) + .transpose() + }; + let start = sample("start_sample")?; + let end = sample("end_sample")?; + let during = start + .as_ref() + .zip(end.as_ref()) + .and_then(|(s, e)| during(s, e)); + let reasons = slow_host_reasons(start.as_ref(), during.as_ref(), cpu_psi_pct); + let u = |v: Option| v.and_then(|v| u64::try_from(v).ok()); + Ok(JobHostView { + lease_id: r.get("lease_id"), + drone: r.get("drone"), + class: enum_from(&r.get::("class"))?, + scope: enum_from(&r.get::("scope"))?, + repo: r.get("repo"), + run_id: u(r.get("run_id")), + job_id: u(r.get("job_id")), + job_name: r.get("job_name"), + leased_at: r.get("leased_at"), + finished_at: r.get("finished_at"), + duration_secs: u(r.get("duration_secs")), + exit_code: r + .get::, _>("exit_code") + .and_then(|v| i32::try_from(v).ok()), + start, + end, + during, + slow_host: !reasons.is_empty(), + slow_host_reasons: reasons, + }) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::store::LeaseStatus; + + #[derive(Default)] + struct FakeHost { + files: HashMap, + df: HashMap, + } + + impl Host for FakeHost { + fn read_to_string(&self, path: &Path) -> Option { + self.files.get(path.to_str()?).cloned() + } + fn exists(&self, path: &Path) -> bool { + path.to_str().is_some_and(|p| self.files.contains_key(p)) + } + fn run(&self, program: &str, args: &[&str]) -> Option { + (program == "df") + .then(|| self.df.get(*args.last()?).cloned()) + .flatten() + } + fn os(&self) -> &'static str { + "linux" + } + fn arch(&self) -> &'static str { + "x86_64" + } + fn available_parallelism(&self) -> Option { + Some(16) + } + } + + fn df(avail_kib: u64) -> String { + format!("fs blocks used avail cap mount\nfs 999999999 0 {avail_kib} 0% /x") + } + + /// An LXC build box shared with interactive builds, overloaded. + fn lxc_host(root_usage: u64) -> FakeHost { + let files = [ + ("/proc/loadavg", "93.40 78.56 69.31 90/3277 3570264\n".to_owned()), + ( + "/proc/pressure/cpu", + "some avg10=53.99 avg60=52.14 avg300=47.57 total=1\nfull avg10=23.75 avg60=25.01 avg300=21.45 total=1\n".to_owned(), + ), + ( + "/proc/pressure/memory", + "some avg10=1.84 avg60=0.55 avg300=0.77 total=1\nfull avg10=0.00 avg60=0.00 avg300=0.21 total=1\n".to_owned(), + ), + ( + "/sys/fs/cgroup/cpu.pressure", + "some avg10=54.89 avg60=52.37 avg300=47.65 total=1\nfull avg10=23.61 avg60=24.95 avg300=21.46 total=1\n".to_owned(), + ), + ( + "/sys/fs/cgroup/cpu.stat", + format!("usage_usec {root_usage}\nuser_usec 1\nnr_throttled 0\nthrottled_usec 0\n"), + ), + ("/proc/self/cgroup", "0::/system.slice/borg-drone.service\n".to_owned()), + ( + "/sys/fs/cgroup/system.slice/borg-drone.service/cpu.stat", + "usage_usec 27600656603\nthrottled_usec 5\n".to_owned(), + ), + ("/proc/meminfo", "MemTotal: 1 kB\nMemAvailable: 4194304 kB\n".to_owned()), + ]; + FakeHost { + files: files.into_iter().map(|(k, v)| (k.to_owned(), v)).collect(), + df: HashMap::from([ + ("/work".to_owned(), df(30 * 1024 * 1024)), + ("/cache".to_owned(), df(5 * 1024 * 1024)), + ]), + } + } + + #[test] + fn psi_parses_some_and_full_and_rejects_empty() { + let p = Psi::parse("some avg10=53.99 avg60=52.14 avg300=1 total=2\nfull avg10=23.75 avg60=25.01 avg300=21.45 total").unwrap(); + assert_eq!(p.some_avg10, Some(53.99)); + assert_eq!(p.some_avg60, Some(52.14)); + assert_eq!(p.full_avg10, Some(23.75)); + assert_eq!(p.full_avg60, Some(25.01)); + assert_eq!(Psi::parse("nothing here"), None); + assert_eq!( + Psi::parse("some avg10=500 avg60=nan").unwrap().some_avg10, + Some(100.0) + ); + } + + #[test] + fn sample_reads_load_pressure_cgroups_disk_and_meter() { + let meter = CpuMeter::default(); + // A heartbeat ten seconds earlier primes the meter. + assert_eq!(observe_cpu(&lxc_host(1_000_000_000), &meter, 0), None); + // 120 core-seconds in 10 s: 12 cores busy. + let s = sample( + &lxc_host(1_120_000_000), + Path::new("/work"), + Path::new("/cache"), + 2, + &meter, + 10_000, + ); + assert_eq!(s.at, 10); + assert_eq!(s.cpus, Some(16)); + assert_eq!(s.load, Some([93.40, 78.56, 69.31])); + assert_eq!(s.cpu_psi.unwrap().some_avg10, Some(53.99)); + assert_eq!(s.cgroup_cpu_psi.unwrap().some_avg10, Some(54.89)); + assert_eq!(s.cpu_some_avg10(), Some(54.89)); + assert_eq!(s.memory_psi.unwrap().some_avg10, Some(1.84)); + assert_eq!(s.io_psi, None); + assert_eq!(s.host_cpu_usage_usec, Some(1_120_000_000)); + assert_eq!(s.host_cpu_busy_cores, Some(12.0)); + assert_eq!(s.drone_cpu_usage_usec, Some(27_600_656_603)); + assert_eq!(s.drone_cpu_throttled_usec, Some(5)); + assert_eq!(s.mem_available_mib, Some(4096)); + assert_eq!(s.work_free_mib, Some(30 * 1024)); + assert_eq!(s.cache_free_mib, Some(5 * 1024)); + assert_eq!(s.active_leases, 2); + // An empty host degrades every field to None. + let bare = sample( + &FakeHost::default(), + Path::new("/work"), + Path::new("/cache"), + 1, + &CpuMeter::default(), + 5_000, + ); + assert_eq!(bare.load, None); + assert_eq!(bare.cpu_psi, None); + assert_eq!(bare.host_cpu_busy_cores, None); + assert_eq!(bare.drone_cpu_usage_usec, None); + } + + #[test] + fn meter_keeps_the_last_rate_for_close_readings_and_resets_on_counter_reset() { + let m = CpuMeter::default(); + assert_eq!(m.observe(0, 0), None); + assert_eq!(m.observe(10_000, 40_000_000), Some(4.0)); + // A lease starting 500 ms after the heartbeat reuses its rate. + assert_eq!(m.observe(10_500, 40_100_000), Some(4.0)); + assert_eq!(m.observe(20_000, 120_000_000), Some(8.0)); + // Counter went backwards (cgroup recreated): no rate until next reading. + assert_eq!(m.observe(30_000, 5), None); + assert_eq!(m.observe(40_000, 10_000_005), Some(1.0)); + } + + #[test] + fn own_cgroup_rejects_escapes() { + let mut h = FakeHost::default(); + h.files + .insert("/proc/self/cgroup".into(), "0::/../etc\n".into()); + assert_eq!(own_cgroup(&h), None); + h.files.insert("/proc/self/cgroup".into(), "0::/\n".into()); + assert_eq!(own_cgroup(&h), Some(Path::new("/sys/fs/cgroup").into())); + } + + #[test] + fn slow_host_uses_pressure_and_ignores_lxc_load_when_pressure_exists() { + let start = HostSample { + at: 0, + cpus: Some(16), + load: Some([93.4, 78.0, 69.0]), + cgroup_cpu_psi: Some(Psi { + some_avg10: Some(10.0), + ..Psi::default() + }), + ..HostSample::default() + }; + // Pressure below threshold: lxcfs load average alone is not evidence. + assert!(slow_host_reasons(Some(&start), None, 30.0).is_empty()); + let hot = HostSample { + cgroup_cpu_psi: Some(Psi { + some_avg10: Some(54.9), + ..Psi::default() + }), + ..start.clone() + }; + let r = slow_host_reasons(Some(&hot), None, 30.0); + assert_eq!(r, vec!["cpu pressure 54.9% at start (threshold 30%)"]); + // No pressure at all: fall back to load per cpu. + let no_psi = HostSample { + cgroup_cpu_psi: None, + ..start.clone() + }; + let r = slow_host_reasons(Some(&no_psi), None, 30.0); + assert_eq!(r.len(), 1); + assert!(r[0].starts_with("load 93.4 on 16 cpus"), "{r:?}"); + // A quiet burst drone is not slow. + let quiet = HostSample { + at: 0, + cpus: Some(56), + load: Some([20.0, 20.0, 20.0]), + cpu_psi: Some(Psi { + some_avg10: Some(0.5), + ..Psi::default() + }), + host_cpu_busy_cores: Some(20.0), + work_free_mib: Some(500_000), + cache_free_mib: Some(500_000), + ..HostSample::default() + }; + assert!(slow_host_reasons(Some(&quiet), None, 30.0).is_empty()); + let busy = During { + secs: 100, + host_cpu_busy_cores: Some(55.0), + drone_cpu_busy_cores: None, + drone_cpu_throttled_secs: None, + }; + assert_eq!( + slow_host_reasons(Some(&quiet), Some(&busy), 30.0), + vec!["55.0 of 56 cpus busy on average during the job"] + ); + assert!(slow_host_reasons(None, None, 30.0).is_empty()); + } + + #[test] + fn sanitize_drops_non_finite_and_future_times() { + let s = HostSample { + at: 999, + load: Some([f64::NAN, 1.0, 1.0]), + host_cpu_busy_cores: Some(-3.0), + cpu_psi: Some(Psi { + some_avg10: Some(f64::INFINITY), + some_avg60: Some(250.0), + ..Psi::default() + }), + ..HostSample::default() + } + .sanitized(100); + assert_eq!(s.at, 100); + assert_eq!(s.load, None); + assert_eq!(s.host_cpu_busy_cores, Some(0.0)); + assert_eq!(s.cpu_psi.unwrap().some_avg10, None); + assert_eq!(s.cpu_psi.unwrap().some_avg60, Some(100.0)); + } + + fn lease(id: &str, created_at: i64) -> LeaseRow { + LeaseRow { + id: id.into(), + drone: "shared-box".into(), + class: Class::Heavy, + scope: Scope::Main, + labels: vec![], + runner_name: format!("borg-{id}"), + runner_id: 1, + jit_config: None, + state: LeaseStatus::Running, + created_at, + expires_at: created_at + 600, + } + } + + #[tokio::test] + async fn rows_bind_jobs_and_results_in_either_order_and_prune() { + let s = Store::open_in_memory().await.unwrap(); + let (a, b, c) = (lease("a", 100), lease("b", 200), lease("c", 300)); + let result = LeaseResult { + exit_code: 1, + duration_secs: 90, + peak_rss_mib: None, + }; + let start = HostSample { + at: 110, + cpus: Some(16), + cpu_psi: Some(Psi { + some_avg10: Some(54.0), + ..Psi::default() + }), + host_cpu_usage_usec: Some(0), + drone_cpu_usage_usec: Some(0), + ..HostSample::default() + }; + let end = HostSample { + at: 210, + host_cpu_usage_usec: Some(1_500_000_000), + drone_cpu_usage_usec: Some(400_000_000), + ..start.clone() + }; + let report = LeaseHostReport { + start: Some(start), + end: Some(end), + }; + // a: job seen first, result later. + s.bind_lease_host_job(&a, "o/r", 7, 70, "e2e (shard 3)") + .await + .unwrap(); + s.record_lease_host(&a, &result, Some(&report), 220) + .await + .unwrap(); + // b: result first (old drone, no samples), job later. + s.record_lease_host(&b, &result, None, 230).await.unwrap(); + s.bind_lease_host_job(&b, "o/r", 7, 71, "lint") + .await + .unwrap(); + // c: another run. + s.bind_lease_host_job(&c, "o/r", 8, 80, "x").await.unwrap(); + + let v = s.job_host(70, 30.0).await.unwrap().unwrap(); + assert_eq!(v.drone, "shared-box"); + assert_eq!((v.run_id, v.job_id), (Some(7), Some(70))); + assert_eq!(v.job_name.as_deref(), Some("e2e (shard 3)")); + assert_eq!( + (v.exit_code, v.duration_secs, v.finished_at), + (Some(1), Some(90), Some(220)) + ); + let d = v.during.unwrap(); + assert_eq!(d.secs, 100); + assert_eq!(d.host_cpu_busy_cores, Some(15.0)); + assert_eq!(d.drone_cpu_busy_cores, Some(4.0)); + assert!(v.slow_host); + assert_eq!(v.slow_host_reasons.len(), 2, "{:?}", v.slow_host_reasons); + + let old = s.job_host(71, 30.0).await.unwrap().unwrap(); + assert_eq!(old.start, None); + assert!(!old.slow_host); + assert_eq!(old.exit_code, Some(1)); + + let run = s.run_hosts(7, 30.0).await.unwrap(); + assert_eq!( + run.iter().map(|v| v.lease_id.as_str()).collect::>(), + ["a", "b"] + ); + assert!(s.job_host(999, 30.0).await.unwrap().is_none()); + + // Re-binding the same job is a no-op; a retried result keeps samples. + s.bind_lease_host_job(&a, "o/r", 7, 70, "e2e (shard 3)") + .await + .unwrap(); + s.record_lease_host(&a, &result, None, 240).await.unwrap(); + assert!(s.job_host(70, 30.0).await.unwrap().unwrap().start.is_some()); + + // Age prunes a; the row cap then keeps only the newest one (c). + assert_eq!(s.prune_lease_hosts(150, 1).await.unwrap(), 2); + assert!(s.job_host(70, 30.0).await.unwrap().is_none()); + assert!(s.job_host(71, 30.0).await.unwrap().is_none()); + assert!(s.job_host(80, 30.0).await.unwrap().is_some()); + } +} diff --git a/src/lib.rs b/src/lib.rs index f54d334..68dcafe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,7 @@ pub mod explain; pub mod github; pub mod health; pub mod homeostasis; +pub mod hostdiag; pub mod janitor; pub mod lead; pub mod mirror_auth; diff --git a/src/preempt.rs b/src/preempt.rs index ddb3628..e3a9dd9 100644 --- a/src/preempt.rs +++ b/src/preempt.rs @@ -210,6 +210,7 @@ mod tests { blocked_classes: Vec::new(), trust: crate::classes::Trust::Shared, cpu_psi_pm: None, + cpu_load_pm: None, } } diff --git a/src/probe/mod.rs b/src/probe/mod.rs index 76b9904..586ee98 100644 --- a/src/probe/mod.rs +++ b/src/probe/mod.rs @@ -89,10 +89,24 @@ pub fn env_facts(host: &dyn Host, now: i64) -> EnvFacts { observed_at: now, mem_available_mib: mem_available_mib(host), cpu_psi_some_avg10_pm: cpu_psi_pm(host), + load_per_cpu_pm: load_per_cpu_pm(host), ..EnvFacts::default() } } +/// The 1-minute load average per usable CPU, in permille. +fn load_per_cpu_pm(host: &dyn Host) -> Option { + let load = host + .read_to_string(Path::new("/proc/loadavg"))? + .split_whitespace() + .next()? + .parse::() + .ok() + .filter(|l| l.is_finite() && *l >= 0.0)?; + let cpus = f64::from(host.available_parallelism()?.max(1)); + Some((load / cpus * 1000.0).round().min(f64::from(u32::MAX)) as u32) +} + /// How long a tools probe is reused before the tools are looked up again. const TOOLS_TTL_SECS: i64 = 60; /// Longest version string reported for a tool. @@ -162,7 +176,7 @@ impl ToolsCache { } } -fn mem_available_mib(host: &dyn Host) -> Option { +pub(crate) fn mem_available_mib(host: &dyn Host) -> Option { let meminfo = host.read_to_string(Path::new("/proc/meminfo"))?; let kib = meminfo .lines() @@ -510,6 +524,18 @@ mod tests { assert!(exec.calls.lock().unwrap().len() > n); } + #[test] + fn load_per_cpu_is_permille_of_usable_cpus() { + let mut h = linux(); + assert_eq!(load_per_cpu_pm(&h), None); + h.files + .insert("/proc/loadavg", "93.40 78.56 69.31 90/3277 3570264\n"); + assert_eq!(load_per_cpu_pm(&h), Some(5838)); + assert_eq!(env_facts(&h, 1).load_per_cpu_pm, Some(5838)); + h.files.insert("/proc/loadavg", "garbage"); + assert_eq!(load_per_cpu_pm(&h), None); + } + #[test] fn psi_is_the_higher_of_host_and_cgroup_pressure() { let mut h = linux(); diff --git a/src/queen.rs b/src/queen.rs index b983335..15ff14d 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -595,12 +595,15 @@ impl Queen { None, now_ts, ); - if let Some(a) = &self.cfg.admission { - state.cpu_psi_pm = d + if let Some(a) = &self.cfg.admission + && let Some(f) = d .env_facts .as_ref() .filter(|f| now_ts - f.observed_at <= a.facts_max_age_secs) - .and_then(|f| f.cpu_psi_some_avg10_pm); + { + state.cpu_psi_pm = f.cpu_psi_some_avg10_pm; + state.cpu_load_pm = + crate::allocator::load_pressure_pm(f.load_per_cpu_pm, f.cpu_busy_pm); } state.profiles = profiles .iter() @@ -963,6 +966,16 @@ impl Queen { if let Err(e) = self.store.prune_telemetry(now_ts - keep).await { tracing::warn!("prune telemetry failed: {e}"); } + if let Err(e) = self + .store + .prune_lease_hosts( + now_ts - crate::hostdiag::KEEP_SECS, + crate::hostdiag::MAX_ROWS, + ) + .await + { + tracing::warn!("prune lease host diagnostics failed: {e}"); + } // Grants and the snapshot are published before optional feedback. Bound // this pass so unavailable job history cannot stall subsequent ticks. match tokio::time::timeout( @@ -1366,11 +1379,10 @@ impl Queen { groups[&(u.platform.as_str(), u.rust_toolchain.as_deref())] .iter() .filter(|d| { - let state = AdmissionState::of(d.cpu_psi_pm, admission); d.accepting && d.trust.admits(u.scope) - && state != AdmissionState::Closed - && (state != AdmissionState::PrBlocked || u.scope != Scope::Pr) + && AdmissionState::of(d.admission_pressure_pm(), admission) + .admits(u.class, u.scope) && !d.blocked_classes.contains(&u.class) && max_slots(d, self.cfg.classes.spec(u.class)) > 0 && (d.bids.is_empty() @@ -1641,6 +1653,9 @@ impl Queen { && let Some(lease) = self.store.lease_for_runner(id, name).await? { self.store.remember_job(&lease, repo, job.id, now).await?; + self.store + .bind_lease_host_job(&lease, repo, job.run_id, job.id, &job.name) + .await?; } } let mut pending = self.store.pending_jobs().await?; @@ -1687,6 +1702,9 @@ impl Queen { { continue; } + self.store + .bind_lease_host_job(&lease, &repo, job.run_id, job.id, &job.name) + .await?; if let Some((telemetry, duration, finished)) = job.telemetry(&repo) { self.store .record_github_telemetry(&lease, &telemetry, duration, finished) diff --git a/src/store.rs b/src/store.rs index fff1d44..6e6e150 100644 --- a/src/store.rs +++ b/src/store.rs @@ -178,6 +178,27 @@ CREATE TABLE IF NOT EXISTS job_telemetry ( finished_at INTEGER NOT NULL ); CREATE INDEX IF NOT EXISTS job_telemetry_finished ON job_telemetry(finished_at); +-- Host load at lease start and end, with the GitHub job the lease ran; see +-- hostdiag.rs. Pruned after seven days and to a row cap. +CREATE TABLE IF NOT EXISTS lease_hosts ( + lease_id TEXT PRIMARY KEY, + drone TEXT NOT NULL, + class TEXT NOT NULL, + scope TEXT NOT NULL, + leased_at INTEGER NOT NULL, + repo TEXT, + run_id INTEGER, + job_id INTEGER, + job_name TEXT, + finished_at INTEGER, + exit_code INTEGER, + duration_secs INTEGER, + start_sample TEXT, + end_sample TEXT +); +CREATE INDEX IF NOT EXISTS lease_hosts_job ON lease_hosts(job_id); +CREATE INDEX IF NOT EXISTS lease_hosts_run ON lease_hosts(run_id); +CREATE INDEX IF NOT EXISTS lease_hosts_leased ON lease_hosts(leased_at); CREATE TABLE IF NOT EXISTS preemptions ( lease_id TEXT PRIMARY KEY, drone TEXT NOT NULL, diff --git a/tests/drone.rs b/tests/drone.rs index 75da712..2629fa7 100644 --- a/tests/drone.rs +++ b/tests/drone.rs @@ -198,6 +198,19 @@ async fn runs_a_granted_lease_cleans_up_and_reports_the_result() { .is_empty(), "result reported, lease done" ); + let diag = rig + .h + .queen + .store() + .lease_host(&lease_id, 30.0) + .await + .unwrap() + .expect("host samples stored with the result"); + assert_eq!(diag.drone, "d1"); + assert_eq!(diag.exit_code, Some(0)); + let (start, end) = (diag.start.unwrap(), diag.end.unwrap()); + assert_eq!((start.cpus, start.active_leases), (Some(16), 1)); + assert!(end.at >= start.at); let calls = rig.exec.0.lock().unwrap().clone(); assert!( calls diff --git a/tests/queen.rs b/tests/queen.rs index a3dfe63..f838da1 100644 --- a/tests/queen.rs +++ b/tests/queen.rs @@ -785,6 +785,168 @@ async fn observed_running_job_is_collected_after_workflow_leaves_the_queue() { assert!(h.queen.store().pending_jobs().await.unwrap().is_empty()); } +#[tokio::test] +async fn job_and_run_endpoints_show_the_drone_and_its_host_load() { + let h = Harness::start( + json!({"admission": {"psi_soft_pm": 300, "psi_hard_pm": 600, "facts_max_age_secs": 180}}), + ) + .await; + let token = h.enroll("shared-box").await; + h.heartbeat(&token, &[]).await; + let lease_id = lease_for(&h, &token, "borg-ci-heavy", 4242).await; + let lease = h.queen.store().lease(&lease_id).await.unwrap().unwrap(); + let start = OffsetDateTime::now_utc() - time::Duration::seconds(45); + let ts = start + .format(&time::format_description::well_known::Rfc3339) + .unwrap(); + h.state().jobs = vec![json!({ + "id": 4242, "run_id": 1, "name": "ui-e2e (shard 3)", "status": "in_progress", + "conclusion": null, "labels": ["borg-ci-heavy"], + "runner_id": lease.runner_id, "runner_name": lease.runner_name, + "created_at": ts, "started_at": ts, "completed_at": null, "steps": [] + })]; + h.tick().await; + + // Seen running, no result yet: the drone is known, samples are not. + let job: Value = h + .http + .get(format!("{}/jobs/4242", h.url)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(job["drone"], "shared-box"); + // The mock serves this job under queued run 1. + assert_eq!(job["run_id"], 1); + assert_eq!(job["job_name"], "ui-e2e (shard 3)"); + assert_eq!(job["start"], Value::Null); + assert_eq!(job["slow_host"], false); + + let status = post_result( + &h, + &token, + &lease_id, + json!({"exit_code": 1, "duration_secs": 100, "host": { + "start": {"at": 1000, "cpus": 16, "load": [77.0, 60.0, 50.0], + "cgroup_cpu_psi": {"some_avg10": 54.9, "some_avg60": 52.4}, + "host_cpu_usage_usec": 0, "drone_cpu_usage_usec": 0, + "work_free_mib": 30000, "cache_free_mib": 30000, "active_leases": 2}, + "end": {"at": 1100, "cpus": 16, "host_cpu_usage_usec": 1550000000u64, + "drone_cpu_usage_usec": 300000000, "active_leases": 1}}}), + ) + .await; + assert_eq!(status, 204); + + let job: Value = h + .http + .get(format!("{}/jobs/4242", h.url)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(job["exit_code"], 1); + assert_eq!(job["duration_secs"], 100); + assert_eq!(job["start"]["active_leases"], 2); + assert_eq!(job["during"]["host_cpu_busy_cores"], 15.5); + assert_eq!(job["during"]["drone_cpu_busy_cores"], 3.0); + assert_eq!(job["slow_host"], true); + let reasons = job["slow_host_reasons"].as_array().unwrap(); + assert!( + reasons[0] + .as_str() + .unwrap() + .starts_with("cpu pressure 54.9% at start"), + "{reasons:?}" + ); + + let run: Value = h + .http + .get(format!("{}/runs/1", h.url)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(run["jobs"].as_array().unwrap().len(), 1); + assert_eq!(run["jobs"][0]["job_id"], 4242); + + let missing = h + .http + .get(format!("{}/jobs/1", h.url)) + .send() + .await + .unwrap(); + assert_eq!(missing.status(), 404); + let empty: Value = h + .http + .get(format!("{}/runs/2", h.url)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(empty["jobs"], json!([])); +} + +#[tokio::test] +async fn admission_falls_back_to_load_without_psi_and_light_only_sheds_main_heavy() { + let h = Harness::start(json!({"admission": { + "psi_soft_pm": 300, "psi_light_only_pm": 500, "psi_hard_pm": 600, "facts_max_age_secs": 180}})) + .await; + let token = h.enroll("no-psi").await; + let now = OffsetDateTime::now_utc().unix_timestamp(); + // No PSI, load 93 on 16 saturated CPUs: estimated 828 permille, closed. + h.heartbeat_with( + &token, + json!({"env_facts": {"observed_at": now, "load_per_cpu_pm": 5838, "cpu_busy_pm": 990}}), + ) + .await; + h.queue_job(1, "borg-ci-heavy", 5); + h.tick().await; + let s = h.status().await; + let a = &s["drones"][0]["admission"]; + assert_eq!(a["state"], "closed", "{a}"); + assert_eq!(a["cpu_psi_pm"], Value::Null); + assert_eq!(a["cpu_load_pm"], 828); + assert!( + h.state().mints.is_empty(), + "an overloaded drone without PSI gets no main heavy" + ); + assert!( + h.metrics() + .await + .contains("borg_drone_cpu_load_pm{drone=\"no-psi\"} 828") + ); + + // With PSI 540 the drone is light-only: main light placed, main heavy not. + h.heartbeat_with( + &token, + json!({"env_facts": {"observed_at": now, "cpu_psi_some_avg10_pm": 540, + "load_per_cpu_pm": 5838, "cpu_busy_pm": 990}}), + ) + .await; + h.queue_job(2, "borg-ci-light", 5); + h.tick().await; + assert_eq!( + h.status().await["drones"][0]["admission"]["state"], + "light_only" + ); + let minted: Vec = h + .state() + .mints + .iter() + .map(|m| m["labels"].to_string()) + .collect(); + assert_eq!(minted.len(), 1, "{minted:?}"); + assert!(minted[0].contains("borg-ci-light"), "{minted:?}"); +} + #[tokio::test] async fn github_telemetry_requires_exact_runner_identity_and_preserves_drone_evidence() { let h = Harness::start(json!({})).await; From 46f67392af97e92ea6bda500b272812fd44479ae Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 17 Oct 2026 01:51:47 +0000 Subject: [PATCH 23/27] chore: project borg from Mono 6b41db420e1a --- .repository-projection.json | 6 +-- docs/operations.md | 17 +++++- src/allocator.rs | 100 ++++++++++++++++++++++++++++++++++-- src/bid.rs | 9 ++++ src/explain.rs | 3 +- 5 files changed, 126 insertions(+), 9 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 36e2b12..d6a2d76 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "54af7346321ae5f3a1aeb1ae703bd34e80c204ec", + "sourceSha": "6b41db420e1a0d7faafc56c1b313e0665c19e95d", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "ec3cb38cabf0371e8e514cab74a1d0b494c02395", + "priorProjectedBase": "5381edcccfcd725e2e26f0f61ce52c169c0aaaf4", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", - "contentDigest": "cd2134012915159bfbe895ad883939b40f10ccea2675fd6869366f7c13de4979", + "contentDigest": "b572db642d7e2ac167a63d2e052032cc18be7af05571d2d24a3188d404c03ea0", "publicationEligible": true } diff --git a/docs/operations.md b/docs/operations.md index c3e72c6..f81b8cb 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -321,8 +321,9 @@ The allocator drops a bid with no `available_slots` or above the current pressure, and scores the rest with integers: `3*cache_affinity - 3*resource_pressure - 2*start - 1*(cost-1000)`, where `start` is `estimated_start_ms` as permille of 60 s, all scaled by -`confidence/1000`. The highest score wins; ties go to the tightest memory fit, -then the drone name. A drone that sends no bids, or only expired ones, is +`confidence/1000`. A drone's total score is that bid score plus its learned +profile term (see [Learned profiles](#learned-profiles)) plus a host term. The +highest score wins; ties go to the tightest memory fit, then the drone name. A drone that sends no bids, or only expired ones, is treated as before (always eligible, neutral score); a drone that bids but not for a key declines that key. GitHub still chooses which queued job a runner takes: bids steer only which host receives a freshly minted runner. @@ -570,6 +571,18 @@ proportion to the evidence, so proven faster and more reliable hosts are preferred when placing a runner. See `profiles` in the queen config and `profiles[]` in `/status`. +The host term covers how loaded a drone is right now: +`-1 * pressure + 250 * free_vcpu / budget_vcpu` for medium and heavy work, and +`-1 * pressure` for light work. Here `pressure` is the drone's fresh CPU +pressure in permille: PSI, or the load estimate without PSI (see +[Admission by CPU pressure](#admission-by-cpu-pressure)), and 0 when unknown or +when admission is off. With both terms, main heavy and medium work goes to +drones that are fast and idle first. A host that is historically slow (speed +550 permille costs 450 points) or loaded (PSI 450 permille costs 450 points) +gets work only when nothing better has room, and above the admission +thresholds it gets none. Light work keeps packing onto the tightest fit. +`/explain` shows each candidate's `host` term next to `bid` and `profile`. + ### Warm-floor controller Optional (`homeostasis.enabled`). A damped feedback loop raises the warm floor of diff --git a/src/allocator.rs b/src/allocator.rs index 5900680..74638f4 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -381,9 +381,29 @@ pub struct Candidate { pub bid: Option, pub bid_score: i64, pub profile_score: i64, + /// Current load: minus the drone's CPU pressure, plus its idle share for + /// medium and heavy work. See [`host_term`]. + #[serde(default)] + pub host_score: i64, pub score: Option, } +/// The placement term for how loaded a drone is right now: `-host_pressure * +/// pressure` (PSI, or the load estimate without PSI; unknown counts as 0) and, +/// for medium and heavy work, `+idle_share * free vCPU share`. Light work packs +/// as before. Together with the learned speed in the profile term, main heavy +/// and medium work goes to fast, idle drones first and to slow or loaded ones +/// only when nothing better has room. +pub fn host_term(drone: &DroneState, free_vcpu: u32, class: Class, w: &ScoreWeights) -> i64 { + let pressure = i64::from(drone.admission_pressure_pm().unwrap_or(0).min(1000)); + let idle = if class == Class::Light || drone.vcpu == 0 { + 0 + } else { + i64::from(free_vcpu.min(drone.vcpu)) * 1000 / i64::from(drone.vcpu) + }; + -w.host_pressure * pressure + w.idle_share * idle / 1000 +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PlacementDecision { pub class: Class, @@ -478,6 +498,7 @@ struct BidEvaluation { bid: Option, bid_score: i64, profile_score: i64, + host_score: i64, rejection_reasons: Vec, score: Option, } @@ -496,6 +517,7 @@ impl Free<'_> { .iter() .find(|p| p.class == class) .map_or(0, |p| profile_term(p, weights)); + let host_score = host_term(self.drone, self.vcpu, class, weights); let bid = self .drone .bids @@ -522,9 +544,12 @@ impl Free<'_> { bid, bid_score, profile_score, - score: rejection_reasons - .is_empty() - .then_some(bid_score.saturating_add(profile_score)), + host_score, + score: rejection_reasons.is_empty().then_some( + bid_score + .saturating_add(profile_score) + .saturating_add(host_score), + ), rejection_reasons, } } @@ -616,6 +641,7 @@ impl Free<'_> { bid: evaluated.bid, bid_score: evaluated.bid_score, profile_score: evaluated.profile_score, + host_score: evaluated.host_score, } } @@ -1124,6 +1150,74 @@ mod tests { } } + fn profile(class: Class, speed: u32) -> ProfileScore { + ProfileScore { + class, + speed_permille: speed, + speed_confidence_permille: 1000, + reliability_permille: 1000, + } + } + + #[test] + fn main_heavy_prefers_fast_idle_drones_and_slow_or_loaded_ones_take_overflow() { + let fast = DroneState { + profiles: vec![profile(Class::Heavy, 1100)], + ..host("fast") + }; + // Ids sort before "fast", so ties would go to them. + // As fast, but its host is under CPU pressure right now. + let loaded = DroneState { + cpu_psi_pm: Some(450), + profiles: vec![profile(Class::Heavy, 1100)], + ..host("a-loaded") + }; + // Idle, but historically half as fast. + let slow = DroneState { + profiles: vec![profile(Class::Heavy, 550)], + ..host("0-slow") + }; + let drones = vec![loaded.clone(), slow.clone(), fast.clone()]; + let first = |n| { + let a = allocate(&input(&[(Class::Heavy, Scope::Main, n)], drones.clone())); + a.grants + .iter() + .map(|g| g.drone_id.clone()) + .collect::>() + }; + // One default heavy fits per 28 GiB host. Overflow goes to the less + // loaded of the rest; the historically slow host comes last. + assert_eq!(first(1), ["fast"]); + assert_eq!(first(2), ["fast", "a-loaded"]); + assert_eq!(first(3), ["fast", "a-loaded", "0-slow"]); + + // Same speed: the idle host wins medium/heavy, light still packs. + let busy = host("busy"); + let idle = host("idle"); + let mut i = input( + &[(Class::Medium, Scope::Main, 1)], + vec![busy.clone(), idle.clone()], + ); + i.leases = vec![lease( + "busy", + Class::Medium, + Scope::Main, + LeaseState::Running, + )]; + assert_eq!(allocate(&i).grants[0].drone_id, "idle"); + i.demand = vec![Demand { + class: Class::Light, + scope: Scope::Main, + count: 1, + }]; + assert_eq!(allocate(&i).grants[0].drone_id, "busy"); + + let w = ScoreWeights::default(); + assert_eq!(host_term(&loaded, 16, Class::Heavy, &w), -450 + 250); + assert_eq!(host_term(&slow, 8, Class::Heavy, &w), 125); + assert_eq!(host_term(&loaded, 16, Class::Light, &w), -450); + } + #[test] fn load_estimate_approximates_waiting_share_and_ignores_idle_cpus() { assert_eq!(load_pressure_pm(None, Some(1000)), None); diff --git a/src/bid.rs b/src/bid.rs index 3b807ab..78f683c 100644 --- a/src/bid.rs +++ b/src/bid.rs @@ -98,6 +98,13 @@ pub struct ScoreWeights { pub profile_speed: i64, /// Applied to a learned profile's reliability shortfall (permille). pub profile_reliability: i64, + /// Applied to the drone's current CPU pressure (permille; PSI, or the + /// load estimate without PSI): a loaded host scores lower. + pub host_pressure: i64, + /// Applied to the share of the drone's vCPU budget still free (permille, + /// divided by 1000) for medium and heavy work: an idle host scores higher + /// than one already running builds. + pub idle_share: i64, } impl Default for ScoreWeights { @@ -109,6 +116,8 @@ impl Default for ScoreWeights { cost: 1, profile_speed: 1, profile_reliability: 3, + host_pressure: 1, + idle_share: 250, } } } diff --git a/src/explain.rs b/src/explain.rs index 0261c1b..b67fda1 100644 --- a/src/explain.rs +++ b/src/explain.rs @@ -109,7 +109,7 @@ pub fn render_text(report: &Explanation) -> String { for host in &d.candidates { let _ = writeln!( out, - " {}: {} / {} MiB, {} / {} vCPU free; lease slots {:?}; score {:?} (bid {}, profile {}); {:?}; reservation {:?}", + " {}: {} / {} MiB, {} / {} vCPU free; lease slots {:?}; score {:?} (bid {}, profile {}, host {}); {:?}; reservation {:?}", host.drone_id, host.room.mem_mib, host.budget.mem_mib, @@ -119,6 +119,7 @@ pub fn render_text(report: &Explanation) -> String { host.score, host.bid_score, host.profile_score, + host.host_score, host.rejection_reasons, host.reserved_for ); From a1bad3ebce7c1689dc8ab8717c404708b3333a32 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 17 Oct 2026 06:18:06 +0000 Subject: [PATCH 24/27] chore: project borg from Mono ee64d86fe58c --- .repository-projection.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index d6a2d76..c2d7b43 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,9 +3,9 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "6b41db420e1a0d7faafc56c1b313e0665c19e95d", + "sourceSha": "ee64d86fe58c1b340ccc4d467ecac12b2b27e321", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "5381edcccfcd725e2e26f0f61ce52c169c0aaaf4", + "priorProjectedBase": "99325df741fa71b9ba934b66470e4dad413bc83d", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", "contentDigest": "b572db642d7e2ac167a63d2e052032cc18be7af05571d2d24a3188d404c03ea0", From e0ec7c19eef2a56387e5339001512750f37e15d4 Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 17 Oct 2026 06:41:19 +0000 Subject: [PATCH 25/27] chore: project borg from Mono 6077c90a491c --- .repository-projection.json | 6 +- src/allocator.rs | 19 +++--- src/api.rs | 36 ++++++------ src/bid.rs | 4 +- src/config.rs | 26 ++++---- src/drone/caches/gc.rs | 8 +-- src/envcontract.rs | 4 +- src/hostdiag.rs | 114 +++++++++++++++++++----------------- src/janitor.rs | 4 +- src/queen.rs | 6 +- src/store.rs | 2 +- tests/hostdiag_wire.rs | 33 +++++++++++ tests/suite.rs | 1 + 13 files changed, 152 insertions(+), 111 deletions(-) create mode 100644 tests/hostdiag_wire.rs diff --git a/.repository-projection.json b/.repository-projection.json index c2d7b43..e82230f 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "ee64d86fe58c1b340ccc4d467ecac12b2b27e321", + "sourceSha": "6077c90a491c4384f1915a0c2f7bd244bf30982e", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "99325df741fa71b9ba934b66470e4dad413bc83d", + "priorProjectedBase": "c6d87f59583ab6b292262558a7086011a7e527f1", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", - "contentDigest": "b572db642d7e2ac167a63d2e052032cc18be7af05571d2d24a3188d404c03ea0", + "contentDigest": "7036b28c11658bd46145527ba700dd786cffd76a53418a3c5b3ec9f3ce8c4a06", "publicationEligible": true } diff --git a/src/allocator.rs b/src/allocator.rs index 74638f4..2bba998 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -89,7 +89,7 @@ pub struct DroneState { /// (see [`load_pressure_pm`]). Admission uses it only when `cpu_psi_pm` /// is `None`. #[serde(default)] - pub cpu_load_pm: Option, + pub(crate) cpu_load_pm: Option, } /// CPU pressure (permille) estimated from the 1-minute load average per CPU @@ -98,7 +98,10 @@ pub struct DroneState { /// approximates PSI `cpu some`. When the root cgroup shows the CPUs under 90% /// busy (`cpu_busy_pm`), the run queue is not CPU contention (I/O wait, or an /// lxcfs load average) and the estimate is 0. -pub fn load_pressure_pm(load_per_cpu_pm: Option, cpu_busy_pm: Option) -> Option { +pub(crate) fn load_pressure_pm( + load_per_cpu_pm: Option, + cpu_busy_pm: Option, +) -> Option { let load = u64::from(load_per_cpu_pm?); if load <= 1000 || cpu_busy_pm.is_some_and(|b| b < 900) { return Some(0); @@ -112,7 +115,7 @@ pub struct AdmissionThresholds { pub psi_soft_pm: u16, /// From this pressure only light main and deploy grants; at or above /// `psi_hard_pm` it has no effect. - pub psi_light_only_pm: u16, + pub(crate) psi_light_only_pm: u16, pub psi_hard_pm: u16, } @@ -155,7 +158,7 @@ impl AdmissionState { } } - pub fn admits(self, class: Class, scope: Scope) -> bool { + pub(crate) fn admits(self, class: Class, scope: Scope) -> bool { match self { Self::Open => true, Self::PrBlocked => scope != Scope::Pr, @@ -169,7 +172,7 @@ impl AdmissionState { impl DroneState { /// The pressure admission uses: PSI when measured, else the load estimate. - pub fn admission_pressure_pm(&self) -> Option { + pub(crate) fn admission_pressure_pm(&self) -> Option { self.cpu_psi_pm.or(self.cpu_load_pm) } } @@ -183,7 +186,7 @@ pub struct DroneAdmission { /// The load-average estimate; drives `state` only when `cpu_psi_pm` is /// `None`. #[serde(default)] - pub cpu_load_pm: Option, + pub(crate) cpu_load_pm: Option, } impl DroneState { @@ -384,7 +387,7 @@ pub struct Candidate { /// Current load: minus the drone's CPU pressure, plus its idle share for /// medium and heavy work. See [`host_term`]. #[serde(default)] - pub host_score: i64, + pub(crate) host_score: i64, pub score: Option, } @@ -394,7 +397,7 @@ pub struct Candidate { /// as before. Together with the learned speed in the profile term, main heavy /// and medium work goes to fast, idle drones first and to slow or loaded ones /// only when nothing better has room. -pub fn host_term(drone: &DroneState, free_vcpu: u32, class: Class, w: &ScoreWeights) -> i64 { +fn host_term(drone: &DroneState, free_vcpu: u32, class: Class, w: &ScoreWeights) -> i64 { let pressure = i64::from(drone.admission_pressure_pm().unwrap_or(0).min(1000)); let idle = if class == Class::Light || drone.vcpu == 0 { 0 diff --git a/src/api.rs b/src/api.rs index 701d667..e0e5d06 100644 --- a/src/api.rs +++ b/src/api.rs @@ -159,10 +159,10 @@ pub struct CacheStat { pub misses: Option, /// Whether `size_bytes` covers the whole cache. #[serde(default)] - pub size_status: SizeStatus, + pub(crate) size_status: SizeStatus, /// Whether `hits` and `misses` are real counts. #[serde(default)] - pub counters: Counters, + pub(crate) counters: Counters, } /// How complete a cache's `size_bytes` is. @@ -278,7 +278,7 @@ pub struct ResultRequest { /// Host load when the runner started and when it exited; see /// [`crate::hostdiag`]. #[serde(default, skip_serializing_if = "Option::is_none")] - pub host: Option, + pub(crate) host: Option, } #[derive(Debug)] @@ -736,12 +736,12 @@ pub struct DroneStatus { /// `alive` (heartbeat within `drone_stale_secs`) or `stale`. A stale /// drone holds no schedulable capacity whatever its `health` says; /// `health` is the historical failure evidence only. - pub availability: Availability, + availability: Availability, /// Seconds since the last heartbeat; `None` if it never sent one. - pub last_seen_age_secs: Option, + last_seen_age_secs: Option, /// Seconds until a stale drone is removed (`drone_expiry`); `None` while /// alive or when it is never removed. - pub expires_in_secs: Option, + expires_in_secs: Option, pub accepting: bool, /// Enforced independently of the optional environment contract. pub disk_admission: crate::envcontract::DiskAdmission, @@ -790,18 +790,18 @@ pub enum Availability { #[derive(Debug, Serialize)] pub struct CapacityStatus { /// Slots the alive, accepting drones could hold if empty. - pub schedulable_slots: u32, + schedulable_slots: u32, /// Slots free right now on alive, accepting drones. - pub free_slots: u32, + free_slots: u32, /// Jobs borg could not place at the last pass (`unmet_placeable`). - pub unmet_main: u32, - pub unmet_pr: u32, - pub unmet_deploy: u32, + unmet_main: u32, + unmet_pr: u32, + unmet_deploy: u32, /// Since when main demand has been unmet with schedulable slots at or /// under `capacity_alarm.max_slots`. - pub starved_since: Option, + starved_since: Option, /// The `capacity_starved` health event is raised and not yet restored. - pub alarm: bool, + alarm: bool, } /// What CPU pressure currently lets this drone take. @@ -811,7 +811,7 @@ pub struct DroneAdmissionView { pub cpu_psi_pm: Option, /// Pressure estimated from the load average; `state` uses it only when /// `cpu_psi_pm` is null. - pub cpu_load_pm: Option, + cpu_load_pm: Option, } /// One drone's build environment contract. @@ -857,9 +857,9 @@ pub struct Status { pub queue: Vec, pub drones: Vec, /// Per class: schedulable slots on alive drones next to unmet demand. - pub capacity: BTreeMap, + capacity: BTreeMap, /// Names of drones whose heartbeat is older than `drone_stale_secs`. - pub stale_drones: Vec, + stale_drones: Vec, /// Demand the allocator could not place. Counts jobs that non-borg runners /// may serve; see `unmet_placeable`. pub unmet: Vec, @@ -1362,8 +1362,8 @@ async fn job_host( #[derive(Debug, Serialize, Deserialize)] pub struct RunHosts { - pub run_id: u64, - pub jobs: Vec, + run_id: u64, + jobs: Vec, } /// `GET /runs/{run_id}`: every job of a workflow run that ran on a borg diff --git a/src/bid.rs b/src/bid.rs index 78f683c..c362166 100644 --- a/src/bid.rs +++ b/src/bid.rs @@ -100,11 +100,11 @@ pub struct ScoreWeights { pub profile_reliability: i64, /// Applied to the drone's current CPU pressure (permille; PSI, or the /// load estimate without PSI): a loaded host scores lower. - pub host_pressure: i64, + pub(crate) host_pressure: i64, /// Applied to the share of the drone's vCPU budget still free (permille, /// divided by 1000) for medium and heavy work: an idle host scores higher /// than one already running builds. - pub idle_share: i64, + pub(crate) idle_share: i64, } impl Default for ScoreWeights { diff --git a/src/config.rs b/src/config.rs index dae08b5..afd16ed 100644 --- a/src/config.rs +++ b/src/config.rs @@ -90,12 +90,12 @@ pub struct JanitorSettings { /// pull-request runs are cancelled once the pull request is closed or /// merged. Leave out workflows that run on the `closed` event. #[serde(default)] - pub closed_pr_workflows: Vec, + pub(crate) closed_pr_workflows: Vec, /// Branches whose queued `push` runs are cancelled when a newer run of the /// same workflow on the branch is queued or running. Running push runs /// are never cancelled. #[serde(default)] - pub superseded_push_branches: Vec, + pub(crate) superseded_push_branches: Vec, } fn yes() -> bool { @@ -116,11 +116,11 @@ fn default_janitor_max() -> u32 { #[serde(deny_unknown_fields, default)] pub struct DroneExpirySettings { /// Remove a drone silent this long; 0 never removes. - pub after_secs: i64, + after_secs: i64, /// Name prefixes of short-lived cloud drones (for example `gcp-burst-`). - pub burst_prefixes: Vec, + burst_prefixes: Vec, /// Remove a drone matching `burst_prefixes` silent this long; 0 never. - pub burst_after_secs: i64, + burst_after_secs: i64, } impl Default for DroneExpirySettings { @@ -135,7 +135,7 @@ impl Default for DroneExpirySettings { impl DroneExpirySettings { /// The silence after which `name` is removed, if ever. - pub fn expiry_for(&self, name: &str) -> Option { + pub(crate) fn expiry_for(&self, name: &str) -> Option { let secs = if self .burst_prefixes .iter() @@ -155,10 +155,10 @@ impl DroneExpirySettings { #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields, default)] pub struct CapacityAlarmSettings { - pub enabled: bool, - pub classes: Vec, - pub max_slots: u32, - pub after_secs: i64, + pub(crate) enabled: bool, + pub(crate) classes: Vec, + pub(crate) max_slots: u32, + pub(crate) after_secs: i64, } impl Default for CapacityAlarmSettings { @@ -207,7 +207,7 @@ pub struct AdmissionSettings { /// and `psi_hard_pm`. Unset, the drone goes from pull-request blocked /// straight to closed. #[serde(default)] - pub psi_light_only_pm: Option, + psi_light_only_pm: Option, /// From this pressure a drone takes no new grants at all. pub psi_hard_pm: u16, /// A drone's reported pressure older than this counts as unknown. @@ -267,10 +267,10 @@ pub struct QueenConfig { pub drone_stale_secs: i64, /// When silent drones are removed; see [`DroneExpirySettings`]. #[serde(default)] - pub drone_expiry: DroneExpirySettings, + pub(crate) drone_expiry: DroneExpirySettings, /// The fleet-capacity alarm; see [`CapacityAlarmSettings`]. #[serde(default)] - pub capacity_alarm: CapacityAlarmSettings, + pub(crate) capacity_alarm: CapacityAlarmSettings, /// Demand older than this triggers the starvation guard. `false` (or JSON /// `null`) disables it; TOML has no null. #[serde(default = "default_starvation", deserialize_with = "secs_or_off")] diff --git a/src/drone/caches/gc.rs b/src/drone/caches/gc.rs index c4d066a..65d256e 100644 --- a/src/drone/caches/gc.rs +++ b/src/drone/caches/gc.rs @@ -42,15 +42,15 @@ pub fn measure(path: &Path) -> (u64, SystemTime) { /// What [`measure_tree`] found. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct Measured { - pub size: u64, - pub newest: SystemTime, + pub(super) size: u64, + pub(super) newest: SystemTime, /// Directories that could not be listed (for example owned by the job /// user and closed to the drone's user): `size` is then a lower bound. - pub unreadable_dirs: u64, + pub(super) unreadable_dirs: u64, } /// [`measure`], also counting directories that could not be listed. -pub fn measure_tree(path: &Path) -> Measured { +pub(super) fn measure_tree(path: &Path) -> Measured { let mut size = 0; let mut newest = SystemTime::UNIX_EPOCH; let mut unreadable_dirs = 0; diff --git a/src/envcontract.rs b/src/envcontract.rs index 06442f5..da9a8ea 100644 --- a/src/envcontract.rs +++ b/src/envcontract.rs @@ -203,11 +203,11 @@ pub struct EnvFacts { /// `/proc/loadavg` 1-minute average per CPU, in permille (1000 is one /// runnable task per CPU). Admission falls back to it without PSI. #[serde(default)] - pub load_per_cpu_pm: Option, + pub(crate) load_per_cpu_pm: Option, /// Share of the host's CPUs busy since the previous heartbeat, from the /// root cgroup's `cpu.stat`, in permille. #[serde(default)] - pub cpu_busy_pm: Option, + pub(crate) cpu_busy_pm: Option, /// For each tool in [`PROBE_TOOLS`], the first line of its `--version` /// output, or `"missing"`; `None` when the drone did not measure it. #[serde(default)] diff --git a/src/hostdiag.rs b/src/hostdiag.rs index 56b9a64..40a34bf 100644 --- a/src/hostdiag.rs +++ b/src/hostdiag.rs @@ -22,22 +22,22 @@ use crate::probe::{self, Host}; use crate::store::{LeaseResult, LeaseRow, Store, StoreError}; /// Rows are kept this long after the lease was granted. -pub const KEEP_SECS: i64 = 7 * 86_400; +pub(crate) const KEEP_SECS: i64 = 7 * 86_400; /// At most this many rows are kept; the oldest go first. -pub const MAX_ROWS: i64 = 50_000; +pub(crate) const MAX_ROWS: i64 = 50_000; /// CPU pressure (`some avg10`, percent) at job start that marks a slow host /// when the queen has no `admission` block. -pub const SLOW_CPU_PSI_PCT: f64 = 30.0; +pub(crate) const SLOW_CPU_PSI_PCT: f64 = 30.0; /// Memory pressure (`some avg10`, percent) at job start that marks a slow host. -pub const SLOW_MEM_PSI_PCT: f64 = 10.0; +const SLOW_MEM_PSI_PCT: f64 = 10.0; /// Load average per CPU that marks a slow host. Only used when the host has no /// CPU pressure reading, because an LXC's load average is computed by lxcfs and /// over-reports. -pub const SLOW_LOAD_PER_CPU: f64 = 1.5; +const SLOW_LOAD_PER_CPU: f64 = 1.5; /// Share of the host's CPUs busy that marks a slow host. -pub const SLOW_BUSY_SHARE: f64 = 0.9; +const SLOW_BUSY_SHARE: f64 = 0.9; /// Free disk on the work or cache volume that marks a slow host. -pub const SLOW_FREE_DISK_MIB: u64 = 10 * 1024; +const SLOW_FREE_DISK_MIB: u64 = 10 * 1024; const MAX_NAME_CHARS: usize = 120; @@ -45,19 +45,19 @@ const MAX_NAME_CHARS: usize = 120; #[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)] pub struct Psi { #[serde(default, skip_serializing_if = "Option::is_none")] - pub some_avg10: Option, + some_avg10: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub some_avg60: Option, + some_avg60: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub full_avg10: Option, + full_avg10: Option, #[serde(default, skip_serializing_if = "Option::is_none")] - pub full_avg60: Option, + full_avg60: Option, } impl Psi { /// Parse `some avg10=1.00 avg60=...` / `full ...` lines. `None` when no /// field could be read. - pub fn parse(text: &str) -> Option { + fn parse(text: &str) -> Option { let field = |kind: &str, key: &str| -> Option { text.lines() .find(|l| l.starts_with(kind))? @@ -99,44 +99,44 @@ pub struct HostSample { /// `/proc/loadavg` 1, 5 and 15 minute averages. Inside an LXC this is /// lxcfs's estimate and over-reports; prefer the pressure and busy fields. #[serde(default)] - pub load: Option<[f64; 3]>, + load: Option<[f64; 3]>, /// `/proc/pressure/cpu`. #[serde(default)] - pub cpu_psi: Option, + cpu_psi: Option, /// `/proc/pressure/memory`. #[serde(default)] - pub memory_psi: Option, + memory_psi: Option, /// `/proc/pressure/io`. #[serde(default)] - pub io_psi: Option, + io_psi: Option, /// `/sys/fs/cgroup/cpu.pressure`: the cgroup-namespace root, which is the /// whole container inside an LXC. #[serde(default)] - pub cgroup_cpu_psi: Option, + cgroup_cpu_psi: Option, /// `/sys/fs/cgroup/memory.pressure`. #[serde(default)] - pub cgroup_memory_psi: Option, + cgroup_memory_psi: Option, /// `usage_usec` of `/sys/fs/cgroup/cpu.stat` (whole host or container). #[serde(default)] - pub host_cpu_usage_usec: Option, + host_cpu_usage_usec: Option, /// Cores busy in that cgroup, averaged since the drone's previous sample /// (at most one heartbeat ago). #[serde(default)] - pub host_cpu_busy_cores: Option, + host_cpu_busy_cores: Option, /// `usage_usec` of the drone's own cgroup; jobs run inside it. #[serde(default)] - pub drone_cpu_usage_usec: Option, + drone_cpu_usage_usec: Option, /// `throttled_usec` of the drone's own cgroup. #[serde(default)] - pub drone_cpu_throttled_usec: Option, + drone_cpu_throttled_usec: Option, #[serde(default)] - pub mem_available_mib: Option, + mem_available_mib: Option, /// Free space on the filesystem holding the drone's `work_root`. #[serde(default)] - pub work_free_mib: Option, + work_free_mib: Option, /// Free space on the filesystem holding the drone's `cache_root`. #[serde(default)] - pub cache_free_mib: Option, + cache_free_mib: Option, /// Leases the drone held at this moment, this one included. #[serde(default)] pub active_leases: u32, @@ -144,7 +144,7 @@ pub struct HostSample { impl HostSample { /// Clamp what a drone sent: finite non-negative numbers, no future time. - pub fn sanitized(mut self, now: i64) -> Self { + fn sanitized(mut self, now: i64) -> Self { let num = |v: f64| v.is_finite().then(|| v.max(0.0)); self.at = self.at.min(now); self.load = self @@ -164,7 +164,7 @@ impl HostSample { } /// The higher CPU `some avg10` of the host file and the cgroup file. - pub fn cpu_some_avg10(&self) -> Option { + fn cpu_some_avg10(&self) -> Option { [self.cpu_psi, self.cgroup_cpu_psi] .into_iter() .flatten() @@ -185,9 +185,9 @@ impl HostSample { #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] pub struct LeaseHostReport { #[serde(default)] - pub start: Option, + pub(crate) start: Option, #[serde(default)] - pub end: Option, + pub(crate) end: Option, } /// Turns successive `cpu.stat` readings into a busy-cores rate. @@ -208,7 +208,7 @@ impl CpuMeter { /// Record a reading and return the latest rate. A reading that arrives /// within two seconds of the stored one returns the stored rate, so a lease /// that starts right after a heartbeat still gets a measured value. - pub fn observe(&self, at_ms: i64, usage_usec: u64) -> Option { + fn observe(&self, at_ms: i64, usage_usec: u64) -> Option { let mut slot = self.0.lock().ok()?; let next = match *slot { Some(prev) if at_ms - prev.at_ms < METER_MIN_GAP_MS && at_ms >= prev.at_ms => { @@ -266,14 +266,14 @@ fn loadavg(host: &dyn Host) -> Option<[f64; 3]> { /// Feed the meter from the root cgroup and return the busy-cores rate. Called /// on every heartbeat so a lease start has a recent rate. -pub fn observe_cpu(host: &dyn Host, meter: &CpuMeter, at_ms: i64) -> Option { +pub(crate) fn observe_cpu(host: &dyn Host, meter: &CpuMeter, at_ms: i64) -> Option { let (usage, _) = cpu_stat(host, Path::new("/sys/fs/cgroup")); meter.observe(at_ms, usage?) } /// Sample the host for a lease start or end. Every field degrades to `None` /// on its own; this never fails. -pub fn sample( +pub(crate) fn sample( host: &dyn Host, work_root: &Path, cache_root: &Path, @@ -311,15 +311,15 @@ pub fn sample( /// Averages over the job, from the start and end samples. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct During { - pub secs: i64, + secs: i64, /// Cores busy on the whole host (or container). #[serde(skip_serializing_if = "Option::is_none")] - pub host_cpu_busy_cores: Option, + host_cpu_busy_cores: Option, /// Cores used by the drone's cgroup: this job and any other lease. #[serde(skip_serializing_if = "Option::is_none")] - pub drone_cpu_busy_cores: Option, + drone_cpu_busy_cores: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub drone_cpu_throttled_secs: Option, + drone_cpu_throttled_secs: Option, } fn during(start: &HostSample, end: &HostSample) -> Option { @@ -343,7 +343,7 @@ fn during(start: &HostSample, end: &HostSample) -> Option { /// Why a job's host was overloaded, or nothing. `cpu_psi_pct` is the CPU /// pressure threshold (the queen's admission soft threshold when set). -pub fn slow_host_reasons( +fn slow_host_reasons( start: Option<&HostSample>, during: Option<&During>, cpu_psi_pct: f64, @@ -404,27 +404,27 @@ pub fn slow_host_reasons( /// One lease as `GET /jobs/{id}` and `GET /runs/{id}` show it. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct JobHostView { - pub lease_id: String, + lease_id: String, pub drone: String, - pub class: Class, - pub scope: Scope, - pub repo: Option, - pub run_id: Option, - pub job_id: Option, - pub job_name: Option, + class: Class, + scope: Scope, + repo: Option, + run_id: Option, + job_id: Option, + job_name: Option, /// When the queen granted the lease. - pub leased_at: i64, + leased_at: i64, /// When the drone reported the result. - pub finished_at: Option, + finished_at: Option, /// Runner wall time the drone reported. - pub duration_secs: Option, + duration_secs: Option, pub exit_code: Option, pub start: Option, pub end: Option, - pub during: Option, + during: Option, /// The drone was overloaded when the job started or while it ran. - pub slow_host: bool, - pub slow_host_reasons: Vec, + slow_host: bool, + slow_host_reasons: Vec, } fn to_i64(v: u64) -> i64 { @@ -461,7 +461,7 @@ const SELECT: &str = "SELECT lease_id, drone, class, scope, repo, run_id, job_id impl Store { /// Record a drone's result and host samples for its lease. - pub async fn record_lease_host( + pub(crate) async fn record_lease_host( &self, lease: &LeaseRow, result: &LeaseResult, @@ -498,7 +498,7 @@ impl Store { } /// Attach the GitHub job a lease's runner picked up. - pub async fn bind_lease_host_job( + pub(crate) async fn bind_lease_host_job( &self, lease: &LeaseRow, repo: &str, @@ -533,7 +533,7 @@ impl Store { } /// The lease that ran GitHub job `job_id` (the newest, if several did). - pub async fn job_host( + pub(crate) async fn job_host( &self, job_id: u64, cpu_psi_pct: f64, @@ -563,7 +563,7 @@ impl Store { } /// Every lease that ran a job of workflow run `run_id`, oldest first. - pub async fn run_hosts( + pub(crate) async fn run_hosts( &self, run_id: u64, cpu_psi_pct: f64, @@ -579,7 +579,11 @@ impl Store { } /// Drop rows granted before `before`, then all but the newest `max_rows`. - pub async fn prune_lease_hosts(&self, before: i64, max_rows: i64) -> Result { + pub(crate) async fn prune_lease_hosts( + &self, + before: i64, + max_rows: i64, + ) -> Result { // sqlx-runtime-query: standalone-buildable SQLite store; the query! macros need a build-time database let old = sqlx::query("DELETE FROM lease_hosts WHERE leased_at < ?") .bind(before) diff --git a/src/janitor.rs b/src/janitor.rs index 691908c..675d6e5 100644 --- a/src/janitor.rs +++ b/src/janitor.rs @@ -108,9 +108,9 @@ impl PassReport { pub struct JanitorView { pub passes: u64, /// Unix time the queen started; `total` counts from here. - pub since: i64, + since: i64, /// Unix time of the last pass; `None` before the first. - pub last_pass_at: Option, + last_pass_at: Option, /// What the last pass alone found. Runs still queued are found again by /// every pass, so this is the current state. pub last: PassReport, diff --git a/src/queen.rs b/src/queen.rs index 15ff14d..317fb0f 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -96,12 +96,12 @@ pub struct Snapshot { /// `unmet` minus what idle non-borg runners with the label can take. pub unmet_placeable: Vec, /// Per class, slots the fresh accepting drones could hold if empty. - pub schedulable_slots: BTreeMap, + pub(crate) schedulable_slots: BTreeMap, /// Per class, since when main demand has been unmet with schedulable /// slots at or under `capacity_alarm.max_slots`. - pub starved_since: BTreeMap, + pub(crate) starved_since: BTreeMap, /// Classes whose `capacity_starved` event is raised and not yet restored. - pub starved: BTreeSet, + pub(crate) starved: BTreeSet, } /// What the fleet watch remembers between passes. diff --git a/src/store.rs b/src/store.rs index 6e6e150..2e47b42 100644 --- a/src/store.rs +++ b/src/store.rs @@ -808,7 +808,7 @@ impl Store { /// Append a fleet or drone event (stale, expired, capacity) to the /// health audit trail that `/status` shows as `health_events`. - pub async fn record_health_event( + pub(crate) async fn record_health_event( &self, at: i64, drone: &str, diff --git a/tests/hostdiag_wire.rs b/tests/hostdiag_wire.rs new file mode 100644 index 0000000..db2ce0e --- /dev/null +++ b/tests/hostdiag_wire.rs @@ -0,0 +1,33 @@ +//! Host telemetry stays readable through the public wire types. + +#[test] +fn private_host_diagnostic_fields_preserve_the_wire_contract() { + let wire = serde_json::json!({ + "start": { + "at": 100, + "cpus": 16, + "load": [1.0, 2.0, 3.0], + "cpu_psi": { + "some_avg10": 1.0, + "some_avg60": 2.0, + "full_avg10": 3.0, + "full_avg60": 4.0 + }, + "memory_psi": null, + "io_psi": null, + "cgroup_cpu_psi": null, + "cgroup_memory_psi": null, + "host_cpu_usage_usec": 1000, + "host_cpu_busy_cores": 2.0, + "drone_cpu_usage_usec": 500, + "drone_cpu_throttled_usec": 1, + "mem_available_mib": 4096, + "work_free_mib": 20000, + "cache_free_mib": 30000, + "active_leases": 2 + }, + "end": null + }); + let report: borg::hostdiag::LeaseHostReport = serde_json::from_value(wire.clone()).unwrap(); + assert_eq!(serde_json::to_value(report).unwrap(), wire); +} diff --git a/tests/suite.rs b/tests/suite.rs index e6bfa58..8417fb7 100644 --- a/tests/suite.rs +++ b/tests/suite.rs @@ -14,6 +14,7 @@ mod drone; mod explain; mod explain_cli; mod health; +mod hostdiag_wire; mod janitor; mod job_user; mod placement_batch; From 225f431a5c68aef04c1a6e72f194b3c5a0c84bbd Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sat, 17 Oct 2026 07:01:19 +0000 Subject: [PATCH 26/27] chore: project borg from Mono 9604c2352669 --- .repository-projection.json | 6 +++--- src/drone/caches/gc.rs | 2 +- src/queen.rs | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index e82230f..9535a9e 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,11 +3,11 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "6077c90a491c4384f1915a0c2f7bd244bf30982e", + "sourceSha": "9604c2352669ef0ca5a594ac6b92271a961a6ae3", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "c6d87f59583ab6b292262558a7086011a7e527f1", + "priorProjectedBase": "8369d49ae136c5124c1322a4ebc0f2b8b233b597", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", - "contentDigest": "7036b28c11658bd46145527ba700dd786cffd76a53418a3c5b3ec9f3ce8c4a06", + "contentDigest": "3789a06e91280019b7f1f29b4ed4df6224db3eed83918da5d55b912696bef2ed", "publicationEligible": true } diff --git a/src/drone/caches/gc.rs b/src/drone/caches/gc.rs index 65d256e..9cb919b 100644 --- a/src/drone/caches/gc.rs +++ b/src/drone/caches/gc.rs @@ -43,7 +43,7 @@ pub fn measure(path: &Path) -> (u64, SystemTime) { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct Measured { pub(super) size: u64, - pub(super) newest: SystemTime, + newest: SystemTime, /// Directories that could not be listed (for example owned by the job /// user and closed to the drone's user): `size` is then a lower bound. pub(super) unreadable_dirs: u64, diff --git a/src/queen.rs b/src/queen.rs index 317fb0f..c91af32 100644 --- a/src/queen.rs +++ b/src/queen.rs @@ -96,7 +96,7 @@ pub struct Snapshot { /// `unmet` minus what idle non-borg runners with the label can take. pub unmet_placeable: Vec, /// Per class, slots the fresh accepting drones could hold if empty. - pub(crate) schedulable_slots: BTreeMap, + schedulable_slots: BTreeMap, /// Per class, since when main demand has been unmet with schedulable /// slots at or under `capacity_alarm.max_slots`. pub(crate) starved_since: BTreeMap, From 1d12608685c7d3e2e43f78609ad6f98d9347629d Mon Sep 17 00:00:00 2001 From: dx-corp projector Date: Sun, 18 Oct 2026 01:13:58 +0000 Subject: [PATCH 27/27] chore: project borg from Mono 2ba0ca8a9a83 --- .repository-projection.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.repository-projection.json b/.repository-projection.json index 9535a9e..e8bbf6f 100644 --- a/.repository-projection.json +++ b/.repository-projection.json @@ -3,9 +3,9 @@ "projection": "borg", "projectionSchemaVersion": 1, "sourceRepository": "dx-corp/mono", - "sourceSha": "9604c2352669ef0ca5a594ac6b92271a961a6ae3", + "sourceSha": "2ba0ca8a9a83f56b508ad59c7658a48760ce56f2", "destinationRepository": "dx-corp/borg", - "priorProjectedBase": "8369d49ae136c5124c1322a4ebc0f2b8b233b597", + "priorProjectedBase": "06501edada175683601d3a0b36812390b3d786c5", "definitionDigest": "78ba8d1b5794c261ff0afba6d99abc10b9297ea0dd94be4397fb1cbf63cf0c6e", "toolDigest": "61f64178999b764105e6e5c5bc7c5971c164909e", "contentDigest": "3789a06e91280019b7f1f29b4ed4df6224db3eed83918da5d55b912696bef2ed",