mod deadline;
mod probes;
mod scripted;
mod stream;
pub use stream::{Finished, OutputEvent, OutputEvents, OutputLine, StdoutLines};
pub(crate) use scripted::{ScriptedProc, ScriptedResultInfo, split_pump_lines};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::{Duration, Instant, SystemTime};
use crate::sync::atomic::AtomicU8;
use tokio::io::AsyncReadExt;
use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout};
use tokio::task::JoinHandle;
use crate::buffer::{OutputBufferPolicy, OverflowMode, clamp_dropoldest_tail, push_capped_bytes};
use crate::error::Error;
use crate::error::Result;
use crate::group::ProcessGroup;
use crate::pump::{SharedLines, StreamConfig, pump_lines_core};
use crate::result::{Outcome, ProcessResult};
use crate::stdin::ProcessStdin;
use crate::sys::pid_gate::PidGate;
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_config: StreamConfig,
pub stderr_config: StreamConfig,
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>,
#[cfg(feature = "stats")]
proc_identity: Option<crate::sys::ProcIdentity>,
stdout_config: StreamConfig,
stderr_config: StreamConfig,
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>,
pid_gate: Arc<PidGate>,
cancel_token: Option<tokio_util::sync::CancellationToken>,
cancel_task: Option<JoinHandle<()>>,
cancel_at_exit: Option<bool>,
started: Instant,
deadline_anchor: tokio::time::Instant,
start_time: SystemTime,
scripted_result: Option<ScriptedResultInfo>,
}
type OutputReader = Box<dyn tokio::io::AsyncRead + Send + Unpin>;
enum Backend {
Real(Box<RealProc>),
Scripted(Box<ScriptedProc>),
}
struct RealProc {
child: Option<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<()>>>,
}
impl RealProc {
fn child_mut(&mut self) -> &mut Child {
self.child
.as_mut()
.expect("child is present until Drop extracts it")
}
}
impl Backend {
fn own_group(&self) -> Option<&Arc<ProcessGroup>> {
match self {
Backend::Real(real) => real.own_group.as_ref(),
Backend::Scripted(s) => s.own_group(),
}
}
fn scripted_kill(&self) -> Option<scripted::ScriptedKill> {
match self {
Backend::Real(_) => None,
Backend::Scripted(s) => Some(s.kill_handle()),
}
}
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.take_stdout_reader(),
}
}
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.take_stderr_reader(),
}
}
}
impl RunningProcess {
pub(crate) fn from_spawned(s: Spawned) -> Self {
Self {
program: s.program,
backend: Backend::Real(Box::new(RealProc {
child: Some(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,
#[cfg(feature = "stats")]
proc_identity: s.pid.and_then(crate::sys::process_identity),
stdout_config: s.stdout_config,
stderr_config: s.stderr_config,
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)),
pid_gate: Arc::new(PidGate::new(s.pid)),
cancel_token: s.cancel_token,
cancel_task: None,
cancel_at_exit: None,
started: Instant::now(),
deadline_anchor: tokio::time::Instant::now(),
start_time: SystemTime::now(),
scripted_result: None,
}
}
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 own_group_handle(&self) -> Option<Arc<ProcessGroup>> {
self.backend.own_group().cloned()
}
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 gate = self.pid_gate.clone();
self.cancel_task = Some(tokio::spawn(async move {
token.cancelled().await;
if gate.is_retired() {
return;
}
if let Some(g) = group_weak.and_then(|w| w.upgrade()) {
let _ = g.kill_all();
}
crate::sys::pid_gate::force_kill(&gate);
}));
}
}
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, self.proc_identity).cpu_time)
}
#[cfg(feature = "stats")]
pub fn peak_memory_bytes(&self) -> Option<u64> {
self.pid
.and_then(|pid| crate::sys::process_metrics(pid, self.proc_identity).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(crate::error::stdout_not_piped_error(&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, total_lines, total_bytes, duration) = match self.scripted_result {
Some(rec) => (
rec.truncated,
rec.total_lines,
rec.total_bytes,
rec.duration,
),
None => {
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());
(truncated, total_lines, total_bytes, 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_config.clone(),
stderr_sink.clone(),
))
});
self.stderr_sink = Some(stderr_sink.clone());
let stdout_cap = self.buffer.max_bytes;
let stdout_mode = self.buffer.overflow;
let signals = RawStdoutSignals {
seen: Arc::new(AtomicUsize::new(0)),
overflowed: Arc::new(AtomicBool::new(false)),
truncated: Arc::new(AtomicBool::new(false)),
read_error: Arc::new(std::sync::Mutex::new(None)),
};
let stdout_pipe = self.backend.take_stdout_reader();
let out_buf = Arc::new(std::sync::Mutex::new(Vec::new()));
self.stdout_pump = stdout_pipe.map(|pipe| {
tokio::spawn(pump_raw_bytes(
pipe,
out_buf.clone(),
stdout_cap,
stdout_mode,
signals.clone(),
))
});
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 mut stdout = std::mem::take(&mut *out_buf.lock().expect("stdout buffer poisoned"));
clamp_dropoldest_tail(&mut stdout, stdout_cap, stdout_mode);
join_pumps(self.stderr_pump.take().into_iter().collect()).await;
self.finalize_stdin_task().await;
let outcome = self.checked_outcome(outcome)?;
if signals.overflowed.load(Ordering::Relaxed) {
return Err(crate::Error::OutputTooLarge {
program: self.program.clone(),
max_lines: None,
max_bytes: self.buffer.max_bytes,
total_lines: 0,
total_bytes: signals.seen.load(Ordering::Relaxed),
});
}
if stderr_sink.overflowed() {
return Err(crate::Error::OutputTooLarge {
program: self.program.clone(),
max_lines: self.buffer.max_lines,
max_bytes: self.buffer.max_bytes,
total_lines: stderr_sink.count(),
total_bytes: stderr_sink.seen_bytes(),
});
}
if let Some(source) = signals
.read_error
.lock()
.expect("stdout read-error slot poisoned")
.take()
{
return Err(Error::Io(source));
}
if let Some(source) = stderr_sink.take_read_error() {
return Err(Error::Io(source));
}
let stderr_lines = stderr_sink.drain();
let truncated = signals.truncated.load(Ordering::Relaxed) || 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(),
signals
.seen
.load(Ordering::Relaxed)
.saturating_add(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.deadline_anchor.elapsed() >= limit
{
let _ = deadline::claim_timed_out(&self.timeout_state);
}
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};
let every = every.max(Duration::from_millis(1));
let started = self.started;
let acc = Arc::new(Mutex::new(ProfileAcc::default()));
let reaped = Arc::new(AtomicBool::new(false));
let identity = self.proc_identity;
let sampler = self.pid.map(|pid| {
let acc = Arc::clone(&acc);
let reaped = Arc::clone(&reaped);
tokio::spawn(run_profile_sampler(every, reaped, acc, move || {
crate::sys::process_metrics(pid, identity)
}))
});
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 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,
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 = discard_sink_policy();
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));
if matches!(capture, CaptureMode::Discard) {
stdout_sink.start_discarding();
stderr_sink.start_discarding();
}
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;
self.finalize_stdin_task().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(),
max_lines: self.buffer.max_lines,
max_bytes: self.buffer.max_bytes,
total_lines: sink.count(),
total_bytes: sink.seen_bytes(),
});
}
}
}
for sink in [&stdout_sink, &stderr_sink] {
if let Some(source) = sink.take_read_error() {
return Err(Error::Io(source));
}
}
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_config.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_config.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 = Self::classify_stdin_join(task.await);
self.record_stdin_error(observed);
}
async fn finalize_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;
};
let abort = task.abort_handle();
let observed = match tokio::time::timeout(PUMP_TEARDOWN, task).await {
Ok(joined) => Self::classify_stdin_join(joined),
Err(_elapsed) => {
abort.abort();
None
}
};
self.record_stdin_error(observed);
}
fn classify_stdin_join(
joined: std::result::Result<std::io::Result<()>, tokio::task::JoinError>,
) -> Option<std::io::Error> {
match joined {
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}")
})),
}
}
fn record_stdin_error(&mut self, observed: Option<std::io::Error>) {
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_gate.retire();
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 gate = self.pid_gate.clone();
let outcome = match &mut self.backend {
Backend::Real(real) => {
let status = gated_reap(&gate, real.child_mut())
.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) => s.wait_outcome().await,
};
self.pid_gate.retire();
let _ = deadline::claim_exited(&self.timeout_state);
Ok(outcome)
}
async fn drive_to_exit_inner(&mut self) -> Result<ExitCause> {
self.pid_gate.retire();
if let Some(task) = self.deadline_task.take() {
task.abort();
}
let limit = self.timeout;
let token = self.cancel_token.clone();
let started = self.deadline_anchor;
let cancelled = async {
match &token {
Some(token) => token.cancelled().await,
None => std::future::pending::<()>().await,
}
};
let timeout_state = self.timeout_state.clone();
let deadline = async move {
match limit {
Some(limit) => {
deadline::wait_deadline_and_claim(started, limit, &timeout_state).await
}
None => std::future::pending::<bool>().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),
_won = 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"
);
self.teardown_on_timeout().await;
Ok(ExitCause::Exited(Outcome::TimedOut))
}
}
}
async fn kill_tree(&mut self) {
let gate = self.pid_gate.clone();
match &mut self.backend {
Backend::Real(real) => {
let _ = real.child_mut().start_kill();
gate.retire();
if let Some(group) = &real.own_group {
let _ = group.kill_all();
}
let _ = tokio::time::timeout(PUMP_TEARDOWN, real.child_mut().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;
let gate = self.pid_gate.clone();
match &mut self.backend {
Backend::Real(real) => match real.own_group.clone() {
Some(group) => {
let teardown = async move {
let _ = group.graceful_terminate(grace, signal).await;
};
let reap = async {
let r = tokio::time::timeout(
grace.saturating_add(PUMP_TEARDOWN),
real.child_mut().wait(),
)
.await;
gate.retire();
r
};
let _ = tokio::join!(teardown, reap);
}
None => {
#[cfg(unix)]
{
stream::signal_direct_child(real.child_mut().id(), signal);
let reaped_cleanly = matches!(
tokio::time::timeout(grace, real.child_mut().wait()).await,
Ok(Ok(_))
);
if !reaped_cleanly {
let _ = real.child_mut().start_kill();
let _ =
tokio::time::timeout(PUMP_TEARDOWN, real.child_mut().wait()).await;
}
}
#[cfg(not(unix))]
{
let _ = signal;
let _ = real.child_mut().start_kill();
let _ = tokio::time::timeout(PUMP_TEARDOWN, real.child_mut().wait()).await;
}
gate.retire();
}
},
Backend::Scripted(s) => s.kill(),
}
}
fn has_exited_now(&mut self) -> bool {
let gate = self.pid_gate.clone();
let exited = gate.reap_under_lock(|| match &mut self.backend {
Backend::Real(real) => matches!(real.child_mut().try_wait(), Ok(Some(_))),
Backend::Scripted(s) => s.has_exited_now(),
});
if exited {
let _ = deadline::claim_exited(&self.timeout_state);
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_mut().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) {
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();
}
match &mut self.backend {
Backend::Real(real) => {
if let Some(task) = real.stdin_task.take() {
task.abort();
}
if real.own_group.is_none()
&& self.timeout.is_some()
&& self.timeout_grace.is_some()
&& !self.pid_gate.is_retired()
&& let Ok(handle) = tokio::runtime::Handle::try_current()
&& let Some(child) = real.child.take()
{
let gate = self.pid_gate.clone();
handle.spawn(gated_reap_and_retire(gate, child));
} else {
self.pid_gate.retire();
}
}
Backend::Scripted(s) => s.kill(),
}
}
}
pub(crate) fn is_broken_pipe(e: &std::io::Error) -> bool {
e.kind() == std::io::ErrorKind::BrokenPipe || matches!(e.raw_os_error(), Some(109 | 232))
}
fn discard_sink_policy() -> OutputBufferPolicy {
OutputBufferPolicy::bounded(0).with_max_bytes(DISCARD_INFLIGHT_CAP)
}
#[cfg(feature = "stats")]
#[derive(Default)]
struct ProfileAcc {
cpu_time: Option<Duration>,
peak_memory_bytes: Option<u64>,
samples: usize,
}
#[cfg(feature = "stats")]
impl ProfileAcc {
fn fold(&mut self, metrics: crate::sys::ProcMetrics) {
self.samples += 1;
if let Some(cpu) = metrics.cpu_time {
self.cpu_time = Some(cpu);
}
if let Some(peak) = metrics.peak_memory_bytes {
self.peak_memory_bytes =
Some(self.peak_memory_bytes.map_or(peak, |prev| prev.max(peak)));
}
}
}
#[cfg(feature = "stats")]
async fn run_profile_sampler(
every: Duration,
reaped: Arc<AtomicBool>,
acc: Arc<std::sync::Mutex<ProfileAcc>>,
mut source: impl FnMut() -> crate::sys::ProcMetrics,
) {
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 = source();
if reaped.load(Ordering::Acquire) {
break;
}
if let Ok(mut acc) = acc.lock() {
acc.fold(metrics);
}
}
}
async fn gated_reap(
gate: &PidGate,
child: &mut Child,
) -> std::io::Result<std::process::ExitStatus> {
use std::future::Future;
let mut wait = std::pin::pin!(child.wait());
std::future::poll_fn(|cx| {
let mut out = std::task::Poll::Pending;
gate.reap_under_lock(|| match wait.as_mut().poll(cx) {
std::task::Poll::Ready(res) => {
out = std::task::Poll::Ready(res);
true
}
std::task::Poll::Pending => false,
});
out
})
.await
}
async fn gated_reap_and_retire(gate: Arc<PidGate>, mut child: Child) {
let _ = gated_reap(&gate, &mut child).await;
gate.retire();
}
#[derive(Clone)]
struct RawStdoutSignals {
seen: Arc<AtomicUsize>,
overflowed: Arc<AtomicBool>,
truncated: Arc<AtomicBool>,
read_error: Arc<std::sync::Mutex<Option<std::io::Error>>>,
}
async fn pump_raw_bytes<R>(
mut reader: R,
out_buf: Arc<std::sync::Mutex<Vec<u8>>>,
cap: Option<usize>,
mode: OverflowMode,
signals: RawStdoutSignals,
) where
R: tokio::io::AsyncRead + Unpin,
{
let mut chunk = [0u8; 8 * 1024];
loop {
match reader.read(&mut chunk).await {
Ok(0) => break,
Ok(n) => {
signals.seen.fetch_add(n, Ordering::Relaxed);
let mut guard = out_buf.lock().expect("stdout buffer poisoned");
push_capped_bytes(
&mut guard,
&chunk[..n],
cap,
mode,
&signals.overflowed,
&signals.truncated,
);
}
Err(e) if is_broken_pipe(&e) => break,
Err(e) => {
#[cfg(feature = "tracing")]
tracing::warn!(target: "processkit", error = %e, "stdout read error; ending byte capture early");
*signals
.read_error
.lock()
.expect("stdout read-error slot poisoned") = Some(e);
break;
}
}
}
}
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 bare_finish_does_not_error_under_fail_loud_on_uncaptured_stdout() {
let cmd = Command::new("tool").output_buffer(OutputBufferPolicy::fail_loud(1));
let finished = ScriptedRunner::new()
.fallback(Reply::lines(["a", "b", "c", "d"]))
.start(&cmd)
.await
.expect("scripted start")
.finish()
.await
.expect("bare finish must not error under fail_loud");
assert_eq!(finished.outcome, Outcome::Exited(0));
let outcome = ScriptedRunner::new()
.fallback(Reply::lines(["a", "b", "c", "d"]))
.start(&cmd)
.await
.expect("scripted start")
.wait()
.await
.expect("wait must not error under fail_loud either");
assert_eq!(
outcome, finished.outcome,
"bare finish and wait agree on the uncaptured-stdout outcome"
);
}
#[tokio::test]
async fn bare_finish_discards_stdout_but_returns_stderr() {
let finished = ScriptedRunner::new()
.fallback(Reply::fail(0, "e1\ne2\n").with_stdout("o1\no2\n"))
.start(&Command::new("tool"))
.await
.expect("scripted start")
.finish()
.await
.expect("bare finish");
assert_eq!(finished.outcome, Outcome::Exited(0));
assert_eq!(finished.stderr, "e1\ne2");
}
#[tokio::test]
async fn bare_finish_reports_stderr_truncated_when_the_policy_drops_lines() {
let cmd = Command::new("tool").output_buffer(OutputBufferPolicy::bounded(2));
let finished = ScriptedRunner::new()
.fallback(Reply::fail(1, "e1\ne2\ne3\ne4\n"))
.start(&cmd)
.await
.expect("scripted start")
.finish()
.await
.expect("bare finish");
assert!(
finished.stderr_truncated,
"a bounded policy that dropped stderr lines must set stderr_truncated: {finished:?}"
);
let untouched = ScriptedRunner::new()
.fallback(Reply::fail(1, "e1\ne2\ne3\ne4\n"))
.start(&Command::new("tool"))
.await
.expect("scripted start")
.finish()
.await
.expect("bare finish");
assert!(
!untouched.stderr_truncated,
"an unbounded policy must not report stderr truncation: {untouched:?}"
);
}
#[tokio::test]
async fn wait_after_a_dropped_stream_completes() {
let mut run = ScriptedRunner::new()
.fallback(Reply::lines(["a", "b", "c", "d"]))
.start(&Command::new("tool"))
.await
.expect("scripted start");
drop(run.stdout_lines().expect("stdout_lines"));
assert_eq!(
run.wait().await.expect("wait after a dropped stream"),
Outcome::Exited(0)
);
}
#[cfg(feature = "stats")]
#[tokio::test]
async fn profile_after_a_dropped_stream_completes() {
let mut run = ScriptedRunner::new()
.fallback(Reply::lines(["a", "b", "c", "d"]))
.start(&Command::new("tool"))
.await
.expect("scripted start");
drop(run.stdout_lines().expect("stdout_lines"));
let profile = run
.profile(Duration::from_millis(1))
.await
.expect("profile after a dropped stream");
assert_eq!(profile.outcome, Outcome::Exited(0));
}
#[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 a_probe_reap_is_not_flipped_by_a_later_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");
match run
.wait_for(|| async { false }, Duration::from_secs(5))
.await
{
Err(Error::NotReady { .. }) => {}
other => panic!("expected Err(NotReady), got {other:?}"),
}
token.cancel();
let outcome = run
.wait()
.await
.expect("a cancel after the probe's reap observation must not flip it");
assert_eq!(outcome, Outcome::Exited(0));
}
#[tokio::test]
async fn a_probe_reap_with_an_active_cancel_still_surfaces_cancelled() {
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");
token.cancel();
match run
.wait_for(|| async { false }, Duration::from_secs(5))
.await
{
Err(Error::NotReady { .. }) => {}
other => panic!("expected Err(NotReady), got {other:?}"),
}
match run.wait().await {
Err(Error::Cancelled { .. }) => {}
other => panic!("expected Err(Cancelled), got {other:?}"),
}
}
#[tokio::test]
async fn a_consuming_reap_retires_the_pid_gate() {
let run = scripted_handle(&[0]).await;
let gate = run.pid_gate.clone();
assert!(!gate.is_retired(), "a fresh handle's gate is live");
run.wait().await.expect("wait");
assert!(
gate.is_retired(),
"the reap must retire the gate so the watchdogs stand down"
);
}
#[tokio::test]
async fn a_cancelled_reap_retires_the_pid_gate() {
let token = crate::CancellationToken::new();
let run = ScriptedRunner::new()
.fallback(Reply::ok(""))
.start(&Command::new("tool").cancel_on(token.clone()))
.await
.expect("scripted start");
let gate = run.pid_gate.clone();
token.cancel();
let _ = run.wait().await;
assert!(
gate.is_retired(),
"the cancellation reap path must retire the gate"
);
}
#[tokio::test]
async fn a_probe_reap_retires_the_pid_gate() {
let mut run = scripted_handle(&[0]).await; let gate = run.pid_gate.clone();
let _ = run
.wait_for(|| async { false }, Duration::from_secs(5))
.await;
assert!(
gate.is_retired(),
"the probe's reap_under_lock must retire the gate atomically"
);
}
#[tokio::test]
async fn a_wait_any_reap_retires_the_pid_gate() {
let mut run = scripted_handle(&[0]).await; let gate = run.pid_gate.clone();
assert!(!gate.is_retired(), "a fresh handle's gate is live");
let (idx, outcome) = crate::wait_any(&mut [&mut run]).await.expect("wait_any");
assert_eq!((idx, outcome), (0, Outcome::Exited(0)));
assert!(
gate.is_retired(),
"the wait_any (backend_wait) reap must retire the gate so the watchdogs \
stand down"
);
}
#[tokio::test]
async fn a_wait_all_reap_retires_the_pid_gate() {
let mut run = scripted_handle(&[0]).await;
let gate = run.pid_gate.clone();
let outcomes = crate::wait_all(&mut [&mut run]).await.expect("wait_all");
assert_eq!(outcomes, vec![Outcome::Exited(0)]);
assert!(
gate.is_retired(),
"the wait_all reap must retire the gate too"
);
}
fn sleeper_cmd() -> Command {
if cfg!(windows) {
Command::new("cmd").args(["/c", "ping", "-n", "30", "127.0.0.1"])
} else {
Command::new("sleep").arg("30")
}
}
#[tokio::test]
#[ignore = "spawns a real subprocess (shared-group Drop-branch gate retirement)"]
async fn dropping_a_shared_group_handle_without_grace_retires_the_gate() {
let group = crate::group::ProcessGroup::new().expect("a shared process group");
let cmd = sleeper_cmd().timeout(Duration::from_secs(30));
let run = crate::runner::launch(&group, &cmd)
.await
.expect("launch into the shared group");
assert!(
!run.kills_tree_on_drop(),
"a shared-group handle owns no tree — its group does"
);
assert_eq!(
run.timeout_grace, None,
"the branch under test has a timeout but no graceful window"
);
let gate = run.pid_gate.clone();
assert!(
!gate.is_retired(),
"a fresh live handle's gate is not retired"
);
drop(run); assert!(
gate.is_retired(),
"Drop must retire the gate so a mid-poll watchdog's raw kill is a \
linearized no-op, never a SIGKILL on a recycled pid"
);
drop(group);
}
#[tokio::test]
#[ignore = "spawns a real subprocess (own-group Drop-branch gate retirement)"]
async fn dropping_an_own_group_handle_retires_the_gate() {
let cmd = sleeper_cmd().timeout(Duration::from_secs(30));
let run = crate::runner::JobRunner::new()
.start(&cmd)
.await
.expect("start a private-group run");
assert!(
run.kills_tree_on_drop(),
"a private-group handle tears its whole tree down on drop"
);
let gate = run.pid_gate.clone();
assert!(
!gate.is_retired(),
"a fresh live handle's gate is not retired"
);
drop(run); assert!(
gate.is_retired(),
"the own-group Drop branch must retire the gate too"
);
}
#[test] #[ignore = "spawns a real subprocess (no-runtime shared-group+grace Drop gate retirement)"]
fn dropping_a_shared_group_grace_handle_with_no_runtime_retires_the_gate() {
let rt = tokio::runtime::Runtime::new().expect("a test runtime");
let group = crate::group::ProcessGroup::new().expect("a shared process group");
let (run, gate) = rt.block_on(async {
let cmd = sleeper_cmd()
.timeout(Duration::from_secs(30))
.timeout_grace(Duration::from_secs(5));
let run = crate::runner::launch(&group, &cmd)
.await
.expect("launch into the shared group");
let gate = run.pid_gate.clone();
(run, gate)
});
assert!(
!run.kills_tree_on_drop(),
"a shared-group handle owns no tree — its group does (own_group is None)"
);
assert!(
run.timeout.is_some() && run.timeout_grace.is_some(),
"the shape under test has both a timeout and a graceful window"
);
assert!(
!gate.is_retired(),
"a fresh live handle's gate is not retired"
);
assert!(
tokio::runtime::Handle::try_current().is_err(),
"the drop below must run with no current runtime — the scenario under test"
);
drop(run); assert!(
gate.is_retired(),
"Drop must retire the gate even with no runtime current, so a watchdog \
aborted mid-poll can't land a raw kill on the freed/recycled pid"
);
drop(group);
drop(rt);
}
#[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"
);
}
struct RawChunkedReader {
chunks: std::collections::VecDeque<Vec<u8>>,
err_at_end: Option<std::io::Error>,
}
impl RawChunkedReader {
fn new(
chunks: impl IntoIterator<Item = Vec<u8>>,
err_at_end: Option<std::io::Error>,
) -> Self {
Self {
chunks: chunks.into_iter().collect(),
err_at_end,
}
}
}
impl tokio::io::AsyncRead for RawChunkedReader {
fn poll_read(
mut self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
if let Some(chunk) = self.chunks.pop_front() {
let n = chunk.len().min(buf.remaining());
buf.put_slice(&chunk[..n]);
if n < chunk.len() {
self.chunks.push_front(chunk[n..].to_vec());
}
std::task::Poll::Ready(Ok(()))
} else if let Some(err) = self.err_at_end.take() {
std::task::Poll::Ready(Err(err))
} else {
std::task::Poll::Ready(Ok(())) }
}
}
async fn drive_pump_raw_bytes(reader: RawChunkedReader) -> (Vec<u8>, Option<std::io::Error>) {
let out_buf = Arc::new(std::sync::Mutex::new(Vec::new()));
let signals = RawStdoutSignals {
seen: Arc::new(AtomicUsize::new(0)),
overflowed: Arc::new(AtomicBool::new(false)),
truncated: Arc::new(AtomicBool::new(false)),
read_error: Arc::new(std::sync::Mutex::new(None)),
};
pump_raw_bytes(
reader,
out_buf.clone(),
None,
OverflowMode::DropOldest,
signals.clone(),
)
.await;
let bytes = std::mem::take(&mut *out_buf.lock().unwrap());
let err = signals.read_error.lock().unwrap().take();
(bytes, err)
}
#[tokio::test]
async fn pump_raw_bytes_records_a_mid_stream_error_and_keeps_the_prefix() {
let (bytes, err) = drive_pump_raw_bytes(RawChunkedReader::new(
[b"partial".to_vec()],
Some(std::io::Error::other("boom")),
))
.await;
assert_eq!(
bytes, b"partial",
"the prefix read before the error is kept"
);
assert!(
err.is_some(),
"the raw stdout OS read error is recorded for output_bytes to surface as Error::Io"
);
}
#[tokio::test]
async fn pump_raw_bytes_clean_eof_records_no_error() {
let (bytes, err) = drive_pump_raw_bytes(RawChunkedReader::new(
[b"all".to_vec(), b"good".to_vec()],
None,
))
.await;
assert_eq!(bytes, b"allgood");
assert!(err.is_none(), "a clean EOF is a complete capture");
}
#[tokio::test]
async fn pump_raw_bytes_treats_a_broken_pipe_read_as_clean_eof() {
let (bytes, err) = drive_pump_raw_bytes(RawChunkedReader::new(
[b"done".to_vec()],
Some(std::io::Error::from(std::io::ErrorKind::BrokenPipe)),
))
.await;
assert_eq!(bytes, b"done", "the prefix is kept");
assert!(
err.is_none(),
"a broken-pipe read is the normal writer-closed end, not an incomplete capture"
);
}
#[tokio::test]
async fn output_string_surfaces_a_recorded_read_error_as_io() {
let mut run = scripted_handle(&[0]).await; let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
sink.set_read_error(std::io::Error::other("stdout read boom"));
run.stdout_sink = Some(sink);
match run.output_string().await {
Err(Error::Io(e)) => assert_eq!(e.to_string(), "stdout read boom"),
other => panic!("expected Err(Io) for an incomplete capture, got {other:?}"),
}
}
#[tokio::test]
async fn wait_surfaces_a_recorded_read_error_as_io() {
let mut run = scripted_handle(&[0]).await;
let sink = SharedLines::new(&OutputBufferPolicy::unbounded());
sink.set_read_error(std::io::Error::other("stderr read boom"));
run.stderr_sink = Some(sink);
match run.wait().await {
Err(Error::Io(e)) => assert_eq!(e.to_string(), "stderr read boom"),
other => panic!("expected Err(Io) for an incomplete capture, got {other:?}"),
}
}
}
#[cfg(all(test, feature = "stats"))]
mod profile_sampler_tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;
use super::{ProfileAcc, run_profile_sampler};
use crate::sys::ProcMetrics;
fn metrics(cpu_ms: u64, mem: u64) -> ProcMetrics {
ProcMetrics {
cpu_time: Some(Duration::from_millis(cpu_ms)),
peak_memory_bytes: Some(mem),
}
}
#[test]
fn fold_ignores_all_none_readings() {
let mut acc = ProfileAcc::default();
acc.fold(metrics(100, 8192)); acc.fold(ProcMetrics::default()); acc.fold(metrics(200, 4096)); assert_eq!(acc.samples, 3, "every tick is counted, even empty ones");
assert_eq!(
acc.cpu_time,
Some(Duration::from_millis(200)),
"CPU tracks the latest real reading, not the empty one"
);
assert_eq!(
acc.peak_memory_bytes,
Some(8192),
"peak is the max across real readings; an empty reading never lowers it"
);
}
#[tokio::test(start_paused = true)]
async fn sampler_folds_only_identity_matched_readings() {
let reaped = Arc::new(AtomicBool::new(false));
let acc = Arc::new(std::sync::Mutex::new(ProfileAcc::default()));
let calls = Arc::new(AtomicUsize::new(0));
let reaped_src = Arc::clone(&reaped);
let calls_src = Arc::clone(&calls);
let source = move || {
let n = calls_src.fetch_add(1, Ordering::Relaxed);
if n >= 4 {
reaped_src.store(true, Ordering::Release);
}
match n {
0 => metrics(100, 8192), 1 => metrics(200, 4096), _ => ProcMetrics::default(),
}
};
run_profile_sampler(
Duration::from_millis(5),
Arc::clone(&reaped),
Arc::clone(&acc),
source,
)
.await;
let acc = acc.lock().expect("acc mutex");
assert_eq!(
acc.cpu_time,
Some(Duration::from_millis(200)),
"the last identity-matched CPU reading is kept; the stranger's default is ignored"
);
assert_eq!(
acc.peak_memory_bytes,
Some(8192),
"peak reflects only identity-matched readings — the recycled pid never enters it"
);
assert!(
acc.samples >= 2,
"at least the two real readings were sampled"
);
}
}