pub mod line;
pub mod stream;
pub mod text;
use std::collections::VecDeque;
use std::future::Future;
use std::future::poll_fn;
use std::io::Error as IoError;
#[cfg(unix)]
use std::os::fd::AsFd;
use std::pin::Pin;
use std::pin::pin;
use std::process::ExitStatus;
use std::task::Context;
use std::task::Poll;
use std::time::Duration;
use std::time::Instant;
#[cfg(unix)]
use rustix::io::Errno;
#[cfg(unix)]
use rustix::io::read;
#[cfg(unix)]
use rustix::process::Pid;
#[cfg(unix)]
use rustix::process::Signal;
#[cfg(unix)]
use rustix::process::kill_process;
use tokio::io::AsyncRead;
use tokio::io::ReadBuf;
use tokio::process::Child;
use tokio::process::ChildStderr;
use tokio::process::ChildStdout;
use tokio::time::timeout;
pub use self::line::Line;
pub use self::stream::Stream;
pub use self::text::Text;
use crate::error::RunCommandError;
use crate::execution::Capture;
use crate::execution::Execution;
use crate::invocation::Invocation;
use crate::process_id::ProcessId;
const CHUNK: usize = 8 * 1024;
#[cfg(unix)]
const LEFTOVER: usize = 4 * 1024 * 1024;
#[derive(Debug)]
pub struct Run {
invocation: Invocation,
child: Child,
id: Option<ProcessId>,
stdout: Reader<ChildStdout>,
stderr: Reader<ChildStderr>,
lines: VecDeque<Line>,
started: Instant,
}
impl Run {
pub(crate) fn new(invocation: Invocation, mut child: Child, started: Instant) -> Self {
let stdout = Reader::new(Stream::StandardOutput, child.stdout.take());
let stderr = Reader::new(Stream::StandardError, child.stderr.take());
let id = child.id().map(ProcessId::from);
Self {
invocation,
child,
id,
stdout,
stderr,
lines: VecDeque::new(),
started,
}
}
async fn drain(&mut self) -> Result<(), RunCommandError> {
while self.next_line().await?.is_some() {}
Ok(())
}
async fn end(&mut self) -> Result<ExitStatus, RunCommandError> {
self.drain().await?;
self.child
.wait()
.await
.map_err(|source| RunCommandError::IncompleteRun {
invocation: Box::new(self.invocation.clone()),
source,
})
}
async fn exit(&mut self) -> Result<ExitStatus, RunCommandError> {
let Self {
invocation,
child,
stdout,
stderr,
lines,
..
} = self;
let mut wait = pin!(child.wait());
let exited = poll_fn(|context| {
let produced = match poll_streams(stdout, stderr, lines, context) {
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
Poll::Ready(Ok(())) => true,
Poll::Pending => false,
};
if let Poll::Ready(exited) = wait.as_mut().poll(context) {
return Poll::Ready(exited);
}
if produced {
context.waker().wake_by_ref();
}
Poll::Pending
})
.await;
exited.map_err(|source| RunCommandError::IncompleteRun {
invocation: Box::new(invocation.clone()),
source,
})
}
fn finish(self, status: ExitStatus) -> Execution {
Execution::new(
self.invocation,
self.id,
status,
Capture::new(self.stdout.capture),
Capture::new(self.stderr.capture),
self.started.elapsed(),
)
}
pub fn id(&self) -> Option<ProcessId> {
self.id
}
pub fn invocation(&self) -> &Invocation {
&self.invocation
}
fn kill(&mut self) {
let _ = self.child.start_kill();
}
pub async fn next_line(&mut self) -> Result<Option<Line>, RunCommandError> {
loop {
if let Some(line) = self.lines.pop_front() {
return Ok(Some(line));
}
if self.stdout.is_closed() && self.stderr.is_closed() {
return Ok(None);
}
let read = {
let Self {
stdout,
stderr,
lines,
..
} = self;
poll_fn(|context| poll_streams(stdout, stderr, lines, context)).await
};
read.map_err(|source| RunCommandError::IncompleteRun {
invocation: Box::new(self.invocation.clone()),
source,
})?;
}
}
#[cfg(unix)]
fn request_end(&self) -> Request {
let Some(identifier) = self.child.id() else {
return Request::Unsent;
};
let Ok(identifier) = i32::try_from(identifier) else {
return Request::Unsent;
};
let Some(process) = Pid::from_raw(identifier) else {
return Request::Unsent;
};
match kill_process(process, Signal::TERM) {
Ok(()) => Request::Sent,
Err(_) => Request::Unsent,
}
}
#[cfg(not(unix))]
fn request_end(&self) -> Request {
Request::Unsent
}
pub async fn stop(mut self, grace: Duration) -> Result<Execution, RunCommandError> {
let exited = match self.request_end() {
Request::Sent => match timeout(grace, self.exit()).await {
Ok(exited) => Some(exited?),
Err(_) => None,
},
Request::Unsent => None,
};
let status = match exited {
Some(status) => status,
None => {
self.kill();
self.exit().await?
}
};
self.take().await?;
Ok(self.finish(status))
}
#[cfg(unix)]
async fn take(&mut self) -> Result<(), RunCommandError> {
let Self {
invocation,
stdout,
stderr,
lines,
..
} = self;
let taken = stdout.take(lines).and_then(|()| stderr.take(lines));
taken.map_err(|source| RunCommandError::IncompleteRun {
invocation: Box::new(invocation.clone()),
source,
})
}
#[cfg(not(unix))]
async fn take(&mut self) -> Result<(), RunCommandError> {
self.drain().await
}
pub async fn wait(mut self) -> Result<Execution, RunCommandError> {
let status = self.end().await?;
Ok(self.finish(status))
}
pub async fn wait_for_end(&mut self) -> Result<(), RunCommandError> {
self.end().await?;
Ok(())
}
}
#[derive(Debug)]
enum Request {
Sent,
Unsent,
}
fn poll_streams(
stdout: &mut Reader<ChildStdout>,
stderr: &mut Reader<ChildStderr>,
lines: &mut VecDeque<Line>,
context: &mut Context<'_>,
) -> Poll<Result<(), IoError>> {
let mut produced = false;
match stdout.poll(context, lines) {
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
Poll::Ready(Ok(())) => produced = true,
Poll::Pending => {}
}
match stderr.poll(context, lines) {
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
Poll::Ready(Ok(())) => produced = true,
Poll::Pending => {}
}
if produced {
Poll::Ready(Ok(()))
} else {
Poll::Pending
}
}
#[derive(Debug)]
struct Reader<S> {
stream: Stream,
source: Option<S>,
pending: Vec<u8>,
searched: usize,
capture: Vec<u8>,
}
impl<S: AsyncRead + Unpin> Reader<S> {
fn new(stream: Stream, source: Option<S>) -> Self {
Self {
stream,
source,
pending: Vec::new(),
searched: 0,
capture: Vec::new(),
}
}
fn accept(&mut self, read: &[u8], lines: &mut VecDeque<Line>) {
self.capture.extend_from_slice(read);
self.pending.extend_from_slice(read);
self.split(lines);
}
fn close(&mut self, lines: &mut VecDeque<Line>) {
self.source = None;
if !self.pending.is_empty() {
let bytes = std::mem::take(&mut self.pending);
self.searched = 0;
lines.push_back(self.line(&bytes));
}
}
fn is_closed(&self) -> bool {
self.source.is_none()
}
fn line(&self, bytes: &[u8]) -> Line {
Line::new(self.stream, String::from_utf8_lossy(bytes).into_owned())
}
fn poll(
&mut self,
context: &mut Context<'_>,
lines: &mut VecDeque<Line>,
) -> Poll<Result<(), IoError>> {
let Some(source) = self.source.as_mut() else {
return Poll::Pending;
};
let mut bytes = [0; CHUNK];
let mut buffer = ReadBuf::new(&mut bytes);
match Pin::new(source).poll_read(context, &mut buffer) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
Poll::Ready(Ok(())) => {}
}
let read = buffer.filled();
if read.is_empty() {
self.close(lines);
return Poll::Ready(Ok(()));
}
self.accept(read, lines);
Poll::Ready(Ok(()))
}
#[cfg(unix)]
fn take(&mut self, lines: &mut VecDeque<Line>) -> Result<(), IoError>
where
S: AsFd,
{
let mut bytes = [0; CHUNK];
let mut taken = 0;
while taken < LEFTOVER {
let Some(source) = self.source.as_ref() else {
return Ok(());
};
match read(source, &mut bytes) {
Ok(0) => self.close(lines),
Ok(count) => {
taken += count;
self.accept(&bytes[..count], lines);
}
Err(Errno::INTR) => {}
Err(error) if error == Errno::AGAIN || error == Errno::WOULDBLOCK => {
return Ok(());
}
Err(error) => return Err(error.into()),
}
}
Ok(())
}
fn split(&mut self, lines: &mut VecDeque<Line>) {
while let Some(offset) = self.pending[self.searched..]
.iter()
.position(|byte| *byte == b'\n')
{
let end = self.searched + offset;
let mut bytes: Vec<u8> = self.pending.drain(..=end).collect();
bytes.pop();
if bytes.last() == Some(&b'\r') {
bytes.pop();
}
self.searched = 0;
lines.push_back(self.line(&bytes));
}
self.searched = self.pending.len();
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::missing_panics_doc)]
use super::*;
#[test]
fn run_is_send_and_sync() {
fn assert_send_and_sync<T: Send + Sync>() {}
assert_send_and_sync::<Run>();
}
}