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-agent transient guard (dashboard pill + crash-watch
//! suppression) that the sync queue core can't hold itself. The guard set is
//! *reconciled* each loop from [`super::JobQueue::running_transients`], which
//! reports what is **running right now** under each held agent lease — so the
//! label tracks the DAG's progress (signal → swap → reconcile) instead of
//! repeating one intent the template declared before any of it started.
//!
//! 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`) is declared onto the builder each node is handed, and
//! inserted as part of completing that node. Completion itself is not this
//! module's job any more: it happens *inside* the future
//! [`hive_jobq::scheduler::Scheduler::claim_next`] hands back, so a node that
//! ran but was never completed is not an expressible state here.

use std::collections::HashMap;
use std::sync::Arc;

use super::exec;
use crate::coordinator::Coordinator;

/// 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();
    // Last derived pill set we published, keyed by agent (its lease is cap-1,
    // so one pill each). Purely the previous value of a *derived* quantity —
    // it exists to spot transitions, since the dashboard wants edges
    // (`TransientSet` / `TransientCleared`) and the crash watcher wants the
    // moment of the clear. Nothing owns a pill; nothing can leak one.
    let mut transients: HashMap<String, (String, bool)> = 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);
        // Claim exactly one node and get back the work that runs it. `Some`
        // means something started, so there may be more runnable right now —
        // loop again immediately. `None` means nothing is runnable and the
        // loop parks below. That decision is the whole reason the crate hands
        // back a task rather than an id.
        let runner = {
            // Two handles, deliberately: `sched` is the scheduler the crate
            // locks, `node_coord` is what the node's own future captures. One
            // binding can't do both — passing `coord.job_queue.sched()` borrows
            // `coord` for the whole call while the `move` closure wants to take
            // it.
            let sched = Arc::clone(coord.job_queue.sched());
            let node_coord = Arc::clone(&coord);
            hive_jobq::scheduler::Scheduler::claim_next(&sched, move |id, kind, job| {
                let coord = node_coord;
                async move {
                    tracing::info!(
                        dag = coord.job_queue.dag_of(id).unwrap_or_default(),
                        node = id.get(),
                        kind = kind.as_str(),
                        agent = %kind.agent(),
                        "job_queue: node running"
                    );
                    let (grown, result) = exec::run_node(&coord, job, id, &kind).await;
                    match &result {
                        Ok(()) => tracing::info!(node = id.get(), "job_queue: node done"),
                        Err(e) => tracing::warn!(
                            node = id.get(),
                            kind = kind.as_str(),
                            agent = %kind.agent(),
                            error = %format!("{e:#}"),
                            grown_nodes = !grown.is_empty(),
                            "job_queue: node failed"
                        ),
                    }
                    // Growth on a failed node is dropped by `complete_growing`,
                    // not here: failure cancel-cascades inside jobq, so that
                    // rule is the crate's to enforce and this loop does not get
                    // to forget it.
                    (
                        grown,
                        super::outcome_of(result.map_err(|e| format!("{e:#}"))),
                    )
                }
            })
        };
        if let Some(runner) = runner {
            let done_coord = Arc::clone(&coord);
            tokio::spawn(async move {
                // Completion happens inside `runner` — it cannot be forgotten
                // here, which is why there is no completion channel any more.
                let (id, grew) = runner.await;
                if let Err(e) = grew {
                    tracing::warn!(node = id.get(), error = %e, "job_queue: grown job rejected");
                }
                done_coord.emit_rebuild_queue_snapshot();
                // Wake the loop: this node's completion may have unblocked
                // dependents. Previously the completion channel did this.
                done_coord.job_queue.notify.notify_one();
            });
            // 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;
                }
            }
            () = coord.job_queue.notify.notified() => {}
        }
    }
}

/// Publish the transitions between the previously-derived pill set and the
/// current one. `prev` is last loop's derived value, keyed by agent (an agent's
/// lease is cap-1, so at most one pill each).
///
/// The pill set itself isn't owned or stored here — it is
/// [`super::JobQueue::running_transients`], recomputed from the graph. What this
/// publishes is the **edges**, which a derived read can't express on its own:
/// the dashboard wants `TransientSet` / `TransientCleared` events, and the crash
/// watcher wants the *moment* a pill cleared (its grace window is what keeps an
/// operator stop from reading as a crash).
///
/// There used to be an RAII `TransientGuard` per pill here, and a hazard note
/// about dropping stale guards before creating new ones or a same-agent label
/// change would clear the pill it had just set. Both are gone: a guard exists so
/// a cancelled future can't *leak* an imperatively-set transient, and a derived
/// set has nothing to leak — a node that stops running simply stops appearing.
/// (Destroy and migration still take guards; they have no node behind them.)
///
/// `deliberate_stop` rides along per node ([`NodeKind::takes_container_down`])
/// rather than being blanket-`true` for anything holding a lease: `Create` and
/// `Start` hold the agent's lease too, and a container vanishing *while
/// starting* is a real crash that must keep reporting as one.
///
/// [`NodeKind::takes_container_down`]: super::NodeKind::takes_container_down
fn reconcile_transients(coord: &Arc<Coordinator>, prev: &mut HashMap<String, (String, bool)>) {
    let running = coord.job_queue.running_transients();

    // Cleared: in `prev`, gone (or relabelled) now. Emitted before the sets
    // below so a same-agent label change reads as clear-then-set rather than
    // two overlapping pills. `deliberate_stop` is carried in `prev` precisely
    // so it is still available *here* — the node it came from is, by
    // definition, no longer running to be asked.
    prev.retain(|agent, (label, deliberate)| {
        let still = running
            .iter()
            .any(|t| &t.agent == agent && &t.label == label);
        if !still {
            coord.emit_transient_cleared(agent, *deliberate);
        }
        still
    });

    for t in running {
        if prev.get(&t.agent).map(|(l, _)| l) == Some(&t.label) {
            continue;
        }
        coord.emit_transient_set(&t.agent, t.label.clone());
        prev.insert(t.agent, (t.label, t.takes_container_down));
    }
}