dora-cli 1.0.0-rc.2

use std::{
    net::SocketAddr,
    path::PathBuf,
    process::Command,
    time::{Duration, Instant},
};

use clap::Args;
use eyre::Context;

use crate::{
    command::{
        Executable, default_tracing,
        up::{detach_process, dora_executable_path},
    },
    common::{connect_to_coordinator, connect_with_retry},
};

use super::config::ClusterConfig;
use super::{
    format_daemon_port_arg, format_labels_arg, format_zenoh_peer_arg, query_connected_daemons,
    run_ssh, ssh_target,
};

/// Bring up a multi-machine cluster from a cluster.yml file.
///
/// Starts the coordinator locally, then SSH-es into each machine to
/// start a daemon.
///
/// Examples:
///
///   dora cluster up cluster.yml
#[derive(Debug, Args)]
#[clap(verbatim_doc_comment)]
pub struct Up {
    /// Path to the cluster configuration file
    #[clap(value_name = "PATH", value_hint = clap::ValueHint::FilePath)]
    config: PathBuf,
}

impl Executable for Up {
    fn execute(self) -> eyre::Result<()> {
        default_tracing()?;
        let config = ClusterConfig::load(&self.config)?;
        let coordinator_addr: SocketAddr =
            (config.coordinator.addr, config.coordinator.port).into();

        // 1. Connect to existing coordinator or start a new one
        let session = match connect_to_coordinator(coordinator_addr) {
            Ok(s) => {
                println!("Coordinator already running at {coordinator_addr}");
                s
            }
            Err(_) => {
                start_coordinator(config.coordinator.port)?;
                connect_with_retry(coordinator_addr, Duration::from_secs(10)).map_err(|err| {
                    eyre::eyre!("timed out waiting for coordinator at {coordinator_addr}: {err}")
                })?
            }
        };

        // 2. SSH into each machine to start a daemon
        let zenoh_peer_arg = format_zenoh_peer_arg(config.zenoh_peer.as_deref());
        let mut ssh_failures: Vec<(String, String)> = Vec::new();
        for machine in &config.machines {
            let target = ssh_target(machine);
            let labels_arg = format_labels_arg(&machine.labels);
            let daemon_port_arg = format_daemon_port_arg(machine.daemon_port);
            let remote_cmd = format!(
                "nohup dora daemon --machine-id {id} --coordinator-addr {addr} --coordinator-port {port}{daemon_port_arg}{zenoh_peer_arg}{labels} --quiet > /tmp/dora-daemon-{id}.log 2>&1 &",
                id = machine.id,
                addr = config.coordinator.addr,
                port = config.coordinator.port,
                labels = labels_arg,
            );

            println!("Starting daemon on {} ({})", machine.id, target);
            match run_ssh(&target, machine.port, &remote_cmd) {
                Ok(true) => {}
                Ok(false) => {
                    let msg = "ssh command failed".to_string();
                    eprintln!(
                        "  WARNING: failed to start daemon on `{}`: {msg}",
                        machine.id
                    );
                    ssh_failures.push((machine.id.clone(), msg));
                }
                Err(err) => {
                    let msg = format!("{err}");
                    eprintln!(
                        "  WARNING: failed to start daemon on `{}`: {msg}",
                        machine.id
                    );
                    ssh_failures.push((machine.id.clone(), msg));
                }
            }
        }

        // 3. Poll until all (successful) daemons have registered
        let expected: Vec<&str> = config
            .machines
            .iter()
            .filter(|m| !ssh_failures.iter().any(|(id, _)| id == &m.id))
            .map(|m| m.id.as_str())
            .collect();

        let mut missing_daemons: Vec<String> = Vec::new();
        if !expected.is_empty() {
            println!("Waiting for {} daemon(s) to connect...", expected.len());
            let deadline = Instant::now() + Duration::from_secs(30);
            loop {
                let connected = query_connected_daemons(&session)?;
                let all_present = expected.iter().all(|machine_id| {
                    connected
                        .iter()
                        .any(|d| d.daemon_id.matches_machine_id(machine_id))
                });
                if all_present {
                    break;
                }
                if Instant::now() >= deadline {
                    missing_daemons = expected
                        .iter()
                        .copied()
                        .filter(|machine_id| {
                            !connected
                                .iter()
                                .any(|d| d.daemon_id.matches_machine_id(machine_id))
                        })
                        .map(String::from)
                        .collect();
                    eprintln!(
                        "WARNING: timed out waiting for daemon(s): {}",
                        missing_daemons.join(", ")
                    );
                    break;
                }
                std::thread::sleep(Duration::from_millis(500));
            }
        }

        // 4. Report. A partial-up state (some daemons unreachable or never
        // registered) must exit non-zero so callers — scripts, CI, the
        // cluster-e2e job — can react instead of silently treating
        // "Cluster partially up" as success.
        let ok_count = config.machines.len() - ssh_failures.len() - missing_daemons.len();
        if ssh_failures.is_empty() && missing_daemons.is_empty() {
            println!(
                "Cluster is up: coordinator + {} daemon(s)",
                config.machines.len()
            );
            Ok(())
        } else {
            println!(
                "Cluster partially up: coordinator + {ok_count}/{} daemon(s)",
                config.machines.len()
            );
            for (id, reason) in &ssh_failures {
                eprintln!("  {id}: {reason}");
            }
            eyre::bail!(
                "cluster up incomplete: {} ssh failure(s), {} daemon(s) did not register",
                ssh_failures.len(),
                missing_daemons.len()
            )
        }
    }
}

fn start_coordinator(port: u16) -> eyre::Result<()> {
    let path = dora_executable_path()?;
    let mut cmd = Command::new(path);
    cmd.args([
        "coordinator",
        "--interface",
        "0.0.0.0",
        "--port",
        &port.to_string(),
        "--quiet",
    ]);
    detach_process(&mut cmd);
    cmd.spawn().wrap_err("failed to start dora coordinator")?;
    println!("Started coordinator on 0.0.0.0:{port}");
    Ok(())
}