use std::sync::{Arc, Weak};
use std::time::Duration;
use tokio::task::JoinHandle;
use crate::command::Command;
use crate::error::Result;
use crate::group::ProcessGroup;
use crate::result::{Outcome, ProcessResult};
use crate::running::{
Finished, LineCapture, ProcessEvents, RawCapture, RunningProcess, StdoutLines,
};
use crate::sync::atomic::{AtomicU8, Ordering};
const TEARDOWN_DRAIN_GRACE: Duration = Duration::from_millis(500);
const LAST_STAGE_PROBE_MIN: Duration = Duration::from_millis(25);
const LAST_STAGE_PROBE_MAX: Duration = Duration::from_millis(500);
#[must_use = "a Pipeline does nothing until it is run"]
#[derive(Clone)]
pub struct Pipeline {
stages: Vec<Command>,
timeout: Option<Duration>,
cancel_token: Option<tokio_util::sync::CancellationToken>,
}
impl std::fmt::Debug for Pipeline {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Pipeline")
.field("stages", &self.stages.len())
.field("timeout", &self.timeout)
.finish_non_exhaustive()
}
}
struct StageOutcome {
program: String,
outcome: Outcome,
stderr: String,
unchecked: bool,
ok_codes: Vec<i32>,
timeout: Option<Duration>,
torn_down: bool,
stderr_truncated: bool,
}
enum Joined<T> {
Inner(usize, StageOutcome),
Last(ProcessResult<T>, bool),
}
struct Captured<T> {
stdout: T,
stderr: String,
truncated: bool,
total_lines: usize,
total_bytes: usize,
}
trait PipelineCapture: Send + Clone + 'static {
type Tracker: Clone;
fn prepare(process: &mut RunningProcess) -> Result<Self::Tracker>;
fn snapshot(tracker: &Self::Tracker) -> Captured<Self>;
}
impl PipelineCapture for String {
type Tracker = LineCapture;
fn prepare(process: &mut RunningProcess) -> Result<Self::Tracker> {
process.prepare_line_capture()
}
fn snapshot(tracker: &Self::Tracker) -> Captured<Self> {
let (stdout, stderr, truncated, total_lines, total_bytes) = tracker.snapshot();
Captured {
stdout,
stderr,
truncated,
total_lines,
total_bytes,
}
}
}
impl PipelineCapture for Vec<u8> {
type Tracker = RawCapture;
fn prepare(process: &mut RunningProcess) -> Result<Self::Tracker> {
process.prepare_raw_capture()
}
fn snapshot(tracker: &Self::Tracker) -> Captured<Self> {
let (stdout, stderr, truncated, total_lines, total_bytes) = tracker.snapshot();
Captured {
stdout,
stderr,
truncated,
total_lines,
total_bytes,
}
}
}
type LastExitObserver = Box<dyn FnOnce() + Send>;
fn captured_result<T>(result: ProcessResult<T>) -> Captured<T> {
let stderr = result.stderr().to_owned();
let truncated = result.truncated();
let total_lines = result.total_lines();
let total_bytes = result.total_bytes();
Captured {
stdout: result.into_stdout(),
stderr,
truncated,
total_lines,
total_bytes,
}
}
struct LaunchedChain {
stage_groups: Vec<Arc<ProcessGroup>>,
running: Vec<(RunningProcess, bool)>,
started: std::time::Instant,
}
struct DetachedLast {
handle: RunningProcess,
program: String,
ok_codes: Vec<i32>,
timeout: Option<Duration>,
unchecked: bool,
}
impl Pipeline {
pub(crate) fn new(first: Command, second: Command) -> Self {
Pipeline {
stages: vec![first, second],
timeout: None,
cancel_token: None,
}
}
pub fn pipe(mut self, next: Command) -> Self {
self.stages.push(next);
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
pub fn cancel_on(mut self, token: tokio_util::sync::CancellationToken) -> Self {
self.cancel_token = Some(token);
self
}
async fn launch(&self) -> Result<LaunchedChain> {
if let Some(index) = self.stages[..self.stages.len() - 1]
.iter()
.position(Command::wants_pty)
{
return Err(crate::ErrorReason::Unsupported {
operation: format!("pipeline use_pty on non-final stage {}", index + 1),
}
.into());
}
let started = std::time::Instant::now();
let mut stage_groups: Vec<Arc<ProcessGroup>> = Vec::with_capacity(self.stages.len());
let mut running = Vec::with_capacity(self.stages.len());
let mut upstream = None;
for (index, stage) in self.stages.iter().enumerate() {
let mut command = stage.clone();
let non_final = index + 1 < self.stages.len();
if non_final && stage.wants_stderr_merged_in_pipe() {
command.activate_stderr_merge_in_pipe();
}
if non_final && !stage.stdout_is_piped() {
command.pipe_stdout_for_downstream();
}
if let Some(token) = &self.cancel_token
&& command.cancel_token().is_none()
{
command = command.cancel_on(token.clone());
}
if let Some(reader) = upstream.take() {
command.set_pipe_stdin(reader);
}
let group = ProcessGroup::new()?;
let mut process = group.start(&command).await?;
process.attach_group(group);
if let Some(handle) = process.own_group_handle() {
stage_groups.push(handle);
}
if non_final {
upstream = process.take_stdout_pipe();
}
running.push((process, stage.is_unchecked()));
}
Ok(LaunchedChain {
stage_groups,
running,
started,
})
}
fn detach_last(&self, running: &mut Vec<(RunningProcess, bool)>) -> DetachedLast {
let (handle, unchecked) = running.pop().expect("a pipeline has at least two stages");
let stage = self
.stages
.last()
.expect("a pipeline has at least two stages");
DetachedLast {
program: handle.program_name().to_owned(),
ok_codes: stage.ok_codes_vec(),
timeout: stage.configured_timeout(),
unchecked,
handle,
}
}
pub async fn start(&self) -> Result<PipelineSession> {
let LaunchedChain {
stage_groups,
mut running,
started: _,
} = self.launch().await?;
let DetachedLast {
handle: last,
program: last_program,
ok_codes: last_ok_codes,
timeout: last_timeout,
unchecked: last_unchecked,
} = self.detach_last(&mut running);
let inner_count = running.len();
let teardown = tokio_util::sync::CancellationToken::new();
let mut inner_tasks: tokio::task::JoinSet<Result<(usize, StageOutcome)>> =
tokio::task::JoinSet::new();
for (index, ((process, unchecked), stage)) in
running.into_iter().zip(self.stages.iter()).enumerate()
{
let program = process.program_name().to_owned();
let ok_codes = stage.ok_codes_vec();
let timeout = stage.configured_timeout();
let teardown = teardown.clone();
inner_tasks.spawn(async move {
let result = finish_inner_stage(
process,
index,
program,
ok_codes,
timeout,
unchecked,
teardown.clone(),
)
.await;
if result.is_err() {
teardown.cancel();
}
result
});
}
let killer = spawn_group_killer(teardown.clone(), &stage_groups);
let last: SharedLast = Arc::new(std::sync::Mutex::new(Some(last)));
let last_disposition = ExitDisposition::unobserved();
let last_watch = spawn_last_stage_watcher(
&last,
last_ok_codes.clone(),
last_unchecked,
teardown.clone(),
last_disposition.clone(),
);
let chain_state = Arc::new(AtomicU8::new(crate::running::TS_PENDING));
let deadline_task = self.timeout.map(|limit| {
let state = chain_state.clone();
let groups: Vec<Weak<ProcessGroup>> = stage_groups.iter().map(Arc::downgrade).collect();
let anchor = tokio::time::Instant::now();
tokio::spawn(async move {
if crate::running::deadline::wait_deadline_and_claim(anchor, limit, &state).await {
kill_weak_stage_groups(&groups);
}
})
});
Ok(PipelineSession {
last,
last_program,
last_ok_codes,
last_unchecked,
last_timeout,
last_disposition,
last_watch: Some(last_watch),
inner_tasks: Some(inner_tasks),
inner_count,
stage_groups,
teardown,
timeout: self.timeout,
chain_state,
deadline_task,
killer: Some(killer),
})
}
pub async fn output_string(&self) -> Result<ProcessResult<String>> {
self.capture(
|last, at_exit| async move { last.output_string_observing_exit(at_exit).await },
)
.await
}
pub async fn output_bytes(&self) -> Result<ProcessResult<Vec<u8>>> {
self.capture(|last, at_exit| async move { last.output_bytes_observing_exit(at_exit).await })
.await
}
async fn capture<T, C, F>(&self, capture_last: C) -> Result<ProcessResult<T>>
where
T: PipelineCapture,
C: FnOnce(crate::running::RunningProcess, LastExitObserver) -> F,
F: std::future::Future<Output = Result<ProcessResult<T>>> + Send + 'static,
{
let LaunchedChain {
stage_groups,
mut running,
started,
} = self.launch().await?;
let teardown = tokio_util::sync::CancellationToken::new();
let (mut last, last_unchecked) = running.pop().expect("a pipeline has at least two stages");
let last_stage = self
.stages
.last()
.expect("a pipeline has at least two stages");
let last_ok_codes = last_stage.ok_codes_vec();
let last_timeout = last_stage.configured_timeout();
let capture = T::prepare(&mut last)?;
let completed: Arc<std::sync::Mutex<Option<ProcessResult<T>>>> =
Arc::new(std::sync::Mutex::new(None));
let inner_count = running.len();
let mut tasks: tokio::task::JoinSet<Result<Joined<T>>> = tokio::task::JoinSet::new();
for (index, ((process, unchecked), stage)) in
running.into_iter().zip(self.stages.iter()).enumerate()
{
let program = process.program_name().to_owned();
let ok_codes = stage.ok_codes_vec();
let timeout = stage.configured_timeout();
let teardown = teardown.clone();
tasks.spawn(async move {
let (index, outcome) = finish_inner_stage(
process, index, program, ok_codes, timeout, unchecked, teardown,
)
.await?;
Ok(Joined::Inner(index, outcome))
});
}
let last_disposition = ExitDisposition::unobserved();
let last_future = capture_last(last, {
let disposition = last_disposition.clone();
let teardown = teardown.clone();
Box::new(move || {
disposition.latch(teardown.is_cancelled());
})
});
{
let teardown = teardown.clone();
let last_ok_codes = last_ok_codes.clone();
let completed = completed.clone();
tasks.spawn(async move {
let result = last_future.await?;
*completed.lock().expect("pipeline capture result poisoned") = Some(result.clone());
let torn_down = last_disposition.latch(teardown.is_cancelled());
if !torn_down
&& is_checked_failure(result.outcome(), &last_ok_codes, last_unchecked)
{
teardown.cancel();
}
Ok(Joined::Last(result, torn_down))
});
}
let collect = async {
let gather = async {
let joined = drain_unordered(tasks, &teardown)
.await
.map_err(|failure| failure.error)?;
let mut inner_outcomes: Vec<Option<StageOutcome>> =
(0..inner_count).map(|_| None).collect();
let mut last_slot: Option<(ProcessResult<T>, bool)> = None;
for item in joined {
match item {
Joined::Inner(index, outcome) => inner_outcomes[index] = Some(outcome),
Joined::Last(result, torn_down) => last_slot = Some((result, torn_down)),
}
}
let outcomes: Vec<StageOutcome> = inner_outcomes
.into_iter()
.map(|outcome| {
outcome.expect("every inner stage slot is filled when every task succeeded")
})
.collect();
let (last_result, last_torn_down) =
last_slot.expect("last slot is filled when every task succeeded");
Ok::<_, crate::Error>((outcomes, last_result, last_torn_down))
};
tokio::select! {
collected = gather => collected,
() = async {
teardown.cancelled().await;
tokio::time::sleep(TEARDOWN_DRAIN_GRACE).await;
kill_all_stage_groups(&stage_groups);
std::future::pending::<()>().await
} => unreachable!("the teardown killer pends forever after firing"),
}
};
let (mut stages, last_result, last_torn_down) = match self.timeout {
None => collect.await?,
Some(limit) => match tokio::time::timeout(limit, collect).await {
Ok(collected) => collected?,
Err(_elapsed) => {
kill_all_stage_groups(&stage_groups);
let captured = completed
.lock()
.expect("pipeline capture result poisoned")
.take()
.map(captured_result)
.unwrap_or_else(|| T::snapshot(&capture));
let Captured {
stdout,
stderr,
truncated,
total_lines,
total_bytes,
} = captured;
return Ok(ProcessResult::new(
self.pipeline_name(),
stdout,
stderr,
Outcome::TimedOut,
Some(limit),
)
.with_duration(started.elapsed())
.with_truncated(truncated)
.with_overflow_totals(total_lines, total_bytes));
}
},
};
let last_truncated = last_result.truncated();
let (last_total_lines, last_total_bytes) =
(last_result.total_lines(), last_result.total_bytes());
let last_outcome = StageOutcome {
program: last_result.program().to_owned(),
outcome: last_result.outcome(),
stderr: last_result.stderr().to_owned(),
unchecked: last_unchecked,
ok_codes: last_ok_codes,
timeout: last_timeout,
torn_down: last_torn_down,
stderr_truncated: last_truncated,
};
let last_stdout = last_result.into_stdout();
stages.push(last_outcome);
let mut result = pipefail(stages, last_stdout).with_duration(started.elapsed());
if last_truncated {
result = result
.with_truncated(true)
.with_overflow_totals(last_total_lines, last_total_bytes);
}
Ok(result)
}
pub async fn run(&self) -> Result<String> {
let out = self.checked().await?;
self.reject_if_last_truncated(&out)?;
Ok(out.into_stdout().trim_end().to_owned())
}
pub async fn checked(&self) -> Result<ProcessResult<String>> {
self.output_string().await?.ensure_success()
}
pub async fn run_unit(&self) -> Result<()> {
self.output_string().await?.ensure_success().map(drop)
}
pub async fn exit_code(&self) -> Result<i32> {
self.output_string().await?.require_code()
}
pub async fn probe(&self) -> Result<bool> {
let result = self.output_string().await?;
result.probe_bool()
}
pub async fn parse<T, F>(&self, parse: F) -> Result<T>
where
F: FnOnce(&str) -> T,
{
let out = self.checked().await?;
self.reject_if_last_truncated(&out)?;
Ok(parse(out.stdout()))
}
pub async fn try_parse<T, F>(&self, parse: F) -> Result<T>
where
F: FnOnce(&str) -> Result<T>,
{
let out = self.checked().await?;
self.reject_if_last_truncated(&out)?;
parse(out.stdout())
}
fn reject_if_last_truncated(&self, out: &ProcessResult<String>) -> Result<()> {
let policy = self
.stages
.last()
.expect("a pipeline has at least two stages")
.output_buffer_policy();
out.reject_if_truncated(policy.max_lines, policy.max_bytes)
}
fn pipeline_name(&self) -> String {
self.stages
.iter()
.map(|stage| stage.program_name())
.collect::<Vec<_>>()
.join(" | ")
}
}
#[must_use = "a PipelineSession streams a live chain; drop it and the whole chain is killed unread"]
pub struct PipelineSession {
last: SharedLast,
last_program: String,
last_ok_codes: Vec<i32>,
last_unchecked: bool,
last_timeout: Option<Duration>,
last_disposition: ExitDisposition,
last_watch: Option<JoinHandle<()>>,
inner_tasks: Option<tokio::task::JoinSet<Result<(usize, StageOutcome)>>>,
inner_count: usize,
stage_groups: Vec<Arc<ProcessGroup>>,
teardown: tokio_util::sync::CancellationToken,
timeout: Option<Duration>,
chain_state: Arc<AtomicU8>,
deadline_task: Option<JoinHandle<()>>,
killer: Option<JoinHandle<()>>,
}
impl std::fmt::Debug for PipelineSession {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PipelineSession")
.field("last", &self.last_program)
.field("inner_stages", &self.inner_count)
.field("timeout", &self.timeout)
.finish_non_exhaustive()
}
}
impl PipelineSession {
pub fn stdout_lines(&mut self) -> Result<StdoutLines> {
self.with_last(RunningProcess::stdout_lines)
}
pub fn events(&mut self) -> Result<ProcessEvents> {
self.with_last(RunningProcess::events)
}
pub async fn wait_for_line(
&mut self,
predicate: impl Fn(&str) -> bool + Send,
within: Duration,
) -> Result<String> {
let mut last = self.borrow_last();
last.get().wait_for_line(predicate, within).await
}
pub fn pid(&self) -> Option<u32> {
lock_last(&self.last).as_ref().and_then(RunningProcess::pid)
}
pub fn start_kill(&mut self) -> Result<()> {
kill_all_stage_groups(&self.stage_groups);
Ok(())
}
pub async fn finish(mut self) -> Result<Finished> {
let (last, inner_tasks) = self.take_live_parts();
let teardown = self.teardown.clone();
let last_ok_codes = self.last_ok_codes.clone();
let last_unchecked = self.last_unchecked;
let last_fut = {
let teardown = teardown.clone();
let disposition = self.last_disposition.clone();
async move {
let result = last
.finish_observing_exit({
let disposition = disposition.clone();
let teardown = teardown.clone();
move |_outcome| {
disposition.latch(teardown.is_cancelled());
}
})
.await;
let torn_down = disposition.latch(teardown.is_cancelled());
match &result {
Ok(finished) => {
if !torn_down
&& is_checked_failure(finished.outcome, &last_ok_codes, last_unchecked)
{
teardown.cancel();
}
}
Err(_) if !torn_down => teardown.cancel(),
Err(_) => {}
}
(result, torn_down)
}
};
let ((last_res, last_torn_down), inner_res) =
tokio::join!(last_fut, drain_unordered(inner_tasks, &teardown));
self.abort_background();
if self.chain_timed_out() {
return Ok(timeout_finished(
inner_res,
last_res,
&self.last_program,
&self.last_ok_codes,
self.last_unchecked,
self.last_timeout,
last_torn_down,
));
}
let last_finished = last_res?;
let mut inner = inner_res.map_err(|failure| failure.error)?;
inner.sort_by_key(|(index, _)| *index);
let mut stages: Vec<StageOutcome> = inner.into_iter().map(|(_, outcome)| outcome).collect();
stages.push(StageOutcome {
program: self.last_program.clone(),
outcome: last_finished.outcome,
stderr: last_finished.stderr,
unchecked: self.last_unchecked,
ok_codes: self.last_ok_codes.clone(),
timeout: self.last_timeout,
torn_down: last_torn_down,
stderr_truncated: last_finished.stderr_truncated,
});
let folded = pipefail(stages, ());
Ok(Finished {
outcome: folded.outcome(),
stderr: folded.stderr().to_owned(),
stderr_truncated: folded.truncated(),
})
}
fn with_last<R>(&mut self, f: impl FnOnce(&mut RunningProcess) -> R) -> R {
let mut slot = lock_last(&self.last);
f(slot
.as_mut()
.expect("the last stage is live until finish consumes the session"))
}
fn borrow_last(&mut self) -> LastBorrow<'_> {
let handle = lock_last(&self.last)
.take()
.expect("the last stage is live until finish consumes the session");
LastBorrow {
slot: &self.last,
handle: Some(handle),
}
}
#[allow(clippy::type_complexity)]
fn take_live_parts(
&mut self,
) -> (
RunningProcess,
tokio::task::JoinSet<Result<(usize, StageOutcome)>>,
) {
if let Some(task) = self.last_watch.take() {
task.abort();
}
let last = lock_last(&self.last)
.take()
.expect("finish consumes the session exactly once");
let inner_tasks = self
.inner_tasks
.take()
.expect("finish consumes the session exactly once");
(last, inner_tasks)
}
fn chain_timed_out(&self) -> bool {
self.timeout.is_some()
&& !crate::running::deadline::claim_exited(&self.chain_state)
&& self.chain_state.load(Ordering::Acquire) == crate::running::TS_TIMED_OUT
}
fn abort_background(&mut self) {
if let Some(task) = self.deadline_task.take() {
task.abort();
}
if let Some(task) = self.killer.take() {
task.abort();
}
if let Some(task) = self.last_watch.take() {
task.abort();
}
}
}
impl Drop for PipelineSession {
fn drop(&mut self) {
self.abort_background();
}
}
type SharedLast = Arc<std::sync::Mutex<Option<RunningProcess>>>;
fn lock_last(slot: &std::sync::Mutex<Option<RunningProcess>>) -> LastGuard<'_> {
slot.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
type LastGuard<'a> = std::sync::MutexGuard<'a, Option<RunningProcess>>;
struct LastBorrow<'a> {
slot: &'a std::sync::Mutex<Option<RunningProcess>>,
handle: Option<RunningProcess>,
}
impl LastBorrow<'_> {
fn get(&mut self) -> &mut RunningProcess {
self.handle
.as_mut()
.expect("a borrowed last stage is put back only by Drop")
}
}
impl Drop for LastBorrow<'_> {
fn drop(&mut self) {
if let Some(handle) = self.handle.take() {
*lock_last(self.slot) = Some(handle);
}
}
}
fn kill_all_stage_groups(groups: &[Arc<ProcessGroup>]) {
for group in groups {
let _ = group.kill_all();
}
}
fn kill_weak_stage_groups(groups: &[Weak<ProcessGroup>]) {
for group in groups {
if let Some(group) = group.upgrade() {
let _ = group.kill_all();
}
}
}
fn spawn_group_killer(
teardown: tokio_util::sync::CancellationToken,
stage_groups: &[Arc<ProcessGroup>],
) -> JoinHandle<()> {
let groups: Vec<Weak<ProcessGroup>> = stage_groups.iter().map(Arc::downgrade).collect();
tokio::spawn(async move {
teardown.cancelled().await;
tokio::time::sleep(TEARDOWN_DRAIN_GRACE).await;
kill_weak_stage_groups(&groups);
})
}
#[derive(Clone, Debug)]
struct ExitDisposition(Arc<AtomicU8>);
const DISPOSITION_UNOBSERVED: u8 = 0;
const DISPOSITION_CULPRIT: u8 = 1;
const DISPOSITION_VICTIM: u8 = 2;
impl ExitDisposition {
fn unobserved() -> Self {
Self(Arc::new(AtomicU8::new(DISPOSITION_UNOBSERVED)))
}
fn latch(&self, torn_down: bool) -> bool {
let observed = if torn_down {
DISPOSITION_VICTIM
} else {
DISPOSITION_CULPRIT
};
match self.0.compare_exchange(
DISPOSITION_UNOBSERVED,
observed,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => torn_down,
Err(first) => first == DISPOSITION_VICTIM,
}
}
}
fn spawn_last_stage_watcher(
last: &SharedLast,
ok_codes: Vec<i32>,
unchecked: bool,
teardown: tokio_util::sync::CancellationToken,
disposition: ExitDisposition,
) -> JoinHandle<()> {
let slot = Arc::downgrade(last);
tokio::spawn(async move {
let mut delay = LAST_STAGE_PROBE_MIN;
loop {
tokio::time::sleep(delay).await;
delay = (delay * 2).min(LAST_STAGE_PROBE_MAX);
let Some(slot) = slot.upgrade() else {
return;
};
let observed = lock_last(&slot)
.as_mut()
.and_then(RunningProcess::exit_outcome_now);
drop(slot);
let Some(outcome) = observed else {
continue;
};
if !disposition.latch(teardown.is_cancelled())
&& is_checked_failure(outcome, &ok_codes, unchecked)
{
teardown.cancel();
}
return;
}
})
}
#[allow(clippy::too_many_arguments)]
async fn finish_inner_stage(
process: RunningProcess,
index: usize,
program: String,
ok_codes: Vec<i32>,
timeout: Option<Duration>,
unchecked: bool,
teardown: tokio_util::sync::CancellationToken,
) -> Result<(usize, StageOutcome)> {
let disposition = ExitDisposition::unobserved();
let Finished {
outcome,
stderr,
stderr_truncated,
} = process
.finish_observing_exit({
let disposition = disposition.clone();
let teardown = teardown.clone();
move |_outcome| {
disposition.latch(teardown.is_cancelled());
}
})
.await?;
let torn_down = disposition.latch(teardown.is_cancelled());
if !torn_down && is_checked_failure(outcome, &ok_codes, unchecked) {
teardown.cancel();
}
Ok((
index,
StageOutcome {
program,
outcome,
stderr,
unchecked,
ok_codes,
timeout,
torn_down,
stderr_truncated,
},
))
}
fn timeout_finished(
inner_res: InnerDrain,
last_res: Result<Finished>,
last_program: &str,
last_ok_codes: &[i32],
last_unchecked: bool,
last_timeout: Option<Duration>,
last_torn_down: bool,
) -> Finished {
let mut inner = match inner_res {
Ok(inner) => inner,
Err(failure) => failure.completed,
};
inner.sort_by_key(|(index, _)| *index);
let mut stages: Vec<StageOutcome> = inner.into_iter().map(|(_, outcome)| outcome).collect();
if let Ok(last_finished) = last_res {
stages.push(StageOutcome {
program: last_program.to_owned(),
outcome: last_finished.outcome,
stderr: last_finished.stderr,
unchecked: last_unchecked,
ok_codes: last_ok_codes.to_vec(),
timeout: last_timeout,
torn_down: last_torn_down,
stderr_truncated: last_finished.stderr_truncated,
});
}
if stages.is_empty() {
return Finished {
outcome: Outcome::TimedOut,
stderr: String::new(),
stderr_truncated: false,
};
}
let folded = pipefail(stages, ());
Finished {
outcome: Outcome::TimedOut,
stderr: folded.stderr().to_owned(),
stderr_truncated: folded.truncated(),
}
}
fn is_sigpipe(outcome: &Outcome) -> bool {
#[cfg(unix)]
return matches!(outcome, Outcome::Signalled(Some(13)));
#[cfg(not(unix))]
let _ = outcome;
#[cfg(not(unix))]
false
}
fn is_clean_exit(outcome: Outcome, ok_codes: &[i32]) -> bool {
match outcome {
Outcome::Exited(code) => ok_codes.contains(&code),
Outcome::Signalled(_) | Outcome::TimedOut | Outcome::InactivityTimedOut => false,
}
}
fn is_checked_failure(outcome: Outcome, ok_codes: &[i32], unchecked: bool) -> bool {
!unchecked && !is_clean_exit(outcome, ok_codes)
}
fn pipefail<T>(stages: Vec<StageOutcome>, last_stdout: T) -> ProcessResult<T> {
let checked_failures: Vec<_> = stages
.iter()
.filter(|s| !s.unchecked && !is_clean_exit(s.outcome, &s.ok_codes))
.collect();
if let Some(stage) = checked_failures
.iter()
.find(|s| !is_sigpipe(&s.outcome) && !s.torn_down)
.or_else(|| checked_failures.first()) .copied()
{
return ProcessResult::new(
stage.program.clone(),
last_stdout,
stage.stderr.clone(),
stage.outcome,
stage.timeout,
)
.with_ok_codes(stage.ok_codes.clone())
.with_truncated(stage.stderr_truncated);
}
let last = stages.last().expect("a pipeline has at least two stages");
let ok_codes = match last.outcome {
Outcome::Exited(code) if last.unchecked && !last.ok_codes.contains(&code) => vec![code],
_ => last.ok_codes.clone(),
};
ProcessResult::new(
last.program.clone(),
last_stdout,
last.stderr.clone(),
last.outcome,
last.timeout,
)
.with_ok_codes(ok_codes)
.with_truncated(last.stderr_truncated)
}
impl std::ops::BitOr<Command> for Command {
type Output = Pipeline;
fn bitor(self, rhs: Command) -> Pipeline {
self.pipe(rhs)
}
}
impl std::ops::BitOr<Command> for Pipeline {
type Output = Pipeline;
fn bitor(self, rhs: Command) -> Pipeline {
self.pipe(rhs)
}
}
fn join_error(err: tokio::task::JoinError) -> crate::Error {
crate::Error::io(std::io::Error::other(format!(
"pipeline stage task failed: {err}"
)))
}
#[derive(Debug)]
struct PartialDrainError<Item> {
completed: Vec<Item>,
error: crate::Error,
}
type InnerStage = (usize, StageOutcome);
type InnerDrain = std::result::Result<Vec<InnerStage>, PartialDrainError<InnerStage>>;
async fn drain_unordered<Item: 'static>(
mut tasks: tokio::task::JoinSet<Result<Item>>,
teardown: &tokio_util::sync::CancellationToken,
) -> std::result::Result<Vec<Item>, PartialDrainError<Item>> {
let mut collected = Vec::with_capacity(tasks.len());
while let Some(joined) = tasks.join_next().await {
match joined {
Ok(Ok(item)) => collected.push(item),
Ok(Err(err)) => {
teardown.cancel();
return Err(PartialDrainError {
completed: collected,
error: err,
});
}
Err(join_err) => {
teardown.cancel();
return Err(PartialDrainError {
completed: collected,
error: join_error(join_err),
});
}
}
}
Ok(collected)
}
#[cfg(test)]
mod tests {
use super::*;
#[allow(dead_code)]
fn session_and_its_probe_future_stay_send(session: &mut PipelineSession) {
fn assert_send<T: Send>(_: &T) {}
fn assert_send_sync<T: Send + Sync>(_: &T) {}
assert_send_sync(session);
assert_send(&session.wait_for_line(|line| line.is_empty(), Duration::ZERO));
}
fn stage(program: &str, outcome: Outcome) -> StageOutcome {
StageOutcome {
program: program.into(),
outcome,
stderr: String::new(),
unchecked: false,
ok_codes: vec![0],
timeout: None,
torn_down: false,
stderr_truncated: false,
}
}
fn clean(program: &str) -> StageOutcome {
stage(program, Outcome::Exited(0))
}
fn unclean(program: &str, outcome: Outcome, stderr: &str) -> StageOutcome {
StageOutcome {
stderr: stderr.into(),
..stage(program, outcome)
}
}
fn unchecked_fail(program: &str, outcome: Outcome) -> StageOutcome {
StageOutcome {
unchecked: true,
..unclean(program, outcome, "forgiven")
}
}
fn last(outcome: Outcome, unchecked: bool) -> StageOutcome {
StageOutcome {
program: "last".into(),
outcome,
stderr: "last-err".into(),
unchecked,
ok_codes: vec![0],
timeout: None,
torn_down: false,
stderr_truncated: false,
}
}
fn pf(mut inner: Vec<StageOutcome>, last: StageOutcome, stdout: &str) -> ProcessResult<String> {
inner.push(last);
pipefail(inner, stdout.to_owned())
}
#[cfg(feature = "pty")]
#[tokio::test]
async fn non_final_pty_stage_is_rejected_before_any_spawn() {
let error = Command::new("never-spawn-first")
.use_pty()
.pipe(Command::new("never-spawn-second"))
.start()
.await
.expect_err("a PTY cannot provide the next stage's stdin pipe");
assert!(
matches!(
error.reason(),
crate::ErrorReason::Unsupported { operation }
if operation.contains("use_pty") && operation.contains("stage 1")
),
"the wiring error must be typed and identify the non-final stage: {error:?}"
);
}
fn expect_last(outcome: Outcome, stdout: &str) -> ProcessResult<String> {
ProcessResult::new(
"last".into(),
stdout.into(),
"last-err".into(),
outcome,
None,
)
}
#[test]
fn all_clean_inner_stages_let_the_last_stage_speak() {
let ok = pf(
vec![clean("a"), clean("b")],
last(Outcome::Exited(0), false),
"final",
);
assert_eq!(ok, expect_last(Outcome::Exited(0), "final"));
let failing_last = pf(vec![clean("a")], last(Outcome::Exited(3), false), "partial");
assert_eq!(failing_last, expect_last(Outcome::Exited(3), "partial"));
}
#[test]
fn failing_inner_stage_wins_but_stdout_stays_the_chains() {
let result = pf(
vec![clean("a"), unclean("b", Outcome::Exited(2), "b broke")],
last(Outcome::Exited(0), false),
"final",
);
assert_eq!(result.program(), "b", "diagnostics from the failing stage");
assert_eq!(result.code(), Some(2));
assert_eq!(result.stderr(), "b broke");
assert_eq!(
result.stdout(),
"final",
"stdout is what the chain produced — the last stage's"
);
assert!(!result.timed_out());
match result.ensure_success().map_err(|e| e.into_reason()) {
Err(crate::ErrorReason::Exit {
program,
code,
stdout,
stderr,
..
}) => {
assert_eq!(program, "b", "diagnostics from the failing stage");
assert_eq!(code, 2);
assert_eq!(stdout, "final");
assert_eq!(stderr, "b broke");
}
other => panic!("expected ErrorReason::Exit, got {other:?}"),
}
}
#[test]
fn rejected_zero_stage_stays_a_failure_after_attribution() {
let culprit = StageOutcome {
ok_codes: vec![1],
..unclean("check", Outcome::Exited(0), "rejected zero")
};
let result = pf(vec![culprit], last(Outcome::Exited(0), false), "final");
assert_eq!(result.program(), "check", "the rejected-zero stage wins");
assert_eq!(result.code(), Some(0));
assert_eq!(result.stderr(), "rejected zero");
assert!(
!result.is_success(),
"a rejected-zero failure must not report success just because it exited 0"
);
assert!(
result.ensure_success().is_err(),
"the chain must surface the failure, not swallow it as Ok"
);
}
#[test]
fn first_of_several_failures_is_attributed() {
let result = pf(
vec![
unclean("a", Outcome::Exited(1), "first"),
unclean("b", Outcome::Exited(2), "second"),
],
last(Outcome::Exited(0), false),
"out",
);
assert_eq!(result.program(), "a", "pipefail blames the FIRST failure");
assert_eq!(result.code(), Some(1));
assert_eq!(result.stderr(), "first");
match result.ensure_success().map_err(|e| e.into_reason()) {
Err(crate::ErrorReason::Exit { program, .. }) => {
assert_eq!(program, "a", "...and so does the error surface");
}
other => panic!("expected ErrorReason::Exit, got {other:?}"),
}
}
#[test]
fn attributed_inner_stage_truncation_survives_the_fold_not_just_the_last_stage() {
let culprit = StageOutcome {
stderr_truncated: true,
..unclean("b", Outcome::Exited(2), "b broke (clipped)")
};
let result = pf(
vec![clean("a"), culprit],
last(Outcome::Exited(0), false),
"final",
);
assert_eq!(result.program(), "b", "the inner failing stage is blamed");
assert!(
result.truncated(),
"the attributed inner stage's dropped stderr must be visible: {result:?}"
);
let clean_culprit = unclean("b", Outcome::Exited(2), "b broke");
let untruncated = pf(
vec![clean("a"), clean_culprit],
last(Outcome::Exited(0), false),
"final",
);
assert!(
!untruncated.truncated(),
"an inner failure with no dropped stderr must not report truncated: {untruncated:?}"
);
}
#[test]
fn chain_timeout_keeps_pipefail_stderr_and_truncation() {
let mut culprit = stage("culprit", Outcome::Signalled(None));
culprit.stderr = "retained diagnostic".into();
culprit.stderr_truncated = true;
let finished = timeout_finished(
Ok(vec![(0, culprit)]),
Ok(Finished {
outcome: Outcome::Signalled(None),
stderr: String::new(),
stderr_truncated: false,
}),
"last",
&[0],
false,
None,
false,
);
assert_eq!(finished.outcome, Outcome::TimedOut);
assert_eq!(finished.stderr, "retained diagnostic");
assert!(finished.stderr_truncated);
}
#[test]
fn chain_timeout_keeps_completed_stderr_when_another_stage_returns_raw_error() {
let mut culprit = unclean("completed", Outcome::Exited(7), "retained diagnostic");
culprit.stderr_truncated = true;
let finished = timeout_finished(
Err(PartialDrainError {
completed: vec![(0, culprit)],
error: crate::Error::io(std::io::Error::other("raw stage error")),
}),
Ok(Finished {
outcome: Outcome::Signalled(None),
stderr: String::new(),
stderr_truncated: false,
}),
"last",
&[0],
false,
None,
false,
);
assert_eq!(finished.outcome, Outcome::TimedOut);
assert_eq!(finished.stderr, "retained diagnostic");
assert!(finished.stderr_truncated);
}
#[test]
fn all_unchecked_failures_report_success() {
let result = pf(
vec![unchecked_fail("producer", Outcome::Signalled(None))],
last(Outcome::Exited(0), false),
"first line",
);
assert!(result.is_success(), "got {result:?}");
assert_eq!(result.stdout(), "first line");
assert_eq!(result.program(), "last", "the clean last stage speaks");
}
#[test]
fn checked_failure_trumps_unchecked_regardless_of_order() {
let result = pf(
vec![
unchecked_fail("a", Outcome::Exited(141)),
unclean("b", Outcome::Exited(2), "real"),
],
last(Outcome::Exited(0), false),
"out",
);
assert_eq!(result.program(), "b", "unchecked never shields a failure");
assert_eq!(result.code(), Some(2));
let result = pf(
vec![
unclean("a", Outcome::Exited(1), "real"),
unchecked_fail("b", Outcome::Exited(2)),
],
last(Outcome::Exited(0), false),
"out",
);
assert_eq!(result.program(), "a");
assert_eq!(result.code(), Some(1));
}
#[test]
fn attribution_skips_unchecked_to_the_first_checked_failure() {
let result = pf(
vec![
clean("a"),
unchecked_fail("b", Outcome::Exited(1)),
unclean("c", Outcome::Exited(3), "c broke"),
unclean("d", Outcome::Exited(4), "d broke"),
],
last(Outcome::Exited(0), false),
"out",
);
assert_eq!(result.program(), "c", "first CHECKED failure is blamed");
assert_eq!(result.code(), Some(3));
assert_eq!(result.stderr(), "c broke");
}
#[test]
fn unchecked_last_stage_failure_is_forgiven() {
let result = pf(
vec![clean("a")],
last(Outcome::Exited(141), true),
"partial",
);
assert!(result.is_success(), "got {result:?}");
assert_eq!(result.code(), Some(141), "real exit code preserved");
assert_eq!(result.stdout(), "partial", "output is preserved");
assert_eq!(result.stderr(), "last-err", "stderr kept for the curious");
assert!(result.ensure_success().is_ok());
}
#[test]
fn last_stage_ok_codes_are_honoured() {
let mut last_grep = last(Outcome::Exited(1), false);
last_grep.program = "grep".into();
last_grep.ok_codes = vec![0, 1];
let result = pf(vec![clean("a")], last_grep, "matched");
assert!(
result.is_success(),
"exit 1 in the last stage's ok_codes: {result:?}"
);
assert_eq!(result.code(), Some(1), "real code preserved");
assert_eq!(result.program(), "grep");
}
#[test]
fn inner_stage_ok_codes_are_honoured_in_pipefail_cleanliness() {
let mut with_ok = stage("grep", Outcome::Exited(1));
with_ok.ok_codes = vec![0, 1];
let result = pf(vec![with_ok], last(Outcome::Exited(0), false), "out");
assert!(
result.is_success(),
"exit 1 in ok_codes should be clean: {result:?}"
);
assert_eq!(result.program(), "last", "clean inner → last stage speaks");
}
#[test]
fn timed_out_stage_reports_its_own_deadline_not_the_chains() {
let mut timed = unclean("slow", Outcome::TimedOut, "");
timed.timeout = Some(Duration::from_millis(500));
let result = pf(vec![timed], last(Outcome::Exited(0), false), "out");
assert_eq!(result.program(), "slow");
assert!(result.timed_out());
match result.ensure_success().map_err(|e| e.into_reason()) {
Err(crate::ErrorReason::Timeout {
program, timeout, ..
}) => {
assert_eq!(program, "slow");
assert_eq!(
timeout,
Duration::from_millis(500),
"the stage's own deadline, not the chain's 0ns"
);
}
other => panic!("expected ErrorReason::Timeout, got {other:?}"),
}
}
#[cfg(unix)]
#[test]
fn sigpipe_victim_not_blamed_when_downstream_non_sigpipe_failure_exists() {
let sigpipe_victim = unclean("producer", Outcome::Signalled(Some(13)), "pipe broken");
let real_failure = unclean("consumer", Outcome::Exited(2), "consumer broke");
let result = pf(
vec![sigpipe_victim, real_failure],
last(Outcome::Exited(0), false),
"out",
);
assert_eq!(
result.program(),
"consumer",
"downstream non-SIGPIPE culprit, not upstream SIGPIPE victim"
);
assert_eq!(result.code(), Some(2));
}
#[test]
fn torn_down_victim_does_not_steal_blame_from_the_real_failure() {
let torn_upstream = StageOutcome {
torn_down: true,
..unclean(
"upstream",
Outcome::Signalled(Some(9)),
"killed by teardown",
)
};
let culprit = unclean("downstream", Outcome::Exited(3), "the real failure");
let torn_last = StageOutcome {
torn_down: true,
..last(Outcome::Signalled(Some(9)), false)
};
let result = pf(vec![torn_upstream, culprit], torn_last, "");
assert_eq!(
result.program(),
"downstream",
"the failure that triggered teardown wins, not a torn-down victim"
);
assert_eq!(result.code(), Some(3));
match result.ensure_success().map_err(|e| e.into_reason()) {
Err(crate::ErrorReason::Exit { program, code, .. }) => {
assert_eq!(program, "downstream");
assert_eq!(code, 3);
}
other => panic!("expected ErrorReason::Exit, got {other:?}"),
}
}
#[test]
fn all_torn_down_failures_fall_back_to_the_leftmost() {
let first = StageOutcome {
torn_down: true,
..unclean("a", Outcome::Signalled(Some(9)), "first killed")
};
let second = StageOutcome {
torn_down: true,
..unclean("b", Outcome::Signalled(Some(9)), "second killed")
};
let result = pf(vec![first, second], last(Outcome::Exited(0), false), "");
assert_eq!(
result.program(),
"a",
"leftmost victim when all are torn down"
);
assert!(!result.is_success());
}
#[test]
fn a_stage_disposition_is_the_first_observers_and_stays_it() {
let culprit = ExitDisposition::unobserved();
assert!(
!culprit.latch(false),
"no teardown in flight at the exit: a culprit"
);
assert!(
!culprit.latch(true),
"a later observer reads the latched verdict, it never overwrites it"
);
let victim = ExitDisposition::unobserved();
assert!(victim.latch(true), "a teardown was already in flight");
assert!(victim.latch(false), "still a victim on a later read");
}
#[test]
fn a_slow_draining_culprit_is_not_demoted_by_a_later_failure() {
let slow_draining_culprit = unclean("upstream", Outcome::Exited(3), "the real failure");
let later_failure = last(Outcome::Exited(5), false);
let result = pf(vec![slow_draining_culprit], later_failure, "");
assert_eq!(
result.program(),
"upstream",
"the stage that failed first is blamed, however long its drain took"
);
assert_eq!(result.code(), Some(3));
assert_eq!(result.stderr(), "the real failure");
}
#[cfg(unix)]
#[test]
fn a_slow_draining_last_stage_keeps_the_blame_from_a_sigpipe_producer() {
let result = pf(
vec![unclean(
"producer",
Outcome::Signalled(Some(13)),
"producer noise",
)],
last(Outcome::Exited(3), false),
"",
);
assert_eq!(
result.code(),
Some(3),
"the last stage failed on its own and is not a teardown victim"
);
assert_eq!(result.stderr(), "last-err", "with its own diagnostics");
let demoted = StageOutcome {
torn_down: true,
..last(Outcome::Exited(3), false)
};
let regressed = pf(
vec![unclean(
"producer",
Outcome::Signalled(Some(13)),
"producer noise",
)],
demoted,
"",
);
assert!(
matches!(regressed.outcome(), Outcome::Signalled(Some(13))),
"the demotion is user-visible, not cosmetic: {:?}",
regressed.outcome()
);
assert_eq!(regressed.stderr(), "producer noise");
}
#[test]
fn checked_last_stage_failure_still_speaks_verbatim() {
let result = pf(vec![clean("a")], last(Outcome::Exited(3), false), "partial");
assert_eq!(result, expect_last(Outcome::Exited(3), "partial"));
}
#[test]
fn unchecked_never_forgives_a_timeout() {
let result = pf(vec![clean("a")], last(Outcome::TimedOut, true), "");
assert!(result.timed_out());
assert!(!result.is_success());
}
#[test]
fn unchecked_never_forgives_a_signal_kill() {
let result = pf(
vec![clean("a")],
last(Outcome::Signalled(Some(9)), true),
"",
);
assert!(matches!(result.outcome(), Outcome::Signalled(Some(9))));
assert!(!result.is_success());
}
#[test]
fn bitor_chains_like_pipe() {
let chain = Command::new("a") | Command::new("b") | Command::new("c");
assert_eq!(chain.stages.len(), 3, "a | b | c is one three-stage chain");
assert_eq!(chain.pipeline_name(), "a | b | c");
assert!(chain.timeout.is_none());
}
#[test]
fn signal_killed_inner_stage_counts_as_unclean() {
let result = pf(
vec![unclean("a", Outcome::Signalled(None), "killed")],
last(Outcome::Exited(0), false),
"out",
);
assert_eq!(result.program(), "a");
assert_eq!(result.code(), None);
assert_eq!(result.stderr(), "killed");
assert!(!result.timed_out(), "a stage kill is not a chain timeout");
match result.ensure_success().map_err(|e| e.into_reason()) {
Err(crate::ErrorReason::Signalled {
program, signal, ..
}) => {
assert_eq!(program, "a");
assert_eq!(signal, None);
}
other => panic!("expected ErrorReason::Signalled, got {other:?}"),
}
}
async fn quiet_until_teardown(
teardown: tokio_util::sync::CancellationToken,
item: &'static str,
) -> Result<&'static str> {
teardown.cancelled().await;
Ok(item)
}
#[tokio::test]
async fn drain_unordered_wakes_a_quiet_task_when_a_later_task_returns_a_raw_error() {
let teardown = tokio_util::sync::CancellationToken::new();
let mut tasks: tokio::task::JoinSet<Result<&'static str>> = tokio::task::JoinSet::new();
tasks.spawn(quiet_until_teardown(teardown.clone(), "quiet-upstream"));
tasks.spawn(async { Err(crate::Error::io(std::io::Error::other("downstream boom"))) });
let drained =
tokio::time::timeout(Duration::from_secs(5), drain_unordered(tasks, &teardown))
.await
.expect(
"drain_unordered must not hang on the still-pending quiet task \
once the sibling's raw error has fired teardown",
);
match drained.map_err(|e| e.error.into_reason()) {
Err(crate::ErrorReason::Io(err)) => {
assert_eq!(err.to_string(), "downstream boom");
}
other => panic!("expected the downstream stage's own Io error, got {other:?}"),
}
assert!(
teardown.is_cancelled(),
"a raw Err must fire teardown so a quiet sibling is unblocked"
);
}
#[tokio::test]
async fn drain_unordered_fires_teardown_on_a_task_panic_so_a_quiet_sibling_still_resolves() {
let teardown = tokio_util::sync::CancellationToken::new();
let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(2));
let mut tasks: tokio::task::JoinSet<Result<&'static str>> = tokio::task::JoinSet::new();
tasks.spawn({
let teardown = teardown.clone();
let barrier = barrier.clone();
async move {
barrier.wait().await;
quiet_until_teardown(teardown, "quiet-upstream").await
}
});
tasks.spawn(async move {
barrier.wait().await;
panic!("downstream stage task panicked");
});
let drained =
tokio::time::timeout(Duration::from_secs(5), drain_unordered(tasks, &teardown))
.await
.expect(
"drain_unordered must not hang on the still-pending quiet task \
once the sibling's panic has fired teardown",
);
match drained.map_err(|e| e.error.into_reason()) {
Err(crate::ErrorReason::Io(err)) => {
assert!(
err.to_string().contains("pipeline stage task failed"),
"expected the wrapped JoinError, got {err}"
);
}
other => panic!("expected a wrapped JoinError, got {other:?}"),
}
assert!(
teardown.is_cancelled(),
"a task panic must fire teardown so a quiet sibling is unblocked"
);
}
#[tokio::test]
async fn drain_unordered_returns_every_item_when_the_whole_set_finishes_clean() {
let teardown = tokio_util::sync::CancellationToken::new();
let mut tasks: tokio::task::JoinSet<Result<u32>> = tokio::task::JoinSet::new();
for item in [1u32, 2, 3] {
tasks.spawn(async move {
tokio::task::yield_now().await;
Ok(item)
});
}
let mut drained = drain_unordered(tasks, &teardown)
.await
.expect("every task succeeded");
drained.sort_unstable();
assert_eq!(
drained,
vec![1, 2, 3],
"every task's payload survives the drain, completion order aside"
);
assert!(
!teardown.is_cancelled(),
"an all-clean drain must never fire teardown"
);
}
}