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;

/// Pills published on the previous tick: `(agent, label) -> takes_container_down`.
///
/// Keyed by the **pair**, not by agent. An agent can have several pills at once
/// now that [`super::JobQueue::running_transients`] tests status alone — a
/// lease-exempt `Prebuild` for `a` runs happily while another DAG holds `a`'s
/// lease, and both name `a`. Keying by agent would drop one arbitrarily and,
/// worse, lose its `takes_container_down` — which is the crash watcher's input
/// and is stored as the value precisely so it survives to *clear* time, when the
/// node that carried it is already gone.
type TransientSeen = HashMap<(String, String), bool>;

/// 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. 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 = TransientSeen::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, builder| {
                let coord = node_coord;
                async move {
                    let kind_str = <&str>::from(&kind);
                    tracing::info!(
                        dag = coord.job_queue.root_of(id).unwrap_or_default(),
                        node = id.get(),
                        kind = kind_str,
                        agent = %kind.agent(),
                        "job_queue: node running"
                    );
                    let (grown, result) = exec::run_node(&coord, builder, id, &kind).await;
                    match &result {
                        Ok(()) => tracing::info!(node = id.get(), "job_queue: node done"),
                        Err(e) => tracing::warn!(
                            node = id.get(),
                            kind = kind_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 TransientSeen) {
    let running = coord.job_queue.running_transients();

    // Cleared: in `prev`, gone now. Emitted before the sets below so a
    // replacement reads as clear-then-set rather than two overlapping pills.
    // `deliberate_stop` is the value 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, label, *deliberate);
        }
        still
    });

    for t in running {
        let key = (t.agent, t.label);
        if prev.contains_key(&key) {
            continue;
        }
        coord.emit_transient_set(&key.0, key.1.clone());
        prev.insert(key, t.takes_container_down);
    }
}

#[cfg(test)]
mod tests {
    use hive_jobq::scheduler::{Outcome, Scheduler};

    use super::super::{JobQueue, State, templates};

    /// Claim one runnable node through the same seam [`super::run_worker`]
    /// uses. The returned future completes the node when awaited. `None` when
    /// nothing is runnable.
    fn claim(q: &JobQueue) -> Option<impl Future<Output = ()>> {
        Scheduler::claim_next(q.sched(), |_, _, builder| async move {
            (builder, Outcome::Done)
        })
        .map(|run| async move {
            run.await.1.expect("a sweep grows nothing");
        })
    }

    /// Kinds of the nodes currently `Running`, sorted.
    fn running(q: &JobQueue) -> Vec<&'static str> {
        let sched = q.sched().lock().expect("job_queue mutex poisoned");
        let mut kinds: Vec<&'static str> = sched
            .graph()
            .nodes()
            .filter(|n| n.state == State::Running)
            .map(|n| <&'static str>::from(&n.payload))
            .collect();
        kinds.sort_unstable();
        kinds
    }

    fn count(q: &JobQueue, kind: &str) -> usize {
        let sched = q.sched().lock().expect("job_queue mutex poisoned");
        sched
            .graph()
            .nodes()
            .filter(|n| <&str>::from(&n.payload) == kind)
            .count()
    }

    #[tokio::test]
    async fn two_passes_of_one_sweep_never_run_together() {
        let q = JobQueue::new(1);
        for _ in 0..2 {
            q.insert_job(|b| vec![templates::matrix_sweep(b).guid()])
                .expect("insert");
        }
        q.insert_job(|b| vec![templates::knowledge_pull(b).guid()])
            .expect("insert");

        let first = claim(&q).expect("a matrix pass is runnable");
        let other = claim(&q).expect("a different sweep runs alongside it");
        assert!(
            claim(&q).is_none(),
            "the second matrix pass waits for the first"
        );
        assert_eq!(running(&q), ["knowledge_pull", "matrix_sweep"]);

        first.await;
        other.await;
        let second = claim(&q).expect("the second matrix pass runs once the first is done");
        assert_eq!(running(&q), ["matrix_sweep"]);
        second.await;
    }

    /// A tick landing while a pass is already live queues a second pass; the
    /// shared `Resource::KnowledgeTree` dep holds it back until the first
    /// finishes, so the two never run together.
    #[tokio::test]
    async fn a_tick_during_a_live_pull_queues_one_pass_behind_it() {
        let q = JobQueue::new(1);
        q.insert_job(|b| vec![templates::knowledge_pull(b).guid()])
            .expect("insert");
        let live = claim(&q).expect("the first pull is runnable");
        assert_eq!(running(&q), ["knowledge_pull"]);

        q.insert_job(|b| vec![templates::knowledge_pull(b).guid()])
            .expect("insert");
        assert_eq!(count(&q, "knowledge_pull"), 2);
        assert!(
            claim(&q).is_none(),
            "the second pull waits for the resource the first one holds"
        );

        live.await;
        let second = claim(&q).expect("the second pull runs once the first is done");
        assert_eq!(running(&q), ["knowledge_pull"]);
        second.await;
    }
}