forge-orchestration 0.6.0

//! Reconciliation control loop: converge desired state onto actual state.
//!
//! Borg's value is not the placement *decision* — it is the *control loop* that
//! keeps reality matching intent: schedule what is missing, persist the binding,
//! restart what failed, reschedule off lost nodes, and execute autoscaling. The
//! base Forge runtime had none of this: `Forge::run()` loaded state once, served
//! HTTP, and saved on shutdown; nothing ever scheduled a workload.
//!
//! [`Reconciler`] is that loop. It is deliberately *self-contained* and wires
//! into the runtime through the **shared [`StateStore`]** rather than by taking
//! ownership of the runtime's internals:
//!
//! - **Desired** state is the set of [`Job`]s persisted under [`keys::JOBS`]
//!   (exactly what [`crate::runtime::Forge::submit_job`] already writes). The
//!   reconciler reloads them each tick, so newly submitted jobs are picked up
//!   with no shared-mutable-state plumbing.
//! - **Actual** state is the [`Assignment`] map, persisted under
//!   [`keys::ASSIGNMENTS`]. It is the single source of truth; in-memory node
//!   capacity is *derived* from it and rebuilt on [`Reconciler::bootstrap`].
//!
//! ## Correctness properties
//!
//! 1. **No double-scheduling / no capacity leak.** Everything is keyed by the
//!    deterministic workload id `"{job}/{group}/{ordinal}"`. A replica is either
//!    in `bound` (placed) or not; the backoff table is keyed by the same id, so a
//!    failed replica can never be scheduled twice in one pass.
//! 2. **Capacity is always derivable from persisted truth.** A crash mid-tick
//!    loses only unpersisted bindings; on restart [`Reconciler::bootstrap`]
//!    replays the persisted assignments to rebuild node capacity, and any binding
//!    that never got persisted is simply re-derived as missing and rescheduled.
//!    There is no path that double-allocates.
//! 3. **GPU device ids are freed exactly.** [`NodeResources::allocate`] picks GPU
//!    device ids internally; we capture *which* ids it took via an allocate-time
//!    snapshot diff and store them on the [`Assignment`], so
//!    [`NodeResources::release`] frees precisely those GPUs (no leak).
//! 4. **Lost nodes converge.** Removing a node drops its assignments from
//!    `bound`; the next pass re-derives them as missing and reschedules them onto
//!    surviving nodes.
//! 5. **No lock held across `.await`.** Jobs are cloned into a `Vec` before any
//!    async autoscaler/store call; node mutation happens in a later synchronous
//!    phase. The reconciler owns plain `HashMap`s (no `DashMap` guard escapes).

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;
use tracing::{debug, info, warn};

use super::algorithms::{BinPackScheduler, SchedulingAlgorithm};
use super::gang::{GangReservation, GangScheduler};
use super::sim::SimCell;
use super::{NodeResources, ResourceRequirements, Workload};
use crate::autoscaler::{Autoscaler, MetricsSnapshot, ScalingDecision};
use crate::error::Result;
use crate::job::{Job, TaskGroup};
use crate::storage::{keys, store_get_json, store_set_json, BoxedStateStore};
use crate::types::NodeId;

/// Observed lifecycle status of a placed replica.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TaskStatus {
    /// Scheduled and bound, not yet confirmed running.
    Pending,
    /// Confirmed running.
    Running,
    /// Failed; will be released and rescheduled.
    Failed,
}

/// A committed placement of one replica (the unit of *actual* state).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Assignment {
    /// Deterministic workload id: `"{job}/{group}/{ordinal}"`.
    pub workload_id: String,
    /// Owning job id.
    pub job_id: String,
    /// Group name within the job.
    pub group: String,
    /// Replica ordinal within the group.
    pub ordinal: u32,
    /// Node this replica is bound to.
    pub node: NodeId,
    /// GPU device ids reserved for this replica (for exact release).
    pub gpu_ids: Vec<u32>,
    /// Resources reserved (kept so release is exact even if the job changes).
    pub resources: ResourceRequirements,
    /// Observed status.
    pub status: TaskStatus,
}

/// A committed gang placement of a [`SimCell`] (world + agents), with the
/// per-node reservations needed to release it exactly on removal.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimBinding {
    /// Owning sim cell id.
    pub cell_id: String,
    /// `(member_workload_id, node_id)` for every member.
    pub placements: Vec<(String, NodeId)>,
    /// Per-node reservations (exact CPU/mem + GPU device ids).
    pub reservations: Vec<GangReservation>,
}

/// Backoff bookkeeping for a replica that could not be placed.
#[derive(Debug, Clone)]
struct BackoffState {
    attempts: u32,
    next_try: Instant,
}

/// Source of per-job utilization metrics that drive autoscaling.
///
/// Kept as an injectable hook so the reconciler stays decoupled from any
/// particular telemetry pipeline (and so tests are deterministic). Returns
/// `(cpu_util, mem_util)` in `0.0..=1.0` for a job, or `None` if unknown (in
/// which case that job is not autoscaled this pass).
pub trait MetricsSource: Send + Sync {
    /// Current `(cpu_util, mem_util)` for `job_id`, if known.
    fn job_metrics(&self, job_id: &str) -> Option<(f64, f64)>;
}

/// Summary of what a single reconcile pass did (for logging and tests).
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ReconcileReport {
    /// Replicas newly scheduled and bound this pass.
    pub scheduled: usize,
    /// Replicas released (scaled down, job removed, node lost, or failed).
    pub released: usize,
    /// Desired replicas still unplaced (no capacity / backing off).
    pub pending: usize,
    /// Groups whose desired count was changed by autoscaling.
    pub rescaled: usize,
    /// Sim cells newly gang-scheduled this pass.
    pub sim_scheduled: usize,
    /// Sim cells released (removed, or a reservation's node vanished).
    pub sim_released: usize,
}

/// Drives convergence of desired [`Job`]s onto actual [`Assignment`]s.
pub struct Reconciler {
    store: BoxedStateStore,
    autoscaler: Arc<Autoscaler>,
    algorithm: Box<dyn SchedulingAlgorithm>,
    /// Gang scheduler for sim cells (all-or-nothing co-placement).
    gang: GangScheduler,
    metrics_source: Option<Arc<dyn MetricsSource>>,
    /// Registered nodes with live (in-memory) capacity.
    nodes: HashMap<NodeId, NodeResources>,
    /// Actual state: workload_id -> assignment.
    bound: HashMap<String, Assignment>,
    /// Actual state: sim cell id -> gang binding.
    sim_bound: HashMap<String, SimBinding>,
    /// Backoff for replicas that failed to place.
    backoff: HashMap<String, BackoffState>,
    /// Backoff for sim cells whose gang failed to place.
    sim_backoff: HashMap<String, BackoffState>,
    /// Loop interval.
    interval: Duration,
    /// Maximum backoff delay.
    max_backoff: Duration,
}

impl Reconciler {
    /// Create a reconciler over a shared store and autoscaler. Uses
    /// [`BinPackScheduler`] for node scoring by default.
    pub fn new(store: BoxedStateStore, autoscaler: Arc<Autoscaler>) -> Self {
        Self {
            store,
            autoscaler,
            algorithm: Box::new(BinPackScheduler::new()),
            gang: GangScheduler::new(),
            metrics_source: None,
            nodes: HashMap::new(),
            bound: HashMap::new(),
            sim_bound: HashMap::new(),
            backoff: HashMap::new(),
            sim_backoff: HashMap::new(),
            interval: Duration::from_secs(5),
            max_backoff: Duration::from_secs(300),
        }
    }

    /// Use a custom node-scoring algorithm.
    pub fn with_algorithm<A: SchedulingAlgorithm + 'static>(mut self, algorithm: A) -> Self {
        self.algorithm = Box::new(algorithm);
        self
    }

    /// Provide a metrics source to enable autoscaling.
    pub fn with_metrics_source(mut self, source: Arc<dyn MetricsSource>) -> Self {
        self.metrics_source = Some(source);
        self
    }

    /// Set the reconcile loop interval.
    pub fn with_interval(mut self, interval: Duration) -> Self {
        self.interval = interval;
        self
    }

    /// Register (or replace) a schedulable node with fresh capacity.
    pub fn register_node(&mut self, node: NodeResources) {
        self.nodes.insert(node.node_id, node);
    }

    /// Remove a node (e.g. it left the cluster). Its assignments are dropped from
    /// actual state so the next pass reschedules them elsewhere.
    pub fn remove_node(&mut self, node_id: &NodeId) {
        self.nodes.remove(node_id);
    }

    /// Number of currently-bound replicas.
    pub fn bound_count(&self) -> usize {
        self.bound.len()
    }

    /// Snapshot of current assignments (for inspection / handlers).
    pub fn assignments(&self) -> Vec<Assignment> {
        self.bound.values().cloned().collect()
    }

    /// Number of currently-bound sim cells.
    pub fn sim_bound_count(&self) -> usize {
        self.sim_bound.len()
    }

    /// Snapshot of current sim-cell gang bindings.
    pub fn sim_bindings(&self) -> Vec<SimBinding> {
        self.sim_bound.values().cloned().collect()
    }

    /// Load persisted actual state and rebuild node capacity from it.
    ///
    /// Call once after registering nodes and before the loop starts. Registered
    /// nodes must already have fresh (unallocated) capacity; this replays each
    /// persisted assignment onto its node so the in-memory counters match the
    /// persisted truth. Assignments whose node is no longer registered are
    /// dropped (they will be rescheduled).
    pub async fn bootstrap(&mut self) -> Result<()> {
        let persisted: HashMap<String, Assignment> =
            store_get_json(self.store.as_ref(), keys::ASSIGNMENTS)
                .await?
                .unwrap_or_default();

        let mut rebuilt = HashMap::new();
        for (id, mut a) in persisted {
            match self.nodes.get_mut(&a.node) {
                Some(node) => {
                    let before: Vec<u32> = node.gpus_allocated.clone();
                    if node.allocate(&a.resources) {
                        // Re-derive the actual GPU ids taken on replay so future
                        // releases stay exact even across restarts.
                        a.gpu_ids = node
                            .gpus_allocated
                            .iter()
                            .filter(|d| !before.contains(d))
                            .copied()
                            .collect();
                        rebuilt.insert(id, a);
                    } else {
                        warn!(workload = %id, "assignment no longer fits on replay; dropping");
                    }
                }
                None => {
                    debug!(workload = %id, "assignment's node is gone; dropping for reschedule");
                }
            }
        }

        let count = rebuilt.len();
        self.bound = rebuilt;

        // Replay persisted sim-cell gang bindings, rebuilding node capacity from
        // their reservations. A binding that cannot be fully replayed (node gone
        // or no longer fits) is rolled back and dropped so it reschedules.
        let persisted_sim: HashMap<String, SimBinding> =
            store_get_json(self.store.as_ref(), keys::SIM_ASSIGNMENTS)
                .await?
                .unwrap_or_default();
        let mut rebuilt_sim = HashMap::new();
        for (cid, mut binding) in persisted_sim {
            let mut applied: Vec<(NodeId, ResourceRequirements, Vec<u32>)> = Vec::new();
            let mut ok = true;
            for r in &binding.reservations {
                match self.nodes.get_mut(&r.node) {
                    Some(node) => {
                        let before: Vec<u32> = node.gpus_allocated.clone();
                        if node.allocate(&r.resources) {
                            let gpu_ids = node
                                .gpus_allocated
                                .iter()
                                .filter(|d| !before.contains(d))
                                .copied()
                                .collect::<Vec<_>>();
                            applied.push((r.node, r.resources.clone(), gpu_ids));
                        } else {
                            ok = false;
                            break;
                        }
                    }
                    None => {
                        ok = false;
                        break;
                    }
                }
            }
            if ok {
                for (res, (_, _, gpu_ids)) in binding.reservations.iter_mut().zip(applied.iter()) {
                    res.gpu_ids = gpu_ids.clone();
                }
                rebuilt_sim.insert(cid, binding);
            } else {
                for (node_id, req, gpu_ids) in &applied {
                    if let Some(node) = self.nodes.get_mut(node_id) {
                        node.release(req, gpu_ids);
                    }
                }
                debug!(cell = %cid, "sim binding could not be fully replayed; dropping for reschedule");
            }
        }
        let sim_count = rebuilt_sim.len();
        self.sim_bound = rebuilt_sim;

        // Persist the cleaned-up view so a node that vanished while we were down
        // does not leave stale assignments behind.
        self.persist().await?;
        self.persist_sim().await?;
        info!(restored = count, sim_restored = sim_count, "reconciler bootstrapped from store");
        Ok(())
    }

    /// Run the reconcile loop until a shutdown signal is received.
    pub async fn run(mut self, mut shutdown_rx: broadcast::Receiver<()>) -> Result<()> {
        let mut ticker = tokio::time::interval(self.interval);
        info!(interval_secs = self.interval.as_secs(), "reconcile loop started");
        loop {
            tokio::select! {
                _ = ticker.tick() => {
                    match self.reconcile_once().await {
                        Ok(report) => {
                            if report != ReconcileReport::default() {
                                debug!(?report, "reconcile pass");
                            }
                        }
                        Err(e) => warn!(error = %e, "reconcile pass failed"),
                    }
                }
                _ = shutdown_rx.recv() => {
                    info!("reconcile loop stopping; persisting final state");
                    let _ = self.persist().await;
                    return Ok(());
                }
            }
        }
    }

    /// Execute a single convergence pass. Public so it can be driven manually
    /// (and unit-tested) without spawning the loop.
    pub async fn reconcile_once(&mut self) -> Result<ReconcileReport> {
        let mut report = ReconcileReport::default();

        // ---- Phase 0: load desired state (clone out; no store guard held) ----
        let mut jobs = self.load_jobs().await?;

        // ---- Phase 1: autoscaling (async; mutates desired counts) ------------
        // Done before deriving the desired set so a scale decision takes effect
        // this same pass. Holds no node borrow across the await.
        report.rescaled = self.autoscale(&mut jobs).await?;

        // ---- Phase 2: derive the desired replica set -------------------------
        let desired = expand_jobs(&jobs);
        let desired_ids: std::collections::HashSet<&String> = desired.keys().collect();

        // ---- Phase 3: GC — release bound replicas that are no longer desired
        // or whose node has vanished. ------------------------------------------
        let to_release: Vec<String> = self
            .bound
            .keys()
            .filter(|id| {
                !desired_ids.contains(*id) || !self.nodes.contains_key(&self.bound[*id].node)
            })
            .cloned()
            .collect();
        for id in to_release {
            if let Some(a) = self.bound.remove(&id) {
                // Only release capacity if the node still exists; a vanished node
                // takes its capacity with it.
                if self.nodes.contains_key(&a.node) {
                    Self::do_release(&mut self.nodes, &a);
                }
                self.backoff.remove(&id);
                report.released += 1;
            }
        }

        // ---- Phase 4: schedule desired replicas that are not bound -----------
        let now = Instant::now();
        // Deterministic order so packing is stable and tests are reproducible.
        let mut pending_ids: Vec<&String> = desired
            .keys()
            .filter(|id| !self.bound.contains_key(*id))
            .collect();
        pending_ids.sort();

        for id in pending_ids {
            // Respect backoff.
            if let Some(b) = self.backoff.get(id) {
                if now < b.next_try {
                    report.pending += 1;
                    continue;
                }
            }

            let spec = &desired[id];
            match Self::try_place(&mut self.nodes, self.algorithm.as_ref(), &spec.workload) {
                Some((node_id, gpu_ids)) => {
                    let assignment = Assignment {
                        workload_id: spec.workload.id.clone(),
                        job_id: spec.job_id.clone(),
                        group: spec.group.clone(),
                        ordinal: spec.ordinal,
                        node: node_id,
                        gpu_ids,
                        resources: spec.workload.resources.clone(),
                        status: TaskStatus::Pending,
                    };
                    self.bound.insert(id.clone(), assignment);
                    self.backoff.remove(id);
                    report.scheduled += 1;
                }
                None => {
                    self.bump_backoff(id, now);
                    report.pending += 1;
                }
            }
        }

        // ---- Phase 5: persist actual state (durable truth) -------------------
        if report.scheduled > 0 || report.released > 0 {
            self.persist().await?;
        }

        // ---- Phase 6: sim cells — gang-scheduled, deadline-ordered -----------
        self.reconcile_sim_cells(now, &mut report).await?;

        Ok(report)
    }

    /// Converge desired [`SimCell`]s (loaded from the store) onto gang
    /// placements: release cells no longer desired (or whose node vanished),
    /// then gang-schedule the rest most-urgent-tick-first, with backoff. This is
    /// the SimCell -> gang -> reconcile path driven by the live loop.
    async fn reconcile_sim_cells(
        &mut self,
        now: Instant,
        report: &mut ReconcileReport,
    ) -> Result<()> {
        let cells = self.load_sim_cells().await?;
        let desired: std::collections::HashSet<&String> = cells.iter().map(|c| &c.id).collect();

        // GC: release bindings no longer desired, or with a vanished node.
        let to_release: Vec<String> = self
            .sim_bound
            .keys()
            .filter(|cid| {
                !desired.contains(*cid)
                    || self.sim_bound[*cid]
                        .reservations
                        .iter()
                        .any(|r| !self.nodes.contains_key(&r.node))
            })
            .cloned()
            .collect();
        for cid in to_release {
            if let Some(binding) = self.sim_bound.remove(&cid) {
                for r in &binding.reservations {
                    if let Some(node) = self.nodes.get_mut(&r.node) {
                        node.release(&r.resources, &r.gpu_ids);
                    }
                }
                self.sim_backoff.remove(&cid);
                report.sim_released += 1;
            }
        }

        // Schedule unbound cells, most-urgent tick deadline first.
        let mut pending: Vec<&SimCell> = cells
            .iter()
            .filter(|c| !self.sim_bound.contains_key(&c.id))
            .collect();
        pending.sort_by(|a, b| a.next_deadline.cmp(&b.next_deadline).then(a.id.cmp(&b.id)));

        for cell in pending {
            if let Some(b) = self.sim_backoff.get(&cell.id) {
                if now < b.next_try {
                    continue;
                }
            }
            let group = cell.gang_group();
            let decision = self.gang.schedule_gang(&group, &mut self.nodes);
            if decision.committed {
                self.sim_bound.insert(
                    cell.id.clone(),
                    SimBinding {
                        cell_id: cell.id.clone(),
                        placements: decision.placements,
                        reservations: decision.reservations,
                    },
                );
                self.sim_backoff.remove(&cell.id);
                report.sim_scheduled += 1;
            } else {
                Self::backoff_bump(&mut self.sim_backoff, &cell.id, now, self.max_backoff);
            }
        }

        if report.sim_scheduled > 0 || report.sim_released > 0 {
            self.persist_sim().await?;
        }
        Ok(())
    }

    /// Mark a running replica as failed; it is released and rescheduled next pass.
    pub fn mark_failed(&mut self, workload_id: &str) {
        if let Some(a) = self.bound.remove(workload_id) {
            if self.nodes.contains_key(&a.node) {
                Self::do_release(&mut self.nodes, &a);
            }
        }
    }

    // ---- internals --------------------------------------------------------

    async fn load_jobs(&self) -> Result<Vec<Job>> {
        let mut jobs = Vec::new();
        for key in self.store.list_prefix(keys::JOBS).await? {
            if let Some(job) = store_get_json::<Job>(self.store.as_ref(), &key).await? {
                jobs.push(job);
            }
        }
        Ok(jobs)
    }

    /// Evaluate autoscaling per autoscalable group; mutate desired counts in
    /// `jobs` and persist changed jobs. Returns the number of groups rescaled.
    async fn autoscale(&self, jobs: &mut [Job]) -> Result<usize> {
        let Some(source) = self.metrics_source.clone() else {
            return Ok(0);
        };
        let mut rescaled = 0;
        for job in jobs.iter_mut() {
            let Some((cpu, mem)) = source.job_metrics(&job.id) else {
                continue;
            };
            let mut changed = false;
            for group in &mut job.groups {
                // Only groups with headroom to move are autoscalable.
                if group.scaling.min >= group.scaling.max {
                    continue;
                }
                let key = format!("{}/{}", job.id, group.name);
                let snap = MetricsSnapshot::new(cpu, mem, group.scaling.desired);
                // Autoscaler::evaluate enforces hysteresis internally and only
                // records the cooldown when it actually returns a scaling action.
                let decision = self.autoscaler.evaluate(&key, snap).await;
                let new_desired = match decision {
                    ScalingDecision::ScaleUp(n) => {
                        (group.scaling.desired + n).min(group.scaling.max)
                    }
                    ScalingDecision::ScaleDown(n) => group
                        .scaling
                        .desired
                        .saturating_sub(n)
                        .max(group.scaling.min),
                    ScalingDecision::ScaleTo(c) => c.clamp(group.scaling.min, group.scaling.max),
                    ScalingDecision::NoChange => group.scaling.desired,
                };
                if new_desired != group.scaling.desired {
                    info!(job = %job.id, group = %group.name, from = group.scaling.desired, to = new_desired, "autoscale");
                    group.scaling.desired = new_desired;
                    changed = true;
                    rescaled += 1;
                }
            }
            if changed {
                // Persist the new desired so it is durable and visible to the
                // runtime's own view of the job on its next read.
                store_set_json(self.store.as_ref(), &keys::job(&job.id), job).await?;
            }
        }
        Ok(rescaled)
    }

    async fn persist(&self) -> Result<()> {
        store_set_json(self.store.as_ref(), keys::ASSIGNMENTS, &self.bound).await
    }

    async fn persist_sim(&self) -> Result<()> {
        store_set_json(self.store.as_ref(), keys::SIM_ASSIGNMENTS, &self.sim_bound).await
    }

    /// Load desired sim cells persisted under [`keys::SIMCELLS`].
    async fn load_sim_cells(&self) -> Result<Vec<SimCell>> {
        let mut cells = Vec::new();
        for key in self.store.list_prefix(keys::SIMCELLS).await? {
            if let Some(cell) = store_get_json::<SimCell>(self.store.as_ref(), &key).await? {
                cells.push(cell);
            }
        }
        Ok(cells)
    }

    /// Exponential backoff bump on an arbitrary backoff map (jobs or sim cells).
    fn backoff_bump(
        map: &mut HashMap<String, BackoffState>,
        id: &str,
        now: Instant,
        max_backoff: Duration,
    ) {
        let entry = map
            .entry(id.to_string())
            .or_insert(BackoffState { attempts: 0, next_try: now });
        entry.attempts = entry.attempts.saturating_add(1);
        let secs = 1u64.checked_shl(entry.attempts.min(20)).unwrap_or(u64::MAX);
        entry.next_try = now + Duration::from_secs(secs).min(max_backoff);
    }

    fn bump_backoff(&mut self, id: &str, now: Instant) {
        let entry = self
            .backoff
            .entry(id.to_string())
            .or_insert(BackoffState { attempts: 0, next_try: now });
        entry.attempts = entry.attempts.saturating_add(1);
        // Exponential: 1s, 2s, 4s, ... capped at max_backoff.
        let secs = 1u64.checked_shl(entry.attempts.min(20)).unwrap_or(u64::MAX);
        let delay = Duration::from_secs(secs).min(self.max_backoff);
        entry.next_try = now + delay;
    }

    /// Place a workload on the best-scoring node that fits, committing the
    /// allocation. Returns the node and the GPU device ids reserved (captured via
    /// an allocate-time snapshot diff so release is exact).
    fn try_place(
        nodes: &mut HashMap<NodeId, NodeResources>,
        algorithm: &dyn SchedulingAlgorithm,
        workload: &Workload,
    ) -> Option<(NodeId, Vec<u32>)> {
        let req = &workload.resources;
        let best_id = nodes
            .values()
            .filter(|n| n.can_fit(req))
            .max_by(|a, b| {
                algorithm
                    .score(workload, a)
                    .partial_cmp(&algorithm.score(workload, b))
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|n| n.node_id)?;

        let node = nodes.get_mut(&best_id)?;
        let before: Vec<u32> = node.gpus_allocated.clone();
        if node.allocate(req) {
            let gpu_ids = node
                .gpus_allocated
                .iter()
                .filter(|d| !before.contains(d))
                .copied()
                .collect();
            Some((best_id, gpu_ids))
        } else {
            None
        }
    }

    fn do_release(nodes: &mut HashMap<NodeId, NodeResources>, a: &Assignment) {
        if let Some(node) = nodes.get_mut(&a.node) {
            node.release(&a.resources, &a.gpu_ids);
        }
    }
}

/// A desired replica derived from a [`Job`].
struct DesiredReplica {
    workload: Workload,
    job_id: String,
    group: String,
    ordinal: u32,
}

/// Expand all jobs into the desired replica set, keyed by deterministic id.
fn expand_jobs(jobs: &[Job]) -> HashMap<String, DesiredReplica> {
    let mut out = HashMap::new();
    for job in jobs {
        for group in &job.groups {
            let resources = group_resources(group);
            for ordinal in 0..group.scaling.desired {
                let id = format!("{}/{}/{}", job.id, group.name, ordinal);
                let workload = Workload::new(id.clone(), format!("{}-{}", job.name, group.name))
                    .with_resources(resources.clone())
                    .with_priority(job.priority as i32);
                out.insert(
                    id,
                    DesiredReplica {
                        workload,
                        job_id: job.id.clone(),
                        group: group.name.clone(),
                        ordinal,
                    },
                );
            }
        }
    }
    out
}

/// Sum a group's task resources into the scheduler's [`ResourceRequirements`].
///
/// Forge `Resources` are in MHz / MB; the scheduler bins in millicores / MB.
/// The units only need to be self-consistent across nodes and workloads (they
/// are: `NodeResources` capacity is expressed in the same scale), so we map
/// MHz -> the scheduler's CPU unit 1:1.
fn group_resources(group: &TaskGroup) -> ResourceRequirements {
    let mut cpu: u64 = 0;
    let mut memory: u64 = 0;
    let mut gpu: u32 = 0;
    for task in &group.tasks {
        cpu += task.resources.cpu as u64;
        memory += task.resources.memory as u64;
        gpu += task.resources.gpu.unwrap_or(0);
    }
    let mut req = ResourceRequirements::new().cpu(cpu).memory(memory);
    if gpu > 0 {
        // Per-GPU memory is not modeled on Forge `Resources`; leave at 0 so any
        // GPU satisfies the count constraint until a richer model is added.
        req = req.gpu(gpu, 0);
    }
    req
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::autoscaler::AutoscalerConfig;
    use crate::job::{Job, Task};
    use crate::storage::MemoryStore;
    use crate::types::{GpuResources, NodeId};

    fn store() -> BoxedStateStore {
        Arc::new(MemoryStore::new())
    }

    fn autoscaler() -> Arc<Autoscaler> {
        Arc::new(Autoscaler::new(AutoscalerConfig::default().hysteresis_secs(0)).unwrap())
    }

    async fn submit(store: &BoxedStateStore, job: &Job) {
        store_set_json(store.as_ref(), &keys::job(&job.id), job)
            .await
            .unwrap();
    }

    fn job_with_desired(name: &str, group: &str, cpu: u32, mem: u32, desired: u32) -> Job {
        let mut job = Job::new(name).with_group(group, Task::new("t").resources(cpu, mem));
        job.groups[0].scaling = crate::job::ScalingConfig::new(1, 10).with_desired(desired);
        job
    }

    #[tokio::test]
    async fn schedules_all_desired_replicas() {
        let store = store();
        let job = job_with_desired("svc", "api", 1000, 1024, 3);
        submit(&store, &job).await;

        let mut rec = Reconciler::new(store, autoscaler());
        rec.register_node(NodeResources::new(NodeId::new(), 8000, 8192));

        let report = rec.reconcile_once().await.unwrap();
        assert_eq!(report.scheduled, 3);
        assert_eq!(rec.bound_count(), 3);

        // Idempotent: a second pass changes nothing.
        let report2 = rec.reconcile_once().await.unwrap();
        assert_eq!(report2.scheduled, 0);
        assert_eq!(report2.released, 0);
        assert_eq!(rec.bound_count(), 3);
    }

    #[tokio::test]
    async fn scale_down_releases_and_restores_capacity() {
        let store = store();
        let mut job = job_with_desired("svc", "api", 1000, 1024, 3);
        submit(&store, &job).await;

        let mut rec = Reconciler::new(store.clone(), autoscaler());
        let node_id = NodeId::new();
        rec.register_node(NodeResources::new(node_id, 8000, 8192));
        rec.reconcile_once().await.unwrap();
        assert_eq!(rec.bound_count(), 3);
        let used_at_3 = rec.nodes[&node_id].cpu_allocated;
        assert_eq!(used_at_3, 3000);

        // Scale the job down to 1 and persist.
        job.groups[0].scaling.desired = 1;
        submit(&store, &job).await;

        let report = rec.reconcile_once().await.unwrap();
        assert_eq!(report.released, 2);
        assert_eq!(rec.bound_count(), 1);
        assert_eq!(rec.nodes[&node_id].cpu_allocated, 1000);
    }

    #[tokio::test]
    async fn node_loss_reschedules_onto_survivor() {
        let store = store();
        let job = job_with_desired("svc", "api", 1000, 1024, 2);
        submit(&store, &job).await;

        let mut rec = Reconciler::new(store, autoscaler());
        let dead = NodeId::new();
        rec.register_node(NodeResources::new(dead, 8000, 8192));
        rec.reconcile_once().await.unwrap();
        assert!(rec.assignments().iter().all(|a| a.node == dead));

        // The node leaves; a fresh one joins.
        let survivor = NodeId::new();
        rec.remove_node(&dead);
        rec.register_node(NodeResources::new(survivor, 8000, 8192));

        let report = rec.reconcile_once().await.unwrap();
        // Both replicas released off the dead node and rescheduled onto survivor.
        assert_eq!(rec.bound_count(), 2);
        assert!(rec.assignments().iter().all(|a| a.node == survivor));
        assert!(report.scheduled >= 2);
    }

    #[tokio::test]
    async fn insufficient_capacity_backs_off_without_overallocating() {
        let store = store();
        // 4 replicas of 1000 cpu, node only fits 2.
        let job = job_with_desired("svc", "api", 1000, 1024, 4);
        submit(&store, &job).await;

        let mut rec = Reconciler::new(store, autoscaler());
        let node_id = NodeId::new();
        rec.register_node(NodeResources::new(node_id, 2000, 8192));

        let report = rec.reconcile_once().await.unwrap();
        assert_eq!(report.scheduled, 2);
        assert_eq!(report.pending, 2);
        assert_eq!(rec.bound_count(), 2);
        // Never allocated beyond capacity.
        assert!(rec.nodes[&node_id].cpu_allocated <= 2000);
    }

    #[tokio::test]
    async fn gpu_ids_are_freed_exactly_on_release() {
        let store = store();
        let mut job = Job::new("train").with_group(
            "worker",
            Task::new("w").with_resources(crate::job::Resources::new(500, 1024).with_gpu(1)),
        );
        job.groups[0].scaling = crate::job::ScalingConfig::new(1, 10).with_desired(2);
        submit(&store, &job).await;

        let mut rec = Reconciler::new(store.clone(), autoscaler());
        let node_id = NodeId::new();
        let node = NodeResources::new(node_id, 8000, 8192)
            .with_gpu(GpuResources::new(0, "A100", 40960))
            .with_gpu(GpuResources::new(1, "A100", 40960));
        rec.register_node(node);

        rec.reconcile_once().await.unwrap();
        assert_eq!(rec.bound_count(), 2);
        assert_eq!(rec.nodes[&node_id].gpus_allocated.len(), 2);
        // Each assignment recorded exactly one distinct GPU id.
        let mut ids: Vec<u32> = rec.assignments().iter().flat_map(|a| a.gpu_ids.clone()).collect();
        ids.sort();
        assert_eq!(ids, vec![0, 1]);

        // Scale to 0 and confirm GPUs are returned, not leaked.
        job.groups[0].scaling.desired = 0;
        submit(&store, &job).await;
        rec.reconcile_once().await.unwrap();
        assert_eq!(rec.bound_count(), 0);
        assert_eq!(rec.nodes[&node_id].gpus_allocated.len(), 0);
    }

    #[tokio::test]
    async fn bootstrap_rebuilds_actual_state_and_capacity() {
        let store = store();
        let job = job_with_desired("svc", "api", 1000, 1024, 3);
        submit(&store, &job).await;

        let node_id = NodeId::new();
        {
            let mut rec = Reconciler::new(store.clone(), autoscaler());
            rec.register_node(NodeResources::new(node_id, 8000, 8192));
            rec.reconcile_once().await.unwrap();
            assert_eq!(rec.bound_count(), 3);
        }

        // A fresh reconciler with a fresh (empty-capacity) node rebuilds from the
        // persisted assignments.
        let mut rec2 = Reconciler::new(store, autoscaler());
        rec2.register_node(NodeResources::new(node_id, 8000, 8192));
        rec2.bootstrap().await.unwrap();
        assert_eq!(rec2.bound_count(), 3);
        assert_eq!(rec2.nodes[&node_id].cpu_allocated, 3000);
    }

    struct FixedMetrics(f64, f64);
    impl MetricsSource for FixedMetrics {
        fn job_metrics(&self, _job_id: &str) -> Option<(f64, f64)> {
            Some((self.0, self.1))
        }
    }

    #[tokio::test]
    async fn autoscaling_increases_desired_and_persists() {
        let store = store();
        let job = job_with_desired("svc", "api", 500, 512, 2);
        let job_id = job.id.clone();
        submit(&store, &job).await;

        let mut rec = Reconciler::new(store.clone(), autoscaler())
            .with_metrics_source(Arc::new(FixedMetrics(0.95, 0.95)));
        rec.register_node(NodeResources::new(NodeId::new(), 16000, 32768));

        let report = rec.reconcile_once().await.unwrap();
        assert_eq!(report.rescaled, 1, "high utilization should scale up");

        // Desired was persisted back to the job, and the extra replica is bound.
        let reloaded: Job = store_get_json(store.as_ref(), &keys::job(&job_id))
            .await
            .unwrap()
            .unwrap();
        assert_eq!(reloaded.groups[0].scaling.desired, 3);
        assert_eq!(rec.bound_count(), 3);
    }

    #[tokio::test]
    async fn schedules_and_releases_sim_cells_via_gang() {
        use crate::scheduler::sim::{AgentPolicy, CoPlacement, SimCell, SimWorld};
        use crate::storage::StateStore;

        let store = store();
        // Sim cell: CPU world + 2 GPU agents, interconnect-local (gang, one node).
        let cell = SimCell::new("alpha", SimWorld::cpu(1000, 2048), Duration::from_millis(50))
            .with_agent(AgentPolicy::gpu("p0", 200, 256, 4096))
            .with_agent(AgentPolicy::gpu("p1", 200, 256, 4096))
            .with_co_placement(CoPlacement::InterconnectLocalGpu);
        store_set_json(store.as_ref(), &keys::simcell("alpha"), &cell)
            .await
            .unwrap();

        let mut rec = Reconciler::new(store.clone(), autoscaler());
        let node_id = NodeId::new();
        rec.register_node(
            NodeResources::new(node_id, 8000, 16384)
                .with_gpu(GpuResources::new(0, "A100", 8192))
                .with_gpu(GpuResources::new(1, "A100", 8192)),
        );

        // SimCell -> gang -> reconcile: all 3 members co-placed on the one node.
        let report = rec.reconcile_once().await.unwrap();
        assert_eq!(report.sim_scheduled, 1);
        assert_eq!(rec.sim_bound_count(), 1);
        assert_eq!(rec.nodes[&node_id].gpus_allocated.len(), 2);
        assert_eq!(rec.nodes[&node_id].cpu_allocated, 1000 + 200 + 200);

        // Idempotent: a second pass changes nothing.
        let r2 = rec.reconcile_once().await.unwrap();
        assert_eq!(r2.sim_scheduled, 0);
        assert_eq!(r2.sim_released, 0);

        // Remove the cell -> released; CPU and GPU devices freed exactly.
        store.delete(&keys::simcell("alpha")).await.unwrap();
        let r3 = rec.reconcile_once().await.unwrap();
        assert_eq!(r3.sim_released, 1);
        assert_eq!(rec.sim_bound_count(), 0);
        assert_eq!(rec.nodes[&node_id].gpus_allocated.len(), 0);
        assert_eq!(rec.nodes[&node_id].cpu_allocated, 0);
    }
}