diff --git a/Cargo.lock b/Cargo.lock index f5de553e..0509a77e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1756,8 +1756,10 @@ name = "hive-agent-sock" version = "0.1.0" dependencies = [ "hive-sh4re", + "libc", "serde", "serde_json", + "tempfile", "tracing", ] diff --git a/docs/scheduler/coordinator.md b/docs/scheduler/coordinator.md index 51654ebf..c40fc929 100644 --- a/docs/scheduler/coordinator.md +++ b/docs/scheduler/coordinator.md @@ -479,6 +479,20 @@ changed cap reaches them at the next container start without a rebuild. `hyperhive-memory-cap-check.service` logs a warning naming the path in the container's journal. The container's `MemoryMax=` still applies. +When the container reaches its cap, the kernel kills in this order: + +1. a subagent run (`oom_score_adj` +1000); +2. a bash task (0); +3. the harness and the agent's own `claude` (-500). + +Each subagent run and each bash task has its own cgroup with +`memory.oom.group=1`, so the kernel kills all of a run or task and nothing +outside it. Inside the container these are +`/sys/fs/cgroup/system.slice/hive-subagent-daemon.service/runs/-/` +and `/sys/fs/cgroup/system.slice/hive-bash-daemon.service/tasks//`; the +daemon itself sits in `daemon/` beside them. A daemon that can't set these +up logs one warning at start and runs everything in its own cgroup. +
Implementation: swarm caps swarm-controller keeps the caps in the `agent-limits` KV bucket, one key per diff --git a/hive-agent-sock/Cargo.toml b/hive-agent-sock/Cargo.toml index b0585ffd..240baed8 100644 --- a/hive-agent-sock/Cargo.toml +++ b/hive-agent-sock/Cargo.toml @@ -8,7 +8,12 @@ readme = "README.md" workspace = true [dependencies] +# `run_cgroup`: `write(2)` between fork and exec. +libc.workspace = true serde.workspace = true serde_json.workspace = true tracing.workspace = true hive-sh4re.workspace = true + +[dev-dependencies] +tempfile = "3" diff --git a/hive-agent-sock/README.md b/hive-agent-sock/README.md index 4178050d..e2980965 100644 --- a/hive-agent-sock/README.md +++ b/hive-agent-sock/README.md @@ -2,7 +2,7 @@ What's common across the in-container "agent plugins" — the matrix / bash / subagent MCP daemons and `forge_notify` today, plus any user-configured MCP -server declared in an agent's `agent.nix`. Two things live here: +server declared in an agent's `agent.nix`. Three things live here: ## Wire types for the in-agent socket @@ -38,3 +38,12 @@ entry into the shape claude expects in a `availableToSubagents`) — both already depend on this crate for the socket types above, so a new plugin-facing shared type belongs here rather than in a crate of its own. + +## `run_cgroup`: one cgroup per run + +`hive-bash-mcp` puts each bash task in `/tasks//` and +`hive-subagent-mcp` each subagent run in `/runs/-/`, both with +`memory.oom.group=1`, so the OOM killer takes one whole run and leaves the +daemon alone. Both daemons are delegated their cgroup with `Delegate=memory +pids` and `DelegateSubgroup=daemon` (`nix/agent-modules/mcp.nix`); outside such +a unit the module logs one warning and runs stay in the daemon's cgroup. diff --git a/hive-agent-sock/src/lib.rs b/hive-agent-sock/src/lib.rs index 0edec714..df42b873 100644 --- a/hive-agent-sock/src/lib.rs +++ b/hive-agent-sock/src/lib.rs @@ -20,6 +20,9 @@ //! `hive-agent` (the main session's `--mcp-config`) and `hive-subagent-mcp` //! (the subagent-eligible subset) — not a socket protocol at all, but this //! crate is the natural home since both consumers already depend on it. +//! +//! **`run_cgroup`**: the per-run cgroups `hive-bash-mcp` and +//! `hive-subagent-mcp` put their tasks and subagent runs in. use serde::{Deserialize, Serialize}; @@ -27,6 +30,7 @@ use hive_sh4re::inbox::LooseEnd; pub mod extra_mcp; pub mod paths; +pub mod run_cgroup; /// In-container path of the harness-served in-agent socket. The harness /// binds it on boot; the in-container producers dial it for todo ops. diff --git a/hive-agent-sock/src/run_cgroup.rs b/hive-agent-sock/src/run_cgroup.rs new file mode 100644 index 00000000..eca80225 --- /dev/null +++ b/hive-agent-sock/src/run_cgroup.rs @@ -0,0 +1,356 @@ +//! One cgroup per run, under a daemon unit that systemd delegated with +//! `Delegate=memory pids` and `DelegateSubgroup=daemon` +//! (`nix/agent-modules/mcp.nix`). The daemon process lives in the unit's +//! `daemon/` leaf; each run gets `///` with +//! `memory.oom.group=1`, so the OOM killer takes a whole run, never part of +//! one and never the daemon. +//! +//! systemd enables no controller below a service's own cgroup, so +//! [`RunCgroups::init`] writes `+memory +pids` into the unit's and the +//! `/` node's `cgroup.subtree_control` itself. Without that the run +//! cgroups would have no `memory.oom.group` file. +//! +//! Two ways into a run's cgroup, both before the run's program executes, so +//! everything it later forks is born inside: +//! - [`enter_in_pre_exec`], from a `Command::pre_exec` hook, for a spawn +//! this crate's caller builds itself; +//! - [`RunCgroup::entry`], for a spawn built by a library with no such hook: +//! it runs this daemon's own binary, whose `main` calls [`exec_entry`] +//! first thing, which enters the cgroup and then execs the real program. + +use std::ffi::OsStr; +use std::fs::File; +use std::io; +use std::os::fd::RawFd; +use std::os::unix::process::CommandExt; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +/// The leaf systemd puts the daemon process in. Must equal the unit's +/// `DelegateSubgroup=` in `nix/agent-modules/mcp.nix`. +const DAEMON_SUBGROUP: &str = "daemon"; + +const CGROUP_ROOT: &str = "/sys/fs/cgroup"; +const CONTROLLERS: &str = "+memory +pids"; + +/// The program [`exec_entry`] execs. Its presence is what makes a process +/// an entry wrapper rather than the daemon. +const ENTRY_PROGRAM_ENV: &str = "HIVE_RUN_ENTRY_PROGRAM"; +const ENTRY_CGROUP_ENV: &str = "HIVE_RUN_ENTRY_CGROUP"; +const ENTRY_OOM_SCORE_ADJ_ENV: &str = "HIVE_RUN_ENTRY_OOM_SCORE_ADJ"; + +/// How long a run's cgroup may stay populated after the run ended before +/// removal is left to the next daemon start. +const RELEASE_GRACE: Duration = Duration::from_secs(30); +const RELEASE_POLL: Duration = Duration::from_millis(250); + +/// The `//` node run cgroups are created under. +#[derive(Debug)] +pub struct RunCgroups { + parent: PathBuf, +} + +impl RunCgroups { + /// Set up `//` for this process's unit. `None`, with one + /// warning, when this process is not in a delegated `daemon/` leaf on + /// cgroup v2 (a test, a dev shell, a unit without `Delegate=`): runs then + /// stay in the daemon's own cgroup. + #[must_use] + pub fn init(kind: &str) -> Option { + let result = std::fs::read_to_string("/proc/self/cgroup") + .and_then(|own| Self::init_at(Path::new(CGROUP_ROOT), &own, kind)); + match result { + Ok(cgroups) => { + tracing::info!(parent = %cgroups.parent.display(), "run cgroups ready"); + Some(cgroups) + } + Err(e) => { + tracing::warn!( + error = %e, + "run cgroups unavailable; runs share the daemon's cgroup" + ); + None + } + } + } + + /// [`RunCgroups::init`] against cgroup2 mounted at `root`, with + /// `proc_self_cgroup` the content of `/proc/self/cgroup`. + fn init_at(root: &Path, proc_self_cgroup: &str, kind: &str) -> io::Result { + let own = proc_self_cgroup + .lines() + .find_map(|line| line.strip_prefix("0::")) + .ok_or_else(|| io::Error::other("no cgroup v2 entry in /proc/self/cgroup"))?; + let own = root.join(own.trim().trim_start_matches('/')); + if own.file_name().and_then(|n| n.to_str()) != Some(DAEMON_SUBGROUP) { + return Err(io::Error::other(format!( + "{} is not a delegated `{DAEMON_SUBGROUP}` leaf", + own.display() + ))); + } + let unit = own + .parent() + .ok_or_else(|| io::Error::other("daemon leaf has no parent"))?; + std::fs::write(unit.join("cgroup.subtree_control"), CONTROLLERS)?; + let parent = unit.join(kind); + match std::fs::create_dir(&parent) { + Ok(()) => {} + Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {} + Err(e) => return Err(e), + } + std::fs::write(parent.join("cgroup.subtree_control"), CONTROLLERS)?; + // Runs that outlived the previous daemon's grace leave empty dirs; + // a populated one belongs to a process still alive and stays. + for entry in std::fs::read_dir(&parent)?.flatten() { + if entry.file_type().is_ok_and(|t| t.is_dir()) { + let _ = std::fs::remove_dir(entry.path()); + } + } + Ok(Self { parent }) + } + + /// Create the run cgroup `name`, with `memory.oom.group` set. + /// + /// # Errors + /// + /// `AlreadyExists` when a populated cgroup of that name is still there, + /// or the underlying failure creating it or writing `memory.oom.group`. + pub fn create(&self, name: &str) -> io::Result { + let dir = self.parent.join(name); + match std::fs::create_dir(&dir) { + Err(e) if e.kind() == io::ErrorKind::AlreadyExists => { + std::fs::remove_dir(&dir).map_err(|_| e)?; + std::fs::create_dir(&dir)?; + } + other => other?, + } + let cgroup = RunCgroup { dir }; + std::fs::write(cgroup.dir.join("memory.oom.group"), "1")?; + Ok(cgroup) + } +} + +/// One run's cgroup. Dropping it removes the directory once the run's last +/// process has left, waiting up to 30 seconds on a background thread. +#[derive(Debug)] +pub struct RunCgroup { + dir: PathBuf, +} + +impl RunCgroup { + /// The run's `cgroup.procs`, open for [`enter_in_pre_exec`]. Keep it + /// open until the spawn has returned. + /// + /// # Errors + /// + /// The failure opening the file. + pub fn procs(&self) -> io::Result { + File::options() + .write(true) + .open(self.dir.join("cgroup.procs")) + } + + /// The program and extra environment that run `program` inside this + /// cgroup with `oom_score_adj`: this daemon's own binary, whose `main` + /// starts with [`exec_entry`]. + #[must_use] + pub fn entry(&self, program: &Path, oom_score_adj: i32) -> (PathBuf, Vec<(String, String)>) { + ( + PathBuf::from("/proc/self/exe"), + vec![ + ( + ENTRY_PROGRAM_ENV.to_owned(), + program.to_string_lossy().into_owned(), + ), + ( + ENTRY_CGROUP_ENV.to_owned(), + self.dir.to_string_lossy().into_owned(), + ), + ( + ENTRY_OOM_SCORE_ADJ_ENV.to_owned(), + oom_score_adj.to_string(), + ), + ], + ) + } +} + +impl Drop for RunCgroup { + fn drop(&mut self) { + match std::fs::remove_dir(&self.dir) { + Err(e) if e.raw_os_error() == Some(libc::EBUSY) => { + let dir = std::mem::take(&mut self.dir); + std::thread::spawn(move || release_when_empty(&dir)); + } + Err(e) if e.kind() != io::ErrorKind::NotFound => { + tracing::warn!(dir = %self.dir.display(), error = %e, "removing run cgroup failed"); + } + _ => {} + } + } +} + +/// Retry removing a still-populated run cgroup until it empties or +/// [`RELEASE_GRACE`] runs out; past that the next daemon start removes it. +fn release_when_empty(dir: &Path) { + let deadline = std::time::Instant::now() + RELEASE_GRACE; + while std::time::Instant::now() < deadline { + std::thread::sleep(RELEASE_POLL); + match std::fs::remove_dir(dir) { + Err(e) if e.raw_os_error() == Some(libc::EBUSY) => {} + _ => return, + } + } + tracing::info!(dir = %dir.display(), "run cgroup still populated after the run ended; left in place"); +} + +/// Move the calling process into the cgroup whose `cgroup.procs` is `procs`. +/// Only `write(2)`, so it is safe between fork and exec. +/// +/// # Errors +/// +/// The `write(2)` failure. +pub fn enter_in_pre_exec(procs: RawFd) -> io::Result<()> { + // SAFETY: a plain write of a static buffer to a caller-owned fd. + if unsafe { libc::write(procs, b"0".as_ptr().cast(), 1) } < 0 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +/// When this process was started by [`RunCgroup::entry`]: enter the run's +/// cgroup, set its `oom_score_adj`, and exec the run's program with this +/// process's arguments; on failure print the reason and exit 127. Otherwise +/// return, doing nothing. Call it first in `main`. +pub fn exec_entry() { + let Some(program) = std::env::var_os(ENTRY_PROGRAM_ENV) else { + return; + }; + let err = enter_and_exec(&program); + eprintln!("run cgroup entry failed: {err}"); + std::process::exit(127); +} + +fn enter_and_exec(program: &OsStr) -> io::Error { + if let Some(cgroup) = std::env::var_os(ENTRY_CGROUP_ENV) + && let Err(e) = std::fs::write(Path::new(&cgroup).join("cgroup.procs"), "0") + { + return e; + } + if let Some(adj) = std::env::var_os(ENTRY_OOM_SCORE_ADJ_ENV) + && let Err(e) = std::fs::write("/proc/self/oom_score_adj", adj.as_encoded_bytes()) + { + return e; + } + std::process::Command::new(program) + .args(std::env::args_os().skip(1)) + .env_remove(ENTRY_PROGRAM_ENV) + .env_remove(ENTRY_CGROUP_ENV) + .env_remove(ENTRY_OOM_SCORE_ADJ_ENV) + .exec() +} + +#[cfg(test)] +mod tests { + use super::*; + + const KIND: &str = "runs"; + + fn unit_tree() -> (tempfile::TempDir, PathBuf) { + let root = tempfile::tempdir().expect("tempdir"); + let unit = root.path().join("system.slice/x.service"); + std::fs::create_dir_all(unit.join(DAEMON_SUBGROUP)).expect("daemon leaf"); + (root, unit) + } + + const OWN: &str = "0::/system.slice/x.service/daemon\n"; + + #[test] + fn init_enables_controllers_on_the_unit_and_the_kind_node() { + let (root, unit) = unit_tree(); + let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init"); + assert_eq!(cgroups.parent, unit.join(KIND)); + for node in [&unit, &unit.join(KIND)] { + let control = + std::fs::read_to_string(node.join("cgroup.subtree_control")).expect("written"); + assert_eq!(control, CONTROLLERS); + } + } + + #[test] + fn init_refuses_a_process_outside_a_daemon_leaf() { + let (root, unit) = unit_tree(); + let err = RunCgroups::init_at(root.path(), "0::/system.slice/x.service\n", KIND) + .expect_err("not a daemon leaf"); + assert!(err.to_string().contains(DAEMON_SUBGROUP), "{err}"); + assert!(!unit.join("cgroup.subtree_control").exists()); + } + + #[test] + fn init_refuses_cgroup_v1() { + let (root, _unit) = unit_tree(); + assert!(RunCgroups::init_at(root.path(), "1:memory:/x\n", KIND).is_err()); + } + + #[test] + fn init_removes_empty_leftover_runs_and_keeps_populated_ones() { + let (root, unit) = unit_tree(); + let empty = unit.join(KIND).join("old-0"); + let populated = unit.join(KIND).join("old-1"); + std::fs::create_dir_all(&empty).expect("empty"); + std::fs::create_dir_all(&populated).expect("populated"); + std::fs::write(populated.join("cgroup.procs"), "42").expect("procs"); + RunCgroups::init_at(root.path(), OWN, KIND).expect("init"); + assert!(!empty.exists()); + assert!(populated.exists()); + } + + #[test] + fn create_sets_oom_group_and_drop_removes_the_dir() { + let (root, unit) = unit_tree(); + let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init"); + let run = cgroups.create("alpha-0").expect("create"); + let dir = unit.join(KIND).join("alpha-0"); + assert_eq!(run.dir, dir); + assert_eq!( + std::fs::read_to_string(dir.join("memory.oom.group")).expect("written"), + "1" + ); + // cgroupfs drops a cgroup's interface files with the directory; a + // plain tempdir needs them gone first. + std::fs::remove_file(dir.join("memory.oom.group")).expect("rm"); + drop(run); + assert!(!dir.exists()); + } + + #[test] + fn create_replaces_an_empty_leftover_of_the_same_name() { + let (root, unit) = unit_tree(); + let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init"); + std::fs::create_dir(unit.join(KIND).join("t1")).expect("leftover"); + cgroups.create("t1").expect("create over an empty leftover"); + } + + #[test] + fn entry_names_this_binary_and_the_run() { + let (root, _unit) = unit_tree(); + let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init"); + let run = cgroups.create("alpha-0").expect("create"); + let (program, env) = run.entry(Path::new("claude"), 1000); + assert_eq!(program, Path::new("/proc/self/exe")); + let get = |k: &str| env.iter().find(|(n, _)| n == k).map(|(_, v)| v.as_str()); + assert_eq!(get(ENTRY_PROGRAM_ENV), Some("claude")); + assert_eq!(get(ENTRY_CGROUP_ENV), run.dir.to_str()); + assert_eq!(get(ENTRY_OOM_SCORE_ADJ_ENV), Some("1000")); + } + + #[test] + fn enter_in_pre_exec_writes_zero_to_the_fd() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("cgroup.procs"); + std::fs::write(&path, "").expect("create"); + let procs = File::options().write(true).open(&path).expect("open"); + enter_in_pre_exec(std::os::fd::AsRawFd::as_raw_fd(&procs)).expect("write"); + assert_eq!(std::fs::read_to_string(&path).expect("read"), "0"); + } +} diff --git a/hive-bash-mcp/src/main.rs b/hive-bash-mcp/src/main.rs index e928b406..8666dee3 100644 --- a/hive-bash-mcp/src/main.rs +++ b/hive-bash-mcp/src/main.rs @@ -39,9 +39,11 @@ async fn main() -> Result<()> { "hive-bash-daemon starting" ); + let cgroups = hive_agent_sock::run_cgroup::RunCgroups::init("tasks"); + // Start the background runner loop — scans for pending tasks and // spawns them, pushing todos to the harness on task transitions. - hive_bash_mcp::runner::spawn_loop(todo_socket); + hive_bash_mcp::runner::spawn_loop(todo_socket, cgroups); // Serve the MCP tools over streamable-http forever. hive_bash_mcp::mcp::serve_http(cli.http).await diff --git a/hive-bash-mcp/src/runner.rs b/hive-bash-mcp/src/runner.rs index 00879352..02f52a2b 100644 --- a/hive-bash-mcp/src/runner.rs +++ b/hive-bash-mcp/src/runner.rs @@ -34,6 +34,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use anyhow::{Result, bail}; use hive_agent_sock::Request as TodoReq; +use hive_agent_sock::run_cgroup::{self, RunCgroup, RunCgroups}; use tokio::io::AsyncWriteExt as _; use tokio::sync::Notify; @@ -431,13 +432,14 @@ pub fn kill_task(id: &str, force: bool) -> (bool, bool) { /// Spawn the background runner loop as a detached tokio task. `socket` is /// the harness's in-agent todo socket (`HIVE_AGENT_SOCKET`), where task /// transitions are pushed as todos. Call once at daemon startup. -pub fn spawn_loop(socket: PathBuf) { +pub fn spawn_loop(socket: PathBuf, cgroups: Option) { + let cgroups = Arc::new(cgroups); tokio::spawn(async move { - run_loop(socket).await; + run_loop(socket, cgroups).await; }); } -async fn run_loop(socket: PathBuf) { +async fn run_loop(socket: PathBuf, cgroups: Arc>) { if let Err(e) = std::fs::create_dir_all(paths::tasks_dir()) { tracing::warn!(error = ?e, "bash_runner: create tasks dir failed"); } @@ -450,7 +452,7 @@ async fn run_loop(socket: PathBuf) { let claimed: Arc>> = Arc::new(Mutex::new(HashSet::new())); loop { - poll_once(&socket, &claimed); + poll_once(&socket, &claimed, &cgroups); tokio::time::sleep(POLL_INTERVAL).await; } } @@ -541,7 +543,11 @@ async fn mark_interrupted(socket: &Path) { } } -fn poll_once(socket: &Path, claimed: &Arc>>) { +fn poll_once( + socket: &Path, + claimed: &Arc>>, + cgroups: &Arc>, +) { let Ok(rd) = std::fs::read_dir(paths::tasks_dir()) else { return; }; @@ -567,8 +573,9 @@ fn poll_once(socket: &Path, claimed: &Arc>>) { claimed.lock().unwrap().insert(id.clone()); let socket = socket.to_path_buf(); let claimed = claimed.clone(); + let cgroups = Arc::clone(cgroups); tokio::spawn(async move { - run_task(task, &socket).await; + run_task(task, &socket, cgroups.as_ref().as_ref()).await; claimed.lock().unwrap().remove(&id); }); } @@ -578,7 +585,7 @@ fn poll_once(socket: &Path, claimed: &Arc>>) { // Task execution // --------------------------------------------------------------------------- -async fn run_task(mut task: TaskFile, socket: &Path) { +async fn run_task(mut task: TaskFile, socket: &Path, cgroups: Option<&RunCgroups>) { let id = task.id.clone(); tracing::info!(id = %id, cmd = %task.cmd, "bash_runner: starting task"); @@ -616,6 +623,7 @@ async fn run_task(mut task: TaskFile, socket: &Path) { }, ); + let cgroup = cgroups.and_then(|cgroups| task_cgroup(cgroups, &id)); let outcome = exec_cmd( &task.cmd, &out_path, @@ -623,8 +631,10 @@ async fn run_task(mut task: TaskFile, socket: &Path) { task.timeout_secs, &cancel, &force, + cgroup.as_ref().map(|(_, procs)| procs), ) .await; + drop(cgroup); // Deregister before writing terminal state — the task is no longer killable. running().lock().unwrap().remove(&id); @@ -693,6 +703,20 @@ async fn run_task(mut task: TaskFile, socket: &Path) { .await; } +/// `id`'s cgroup and its open `cgroup.procs`, or `None` (logged) when either +/// fails, in which case the task runs in the daemon's cgroup. +fn task_cgroup(cgroups: &RunCgroups, id: &str) -> Option<(RunCgroup, std::fs::File)> { + let opened = cgroups.create(id).and_then(|cgroup| { + let procs = cgroup.procs()?; + Ok((cgroup, procs)) + }); + opened + .inspect_err(|e| { + tracing::warn!(id = %id, error = %e, "bash_runner: task cgroup unavailable"); + }) + .ok() +} + /// Run `bash -c cmd` in its own process group, streaming output to files. /// `timeout_secs = None` means no timeout — run until natural exit. /// @@ -721,6 +745,7 @@ async fn exec_cmd( timeout_secs: Option, cancel: &Notify, force: &AtomicBool, + cgroup_procs: Option<&std::fs::File>, ) -> Result { use tokio::process::Command; // `/usr/bin/env bash` rather than a bare `bash`: `/usr/bin/env` is at a @@ -734,14 +759,20 @@ async fn exec_cmd( .arg(format!("set -o pipefail; {cmd}")) .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()); - // SAFETY: `setpgid`/`nice` are async-signal-safe and only touch the - // calling (child) process before exec — no allocations, no locks. + let cgroup_procs = cgroup_procs.map(std::os::fd::AsRawFd::as_raw_fd); + // SAFETY: `setpgid`/`nice`/`write` are async-signal-safe and only touch + // the calling (child) process before exec — no allocations, no locks. // `setpgid(0, 0)` makes the child its own process-group leader so the // whole subtree can be signalled as a group (`kill(-pgid, …)`). + // Entering the task's cgroup before exec means nothing the task forks + // starts outside it. unsafe { - cmd_builder.pre_exec(|| { + cmd_builder.pre_exec(move || { libc::setpgid(0, 0); libc::nice(10); + if let Some(procs) = cgroup_procs { + run_cgroup::enter_in_pre_exec(procs)?; + } Ok(()) }); } diff --git a/hive-subagent-mcp/src/main.rs b/hive-subagent-mcp/src/main.rs index 7f06e651..a8ffc4df 100644 --- a/hive-subagent-mcp/src/main.rs +++ b/hive-subagent-mcp/src/main.rs @@ -22,8 +22,15 @@ struct Cli { http: std::net::SocketAddr, } +fn main() -> Result<()> { + // Each subagent run starts as this binary, which enters the run's cgroup + // and execs the run's program (`RunCgroup::entry`). + hive_agent_sock::run_cgroup::exec_entry(); + serve() +} + #[tokio::main] -async fn main() -> Result<()> { +async fn serve() -> Result<()> { tracing_subscriber::fmt() .with_env_filter( tracing_subscriber::EnvFilter::try_from_env("RUST_LOG") @@ -57,7 +64,8 @@ async fn main() -> Result<()> { hive_subagent_mcp::session::State::new(todo_socket, signal_base) .on_runtime(runtime) .publishing_to(hive_subagent_mcp::swarm_term::Publisher::spawn()) - .reading_models_from(hive_subagent_mcp::swarm_models::Reader::spawn()), + .reading_models_from(hive_subagent_mcp::swarm_models::Reader::spawn()) + .in_run_cgroups(hive_agent_sock::run_cgroup::RunCgroups::init("runs")), ); // Serve the MCP tools over streamable-http forever. No background poll diff --git a/hive-subagent-mcp/src/session.rs b/hive-subagent-mcp/src/session.rs index 64aa940f..907ae83e 100644 --- a/hive-subagent-mcp/src/session.rs +++ b/hive-subagent-mcp/src/session.rs @@ -89,9 +89,11 @@ use std::collections::HashMap; use std::os::unix::process::ExitStatusExt as _; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use std::time::{Duration, Instant}; +use hive_agent_sock::run_cgroup::{RunCgroup, RunCgroups}; use hive_claude::{Attach, Claude, Config, SessionStore}; use hive_runtime::{ AcpCommand, AcpRuntime, PermissionAsk, PermissionPolicy, Runtime as _, RuntimeSpec, TurnStop, @@ -405,6 +407,11 @@ pub struct State { /// The agent's own model list, checked when the swarm sets none: /// [`available_models`] outside tests. local_models: fn() -> Option, + /// Where each run gets its own cgroup, when this daemon was delegated + /// one ([`State::in_run_cgroups`]). + run_cgroups: Option, + /// The `` of the next run cgroup `-`. + run_seq: AtomicU64, } /// Which opaque URL segment belongs to which session — the whole of a @@ -460,6 +467,8 @@ impl State { term: None, models: None, local_models: available_models, + run_cgroups: None, + run_seq: AtomicU64::new(0), } } @@ -478,6 +487,26 @@ impl State { self } + /// Run each subagent run in its own cgroup under `cgroups`, at + /// `oom_score_adj` +1000. + #[must_use] + pub fn in_run_cgroups(mut self, cgroups: Option) -> Self { + self.run_cgroups = cgroups; + self + } + + /// A new cgroup for one run of `name`; `None` without run cgroups, or + /// (logged) when creating it failed, and the run then stays in the + /// daemon's cgroup. + fn run_cgroup(&self, name: &str) -> Option { + let cgroups = self.run_cgroups.as_ref()?; + let n = self.run_seq.fetch_add(1, Ordering::Relaxed); + cgroups + .create(&format!("{name}-{n}")) + .inspect_err(|e| tracing::warn!(name, error = %e, "subagent: run cgroup unavailable")) + .ok() + } + /// Refuse a `model` outside the list in effect for this agent (see /// `check_model`). The swarm key is read only for a named model. async fn allow_model(&self, model: Option<&str>) -> anyhow::Result<()> { @@ -1394,7 +1423,8 @@ fn start_on_runtime( trigger: String, ) -> anyhow::Result { if let SubagentRuntime::Acp { command, sessions } = &state.runtime { - let runtime = acp_runtime(command.clone(), sessions, name); + let cgroup = state.run_cgroup(name); + let runtime = acp_runtime(command.clone(), sessions, name, cgroup.as_ref()); if let Some(archived) = runtime .archive() .map_err(|e| anyhow::anyhow!("archiving the prior `{name}` session failed: {e}"))? @@ -1409,7 +1439,7 @@ fn start_on_runtime( system_prompt_file: system_prompt, ..config }; - return spawn_and_track_acp(state, name, runtime, config, trigger, None); + return spawn_and_track_acp(state, name, runtime, cgroup, config, trigger, None); } let store = build_store(&config)?; if store.find_by_title(name).is_some() { @@ -1608,11 +1638,13 @@ fn continue_on_runtime( ); } state.begin_continue(name, dir); - let runtime = acp_runtime(command.clone(), sessions, name); + let cgroup = state.run_cgroup(name); + let runtime = acp_runtime(command.clone(), sessions, name, cgroup.as_ref()); return spawn_and_track_acp( state, name, runtime, + cgroup, config, prompt, Some(Arc::clone(verdict)), @@ -1736,6 +1768,10 @@ fn spawn_and_track( prompt: String, verdict: Option, ) -> anyhow::Result { + // Every turn of the run, and all they start, share one cgroup, released + // when the background task below ends. + let cgroup = state.run_cgroup(name); + let config = &entered_config(config, cgroup.as_ref()); let running = Claude::spawn(config, attach) .map_err(|e| anyhow::anyhow!("starting the subagent process failed: {e}"))?; state.track(name, running.cancel_handle().into()); @@ -1761,6 +1797,7 @@ fn spawn_and_track( let state = Arc::clone(state); let task_name = name.to_owned(); tokio::spawn(async move { + let _cgroup = cgroup; let mut running = running; let mut prompt = prompt; loop { @@ -1853,11 +1890,12 @@ async fn after_turn( /// `runtime`, whose agent process is spawned by the first turn and killed /// when the run ends and the runtime is dropped. Returns before the agent /// exists, so a spawn failure is the first turn's `Failed` end rather than -/// this call's error. +/// this call's error. The run's `cgroup` is released after the runtime. fn spawn_and_track_acp( state: &Arc, name: &str, runtime: AcpRuntime, + cgroup: Option, config: Config, prompt: String, verdict: Option, @@ -1884,12 +1922,14 @@ fn spawn_and_track_acp( Err(e) => TurnEnd::Failed(format!("ACP error: {e}")), }; let Some(next) = after_turn(&state, &task_name, end, verdict.as_ref()).await else { - return; + break; }; state.track(&task_name, canceller.clone().into()); state.note_event(&task_name); prompt = next; } + drop(runtime); + drop(cgroup); }); Ok(format!("subagent `{name}` started")) @@ -1901,19 +1941,51 @@ fn acp_session_file(sessions: &Path, name: &str) -> PathBuf { sessions.join(format!("subagent-acp-session-{name}")) } +/// A subagent run's `oom_score_adj`: the maximum, so under memory pressure +/// the kernel takes a subagent run before a bash task (0) or the agent's +/// own session (-500, `nix/agent-modules/agent-service.nix`). +const RUN_OOM_SCORE_ADJ: i32 = 1000; + +/// `config` with its program started inside `cgroup`, if there is one. +/// `hive-claude` builds the `Command` itself and offers no `pre_exec`, so +/// the program becomes this daemon's binary as a [`RunCgroup::entry`] +/// wrapper. +fn entered_config(config: &Config, cgroup: Option<&RunCgroup>) -> Config { + let mut config = config.clone(); + if let Some(cgroup) = cgroup { + // `hive-claude`'s own default program. + let program = config.program.take().unwrap_or_else(|| "claude".into()); + let (entry, env) = cgroup.entry(&program, RUN_OOM_SCORE_ADJ); + config.program = Some(entry); + config.env.extend(env); + } + config +} + +/// [`entered_config`] for an ACP agent's command. +fn entered_command(mut command: AcpCommand, cgroup: Option<&RunCgroup>) -> AcpCommand { + if let Some(cgroup) = cgroup { + let (entry, env) = cgroup.entry(Path::new(&command.command), RUN_OOM_SCORE_ADJ); + command.command = entry.to_string_lossy().into_owned(); + command.env.extend(env); + } + command +} + /// The ACP runtime `name`'s run drives, keeping its session id in /// [`acp_session_file`] and answering permission requests per -/// [`acp_permits`]. +/// [`acp_permits`]. Its agent process starts inside `cgroup`. fn acp_runtime( command: AcpCommand, sessions: &Path, name: &str, + cgroup: Option<&RunCgroup>, ) -> AcpRuntime { let groups = hive_sh4re::permissions::effective_tool_groups(); let permit: PermissionPolicy = Arc::new(move |ask: &PermissionAsk<'_>| acp_permits(ask, &groups)); AcpRuntime::new( - command, + entered_command(command, cgroup), acp_session_file(sessions, name), permit, hive_claude::NeverCompact, diff --git a/hive-subagent-mcp/tests/run_entry.rs b/hive-subagent-mcp/tests/run_entry.rs new file mode 100644 index 00000000..b194276f --- /dev/null +++ b/hive-subagent-mcp/tests/run_entry.rs @@ -0,0 +1,28 @@ +//! The daemon binary as a run's entry wrapper (`run_cgroup::RunCgroup::entry` +//! sets these variables; its unit test pins the names). + +use std::process::Command; + +#[test] +fn the_entry_enters_the_cgroup_raises_the_score_and_execs_the_program() { + let cgroup = std::env::temp_dir().join(format!("hive-run-entry-{}", std::process::id())); + std::fs::create_dir_all(&cgroup).expect("fake cgroup"); + let out = Command::new(env!("CARGO_BIN_EXE_hive-subagent-daemon")) + // Stands in for `claude`: prints what it was started with. + .env("HIVE_RUN_ENTRY_PROGRAM", "/bin/sh") + .env("HIVE_RUN_ENTRY_CGROUP", &cgroup) + .env("HIVE_RUN_ENTRY_OOM_SCORE_ADJ", "1000") + .args([ + "-c", + r#"printf '%s|%s|%s' "$(cat /proc/self/oom_score_adj)" "${HIVE_RUN_ENTRY_CGROUP-unset}" "$1""#, + "sh", + "an-arg", + ]) + .output() + .expect("the binary runs"); + let procs = std::fs::read_to_string(cgroup.join("cgroup.procs")); + std::fs::remove_dir_all(&cgroup).ok(); + assert!(out.status.success(), "{out:?}"); + assert_eq!(String::from_utf8_lossy(&out.stdout), "1000|unset|an-arg"); + assert_eq!(procs.expect("cgroup.procs written"), "0"); +} diff --git a/nix/agent-modules/agent-service.nix b/nix/agent-modules/agent-service.nix index 88c4c81b..27d803f0 100644 --- a/nix/agent-modules/agent-service.nix +++ b/nix/agent-modules/agent-service.nix @@ -539,8 +539,8 @@ in # A LOWER OOMScoreAdjust means LESS likely to be killed, so this # negative value puts the harness — and the agent's own `claude`, # which it spawns as a child and which therefore inherits the - # value — last in line under memory pressure, behind - # `hive-subagent-daemon`'s `+500` (see ./mcp.nix). Losing a + # value — last in line under memory pressure, behind each subagent + # run (+1000) and bash task (0) (see ./mcp.nix). Losing a # subagent costs one restartable task; losing this costs the # session that was supervising it. Not `-1000`, which would # exempt the harness entirely and leave the kernel with nothing diff --git a/nix/agent-modules/mcp.nix b/nix/agent-modules/mcp.nix index d8b04f75..77bbd155 100644 --- a/nix/agent-modules/mcp.nix +++ b/nix/agent-modules/mcp.nix @@ -338,6 +338,11 @@ in # take every running sibling task with it. `continue` scopes the # loss to the process the kernel chose. OOMPolicy = "continue"; + # The daemon runs each task in its own `tasks//` cgroup with + # `memory.oom.group=1`, so an OOM kill takes one whole task. The + # subgroup name is `DAEMON_SUBGROUP` in hive-agent-sock's `run_cgroup`. + Delegate = "memory pids"; + DelegateSubgroup = "daemon"; User = userName; Group = userName; }; @@ -436,18 +441,15 @@ in # `start`/`continue`/`status`/`interrupt` with no stdio fallback. Restart = "always"; RestartSec = 3; - # A HIGHER OOMScoreAdjust means MORE likely to be killed: the - # kernel adds it to the badness score it derives from the - # process's memory footprint, then kills the highest scorer. So - # the positive value here and the negative one on `hive-agent` - # put this unit first in line — a subagent turn can be retried, - # the agent's own turn is the thing everything else hangs off. - # `+500` against `hive-agent`'s `-500` is a full half of the - # score range apart, enough that a fat subagent outranks a fatter - # harness rather than merely tying with it. Every nested `claude` - # inherits the value from the daemon, so the ordering covers the - # whole subtree and not just this process. - OOMScoreAdjust = 500; + # The daemon runs each subagent run in its own `runs/-/` + # cgroup with `memory.oom.group=1` and `oom_score_adj` +1000, so an + # OOM kill takes one whole run, ahead of a bash task (0) and + # `hive-agent` (-500). No `OOMScoreAdjust=` here: the daemon stays + # at 0, and a unit value would also be the floor no run could go + # below. The subgroup name is `DAEMON_SUBGROUP` in hive-agent-sock's + # `run_cgroup`. + Delegate = "memory pids"; + DelegateSubgroup = "daemon"; # systemd's default is `stop`, which tears the whole unit down as # soon as the kernel kills *any* process in its cgroup — so one # over-large subagent takes the daemon and every sibling session diff --git a/nix/module-eval/agent-memory.nix b/nix/module-eval/agent-memory.nix index 92be5450..16ade1dd 100644 --- a/nix/module-eval/agent-memory.nix +++ b/nix/module-eval/agent-memory.nix @@ -55,34 +55,45 @@ let && lib.elem "multi-user.target" agentDefault.systemd.services.hyperhive-memory-cap-check.wantedBy; } { - # The sign is the whole property, and it is easy to write backwards: - # a HIGHER OOMScoreAdjust is a MORE likely victim. So the subagent - # daemon must be strictly above zero and the harness strictly below - # it — swap the two and the kernel takes the agent's own turn first, - # which is worse than setting nothing at all. Both nested `claude` - # processes inherit their unit's value, so ordering the units orders - # the sessions. Held as an inequality rather than two constants: what - # must not drift is the order, not the magnitudes. - name = "the OOM killer prefers a subagent over the agent's own session"; + # A HIGHER OOMScoreAdjust is a MORE likely victim. The subagent daemon + # puts each run at +1000 itself and bash tasks stay at 0, so the harness + # must sit below 0 for the agent's own session to go last. Neither + # daemon unit carries a value: it would make the idle daemon a likelier + # victim than a bash task, and a floor no run under it could go below. + name = "the OOM killer takes the agent's own session last"; ok = - let - sub = (agentSubagentDaemon agentDefault).serviceConfig.OOMScoreAdjust; - own = (agentHarness agentDefault).serviceConfig.OOMScoreAdjust; - in - sub > 0 && own < 0 && sub > own; + (agentHarness agentDefault).serviceConfig.OOMScoreAdjust < 0 + && !((agentSubagentDaemon agentDefault).serviceConfig ? OOMScoreAdjust) + && !((agentBashDaemon agentDefault).serviceConfig ? OOMScoreAdjust); } { - # The pair the case above only makes sense with: preferring this unit - # as the victim is an improvement only if losing one subagent isn't - # losing all of them. systemd's default `stop` would take the daemon - # and every sibling session down with whichever process the kernel - # chose, which is the blast radius that made the preference a bad - # trade in the first place. + # Both daemons create a cgroup per run under their own; without the + # delegation they cannot, and without the `daemon` leaf (the name + # hive-agent-sock's `run_cgroup` looks for) they refuse to, so runs + # would share the daemon's cgroup again. + name = "the subagent and bash daemons are delegated their cgroup"; + ok = + lib.all + ( + unit: + unit.serviceConfig.Delegate or null == "memory pids" + && unit.serviceConfig.DelegateSubgroup or null == "daemon" + ) + [ + (agentSubagentDaemon agentDefault) + (agentBashDaemon agentDefault) + ]; + } + { + # Preferring a subagent run as the victim is an improvement only if + # losing one run isn't losing all of them. systemd's default `stop` + # would take the daemon and every sibling session down with whichever + # process the kernel chose. name = "one subagent losing the OOM draw does not stop the daemon"; ok = (agentSubagentDaemon agentDefault).serviceConfig.OOMPolicy == "continue"; } { - # Each bash task is a process in the daemon's cgroup; under `stop` one + # Every bash task runs under the daemon's unit; under `stop` one # OOM-killed task would end every running sibling task. name = "one bash task losing the OOM draw does not stop the daemon"; ok = (agentBashDaemon agentDefault).serviceConfig.OOMPolicy or null == "continue";