launcher/src/health.rs

use anyhow::Result;

use crate::config::{Cluster, Config, Recipe};
use crate::tmux::{SessionHealth, TmuxManager};

/// Health status of a node's tmux session
#[derive(Debug, Clone, PartialEq)]
pub enum SessionStatus {
    Running,
    Stopped,
    Degraded,
    Unknown,
}

/// Health report for a single node
#[derive(Debug, Clone)]
pub struct NodeHealth {
    pub name: String,
    pub ib_ip: String,
    pub rank: u32,
    pub session_exists: bool,
    pub process_running: bool,
    pub port_accessible: bool,
    pub status: SessionStatus,
}

/// Overall health report for a cluster
#[derive(Debug, Clone)]
pub struct HealthReport {
    pub nodes: Vec<NodeHealth>,
    pub overall: HealthStatus,
}

/// Overall cluster health status
#[derive(Debug, Clone, PartialEq)]
pub enum HealthStatus {
    Healthy,
    Degraded,
    Unhealthy,
}

/// Checks the health of all tmux sessions for a recipe across a cluster.
/// Uses deep health checks: session existence, process liveness, GPU status, API reachability.
/// Returns a [`HealthReport`] with per-node details — no output is printed by this function.
pub async fn check_cluster_health(
    manager: &TmuxManager,
    config: &Config,
    cluster: &Cluster,
    recipe: &Recipe,
    dry_run: bool,
) -> Result<HealthReport> {
    let mut nodes: Vec<NodeHealth> = Vec::new();

    let mut handles = Vec::new();
    for node in &cluster.nodes {
        let sess = manager.session_name(recipe, node.rank);
        let config = config.clone();
        let node = node.clone();
        let cluster = cluster.clone();
        let recipe = recipe.clone();
        let dry = dry_run;
        let manager = manager.clone();
        handles.push(tokio::spawn(async move {
            let health = manager
                .deep_check_session(&config, &node, &cluster, &sess, &recipe, dry)
                .await;
            (node.name.clone(), node.ib_ip.clone(), node.rank, health)
        }));
    }

    let results = futures::future::join_all(handles).await;
    for result in results {
        match result {
            Ok((name, ib_ip, rank, Ok(health))) => {
                let status = classify_status(&health);
                nodes.push(NodeHealth {
                    name,
                    ib_ip,
                    rank,
                    session_exists: health.session_exists,
                    process_running: health.process_alive,
                    port_accessible: health.api_ok,
                    status,
                });
            }
            Ok((name, ib_ip, rank, Err(_))) => {
                nodes.push(NodeHealth {
                    name,
                    ib_ip,
                    rank,
                    session_exists: false,
                    process_running: false,
                    port_accessible: false,
                    status: SessionStatus::Unknown,
                });
            }
            Err(_) => {}
        }
    }

    let overall = compute_overall_status(&nodes);

    Ok(HealthReport { nodes, overall })
}

/// Classifies overall session status from detailed health data.
pub fn classify_status(health: &SessionHealth) -> SessionStatus {
    if health.is_healthy() {
        SessionStatus::Running
    } else if health.is_degraded() || (health.session_exists && !health.gpu_ok) {
        SessionStatus::Degraded
    } else {
        SessionStatus::Stopped
    }
}

/// Computes overall cluster health from per-node statuses.
fn compute_overall_status(nodes: &[NodeHealth]) -> HealthStatus {
    let has_unknown = nodes.iter().any(|n| n.status == SessionStatus::Unknown);
    let has_stopped = nodes.iter().any(|n| n.status == SessionStatus::Stopped);
    let has_degraded = nodes.iter().any(|n| n.status == SessionStatus::Degraded);
    let all_running = nodes.iter().all(|n| n.status == SessionStatus::Running);

    if all_running {
        HealthStatus::Healthy
    } else if !has_unknown && !has_stopped && has_degraded {
        HealthStatus::Degraded
    } else {
        HealthStatus::Unhealthy
    }
}