From 15a9d5b6522c8c0ba81ec204055483a24bb954d4 Mon Sep 17 00:00:00 2001 From: atlas Date: Mon, 27 Jul 2026 19:39:44 +0200 Subject: [PATCH] feat(#2454): drain agents before stopping them in the boot sweep MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A host restart brings hive-c0re up and the startup sweep rebuilds every stale agent. Until now that stop was mechanical: `StopForUpdate` hung straight off `Prebuild`, so an agent that was mid-turn when the host went down had its turn cut off rather than finished. The sweep now builds the same `Signal` -> `Drain` -> `StopForUpdate` chain a graceful `hivectl restart` already uses, reusing the existing nodes and `GRACEFUL_STOP_TIMEOUT` unchanged. Cost is bounded: the per-agent drains overlap, so the sweep waits one timeout in total rather than one per agent. `Signal` parents the rest of the stop instead of sitting beside it. All three of `Signal` / `Drain` / `StopForUpdate` declare the agent lease, and a resource is held across its holder's whole subtree — as siblings each would take the lease separately, leaving a window between them for another DAG to claim the agent mid-bounce. Scope is the boot sweep alone: a manual rebuild, a meta-update cascade child and a deploy all still stop mechanically, and a test pins that shape. `rebuild_nodes` takes a `RebuildOpts` struct rather than a second positional `bool`, which two adjacent flags would have made easy to swap at a call site. Its callers no longer hard-code the subgraph's length either: the `EmitRebuilt` tails and `FinalizeDeploy` used literal indices that silently encoded "this builder emits exactly six nodes with `Reconcile` last", which a variable-length subgraph turns into a wrong-node edge rather than a compile error. They read the index off the emitted list now. --- hive-c0re/src/job_queue/exec.rs | 28 +++++- hive-c0re/src/job_queue/submit.rs | 11 ++- hive-c0re/src/job_queue/templates.rs | 133 +++++++++++++++++++++------ hive-c0re/src/job_queue/tests.rs | 102 +++++++++++++++++++- 4 files changed, 236 insertions(+), 38 deletions(-) diff --git a/hive-c0re/src/job_queue/exec.rs b/hive-c0re/src/job_queue/exec.rs index b69e29a2..3d8ef461 100644 --- a/hive-c0re/src/job_queue/exec.rs +++ b/hive-c0re/src/job_queue/exec.rs @@ -322,10 +322,25 @@ async fn run_meta_lock( // (rooted on this `MetaLock`, so they build against the post-bump // lock), rather than fanning out child DAGs. `relock = true` — a // boot sweep relocks per-agent like a manual rebuild. + // + // `graceful = true` here and nowhere else: a boot sweep stops agents + // that were already mid-turn when the host came up, so they get their + // drain window rather than being cut off. The per-agent drains overlap, + // so the sweep's cost ceiling is one `GRACEFUL_STOP_TIMEOUT` in total, + // not one per agent. let append_subgraph = fanout .unwrap_or_default() .iter() - .map(|agent| super::templates::rebuild_nodes(agent, true, 0)) + .map(|agent| { + super::templates::rebuild_nodes( + agent, + super::templates::RebuildOpts { + relock: true, + graceful: true, + }, + 0, + ) + }) .collect(); return Ok(NodeOutput { append_subgraph }); } @@ -345,7 +360,16 @@ async fn run_meta_lock( // branch encoded). let append_subgraph = cascade .iter() - .map(|agent| super::templates::rebuild_nodes(agent, false, 0)) + .map(|agent| { + super::templates::rebuild_nodes( + agent, + super::templates::RebuildOpts { + relock: false, + graceful: false, + }, + 0, + ) + }) .collect(); Ok(NodeOutput { append_subgraph }) } diff --git a/hive-c0re/src/job_queue/submit.rs b/hive-c0re/src/job_queue/submit.rs index 71a2fd53..528f3bf6 100644 --- a/hive-c0re/src/job_queue/submit.rs +++ b/hive-c0re/src/job_queue/submit.rs @@ -25,7 +25,7 @@ use std::sync::Arc; use super::model::{DagSpec, Dep, NodeKind, NodeSpec}; -use super::templates::{after_ok, child, node, rebuild_nodes}; +use super::templates::{RebuildOpts, after_ok, child, node, rebuild_nodes}; use super::{Source, templates}; use crate::coordinator::{Coordinator, TransientKind}; use crate::lifecycle; @@ -98,7 +98,14 @@ fn start_chain(agent: &str, running: bool, stale: bool) -> Vec { // `MetaSync` root deps `after_ok(0)` = the head). `MetaSync`, // `Prebuild` + `Reconcile` are their own group roots (top-level, per // `rebuild_nodes`). - n.extend(rebuild_nodes(agent, true, 1)); + n.extend(rebuild_nodes( + agent, + RebuildOpts { + relock: true, + graceful: false, + }, + 1, + )); } else { n.push(child( 0, diff --git a/hive-c0re/src/job_queue/templates.rs b/hive-c0re/src/job_queue/templates.rs index 08608491..2245bd23 100644 --- a/hive-c0re/src/job_queue/templates.rs +++ b/hive-c0re/src/job_queue/templates.rs @@ -19,6 +19,7 @@ //! //! ```text //! rebuild(a): MetaSync(a) → Prebuild(a) → StopForUpdate(a) → Swap(a) →(ok) PostSwap(a) →(any) Reconcile(a) +//! rebuild(a) graceful: … Prebuild(a) → Signal(a) → Drain(a) → StopForUpdate(a) → … [boot sweep only] //! spawn(a): Provision(a) → Create(a) → WriteDropin(a) → Reconcile(a) [wanted=Up at approve] //! perm-change(a): WritePermFile(a) → «rebuild subgraph» //! meta-update(inp): MetaLock(inp) →«in-DAG rebuild subgraph per affected a» @@ -185,6 +186,19 @@ pub(crate) fn child(parent: u64, kind: NodeKind, deps: Vec) -> NodeSpec { } } +/// Knobs for [`rebuild_nodes`]. A struct rather than two positional `bool`s so +/// a call site cannot silently swap them. +#[derive(Debug, Clone, Copy)] +pub(crate) struct RebuildOpts { + /// Re-lock the meta flake inside `MetaSync`. + pub relock: bool, + /// Give the agent its `Signal` → `Drain` window to finish the turn in + /// flight before the container is stopped, instead of stopping it + /// outright. Costs up to one `GRACEFUL_STOP_TIMEOUT` per subgraph, and + /// those overlap across agents. + pub graceful: bool, +} + /// The rebuild node subtree (nested, three group roots). `base` is the spec /// index of the first node (`MetaSync`). Structure: /// - `MetaSync` (base+0, **root**): the meta-repo preamble (dir prep, agent @@ -196,14 +210,18 @@ pub(crate) fn child(parent: u64, kind: NodeKind, deps: Vec) -> NodeSpec { /// - `Prebuild` (base+1, **root**): `AfterOk` `MetaSync`. Owns the build slot /// for the whole mechanical subtree below it. Lease-exempt — the nix build /// overlaps other DAGs on the same agent. -/// - `StopForUpdate` (base+2, child of `Prebuild`): owns the agent lease. Runs -/// once `Prebuild` reaches `Finishing` (parent gate). -/// - `Swap` (base+3, child of `StopForUpdate`): borrows the agent lease from its -/// parent and the build slot from grand-ancestor `Prebuild` — both continuous. -/// - `PostSwap` (base+4, child of `StopForUpdate`): the swap's Ok-only +/// - the **stop root** (base+2, child of `Prebuild`): owns the agent lease and +/// runs once `Prebuild` reaches `Finishing` (parent gate). Non-graceful that +/// is `StopForUpdate` itself; graceful it is `Signal`, with `Drain` and then +/// `StopForUpdate` as its children so the lease stays continuous across the +/// whole stop — siblings would each take the lease separately and leave a gap +/// another DAG could claim the agent in, mid-bounce. +/// - `Swap` (child of `StopForUpdate`): borrows the agent lease from its +/// ancestors and the build slot from `Prebuild` — both continuous. +/// - `PostSwap` (child of `StopForUpdate`): the swap's Ok-only /// bookkeeping tail (rev marker, forge/matrix sync, kick, rescan), `AfterOk` /// its sibling `Swap`. -/// - `Reconcile` (base+5, **root**): `AfterAny` `Prebuild`, which rolls up +/// - `Reconcile` (**last, root**): `AfterAny` `Prebuild`, which rolls up /// terminal only once its whole mechanical subtree (SFU→Swap→PostSwap) has /// settled — so `Reconcile` runs after the swap regardless of outcome, and as /// a top-level root it survives the cancel-cascade of a failed `Prebuild` @@ -211,9 +229,10 @@ pub(crate) fn child(parent: u64, kind: NodeKind, deps: Vec) -> NodeSpec { /// cancel-cascades `Prebuild`, i.e. terminal, so the tail still runs). It /// takes a fresh lease; the tiny gap is harmless — `Reconcile` converges to /// the persisted `wanted` idempotently. -pub(crate) fn rebuild_nodes(agent: &str, relock: bool, base: u64) -> Vec { +pub(crate) fn rebuild_nodes(agent: &str, opts: RebuildOpts, base: u64) -> Vec { let a = || agent.to_owned(); - vec![ + let RebuildOpts { relock, graceful } = opts; + let mut nodes = vec![ node( NodeKind::MetaSync { agent: a(), relock }, if base == 0 { @@ -223,21 +242,43 @@ pub(crate) fn rebuild_nodes(agent: &str, relock: bool, base: u64) -> Vec Vec Vec { - let mut nodes = rebuild_nodes(agent, false, 0); + let mut nodes = rebuild_nodes( + agent, + RebuildOpts { + relock: false, + graceful: false, + }, + 0, + ); + let reconcile = reconcile_index(&nodes, 0); nodes.push(node( NodeKind::FinalizeDeploy { agent: agent.to_owned(), @@ -276,7 +325,7 @@ pub(crate) fn deploy_rebuild_nodes(agent: &str) -> Vec { when: DepWhen::AFTER_OK, }, Dep { - on: 5, + on: reconcile, when: DepWhen::AFTER_OK, }, ], @@ -284,6 +333,13 @@ pub(crate) fn deploy_rebuild_nodes(agent: &str) -> Vec { nodes } +/// Spec index of the `Reconcile` root a [`rebuild_nodes`] subgraph ends on, +/// for callers that gate a tail on it. Read off the emitted list rather than +/// hard-coded, because the subgraph's length depends on [`RebuildOpts`]. +fn reconcile_index(rebuild: &[NodeSpec], base: u64) -> u64 { + base + u64::try_from(rebuild.len()).unwrap_or(0).saturating_sub(1) +} + /// One uniform rebuild shape — no `was_running` branch. `StopForUpdate` /// noops when already down; the tail `Reconcile` auto-noops the start /// when `wanted = Offline` (a rebuild of a deliberately-stopped agent @@ -296,8 +352,17 @@ pub(crate) fn deploy_rebuild_nodes(agent: &str) -> Vec { /// node. Edging `Reconcile` alone would not do: it is `AfterAny` `Prebuild`, so /// it reaches `Done` even after a failed swap and the tail would report success. pub fn rebuild(agent: &str, source: Source, reason: String, relock: bool) -> DagSpec { - let mut nodes = rebuild_nodes(agent, relock, 0); - nodes.extend(emit_rebuilt_tails(agent, &[0, 1, 5], 6)); + let mut nodes = rebuild_nodes( + agent, + RebuildOpts { + relock, + graceful: false, + }, + 0, + ); + let reconcile = reconcile_index(&nodes, 0); + let tail_base = u64::try_from(nodes.len()).unwrap_or(0); + nodes.extend(emit_rebuilt_tails(agent, &[0, 1, reconcile], tail_base)); DagSpec { source, reason, @@ -433,8 +498,18 @@ pub fn perm_change(agent: &str, source: Source, reason: String, payload: PermPay }, Vec::new(), )]; - nodes.extend(rebuild_nodes(agent, true, 1)); - nodes.extend(emit_rebuilt_tails(agent, &[0, 1, 2, 6], 7)); + let rebuild = rebuild_nodes( + agent, + RebuildOpts { + relock: true, + graceful: false, + }, + 1, + ); + let reconcile = reconcile_index(&rebuild, 1); + nodes.extend(rebuild); + let tail_base = u64::try_from(nodes.len()).unwrap_or(0); + nodes.extend(emit_rebuilt_tails(agent, &[0, 1, 2, reconcile], tail_base)); DagSpec { source, reason, diff --git a/hive-c0re/src/job_queue/tests.rs b/hive-c0re/src/job_queue/tests.rs index 0088e837..35753a8e 100644 --- a/hive-c0re/src/job_queue/tests.rs +++ b/hive-c0re/src/job_queue/tests.rs @@ -234,6 +234,82 @@ fn rebuild_chain_claims_in_dep_order() { assert_eq!(state_of(&q, id), State::Done); } +/// The boot sweep's graceful shape: the agent gets `Signal` → `Drain` to +/// finish its turn before `StopForUpdate` takes the container down. `Signal` +/// *parents* the rest of the stop rather than sitting beside it, so the agent +/// lease is held continuously across the whole bounce — as siblings, each of +/// `Signal` / `Drain` / `StopForUpdate` would acquire the lease separately and +/// leave a window for another DAG to claim the agent mid-stop. +#[test] +fn graceful_rebuild_chain_drains_before_stopping() { + let q = JobQueue::new(1); + let spec = DagSpec { + source: Source::AutoUpdate, + reason: "sweep".to_owned(), + approval_id: None, + inputs: Vec::new(), + transient: None, + nodes: templates::rebuild_nodes( + "agent-a", + templates::RebuildOpts { + relock: true, + graceful: true, + }, + 0, + ), + }; + let id = submit(&q, spec); + for expected in [ + "meta_sync", + "prebuild", + "signal", + "drain", + "stop_for_update", + "swap", + "post_swap", + "reconcile", + ] { + let c = claim_one(&q); + assert_eq!(c.dag_id, id); + assert_eq!(c.kind.as_str(), expected); + assert!( + q.claim_ready().is_empty(), + "chain must serialize: nothing ready while {expected} runs" + ); + q.complete_node(id, c.node_id, Ok(())); + } + assert_eq!(state_of(&q, id), State::Done); +} + +/// The non-graceful shape is the default everywhere except the boot sweep: +/// a manual rebuild, a meta-update cascade child and a deploy must NOT spend a +/// drain window, so `StopForUpdate` still hangs straight off `Prebuild`. +#[test] +fn non_graceful_rebuild_has_no_signal_or_drain() { + let kinds: Vec = templates::rebuild_nodes( + "agent-a", + templates::RebuildOpts { + relock: true, + graceful: false, + }, + 0, + ) + .iter() + .map(|n| n.kind.as_str().to_owned()) + .collect(); + assert_eq!( + kinds, + vec![ + "meta_sync", + "prebuild", + "stop_for_update", + "swap", + "post_swap", + "reconcile" + ] + ); +} + // ---- build slots ---- #[test] @@ -644,10 +720,19 @@ fn append_subgraph_roots_on_emitter_and_rebases_local_deps() { let emitter = claim_one(&q); assert_eq!(emitter.kind.as_str(), "meta_lock"); // Two independent per-agent subgraphs — the REAL production shape the - // sweep MetaLock grows (`rebuild_nodes(_, true, 0)`: root MetaSync → root - // Prebuild → StopForUpdate → Swap → Reconcile, local 0-based deps), so this - // test tracks any drift in that builder's root-first (`base = 0`) shape. - let subgraph = |agent: &str| templates::rebuild_nodes(agent, true, 0); + // sweep MetaLock grows: root MetaSync → root Prebuild → Signal → Drain → + // StopForUpdate → Swap → Reconcile, local 0-based deps. `graceful` must + // match the sweep arm of `run_meta_lock` or this stops tracking production. + let subgraph = |agent: &str| { + templates::rebuild_nodes( + agent, + templates::RebuildOpts { + relock: true, + graceful: true, + }, + 0, + ) + }; // Must append BEFORE completing the emitter (the documented contract). q.append_subgraph(id, &subgraph("a"), emitter.node_id); q.append_subgraph(id, &subgraph("b"), emitter.node_id); @@ -715,7 +800,14 @@ fn meta_update_carries_rebuilding_transient_and_grows_cascade_in_dag() { for agent in ["alice", "bob"] { q.append_subgraph( id, - &templates::rebuild_nodes(agent, false, 0), + &templates::rebuild_nodes( + agent, + templates::RebuildOpts { + relock: false, + graceful: false, + }, + 0, + ), meta_lock.node_id, ); }