Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

exec: Execute batches before they get too long #1020

Merged
merged 1 commit into from
May 28, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 38 additions & 4 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
31 changes: 20 additions & 11 deletions src/exec/command.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -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()
Expand All @@ -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
}
}
}
13 changes: 1 addition & 12 deletions src/exec/job.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
157 changes: 112 additions & 45 deletions src/exec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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()?;
}
Comment on lines +184 to +189
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice ❤️


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(())
}
}

Expand Down Expand Up @@ -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)
}
}

Expand Down
Loading