parallelize graceful agent drains, serialize container stops on fast lane; unify shutdown+checkpoint+compact prompt

This commit is contained in:
damocles 2026-06-29 19:19:20 +02:00 committed by mara
commit fc42f97691
5 changed files with 189 additions and 60 deletions

View file

@ -43,8 +43,11 @@ pub enum QueueKind {
/// actions for different agents never race on the shared JSON file.
PermChange,
/// Gracefully stop a container: signal the harness to run one
/// stop-checkpoint turn (flush durable `/state`), wait for it to drain,
/// then `nixos-container stop`. Falls back to a hard stop on timeout.
/// stop-checkpoint turn (flush durable `/state`), then hand the drain-wait
/// to a detached watcher (freeing the build lane) which, once the agent
/// drains or `GRACEFUL_STOP_TIMEOUT` elapses, enqueues a fast-lane `Stop`
/// for the actual `nixos-container stop`. The build worker only does the
/// cheap signal, so whole-hive graceful stops overlap every agent's drain.
GracefulStop,
/// Start a stopped container (`lifecycle::start`). Routed through the
/// queue so the dashboard shows a visible queued→running transient — a
@ -77,7 +80,8 @@ impl QueueKind {
/// serial fast worker concurrently with the build lane (so a stop/start
/// never waits behind another container's slow build). `GracefulStop`
/// and `Restart` are deliberately NOT fast — they go through the build
/// lane (`GracefulStop` holds the worker while the harness drains;
/// lane (`GracefulStop` does the cheap harness signal then detaches the
/// drain-wait, enqueueing a fast-lane `Stop` for the real container stop;
/// `Restart` is a stop+start).
pub fn is_fast(self) -> bool {
matches!(self, QueueKind::Start | QueueKind::Stop)
@ -752,10 +756,12 @@ fn lane_clear(entries: &VecDeque<QueueEntry>, e: &QueueEntry) -> bool {
/// signal the worker exits after its current entry finishes; pending
/// `Queued` entries are dropped (they'll either be replayed by the
/// startup sweep on next boot or left for an operator to re-queue).
/// Max time the `GracefulStop` worker waits for the harness to run its
/// Max time the `GracefulStop` drain watcher waits for the harness to run its
/// stop-checkpoint turn + drain before falling back to a hard container stop.
/// Generous — a checkpoint turn can take a while — but bounded so a wedged
/// agent never blocks the stop indefinitely.
/// agent never blocks the stop indefinitely. The wait runs in a detached
/// watcher task (not the build worker), so a whole-hive graceful stop overlaps
/// every agent's drain instead of serialising N × this timeout.
const GRACEFUL_STOP_TIMEOUT: std::time::Duration = std::time::Duration::from_mins(3);
/// Run one claimed queue entry to completion: snapshot, dispatch, mark
@ -1000,7 +1006,7 @@ async fn dispatch(
crate::auto_update::current_flake_rev(&coord.hyperhive_flake).unwrap_or_default();
crate::auto_update::rebuild_agent(coord, name, &current_rev, Some(entry.id), true).await
}
(QueueKind::GracefulStop, _) => run_graceful_stop(coord, entry).await,
(QueueKind::GracefulStop, _) => run_graceful_stop(coord, entry),
(QueueKind::Start, _) => run_start(coord, entry).await,
(QueueKind::Stop, _) => run_stop(coord, entry).await,
}
@ -1044,40 +1050,69 @@ async fn run_stop(
/// Run one `GracefulStop` entry: signal the harness to quiesce (it returns
/// `GracefulStop` on its next `Recv`, runs one stop-checkpoint turn to flush
/// durable `/state`, then exits), wait for it to drain — bounded by
/// `GRACEFUL_STOP_TIMEOUT` so a wedged agent can't block forever — then stop
/// the container with the same teardown as a plain kill.
async fn run_graceful_stop(
/// durable `/state`, then exits), then hand the drain-wait + container stop to
/// a detached watcher and return — freeing the build lane immediately.
///
/// This is the concurrency split: the build worker only does the cheap signal,
/// so a whole-hive graceful stop signals every agent up front and their
/// checkpoint drains overlap. The watcher waits for this agent's drain
/// (bounded by `GRACEFUL_STOP_TIMEOUT` so a wedged agent can't block forever),
/// then enqueues a fast-lane `Stop` for the actual `nixos-container stop`.
/// Routing the real stop through the fast lane means the container stops
/// serialise there (one stop at a time) while the drains ran in parallel.
// Returns `Result` for symmetry with the other `run_*` dispatch arms even
// though the fallible drain now lives in the detached watcher and this body
// is infallible.
#[allow(clippy::unnecessary_wraps)]
fn run_graceful_stop(
coord: &std::sync::Arc<crate::coordinator::Coordinator>,
entry: &QueueEntry,
) -> anyhow::Result<()> {
let name = &entry.agent;
let _guard = coord.transient_guard(name, crate::coordinator::TransientKind::Stopping);
let name = entry.agent.clone();
let parent_id = entry.id;
let source = entry.source;
// Signal the harness; the kick breaks an idle long-poll so it's seen promptly.
coord.set_queue_step(Some(entry.id), "graceful stop: signalling agent");
coord.mark_graceful_stop(name);
coord.kick_agent(name, "graceful stop requested");
// Wait for the harness to drain (it clears the flag via `GracefulStopComplete`)
// or fall back to a hard stop after the timeout. The single queue worker is
// intentionally held for the duration — graceful stops are infrequent.
coord.set_queue_step(Some(entry.id), "graceful stop: waiting for agent to drain");
let deadline = std::time::Instant::now() + GRACEFUL_STOP_TIMEOUT;
while coord.is_graceful_stop_pending(name) {
if std::time::Instant::now() >= deadline {
tracing::warn!(agent = %name, "graceful stop: drain timed out — hard-stopping");
break;
coord.mark_graceful_stop(&name);
coord.kick_agent(&name, "graceful stop requested");
// Detached watcher: wait for the drain (or timeout), then enqueue the
// container stop on the fast lane. The build entry itself is now Done —
// the dashboard groups the follow-up `Stop` under it via `parent_id`.
let coord = std::sync::Arc::clone(coord);
tokio::spawn(async move {
// Hold the `Stopping` transient across the drain so the dashboard keeps
// showing the agent quiescing; dropped before the fast `Stop` is
// enqueued (its `run_stop` re-establishes the transient) so the two
// never clobber each other's clear-on-drop.
let guard = coord.transient_guard(&name, crate::coordinator::TransientKind::Stopping);
// Wait for the harness to drain (it clears the flag via
// `GracefulStopComplete`) or fall back to a hard stop after the timeout.
let deadline = std::time::Instant::now() + GRACEFUL_STOP_TIMEOUT;
while coord.is_graceful_stop_pending(&name) {
if std::time::Instant::now() >= deadline {
tracing::warn!(agent = %name, "graceful stop: drain timed out — hard-stopping");
break;
}
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
}
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
}
coord.clear_graceful_stop(name);
// Stop the container — same teardown as a plain kill.
coord.set_queue_step(Some(entry.id), "nixos-container stop");
crate::lifecycle::kill(name).await?;
coord.unregister_agent(name);
coord.notify_manager(&hive_sh4re::HelperEvent::Killed {
agent: name.clone(),
coord.clear_graceful_stop(&name);
drop(guard);
// Enqueue the actual container stop on the fast lane (same teardown as a
// plain kill — `run_stop`). `parent_id` links it to the graceful entry
// for dashboard grouping. `enqueue_full` nudges the fast worker itself.
coord.rebuild_queue.enqueue_full(FullEnqueue {
kind: QueueKind::Stop,
agent: name.clone(),
source,
reason: format!("container stop after graceful drain of {name}"),
parent_id: Some(parent_id),
inputs: Vec::new(),
approval_id: None,
perm_payload: None,
depends_on: Vec::new(),
});
coord.emit_rebuild_queue_snapshot();
});
coord.rescan_containers_and_emit().await;
Ok(())
}