zc2 0.0.29

//! Pure reconcile decisions for `zc agent` (spec §6.5): given what the owner
//! asked for and what is running, what to do next. No processes, no I/O.

use crate::agent::summary::Problem;
use std::collections::BTreeSet;
use std::path::PathBuf;
use std::time::{Duration, Instant};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Desired {
    pub sharing: bool,
    pub workers: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BrokerObs {
    /// Nothing on the broker port.
    Stopped,
    /// Our broker process is alive but `/health` does not answer yet.
    Starting,
    Running,
    /// We sent it SIGTERM and it is flushing its WAL.
    Stopping,
    /// A zakuro broker we did not start (e.g. a manual `zc up -d`).
    Unmanaged,
    /// Something that is not a zakuro broker holds the broker port.
    PortBlocked,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WorkerObs {
    pub index: u32,
    /// A process is alive for this slot. False: it crashed (or exited while draining).
    pub running: bool,
    pub draining: bool,
    pub active_requests: u32,
    /// Draining for longer than the drain timeout.
    pub drain_expired: bool,
    /// For a crashed slot: its restart backoff has elapsed.
    pub restart_ready: bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Observed {
    pub broker: BrokerObs,
    pub workers: Vec<WorkerObs>,
    pub prerequisites_ok: bool,
    /// The broker's own crash-restart backoff
    /// (`restart_backoff(consecutive broker failures)`) has elapsed. The
    /// supervisor (Task 10) computes this; `plan_actions` only gates on it.
    pub broker_restart_ready: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Action {
    StartBroker,
    StopBroker,
    StartWorker(u32),
    DrainWorker(u32),
    /// Stop the slot's process (if any), `DELETE` it from the broker, forget the slot.
    KillWorker(u32),
    RestartWorker(u32),
}

pub fn plan_actions(desired: &Desired, observed: &Observed) -> Vec<Action> {
    use Action::*;
    let mut actions = Vec::new();
    if matches!(
        observed.broker,
        BrokerObs::Unmanaged | BrokerObs::PortBlocked
    ) {
        return actions; // never touch a broker we did not start
    }

    // Drains end when the worker is idle, out of time, or already gone.
    for w in observed.workers.iter().filter(|w| w.draining) {
        if !w.running || w.active_requests == 0 || w.drain_expired {
            actions.push(KillWorker(w.index));
        }
    }

    let live: Vec<&WorkerObs> = observed
        .workers
        .iter()
        .filter(|w| w.running && !w.draining)
        .collect();
    let mut crashed: Vec<&WorkerObs> = observed
        .workers
        .iter()
        .filter(|w| !w.running && !w.draining)
        .collect();
    crashed.sort_by_key(|w| w.index);

    if !desired.sharing {
        actions.extend(live.iter().map(|w| DrainWorker(w.index)));
        actions.extend(crashed.iter().map(|w| KillWorker(w.index)));
        // Pausing stops the broker once no worker process is left, whether
        // it is already `Running` or still `Starting`.
        if matches!(observed.broker, BrokerObs::Running | BrokerObs::Starting)
            && !observed.workers.iter().any(|w| w.running)
        {
            actions.push(StopBroker);
        }
        return actions;
    }

    match observed.broker {
        BrokerObs::Stopped => {
            // Wait out the broker's own crash-restart backoff before retrying.
            if observed.broker_restart_ready {
                actions.push(StartBroker);
            }
            return actions;
        }
        BrokerObs::Starting | BrokerObs::Stopping => return actions,
        _ => {}
    }

    let want = desired.workers as usize;
    if live.len() >= want {
        // Highest-index idle first; if none is idle, highest index.
        let mut pick = live.clone();
        pick.sort_by_key(|w| (w.active_requests != 0, std::cmp::Reverse(w.index)));
        actions.extend(
            pick.iter()
                .take(live.len() - want)
                .map(|w| DrainWorker(w.index)),
        );
        actions.extend(crashed.iter().map(|w| KillWorker(w.index)));
        return actions;
    }
    if !observed.prerequisites_ok {
        return actions;
    }

    let mut need = want - live.len();
    for w in &crashed {
        if need == 0 {
            actions.push(KillWorker(w.index));
            continue;
        }
        if w.restart_ready {
            actions.push(RestartWorker(w.index));
        }
        need -= 1; // a slot in backoff still holds its place
    }
    let taken: BTreeSet<u32> = observed.workers.iter().map(|w| w.index).collect();
    let mut i = 0;
    while need > 0 {
        if !taken.contains(&i) {
            actions.push(StartWorker(i));
            need -= 1;
        }
        i += 1;
    }
    actions
}

/// 1, 2, 4 … 60 s for the 1st, 2nd, 3rd … consecutive failure.
pub fn restart_backoff(consecutive_failures: u32) -> Duration {
    let n = consecutive_failures.max(1) - 1;
    Duration::from_secs((1u64 << n.min(6)).min(60))
}

/// Three failures within five minutes.
pub fn is_crashloop(failures: &[Instant], now: Instant) -> bool {
    failures
        .iter()
        .filter(|t| now.saturating_duration_since(**t) <= Duration::from_secs(300))
        .count()
        >= 3
}

pub fn find_in_path(name: &str, path_var: &str) -> Option<PathBuf> {
    path_var
        .split(':')
        .filter(|d| !d.is_empty())
        .map(|d| PathBuf::from(d).join(name))
        .find(|p| p.is_file())
}

/// Spec §6.5: credentials, `uv` and the zakuro dir must exist before a worker
/// starts. A replaced worker command (tests, `ZAKURO_AGENT_WORKER_CMD`) needs
/// neither `uv` nor the zakuro dir.
///
/// `zakuro_dir_detail` is the `zakuro_dir_missing` problem's `detail`: `Some`
/// when a `ZAKURO_WORKER_DIR` setting pointed at an invalid path (naming it),
/// `None` otherwise. Ignored when `zakuro_dir_found` is true.
pub fn prerequisite_problems(
    uv_found: bool,
    zakuro_dir_found: bool,
    zakuro_dir_detail: Option<String>,
    logged_in: bool,
    worker_cmd_overridden: bool,
) -> Vec<Problem> {
    let mut out = Vec::new();
    if !logged_in {
        out.push(Problem::new("not_logged_in"));
    }
    if !worker_cmd_overridden {
        if !uv_found {
            out.push(Problem::new("uv_missing"));
        }
        if !zakuro_dir_found {
            let mut p = Problem::new("zakuro_dir_missing");
            if let Some(detail) = zakuro_dir_detail {
                p = p.with_detail(detail);
            }
            out.push(p);
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::Action::*;
    use super::*;

    fn w(index: u32) -> WorkerObs {
        WorkerObs {
            index,
            running: true,
            draining: false,
            active_requests: 0,
            drain_expired: false,
            restart_ready: true,
        }
    }
    fn busy(index: u32) -> WorkerObs {
        WorkerObs {
            active_requests: 1,
            ..w(index)
        }
    }
    fn draining(index: u32, active: u32, expired: bool) -> WorkerObs {
        WorkerObs {
            draining: true,
            active_requests: active,
            drain_expired: expired,
            ..w(index)
        }
    }
    fn crashed(index: u32, ready: bool) -> WorkerObs {
        WorkerObs {
            running: false,
            restart_ready: ready,
            ..w(index)
        }
    }
    fn obs(broker: BrokerObs, workers: Vec<WorkerObs>) -> Observed {
        Observed {
            broker,
            workers,
            prerequisites_ok: true,
            broker_restart_ready: true,
        }
    }
    fn want(sharing: bool, workers: u32) -> Desired {
        Desired { sharing, workers }
    }

    #[test]
    fn idle_and_not_sharing_does_nothing() {
        assert_eq!(
            plan_actions(&want(false, 1), &obs(BrokerObs::Stopped, vec![])),
            vec![]
        );
    }

    #[test]
    fn sharing_starts_the_broker_first_then_waits_for_it() {
        assert_eq!(
            plan_actions(&want(true, 2), &obs(BrokerObs::Stopped, vec![])),
            vec![StartBroker]
        );
        assert_eq!(
            plan_actions(&want(true, 2), &obs(BrokerObs::Starting, vec![])),
            vec![]
        );
    }

    #[test]
    fn workers_start_at_the_lowest_free_indices() {
        assert_eq!(
            plan_actions(&want(true, 3), &obs(BrokerObs::Running, vec![w(1)])),
            vec![StartWorker(0), StartWorker(2)]
        );
    }

    #[test]
    fn missing_prerequisites_start_nothing() {
        let o = Observed {
            broker: BrokerObs::Running,
            workers: vec![],
            prerequisites_ok: false,
            broker_restart_ready: true,
        };
        assert_eq!(plan_actions(&want(true, 2), &o), vec![]);
    }

    #[test]
    fn scale_down_drains_the_highest_idle_workers_first() {
        assert_eq!(
            plan_actions(
                &want(true, 1),
                &obs(BrokerObs::Running, vec![w(0), busy(1), w(2)])
            ),
            vec![DrainWorker(2), DrainWorker(0)]
        );
    }

    #[test]
    fn scale_down_with_everyone_busy_drains_the_highest_index() {
        assert_eq!(
            plan_actions(
                &want(true, 1),
                &obs(BrokerObs::Running, vec![busy(0), busy(1)])
            ),
            vec![DrainWorker(1)]
        );
    }

    #[test]
    fn a_drain_ends_when_idle_out_of_time_or_already_gone() {
        let dead_draining = WorkerObs {
            running: false,
            ..draining(3, 1, false)
        };
        assert_eq!(
            plan_actions(
                &want(true, 0),
                &obs(
                    BrokerObs::Running,
                    vec![
                        draining(0, 0, false),
                        draining(1, 2, false),
                        draining(2, 2, true),
                        dead_draining
                    ]
                )
            ),
            vec![KillWorker(0), KillWorker(2), KillWorker(3)]
        );
    }

    #[test]
    fn pause_drains_everyone_then_stops_the_broker() {
        let off = want(false, 2);
        assert_eq!(
            plan_actions(&off, &obs(BrokerObs::Running, vec![w(0), w(1)])),
            vec![DrainWorker(0), DrainWorker(1)]
        );
        assert_eq!(
            plan_actions(&off, &obs(BrokerObs::Running, vec![draining(0, 1, false)])),
            vec![]
        );
        assert_eq!(
            plan_actions(&off, &obs(BrokerObs::Running, vec![])),
            vec![StopBroker]
        );
    }

    /// A broker that is still starting (no worker process left) is stopped
    /// on pause exactly like a running one — the brief's original tests only
    /// covered `BrokerObs::Running`.
    #[test]
    fn pause_stops_the_broker_while_it_is_still_starting() {
        let off = want(false, 2);
        assert_eq!(
            plan_actions(&off, &obs(BrokerObs::Starting, vec![w(0), w(1)])),
            vec![DrainWorker(0), DrainWorker(1)]
        );
        assert_eq!(
            plan_actions(&off, &obs(BrokerObs::Starting, vec![draining(0, 1, false)])),
            vec![]
        );
        assert_eq!(
            plan_actions(&off, &obs(BrokerObs::Starting, vec![])),
            vec![StopBroker]
        );
    }

    #[test]
    fn a_crashed_worker_restarts_only_after_its_backoff() {
        assert_eq!(
            plan_actions(
                &want(true, 1),
                &obs(BrokerObs::Running, vec![crashed(0, true)])
            ),
            vec![RestartWorker(0)]
        );
        assert_eq!(
            plan_actions(
                &want(true, 1),
                &obs(BrokerObs::Running, vec![crashed(0, false)])
            ),
            vec![],
            "a slot in backoff holds its place: no replacement at another index"
        );
    }

    #[test]
    fn crashed_slots_beyond_the_target_are_cleaned_up() {
        assert_eq!(
            plan_actions(
                &want(true, 1),
                &obs(BrokerObs::Running, vec![w(0), crashed(1, false)])
            ),
            vec![KillWorker(1)]
        );
    }

    #[test]
    fn an_unmanaged_or_blocked_broker_is_never_touched() {
        assert_eq!(
            plan_actions(&want(true, 3), &obs(BrokerObs::Unmanaged, vec![w(0)])),
            vec![]
        );
        assert_eq!(
            plan_actions(&want(false, 0), &obs(BrokerObs::PortBlocked, vec![])),
            vec![]
        );
    }

    /// `StartBroker` waits for the broker's own crash-restart backoff. The
    /// supervisor (Task 10) will set `broker_restart_ready` from
    /// `restart_backoff(consecutive broker failures)`.
    #[test]
    fn sharing_waits_out_the_broker_restart_backoff_before_starting_it() {
        let not_ready = Observed {
            broker: BrokerObs::Stopped,
            workers: vec![],
            prerequisites_ok: true,
            broker_restart_ready: false,
        };
        assert_eq!(plan_actions(&want(true, 1), &not_ready), vec![]);

        let ready = Observed {
            broker_restart_ready: true,
            ..not_ready
        };
        assert_eq!(plan_actions(&want(true, 1), &ready), vec![StartBroker]);
    }

    #[test]
    fn restart_backoff_doubles_up_to_a_minute() {
        let secs: Vec<u64> = [1, 2, 3, 4, 7, 20]
            .iter()
            .map(|n| restart_backoff(*n).as_secs())
            .collect();
        assert_eq!(secs, vec![1, 2, 4, 8, 60, 60]);
    }

    #[test]
    fn crashloop_is_three_failures_within_five_minutes() {
        let t = Instant::now();
        let at = |s: u64| t + Duration::from_secs(s);
        assert!(!is_crashloop(&[at(0), at(10)], at(20)));
        assert!(is_crashloop(&[at(0), at(10), at(20)], at(30)));
        assert!(
            !is_crashloop(&[at(0), at(10), at(20)], at(400)),
            "old failures age out"
        );
    }

    #[test]
    fn prerequisites() {
        let dir = crate::agent::files::tests::tmp("path");
        std::fs::create_dir_all(&dir).unwrap();
        std::fs::write(dir.join("uv"), b"#!/bin/sh\n").unwrap();
        let path = format!("/nonexistent:{}", dir.display());
        assert_eq!(find_in_path("uv", &path), Some(dir.join("uv")));
        assert_eq!(find_in_path("nope", &path), None);

        let codes = |p: Vec<crate::agent::summary::Problem>| {
            p.into_iter().map(|p| p.code).collect::<Vec<_>>()
        };
        assert_eq!(
            codes(prerequisite_problems(true, true, None, true, false)),
            Vec::<String>::new()
        );
        assert_eq!(
            codes(prerequisite_problems(false, false, None, false, false)),
            vec!["not_logged_in", "uv_missing", "zakuro_dir_missing"]
        );
        assert_eq!(
            codes(prerequisite_problems(false, false, None, true, true)),
            Vec::<String>::new(),
            "an overridden worker command needs neither uv nor the zakuro dir"
        );
    }

    /// The `zakuro_dir_missing` problem carries the invalid setting in its
    /// `detail`, or `null` when nothing was set at all.
    #[test]
    fn zakuro_dir_missing_detail_names_the_bad_setting() {
        let find = |detail: Option<String>| {
            prerequisite_problems(true, false, detail, true, false)
                .into_iter()
                .find(|p| p.code == "zakuro_dir_missing")
                .expect("zakuro_dir_missing is reported")
        };
        assert_eq!(
            find(Some(
                "ZAKURO_WORKER_DIR=/bad has no zakuro/worker/server.py".into()
            ))
            .detail,
            Some("ZAKURO_WORKER_DIR=/bad has no zakuro/worker/server.py".into())
        );
        assert_eq!(find(None).detail, None);
    }
}