exec: Execute batches before they get too long

Fixes sharkdp#410.

Author: tavianator

Cargo.lock

@@ -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"

Cargo.toml

@@ -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,20 +85,23 @@ 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);
             }
         }
     }
     output_buffer.write();
     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
+        }
     }
 }

src/exec/command.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)
 }

src/exec/job.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,119 @@ 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)
             }
         }
-        ExitCode::Success
+    }
+}
+
+/// 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()?;
+        }
+
+        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 +294,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)
     }
 }