consortium_cli/

claw.rs

//! claw — parallel cluster command execution (replaces clush).
//!
//! Usage:
//!   claw -w node[1-5] echo hello        # run on nodes
//!   claw -w node[1-5] -b 'uname -r'     # gather identical output
//!   claw -w node[1-5] -l root uptime     # run as root
//!   claw --copy /etc/hosts -w node[1-5]  # copy file to nodes
//!   echo 'script' | claw -w node[1-5]    # pipe stdin to command

use std::collections::HashMap;
use std::io::{self, Read, Write};
use std::process;
use std::time::{Duration, Instant};

use clap::Parser;

use crate::display;
use crate::fold;
use crate::groups;
use crate::output::{CliOutput, OutputArgs};
use consortium::node_set::NodeSet;
use consortium::task::{Task, TaskError};
use consortium::worker::exec::ExecWorker;
use consortium::worker::ssh::SshOptions;
use consortium::worker::EventHandler;

/// claw — execute commands in parallel across cluster nodes.
///
/// Runs shell commands on multiple nodes simultaneously with
/// configurable fanout, output gathering, and timeout control.
#[derive(Parser)]
#[command(name = "claw", version, about)]
pub struct Args {
    // ── Node selection ─────────────────────────────────────────────────
    /// Target nodes (bracket notation, e.g., node[1-5]).
    #[arg(short = 'w', long = "nodes")]
    nodes: Option<String>,

    /// Exclude nodes from the target set.
    #[arg(short = 'x', long = "exclude")]
    exclude: Vec<String>,

    /// Use all nodes from the default group.
    #[arg(short = 'a', long = "all")]
    all: bool,

    /// Use nodes from a named group.
    #[arg(short = 'g', long = "group")]
    group: Vec<String>,

    /// Exclude nodes from a named group.
    #[arg(short = 'X', long = "exclude-group")]
    exclude_group: Vec<String>,

    /// Pick N random nodes.
    #[arg(long = "pick")]
    pick: Option<usize>,

    /// Read nodes from a file (one per line).
    #[arg(long = "hostfile")]
    hostfile: Option<String>,

    // ── Execution ──────────────────────────────────────────────────────
    /// Maximum concurrent connections (fanout).
    #[arg(short = 'f', long = "fanout")]
    fanout: Option<usize>,

    /// Remote user for SSH connections.
    #[arg(short = 'l', long = "user")]
    user: Option<String>,

    /// SSH connect timeout in seconds.
    #[arg(short = 't', long = "connect-timeout")]
    connect_timeout: Option<u64>,

    /// Command timeout in seconds.
    #[arg(short = 'u', long = "command-timeout")]
    command_timeout: Option<u64>,

    /// Worker type: exec or ssh (default: ssh for remote, exec for localhost).
    #[arg(short = 'R', long = "worker")]
    worker: Option<String>,

    /// SSH options to pass through (e.g., "-o BatchMode=yes").
    #[arg(short = 'o', long = "options")]
    options: Option<String>,

    /// Remote shell command (default: ssh).
    #[arg(long = "remote", default_value = "ssh")]
    remote: String,

    /// Topology file for tree mode.
    #[arg(long = "topology")]
    topology: Option<String>,

    // ── Output ─────────────────────────────────────────────────────────
    /// Gather identical output together (dshbak mode).
    #[arg(short = 'b', long = "dshbak")]
    dshbak: bool,

    /// Line mode — no output gathering.
    #[arg(short = 'L', long = "line")]
    line: bool,

    /// Disable node labels on output.
    #[arg(short = 'N', long = "label")]
    no_label: bool,

    /// Quiet mode — suppress output, just return exit code.
    #[arg(short = 'q', long = "quiet")]
    quiet: bool,

    /// Return the maximum exit code from all nodes.
    #[arg(short = 'S', long = "maxrc")]
    maxrc: bool,

    /// Fold along these axis only (axis 1..n for nD nodeset).
    #[arg(long = "axis", value_name = "RANGESET")]
    axis: Option<String>,

    #[command(flatten)]
    output: OutputArgs,

    // ── File transfer ──────────────────────────────────────────────────
    /// Copy a file to all target nodes.
    #[arg(long = "copy")]
    copy: Option<String>,

    /// Reverse copy — copy file from nodes to local.
    #[arg(long = "rcopy")]
    rcopy: Option<String>,

    /// Destination directory for copy operations.
    #[arg(long = "dest")]
    dest: Option<String>,

    // ── Positional ─────────────────────────────────────────────────────
    /// Command and arguments to execute.
    command: Vec<String>,
}

pub fn run() {
    let args = Args::parse();
    let exit_code = match execute(args) {
        Ok(code) => code,
        Err(e) => {
            eprintln!("claw: {e}");
            1
        }
    };
    process::exit(exit_code);
}

fn execute(args: Args) -> anyhow::Result<i32> {
    let cli_out = CliOutput::from_args(&args.output);

    // User-specified nD-nodeset fold axis for output display (#356).
    // An empty axis list behaves like Python's empty DEFAULTS.fold_axis
    // tuple (falsy → fold along all axes).
    let fold_axes = match &args.axis {
        Some(axis) => {
            Some(fold::parse_fold_axis(axis).map_err(|e| anyhow::anyhow!("Parse error: {e}"))?)
        }
        None => None,
    };
    let fold_axis: Option<&[i64]> = fold_axes.as_deref().filter(|v| !v.is_empty());

    let stdout = io::stdout();
    let mut out = stdout.lock();

    // ── Resolve target nodes ───────────────────────────────────────────
    let target_nodes = resolve_nodes(&args)?;
    if target_nodes.is_empty() {
        anyhow::bail!("no target nodes specified (use -w, -a, -g, or --hostfile)");
    }

    if cli_out.is_verbose() {
        cli_out.info(format!(
            "claw: targeting {} node(s): {}",
            target_nodes.len(),
            target_nodes
        ));
    }

    // ── Build the command string ───────────────────────────────────────
    let command = if !args.command.is_empty() {
        args.command.join(" ")
    } else if !is_terminal::is_terminal(io::stdin()) {
        // Read command from stdin
        let mut buf = String::new();
        io::stdin().read_to_string(&mut buf)?;
        buf.trim_end().to_string()
    } else {
        anyhow::bail!("no command specified");
    };

    if command.is_empty() {
        anyhow::bail!("empty command");
    }

    // ── Build SSH options if needed ────────────────────────────────────
    let _ssh_opts = build_ssh_options(&args);

    // ── Determine worker type ──────────────────────────────────────────
    let use_ssh = match args.worker.as_deref() {
        Some("exec") => false,
        Some("ssh") => true,
        Some(other) => anyhow::bail!("unknown worker type: {other} (use 'exec' or 'ssh')"),
        None => {
            // Auto-detect: if all nodes are localhost, use exec
            let all_local = target_nodes
                .iter()
                .all(|n| n == "localhost" || n == "127.0.0.1" || n == "::1");
            !all_local
        }
    };

    // ── Build the command with SSH wrapping if needed ───────────────────
    let node_list: Vec<String> = target_nodes.iter().collect();
    let fanout = args.fanout.unwrap_or(64);
    let timeout = args.command_timeout.map(Duration::from_secs);

    let effective_command = if use_ssh {
        // Build SSH command: ssh [options] [user@]%h command
        let mut ssh_cmd = args.remote.clone();
        if let Some(ref opts) = args.options {
            ssh_cmd.push(' ');
            ssh_cmd.push_str(opts);
        }
        if let Some(ref ct) = args.connect_timeout {
            ssh_cmd.push_str(&format!(" -o ConnectTimeout={ct}"));
        }
        if let Some(ref user) = args.user {
            ssh_cmd.push_str(&format!(" {user}@%h"));
        } else {
            ssh_cmd.push_str(" %h");
        }
        ssh_cmd.push(' ');
        ssh_cmd.push_str(&shell_escape(&command));
        ssh_cmd
    } else {
        command.clone()
    };

    // ── Create and run the task ────────────────────────────────────────
    let mut task = Task::new();
    let worker =
        ExecWorker::new(node_list.clone(), effective_command, fanout, timeout).with_stderr(true);

    // Set up progress tracking (only when stderr is a TTY and >1 node)
    let num_nodes = node_list.len();
    let show_progress = !args.quiet && is_terminal::is_terminal(io::stderr()) && num_nodes > 1;

    // In standard (non-gather, non-quiet) mode, stream output lines as they
    // arrive (GH#528/GH#597 line-buffering port): a LineStreamHandler prints
    // each completed line immediately instead of buffering until exit.
    let stream_lines = !args.dshbak && !args.quiet;

    let progress = if show_progress {
        let (bar, state, handler) = display::create_progress(num_nodes);
        (Some(bar), Some(state), Some(handler))
    } else {
        (None, None, None)
    };

    let user_handler: Option<Box<dyn EventHandler>> = if stream_lines {
        let mut handler = display::LineStreamHandler::new(!args.no_label);
        if let Some(progress_handler) = progress.2 {
            handler = handler.with_inner(Box::new(progress_handler));
        }
        Some(Box::new(handler))
    } else {
        progress.2.map(|h| Box::new(h) as Box<dyn EventHandler>)
    };

    task.schedule(Box::new(worker), user_handler, false);
    let pb = progress.0;
    let state = progress.1;

    let start = Instant::now();
    let task_timeout = args.command_timeout.map(|t| Duration::from_secs(t + 5));
    match task.run(task_timeout) {
        Ok(()) => {}
        Err(TaskError::Timeout) => {
            if !args.quiet {
                eprintln!("claw: command timed out");
            }
        }
        Err(e) => return Err(e.into()),
    }

    // Finalize progress display
    if let (Some(bar), Some(state)) = (&pb, &state) {
        display::finish_progress(bar, state, start.elapsed());
    }

    // ── Collect and display output ─────────────────────────────────────
    let max_rc = display_results(&task, &args, fold_axis, stream_lines, &mut out)?;

    if args.maxrc {
        Ok(max_rc.unwrap_or(0))
    } else if max_rc.map_or(false, |rc| rc != 0) {
        Ok(1)
    } else {
        Ok(0)
    }
}

/// Resolve target nodes from all sources.
fn resolve_nodes(args: &Args) -> anyhow::Result<NodeSet> {
    let mut ns = NodeSet::new();

    // -w / --nodes
    if let Some(ref pattern) = args.nodes {
        let parsed = NodeSet::parse(pattern)?;
        ns.update(&parsed);
    }

    // -g / --group
    if !args.group.is_empty() || args.all {
        let mut config = groups::load_group_resolver()?;
        let resolver = config.resolver()?;

        if args.all {
            // -a: all nodes from default group source
            let all_nodes = resolver.all_nodes(None)?;
            for node in &all_nodes {
                let _ = ns.update_str(node);
            }
        }

        for group_name in &args.group {
            let nodes = resolver.group_nodes(group_name, None)?;
            for node in &nodes {
                let _ = ns.update_str(node);
            }
        }

        // -X / --exclude-group
        for group_name in &args.exclude_group {
            let nodes = resolver.group_nodes(group_name, None)?;
            for node in &nodes {
                let exclude_ns = NodeSet::parse(node)?;
                ns.difference_update(&exclude_ns);
            }
        }
    }

    // --hostfile
    if let Some(ref path) = args.hostfile {
        let content = std::fs::read_to_string(path)?;
        for line in content.lines() {
            let trimmed = line.trim();
            if !trimmed.is_empty() && !trimmed.starts_with('#') {
                let _ = ns.update_str(trimmed);
            }
        }
    }

    // Excludes (-x)
    for pat in &args.exclude {
        let other = NodeSet::parse(pat)?;
        ns.difference_update(&other);
    }

    // Pick random subset
    if let Some(n) = args.pick {
        let all: Vec<String> = ns.iter().collect();
        if n < all.len() {
            let mut new_ns = NodeSet::new();
            // Simple sequential pick (deterministic for now)
            for node in all.iter().take(n) {
                let _ = new_ns.update_str(node);
            }
            ns = new_ns;
        }
    }

    Ok(ns)
}

/// Build SSH options from CLI args.
fn build_ssh_options(args: &Args) -> SshOptions {
    let mut opts = SshOptions::default();
    if let Some(ref user) = args.user {
        opts.user = Some(user.clone());
    }
    if let Some(ct) = args.connect_timeout {
        opts.connect_timeout = Some(ct as u32);
    }
    if let Some(ref _extra) = args.options {
        // Parse "-o Key=Value" style options
        opts.ssh_path = args.remote.clone(); // String, not Option
                                             // Store raw options for passthrough
    }
    opts
}

/// Display results from a completed task.
fn display_results(
    task: &Task,
    args: &Args,
    fold_axis: Option<&[i64]>,
    streamed: bool,
    out: &mut impl Write,
) -> anyhow::Result<Option<i32>> {
    if args.quiet {
        return Ok(task.max_retcode());
    }

    if args.dshbak {
        display_dshbak(task, fold_axis, out)?;
    } else if !streamed {
        display_standard(task, args, out)?;
    }
    // When `streamed`, standard output was already printed line by line
    // during execution by the LineStreamHandler.

    // Show return codes for non-zero exits
    let retcodes = task.iter_retcodes(None);
    for (rc, nodes) in &retcodes {
        if *rc != 0 {
            let mut ns = NodeSet::new();
            for node in nodes {
                let _ = ns.update_str(node);
            }
            eprintln!("claw: {ns}: exited with return code {rc}");
        }
    }

    // Show timeouts
    let timeout_count = task.num_timeout();
    if timeout_count > 0 {
        let timeout_nodes: Vec<&str> = task.iter_keys_timeout().collect();
        let mut ns = NodeSet::new();
        for node in &timeout_nodes {
            let _ = ns.update_str(node);
        }
        eprintln!("claw: {ns}: command timeout");
    }

    Ok(task.max_retcode())
}

/// Display output in dshbak (gathered) mode.
///
/// Groups nodes by identical output, then prints one header block per group.
/// Groups sort like upstream clush (`bufnodeset_cmpkey`): larger groups
/// first, then by first node name. Headers fold with the nD-aware folder
/// honoring the `--axis` constraint (#356) and carry the ` (N)` node count
/// like upstream `Display.format_header()`.
fn display_dshbak(
    task: &Task,
    fold_axis: Option<&[i64]>,
    out: &mut impl Write,
) -> anyhow::Result<()> {
    // Group nodes by their output
    let mut output_to_nodes: HashMap<Vec<u8>, Vec<String>> = HashMap::new();

    // Iterate over all nodes that ran
    let retcodes = task.iter_retcodes(None);
    for (_rc, nodes) in &retcodes {
        for node in nodes {
            let buf = task.node_buffer(node).unwrap_or_default();
            output_to_nodes
                .entry(buf)
                .or_default()
                .push(node.to_string());
        }
    }

    let mut groups: Vec<(Vec<u8>, Vec<String>)> = output_to_nodes.into_iter().collect();
    for (_buf, names) in groups.iter_mut() {
        names.sort();
    }
    // larger nodeset first, then sorted by first node index
    groups.sort_by(|a, b| b.1.len().cmp(&a.1.len()).then_with(|| a.1[0].cmp(&b.1[0])));

    for (output, names) in &groups {
        let folded = fold::fold_nodes(names, fold_axis);
        display::print_gathered_header(&folded, names.len(), out)?;
        out.write_all(output)?;
        if !output.is_empty() && !output.ends_with(b"\n") {
            writeln!(out)?;
        }
    }

    Ok(())
}

/// Display output in standard (labeled line) mode.
fn display_standard(task: &Task, args: &Args, out: &mut impl Write) -> anyhow::Result<()> {
    let retcodes = task.iter_retcodes(None);

    for (_rc, nodes) in &retcodes {
        for node in nodes {
            let buf = task.node_buffer(node).unwrap_or_default();
            if buf.is_empty() {
                continue;
            }
            let text = String::from_utf8_lossy(&buf);
            for line in text.lines() {
                if args.no_label {
                    writeln!(out, "{line}")?;
                } else {
                    display::print_line_with_label(node, line, out)?;
                }
            }
        }
    }

    Ok(())
}

/// Shell-escape a command string for SSH passthrough.
fn shell_escape(s: &str) -> String {
    // If it contains no special chars, return as-is
    if s.chars()
        .all(|c| c.is_alphanumeric() || matches!(c, '-' | '_' | '.' | '/' | '=' | ':' | ','))
    {
        return s.to_string();
    }
    // Single-quote escape
    format!("'{}'", s.replace('\'', "'\\''"))
}