feat(#2454): drain agents before stopping them in the boot sweep

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.
This commit is contained in:
atlas 2026-07-27 19:39:44 +02:00 committed by mara
commit 15a9d5b652
4 changed files with 236 additions and 38 deletions

View file

@ -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<Dep>) -> 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<Dep>) -> 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<Dep>) -> 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<NodeSpec> {
pub(crate) fn rebuild_nodes(agent: &str, opts: RebuildOpts, base: u64) -> Vec<NodeSpec> {
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<NodeSpe
},
),
node(NodeKind::Prebuild { agent: a() }, after_ok(base)),
child(base + 1, NodeKind::StopForUpdate { agent: a() }, Vec::new()),
child(base + 2, NodeKind::Swap { agent: a() }, Vec::new()),
child(
base + 2,
NodeKind::PostSwap { agent: a() },
after_ok(base + 3),
),
node(
NodeKind::Reconcile { agent: a() },
vec![Dep {
on: base + 1,
when: DepWhen::AFTER_ANY,
}],
),
]
];
// The stop root hangs off `Prebuild` and owns the agent lease for
// everything below it.
let stop_root = base + 2;
if graceful {
nodes.push(child(base + 1, NodeKind::Signal { agent: a() }, Vec::new()));
// `Drain` is a *child* of `Signal`, so the parent gate already orders
// it — a child must not dep on its own parent (dep-scope).
nodes.push(child(stop_root, NodeKind::Drain { agent: a() }, Vec::new()));
nodes.push(child(
stop_root,
NodeKind::StopForUpdate { agent: a() },
after_ok(stop_root + 1),
));
} else {
nodes.push(child(
base + 1,
NodeKind::StopForUpdate { agent: a() },
Vec::new(),
));
}
// Index of `StopForUpdate`, which parents the swap pair either way.
let sfu = if graceful { stop_root + 2 } else { stop_root };
nodes.push(child(sfu, NodeKind::Swap { agent: a() }, Vec::new()));
nodes.push(child(
sfu,
NodeKind::PostSwap { agent: a() },
after_ok(sfu + 1),
));
nodes.push(node(
NodeKind::Reconcile { agent: a() },
vec![Dep {
on: base + 1,
when: DepWhen::AFTER_ANY,
}],
));
nodes
}
/// The rebuild subgraph a [`NodeKind::DeployApply`] grows into its own DAG once
@ -265,7 +306,15 @@ pub(crate) fn rebuild_nodes(agent: &str, relock: bool, base: u64) -> Vec<NodeSpe
/// `MetaSync` and `FinalizeDeploy` declare is re-entered from the ancestor
/// already holding it rather than deadlocking against it.
pub(crate) fn deploy_rebuild_nodes(agent: &str) -> Vec<NodeSpec> {
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<NodeSpec> {
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<NodeSpec> {
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<NodeSpec> {
/// 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,