Skip to main content

consortium_cli/
molt.rs

1//! molt — aggregate cluster output (replaces clubak).
2//!
3//! Reads "node:output" lines from stdin or files and aggregates them.
4//!
5//! Usage:
6//!   some_command | molt -b        # gather nodes with identical output
7//!   some_command | molt -L        # line mode, ordered by node
8//!   molt -b < output.txt         # from file
9
10use std::collections::HashMap;
11use std::io::{self, BufRead, Write};
12use std::process;
13
14use clap::{Parser, ValueEnum};
15
16use crate::display;
17use crate::fold;
18use crate::output::{CliOutput, OutputArgs};
19use consortium::node_set::NodeSet;
20
21/// Whether input keys are interpreted as nodesets (upstream clubak
22/// `--interpret-keys`; THREE_CHOICES).
23#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq, Default)]
24enum InterpretKeys {
25    /// Keys are used verbatim, never parsed as nodesets.
26    Never,
27    /// Keys must parse as nodesets; a parse error is fatal.
28    Always,
29    /// Keys are parsed as nodesets until one fails, then never again.
30    #[default]
31    Auto,
32}
33
34/// molt — aggregate and display cluster command output.
35///
36/// Reads "key: output" lines from stdin (like dsh/pdsh output) and
37/// groups nodes that produced identical output together.
38#[derive(Parser)]
39#[command(name = "molt", version, about)]
40pub struct Args {
41    /// Gather nodes with identical output (dshbak mode).
42    #[arg(short = 'b', long = "dshbak")]
43    dshbak: bool,
44
45    /// Line mode — disable header, order by node name.
46    #[arg(short = 'L', long = "line")]
47    line: bool,
48
49    /// Key separator (default: ':').
50    #[arg(short = 'S', long = "separator", default_value = ":")]
51    separator: String,
52
53    /// Message tree trace mode.
54    #[arg(short = 'T', long = "tree")]
55    tree: bool,
56
57    /// No group source prefix in regroup output.
58    #[arg(short = 'G', long = "groupbase")]
59    groupbase: bool,
60
61    /// Fold along these axis only (axis 1..n for nD nodeset).
62    #[arg(long = "axis", value_name = "RANGESET")]
63    axis: Option<String>,
64
65    /// Whether to interpret keys (never, always or auto).
66    #[arg(long = "interpret-keys", value_name = "WHEN", default_value = "auto")]
67    interpret_keys: InterpretKeys,
68
69    /// Input files (reads stdin if none given).
70    files: Vec<String>,
71
72    #[command(flatten)]
73    output: OutputArgs,
74}
75
76pub fn run() {
77    let args = Args::parse();
78    let _out = CliOutput::from_args(&args.output);
79    if let Err(e) = execute(args) {
80        eprintln!("molt: {e}");
81        process::exit(1);
82    }
83}
84
85fn execute(args: Args) -> anyhow::Result<()> {
86    // User-specified nD-nodeset fold axis for output display (#356).
87    // An empty axis list behaves like Python's empty DEFAULTS.fold_axis
88    // tuple (falsy → fold along all axes).
89    let fold_axes = match &args.axis {
90        Some(axis) => {
91            Some(fold::parse_fold_axis(axis).map_err(|e| anyhow::anyhow!("Parse error: {e}"))?)
92        }
93        None => None,
94    };
95    let fold_axis: Option<&[i64]> = fold_axes.as_deref().filter(|v| !v.is_empty());
96
97    let stdout = io::stdout();
98    let mut out = stdout.lock();
99
100    // Read all input lines
101    let lines = read_input(&args)?;
102
103    if args.dshbak {
104        run_dshbak(&lines, &args, fold_axis, &mut out)?;
105    } else if args.line {
106        run_line_mode(&lines, &args, &mut out)?;
107    } else {
108        // Default: pass through with node labels
109        run_passthrough(&lines, &mut out)?;
110    }
111
112    Ok(())
113}
114
115/// Read input from files or stdin.
116fn read_input(args: &Args) -> anyhow::Result<Vec<String>> {
117    let mut lines = Vec::new();
118
119    if args.files.is_empty() {
120        let stdin = io::stdin();
121        for line in stdin.lock().lines() {
122            lines.push(line?);
123        }
124    } else {
125        for path in &args.files {
126            let file = std::fs::File::open(path)?;
127            let reader = io::BufReader::new(file);
128            for line in reader.lines() {
129                lines.push(line?);
130            }
131        }
132    }
133
134    Ok(lines)
135}
136
137/// Parse a line into (key, message) using the separator.
138///
139/// Like upstream clubak, the key is stripped but the content is kept
140/// verbatim (`line.split(separator, 1)`), so a leading space after the
141/// separator is preserved in gathered output.
142fn parse_line<'a>(line: &'a str, separator: &str) -> Option<(&'a str, &'a str)> {
143    let idx = line.find(separator)?;
144    let key = line[..idx].trim();
145    let msg = &line[idx + separator.len()..];
146    if key.is_empty() {
147        return None;
148    }
149    Some((key, msg))
150}
151
152/// Expand an input key into node names according to `--interpret-keys`.
153///
154/// Ports clubak's key handling: in `always`/`auto` mode the key is parsed
155/// as a NodeSet (nodeset keys expand to several nodes); in `auto` mode a
156/// parse failure permanently switches interpretation off (`enable_nodeset_key
157/// = False` upstream), here modeled by flipping the mode in place.
158fn expand_key(key: &str, mode: &mut InterpretKeys) -> anyhow::Result<Vec<String>> {
159    if *mode == InterpretKeys::Never {
160        return Ok(vec![key.to_string()]);
161    }
162    match NodeSet::parse(key) {
163        Ok(ns) => {
164            let nodes: Vec<String> = ns.iter().collect();
165            if nodes.is_empty() {
166                Ok(vec![key.to_string()])
167            } else {
168                Ok(nodes)
169            }
170        }
171        Err(e) => {
172            if *mode == InterpretKeys::Always {
173                return Err(anyhow::anyhow!("Parse error: {e}"));
174            }
175            // auto => switch off
176            *mode = InterpretKeys::Never;
177            Ok(vec![key.to_string()])
178        }
179    }
180}
181
182/// dshbak mode: group nodes with identical output.
183///
184/// One header block per distinct full output buffer, like upstream
185/// `print_gather(nodeset, tree[node])`. Groups sort like upstream
186/// (`nodeset_cmpkey`): larger groups first, then by first node name.
187/// Headers fold with the nD-aware folder honoring the `--axis` constraint
188/// (#356) and carry the ` (N)` node count like upstream
189/// `Display.format_header()`.
190fn run_dshbak(
191    lines: &[String],
192    args: &Args,
193    fold_axis: Option<&[i64]>,
194    out: &mut impl Write,
195) -> anyhow::Result<()> {
196    let mut interpret = args.interpret_keys;
197
198    // Aggregate per-node message buffers (joined with '\n' like MsgTree).
199    let mut node_bufs: HashMap<String, Vec<Vec<u8>>> = HashMap::new();
200    for line in lines {
201        if let Some((key, msg)) = parse_line(line, &args.separator) {
202            for node in expand_key(key, &mut interpret)? {
203                node_bufs
204                    .entry(node)
205                    .or_default()
206                    .push(msg.as_bytes().to_vec());
207            }
208        }
209    }
210
211    // Group nodes by identical full output buffer.
212    let mut groups: HashMap<Vec<u8>, Vec<String>> = HashMap::new();
213    for (node, msgs) in node_bufs {
214        groups.entry(msgs.join(&b"\n"[..])).or_default().push(node);
215    }
216
217    let mut groups: Vec<(Vec<u8>, Vec<String>)> = groups.into_iter().collect();
218    for (_buf, names) in groups.iter_mut() {
219        names.sort();
220    }
221    // larger nodeset first, then sorted by first node index
222    groups.sort_by(|a, b| b.1.len().cmp(&a.1.len()).then_with(|| a.1[0].cmp(&b.1[0])));
223
224    for (message, names) in &groups {
225        let folded = if interpret == InterpretKeys::Never {
226            // raw keys: no nodeset folding, --axis is a no-op
227            names.join(",")
228        } else {
229            fold::fold_nodes(names, fold_axis)
230        };
231        display::print_gathered_header(&folded, names.len(), out)?;
232        out.write_all(message)?;
233        writeln!(out)?;
234    }
235
236    Ok(())
237}
238
239/// Line mode: output ordered by node name.
240fn run_line_mode(lines: &[String], args: &Args, out: &mut impl Write) -> anyhow::Result<()> {
241    // Collect per-node output, then sort by node name
242    let mut node_output: std::collections::BTreeMap<String, Vec<String>> =
243        std::collections::BTreeMap::new();
244
245    let mut interpret = args.interpret_keys;
246    for line in lines {
247        if let Some((key, msg)) = parse_line(line, &args.separator) {
248            for node in expand_key(key, &mut interpret)? {
249                node_output.entry(node).or_default().push(msg.to_string());
250            }
251        }
252    }
253
254    for (node, msgs) in &node_output {
255        for msg in msgs {
256            display::print_line_with_label(node, msg, out)?;
257        }
258    }
259
260    Ok(())
261}
262
263/// Passthrough: just print lines as-is.
264fn run_passthrough(lines: &[String], out: &mut impl Write) -> anyhow::Result<()> {
265    for line in lines {
266        writeln!(out, "{line}")?;
267    }
268    Ok(())
269}