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

//! Generic job-DAG queue + desired-state reconciliation — replaces the
//! old flat `rebuild_queue`. Jobs are nodes in per-request DAGs (see
//! [`templates`]); the special cases (graceful-stop watcher thread,
//! deferred-start follow-up, meta-update cascade) collapse into DAG
//! *shapes* over a shared set of primitive nodes ([`model::NodeKind`]).
//!
//! Concurrency is gated by two resource classes:
//! 1. **Build slots** — N permits (`services.hyperhive.c0re.buildSlots`,
//!    default 1) held by nix-heavy nodes for the node's duration.
//! 2. **Per-agent lifecycle lease** — keyed on the *node's* agent and
//!    globally exclusive per agent across all DAGs: acquired before a
//!    container-affecting node runs, held (by the owning DAG) until no
//!    live node of that DAG still targets the agent, so two DAGs never
//!    interleave container ops on the same agent. A DAG spanning multiple
//!    agents holds one lease per agent it touches.
//!
//! The meta *repo* is serialized by `meta::META_LOCK` inside the
//! executors themselves. Per-agent power *intent* (`wanted`) lives in
//! the durable [`crate::power`] store; the DAGs are the reconcile
//! mechanism. Design + rationale: `docs/coordinator.md::Job queue`.

pub mod exec;
pub mod model;
pub mod scheduler;
pub mod submit;
pub mod templates;
#[cfg(test)]
mod tests;

use std::collections::{HashMap, VecDeque};
use std::sync::Mutex;

use hive_sh4re::wire_time::now_unix;
use tokio::sync::Notify;

pub use model::{
    Dag, DagSpec, DagView, DepWhen, Node, NodeId, NodeKind, PermPayload, Source, State, Template,
};

/// How many terminal DAGs (`Done` / `Failed` / `Cancelled`) to retain
/// per template in the snapshot, matching the old per-kind history cap.
const MAX_HISTORY_PER_TEMPLATE: usize = 5;

/// Terminal DAGs younger than this are exempt from the per-template
/// history cap. A broad `hivectl stop`/`start` submits many
/// same-template DAGs that can all settle within one poll interval —
/// without the grace, the cap would evict some before the ~1s
/// `QueueDag` poller ever observes their terminal state, silently
/// swallowing failures.
const HISTORY_GRACE_SECS: i64 = 300;

/// Cap on stored node error strings.
const MAX_ERROR_LEN: usize = 2_000;

/// A node claimed for execution — everything the executor needs,
/// snapshotted at claim time.
#[derive(Debug, Clone)]
pub struct Claim {
    pub dag_id: u64,
    pub node_id: NodeId,
    pub kind: NodeKind,
    /// The agent this node targets (its own, not a DAG-level field). The
    /// executor operates on this agent's container; the lease is keyed on
    /// it.
    pub agent: String,
    pub template: Template,
    pub source: Source,
    pub approval_id: Option<i64>,
    pub inputs: Vec<String>,
    pub perm_payload: Option<PermPayload>,
    /// True when claiming this node newly acquired its agent's lease —
    /// the scheduler creates the per-`(dag, agent)` transient guard on
    /// this edge.
    pub lease_acquired: bool,
    /// Transient pill kind for the lease window (from the spec).
    pub transient: Option<crate::coordinator::TransientKind>,
}

/// Summary of a DAG that just reached its terminal roll-up state —
/// input to the approval-resolution hook and the lease/transient
/// release.
#[derive(Debug, Clone)]
pub struct TerminalDag {
    pub dag_id: u64,
    pub template: Template,
    /// Distinct agents this DAG's nodes targeted (one for a single-agent
    /// DAG). The cancel-revert hook walks these to snap each agent's
    /// power intent back on a cancelled power-op DAG.
    pub agents: Vec<String>,
    pub approval_id: Option<i64>,
    pub state: State,
    /// First failed node's error when `state == Failed`.
    pub error: Option<String>,
}

#[derive(Debug, Default)]
struct Inner {
    dags: VecDeque<Dag>,
    next_id: u64,
    build_slots: usize,
    slots_used: usize,
    /// agent → dag id currently holding that agent's lifecycle lease.
    leases: HashMap<String, u64>,
    /// Terminal roll-ups not yet consumed by the scheduler
    /// ([`JobQueue::drain_terminal`]). Fed by every path that settles
    /// state — node completion AND the cancel surfaces — so the
    /// terminal hooks (approval resolution, intent revert, transient
    /// release) fire exactly once per DAG no matter how it ended.
    pending_terminal: Vec<TerminalDag>,
}

/// The queue. Lives on `Coordinator` (one per hive-c0re process); a
/// single scheduler task ([`scheduler::run_worker`]) drives it —
/// concurrency comes from the build-slot count, not multiple workers.
#[derive(Debug)]
pub struct JobQueue {
    inner: Mutex<Inner>,
    /// Wakes the scheduler when something new arrives or state changed.
    pub(crate) notify: Notify,
}

impl Default for JobQueue {
    fn default() -> Self {
        Self::new(1)
    }
}

impl JobQueue {
    pub fn new(build_slots: usize) -> Self {
        Self {
            inner: Mutex::new(Inner {
                build_slots: build_slots.max(1),
                ..Inner::default()
            }),
            notify: Notify::new(),
        }
    }

    /// Submit a DAG. Validates the spec (cycle rejection) and returns the
    /// newly-allocated DAG id.
    ///
    /// Submit-time dedup was removed with the agent-per-node refactor
    /// (a multi-agent DAG has no single agent to key a dedup on) — every
    /// submit now enqueues a fresh DAG. Whether any dedup needs
    /// reintroducing (and in what form) is tracked as a follow-up; see the
    /// dedup re-evaluation issue.
    pub fn submit(&self, spec: DagSpec) -> anyhow::Result<u64> {
        templates::validate(&spec)?;
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let id = Self::push_dag(&mut inner, spec);
        drop(inner);
        self.notify.notify_one();
        Ok(id)
    }

    /// Append fan-out children under a parent DAG (meta-update / sweep
    /// cascade); returns the child ids created. No dedup (see
    /// [`Self::submit`]).
    pub fn append_children(&self, specs: Vec<DagSpec>) -> Vec<u64> {
        let mut ids = Vec::with_capacity(specs.len());
        for spec in specs {
            match self.submit(spec) {
                Ok(id) => ids.push(id),
                Err(e) => tracing::error!(error = ?e, "job_queue: invalid fan-out child spec"),
            }
        }
        ids
    }

    /// Append a node into a *live* (non-terminal) DAG at runtime,
    /// depending `AfterOk` on `dep_on` (the node that emitted it). Lets a
    /// planner node — e.g. [`NodeKind::Reconcile`] — fan a mechanical
    /// sub-step ([`NodeKind::Start`] / [`NodeKind::Stop`]) out as a
    /// first-class node in the *same* DAG.
    ///
    /// Must be called *before* the emitting node's [`Self::complete_node`]
    /// so the DAG doesn't roll terminal with the new node still pending —
    /// that keeps the lease-window transient held across the sub-step and
    /// lets the appended node's `AfterOk` dep resolve as soon as the
    /// emitter settles `Done`. No-op (returns `None`) if the DAG is gone.
    pub fn append_node(&self, dag_id: u64, kind: NodeKind, dep_on: NodeId) -> Option<NodeId> {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let dag = inner.dags.iter_mut().find(|d| d.id == dag_id)?;
        // The appended sub-step targets the same agent as the node that
        // emitted it (a `Reconcile` fanning out its `Start`/`Stop` acts on
        // the same container), so inherit `dep_on`'s agent.
        let agent = dag.node(dep_on)?.agent.clone();
        let new_id: NodeId = u32::try_from(dag.nodes.len()).unwrap_or(u32::MAX);
        dag.nodes.push(Node {
            id: new_id,
            agent,
            kind,
            deps: vec![model::Dep {
                on: dep_on,
                when: DepWhen::AfterOk,
            }],
            state: State::Queued,
            step: None,
            build_log_id: None,
            started_at: None,
            finished_at: None,
            error: None,
        });
        drop(inner);
        self.notify.notify_one();
        Some(new_id)
    }

    fn push_dag(inner: &mut Inner, spec: DagSpec) -> u64 {
        inner.next_id += 1;
        let id = inner.next_id;
        let nodes = spec
            .nodes
            .into_iter()
            .enumerate()
            .map(|(i, n)| Node {
                id: u32::try_from(i).unwrap_or(u32::MAX),
                agent: n.agent,
                kind: n.kind,
                deps: n.deps,
                state: State::Queued,
                step: None,
                build_log_id: None,
                started_at: None,
                finished_at: None,
                error: None,
            })
            .collect();
        inner.dags.push_back(Dag {
            id,
            template: spec.template,
            source: spec.source,
            reason: spec.reason,
            parent_id: spec.parent_id,
            approval_id: spec.approval_id,
            inputs: spec.inputs,
            perm_payload: spec.perm_payload,
            transient: spec.transient,
            created_at: now_unix(),
            nodes,
            terminal_reported: false,
        });
        id
    }

    /// Claim every currently-ready node, acquiring resources, and mark
    /// them `Running`. A node is ready when it's `Queued`, every dep is
    /// satisfied (`AfterOk`: dep `Done`; `AfterAny`: dep terminal), and
    /// its resources are free (build slot; agent lease free or already
    /// held by this DAG). Iteration is in DAG-submit order, so
    /// simultaneously-ready nodes compete FIFO — bulk operations drain
    /// predictably.
    pub fn claim_ready(&self) -> Vec<Claim> {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        Self::propagate_cancellations(&mut inner);
        let mut claims = Vec::new();
        let inner = &mut *inner;
        for di in 0..inner.dags.len() {
            // Split-borrow dance: deps are checked against the same
            // DAG's other nodes, so snapshot the states first.
            let dag = &inner.dags[di];
            let dag_id = dag.id;
            let ready_ids: Vec<NodeId> = dag
                .nodes
                .iter()
                .filter(|n| n.state == State::Queued && Self::deps_satisfied(dag, n))
                .map(|n| n.id)
                .collect();
            for node_id in ready_ids {
                let dag = &inner.dags[di];
                let node = dag.node(node_id).expect("node id from same dag");
                let needs_slot = node.kind.needs_build_slot();
                let needs_lease = node.kind.needs_lease();
                // The lifecycle lease is keyed on the *node's* agent, held
                // by this DAG (dag_id) — still globally exclusive per agent
                // across all DAGs. A multi-agent DAG acquires one lease per
                // agent it touches; each is released in `settle` once no
                // live node of this DAG still targets that agent.
                let node_agent = node.agent.clone();
                if needs_slot && inner.slots_used >= inner.build_slots {
                    continue;
                }
                let mut lease_acquired = false;
                if needs_lease {
                    match inner.leases.get(node_agent.as_str()) {
                        Some(&holder) if holder != dag_id => continue,
                        Some(_) => {}
                        None => {
                            inner.leases.insert(node_agent.clone(), dag_id);
                            lease_acquired = true;
                        }
                    }
                }
                if needs_slot {
                    inner.slots_used += 1;
                }
                let dag = &mut inner.dags[di];
                let claim = Claim {
                    dag_id,
                    node_id,
                    kind: dag.node(node_id).expect("node").kind.clone(),
                    agent: node_agent,
                    template: dag.template,
                    source: dag.source,
                    approval_id: dag.approval_id,
                    inputs: dag.inputs.clone(),
                    perm_payload: dag.perm_payload.clone(),
                    lease_acquired,
                    transient: dag.transient,
                };
                let node = dag.node_mut(node_id).expect("node");
                node.state = State::Running;
                node.started_at = Some(now_unix());
                claims.push(claim);
            }
        }
        claims
    }

    fn deps_satisfied(dag: &Dag, node: &Node) -> bool {
        node.deps.iter().all(|dep| {
            dag.node(dep.on).is_some_and(|d| match dep.when {
                DepWhen::AfterOk => d.state == State::Done,
                DepWhen::AfterAny => d.state.is_terminal(),
            })
        })
    }

    /// Cancel-downstream: a `Queued` node with an `AfterOk` dep that
    /// `Failed` / `Cancelled` becomes `Cancelled` itself. Loops to a
    /// fixpoint so the cancellation cascades through chains.
    fn propagate_cancellations(inner: &mut Inner) {
        for dag in &mut inner.dags {
            loop {
                let doomed: Vec<NodeId> = dag
                    .nodes
                    .iter()
                    .filter(|n| {
                        n.state == State::Queued
                            && n.deps.iter().any(|dep| {
                                dep.when == DepWhen::AfterOk
                                    && dag.node(dep.on).is_some_and(|d| {
                                        matches!(d.state, State::Failed | State::Cancelled)
                                    })
                            })
                    })
                    .map(|n| n.id)
                    .collect();
                if doomed.is_empty() {
                    break;
                }
                let now = now_unix();
                for id in doomed {
                    if let Some(n) = dag.node_mut(id) {
                        n.state = State::Cancelled;
                        n.finished_at = Some(now);
                    }
                }
            }
        }
    }

    /// Mark a claimed node terminal, release its build slot, cascade
    /// cancellations, and settle terminal DAGs (lease release + history
    /// trim; the terminal roll-up lands in the [`Self::drain_terminal`]
    /// buffer). `error` is stored (truncated) when `result` is `Err`.
    pub fn complete_node(&self, dag_id: u64, node_id: NodeId, result: Result<(), String>) {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        if let Some(dag) = inner.dags.iter_mut().find(|d| d.id == dag_id)
            && let Some(node) = dag.node_mut(node_id)
            && node.state == State::Running
        {
            let needs_slot = node.kind.needs_build_slot();
            node.finished_at = Some(now_unix());
            node.step = None;
            match result {
                Ok(()) => node.state = State::Done,
                Err(e) => {
                    node.state = State::Failed;
                    let mut msg = e;
                    if msg.len() > MAX_ERROR_LEN {
                        msg.truncate(
                            (0..=MAX_ERROR_LEN)
                                .rev()
                                .find(|i| msg.is_char_boundary(*i))
                                .unwrap_or(0),
                        );
                        msg.push('…');
                    }
                    node.error = Some(msg);
                }
            }
            if needs_slot {
                inner.slots_used = inner.slots_used.saturating_sub(1);
            }
        }
        Self::settle(&mut inner);
        drop(inner);
        self.notify.notify_one();
    }

    /// Take the terminal roll-ups accumulated since the last drain.
    /// The scheduler calls this after every wakeup and runs the
    /// terminal hooks on each entry.
    pub fn drain_terminal(&self) -> Vec<TerminalDag> {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        std::mem::take(&mut inner.pending_terminal)
    }

    /// Propagate cancellations, release the leases of newly-terminal
    /// DAGs, buffer each terminal roll-up exactly once (the
    /// `terminal_reported` flag) for [`Self::drain_terminal`], and trim
    /// history.
    fn settle(inner: &mut Inner) {
        Self::propagate_cancellations(inner);
        let mut freed: Vec<String> = Vec::new();
        let mut reports: Vec<TerminalDag> = Vec::new();
        for dag in &mut inner.dags {
            if !dag.is_terminal() || dag.terminal_reported {
                continue;
            }
            dag.terminal_reported = true;
            // Free every agent-lease this DAG holds (one per distinct
            // agent it touched). Single-agent DAGs release their one lease;
            // a future multi-agent DAG releases all of them at terminal.
            // (Per-agent early release — freeing an agent's lease the moment
            // that agent's subgraph is terminal rather than at whole-DAG
            // terminal — is a refinement for the multi-agent-emission
            // follow-up, where it actually matters.)
            let agents = dag.agents();
            for agent in &agents {
                if inner.leases.get(agent.as_str()) == Some(&dag.id) {
                    freed.push(agent.clone());
                }
            }
            reports.push(TerminalDag {
                dag_id: dag.id,
                template: dag.template,
                agents,
                approval_id: dag.approval_id,
                state: dag.rollup(),
                error: dag.first_error().map(str::to_owned),
            });
        }
        inner.pending_terminal.append(&mut reports);
        for agent in freed {
            inner.leases.remove(&agent);
        }
        Self::trim_history(inner, now_unix() - HISTORY_GRACE_SECS);
    }

    /// Cancel a DAG that hasn't started yet (roll-up `Queued`): every
    /// node flips to `Cancelled`. No-op (returns `false`) once any node
    /// is running or terminal — an in-flight nix build isn't
    /// interruptible, matching the old queue's rule.
    pub fn cancel(&self, dag_id: u64) -> bool {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let Some(dag) = inner.dags.iter_mut().find(|d| d.id == dag_id) else {
            return false;
        };
        if dag.rollup() != State::Queued {
            return false;
        }
        let now = now_unix();
        for n in &mut dag.nodes {
            n.state = State::Cancelled;
            n.finished_at = Some(now);
        }
        // Settle buffers the terminal roll-up; the notify wakes the
        // scheduler, which drains it and fires the terminal hooks
        // (approval resolution, power-intent revert).
        Self::settle(&mut inner);
        drop(inner);
        self.notify.notify_one();
        true
    }

    /// Cancel every still-fully-queued child DAG of `parent`. Running
    /// children are left alone. Returns the count of cancelled DAGs.
    pub fn cancel_children(&self, parent: u64) -> usize {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let now = now_unix();
        let mut count = 0;
        for dag in &mut inner.dags {
            if dag.parent_id == Some(parent) && dag.rollup() == State::Queued {
                for n in &mut dag.nodes {
                    n.state = State::Cancelled;
                    n.finished_at = Some(now);
                }
                count += 1;
            }
        }
        if count > 0 {
            Self::settle(&mut inner);
            drop(inner);
            self.notify.notify_one();
        }
        count
    }

    /// Set the step label on a `Running` node. Returns `true` when the
    /// label actually changed (callers emit a snapshot only then).
    pub fn set_step(&self, dag_id: u64, node_id: NodeId, step: &str) -> bool {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let Some(node) = inner
            .dags
            .iter_mut()
            .find(|d| d.id == dag_id)
            .and_then(|d| d.node_mut(node_id))
        else {
            return false;
        };
        if node.state != State::Running || node.step.as_deref() == Some(step) {
            return false;
        }
        node.step = Some(step.to_owned());
        true
    }

    /// Set the step label on the DAG's currently-running node —
    /// compatibility surface for the opaque approval pipeline, whose
    /// callbacks only know the DAG id. Single-node approval DAGs make
    /// this exact.
    pub fn set_step_running(&self, dag_id: u64, step: &str) -> bool {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let Some(node) = inner
            .dags
            .iter_mut()
            .find(|d| d.id == dag_id)
            .and_then(|d| d.nodes.iter_mut().find(|n| n.state == State::Running))
        else {
            return false;
        };
        if node.step.as_deref() == Some(step) {
            return false;
        }
        node.step = Some(step.to_owned());
        true
    }

    /// Link a `build_logs` row to a specific `Running` node.
    pub fn set_build_log_id(&self, dag_id: u64, node_id: NodeId, log_id: i64) -> bool {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let Some(node) = inner
            .dags
            .iter_mut()
            .find(|d| d.id == dag_id)
            .and_then(|d| d.node_mut(node_id))
        else {
            return false;
        };
        if node.state != State::Running {
            return false;
        }
        node.build_log_id = Some(log_id);
        true
    }

    /// Link a `build_logs` row to the DAG's currently-running node —
    /// DAG-id-only compatibility surface (approval pipeline callbacks).
    pub fn set_build_log_id_running(&self, dag_id: u64, log_id: i64) -> bool {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        let Some(node) = inner
            .dags
            .iter_mut()
            .find(|d| d.id == dag_id)
            .and_then(|d| d.nodes.iter_mut().find(|n| n.state == State::Running))
        else {
            return false;
        };
        node.build_log_id = Some(log_id);
        true
    }

    /// Snapshot every DAG for `/api/state` + `RebuildQueueChanged`.
    pub fn snapshot(&self) -> Vec<DagView> {
        let inner = self.inner.lock().expect("job_queue mutex poisoned");
        inner.dags.iter().map(Dag::view).collect()
    }

    /// Number of live (non-terminal) DAGs — used by tests and
    /// diagnostics.
    #[cfg(test)]
    pub fn live_count(&self) -> usize {
        let inner = self.inner.lock().expect("job_queue mutex poisoned");
        inner.dags.iter().filter(|d| !d.is_terminal()).count()
    }

    /// Keep only the newest `MAX_HISTORY_PER_TEMPLATE` terminal DAGs
    /// per template. Never evicted: live DAGs; terminal parents with
    /// live children (a fan-out parent is terminal the moment its
    /// `MetaLock` completes — evicting it while cascade rebuilds run
    /// would orphan their dashboard group); and terminal DAGs that
    /// finished after `grace_cutoff` (see [`HISTORY_GRACE_SECS`]).
    fn trim_history(inner: &mut Inner, grace_cutoff: i64) {
        let live_parents: std::collections::HashSet<u64> = inner
            .dags
            .iter()
            .filter(|d| !d.is_terminal())
            .filter_map(|d| d.parent_id)
            .collect();
        let mut counts: HashMap<Template, usize> = HashMap::new();
        let kept: Vec<Dag> = inner
            .dags
            .iter()
            .rev()
            .filter(|d| {
                if !d.is_terminal() || live_parents.contains(&d.id) {
                    return true;
                }
                let finished = d.nodes.iter().filter_map(|n| n.finished_at).max();
                if finished.is_none_or(|t| t > grace_cutoff) {
                    return true;
                }
                let n = counts.entry(d.template).or_insert(0);
                *n += 1;
                *n <= MAX_HISTORY_PER_TEMPLATE
            })
            .cloned()
            .collect();
        inner.dags = kept.into_iter().rev().collect();
    }

    /// Test hook: trim with the grace window disabled, so eviction
    /// behavior is assertable without aging real timestamps.
    #[cfg(test)]
    pub(crate) fn trim_ignoring_grace(&self) {
        let mut inner = self.inner.lock().expect("job_queue mutex poisoned");
        Self::trim_history(&mut inner, i64::MAX);
    }
}