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

//! DAG shape builders — every operation as a template over the shared
//! node primitives — plus submit-time cycle validation (petgraph is
//! confined to this validation; the runtime store stays the plain
//! `Vec<Node>` + `deps`).
//!
//! Every node carries its own `agent` (there is no DAG-level agent) — the
//! `node` helper stamps the template's agent onto each. `restart` takes an
//! agent *list* and stamps each agent onto its own subgraph, so a hive-wide
//! `hivectl restart` is ONE DAG with N independent per-agent subgraphs
//! (each a root chain, run concurrently on its own lease) rather than N
//! separate DAGs. The other templates are still single-agent.
//!
//! The power ops write the durable `wanted` intent via a head
//! `SetWanted(w)` node (not a pre-submit side effect); it holds the agent
//! lease so intent+reconcile is atomic per-agent.
//!
//! ```text
//! rebuild(a):        Prebuild(a) → StopForUpdate(a) → Swap(a) →(any) Reconcile(a)
//! graceful-stop(a):  SetWanted(a,Off) → Signal(a) → Drain(a) → Reconcile(a)
//! restart(a):        SetWanted(a,Up) → StopForUpdate(a) → Reconcile(a)
//! graceful-restart(a): SetWanted(a,Up) → Signal(a) → Drain(a) → StopForUpdate(a) → Reconcile(a)
//! start(a):          SetWanted(a,Up) → Reconcile(a)
//! stop(a):           SetWanted(a,Off) → Reconcile(a)
//! stale-start(a):    SetWanted(a,Up) → «rebuild subgraph»   (rev stale; prebuild noops, agent down)
//! spawn(a):          Provision(a) → Create(a) → WriteDropin(a) → Reconcile(a)  [wanted=Up at approve]
//! perm-change(a):    WritePermFile(a) → «rebuild subgraph»
//! meta-update(inp):  MetaLock(inp) → «fan-out rebuild(a) per affected a»
//! startup sweep:     MetaLock(hyperhive, non-fatal) → «fan-out rebuild(stale a)»
//! ```

use anyhow::{Result, bail};

use super::model::{DagSpec, Dep, DepWhen, NodeKind, NodeSpec, PermPayload, Source, Template};
use crate::coordinator::TransientKind;

/// After-ok edge on the previous node — the common chain link.
fn after_ok(on: u32) -> Vec<Dep> {
    vec![Dep {
        on,
        when: DepWhen::AfterOk,
    }]
}

/// Build one node targeting `agent`. The single place templates stamp a
/// node's agent, so a whole template is single-agent by passing the same
/// `agent` to every call.
fn node(agent: &str, kind: NodeKind, deps: Vec<Dep>) -> NodeSpec {
    NodeSpec {
        agent: agent.to_owned(),
        kind,
        deps,
    }
}

/// The rebuild node chain. `Reconcile` deps on `Swap` with `AfterAny`:
/// it must run even when the profile swap failed, so a previously-up
/// agent comes back on its old config (today's recovery-start). This
/// is the only `AfterAny` edge in v1.
fn rebuild_nodes(agent: &str, relock: bool, base: u32) -> Vec<NodeSpec> {
    vec![
        node(
            agent,
            NodeKind::Prebuild { relock },
            if base == 0 {
                Vec::new()
            } else {
                after_ok(base - 1)
            },
        ),
        node(agent, NodeKind::StopForUpdate, after_ok(base)),
        node(agent, NodeKind::Swap, after_ok(base + 1)),
        node(
            agent,
            NodeKind::Reconcile,
            vec![Dep {
                on: base + 2,
                when: DepWhen::AfterAny,
            }],
        ),
    ]
}

/// One uniform rebuild shape — no `was_running` branch. `StopForUpdate`
/// noops when already down; the tail `Reconcile` auto-noops the start
/// when `wanted = Offline` (a rebuild of a deliberately-stopped agent
/// leaves it stopped). `relock = false` only for meta-update cascade
/// children.
pub fn rebuild(
    agent: &str,
    source: Source,
    reason: String,
    parent_id: Option<u64>,
    relock: bool,
) -> DagSpec {
    DagSpec {
        template: Template::Rebuild,
        source,
        reason,
        parent_id,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Rebuilding),
        nodes: rebuild_nodes(agent, relock, 0),
    }
}

/// Approval-driven deploy (`ApplyCommit` / `MergeConfigPr`): the whole
/// two-phase pipeline stays one opaque node in v1 (design doc §9) —
/// wire-visible as a `rebuild` card like today.
pub fn approval_deploy(agent: &str, approval_id: i64, reason: String) -> DagSpec {
    DagSpec {
        template: Template::Rebuild,
        source: Source::Approval,
        reason,
        parent_id: None,
        approval_id: Some(approval_id),
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Rebuilding),
        nodes: vec![node(agent, NodeKind::ApprovalDeploy, Vec::new())],
    }
}

/// Graceful stop: cheap `Signal` fires immediately (no build slot), the
/// `Drain` awaits the harness checkpoint (bounded), and the tail
/// `Reconcile` performs the actual container stop. The head `SetWanted`
/// node writes `wanted = Offline` as part of the DAG (was a pre-submit
/// side effect). A whole-hive graceful stop therefore signals every agent
/// up front and overlaps every drain, replacing the old detached-watcher
/// thread structurally.
pub fn graceful_stop(agent: &str, source: Source, reason: String) -> DagSpec {
    DagSpec {
        template: Template::GracefulStop,
        source,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Stopping),
        nodes: vec![
            node(agent, NodeKind::SetWanted { up: false }, Vec::new()),
            node(agent, NodeKind::Signal, after_ok(0)),
            node(agent, NodeKind::Drain, after_ok(1)),
            node(agent, NodeKind::Reconcile, after_ok(2)),
        ],
    }
}

/// Restart one or more agents in a **single** DAG. Each agent gets an
/// independent subgraph — a head `SetWanted(Up)` (a root: no cross-agent
/// dep) then its restart chain — so all agents' restarts run concurrently,
/// each taking its own agent lease. `graceful` inserts `Signal → Drain`
/// before the mechanical `StopForUpdate` (each agent's subgraph is still one
/// atomic restart). One agent = the ordinary single-agent restart; many =
/// a hive-wide `hivectl restart` as one DAG instead of N separate ones.
pub fn restart(agents: &[String], graceful: bool, source: Source, reason: String) -> DagSpec {
    let mut nodes = Vec::new();
    for agent in agents {
        let base = u32::try_from(nodes.len()).unwrap_or(u32::MAX);
        // Head of this agent's subgraph — a root (empty deps), so the N
        // per-agent subgraphs are independent and run concurrently.
        nodes.push(node(agent, NodeKind::SetWanted { up: true }, Vec::new()));
        if graceful {
            nodes.push(node(agent, NodeKind::Signal, after_ok(base)));
            nodes.push(node(agent, NodeKind::Drain, after_ok(base + 1)));
            nodes.push(node(agent, NodeKind::StopForUpdate, after_ok(base + 2)));
            nodes.push(node(agent, NodeKind::Reconcile, after_ok(base + 3)));
        } else {
            nodes.push(node(agent, NodeKind::StopForUpdate, after_ok(base)));
            nodes.push(node(agent, NodeKind::Reconcile, after_ok(base + 1)));
        }
    }
    DagSpec {
        template: if graceful {
            Template::GracefulRestart
        } else {
            Template::Restart
        },
        source,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Restarting),
        nodes,
    }
}

/// Boot-time reconcile: a single `Reconcile` node that converges observed
/// power state to the persisted intent — `wanted` is untouched (no
/// `SetWanted`), unlike the operator `start`/`stop` templates. Used only
/// by the boot sweep now.
pub fn reconcile_only(
    template: Template,
    agent: &str,
    source: Source,
    reason: String,
    transient: Option<TransientKind>,
) -> DagSpec {
    DagSpec {
        template,
        source,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient,
        nodes: vec![node(agent, NodeKind::Reconcile, Vec::new())],
    }
}

/// Start: write `wanted = Up` (head `SetWanted`), then reconcile (which
/// starts the container). The intent write is a DAG node now, not a
/// pre-submit side effect.
pub fn start(agent: &str, source: Source, reason: String) -> DagSpec {
    DagSpec {
        template: Template::Start,
        source,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Starting),
        nodes: vec![
            node(agent, NodeKind::SetWanted { up: true }, Vec::new()),
            node(agent, NodeKind::Reconcile, after_ok(0)),
        ],
    }
}

/// Stop: write `wanted = Offline` (head `SetWanted`), then reconcile
/// (kill + unregister + `Killed` event).
pub fn stop(agent: &str, source: Source, reason: String) -> DagSpec {
    DagSpec {
        template: Template::Stop,
        source,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Stopping),
        nodes: vec![
            node(agent, NodeKind::SetWanted { up: false }, Vec::new()),
            node(agent, NodeKind::Reconcile, after_ok(0)),
        ],
    }
}

/// Stale start: a `start` whose rev marker is stale, so it rebuilds
/// before coming up — `SetWanted(Up)` → «rebuild subgraph» → `Reconcile`
/// (the rebuild's tail `Reconcile` starts it, since `wanted = Up`). The
/// rebuild nodes are the same `rebuild_nodes` chain a manual rebuild uses
/// (reused, not a variant); only the leading `SetWanted(Up)` intent
/// differs. Shows as a `Rebuild` on the dashboard like the old
/// submit-time upgrade did.
pub fn stale_start(agent: &str, source: Source, reason: String) -> DagSpec {
    let mut nodes = vec![node(agent, NodeKind::SetWanted { up: true }, Vec::new())];
    nodes.extend(rebuild_nodes(agent, true, 1));
    DagSpec {
        template: Template::Rebuild,
        source,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Rebuilding),
        nodes,
    }
}

/// First-deploy spawn (approval-driven): `Provision` (proposed/applied
/// repos, state subvolume, meta registration) then `Create`
/// (`nixos-container create`), drop-in write, then `Reconcile` starts
/// the container (`wanted = Up` written at approve time).
pub fn spawn(agent: &str, approval_id: i64, reason: String) -> DagSpec {
    DagSpec {
        template: Template::Spawn,
        source: Source::Approval,
        reason,
        parent_id: None,
        approval_id: Some(approval_id),
        inputs: Vec::new(),
        perm_payload: None,
        transient: Some(TransientKind::Spawning),
        nodes: vec![
            node(agent, NodeKind::Provision, Vec::new()),
            node(agent, NodeKind::Create, after_ok(0)),
            node(agent, NodeKind::WriteDropin, after_ok(1)),
            node(agent, NodeKind::Reconcile, after_ok(2)),
        ],
    }
}

/// Perm change: commit the JSON file(s), then the rebuild subgraph so
/// the updated `HIVE_TOOL_GROUPS` / `HIVE_CAPABILITIES` env var takes
/// effect in the container.
pub fn perm_change(agent: &str, source: Source, reason: String, payload: PermPayload) -> DagSpec {
    let mut nodes = vec![node(agent, NodeKind::WritePermFile, Vec::new())];
    nodes.extend(rebuild_nodes(agent, true, 1));
    DagSpec {
        template: Template::PermChange,
        source,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: Some(payload),
        transient: Some(TransientKind::Rebuilding),
        nodes,
    }
}

/// Meta-input lock bump. Child `Rebuild` DAGs fan out on completion —
/// appended *after* the bump lands so their prebuilds run against the
/// post-bump lock (and so a failed bump simply fans out nothing,
/// replacing the old pre-enqueue + `cancel_children` dance).
pub fn meta_update(
    inputs: Vec<String>,
    source: Source,
    reason: String,
    approval_id: Option<i64>,
) -> DagSpec {
    DagSpec {
        template: Template::MetaUpdate,
        source,
        reason,
        parent_id: None,
        approval_id,
        inputs,
        perm_payload: None,
        transient: None,
        nodes: vec![node(
            "hyperhive",
            NodeKind::MetaLock {
                sweep: false,
                fanout: None,
            },
            Vec::new(),
        )],
    }
}

/// Boot-time root anchor DAG: a single [`NodeKind::Noop`] node that groups
/// this boot's `StartupSweep` + per-agent `Reconcile` child DAGs (linked via
/// `parent_id`) into one tree so the dashboard renders the boot as one entry.
/// Holds no lease and does no work — the children it anchors still run
/// concurrently. `auto_update::run` submits it first (when there's any boot
/// work), then parents the sweep + reconciles onto its id.
pub fn boot_root(reason: String) -> DagSpec {
    DagSpec {
        template: Template::Boot,
        source: Source::AutoUpdate,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: None,
        nodes: vec![node("hyperhive", NodeKind::Noop, Vec::new())],
    }
}

/// Boot-time sweep parent: bump meta's hyperhive input (non-fatal),
/// then fan out `Rebuild` children for the precomputed stale agent
/// list (topology-sorted by the caller).
pub fn startup_sweep(reason: String, stale_agents: Vec<String>) -> DagSpec {
    DagSpec {
        template: Template::StartupSweep,
        source: Source::AutoUpdate,
        reason,
        parent_id: None,
        approval_id: None,
        inputs: Vec::new(),
        perm_payload: None,
        transient: None,
        nodes: vec![node(
            "hyperhive",
            NodeKind::MetaLock {
                sweep: true,
                fanout: Some(stale_agents),
            },
            Vec::new(),
        )],
    }
}

/// Validate a spec before it enters the queue: node ids are dense
/// (index = id), deps reference existing nodes, and the dep graph is
/// acyclic (petgraph `toposort`). Rejecting cycles here fixes the old
/// queue's documented "circular dep silently deadlocks forever" caveat.
pub fn validate(spec: &DagSpec) -> Result<()> {
    if spec.nodes.is_empty() {
        bail!("dag spec {:?} has no nodes", spec.template);
    }
    let n = spec.nodes.len();
    let mut graph = petgraph::graph::DiGraph::<u32, ()>::new();
    let idx: Vec<_> = (0..n)
        .map(|i| graph.add_node(u32::try_from(i).unwrap_or(u32::MAX)))
        .collect();
    for (i, node) in spec.nodes.iter().enumerate() {
        for dep in &node.deps {
            let Some(&dep_idx) = idx.get(dep.on as usize) else {
                bail!(
                    "dag spec {:?} node {i} depends on unknown node {}",
                    spec.template,
                    dep.on
                );
            };
            graph.add_edge(dep_idx, idx[i], ());
        }
    }
    if petgraph::algo::toposort(&graph, None).is_err() {
        bail!("dag spec {:?} contains a dependency cycle", spec.template);
    }
    Ok(())
}