aria_primitives/process/
run.rs

1use std::fmt;
2use std::io::{self, Read, Write};
3use std::process::{Command, Output, Stdio};
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::thread;
6use std::time::{Duration, Instant};
7
8use super::{ChildProcess, ChildProcessError, WaitOutcome};
9use crate::cancel::CancelToken;
10
11#[derive(Debug, Clone, Copy)]
12pub struct Bounds {
13    pub timeout: Duration,
14    pub stdout: Capture,
15    pub stderr: Capture,
16}
17
18/// What to keep of a pipe, in bytes.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum Capture {
21    /// Past the limit, the child is stopped and the run fails.
22    FailAbove(usize),
23    /// Keeps the first bytes and discards the rest as it arrives.
24    KeepHead(usize),
25    KeepTail(usize),
26}
27
28impl Capture {
29    fn limit(self) -> usize {
30        match self {
31            Self::FailAbove(limit) | Self::KeepHead(limit) | Self::KeepTail(limit) => limit,
32        }
33    }
34}
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum Pipe {
38    Stdout,
39    Stderr,
40}
41
42impl fmt::Display for Pipe {
43    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
44        f.write_str(match self {
45            Self::Stdout => "stdout",
46            Self::Stderr => "stderr",
47        })
48    }
49}
50
51#[derive(Debug, thiserror::Error)]
52pub enum RunError {
53    #[error("failed to start the child process")]
54    Spawn(#[source] io::Error),
55
56    #[error("failed to write the child process's input")]
57    Write(#[source] io::Error),
58
59    #[error("failed to read the child process's output")]
60    Read(#[source] io::Error),
61
62    #[error("failed to wait for the child process")]
63    Wait(#[source] ChildProcessError),
64
65    #[error("the child process timed out after {0:?}")]
66    TimedOut(Duration),
67
68    #[error("the child process was cancelled")]
69    Cancelled,
70
71    #[error("the child process's {pipe} exceeded the output limit of {limit} bytes")]
72    OutputTooLarge { pipe: Pipe, limit: usize },
73}
74
75/// Runs `command` to completion and returns what it wrote, as `bounds`
76/// allows.
77///
78/// Both output pipes are drained while the child runs, so it never blocks on
79/// a full pipe. An empty `input` gives the child `/dev/null`. A child that
80/// exits before reading all of `input` is not a failure; its status says why.
81///
82/// An overflow outranks every other outcome. The deadline bounds the child's
83/// process group; a descendant that left the group and still holds a pipe
84/// delays the return until it exits.
85pub fn run(
86    command: &mut Command,
87    input: &[u8],
88    bounds: Bounds,
89    cancel: &CancelToken,
90) -> Result<Output, RunError> {
91    run_with_listener(command, input, bounds, cancel, |_| {})
92}
93
94/// [`run`], with stdout handed to `on_stdout` as it is read. Under
95/// [`Capture::FailAbove`] and [`Capture::KeepHead`] that is only the bytes within
96/// the limit; under [`Capture::KeepTail`] it is every byte. It gets nothing
97/// past a `FailAbove` limit, nor after a deadline or cancel stops the child.
98pub fn run_with_listener(
99    command: &mut Command,
100    input: &[u8],
101    bounds: Bounds,
102    cancel: &CancelToken,
103    mut on_stdout: impl FnMut(&[u8]) + Send,
104) -> Result<Output, RunError> {
105    let deadline = Instant::now() + bounds.timeout;
106    command
107        .stdin(if input.is_empty() {
108            Stdio::null()
109        } else {
110            Stdio::piped()
111        })
112        .stdout(Stdio::piped())
113        .stderr(Stdio::piped());
114    let mut child = ChildProcess::spawn(command).map_err(RunError::Spawn)?;
115    let stdin = child.take_stdin();
116    let stdout = child.take_stdout().ok_or_else(|| not_piped("stdout"))?;
117    let stderr = child.take_stderr().ok_or_else(|| not_piped("stderr"))?;
118
119    // Set when the run can no longer succeed. The wait then ends and the
120    // child is stopped, instead of running to its deadline.
121    let stop = AtomicBool::new(false);
122    let (waited, stdout, stderr, written) = thread::scope(|scope| {
123        // Owned in here so its drop stops the group before the scope joins
124        // the readers, even while unwinding. A reader can block until the
125        // group is gone.
126        let mut child = child;
127        let stdout = scope.spawn(|| {
128            let _unwinding = StopOnUnwind(&stop);
129            drain(stdout, bounds.stdout, &stop, |bytes| {
130                if !stop.load(Ordering::Relaxed) {
131                    on_stdout(bytes);
132                }
133            })
134        });
135        let stderr = scope.spawn(|| drain(stderr, bounds.stderr, &stop, |_| {}));
136        let writer = stdin.map(|pipe| scope.spawn(|| write_input(pipe, input, &stop)));
137        let waited = child.wait_until(deadline, || {
138            stop.load(Ordering::Relaxed) || cancel.is_cancelled()
139        });
140        if !matches!(waited, Ok(WaitOutcome::Exited(_))) {
141            stop.store(true, Ordering::Relaxed);
142        }
143        drop(child);
144        let written = writer.map_or(Ok(()), joined);
145        (waited, joined(stdout), joined(stderr), written)
146    });
147    outcome(&bounds, waited, stdout, stderr, written)
148}
149
150/// A pipe or input failure raises `stop`, so the wait can end as
151/// `Cancelled`. Checking those failures first reports the real cause.
152fn outcome(
153    bounds: &Bounds,
154    waited: Result<WaitOutcome, ChildProcessError>,
155    stdout: io::Result<Drained>,
156    stderr: io::Result<Drained>,
157    written: io::Result<()>,
158) -> Result<Output, RunError> {
159    for (drained, pipe, capture) in [
160        (&stdout, Pipe::Stdout, bounds.stdout),
161        (&stderr, Pipe::Stderr, bounds.stderr),
162    ] {
163        if drained.as_ref().is_ok_and(|drained| drained.overflowed) {
164            return Err(RunError::OutputTooLarge {
165                pipe,
166                limit: capture.limit(),
167            });
168        }
169    }
170    let stdout = stdout.map_err(RunError::Read)?.kept;
171    let stderr = stderr.map_err(RunError::Read)?.kept;
172    written.map_err(RunError::Write)?;
173    match waited.map_err(RunError::Wait)? {
174        WaitOutcome::Exited(status) => Ok(Output {
175            status,
176            stdout,
177            stderr,
178        }),
179        WaitOutcome::TimedOut => Err(RunError::TimedOut(bounds.timeout)),
180        WaitOutcome::Cancelled => Err(RunError::Cancelled),
181    }
182}
183
184/// Raises `stop` if the thread unwinds, so a panicking listener ends the
185/// wait instead of leaving it to the deadline.
186struct StopOnUnwind<'a>(&'a AtomicBool);
187
188impl Drop for StopOnUnwind<'_> {
189    fn drop(&mut self) {
190        if thread::panicking() {
191            self.0.store(true, Ordering::Relaxed);
192        }
193    }
194}
195
196fn not_piped(stream: &str) -> RunError {
197    RunError::Read(io::Error::other(format!(
198        "the child's {stream} was not piped"
199    )))
200}
201
202fn joined<T>(handle: thread::ScopedJoinHandle<'_, T>) -> T {
203    handle
204        .join()
205        .unwrap_or_else(|panic| std::panic::resume_unwind(panic))
206}
207
208fn write_input(mut pipe: impl Write, input: &[u8], stop: &AtomicBool) -> io::Result<()> {
209    match pipe.write_all(input) {
210        Err(error) if error.kind() == io::ErrorKind::BrokenPipe => Ok(()),
211        Err(error) => {
212            stop.store(true, Ordering::Relaxed);
213            Err(error)
214        }
215        Ok(()) => Ok(()),
216    }
217}
218
219struct Drained {
220    kept: Vec<u8>,
221    overflowed: bool,
222}
223
224/// Bytes past the limit are read and dropped, so the child never blocks on
225/// the pipe or dies of a closed one.
226fn drain(
227    mut pipe: impl Read,
228    capture: Capture,
229    stop: &AtomicBool,
230    mut forward: impl FnMut(&[u8]),
231) -> io::Result<Drained> {
232    let mut kept = Vec::new();
233    let mut overflowed = false;
234    let mut chunk = [0; 8 * 1024];
235    loop {
236        let read = match pipe.read(&mut chunk) {
237            Ok(0) => return Ok(Drained { kept, overflowed }),
238            Ok(read) => read,
239            Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
240            Err(_) if overflowed => return Ok(Drained { kept, overflowed }),
241            Err(error) => {
242                stop.store(true, Ordering::Relaxed);
243                return Err(error);
244            }
245        };
246        let chunk = &chunk[..read];
247        match capture {
248            Capture::FailAbove(limit) | Capture::KeepHead(limit) => {
249                let room = limit - kept.len();
250                let taken = &chunk[..chunk.len().min(room)];
251                kept.extend_from_slice(taken);
252                if !taken.is_empty() {
253                    forward(taken);
254                }
255                if chunk.len() > room && matches!(capture, Capture::FailAbove(_)) {
256                    overflowed = true;
257                    stop.store(true, Ordering::Relaxed);
258                }
259            }
260            Capture::KeepTail(limit) => {
261                kept.extend_from_slice(chunk);
262                forward(chunk);
263                if kept.len() > limit {
264                    kept.drain(..kept.len() - limit);
265                }
266            }
267        }
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    const PATIENCE: Duration = Duration::from_secs(5);
276    const LIMIT: usize = 1024;
277
278    fn shell(script: &str) -> Command {
279        let mut command = Command::new("/bin/sh");
280        command.arg("-c").arg(script);
281        command
282    }
283
284    fn bounds(stdout: Capture, stderr: Capture) -> Bounds {
285        Bounds {
286            timeout: PATIENCE,
287            stdout,
288            stderr,
289        }
290    }
291
292    fn failing() -> Bounds {
293        bounds(Capture::FailAbove(LIMIT), Capture::FailAbove(LIMIT))
294    }
295
296    fn run_script(script: &str, input: &[u8], bounds: Bounds) -> Result<Output, RunError> {
297        run(&mut shell(script), input, bounds, &CancelToken::never())
298    }
299
300    #[test]
301    fn stdin_reaches_the_child_and_stdout_streams_back() {
302        let mut streamed = Vec::new();
303        let output = run_with_listener(
304            &mut shell("cat; printf last"),
305            b"one\ntwo\n",
306            failing(),
307            &CancelToken::never(),
308            |bytes| streamed.extend_from_slice(bytes),
309        )
310        .expect("the run succeeds");
311
312        assert!(output.status.success());
313        assert_eq!(output.stdout, b"one\ntwo\nlast");
314        assert_eq!(streamed, output.stdout);
315    }
316
317    #[test]
318    fn an_empty_input_ends_at_once() {
319        let output = run_script("cat", b"", failing()).expect("the run succeeds");
320
321        assert!(output.stdout.is_empty());
322    }
323
324    #[test]
325    fn a_child_past_a_fail_limit_is_stopped() {
326        for (script, pipe) in [("yes", Pipe::Stdout), ("yes >&2", Pipe::Stderr)] {
327            let started = Instant::now();
328            let error = run_script(script, b"", failing()).expect_err("the output overflows");
329
330            assert!(
331                matches!(error, RunError::OutputTooLarge { pipe: p, limit: LIMIT } if p == pipe),
332                "{error:?}"
333            );
334            assert!(started.elapsed() < PATIENCE);
335        }
336    }
337
338    #[test]
339    fn one_byte_past_the_limit_fails_even_after_a_clean_exit() {
340        let script = format!("yes | head -c {}; exit 0", LIMIT + 1);
341        let error = run_script(&script, b"", failing()).expect_err("the output overflows");
342
343        assert!(
344            matches!(error, RunError::OutputTooLarge { .. }),
345            "{error:?}"
346        );
347    }
348
349    #[test]
350    fn output_exactly_at_the_limit_is_kept() {
351        let script = format!("yes | head -c {LIMIT}");
352        let output = run_script(&script, b"", failing()).expect("the output fits");
353
354        assert_eq!(output.stdout.len(), LIMIT);
355    }
356
357    #[test]
358    fn keep_head_truncates_and_lets_the_child_finish() {
359        // The shell writes the lines itself, so a closed pipe would kill it
360        // before it reached `exit 7`.
361        let script = "i=0; while [ $i -lt 2000 ]; do echo 0123456789012345678901234567890123456789; \
362                      i=$((i+1)); done; exit 7";
363        let output = run_script(
364            script,
365            b"",
366            bounds(Capture::KeepHead(LIMIT), Capture::FailAbove(LIMIT)),
367        )
368        .expect("the run succeeds");
369
370        assert_eq!(output.status.code(), Some(7));
371        assert_eq!(output.stdout.len(), LIMIT);
372    }
373
374    #[test]
375    fn keep_tail_keeps_the_last_bytes() {
376        let output = run_script(
377            "yes | head -c 100000 >&2; printf END >&2",
378            b"",
379            bounds(Capture::FailAbove(LIMIT), Capture::KeepTail(3)),
380        )
381        .expect("the run succeeds");
382
383        assert_eq!(output.stderr, b"END");
384    }
385
386    #[test]
387    fn on_stdout_gets_what_the_capture_lets_through() {
388        let script = "yes | head -c 100000";
389        for (capture, expected) in [
390            (Capture::KeepHead(LIMIT), LIMIT),
391            (Capture::KeepTail(3), 100_000),
392        ] {
393            let mut streamed = 0;
394            run_with_listener(
395                &mut shell(script),
396                b"",
397                bounds(capture, Capture::FailAbove(LIMIT)),
398                &CancelToken::never(),
399                |bytes| streamed += bytes.len(),
400            )
401            .expect("the run succeeds");
402
403            assert_eq!(streamed, expected, "{capture:?}");
404        }
405    }
406
407    #[test]
408    fn a_child_that_ignores_its_input_still_succeeds() {
409        let output = run_script("exit 0", &[b'x'; 1_000_000], failing()).expect("the run succeeds");
410
411        assert!(output.status.success());
412    }
413
414    #[test]
415    fn a_background_descendant_does_not_hold_the_run() {
416        let started = Instant::now();
417        let output = run_script("sleep 30 & exit 0", b"", failing()).expect("the run succeeds");
418
419        assert!(output.status.success());
420        assert!(started.elapsed() < PATIENCE);
421    }
422
423    #[test]
424    fn the_deadline_and_a_cancel_end_the_run_distinctly() {
425        let timeout = Duration::from_millis(100);
426        let started = Instant::now();
427        let error = run_script(
428            "sleep 30",
429            b"",
430            Bounds {
431                timeout,
432                ..failing()
433            },
434        )
435        .expect_err("the deadline passes");
436        assert!(
437            matches!(error, RunError::TimedOut(t) if t == timeout),
438            "{error:?}"
439        );
440        assert!(started.elapsed() < PATIENCE);
441
442        let cancel = CancelToken::new();
443        cancel.cancel();
444        let started = Instant::now();
445        let error =
446            run(&mut shell("sleep 30"), b"", failing(), &cancel).expect_err("the run is cancelled");
447        assert!(matches!(error, RunError::Cancelled), "{error:?}");
448        assert!(started.elapsed() < PATIENCE);
449    }
450
451    #[test]
452    fn a_panicking_cancel_check_still_stops_the_child() {
453        let started = Instant::now();
454        let unwound = std::panic::catch_unwind(|| {
455            run(
456                &mut shell("sleep 30"),
457                b"",
458                failing(),
459                &CancelToken::from_fn(|| panic!("the check fails")),
460            )
461        });
462
463        assert!(unwound.is_err());
464        assert!(started.elapsed() < PATIENCE);
465    }
466
467    #[test]
468    fn a_panicking_listener_still_stops_the_child() {
469        let started = Instant::now();
470        let unwound = std::panic::catch_unwind(|| {
471            run_with_listener(
472                &mut shell("echo hi; sleep 30"),
473                b"",
474                failing(),
475                &CancelToken::never(),
476                |_| panic!("the listener fails"),
477            )
478        });
479
480        assert!(unwound.is_err());
481        assert!(started.elapsed() < PATIENCE);
482    }
483}
aria_primitives/process/run.rs

aria_primitives/process/
run.rs
