use crate::errors::*;
use crate::jobserver::{Acquired, Client};
use async_trait::async_trait;
use std::boxed::Box;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::io;
use std::path::Path;
use std::process::{Command, ExitStatus, Output, Stdio};
use std::sync::{Arc, Mutex};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::process::{ChildStderr, ChildStdin, ChildStdout};
#[async_trait]
pub trait CommandChild {
type I: AsyncWrite + Unpin + Sync + Send + 'static;
type O: AsyncRead + Unpin + Sync + Send + 'static;
type E: AsyncRead + Unpin + Sync + Send + 'static;
fn take_stdin(&mut self) -> Option<Self::I>;
fn take_stdout(&mut self) -> Option<Self::O>;
fn take_stderr(&mut self) -> Option<Self::E>;
async fn wait(self) -> io::Result<ExitStatus>;
async fn wait_with_output(self) -> io::Result<Output>;
}
#[async_trait]
pub trait RunCommand: fmt::Debug + Send {
type C: CommandChild + Send + 'static;
fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Self;
fn args<S: AsRef<OsStr>>(&mut self, args: &[S]) -> &mut Self;
fn env<K, V>(&mut self, key: K, val: V) -> &mut Self
where
K: AsRef<OsStr>,
V: AsRef<OsStr>;
fn envs<I, K, V>(&mut self, vars: I) -> &mut Self
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<OsStr>,
V: AsRef<OsStr>;
fn env_clear(&mut self) -> &mut Self;
fn current_dir<P: AsRef<Path>>(&mut self, dir: P) -> &mut Self;
fn stdin(&mut self, cfg: Stdio) -> &mut Self;
fn stdout(&mut self, cfg: Stdio) -> &mut Self;
fn stderr(&mut self, cfg: Stdio) -> &mut Self;
async fn spawn(&mut self) -> Result<Self::C>;
}
pub trait CommandCreator {
type Cmd: RunCommand;
fn new(client: &Client) -> Self;
fn new_command<S: AsRef<OsStr>>(&mut self, program: S) -> Self::Cmd;
}
pub trait CommandCreatorSync: Clone + Send + Sync + 'static {
type Cmd: RunCommand;
fn new(client: &Client) -> Self;
fn new_command_sync<S: AsRef<OsStr>>(&mut self, program: S) -> Self::Cmd;
}
pub struct Child {
inner: tokio::process::Child,
token: Acquired,
}
#[async_trait]
impl CommandChild for Child {
type I = ChildStdin;
type O = ChildStdout;
type E = ChildStderr;
fn take_stdin(&mut self) -> Option<ChildStdin> {
self.inner.stdin.take()
}
fn take_stdout(&mut self) -> Option<ChildStdout> {
self.inner.stdout.take()
}
fn take_stderr(&mut self) -> Option<ChildStderr> {
self.inner.stderr.take()
}
async fn wait(self) -> io::Result<ExitStatus> {
let Child { mut inner, token } = self;
inner.wait().await.inspect(|_ret| {
drop(token);
})
}
async fn wait_with_output(self) -> io::Result<Output> {
let Child { inner, token } = self;
inner.wait_with_output().await.inspect(|_ret| {
drop(token);
})
}
}
pub struct AsyncCommand {
inner: Option<Command>,
jobserver: Client,
}
impl AsyncCommand {
pub fn new<S: AsRef<OsStr>>(program: S, jobserver: Client) -> AsyncCommand {
AsyncCommand {
inner: Some(Command::new(program)),
jobserver,
}
}
fn inner(&mut self) -> &mut Command {
self.inner.as_mut().expect("can't reuse commands")
}
}
#[async_trait]
impl RunCommand for AsyncCommand {
type C = Child;
fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut AsyncCommand {
self.inner().arg(arg);
self
}
fn args<S: AsRef<OsStr>>(&mut self, args: &[S]) -> &mut AsyncCommand {
self.inner().args(args);
self
}
fn env<K, V>(&mut self, key: K, val: V) -> &mut AsyncCommand
where
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
self.inner().env(key, val);
self
}
fn envs<I, K, V>(&mut self, vars: I) -> &mut Self
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
self.inner().envs(vars);
self
}
fn env_clear(&mut self) -> &mut AsyncCommand {
self.inner().env_clear();
self
}
fn current_dir<P: AsRef<Path>>(&mut self, dir: P) -> &mut AsyncCommand {
self.inner().current_dir(dir);
self
}
fn stdin(&mut self, cfg: Stdio) -> &mut AsyncCommand {
self.inner().stdin(cfg);
self
}
fn stdout(&mut self, cfg: Stdio) -> &mut AsyncCommand {
self.inner().stdout(cfg);
self
}
fn stderr(&mut self, cfg: Stdio) -> &mut AsyncCommand {
self.inner().stderr(cfg);
self
}
async fn spawn(&mut self) -> Result<Child> {
let mut inner = self.inner.take().unwrap();
inner.env_remove("MAKEFLAGS");
inner.env_remove("MFLAGS");
inner.env_remove("CARGO_MAKEFLAGS");
self.jobserver.configure(&mut inner);
let token = self.jobserver.acquire().await?;
let mut inner = tokio::process::Command::from(inner);
let child = inner
.spawn()
.with_context(|| format!("failed to spawn {:?}", inner))?;
Ok(Child {
inner: child,
token,
})
}
}
impl fmt::Debug for AsyncCommand {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.inner.fmt(f)
}
}
#[derive(Clone)]
pub struct ProcessCommandCreator {
jobserver: Client,
}
impl CommandCreator for ProcessCommandCreator {
type Cmd = AsyncCommand;
fn new(client: &Client) -> ProcessCommandCreator {
ProcessCommandCreator {
jobserver: client.clone(),
}
}
fn new_command<S: AsRef<OsStr>>(&mut self, program: S) -> AsyncCommand {
AsyncCommand::new(program, self.jobserver.clone())
}
}
impl CommandCreatorSync for ProcessCommandCreator {
type Cmd = AsyncCommand;
fn new(client: &Client) -> ProcessCommandCreator {
CommandCreator::new(client)
}
fn new_command_sync<S: AsRef<OsStr>>(&mut self, program: S) -> AsyncCommand {
self.new_command(program)
}
}
#[cfg(unix)]
use std::os::unix::process::ExitStatusExt;
#[cfg(windows)]
use std::os::windows::process::ExitStatusExt;
#[cfg(unix)]
pub type ExitStatusValue = i32;
#[cfg(windows)]
pub type ExitStatusValue = u32;
#[allow(dead_code)]
pub fn exit_status(v: ExitStatusValue) -> ExitStatus {
ExitStatus::from_raw(v)
}
#[allow(dead_code)]
#[derive(Debug)]
pub struct MockChild {
pub stdin: Option<io::Cursor<Vec<u8>>>,
pub stdout: Option<io::Cursor<Vec<u8>>>,
pub stderr: Option<io::Cursor<Vec<u8>>>,
pub wait_result: Option<io::Result<ExitStatus>>,
}
impl MockChild {
#[allow(dead_code)]
pub fn new<T: AsRef<[u8]>, U: AsRef<[u8]>>(
status: ExitStatus,
stdout: T,
stderr: U,
) -> MockChild {
MockChild {
stdin: Some(io::Cursor::new(vec![])),
stdout: Some(io::Cursor::new(stdout.as_ref().to_vec())),
stderr: Some(io::Cursor::new(stderr.as_ref().to_vec())),
wait_result: Some(Ok(status)),
}
}
#[allow(dead_code)]
pub fn with_error(err: io::Error) -> MockChild {
MockChild {
stdin: None,
stdout: None,
stderr: None,
wait_result: Some(Err(err)),
}
}
}
#[async_trait]
impl CommandChild for MockChild {
type I = io::Cursor<Vec<u8>>;
type O = io::Cursor<Vec<u8>>;
type E = io::Cursor<Vec<u8>>;
fn take_stdin(&mut self) -> Option<io::Cursor<Vec<u8>>> {
self.stdin.take()
}
fn take_stdout(&mut self) -> Option<io::Cursor<Vec<u8>>> {
self.stdout.take()
}
fn take_stderr(&mut self) -> Option<io::Cursor<Vec<u8>>> {
self.stderr.take()
}
async fn wait(mut self) -> io::Result<ExitStatus> {
self.wait_result.take().unwrap()
}
async fn wait_with_output(self) -> io::Result<Output> {
let MockChild {
stdout,
stderr,
wait_result,
..
} = self;
wait_result.unwrap().map(|status| Output {
status,
stdout: stdout.map(|c| c.into_inner()).unwrap_or_else(Vec::new),
stderr: stderr.map(|c| c.into_inner()).unwrap_or_else(Vec::new),
})
}
}
pub enum ChildOrCall {
Child(Result<MockChild>),
Call(Box<dyn Fn(&[OsString]) -> Result<MockChild> + Send>),
}
impl fmt::Debug for ChildOrCall {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
ChildOrCall::Child(ref r) => write!(f, "ChildOrCall::Child({:?}", r),
ChildOrCall::Call(_) => write!(f, "ChildOrCall::Call(...)"),
}
}
}
#[allow(dead_code)]
#[derive(Debug)]
pub struct MockCommand {
pub child: Option<ChildOrCall>,
pub args: Vec<OsString>,
}
#[async_trait]
impl RunCommand for MockCommand {
type C = MockChild;
fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut MockCommand {
self.args.push(arg.as_ref().to_owned());
self
}
fn args<S: AsRef<OsStr>>(&mut self, args: &[S]) -> &mut MockCommand {
self.args.extend(args.iter().map(|a| a.as_ref().to_owned()));
self
}
fn env<K, V>(&mut self, _key: K, _val: V) -> &mut MockCommand
where
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
self
}
fn envs<I, K, V>(&mut self, _vars: I) -> &mut Self
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
self
}
fn env_clear(&mut self) -> &mut MockCommand {
self
}
fn current_dir<P: AsRef<Path>>(&mut self, _dir: P) -> &mut MockCommand {
self
}
fn stdin(&mut self, _cfg: Stdio) -> &mut MockCommand {
self
}
fn stdout(&mut self, _cfg: Stdio) -> &mut MockCommand {
self
}
fn stderr(&mut self, _cfg: Stdio) -> &mut MockCommand {
self
}
async fn spawn(&mut self) -> Result<MockChild> {
match self.child.take().unwrap() {
ChildOrCall::Child(c) => c,
ChildOrCall::Call(f) => f(&self.args),
}
}
}
#[allow(dead_code)]
pub struct MockCommandCreator {
pub children: Vec<ChildOrCall>,
}
impl MockCommandCreator {
#[allow(dead_code)]
pub fn next_command_spawns(&mut self, child: Result<MockChild>) {
self.children.push(ChildOrCall::Child(child));
}
#[allow(dead_code)]
pub fn next_command_calls<C>(&mut self, call: C)
where
C: Fn(&[OsString]) -> Result<MockChild> + Send + 'static,
{
self.children.push(ChildOrCall::Call(Box::new(call)));
}
}
impl CommandCreator for MockCommandCreator {
type Cmd = MockCommand;
fn new(_client: &Client) -> MockCommandCreator {
MockCommandCreator {
children: Vec::new(),
}
}
fn new_command<S: AsRef<OsStr>>(&mut self, _program: S) -> MockCommand {
assert!(
!self.children.is_empty(),
"Too many calls to MockCommandCreator::new_command, or not enough to MockCommandCreator::new_command_spawns!"
);
MockCommand {
child: Some(self.children.remove(0)),
args: vec![],
}
}
}
impl<T: CommandCreator + 'static + Send> CommandCreatorSync for Arc<Mutex<T>> {
type Cmd = T::Cmd;
fn new(client: &Client) -> Arc<Mutex<T>> {
Arc::new(Mutex::new(T::new(client)))
}
fn new_command_sync<S: AsRef<OsStr>>(&mut self, program: S) -> T::Cmd {
self.lock().unwrap().new_command(program)
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::jobserver::Client;
use crate::test::utils::*;
use std::ffi::OsStr;
use std::io;
use std::process::{ExitStatus, Output};
use std::sync::{Arc, Mutex};
use std::thread;
fn spawn_command<T: CommandCreator, S: AsRef<OsStr>>(
creator: &mut T,
program: S,
) -> Result<<<T as CommandCreator>::Cmd as RunCommand>::C> {
creator.new_command(program).spawn().wait()
}
fn spawn_wait_command<T: CommandCreator, S: AsRef<OsStr>>(
creator: &mut T,
program: S,
) -> Result<ExitStatus> {
Ok(spawn_command(creator, program)?.wait().wait()?)
}
fn spawn_output_command<T: CommandCreator, S: AsRef<OsStr>>(
creator: &mut T,
program: S,
) -> Result<Output> {
Ok(spawn_command(creator, program)?.wait_with_output().wait()?)
}
fn spawn_on_thread<T: CommandCreatorSync + Send + 'static>(
mut t: T,
really: bool,
) -> ExitStatus {
thread::spawn(move || {
if really {
t.new_command_sync("foo")
.spawn()
.wait()
.unwrap()
.wait()
.wait()
.unwrap()
} else {
exit_status(1)
}
})
.join()
.unwrap()
}
#[test]
fn test_mock_command_wait() {
let client = Client::new_num(1);
let mut creator = MockCommandCreator::new(&client);
creator.next_command_spawns(Ok(MockChild::new(exit_status(0), "hello", "error")));
assert_eq!(
0,
spawn_wait_command(&mut creator, "foo")
.unwrap()
.code()
.unwrap()
);
}
#[test]
#[should_panic]
fn test_unexpected_new_command() {
let client = Client::new_num(1);
let mut creator = MockCommandCreator::new(&client);
creator.new_command("foo").spawn().wait().unwrap();
}
#[test]
fn test_mock_command_output() {
let client = Client::new_num(1);
let mut creator = MockCommandCreator::new(&client);
creator.next_command_spawns(Ok(MockChild::new(exit_status(0), "hello", "error")));
let output = spawn_output_command(&mut creator, "foo").unwrap();
assert_eq!(0, output.status.code().unwrap());
assert_eq!(b"hello".to_vec(), output.stdout);
assert_eq!(b"error".to_vec(), output.stderr);
}
#[test]
fn test_mock_command_calls() {
let client = Client::new_num(1);
let mut creator = MockCommandCreator::new(&client);
creator.next_command_calls(|_| Ok(MockChild::new(exit_status(0), "hello", "error")));
let output = spawn_output_command(&mut creator, "foo").unwrap();
assert_eq!(0, output.status.code().unwrap());
assert_eq!(b"hello".to_vec(), output.stdout);
assert_eq!(b"error".to_vec(), output.stderr);
}
#[test]
fn test_mock_spawn_error() {
let client = Client::new_num(1);
let mut creator = MockCommandCreator::new(&client);
creator.next_command_spawns(Err(anyhow!("error")));
let e = spawn_command(&mut creator, "foo").err().unwrap();
assert_eq!("error", e.to_string());
}
#[test]
fn test_mock_wait_error() {
let client = Client::new_num(1);
let mut creator = MockCommandCreator::new(&client);
creator.next_command_spawns(Ok(MockChild::with_error(io::Error::other("error"))));
let e = spawn_wait_command(&mut creator, "foo").err().unwrap();
assert_eq!("error", e.to_string());
}
#[test]
fn test_mock_command_sync() {
let client = Client::new_num(1);
let creator = Arc::new(Mutex::new(MockCommandCreator::new(&client)));
next_command(
&creator,
Ok(MockChild::new(exit_status(0), "hello", "error")),
);
assert_eq!(exit_status(0), spawn_on_thread(creator, true));
}
}