use crate::{
Command, ExecutionCompletion, ExecutorError, ExecutorResult, Input, InputCompletion, Output,
};
use alloc::borrow::ToOwned;
use std::{ffi::OsStr, io::ErrorKind, process::Stdio};
use tokio::process::Child;
#[derive(Debug)]
pub struct Executor(Command, crate::BatchOptions);
impl Executor {
pub fn new(program: impl AsRef<OsStr>) -> Self {
let libexec_path = asimov_env::paths::asimov_root()
.join("libexec")
.join(program.as_ref());
let mut command = if libexec_path.exists() {
Command::new(libexec_path)
} else {
Command::new(program)
};
command.env("NO_COLOR", "1"); command.stdin(Stdio::null());
command.stdout(Stdio::null());
command.stderr(Stdio::null());
command.kill_on_drop(true);
Self(command, crate::BatchOptions::default())
}
#[must_use]
pub fn with_batching(mut self, options: crate::BatchOptions) -> Self {
self.1 = options;
self
}
pub fn batch_options(&self) -> crate::BatchOptions {
self.1
}
pub fn command(&mut self) -> &mut Command {
&mut self.0
}
pub fn ignore_stdin(&mut self) {
self.0.stdin(Stdio::null());
}
pub fn ignore_stdout(&mut self) {
self.0.stdout(Stdio::null());
}
pub fn ignore_stderr(&mut self) {
self.0.stderr(Stdio::null());
}
pub fn capture_stdout(&mut self) {
self.0.stdout(Stdio::piped());
}
pub fn capture_stderr(&mut self) {
self.0.stderr(Stdio::piped());
}
pub async fn execute(&mut self) -> ExecutorResult {
let process = self.spawn().await?;
self.wait(process).await
}
pub async fn execute_with_input(&mut self, input: &mut Input) -> ExecutorResult {
self.execute_with_io(input, &mut Output::Captured).await
}
pub async fn execute_with_io(
&mut self,
input: &mut Input,
output: &mut Output,
) -> ExecutorResult {
self.execute_with_io_completion(input, output)
.await?
.into_result()
}
pub async fn execute_with_io_completion(
&mut self,
input: &mut Input,
output: &mut Output,
) -> Result<ExecutionCompletion, ExecutorError> {
communicate(self.spawn().await?, input, output).await
}
pub async fn spawn(&mut self) -> Result<Child, ExecutorError> {
match self.0.spawn() {
Ok(process) => Ok(process),
Err(err) if err.kind() == ErrorKind::NotFound => {
let program = self.0.as_std().get_program().to_owned();
return Err(ExecutorError::MissingProgram(program));
},
Err(err) => return Err(ExecutorError::SpawnFailure(err)),
}
}
pub async fn wait(&mut self, process: Child) -> ExecutorResult {
communicate(process, &mut Input::Ignored, &mut Output::Captured)
.await?
.into_result()
}
}
pub(crate) async fn communicate(
mut process: Child,
input: &mut Input,
output: &mut Output,
) -> Result<ExecutionCompletion, ExecutorError> {
communicate_child(&mut process, input, output).await
}
pub(crate) async fn communicate_child(
process: &mut Child,
input: &mut Input,
output: &mut Output,
) -> Result<ExecutionCompletion, ExecutorError> {
use crate::completion::InputFailure;
use alloc::vec::Vec;
use tokio::io::AsyncReadExt;
let stdin = process.stdin.take();
let stdout = process.stdout.take();
let stderr = process.stderr.take();
let supervise = async {
let feed = input.write_to(stdin);
tokio::pin!(feed);
let input = tokio::select! {
biased;
result = &mut feed => match result {
Ok(()) => InputCompletion::Complete,
Err(InputFailure::Source(error)) => {
process.start_kill()?;
InputCompletion::SourceFailed(error)
},
Err(InputFailure::Write(error)) => {
if error.kind() != ErrorKind::BrokenPipe {
process.start_kill()?;
}
InputCompletion::WriteFailed(error)
},
},
status = process.wait() => {
return Ok::<_, std::io::Error>((status?, InputCompletion::Interrupted));
},
};
Ok((process.wait().await?, input))
};
let read_stderr = async {
let mut bytes = Vec::new();
if let Some(mut stderr) = stderr {
stderr.read_to_end(&mut bytes).await?;
}
Ok::<_, std::io::Error>(bytes)
};
let result = tokio::try_join!(supervise, output.read_from(stdout), read_stderr);
match result {
Ok(((status, input), stdout, stderr)) => {
#[cfg(feature = "tracing")]
tracing::trace!("The command exited with: {}", status);
Ok(ExecutionCompletion {
output: std::process::Output {
status,
stdout,
stderr,
},
input,
})
},
Err(error) => {
let _ = process.start_kill();
let _ = process.wait().await;
Err(error.into())
},
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_success() {
let mut runner = Executor::new("curl");
runner.command().arg("https://www.google.com");
let result = runner.execute().await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_missing_program() {
let mut runner = Executor::new("this-command-does-not-exist");
let result = runner.execute().await;
assert!(matches!(result, Err(ExecutorError::MissingProgram(_))));
}
#[cfg(unix)]
#[tokio::test]
async fn test_spawn_failure() {
let mut runner = Executor::new("/dev/null");
let result = runner.execute().await;
assert!(matches!(result, Err(ExecutorError::SpawnFailure(_))));
}
#[tokio::test]
async fn test_unexpected_failure() {
let mut runner = Executor::new("curl");
let result = runner.execute().await;
assert!(matches!(
result,
Err(ExecutorError::UnexpectedFailure(_, _))
));
}
}