diff --git a/Changelog.md b/Changelog.md index 9240055..74588ba 100644 --- a/Changelog.md +++ b/Changelog.md @@ -1,6 +1,16 @@ # Changelog - 0.6.0 + - `each --threads=N` runs up to N commands in parallel, as a threadpool. Each command's output + is buffered and written in one piece, so parallel commands can't mangle each other's output. + Results are printed in the order the commands finish; `each --sequential` prints them in the + order of the input instead, waiting for earlier lines. + - `each` draws a progress bar on stderr while more than one thread is running, leaving stdout + free for piping + - **breaking**: a command failing under `each` no longer aborts the run. The failure is reported + on stderr, the remaining lines still run, and `string` exits 1 at the end. + - errors are now printed instead of being swallowed by a silent `exit(1)`, and the output of a + failed command is actually included in the message (it was always empty before) - **breaking**: `contains`, `starts-with` and `ends-with` are no longer predicates over the whole input. Instead of exiting 0/1 they now filter the input line by line and print every matching line. Shell conditions relying on the exit code (`string contains foo && ...`) need to be diff --git a/README.md b/README.md index 48035fd..84fa230 100644 --- a/README.md +++ b/README.md @@ -12,21 +12,24 @@ Cli for common string operations. Takes input from stdin. Usage: string Commands: - case Transform upper- or lowercase - reverse Reverse order of lines - substr Extract part of a given string - split Split up a string by a separator and print the parts on separate lines - length Returns the length the input string - replace Replace all matching characters - line Pick a single line by index - interleave Interleave input and only print every nth line - distinct Output the set of input strings without repetitions, in order - trim Trim whitespace on lines and ignore empty ones - chars Prints all chars on separate lines - template Useful for templating, replace sections of input with the output of a shell command or script - Map each line of input to a subcommand. - each - help Print this message or the help of the given subcommand(s) + case Transform upper- or lowercase + reverse Reverse order of lines + substr Extract part of a given string + split Split up a string by a separator and print the parts on separate lines + join Join lines with a separator into a single string + contains Print all lines containing the given string + starts-with Print all lines starting with the given prefix, ignoring leading whitespace + ends-with Print all lines ending with the given suffix, ignoring trailing whitespace + length Returns the length the input string + replace Replace all matching characters + line Pick a single line by index + interleave Interleave input and only print every nth line + distinct Output the set of input strings without repetitions, in order + trim Trim whitespace on lines and ignore empty ones + chars Prints all chars on separate lines + template Useful for templating, replace sections of input with the output of a shell command or script + each Map each line of input to a subcommand + help Print this message or the help of the given subcommand(s) Options: -h, --help Print help @@ -40,9 +43,11 @@ This is mostly because there are thousands of ways to do the tasks `shell-string More than anything I hated finding some solution for file templating over and over again. I wrote `shell-string` to never again have to think about what the best way of templating a file is. It's always this, period. -## Why is `shell-string` good for templating files? +## Template Files +`shell-string` is good for templating files. -Because you practically have no restrictions. +It's a very simple and clean solution where +you practically have no restrictions. You need to just drop in some environment variables? Easy, just write `{{ echo $MY_VAR }}` into the template. Is complex logic needed? You could write `{{ console.log(crazyStuff()) }}` and you're golden. Just execute with `--shell=node`. You want to use `haskell` in your template files? Use `--shell=ghci`! @@ -81,6 +86,65 @@ which means - `| string template`: The `|` means "don't print this in a terminal, pipe it to another programm" and that programm is `string` in `template` mode. - `> deployment.yaml`: Write the output of this into a file called `deployment.yaml`. If the file existed, empty it beforehand. +## Use It As A Threadpool + +`string each` runs a command once per line of input. Waiting for those commands one after another +is pure waste whenever they sit on the network instead of the CPU, so `--threads` turns `each` into +a threadpool: + +```sh +cat urls.txt | string each --threads=12 -- curl -s {} +``` + +Twelve requests are in flight at any moment, and `string` keeps the pool full: whenever a command +finishes, the next line is handed to the thread that just became free. Nothing is scheduled up +front, so one slow url doesn't leave eleven threads idling. + +The thing that makes this usable rather than a mess is that **output is never interleaved**. The +output of a command is collected in full and written in one piece, so twelve `curl`s can't scribble +over each other halfway through a line. What you get is the same output you'd get from a sequential +run, just sooner. + +### Watching it work + +Because output is buffered, a long run would otherwise look like a hung terminal. So while more than +one thread is running, a progress bar is drawn — on stderr, never stdout: + +``` +[####################----------] 812/1200 (67%) +``` + +That split is deliberate: `... | string each -t 12 -- curl -s {} > results.txt` shows you the bar in +the terminal while `results.txt` receives nothing but results. + +### Ordering + +Results appear in the order the commands _finish_, which is what you want when you're watching them +come in. When you'd rather have them line up with your input, ask for it: + +```sh +cat urls.txt | string each --threads=12 --sequential -- curl -s {} +``` + +`--sequential` holds finished results back until every earlier line is done, so the output matches +the input line for line. The work still runs on all twelve threads — only the printing waits. + +### When something fails + +One broken url shouldn't throw away the other 1199 responses. A command that exits non-zero is +reported on stderr, naming the line it came from, and the run carries on: + +``` +line 47: error executing command `curl`. +Process terminated with exit code exit status: 6. +Program output: +curl: (6) Could not resolve host: exmaple.com +``` + +At the very end `string` prints how many commands failed and exits 1, so `&&` in a script still +does the right thing — you just get all the successful output, and a full list of what went wrong, +instead of everything stopping at the first problem. + ## Installation Given cargo is installed on your machine execute diff --git a/src/cli.rs b/src/cli.rs index fbc598f..0fde26b 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -118,6 +118,13 @@ pub enum StringCommand { /// Name of placeholder in command to be replaced with input. Default is "{}" #[arg(short = 'v', long = "var", default_value = "{}")] var: String, + /// Number of commands to run in parallel. The default of 1 runs them one after another + #[arg(short = 't', long = "threads", default_value_t = 1)] + threads: usize, + /// Print the results in the order of the input, instead of the order they finish in. + /// Waits for earlier lines, so a single slow command holds back everything behind it + #[arg(long = "sequential", default_value_t = false)] + sequential: bool, /// Command to be executed. Pass "string each -- " so you can pass flags to the command. command: Vec, }, diff --git a/src/exec.rs b/src/exec.rs index 86bd4c4..d0f1a80 100644 --- a/src/exec.rs +++ b/src/exec.rs @@ -1,39 +1,63 @@ +use anyhow::{bail, Context}; use std::{io::Write, process::Stdio}; -pub fn execute(command: &[String], stdin_text: Option<&str>) -> String { +pub fn execute(command: &[String], stdin_text: Option<&str>) -> anyhow::Result { let command_name = &command[0]; - let mut command = std::process::Command::new(&command[0]) + let mut child = std::process::Command::new(&command[0]) .args(&command[1..]) .stdin(Stdio::piped()) .stdout(Stdio::piped()) + // captured, so the output of parallel commands can't smear over each other or the + // progress bar. It is printed as part of the error message if the command fails. + .stderr(Stdio::piped()) .spawn() - .expect(&format!("failed to spawn process {:?}", command)); - - if let Some(stdin_text) = stdin_text { - let stdin = command - .stdin - .as_mut() - .expect("failed to open stdin of command"); - stdin - .write_all(stdin_text.as_bytes()) - .expect("failed to pipe command into shell"); - } + .with_context(|| format!("failed to spawn process {:?}", command))?; + + let output = match stdin_text { + None => child.wait_with_output(), + // Feeding the command from a second thread, because writing the input and reading the + // output have to happen at the same time. Filling up the input pipe while the command + // is blocked on an output pipe nobody reads deadlocks both sides. + Some(stdin_text) => { + let mut stdin = child + .stdin + .take() + .context("failed to open stdin of command")?; + + let (feeding, output) = std::thread::scope(|scope| { + // dropping `stdin` when this thread ends closes the pipe, which is how the + // command gets to see the end of its input + let feeder = scope.spawn(move || stdin.write_all(stdin_text.as_bytes())); + let output = child.wait_with_output(); + + (feeder.join(), output) + }); - let output = command - .wait_with_output() - .expect("failed to aquire programm output"); + match feeding { + Err(_) => bail!("panic while piping input into command `{}`", command_name), + // a command is free to stop reading early, like `head` does. That closes the + // pipe under us, which is not an error of its own — the exit code decides. + Ok(Err(e)) if e.kind() == std::io::ErrorKind::BrokenPipe => {} + Ok(result) => result.context("failed to pipe command into shell")?, + } + + output + } + }; + + let output = output.context("failed to aquire programm output")?; let status: std::process::ExitStatus = output.status; if !status.success() { - eprintln!("error executing command `{}`.\nProcess terminated with exit code {}.\nProgram output:\n{}", + bail!( + "error executing command `{}`.\nProcess terminated with exit code {}.\nProgram output:\n{}", command_name, status, - String::from_utf8(output.stderr).unwrap() + String::from_utf8_lossy(&output.stderr) ); - std::process::exit(1); } - String::from_utf8(output.stdout).expect("programm output was not valid utf-8") + String::from_utf8(output.stdout).context("programm output was not valid utf-8") } #[cfg(test)] @@ -43,7 +67,7 @@ mod test { #[test] fn exec1() { let input = "printf hello"; - let result = execute(&[String::from("sh")], Some(input)); + let result = execute(&[String::from("sh")], Some(input)).unwrap(); let expected = "hello"; assert_eq!(expected, result); @@ -51,9 +75,65 @@ mod test { #[test] fn exec2() { - let result = execute(&[String::from("printf"), String::from("hello")], None); + let result = execute(&[String::from("printf"), String::from("hello")], None).unwrap(); let expected = "hello"; assert_eq!(expected, result); } + + #[test] + fn exec_failure_reports_stderr() { + let command = [ + String::from("sh"), + String::from("-c"), + String::from("echo boom >&2; exit 3"), + ]; + + let err = execute(&command, None).expect_err("command exits 3, so this must fail"); + let message = format!("{err:#}"); + + assert!( + message.contains("boom"), + "stderr is missing from: {}", + message + ); + assert!( + message.contains('3'), + "exit code is missing from: {}", + message + ); + } + + /// Both pipes hold roughly 64kb, so echoing a megabyte back deadlocks any implementation + /// that writes all of stdin before it starts reading stdout. + #[test] + fn large_input_does_not_deadlock() { + let input = "abcdefgh\n".repeat(128 * 1024); + let result = execute(&[String::from("cat")], Some(&input)).unwrap(); + + assert_eq!(result.len(), input.len()); + assert_eq!(result, input); + } + + /// a command that never reads its input leaves us writing into a closed pipe + #[test] + fn input_ignored_by_the_command_is_not_an_error() { + let input = "abcdefgh\n".repeat(128 * 1024); + let command = [ + String::from("sh"), + String::from("-c"), + String::from("echo done"), + ]; + + let result = execute(&command, Some(&input)).unwrap(); + + assert_eq!(result, "done\n"); + } + + #[test] + fn exec_failure_of_missing_binary() { + let command = [String::from("definitely-not-an-existing-binary-42")]; + + assert!(execute(&command, None).is_err()); + } } diff --git a/src/main.rs b/src/main.rs index 39fd0bb..f8db7ba 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,23 +1,26 @@ mod cli; mod exec; +mod pool; +mod progressbar; mod templating; mod util; use templating::template; +use anyhow::bail; use clap::Parser; use itertools::join; use crate::exec::execute; +use crate::progressbar::ProgressBar; fn main() { let command: cli::StringCommand = cli::StringCommand::parse(); let input = util::stdin_as_string(); let mut output = std::io::stdout(); - let result = perform_command(command, input, &mut output); - - if result.is_err() { + if let Err(e) = perform_command(command, input, &mut output) { + eprintln!("{e:#}"); std::process::exit(1); } } @@ -45,6 +48,10 @@ mod tests { buffer: Vec::with_capacity(128), } } + + fn text(&self) -> String { + String::from_utf8_lossy(&self.buffer).into_owned() + } } impl std::io::Write for TestWriter { @@ -353,6 +360,66 @@ mod tests { } } + fn each(threads: usize, sequential: bool, command: &[&str]) -> StringCommand { + Each { + stdin: false, + var: "{}".into(), + threads, + sequential, + command: command.iter().map(|s| s.to_string()).collect(), + } + } + + #[test] + fn each_sequential_keeps_input_order() { + // the first line sleeps longest, so completion order is the reverse of the input + let command = ["sh", "-c", "sleep 0.{}; echo {}"]; + + for threads in [1, 4] { + let mut writer = TestWriter::new(); + perform_command( + each(threads, true, &command), + "3\n2\n1\n".into(), + &mut writer, + ) + .unwrap(); + + assert_eq!(writer, "3\n2\n1\n"); + } + } + + #[test] + fn each_unordered_returns_every_result() { + let command = ["sh", "-c", "sleep 0.{}; echo {}"]; + + let mut writer = TestWriter::new(); + perform_command(each(4, false, &command), "3\n2\n1\n".into(), &mut writer).unwrap(); + + let text = writer.text(); + let mut lines: Vec<&str> = text.lines().collect(); + lines.sort(); + + assert_eq!(lines, ["1", "2", "3"]); + } + + #[test] + fn each_reports_failures_but_finishes_the_rest() { + // "two" is not a number, so `test` exits non-zero for that line only + let command = ["sh", "-c", "test {} -gt 0 && echo {}"]; + + let mut writer = TestWriter::new(); + let res = perform_command(each(1, true, &command), "1\ntwo\n3\n".into(), &mut writer); + + let err = res.expect_err("one command failed, so the run must fail"); + assert!( + format!("{:#}", err).contains("1 of 3"), + "unexpected error: {:#}", + err + ); + // the lines around the failure still made it through + assert_eq!(writer, "1\n3\n"); + } + #[test] fn trim() { let input = " @@ -478,7 +545,7 @@ fn perform_command( end, raw_output, } => { - let result = template(&input, &shell, &begin, &end, !raw_output); + let result = template(&input, &shell, &begin, &end, !raw_output)?; writeln!(output, "{}", result)?; } Chars => { @@ -489,18 +556,50 @@ fn perform_command( Each { stdin, var, + threads, + sequential, ref command, } => { - for line in input.lines() { - let command: Vec<_> = command.iter().map(|s| s.replace(&var, line)).collect(); - let input = if stdin { Some(line) } else { None }; - - let result = execute(&command, input); - if result.ends_with("\n") { - write!(output, "{}", result)?; - } else { - writeln!(output, "{}", result)?; - } + let lines: Vec<&str> = input.lines().collect(); + let mut progress = ProgressBar::new(lines.len(), threads > 1); + let mut failures = 0; + + pool::for_each( + &lines, + threads, + sequential, + // on the worker threads + |line| { + let command: Vec<_> = command.iter().map(|s| s.replace(&var, line)).collect(); + let input = if stdin { Some(*line) } else { None }; + + execute(&command, input) + }, + // back on this thread, so nothing is ever written half way through + |index, result| { + match result { + Ok(result) => { + if result.ends_with("\n") { + write!(output, "{}", result)?; + } else { + writeln!(output, "{}", result)?; + } + } + Err(e) => { + failures += 1; + eprintln!("line {}: {e:#}", index + 1); + } + } + + // flush before redrawing, so the bar stays the last thing on the screen + output.flush()?; + progress.tick(); + Ok(()) + }, + )?; + + if failures > 0 { + bail!("{} of {} commands failed", failures, lines.len()); } } }; diff --git a/src/pool.rs b/src/pool.rs new file mode 100644 index 0000000..112a90a --- /dev/null +++ b/src/pool.rs @@ -0,0 +1,216 @@ +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::mpsc; + +/// Run `job` for every item, on up to `threads` worker threads. +/// +/// `on_result` is invoked on the *calling* thread, so it may hold on to things that aren't `Send`, +/// like the output stream. Results are handed to it in the order the jobs complete, or, when +/// `ordered` is set, in the order of `items` — which means a result has to wait in memory until +/// every earlier item finished. +pub fn for_each( + items: &[T], + threads: usize, + ordered: bool, + job: impl Fn(&T) -> R + Sync, + mut on_result: impl FnMut(usize, R) -> anyhow::Result<()>, +) -> anyhow::Result<()> +where + T: Sync, + R: Send, +{ + if items.is_empty() { + return Ok(()); + } + + let threads = threads.max(1).min(items.len()); + + // the next item to be picked up. Handing out one item at a time keeps all workers busy, + // even when some items take much longer than others. + let cursor = AtomicUsize::new(0); + let cancelled = AtomicBool::new(false); + let (sender, receiver) = mpsc::channel::<(usize, R)>(); + + std::thread::scope(|scope| { + for _ in 0..threads { + let sender = sender.clone(); + let cursor = &cursor; + let cancelled = &cancelled; + let job = &job; + + scope.spawn(move || loop { + if cancelled.load(Ordering::Relaxed) { + break; + } + + let index = cursor.fetch_add(1, Ordering::Relaxed); + if index >= items.len() { + break; + } + + if sender.send((index, job(&items[index]))).is_err() { + break; + } + }); + } + + // the workers hold their own senders. Dropping ours lets the loop below end + // as soon as the last worker is done. + drop(sender); + + let mut pending = BTreeMap::new(); + let mut next = 0; + + for (index, value) in receiver { + let outcome = if ordered { + pending.insert(index, value); + + let mut outcome = Ok(()); + while let Some(value) = pending.remove(&next) { + outcome = on_result(next, value); + next += 1; + + if outcome.is_err() { + break; + } + } + outcome + } else { + on_result(index, value) + }; + + if let Err(e) = outcome { + // nothing consumes results anymore, so stop handing out work + cancelled.store(true, Ordering::Relaxed); + return Err(e); + } + } + + Ok(()) + }) +} + +#[cfg(test)] +mod test { + use super::*; + use std::time::Duration; + + #[test] + fn every_item_is_processed_exactly_once() { + let items: Vec = (0..100).collect(); + let mut results = Vec::new(); + + for_each( + &items, + 8, + false, + |item| item * 2, + |_, value| { + results.push(value); + Ok(()) + }, + ) + .unwrap(); + + results.sort(); + let expected: Vec = (0..100).map(|i| i * 2).collect(); + assert_eq!(results, expected); + } + + /// items finish in reverse order of their input position, so anything but a deliberate + /// re-ordering would show up here. + #[test] + fn ordered_emits_in_input_order() { + let items: Vec = (0..10).collect(); + let mut indices = Vec::new(); + + for_each( + &items, + 10, + true, + |item| { + std::thread::sleep(Duration::from_millis((10 - item) * 5)); + *item + }, + |index, value| { + indices.push((index, value)); + Ok(()) + }, + ) + .unwrap(); + + let expected: Vec<(usize, u64)> = (0..10).map(|i| (i, i as u64)).collect(); + assert_eq!(indices, expected); + } + + #[test] + fn unordered_emits_everything_regardless_of_order() { + let items: Vec = (0..10).collect(); + let mut values = Vec::new(); + + for_each( + &items, + 10, + false, + |item| { + std::thread::sleep(Duration::from_millis((10 - item) * 5)); + *item + }, + |_, value| { + values.push(value); + Ok(()) + }, + ) + .unwrap(); + + values.sort(); + assert_eq!(values, (0..10).collect::>()); + } + + #[test] + fn single_thread_keeps_input_order() { + let items: Vec = (0..20).collect(); + let mut indices = Vec::new(); + + for_each( + &items, + 1, + false, + |item| *item, + |index, _| { + indices.push(index); + Ok(()) + }, + ) + .unwrap(); + + assert_eq!(indices, (0..20).collect::>()); + } + + #[test] + fn a_failing_consumer_stops_the_run() { + let items: Vec = (0..1000).collect(); + let mut seen = 0; + + let result = for_each( + &items, + 4, + false, + |item| *item, + |_, _| { + seen += 1; + anyhow::bail!("stop right there") + }, + ); + + assert!(result.is_err()); + assert_eq!(seen, 1); + } + + #[test] + fn empty_input_is_fine() { + let items: Vec = Vec::new(); + + for_each(&items, 8, false, |item| *item, |_, _| Ok(())).unwrap(); + } +} diff --git a/src/progressbar.rs b/src/progressbar.rs new file mode 100644 index 0000000..80103ab --- /dev/null +++ b/src/progressbar.rs @@ -0,0 +1,59 @@ +const WIDTH: usize = 30; + +/// A progress bar +/// Rendered on stderr, +/// so it stays out of the way of the results on stdout. +pub struct ProgressBar { + items: usize, + done: usize, + enabled: bool, +} + +impl ProgressBar { + pub fn new(items: usize, enabled: bool) -> Self { + ProgressBar { + items, + done: 0, + enabled: enabled && items > 0, + } + } + + /// Count one finished item and redraw the bar. + pub fn tick(&mut self) { + self.done += 1; + self.draw(); + } + + fn draw(&mut self) { + if !self.enabled { + return; + } + + let done = self.done.min(self.items); + let filled = done * WIDTH / self.items; + + // eprint! rather than a raw stderr handle, so the test harness captures it + eprintln!( + "\r[{}{}] {}/{}", + "#".repeat(filled), + "-".repeat(WIDTH - filled), + done, + self.items, + ); + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn counts_up_to_completion() { + let mut progress = ProgressBar::new(4, true); + for _ in 0..4 { + progress.tick(); + } + + assert_eq!(progress.done, 4); + } +} diff --git a/src/templating.rs b/src/templating.rs index b285498..23f855b 100644 --- a/src/templating.rs +++ b/src/templating.rs @@ -7,10 +7,15 @@ use nom::{ IResult, }; -pub fn template(input: &str, shell: &[String], begin: &str, end: &str, trim: bool) -> String { +pub fn template( + input: &str, + shell: &[String], + begin: &str, + end: &str, + trim: bool, +) -> anyhow::Result { if shell.len() == 0 { - eprintln!("must specify a shell"); - std::process::exit(1); + anyhow::bail!("must specify a shell"); } // 1 split text content and commands // 2 map commands to their execution output @@ -25,13 +30,13 @@ pub fn template(input: &str, shell: &[String], begin: &str, end: &str, trim: boo for c in ast { buffer.push_str(c.text); if let Some(cmd) = c.command { - let output = execute(shell, Some(cmd)); + let output = execute(shell, Some(cmd))?; let output = if trim { output.trim() } else { &output }; buffer.push_str(output); } } - buffer + Ok(buffer) } #[derive(PartialEq, Debug)] @@ -104,7 +109,7 @@ mod test { #[test] fn template1() { let input = "hello (echo world)"; - let result = template(input, &["sh".to_string()], "(", ")", true); + let result = template(input, &["sh".to_string()], "(", ")", true).unwrap(); let expected = "hello world"; assert_eq!(expected, result); @@ -113,7 +118,7 @@ mod test { #[test] fn template2() { let input = "Hey (echo VSauce), (echo Michael) here!"; - let result = template(input, &["sh".to_string()], "(", ")", true); + let result = template(input, &["sh".to_string()], "(", ")", true).unwrap(); let expected = "Hey VSauce, Michael here!"; assert_eq!(expected, result); @@ -122,7 +127,7 @@ mod test { #[test] fn template3() { let input = "Hey { echo VSauce }, { echo Michael } here!"; - let result = template(input, &["sh".to_string()], "{", "}", true); + let result = template(input, &["sh".to_string()], "{", "}", true).unwrap(); let expected = "Hey VSauce, Michael here!"; assert_eq!(expected, result); @@ -131,7 +136,7 @@ mod test { #[test] fn template4() { let input = "complex calculation: ^console.log(14)^"; - let result = template(input, &["node".to_string()], "^", "^", true); + let result = template(input, &["node".to_string()], "^", "^", true).unwrap(); let expected = "complex calculation: 14"; assert_eq!(expected, result);