exec: Execute batches before they get too long

Fixes #410.

6 files changed, 215 insertions, 81 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 7553fb8..339611f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -27,6 +27,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "08f9b8508dccb7687a1d6c4ce66b2b0ecef467c94667de27d8d7fe1f8d2a9cdc"
 
 [[package]]
+name = "argmax"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "932fb17af6e53a41ce7312f1ae1ba2a6f3f613fe36f38ad655b212906eb9657f"
+dependencies = [
+ "lazy_static",
+ "libc",
+ "nix 0.23.1",
+]
+
+[[package]]
 name = "atty"
 version = "0.2.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -134,7 +145,7 @@ version = "3.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b37feaa84e6861e00a1f5e5aa8da3ee56d605c9992d33e082786754828e20865"
 dependencies = [
- "nix",
+ "nix 0.24.1",
  "winapi",
 ]
 
@@ -171,6 +182,7 @@ version = "8.3.2"
 dependencies = [
  "ansi_term",
  "anyhow",
+ "argmax",
  "atty",
  "chrono",
  "clap",
@@ -185,7 +197,7 @@ dependencies = [
  "jemallocator",
  "libc",
  "lscolors",
- "nix",
+ "nix 0.24.1",
  "normpath",
  "num_cpus",
  "once_cell",
@@ -329,9 +341,9 @@ checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646"
 
 [[package]]
 name = "libc"
-version = "0.2.125"
+version = "0.2.126"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5916d2ae698f6de9bfb891ad7a8d65c09d232dc58cc4ac433c7da3b2fd84bc2b"
+checksum = "349d5a591cd28b49e1d1037471617a32ddcda5731b99419008085f72d5a53836"
 
 [[package]]
 name = "log"
@@ -358,6 +370,28 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "308cc39be01b73d0d18f82a0e7b2a3df85245f84af96fdddc5d202d27e47b86a"
 
 [[package]]
+name = "memoffset"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "nix"
+version = "0.23.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f866317acbd3a240710c63f065ffb1e4fd466259045ccb504130b7f668f35c6"
+dependencies = [
+ "bitflags",
+ "cc",
+ "cfg-if",
+ "libc",
+ "memoffset",
+]
+
+[[package]]
 name = "nix"
 version = "0.24.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 02014bd..7f94821 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -36,6 +36,7 @@ version_check = "0.9"
 
 [dependencies]
 ansi_term = "0.12"
+argmax = "0.3.0"
 atty = "0.2"
 ignore = "0.4.3"
 num_cpus = "1.13"
diff --git a/src/exec/command.rs b/src/exec/command.rs
index 935b6f4..a642a40 100644
--- a/src/exec/command.rs
+++ b/src/exec/command.rs
@@ -1,8 +1,9 @@
 use std::io;
 use std::io::Write;
-use std::process::Command;
 use std::sync::Mutex;
 
+use argmax::Command;
+
 use crate::error::print_error;
 use crate::exit_codes::ExitCode;
 
@@ -50,13 +51,18 @@ impl<'a> OutputBuffer<'a> {
 }
 
 /// Executes a command.
-pub fn execute_commands<I: Iterator<Item = Command>>(
+pub fn execute_commands<I: Iterator<Item = io::Result<Command>>>(
     cmds: I,
     out_perm: &Mutex<()>,
     enable_output_buffering: bool,
 ) -> ExitCode {
     let mut output_buffer = OutputBuffer::new(out_perm);
-    for mut cmd in cmds {
+    for result in cmds {
+        let mut cmd = match result {
+            Ok(cmd) => cmd,
+            Err(e) => return handle_cmd_error(None, e),
+        };
+
         // Spawn the supplied command.
         let output = if enable_output_buffering {
             cmd.output()
@@ -79,7 +85,7 @@ pub fn execute_commands<I: Iterator<Item = Command>>(
             }
             Err(why) => {
                 output_buffer.write();
-                return handle_cmd_error(&cmd, why);
+                return handle_cmd_error(Some(&cmd), why);
             }
         }
     }
@@ -87,12 +93,15 @@ pub fn execute_commands<I: Iterator<Item = Command>>(
     ExitCode::Success
 }
 
-pub fn handle_cmd_error(cmd: &Command, err: io::Error) -> ExitCode {
-    if err.kind() == io::ErrorKind::NotFound {
-        print_error(format!("Command not found: {:?}", cmd));
-        ExitCode::GeneralError
-    } else {
-        print_error(format!("Problem while executing command: {}", err));
-        ExitCode::GeneralError
+pub fn handle_cmd_error(cmd: Option<&Command>, err: io::Error) -> ExitCode {
+    match (cmd, err) {
+        (Some(cmd), err) if err.kind() == io::ErrorKind::NotFound => {
+            print_error(format!("Command not found: {:?}", cmd));
+            ExitCode::GeneralError
+        }
+        (_, err) => {
+            print_error(format!("Problem while executing command: {}", err));
+            ExitCode::GeneralError
+        }
     }
 }
diff --git a/src/exec/job.rs b/src/exec/job.rs
index 304792b..9b95ac2 100644
--- a/src/exec/job.rs
+++ b/src/exec/job.rs
@@ -62,17 +62,6 @@ pub fn batch(
                 None
             }
         });
-    if limit == 0 {
-        // no limit
-        return cmd.execute_batch(paths);
-    }
 
-    let mut exit_codes = Vec::new();
-    let mut peekable = paths.peekable();
-    while peekable.peek().is_some() {
-        let limited = peekable.by_ref().take(limit);
-        let exit_code = cmd.execute_batch(limited);
-        exit_codes.push(exit_code);
-    }
-    merge_exitcodes(exit_codes)
+    cmd.execute_batch(paths, limit)
 }
diff --git a/src/exec/mod.rs b/src/exec/mod.rs
index d140528..fc26da2 100644
--- a/src/exec/mod.rs
+++ b/src/exec/mod.rs
@@ -5,11 +5,14 @@ mod token;
 
 use std::borrow::Cow;
 use std::ffi::{OsStr, OsString};
+use std::io;
+use std::iter;
 use std::path::{Component, Path, PathBuf, Prefix};
-use std::process::{Command, Stdio};
+use std::process::Stdio;
 use std::sync::{Arc, Mutex};
 
 use anyhow::{bail, Result};
+use argmax::Command;
 use once_cell::sync::Lazy;
 use regex::Regex;
 
@@ -89,20 +92,117 @@ impl CommandSet {
         execute_commands(commands, &out_perm, buffer_output)
     }
 
-    pub fn execute_batch<I>(&self, paths: I) -> ExitCode
+    pub fn execute_batch<I>(&self, paths: I, limit: usize) -> ExitCode
     where
         I: Iterator<Item = PathBuf>,
     {
         let path_separator = self.path_separator.as_deref();
-        let mut paths = paths.collect::<Vec<_>>();
-        paths.sort();
-        for cmd in &self.commands {
-            let exit = cmd.generate_and_execute_batch(&paths, path_separator);
-            if exit != ExitCode::Success {
-                return exit;
+
+        let builders: io::Result<Vec<_>> = self
+            .commands
+            .iter()
+            .map(|c| CommandBuilder::new(c, limit))
+            .collect();
+
+        match builders {
+            Ok(mut builders) => {
+                for path in paths {
+                    for builder in &mut builders {
+                        if let Err(e) = builder.push(&path, path_separator) {
+                            return handle_cmd_error(Some(&builder.cmd), e);
+                        }
+                    }
+                }
+
+                for builder in &mut builders {
+                    if let Err(e) = builder.finish() {
+                        return handle_cmd_error(Some(&builder.cmd), e);
+                    }
+                }
+
+                ExitCode::Success
             }
+            Err(e) => handle_cmd_error(None, e),
+        }
+    }
+}
+
+/// Represents a multi-exec command as it is built.
+#[derive(Debug)]
+struct CommandBuilder {
+    pre_args: Vec<OsString>,
+    path_arg: ArgumentTemplate,
+    post_args: Vec<OsString>,
+    cmd: Command,
+    count: usize,
+    limit: usize,
+}
+
+impl CommandBuilder {
+    fn new(template: &CommandTemplate, limit: usize) -> io::Result<Self> {
+        let mut pre_args = vec![];
+        let mut path_arg = None;
+        let mut post_args = vec![];
+
+        for arg in &template.args {
+            if arg.has_tokens() {
+                path_arg = Some(arg.clone());
+            } else if path_arg == None {
+                pre_args.push(arg.generate("", None));
+            } else {
+                post_args.push(arg.generate("", None));
+            }
+        }
+
+        let cmd = Self::new_command(&pre_args)?;
+
+        Ok(Self {
+            pre_args,
+            path_arg: path_arg.unwrap(),
+            post_args,
+            cmd,
+            count: 0,
+            limit,
+        })
+    }
+
+    fn new_command(pre_args: &[OsString]) -> io::Result<Command> {
+        let mut cmd = Command::new(&pre_args[0]);
+        cmd.stdin(Stdio::inherit());
+        cmd.stdout(Stdio::inherit());
+        cmd.stderr(Stdio::inherit());
+        cmd.try_args(&pre_args[1..])?;
+        Ok(cmd)
+    }
+
+    fn push(&mut self, path: &Path, separator: Option<&str>) -> io::Result<()> {
+        if self.limit > 0 && self.count >= self.limit {
+            self.finish()?;
         }
-        ExitCode::Success
+
+        let arg = self.path_arg.generate(path, separator);
+        if !self
+            .cmd
+            .args_would_fit(iter::once(&arg).chain(&self.post_args))
+        {
+            self.finish()?;
+        }
+
+        self.cmd.try_arg(arg)?;
+        self.count += 1;
+        Ok(())
+    }
+
+    fn finish(&mut self) -> io::Result<()> {
+        if self.count > 0 {
+            self.cmd.try_args(&self.post_args)?;
+            self.cmd.status()?;
+
+            self.cmd = Self::new_command(&self.pre_args)?;
+            self.count = 0;
+        }
+
+        Ok(())
     }
 }
 
@@ -192,45 +292,12 @@ impl CommandTemplate {
     ///
     /// Using the internal `args` field, and a supplied `input` variable, a `Command` will be
     /// build.
-    fn generate(&self, input: &Path, path_separator: Option<&str>) -> Command {
+    fn generate(&self, input: &Path, path_separator: Option<&str>) -> io::Result<Command> {
         let mut cmd = Command::new(self.args[0].generate(&input, path_separator));
         for arg in &self.args[1..] {
-            cmd.arg(arg.generate(&input, path_separator));
-        }
-        cmd
-    }
-
-    fn generate_and_execute_batch(
-        &self,
-        paths: &[PathBuf],
-        path_separator: Option<&str>,
-    ) -> ExitCode {
-        let mut cmd = Command::new(self.args[0].generate("", None));
-        cmd.stdin(Stdio::inherit());
-        cmd.stdout(Stdio::inherit());
-        cmd.stderr(Stdio::inherit());
-
-        let mut has_path = false;
-
-        for arg in &self.args[1..] {
-            if arg.has_tokens() {
-                for path in paths {
-                    cmd.arg(arg.generate(path, path_separator));
-                    has_path = true;
-                }
-            } else {
-                cmd.arg(arg.generate("", None));
-            }
-        }
-
-        if has_path {
-            match cmd.spawn().and_then(|mut c| c.wait()) {
-                Ok(_) => ExitCode::Success,
-                Err(e) => handle_cmd_error(&cmd, e),
-            }
-        } else {
-            ExitCode::Success
+            cmd.try_arg(arg.generate(&input, path_separator))?;
         }
+        Ok(cmd)
     }
 }
 
diff --git a/tests/tests.rs b/tests/tests.rs
index 2b83e64..c012f47 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -1492,10 +1492,49 @@ fn test_exec_batch_multi() {
     }
     let te = TestEnv::new(DEFAULT_DIRS, DEFAULT_FILES);
 
-    te.assert_output(
-        &["foo",  "--exec-batch", "echo", "{}", ";", "--exec-batch", "echo", "{/}"],
-        "./a.foo ./one/b.foo ./one/two/C.Foo2 ./one/two/c.foo ./one/two/three/d.foo ./one/two/three/directory_foo
-        a.foo b.foo C.Foo2 c.foo d.foo directory_foo",
+    let output = te.assert_success_and_get_output(
+        ".",
+        &[
+            "foo",
+            "--exec-batch",
+            "echo",
+            "{}",
+            ";",
+            "--exec-batch",
+            "echo",
+            "{/}",
+        ],
+    );
+    let stdout = std::str::from_utf8(&output.stdout).unwrap();
+    let lines: Vec<_> = stdout
+        .lines()
+        .map(|l| {
+            let mut words: Vec<_> = l.split_whitespace().collect();
+            words.sort();
+            words
+        })
+        .collect();
+
+    assert_eq!(
+        lines,
+        &[
+            [
+                "./a.foo",
+                "./one/b.foo",
+                "./one/two/C.Foo2",
+                "./one/two/c.foo",
+                "./one/two/three/d.foo",
+                "./one/two/three/directory_foo"
+            ],
+            [
+                "C.Foo2",
+                "a.foo",
+                "b.foo",
+                "c.foo",
+                "d.foo",
+                "directory_foo"
+            ],
+        ]
     );
 }
 
@@ -1508,11 +1547,6 @@ fn test_exec_batch_with_limit() {
 
     let te = TestEnv::new(DEFAULT_DIRS, DEFAULT_FILES);
 
-    te.assert_output(
-        &["foo", "--batch-size", "0", "--exec-batch", "echo", "{}"],
-        "./a.foo ./one/b.foo ./one/two/C.Foo2 ./one/two/c.foo ./one/two/three/d.foo ./one/two/three/directory_foo",
-    );
-
     let output = te.assert_success_and_get_output(
         ".",
         &["foo", "--batch-size=2", "--exec-batch", "echo", "{}"],