2f11d69
use crate::command::env::Env;
use crate::command::error::{CmdError, CmdResult};
use std::fmt::{Debug, Display, Formatter};
use std::path::PathBuf;
use std::process::{ExitStatus, Stdio};
use time::UtcDateTime;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::process::{Child, Command as TokioCommand};
pub struct Cmd {
pub caller_info: String,
pub prog: String,
pub args: Vec<String>,
pub(crate) config_args: Vec<String>,
pub dir: PathBuf,
pub cmd: Option<Child>,
pub start: UtcDateTime,
pub env: Env,
pub stdin: Box<dyn AsyncWrite + Unpin + Send + Sync>,
pub stdout: Box<dyn AsyncRead + Unpin + Send + Sync>,
pub stderr: Box<dyn AsyncRead + Unpin + Send + Sync>,
pub timeout: Option<u64>,
pub output_cap: usize,
/// Set once `feed()` has written and closed stdin, so a second `feed()`
/// cannot silently discard data.
pub stdin_fed: bool,
/// Set once the child process has been taken out of `self.cmd` (reaped by
/// `wait()` or handed to the reaper by `cancel_pipe()`). A finished `Cmd`
/// must not be spawned again.
pub finished: bool,
}
pub const DEFAULT_OUTPUT_CAP_BYTES: usize = 64 * 1024 * 1024;
#[derive(Debug, Clone)]
pub struct CmdOutput {
pub status: ExitStatus,
pub stdout: Vec<u8>,
pub stderr: Vec<u8>,
}
impl Clone for Cmd {
fn clone(&self) -> Self {
Self {
caller_info: self.caller_info.clone(),
prog: self.prog.clone(),
args: self.args.clone(),
config_args: self.config_args.clone(),
dir: self.dir.clone(),
cmd: None,
start: self.start,
env: self.env.clone(),
stdin: Box::new(tokio::io::sink()),
stdout: Box::new(tokio::io::empty()),
stderr: Box::new(tokio::io::empty()),
timeout: self.timeout,
output_cap: self.output_cap,
stdin_fed: false,
finished: false,
}
}
}
impl Debug for Cmd {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Cmd")
.field("caller_info", &self.caller_info)
.field("prog", &self.prog)
.field("args", &self.args)
.field("config_args", &self.config_args)
.field("dir", &self.dir)
.field("start", &self.start)
.field("env", &self.env)
.field("timeout", &self.timeout)
.field("cmd", &self.cmd.as_ref().map(|_| "<running>"))
.finish_non_exhaustive()
}
}
impl Display for Cmd {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let mut parts: Vec<&str> = Vec::with_capacity(1 + self.config_args.len() + self.args.len());
parts.push(&self.prog);
for a in &self.config_args {
parts.push(a);
}
for a in &self.args {
parts.push(a);
}
write!(f, "{}", parts.join(" "))
}
}
impl Cmd {
pub fn new(
caller_info: impl Into<String>,
prog: impl Into<String>,
args: Vec<String>,
config_args: Vec<String>,
dir: PathBuf,
env: impl Into<Env>,
timeout: Option<u64>,
) -> Self {
Self {
caller_info: caller_info.into(),
prog: prog.into(),
args,
config_args,
dir,
cmd: None,
start: UtcDateTime::now(),
env: env.into(),
stdin: Box::new(tokio::io::sink()),
stdout: Box::new(tokio::io::empty()),
stderr: Box::new(tokio::io::empty()),
timeout,
output_cap: DEFAULT_OUTPUT_CAP_BYTES,
stdin_fed: false,
finished: false,
}
}
pub fn build_command(&self) -> TokioCommand {
let mut command = TokioCommand::new(&self.prog);
for cfg in &self.config_args {
command.arg(cfg);
}
for arg in &self.args {
command.arg(arg);
}
command.current_dir(&self.dir);
command.stdin(Stdio::piped());
command.stdout(Stdio::piped());
command.stderr(Stdio::piped());
for (k, v) in self.env.iter() {
command.env(k, v);
}
command
}
pub async fn spawn(&mut self) -> CmdResult<()> {
if self.finished {
return Err(CmdError::Custom(format!(
"process `{}` already finished and cannot be spawned again",
self.prog
)));
}
if self.cmd.is_some() {
return Err(CmdError::Custom(format!(
"process `{}` already spawned",
self.prog
)));
}
self.start = UtcDateTime::now();
let mut command = self.build_command();
let mut child = command.spawn().map_err(|source| CmdError::Spawn {
prog: self.prog.clone(),
source,
})?;
let stdin = child
.stdin
.take()
.ok_or_else(|| CmdError::Custom("stdin was not piped".into()))?;
let stdout = child
.stdout
.take()
.ok_or_else(|| CmdError::Custom("stdout was not piped".into()))?;
let stderr = child
.stderr
.take()
.ok_or_else(|| CmdError::Custom("stderr was not piped".into()))?;
self.stdin = Box::new(stdin);
self.stdout = Box::new(stdout);
self.stderr = Box::new(stderr);
self.cmd = Some(child);
Ok(())
}
pub async fn feed(&mut self, input: &[u8]) -> CmdResult<()> {
if self.stdin_fed {
return Err(CmdError::Custom("stdin was already fed".into()));
}
if self.cmd.is_none() {
self.spawn().await?;
}
self.stdin_fed = true;
let mut writer = std::mem::replace(&mut self.stdin, Box::new(tokio::io::sink()));
let write_fut = async {
writer.write_all(input).await?;
writer.shutdown().await
};
if let Some(secs) = self.timeout {
match tokio::time::timeout(std::time::Duration::from_secs(secs), write_fut).await {
Ok(res) => {
res.map_err(|source| CmdError::Io {
path: self.dir.clone(),
source,
})?;
}
Err(_) => {
if let Some(child) = self.cmd.as_mut() {
let _ = child.start_kill();
}
return Err(CmdError::Timeout {
prog: self.prog.clone(),
timeout_ms: secs.saturating_mul(1000),
});
}
}
} else {
write_fut.await.map_err(|source| CmdError::Io {
path: self.dir.clone(),
source,
})?;
}
Ok(())
}
pub async fn wait(&mut self) -> CmdResult<ExitStatus> {
let mut child = self
.cmd
.take()
.ok_or_else(|| CmdError::Custom("process not spawned yet".into()))?;
self.finished = true;
let wait_fut = child.wait();
let status = if let Some(secs) = self.timeout {
match tokio::time::timeout(std::time::Duration::from_secs(secs), wait_fut).await {
Ok(res) => res,
Err(_) => {
// Kill, then reap so the child does not linger as a zombie.
if let Err(source) = child.start_kill() {
return Err(CmdError::Io {
path: self.dir.clone(),
source,
});
}
let _ = child.wait().await;
return Err(CmdError::Timeout {
prog: self.prog.clone(),
timeout_ms: secs.saturating_mul(1000),
});
}
}
} else {
wait_fut.await
};
status.map_err(|source| CmdError::Io {
path: self.dir.clone(),
source,
})
}
pub async fn run(&mut self) -> CmdResult<CmdOutput> {
self.run_inner(false).await
}
/// Run the command but treat a non-zero exit as a *successful* run whose
/// `CmdOutput.status` carries the failure, instead of returning
/// `CmdError::NonZeroExit`.
///
/// Use this when the caller wants to inspect exit codes / stderr itself
/// (e.g. classifying "does not exist" vs real errors).
pub async fn run_soft(&mut self) -> CmdResult<CmdOutput> {
self.run_inner(true).await
}
async fn run_inner(&mut self, soft: bool) -> CmdResult<CmdOutput> {
if self.cmd.is_none() {
self.spawn().await?;
}
let stdout = std::mem::replace(&mut self.stdout, Box::new(tokio::io::empty()));
let stderr = std::mem::replace(&mut self.stderr, Box::new(tokio::io::empty()));
let cap = self.output_cap;
let read_out = read_drain_capped(stdout, cap);
let read_err = read_drain_capped(stderr, cap);
let wait_fut = self.wait();
let (out_res, err_res, wait_res) = tokio::join!(read_out, read_err, wait_fut);
let (stdout, stdout_total) = out_res.map_err(|source| CmdError::Io {
path: self.dir.clone(),
source,
})?;
let (stderr, stderr_total) = err_res.map_err(|source| CmdError::Io {
path: self.dir.clone(),
source,
})?;
if stdout_total > cap as u64 {
return Err(CmdError::PayloadTooLarge {
cap_bytes: cap as u64,
actual_bytes: stdout_total,
});
}
if stderr_total > cap as u64 {
return Err(CmdError::PayloadTooLarge {
cap_bytes: cap as u64,
actual_bytes: stderr_total,
});
}
let status = wait_res?;
if !status.success() && !soft {
return Err(CmdError::NonZeroExit {
prog: self.prog.clone(),
args: self.args.clone(),
code: status.code().unwrap_or(-1),
stderr: String::from_utf8_lossy(&stderr).into_owned(),
});
}
Ok(CmdOutput {
status,
stdout,
stderr,
})
}
pub async fn kill(&mut self) -> CmdResult<()> {
if let Some(child) = self.cmd.as_mut() {
child.start_kill().map_err(|source| CmdError::Io {
path: self.dir.clone(),
source,
})?;
}
Ok(())
}
pub fn try_wait(&mut self) -> CmdResult<Option<ExitStatus>> {
let child = self
.cmd
.as_mut()
.ok_or_else(|| CmdError::Custom("process not spawned yet".into()))?;
child.try_wait().map_err(|source| CmdError::Io {
path: self.dir.clone(),
source,
})
}
pub fn is_running(&self) -> bool {
self.cmd.is_some() && !self.finished
}
#[must_use]
pub fn with_dir(&self, dir: PathBuf) -> Self {
let mut cloned = self.clone();
cloned.dir = dir;
cloned
}
}
/// Read `reader` to completion while capping the retained buffer at `cap`.
///
/// To avoid deadlocking the child (which would block writing once the pipe
/// buffer fills), bytes beyond `cap` are still drained but discarded. Returns
/// the retained prefix and the *total* number of bytes drained, so the caller
/// can distinguish "exactly at cap" from "overflowed".
async fn read_drain_capped<R>(mut reader: R, cap: usize) -> std::io::Result<(Vec<u8>, u64)>
where
R: AsyncRead + Unpin,
{
let mut buf: Vec<u8> = Vec::with_capacity(cap.min(64 * 1024));
let mut total: u64 = 0;
let mut chunk = [0u8; 8192];
loop {
let n = reader.read(&mut chunk).await?;
if n == 0 {
break;
}
total = total.saturating_add(n as u64);
if buf.len() < cap {
let space = cap - buf.len();
buf.extend_from_slice(&chunk[..n.min(space)]);
}
}
Ok((buf, total))
}