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.
//!
//! Also owns the two DAG-lifetime side channels the sync queue core
//! can't hold itself:
//! - the per-DAG transient guard (dashboard pill + crash-watch
//!   suppression), created when a DAG acquires its agent lease and
//!   dropped when the DAG settles terminal;
//! - the `MetaLock` fan-out: appending child `Rebuild` DAGs once the
//!   lock bump lands, so children build against the post-bump lock
//!   (and a failed bump fans out nothing — replacing the old
//!   pre-enqueue + cancel-children dance).

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

use super::exec::{self, NodeOutput};
use super::{Claim, Source, Template, templates};
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 immediately; 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.
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. Keyed per-agent so a multi-agent DAG shows one transient
    // pill per agent it touches.
    let mut transients: HashMap<(u64, String), crate::coordinator::TransientGuard> = HashMap::new();
    loop {
        // Terminal roll-ups can appear without a node completion —
        // the cancel surfaces settle DAGs directly and wake this loop
        // via notify — so drain on every iteration, not just inside
        // handle_completion.
        process_terminals(&coord, &mut transients).await;
        let claims = coord.job_queue.claim_ready();
        if !claims.is_empty() {
            for claim in claims {
                if claim.lease_acquired
                    && let Some(kind) = claim.transient
                {
                    transients.insert(
                        (claim.dag_id, claim.agent.clone()),
                        coord.transient_guard(&claim.agent, kind),
                    );
                }
                tracing::info!(
                    dag = claim.dag_id,
                    node = claim.node_id,
                    kind = claim.kind.as_str(),
                    agent = %claim.agent,
                    template = claim.template.as_str(),
                    "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 });
                });
            }
            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, &mut transients, done).await;
            }
            () = coord.job_queue.notify.notified() => {}
        }
    }
}

async fn handle_completion(
    coord: &Arc<Coordinator>,
    transients: &mut HashMap<(u64, String), crate::coordinator::TransientGuard>,
    done: NodeDone,
) {
    let NodeDone { claim, result } = done;
    match result {
        Ok(output) => {
            tracing::info!(
                dag = claim.dag_id,
                node = claim.node_id,
                "job_queue: node done"
            );
            // Append any in-DAG sub-step nodes (e.g. a `Reconcile`
            // planner's `Start` / `Stop`) BEFORE completing this node, so
            // completing it doesn't roll the DAG terminal while the
            // appended work is still pending — that keeps the lease-window
            // transient held across the sub-step. Each depends `AfterOk`
            // on this node, so it becomes ready the instant this one
            // settles `Done` just below.
            for kind in output.append_nodes {
                coord
                    .job_queue
                    .append_node(claim.dag_id, kind, claim.node_id);
            }
            coord
                .job_queue
                .complete_node(claim.dag_id, claim.node_id, Ok(()));
            if !output.fanout.is_empty() {
                let specs = fanout_specs(&claim, output.fanout);
                coord.job_queue.append_children(specs);
            }
        }
        Err(e) => {
            let msg = format!("{e:#}");
            tracing::warn!(
                dag = claim.dag_id,
                node = claim.node_id,
                kind = claim.kind.as_str(),
                agent = %claim.agent,
                error = %msg,
                "job_queue: node failed"
            );
            coord
                .job_queue
                .complete_node(claim.dag_id, claim.node_id, Err(msg));
        }
    }
    process_terminals(coord, transients).await;
    coord.emit_rebuild_queue_snapshot();
}

/// Drain buffered terminal roll-ups: drop each DAG's lease-window
/// transient guard, then run the terminal hook (approval resolution,
/// `Rebuilt` events, cancelled-power-op intent revert).
async fn process_terminals(
    coord: &Arc<Coordinator>,
    transients: &mut HashMap<(u64, String), crate::coordinator::TransientGuard>,
) {
    for terminal in coord.job_queue.drain_terminal() {
        // Drop every per-agent transient guard this DAG held.
        transients.retain(|(dag_id, _), _| *dag_id != terminal.dag_id);
        exec::on_dag_terminal(coord, &terminal).await;
    }
}

/// Child `Rebuild` specs for a completed `MetaLock` fan-out, grouped
/// under the parent via `parent_id`. Meta-update children skip the
/// per-agent relock (it would revert the bump the parent just
/// committed); sweep children relock like a manual rebuild.
fn fanout_specs(claim: &Claim, agents: Vec<String>) -> Vec<super::DagSpec> {
    let sweep = claim.template == Template::StartupSweep;
    let (source, relock) = if sweep {
        (Source::StartupSweep, true)
    } else {
        (Source::MetaUpdate, false)
    };
    let reason = if sweep {
        "startup sweep".to_owned()
    } else if let Some(approval_id) = claim.approval_id {
        format!("approval #{approval_id} meta input cascade")
    } else {
        "meta-update cascade".to_owned()
    };
    agents
        .into_iter()
        .map(|agent| templates::rebuild(&agent, source, reason.clone(), Some(claim.dag_id), relock))
        .collect()
}