Merge pull request #82 from aaronriekenberg/pipe_mode

Add pipe mode, block size.

Author: aaronriekenberg

Cargo.lock

@@ -15,6 +15,7 @@ repository = "https://github.com/aaronriekenberg/rust-parallel"
 
 [dependencies]
 anyhow = "1"
+bytesize = "2"
 clap = { version = "4", features = ["derive"] }
 indicatif = "0.18"
 itertools = "0.14"

Cargo.toml

@@ -19,6 +19,7 @@ Listed in [Awesome Rust - utilities](https://github.com/rust-unofficial/awesome-
 Similar interface to [GNU Parallel](https://www.gnu.org/software/parallel/parallel_examples.html) or [xargs](https://man7.org/linux/man-pages/man1/xargs.1.html) plus useful features:
 * More than 10x faster than GNU Parallel [in benchmarks](https://github.com/aaronriekenberg/rust-parallel/wiki/Benchmarks)
 * Run commands from [stdin](https://github.com/aaronriekenberg/rust-parallel/wiki/Manual#commands-from-stdin), [input files](https://github.com/aaronriekenberg/rust-parallel/wiki/Manual#reading-multiple-inputs), or [`:::` arguments](https://github.com/aaronriekenberg/rust-parallel/wiki/Manual#commands-from-arguments)
+* [Pipe mode](https://github.com/aaronriekenberg/rust-parallel/wiki/Manual#pipe-mode) to read blocks from stdin and pass to parallel commands automatically.
 * Automatic parallelism to all cpus, or [configure manually](https://github.com/aaronriekenberg/rust-parallel/wiki/Manual#parallelism)
 * Transform inputs with [variables](https://github.com/aaronriekenberg/rust-parallel/wiki/Manual#automatic-variables) or [regular expressions](https://github.com/aaronriekenberg/rust-parallel/wiki/Manual#regular-expression)
 * Prevent [output interleaving](https://github.com/aaronriekenberg/rust-parallel/wiki/Output-Interleaving) and maintain input order with `-k`/`--keep-order`

README.md

@@ -15,6 +15,7 @@ echo '
 1. [Commands from stdin](#commands-from-stdin)
 1. [Command and initial arguments on command line](#command-and-initial-arguments-on-command-line)
 1. [Reading multiple inputs](#reading-multiple-inputs)
+1. [Pipe Mode](#pipe-mode)
 1. [Parallelism](#parallelism)
 1. [Keep Output Order](#keep-output-order)
 1. [Dry run](#dry-run)
@@ -142,6 +143,30 @@ rm -f test
 
 echo '```'
 
+echo '## Pipe Mode
+
+The `--pipe` option can be used to enable pipe mode.
+
+In pipe mode input from stdin is split into blocks and each block is passed to a separate instance of the command via stdin.  Command instances are run in parallel.  
+
+The default block size is 1 MiB, which can be changed with the `--block-size` option.
+
+Here we use `--pipe` to run `wc -l`
+'
+
+echo '```'
+echo '$ cat /usr/share/dict/words | rust-parallel --pipe wc -l'
+cat /usr/share/dict/words | $RUST_PARALLEL --pipe wc -l
+echo '```'
+
+echo 'Here we use pipe mode with with a smaller block size of 500 KiB:'
+
+echo '```'
+echo '$ cat /usr/share/dict/words | rust-parallel --pipe --block-size=500KiB wc -l'
+cat /usr/share/dict/words | $RUST_PARALLEL --pipe --block-size=500KiB wc -l
+echo '```'
+
+
 echo '
 ## Parallelism
 

scripts/generate_manual.sh

@@ -34,6 +34,7 @@ impl Command {
             cmd = ?self.command_and_args.command_path,
             args = ?self.command_and_args.args,
             line = %self.input_line_number,
+            stdin = %self.command_and_args.stdin,
             child_pid,
         ),
         level = "debug")]
@@ -42,13 +43,17 @@ impl Command {
 
         let command_metrics = &context.command_metrics;
 
-        let OwnedCommandAndArgs { command_path, args } = &self.command_and_args;
+        let OwnedCommandAndArgs {
+            command_path,
+            args,
+            stdin,
+        } = &self.command_and_args;
 
         command_metrics.increment_commands_run();
 
         let child_process = match context
             .child_process_factory
-            .spawn(command_path, args)
+            .spawn(command_path, args, stdin.clone())
             .await
         {
             Err(e) => {
@@ -91,8 +96,8 @@ impl std::fmt::Display for Command {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
             f,
-            "cmd={:?},args={:?},line={}",
-            self.command_and_args.command_path, self.command_and_args.args, self.input_line_number,
+            "{},line={}",
+            self.command_and_args, self.input_line_number
         )
     }
 }
@@ -177,10 +182,7 @@ impl CommandService {
             return Ok(());
         };
 
-        let command_and_args = OwnedCommandAndArgs {
-            command_path,
-            args: command_and_args.args,
-        };
+        let command_and_args = command_and_args.with_command_path(command_path);
 
         self.spawn_command(command_and_args, input_line_number)
             .await?;

src/command.rs

@@ -1,9 +1,13 @@
+use bytesize::ByteSize;
+
 use clap::{Parser, ValueEnum};
 
 use tokio::sync::OnceCell;
 
 use tracing::debug;
 
+use std::str::FromStr;
+
 pub const COMMANDS_FROM_ARGS_SEPARATOR: &str = ":::";
 
 /// Execute commands in parallel
@@ -81,6 +85,14 @@ pub struct CommandLineArgs {
     #[arg(short, long)]
     pub keep_order: bool,
 
+    /// Use pipe input mode
+    #[arg(long)]
+    pub pipe: bool,
+
+    /// Block size for pipe input mode, defaults to 1MiB
+    #[arg(long, default_value = "1MiB", value_parser = Self::parse_block_size_bytes)]
+    pub block_size: usize,
+
     /// Path to shell to use for shell mode
     #[arg(long, default_value = Self::default_shell())]
     pub shell_path: String,
@@ -138,6 +150,19 @@ impl CommandLineArgs {
         }
     }
 
+    fn parse_block_size_bytes(s: &str) -> Result<usize, String> {
+        let bytes =
+            ByteSize::from_str(s).map_err(|e| format!("could not parse size `{s}`: {e}"))?;
+
+        if bytes.0 == 0 {
+            return Err("size must be greater than 0".to_string());
+        }
+
+        let bytes =
+            usize::try_from(bytes.0).map_err(|e| format!("size `{s}` is too large: {e}"))?;
+        Ok(bytes)
+    }
+
     fn default_shell() -> &'static str {
         if cfg!(unix) {
             "/bin/bash"

src/command_line_args.rs

@@ -1,14 +1,36 @@
-use std::{collections::VecDeque, path::PathBuf};
+use bytesize::ByteSize;
 
-#[derive(Debug, Eq, PartialEq)]
+use std::{collections::VecDeque, path::PathBuf, sync::Arc};
+
+#[derive(Clone, Debug, Eq, PartialEq, Default)]
+pub struct StdinData(pub Option<Arc<String>>);
+
+impl std::fmt::Display for StdinData {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self.0 {
+            Some(s) => {
+                let size_string = ByteSize::b(s.len().try_into().unwrap_or_default()).to_string();
+                write!(f, "{size_string}")
+            }
+            None => write!(f, "None"),
+        }
+    }
+}
+
+#[derive(Debug, Eq, PartialEq, Default)]
 pub struct OwnedCommandAndArgs {
     pub command_path: PathBuf,
     pub args: Vec<String>,
+    pub stdin: StdinData,
 }
 
 impl std::fmt::Display for OwnedCommandAndArgs {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "cmd={:?},args={:?}", self.command_path, self.args)
+        write!(
+            f,
+            "cmd={:?},args={:?},stdin={}",
+            self.command_path, self.args, self.stdin
+        )
     }
 }
 
@@ -29,6 +51,7 @@ impl TryFrom<VecDeque<String>> for OwnedCommandAndArgs {
         Ok(Self {
             command_path: PathBuf::from(command),
             args: deque.into(),
+            stdin: StdinData(None),
         })
     }
 }
@@ -40,3 +63,15 @@ impl TryFrom<Vec<String>> for OwnedCommandAndArgs {
         Self::try_from(VecDeque::from(vec))
     }
 }
+
+impl OwnedCommandAndArgs {
+    pub fn with_command_path(mut self, command_path: PathBuf) -> Self {
+        self.command_path = command_path;
+        self
+    }
+
+    pub fn with_stdin(mut self, stdin: StdinData) -> Self {
+        self.stdin = stdin;
+        self
+    }
+}

src/common.rs

@@ -46,10 +46,25 @@ impl std::fmt::Display for Input {
     }
 }
 
+#[derive(Debug)]
+pub enum LineNumberOrRange {
+    Single(usize),
+    Range(usize, usize),
+}
+
+impl std::fmt::Display for LineNumberOrRange {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Single(line) => write!(f, "{line}"),
+            Self::Range(start, end) => write!(f, "{start}-{end}"),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct InputLineNumber {
     pub input: Input,
-    pub line_number: usize,
+    pub line_number: LineNumberOrRange,
 }
 
 impl std::fmt::Display for InputLineNumber {
@@ -59,18 +74,22 @@ impl std::fmt::Display for InputLineNumber {
 }
 
 enum InputList {
-    BufferedInputList(Vec<BufferedInput>),
+    Buffered(Vec<BufferedInput>),
 
     CommandLineArgs,
+
+    Pipe,
 }
 
 fn build_input_list(command_line_args: &'static CommandLineArgs) -> InputList {
-    if command_line_args.commands_from_args_mode() {
+    if command_line_args.pipe {
+        InputList::Pipe
+    } else if command_line_args.commands_from_args_mode() {
         InputList::CommandLineArgs
     } else if command_line_args.input_file.is_empty() {
-        InputList::BufferedInputList(vec![BufferedInput::Stdin])
+        InputList::Buffered(vec![BufferedInput::Stdin])
     } else {
-        InputList::BufferedInputList(
+        InputList::Buffered(
             command_line_args
                 .input_file
                 .iter()

src/input.rs

@@ -4,7 +4,7 @@ use tokio::io::{AsyncBufRead, AsyncBufReadExt, BufReader, Split};
 
 use crate::command_line_args::CommandLineArgs;
 
-use super::{BufferedInput, Input, InputLineNumber};
+use super::{BufferedInput, Input};
 
 type AsyncBufReadBox = Box<dyn AsyncBufRead + Unpin + Send>;
 
@@ -54,20 +54,19 @@ impl BufferedInputReader {
         }
     }
 
-    pub async fn next_segment(&mut self) -> anyhow::Result<Option<(InputLineNumber, Vec<u8>)>> {
+    pub async fn next_segment(&mut self) -> anyhow::Result<Option<(Input, usize, Vec<u8>)>> {
         let segment = self.split.next_segment().await?;
 
         match segment {
             None => Ok(None),
             Some(segment) => {
                 self.next_line_number += 1;
 
-                let input_line_number = InputLineNumber {
-                    input: Input::Buffered(self.buffered_input),
-                    line_number: self.next_line_number,
-                };
-
-                Ok(Some((input_line_number, segment)))
+                Ok(Some((
+                    Input::Buffered(self.buffered_input),
+                    self.next_line_number,
+                    segment,
+                )))
             }
         }
     }