mod probes;
mod stream;
pub use stream::{Finished, OutputEvent, OutputEvents, OutputLine, StdoutLines};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::time::{Duration, Instant, SystemTime};
use encoding_rs::Encoding;
use tokio::io::AsyncReadExt;
use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout};
use tokio::sync::Notify;
use tokio::task::{AbortHandle, JoinHandle};
use crate::buffer::OutputBufferPolicy;
use crate::error::Error;
use crate::error::Result;
use crate::group::ProcessGroup;
use crate::pump::{LineHandler, SharedLines, pump_lines_core};
use crate::result::{Outcome, ProcessResult};
use crate::stdin::ProcessStdin;
const PUMP_TEARDOWN: Duration = Duration::from_secs(5);
const DISCARD_INFLIGHT_CAP: usize = 64 << 20;
const TS_PENDING: u8 = 0;
const TS_EXITED: u8 = 1;
pub(crate) const TS_TIMED_OUT: u8 = 2;
enum ExitCause {
Exited(Outcome),
Cancelled,
}
struct FinishedLines {
outcome: Outcome,
stdout_lines: Vec<String>,
stderr_lines: Vec<String>,
}
#[derive(Clone, Copy)]
enum CaptureMode {
Lines,
Discard,
}
pub(crate) struct Spawned {
pub program: String,
pub child: Child,
pub own_group: Option<ProcessGroup>,
pub stdout: Option<ChildStdout>,
pub stderr: Option<ChildStderr>,
pub stdin: Option<ChildStdin>,
pub stdin_task: Option<JoinHandle<std::io::Result<()>>>,
pub timeout: Option<Duration>,
pub timeout_grace: Option<Duration>,
pub timeout_signal: i32,
pub pid: Option<u32>,
pub stdout_encoding: &'static Encoding,
pub stderr_encoding: &'static Encoding,
pub stdout_handler: Option<LineHandler>,
pub stderr_handler: Option<LineHandler>,
pub stdout_tee: Option<crate::pump::TeeSink>,
pub stderr_tee: Option<crate::pump::TeeSink>,
pub buffer: OutputBufferPolicy,
pub ok_codes: Vec<i32>,
pub stdout_piped: bool,
pub cancel_token: Option<tokio_util::sync::CancellationToken>,
}
pub struct RunningProcess {
program: String,
backend: Backend,
timeout: Option<Duration>,
timeout_grace: Option<Duration>,
timeout_signal: i32,
pid: Option<u32>,
stdout_encoding: &'static Encoding,
stderr_encoding: &'static Encoding,
stdout_handler: Option<LineHandler>,
stderr_handler: Option<LineHandler>,
stdout_tee: Option<crate::pump::TeeSink>,
stderr_tee: Option<crate::pump::TeeSink>,
buffer: OutputBufferPolicy,
ok_codes: Vec<i32>,
stdout_sink: Option<Arc<SharedLines>>,
stderr_sink: Option<Arc<SharedLines>>,
stdout_pump: Option<JoinHandle<()>>,
stderr_pump: Option<JoinHandle<()>>,
stdin_error: Option<std::io::Error>,
stdout_piped: bool,
deadline_task: Option<JoinHandle<()>>,
timeout_state: Arc<AtomicU8>,
cancel_token: Option<tokio_util::sync::CancellationToken>,
cancel_task: Option<JoinHandle<()>>,
cancel_at_exit: Option<bool>,
started: Instant,
start_time: SystemTime,
}
type OutputReader = Box<dyn tokio::io::AsyncRead + Send + Unpin>;
enum Backend {
Real(Box<RealProc>),
Scripted(Box<ScriptedProc>),
}
struct RealProc {
child: Child,
own_group: Option<Arc<ProcessGroup>>,
stdout_pipe: Option<ChildStdout>,
stderr_pipe: Option<ChildStderr>,
stdin_pipe: Option<ChildStdin>,
stdin_task: Option<JoinHandle<std::io::Result<()>>>,
}
#[derive(Clone)]
struct ScriptedKill {
killed: Arc<AtomicBool>,
signal: Arc<Notify>,
feeders: Arc<Vec<AbortHandle>>,
}
impl ScriptedKill {
fn fire(&self) {
self.killed.store(true, Ordering::Release);
for feeder in self.feeders.iter() {
feeder.abort();
}
self.signal.notify_one();
}
}
pub(crate) struct ScriptedProc {
stdout: Option<tokio::io::DuplexStream>,
stderr: Option<tokio::io::DuplexStream>,
kill: ScriptedKill,
code: Option<i32>,
timed_out: bool,
signal: Option<i32>,
exit_at: Option<tokio::time::Instant>,
}
impl ScriptedProc {
pub(crate) fn new(
stdout_text: String,
stderr_text: String,
code: Option<i32>,
timed_out: bool,
signal: Option<i32>,
lifetime: Option<Duration>,
line_delay: Option<Duration>,
) -> Self {
let mut feeders = Vec::new();
let mut feed = |text: String| {
let (mut tx, rx) = tokio::io::duplex(64 * 1024);
if text.is_empty() {
return rx; }
let task = tokio::spawn(async move {
use tokio::io::AsyncWriteExt;
match line_delay {
None => {
let _ = tx.write_all(text.as_bytes()).await;
}
Some(delay) => {
for line in text.split_inclusive('\n') {
tokio::time::sleep(delay).await;
if tx.write_all(line.as_bytes()).await.is_err() {
break;
}
}
}
}
});
feeders.push(task.abort_handle());
rx
};
let stdout = feed(stdout_text);
let stderr = feed(stderr_text);
Self {
stdout: Some(stdout),
stderr: Some(stderr),
kill: ScriptedKill {
killed: Arc::new(AtomicBool::new(false)),
signal: Arc::new(Notify::new()),
feeders: Arc::new(feeders),
},
code,
timed_out,
signal,
exit_at: lifetime.map(|d| tokio::time::Instant::now() + d),
}
}
fn kill(&self) {
self.kill.fire();
}
}
impl Backend {
fn own_group(&self) -> Option<&Arc<ProcessGroup>> {
match self {
Backend::Real(real) => real.own_group.as_ref(),
Backend::Scripted(_) => None,
}
}
fn scripted_kill(&self) -> Option<ScriptedKill> {
match self {
Backend::Real(_) => None,
Backend::Scripted(s) => Some(s.kill.clone()),
}
}
fn take_stdout_reader(&mut self) -> Option<OutputReader> {
match self {
Backend::Real(real) => real.stdout_pipe.take().map(|p| Box::new(p) as OutputReader),
Backend::Scripted(s) => s.stdout.take().map(|p| Box::new(p) as OutputReader),
}
}
fn take_stderr_reader(&mut self) -> Option<OutputReader> {
match self {
Backend::Real(real) => real.stderr_pipe.take().map(|p| Box::new(p) as OutputReader),
Backend::Scripted(s) => s.stderr.take().map(|p| Box::new(p) as OutputReader),
}
}
}
impl RunningProcess {
pub(crate) fn from_spawned(s: Spawned) -> Self {
Self {
program: s.program,
backend: Backend::Real(Box::new(RealProc {
child: s.child,
own_group: s.own_group.map(Arc::new),
stdout_pipe: s.stdout,
stderr_pipe: s.stderr,
stdin_pipe: s.stdin,
stdin_task: s.stdin_task,
})),
timeout: s.timeout,
timeout_grace: s.timeout_grace,
timeout_signal: s.timeout_signal,
pid: s.pid,
stdout_encoding: s.stdout_encoding,
stderr_encoding: s.stderr_encoding,
stdout_handler: s.stdout_handler,
stderr_handler: s.stderr_handler,
stdout_tee: s.stdout_tee,
stderr_tee: s.stderr_tee,
buffer: s.buffer,
ok_codes: s.ok_codes,
stdout_sink: None,
stderr_sink: None,
stdout_pump: None,
stderr_pump: None,
stdin_error: None,
stdout_piped: s.stdout_piped,
deadline_task: None,
timeout_state: Arc::new(AtomicU8::new(TS_PENDING)),
cancel_token: s.cancel_token,
cancel_task: None,
cancel_at_exit: None,
started: Instant::now(),
start_time: SystemTime::now(),
}
}
pub(crate) fn from_scripted(command: &crate::command::Command, scripted: ScriptedProc) -> Self {
Self {
program: command.program_name(),
backend: Backend::Scripted(Box::new(scripted)),
timeout: command.configured_timeout(),
timeout_grace: command.configured_timeout_grace(),
timeout_signal: command.timeout_signal_raw(),
pid: None,
stdout_encoding: command.out_encoding(),
stderr_encoding: command.err_encoding(),
stdout_handler: command.stdout_handler(),
stderr_handler: command.stderr_handler(),
stdout_tee: command.stdout_tee_sink(),
stderr_tee: command.stderr_tee_sink(),
buffer: command.output_buffer_policy(),
ok_codes: command.ok_codes_vec(),
stdout_sink: None,
stderr_sink: None,
stdout_pump: None,
stderr_pump: None,
stdin_error: None,
stdout_piped: command.stdout_is_piped(),
deadline_task: None,
timeout_state: Arc::new(AtomicU8::new(TS_PENDING)),
cancel_token: command.cancel_token(),
cancel_task: None,
cancel_at_exit: None,
started: Instant::now(),
start_time: SystemTime::now(),
}
}
pub(crate) fn attach_group(&mut self, group: ProcessGroup) {
if let Backend::Real(real) = &mut self.backend {
real.own_group = Some(Arc::new(group));
}
self.arm_cancel_watchdog();
}
pub(crate) fn arm_cancel_watchdog(&mut self) {
{
if let Some(old) = self.cancel_task.take() {
old.abort();
}
let Some(token) = self.cancel_token.clone() else {
return;
};
let group_weak = self.backend.own_group().map(Arc::downgrade);
let pid = self.pid;
let timeout_state = self.timeout_state.clone();
self.cancel_task = Some(tokio::spawn(async move {
token.cancelled().await;
if timeout_state.load(Ordering::Acquire) != TS_PENDING {
return;
}
if let Some(g) = group_weak.and_then(|w| w.upgrade()) {
let _ = g.kill_all();
}
stream::kill_direct_child(pid);
}));
}
}
fn arm_scripted_deadline(&mut self) {
if self.deadline_task.is_some() {
return;
}
let (Some(limit), Some(kill)) = (self.timeout, self.backend.scripted_kill()) else {
return;
};
let started = self.started;
let timeout_state = self.timeout_state.clone();
self.deadline_task = Some(tokio::spawn(async move {
let remaining = limit
.checked_sub(started.elapsed())
.unwrap_or(Duration::ZERO);
tokio::time::sleep(remaining).await;
if timeout_state
.compare_exchange(
TS_PENDING,
TS_TIMED_OUT,
Ordering::AcqRel,
Ordering::Relaxed,
)
.is_err()
{
return; }
kill.fire();
}));
}
pub(crate) fn take_stdout_pipe(&mut self) -> Option<ChildStdout> {
match &mut self.backend {
Backend::Real(real) => real.stdout_pipe.take(),
Backend::Scripted(_) => None,
}
}
pub(crate) fn program_name(&self) -> &str {
&self.program
}
}
impl std::fmt::Debug for RunningProcess {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RunningProcess")
.field("program", &self.program)
.field("pid", &self.pid)
.field("timeout", &self.timeout)
.finish_non_exhaustive()
}
}
impl RunningProcess {
pub fn pid(&self) -> Option<u32> {
self.pid
}
pub fn start_time(&self) -> SystemTime {
self.start_time
}
pub fn elapsed(&self) -> Duration {
self.started.elapsed()
}
#[cfg(feature = "stats")]
pub fn cpu_time(&self) -> Option<Duration> {
self.pid
.and_then(|pid| crate::sys::process_metrics(pid).cpu_time)
}
#[cfg(feature = "stats")]
pub fn peak_memory_bytes(&self) -> Option<u64> {
self.pid
.and_then(|pid| crate::sys::process_metrics(pid).peak_memory_bytes)
}
pub(crate) fn deadline_arbiter(&self) -> Arc<AtomicU8> {
self.timeout_state.clone()
}
pub fn stdout_line_count(&self) -> usize {
self.stdout_sink.as_ref().map_or(0, |s| s.count())
}
pub fn stderr_line_count(&self) -> usize {
self.stderr_sink.as_ref().map_or(0, |s| s.count())
}
pub fn take_stdin(&mut self) -> Option<ProcessStdin> {
match &mut self.backend {
Backend::Real(real) => real.stdin_pipe.take().map(ProcessStdin::new),
Backend::Scripted(_) => None,
}
}
pub fn kills_tree_on_drop(&self) -> bool {
self.backend.own_group().is_some()
}
fn ensure_stdout_capturable(&self) -> Result<()> {
if self.stdout_piped {
return Ok(());
}
Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"`{}`: stdout is not piped (Command::stdout was set to Inherit/Null), so the \
capture verbs have nothing to read — use StdioMode::Piped to capture it",
self.program
),
)))
}
fn ensure_stdout_streamable(&self) -> Result<()> {
self.ensure_stdout_capturable()?; if self.stdout_sink.is_some() {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"`{}`: stdout was already consumed by an earlier stdout_lines/output_events \
call — stream it once (a second call would yield an empty stream)",
self.program
),
)));
}
Ok(())
}
pub async fn output_string(mut self) -> Result<ProcessResult<String>> {
let finished = self
.finish_lines(CaptureMode::Lines, true, || {})
.await?;
let truncated = self.stdout_sink.as_ref().is_some_and(|s| s.dropped() > 0)
|| self.stderr_sink.as_ref().is_some_and(|s| s.dropped() > 0);
let total_lines = self.stdout_sink.as_ref().map_or(0, |s| s.count())
+ self.stderr_sink.as_ref().map_or(0, |s| s.count());
let total_bytes = self.stdout_sink.as_ref().map_or(0, |s| s.seen_bytes())
+ self.stderr_sink.as_ref().map_or(0, |s| s.seen_bytes());
let duration = self.started.elapsed();
Ok(ProcessResult::new(
self.program.clone(),
finished.stdout_lines.join("\n"),
finished.stderr_lines.join("\n"),
finished.outcome,
self.timeout,
)
.with_duration(duration)
.with_truncated(truncated)
.with_overflow_totals(total_lines, total_bytes)
.with_ok_codes(self.ok_codes.clone()))
}
pub async fn output_bytes(mut self) -> Result<ProcessResult<Vec<u8>>> {
self.ensure_stdout_capturable()?;
if self.stdout_sink.is_some() || self.stderr_sink.is_some() {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"`{}`: output_bytes cannot follow a streaming call (stdout was already \
consumed as lines) — use output_string to collect the streamed lines, or \
call output_bytes without streaming first",
self.program
),
)));
}
let stderr_sink = SharedLines::new(&self.buffer);
self.stderr_pump = self.backend.take_stderr_reader().map(|pipe| {
tokio::spawn(pump_lines_core(
pipe,
self.stderr_encoding,
self.stderr_handler.clone(),
self.stderr_tee.clone(),
stderr_sink.clone(),
))
});
self.stderr_sink = Some(stderr_sink.clone());
let mut stdout_pipe = self.backend.take_stdout_reader();
let out_buf = Arc::new(std::sync::Mutex::new(Vec::new()));
self.stdout_pump = Some({
let out_buf = out_buf.clone();
tokio::spawn(async move {
if let Some(pipe) = &mut stdout_pipe {
let mut chunk = [0u8; 8 * 1024];
loop {
match pipe.read(&mut chunk).await {
Ok(0) => break,
Ok(n) => out_buf
.lock()
.expect("stdout buffer poisoned")
.extend_from_slice(&chunk[..n]),
Err(_e) => {
#[cfg(feature = "tracing")]
tracing::warn!(target: "processkit", error = %_e, "stdout read error; ending byte capture early");
break;
}
}
}
}
})
});
let outcome = self.drive_to_exit().await?;
self.observe_stdin_task().await;
if let Some(out_task) = self.stdout_pump.take() {
let abort = out_task.abort_handle();
if tokio::time::timeout(PUMP_TEARDOWN, out_task).await.is_err() {
abort.abort();
}
}
let stdout = std::mem::take(&mut *out_buf.lock().expect("stdout buffer poisoned"));
join_pumps(self.stderr_pump.take().into_iter().collect()).await;
let outcome = self.checked_outcome(outcome)?;
if stderr_sink.overflowed() {
return Err(crate::Error::OutputTooLarge {
program: self.program.clone(),
line_limit: self.buffer.max_lines,
byte_limit: self.buffer.max_bytes,
total_lines: stderr_sink.count(),
total_bytes: stderr_sink.seen_bytes(),
});
}
let stderr_lines = stderr_sink.drain();
let truncated = stderr_sink.dropped() > 0;
let duration = self.started.elapsed();
Ok(ProcessResult::new(
self.program.clone(),
stdout,
stderr_lines.join("\n"),
outcome,
self.timeout,
)
.with_duration(duration)
.with_truncated(truncated)
.with_overflow_totals(stderr_sink.count(), stderr_sink.seen_bytes())
.with_ok_codes(self.ok_codes.clone()))
}
pub async fn wait(mut self) -> Result<Outcome> {
Ok(self
.finish_lines(CaptureMode::Discard, false, || {})
.await?
.outcome)
}
pub async fn shutdown(mut self, grace: std::time::Duration) -> Result<Outcome> {
let Some(group) = self.backend.own_group().cloned() else {
return Err(Error::Unsupported {
operation: "shutdown (a shared-group handle does not own its group — \
use ProcessGroup::shutdown, or start_kill for just this child)"
.into(),
});
};
if let Some(limit) = self.timeout
&& self.started.elapsed() >= limit
{
let _ = self.timeout_state.compare_exchange(
TS_PENDING,
TS_TIMED_OUT,
Ordering::AcqRel,
Ordering::Relaxed,
);
}
self.timeout = None;
if let Some(task) = self.deadline_task.take() {
task.abort();
}
let (term_result, outcome) = tokio::join!(
group.graceful_terminate(grace, crate::sys::SIGTERM_RAW),
self.wait(),
);
term_result?;
outcome
}
pub(crate) async fn wait_exit(&mut self) -> Result<Outcome> {
let cause = if self.cancel_at_exit.is_some() {
ExitCause::Exited(self.backend_wait().await?)
} else {
let token = self.cancel_token.clone();
let cancelled = async {
match &token {
Some(token) => token.cancelled().await,
None => std::future::pending::<()>().await,
}
};
tokio::select! {
biased; () = cancelled => {
self.kill_tree().await;
ExitCause::Cancelled
}
outcome = self.backend_wait() => ExitCause::Exited(outcome?),
}
};
let outcome = self.on_reaped(cause);
self.observe_stdin_task().await;
self.checked_outcome(outcome)
}
#[cfg(feature = "stats")]
pub async fn profile(mut self, every: Duration) -> Result<crate::stats::RunProfile> {
use std::sync::{Arc, Mutex};
#[derive(Default)]
struct Acc {
cpu_time: Option<Duration>,
peak_memory_bytes: Option<u64>,
samples: usize,
}
let every = every.max(Duration::from_millis(1));
let started = self.started;
let acc = Arc::new(Mutex::new(Acc::default()));
let reaped = Arc::new(AtomicBool::new(false));
let sampler = self.pid.map(|pid| {
let acc = Arc::clone(&acc);
let reaped = Arc::clone(&reaped);
tokio::spawn(async move {
let mut ticker = tokio::time::interval(every);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
ticker.tick().await;
if reaped.load(Ordering::Acquire) {
break;
}
let metrics = crate::sys::process_metrics(pid);
if reaped.load(Ordering::Acquire) {
break;
}
if let Ok(mut acc) = acc.lock() {
acc.samples += 1;
if let Some(cpu) = metrics.cpu_time {
acc.cpu_time = Some(cpu);
}
if let Some(peak) = metrics.peak_memory_bytes {
acc.peak_memory_bytes =
Some(acc.peak_memory_bytes.map_or(peak, |prev| prev.max(peak)));
}
}
}
})
});
struct AbortOnDrop(tokio::task::AbortHandle);
impl Drop for AbortOnDrop {
fn drop(&mut self) {
self.0.abort();
}
}
let _sampler_guard = sampler.as_ref().map(|h| AbortOnDrop(h.abort_handle()));
let outcome = self
.finish_lines(CaptureMode::Discard, false, || {
reaped.store(true, Ordering::Release);
if let Some(task) = &sampler {
task.abort();
}
})
.await?
.outcome;
let exit_code = outcome.code();
let duration = started.elapsed();
let (cpu_time, peak_memory_bytes, samples) = match acc.lock() {
Ok(acc) => (acc.cpu_time, acc.peak_memory_bytes, acc.samples),
Err(_) => (None, None, 0),
};
Ok(crate::stats::RunProfile {
outcome,
exit_code,
duration,
cpu_time,
peak_memory_bytes,
samples,
})
}
async fn finish_lines(
&mut self,
capture: CaptureMode,
expose_counts: bool,
on_exit: impl FnOnce(),
) -> Result<FinishedLines> {
if matches!(capture, CaptureMode::Lines) {
self.ensure_stdout_capturable()?;
}
let discard_policy = OutputBufferPolicy::bounded(0).with_max_bytes(DISCARD_INFLIGHT_CAP);
let sink_policy: &OutputBufferPolicy = match capture {
CaptureMode::Discard => &discard_policy,
CaptureMode::Lines => &self.buffer,
};
let stdout_sink = self
.stdout_sink
.clone()
.unwrap_or_else(|| SharedLines::new(sink_policy));
let stderr_sink = self
.stderr_sink
.clone()
.unwrap_or_else(|| SharedLines::new(sink_policy));
self.spawn_line_pumps(&stdout_sink, &stderr_sink);
if expose_counts {
if self.stdout_sink.is_none() {
self.stdout_sink = Some(stdout_sink.clone());
}
if self.stderr_sink.is_none() {
self.stderr_sink = Some(stderr_sink.clone());
}
}
let outcome = self.drive_to_exit().await;
on_exit();
let outcome = outcome?;
self.observe_stdin_task().await;
let pumps: Vec<_> = [self.stdout_pump.take(), self.stderr_pump.take()]
.into_iter()
.flatten()
.collect();
join_pumps(pumps).await;
let outcome = self.checked_outcome(outcome)?;
if matches!(capture, CaptureMode::Lines) {
for sink in [&stdout_sink, &stderr_sink] {
if sink.overflowed() {
return Err(crate::Error::OutputTooLarge {
program: self.program.clone(),
line_limit: self.buffer.max_lines,
byte_limit: self.buffer.max_bytes,
total_lines: sink.count(),
total_bytes: sink.seen_bytes(),
});
}
}
}
let (stdout_lines, stderr_lines) = match capture {
CaptureMode::Lines => (stdout_sink.drain(), stderr_sink.drain()),
CaptureMode::Discard => (Vec::new(), Vec::new()),
};
Ok(FinishedLines {
outcome,
stdout_lines,
stderr_lines,
})
}
fn spawn_line_pumps(&mut self, stdout_sink: &Arc<SharedLines>, stderr_sink: &Arc<SharedLines>) {
if let Some(pipe) = self.backend.take_stdout_reader() {
self.stdout_pump = Some(tokio::spawn(pump_lines_core(
pipe,
self.stdout_encoding,
self.stdout_handler.clone(),
self.stdout_tee.clone(),
stdout_sink.clone(),
)));
}
if let Some(pipe) = self.backend.take_stderr_reader() {
self.stderr_pump = Some(tokio::spawn(pump_lines_core(
pipe,
self.stderr_encoding,
self.stderr_handler.clone(),
self.stderr_tee.clone(),
stderr_sink.clone(),
)));
}
}
fn checked_outcome(&mut self, outcome: Outcome) -> Result<Outcome> {
if self.cancel_at_exit.unwrap_or(false) {
return Err(Error::Cancelled {
program: self.program.clone(),
});
}
let succeeded = matches!(outcome, Outcome::Exited(code) if self.ok_codes.contains(&code));
if succeeded && let Some(source) = self.stdin_error.take() {
return Err(Error::Stdin {
program: self.program.clone(),
source,
});
}
Ok(outcome)
}
async fn observe_stdin_task(&mut self) {
let task = match &mut self.backend {
Backend::Real(real) => real.stdin_task.take(),
Backend::Scripted(_) => None,
};
let Some(task) = task else {
return;
};
if !task.is_finished() {
if let Backend::Real(real) = &mut self.backend {
real.stdin_task = Some(task);
}
return;
}
let observed = match task.await {
Ok(Ok(())) => None,
Ok(Err(e)) if is_broken_pipe(&e) => None,
Ok(Err(e)) => Some(e),
Err(join_err) => Some(std::io::Error::other(if join_err.is_panic() {
format!("stdin writer task panicked: {join_err}")
} else {
format!("stdin writer task did not complete: {join_err}")
})),
};
if let Some(e) = observed {
#[cfg(feature = "tracing")]
tracing::warn!(
target: "processkit",
program = %self.program,
error = %e,
"stdin writer failed"
);
self.stdin_error = Some(e);
}
}
fn abort_watchdogs(&mut self) {
self.pid = None;
if let Some(task) = self.deadline_task.take() {
task.abort();
}
if let Some(task) = self.cancel_task.take() {
task.abort();
}
}
fn on_reaped(&mut self, cause: ExitCause) -> Outcome {
if self.cancel_at_exit.is_none() {
self.cancel_at_exit = Some(matches!(cause, ExitCause::Cancelled));
}
self.abort_watchdogs();
let outcome = match cause {
ExitCause::Exited(outcome) => outcome,
ExitCause::Cancelled => Outcome::Signalled(None),
};
self.classify_timed_out(outcome)
}
async fn drive_to_exit(&mut self) -> Result<Outcome> {
if let Backend::Real(real) = &mut self.backend {
drop(real.stdin_pipe.take());
}
let cause = if self.cancel_at_exit.is_some() {
ExitCause::Exited(self.backend_wait().await?)
} else {
self.drive_to_exit_inner().await?
};
let outcome = self.on_reaped(cause);
#[cfg(feature = "tracing")]
tracing::debug!(
target: "processkit",
program = %self.program,
outcome = ?outcome,
elapsed_ms = self.started.elapsed().as_millis() as u64,
"process exited"
);
Ok(outcome)
}
fn classify_timed_out(&self, outcome: Outcome) -> Outcome {
if self.timeout_state.load(Ordering::Acquire) == TS_TIMED_OUT {
Outcome::TimedOut
} else {
outcome
}
}
async fn backend_wait(&mut self) -> Result<Outcome> {
let outcome = match &mut self.backend {
Backend::Real(real) => {
let status = real.child.wait().await.map_err(Error::Io)?;
match status.code() {
Some(code) => Outcome::Exited(code),
None => {
#[cfg(unix)]
{
use std::os::unix::process::ExitStatusExt;
Outcome::Signalled(status.signal())
}
#[cfg(not(unix))]
Outcome::Signalled(None)
}
}
}
Backend::Scripted(s) => {
let already_exited =
matches!(s.exit_at, Some(at) if at <= tokio::time::Instant::now());
let classify = |s: &ScriptedProc| match (s.code, s.timed_out) {
(_, true) => Outcome::TimedOut,
(Some(code), false) => Outcome::Exited(code),
(None, false) => Outcome::Signalled(s.signal),
};
if s.kill.killed.load(Ordering::Acquire) && !already_exited {
Outcome::Signalled(None)
} else if already_exited {
classify(s)
} else {
match s.exit_at {
Some(at) => {
tokio::select! {
biased;
() = s.kill.signal.notified() => Outcome::Signalled(None),
() = tokio::time::sleep_until(at) => classify(s),
}
}
None => {
s.kill.signal.notified().await;
Outcome::Signalled(None)
}
}
}
}
};
let _ = self.timeout_state.compare_exchange(
TS_PENDING,
TS_EXITED,
Ordering::AcqRel,
Ordering::Relaxed,
);
Ok(outcome)
}
async fn drive_to_exit_inner(&mut self) -> Result<ExitCause> {
let limit = self.timeout;
let token = self.cancel_token.clone();
let started = self.started;
let watchdog_owns_deadline = self.deadline_task.is_some();
let cancelled = async {
match &token {
Some(token) => token.cancelled().await,
None => std::future::pending::<()>().await,
}
};
let deadline = async move {
match limit {
Some(limit) if !watchdog_owns_deadline => {
let remaining = limit
.checked_sub(started.elapsed())
.unwrap_or(Duration::ZERO);
tokio::time::sleep(remaining).await
}
_ => std::future::pending::<()>().await,
}
};
tokio::select! {
biased; () = cancelled => {
#[cfg(feature = "tracing")]
tracing::debug!(
target: "processkit",
program = %self.program,
"cancellation fired; killing the tree"
);
self.kill_tree().await;
Ok(ExitCause::Cancelled)
}
outcome = self.backend_wait() => outcome.map(ExitCause::Exited),
() = deadline => {
#[cfg(feature = "tracing")]
tracing::warn!(
target: "processkit",
program = %self.program,
timeout_ms = limit.map(|l| l.as_millis() as u64).unwrap_or(0),
"timeout elapsed; killing the tree"
);
let _ = self.timeout_state.compare_exchange(
TS_PENDING,
TS_TIMED_OUT,
Ordering::AcqRel,
Ordering::Relaxed,
);
self.teardown_on_timeout().await;
Ok(ExitCause::Exited(Outcome::TimedOut))
}
}
}
async fn kill_tree(&mut self) {
match &mut self.backend {
Backend::Real(real) => {
let _ = real.child.start_kill();
if let Some(group) = &real.own_group {
let _ = group.kill_all();
}
let _ = tokio::time::timeout(PUMP_TEARDOWN, real.child.wait()).await;
}
Backend::Scripted(s) => s.kill(),
}
}
async fn teardown_on_timeout(&mut self) {
let Some(grace) = self.timeout_grace else {
self.kill_tree().await;
return;
};
let signal = self.timeout_signal;
match &mut self.backend {
Backend::Real(real) => {
let pid = real.child.id();
let own = real.own_group.clone();
let teardown = async {
match &own {
Some(group) => {
let _ = group.graceful_terminate(grace, signal).await;
}
None => {
crate::running::stream::graceful_kill_pid(pid, grace, signal).await;
}
}
};
let reap =
tokio::time::timeout(grace.saturating_add(PUMP_TEARDOWN), real.child.wait());
let _ = tokio::join!(teardown, reap);
}
Backend::Scripted(s) => s.kill(),
}
}
fn has_exited_now(&mut self) -> bool {
let exited = match &mut self.backend {
Backend::Real(real) => matches!(real.child.try_wait(), Ok(Some(_))),
Backend::Scripted(s) => {
s.kill.killed.load(Ordering::Acquire)
|| s.exit_at
.is_some_and(|at| tokio::time::Instant::now() >= at)
}
};
if exited {
let _ = self.timeout_state.compare_exchange(
TS_PENDING,
TS_EXITED,
Ordering::AcqRel,
Ordering::Relaxed,
);
self.abort_watchdogs();
if self.cancel_at_exit.is_none() {
self.cancel_at_exit =
Some(self.cancel_token.as_ref().is_some_and(|t| t.is_cancelled()));
}
}
exited
}
pub fn start_kill(&mut self) -> Result<()> {
match &mut self.backend {
Backend::Real(real) => match real.child.start_kill() {
Ok(()) => {}
Err(e) if e.kind() == std::io::ErrorKind::InvalidInput => {}
Err(e) => return Err(Error::Io(e)),
},
Backend::Scripted(s) => s.kill(),
}
Ok(())
}
}
impl Drop for RunningProcess {
fn drop(&mut self) {
match &mut self.backend {
Backend::Real(real) => {
if let Some(task) = real.stdin_task.take() {
task.abort();
}
}
Backend::Scripted(s) => s.kill(),
}
if let Some(task) = self.deadline_task.take() {
task.abort();
}
if let Some(task) = self.cancel_task.take() {
task.abort();
}
if let Some(task) = self.stdout_pump.take() {
task.abort();
}
if let Some(task) = self.stderr_pump.take() {
task.abort();
}
}
}
fn is_broken_pipe(e: &std::io::Error) -> bool {
e.kind() == std::io::ErrorKind::BrokenPipe || matches!(e.raw_os_error(), Some(109 | 232))
}
async fn join_pumps(tasks: Vec<JoinHandle<()>>) {
if tasks.is_empty() {
return;
}
let aborts: Vec<_> = tasks.iter().map(|t| t.abort_handle()).collect();
let join = async {
for task in tasks {
#[cfg(feature = "tracing")]
if let Err(e) = task.await {
tracing::warn!(target: "processkit", error = %e, "output pump task ended abnormally");
}
#[cfg(not(feature = "tracing"))]
let _ = task.await;
}
};
if tokio::time::timeout(PUMP_TEARDOWN, join).await.is_err() {
#[cfg(feature = "tracing")]
tracing::warn!(
target: "processkit",
timeout_ms = PUMP_TEARDOWN.as_millis() as u64,
aborted = aborts.len(),
"output pumps overran teardown grace; aborting stragglers"
);
for abort in aborts {
abort.abort();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::command::Command;
use crate::doubles::{Reply, ScriptedRunner};
use crate::runner::ProcessRunner;
async fn scripted_handle(ok_codes: &[i32]) -> RunningProcess {
let cmd = Command::new("tool").ok_codes(ok_codes.iter().copied());
ScriptedRunner::new()
.fallback(Reply::ok(""))
.start(&cmd)
.await
.expect("scripted start")
}
#[tokio::test]
async fn stdin_error_surfaces_only_on_a_successful_outcome() {
let mut run = scripted_handle(&[0]).await;
run.stdin_error = Some(std::io::Error::other("boom"));
match run.checked_outcome(Outcome::Exited(0)) {
Err(Error::Stdin { program, source }) => {
assert_eq!(program, "tool");
assert_eq!(source.to_string(), "boom");
}
other => panic!("expected Error::Stdin, got {other:?}"),
}
let mut run = scripted_handle(&[0]).await;
run.stdin_error = Some(std::io::Error::other("boom"));
assert!(matches!(
run.checked_outcome(Outcome::Exited(7)),
Ok(Outcome::Exited(7))
));
let mut run = scripted_handle(&[0]).await;
run.stdin_error = Some(std::io::Error::other("boom"));
assert!(matches!(
run.checked_outcome(Outcome::Signalled(Some(9))),
Ok(Outcome::Signalled(Some(9)))
));
}
#[tokio::test]
async fn stdin_error_respects_ok_codes_widened_success() {
let mut run = scripted_handle(&[0, 3]).await;
run.stdin_error = Some(std::io::Error::other("boom"));
assert!(matches!(
run.checked_outcome(Outcome::Exited(3)),
Err(Error::Stdin { .. })
));
}
#[tokio::test]
async fn no_stdin_error_is_a_clean_passthrough() {
let mut run = scripted_handle(&[0]).await;
assert!(matches!(
run.checked_outcome(Outcome::Exited(0)),
Ok(Outcome::Exited(0))
));
}
#[tokio::test]
async fn output_string_after_partial_stream_is_not_truncated() {
use tokio_stream::StreamExt;
let mut run = ScriptedRunner::new()
.fallback(Reply::lines(["a", "b", "c", "d"]))
.start(&Command::new("tool"))
.await
.expect("scripted start");
{
let mut lines = run.stdout_lines().unwrap();
assert_eq!(lines.next().await.as_deref(), Some("a"));
assert_eq!(lines.next().await.as_deref(), Some("b"));
}
let result = run.output_string().await.expect("output_string");
assert!(
!result.truncated(),
"consumed lines are not truncation under unbounded policy: {result:?}"
);
assert_eq!(
result.stdout(),
"c\nd",
"output_string returns the unconsumed tail"
);
}
#[tokio::test]
async fn output_bytes_after_streaming_errors_instead_of_empty() {
let mut run = ScriptedRunner::new()
.fallback(Reply::lines(["a", "b"]))
.start(&Command::new("tool"))
.await
.expect("scripted start");
drop(run.stdout_lines().unwrap());
let err = run
.output_bytes()
.await
.expect_err("output_bytes after streaming must error, not return empty");
match err {
Error::Io(e) => assert_eq!(e.kind(), std::io::ErrorKind::InvalidInput),
other => panic!("expected Io(InvalidInput), got {other:?}"),
}
}
#[tokio::test]
async fn output_string_after_stream_still_reports_real_truncation() {
let cmd = Command::new("tool").output_buffer(OutputBufferPolicy::bounded(2));
let mut run = ScriptedRunner::new()
.fallback(Reply::lines(["a", "b", "c", "d"]))
.start(&cmd)
.await
.expect("scripted start");
drop(run.stdout_lines().unwrap());
let result = run.output_string().await.expect("output_string");
assert!(
result.truncated(),
"a bounded buffer that dropped lines during streaming must report truncation: {result:?}"
);
}
#[tokio::test]
async fn output_bytes_returns_exact_raw_stdout() {
let result = ScriptedRunner::new()
.fallback(Reply::ok("raw\u{0}bytes\nno trailing newline"))
.start(&Command::new("tool"))
.await
.expect("scripted start")
.output_bytes()
.await
.expect("output_bytes");
assert_eq!(result.stdout(), b"raw\x00bytes\nno trailing newline");
assert!(!result.truncated(), "no policy drop: {result:?}");
}
#[tokio::test]
async fn timed_out_flag_classifies_a_clean_exit_as_timed_out() {
let run = scripted_handle(&[0]).await; run.timeout_state.store(TS_TIMED_OUT, Ordering::Release); let outcome = run.wait().await.expect("wait");
assert_eq!(
outcome,
Outcome::TimedOut,
"a run whose deadline fired must report TimedOut, not the in-grace exit"
);
}
#[tokio::test]
async fn cancellation_beats_the_timed_out_flag() {
let token = crate::CancellationToken::new();
let run = ScriptedRunner::new()
.fallback(Reply::ok(""))
.start(&Command::new("tool").cancel_on(token.clone()))
.await
.expect("scripted start");
run.timeout_state.store(TS_TIMED_OUT, Ordering::Release);
token.cancel();
match run.wait().await {
Err(Error::Cancelled { .. }) => {}
other => panic!("expected Err(Cancelled), got {other:?}"),
}
}
#[tokio::test]
async fn a_natural_wait_any_exit_is_not_flipped_by_a_late_cancel() {
let token = crate::CancellationToken::new();
let mut run = ScriptedRunner::new()
.fallback(Reply::ok("done\n"))
.start(&Command::new("tool").cancel_on(token.clone()))
.await
.expect("scripted start");
let (idx, outcome) = crate::wait_any(&mut [&mut run]).await.expect("wait_any");
assert_eq!((idx, outcome), (0, Outcome::Exited(0)));
token.cancel();
let outcome = run
.wait()
.await
.expect("a late cancel must not flip a natural exit");
assert_eq!(outcome, Outcome::Exited(0));
}
#[tokio::test]
async fn wait_any_classifies_a_timed_out_run() {
let mut run = scripted_handle(&[0]).await; run.timeout_state.store(TS_TIMED_OUT, Ordering::Release); let (idx, outcome) = crate::wait_any(&mut [&mut run]).await.expect("wait_any");
assert_eq!(idx, 0);
assert_eq!(
outcome,
Outcome::TimedOut,
"a timed-out run must report TimedOut through wait_any, not the raw exit"
);
}
#[tokio::test]
async fn natural_reap_claim_beats_a_late_timeout_cas() {
let run = scripted_handle(&[0]).await;
assert!(
run.timeout_state
.compare_exchange(TS_PENDING, TS_EXITED, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
);
assert!(
run.timeout_state
.compare_exchange(
TS_PENDING,
TS_TIMED_OUT,
Ordering::AcqRel,
Ordering::Relaxed
)
.is_err()
);
assert_eq!(
run.classify_timed_out(Outcome::Exited(0)),
Outcome::Exited(0)
);
}
#[tokio::test]
async fn scripted_handle_does_not_kill_a_tree_on_drop() {
let run = scripted_handle(&[0]).await;
assert!(
!run.kills_tree_on_drop(),
"a scripted double has no OS tree to tear down"
);
}
#[tokio::test]
async fn capture_verbs_error_on_a_non_piped_stdout() {
let runner = ScriptedRunner::new().fallback(Reply::ok("ignored"));
let run = runner
.start(&Command::new("tool").stdout(crate::StdioMode::Null))
.await
.unwrap();
match run.output_string().await {
Err(Error::Io(e)) => assert_eq!(e.kind(), std::io::ErrorKind::InvalidInput),
other => panic!("expected Io(InvalidInput), got {other:?}"),
}
let run = runner
.start(&Command::new("tool").stdout(crate::StdioMode::Inherit))
.await
.unwrap();
assert!(matches!(run.output_bytes().await, Err(Error::Io(_))));
let run = ScriptedRunner::new()
.fallback(Reply::ok("hi"))
.start(&Command::new("tool"))
.await
.unwrap();
assert_eq!(run.output_string().await.unwrap().stdout(), "hi");
let run = runner
.start(&Command::new("tool").stdout(crate::StdioMode::Null))
.await
.unwrap();
assert!(
run.wait().await.is_ok(),
"discard verbs do not require a piped stdout"
);
}
}