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

@ -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<String> = 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,
);
}