hyperhive/hive-c0re/src/job_queue/scheduler.rs

//! The single scheduler task that drives all DAGs: claim every ready node (as
//! many as the build slots / leases allow), spawn one executor task per claim,
//! and on any completion re-evaluate. Concurrency comes from the build-slot
//! count, not multiple workers.
//!
//! Owns the per-DAG transient guard (dashboard pill + crash-watch suppression)
//! that the sync queue core can't hold itself. The guard set is *reconciled*
//! from live lease ownership ([`super::JobQueue::held_transients`]) each loop:
//! a `(dag, agent)` pill exists for exactly as long as that agent's lease is
//! held, so it appears when the agent's owner node starts and disappears when
//! its subgraph settles — one pill per agent a DAG touches.
//!
//! Per-DAG terminal work (approval resolution, `Rebuilt`) is not drained here:
//! it runs as the DAG's focused terminal node (`ResolveApproval` /
//! `EmitRebuilt`), dispatched through `exec::run_node` like any other node once
//! the DAG settles.
//!
//! In-DAG growth (a `MetaLock` growing rebuild subgraphs, a `Reconcile` fanning
//! its `Start`/`Stop`) flows through `NodeOutput.append_subgraph`, applied
//! before the emitting node completes — see `handle_completion`.

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use super::Claim;
use super::exec::{self, NodeOutput};
use crate::coordinator::Coordinator;

struct NodeDone {
    claim: Claim,
    result: anyhow::Result<NodeOutput>,
}

/// Scheduler loop. Spawned once at hive-c0re startup from `main.rs`.
///
/// Shutdown semantics: subscribes to `coord.shutdown_rx()`. On a true signal
/// the loop exits — checked explicitly at the top of every iteration (not just
/// in the `select!` below), so a sustained stream of ready work can't defer
/// exit until the queue happens to drain. Already-running node tasks ride the
/// runtime down with the process, and pending `Queued` DAGs are dropped —
/// desired state is re-derived on next boot (boot sweep + reconcile), so the
/// in-memory queue is deliberately not durable.
///
/// That's safe only because every [`super::NodeKind`] is
/// idempotent-convergent: re-running a half-applied node converges instead of
/// double-applying (e.g. a dropped `Swap` re-derives to the same
/// `nixos-container update`, which is itself declarative). That's a property
/// of the node set, not of the queue — it holds today but isn't enforced by
/// the type system, so it's worth re-checking whenever a new `NodeKind` is
/// added. The side-effect tails (`EmitRebuilt`, `ResolveApproval`) sit
/// closest to the edge: dropping or re-running one affects something outside
/// the graph (today, a missed or duplicated notification) rather than
/// reconverging silently.
pub async fn run_worker(coord: Arc<Coordinator>) {
    let mut shutdown = coord.shutdown_rx();
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<NodeDone>();
    // (DAG id, agent) → transient guard held for that agent's lease window.
    let mut transients: HashMap<(u64, String), crate::coordinator::TransientGuard> = HashMap::new();
    loop {
        // Checked every iteration, not just in the `select!` below — a
        // continuous stream of ready claims never reaches the `select!`, so
        // relying on it alone as the only shutdown observation point defers
        // exit until the queue drains. See the doc comment above.
        if *shutdown.borrow() {
            tracing::info!("job_queue: scheduler exiting on shutdown");
            return;
        }
        reconcile_transients(&coord, &mut transients);
        let claims = coord.job_queue.claim_ready();
        if !claims.is_empty() {
            for claim in claims {
                tracing::info!(
                    dag = claim.dag_id,
                    node = claim.node_id.get(),
                    kind = claim.kind.as_str(),
                    agent = %claim.agent,
                    "job_queue: node running"
                );
                let coord = Arc::clone(&coord);
                let tx = tx.clone();
                tokio::spawn(async move {
                    let result = exec::run_node(&coord, &claim).await;
                    // Send failure = scheduler gone (shutdown); drop.
                    let _ = tx.send(NodeDone { claim, result });
                });
            }
            // Newly-started owner nodes now hold their leases — surface the pills.
            reconcile_transients(&coord, &mut transients);
            coord.emit_rebuild_queue_snapshot();
            continue;
        }
        tokio::select! {
            biased;
            res = shutdown.changed() => {
                if res.is_err() || *shutdown.borrow() {
                    tracing::info!("job_queue: scheduler exiting on shutdown");
                    return;
                }
            }
            Some(done) = rx.recv() => {
                handle_completion(&coord, done);
            }
            () = coord.job_queue.notify.notified() => {}
        }
    }
}

fn handle_completion(coord: &Arc<Coordinator>, done: NodeDone) {
    let NodeDone { claim, result } = done;
    match result {
        Ok(output) => {
            tracing::info!(
                dag = claim.dag_id,
                node = claim.node_id.get(),
                "job_queue: node done"
            );
            // Append any in-DAG subgraphs BEFORE completing this node, so
            // completing it doesn't roll the DAG terminal while the appended
            // work is still pending. Each subgraph roots on this node
            // (`AfterOk`), so it becomes ready the instant this one settles
            // `Done` just below — covers both the multi-node case (a `MetaLock`
            // growing per-agent rebuild subgraphs) and the single-node case (a
            // `Reconcile` planner's `Start` / `Stop`).
            for subgraph in &output.append_subgraph {
                coord
                    .job_queue
                    .append_subgraph(claim.dag_id, subgraph, claim.node_id);
            }
            coord.job_queue.complete_node(claim.node_id, Ok(()));
        }
        Err(e) => {
            let msg = format!("{e:#}");
            tracing::warn!(
                dag = claim.dag_id,
                node = claim.node_id.get(),
                kind = claim.kind.as_str(),
                agent = %claim.agent,
                error = %msg,
                "job_queue: node failed"
            );
            coord.job_queue.complete_node(claim.node_id, Err(msg));
        }
    }
    // The next loop iteration re-reconciles the transient pills against the
    // post-completion lease state (a settled subgraph drops its pill).
    coord.emit_rebuild_queue_snapshot();
}

/// Reconcile the transient-guard set against live lease ownership: drop pills
/// whose lease is no longer held, create one for each newly-held `(dag, agent)`.
fn reconcile_transients(
    coord: &Arc<Coordinator>,
    transients: &mut HashMap<(u64, String), crate::coordinator::TransientGuard>,
) {
    let held = coord.job_queue.held_transients();
    let keys: HashSet<(u64, String)> = held.iter().map(|(d, a, _)| (*d, a.clone())).collect();
    transients.retain(|k, _| keys.contains(k));
    for (dag_id, agent, kind) in held {
        transients
            .entry((dag_id, agent.clone()))
            .or_insert_with(|| coord.transient_guard(&agent, kind));
    }
}