From edf8d4102d9473d1b9dcd5bb959530ef3a27704d Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Thu, 12 May 2022 10:41:47 -0400 Subject: [PATCH] exec: Execute batches before they get too long Fixes #410. --- Cargo.lock | 42 ++++++++++-- Cargo.toml | 1 + src/exec/command.rs | 31 +++++---- src/exec/job.rs | 13 +--- src/exec/mod.rs | 157 +++++++++++++++++++++++++++++++------------- tests/tests.rs | 52 ++++++++++++--- 6 files changed, 215 insertions(+), 81 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7553fb8b9..339611f45 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -26,6 +26,17 @@ version = "1.0.57" 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" @@ -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" @@ -357,6 +369,28 @@ version = "2.4.1" 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" diff --git a/Cargo.toml b/Cargo.toml index 02014bd91..7f94821ee 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 935b6f49a..a642a405b 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>( +pub fn execute_commands>>( 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>( } 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>( 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 304792b07..9b95ac24b 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 d1405280b..fc26da2a6 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(&self, paths: I) -> ExitCode + pub fn execute_batch(&self, paths: I, limit: usize) -> ExitCode where I: Iterator, { let path_separator = self.path_separator.as_deref(); - let mut paths = paths.collect::>(); - 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> = 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, + path_arg: ArgumentTemplate, + post_args: Vec, + cmd: Command, + count: usize, + limit: usize, +} + +impl CommandBuilder { + fn new(template: &CommandTemplate, limit: usize) -> io::Result { + 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 { + 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 { 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 2b83e642a..c012f47e1 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", "{}"],