command_stream/zx/
process.rs

1//! [`ProcessPromise`] (a prepared zx command) and [`RunningProcess`] (a
2//! spawned one), plus the runner that drives the child processes.
3//!
4//! Like zx, a failed command is reported by handing back its full
5//! [`ProcessOutput`] as the error value; keeping it unboxed lets callers match
6//! on `Err(out)` and read `out.exit_code` directly.
7#![allow(clippy::result_large_err)]
8
9use std::future::{Future, IntoFuture};
10use std::path::Path;
11use std::pin::Pin;
12use std::process::Stdio;
13use std::sync::atomic::{AtomicBool, Ordering};
14use std::sync::{Arc, Mutex};
15use std::time::{Duration, Instant};
16
17use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
18use tokio::sync::{mpsc, watch};
19
20use super::error::{ZxError, DOCS_URL};
21use super::kill::{kill_tree, signal_name};
22use super::output::{ErrorInfo, ProcessOutput};
23use super::shell::Options;
24
25/// Which stream of a source command feeds a pipe.
26#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
27pub enum PipeFrom {
28    /// Standard output (the default).
29    #[default]
30    Stdout,
31    /// Standard error.
32    Stderr,
33    /// Both, interleaved.
34    Stdall,
35}
36
37type BoxedRun = Pin<Box<dyn Future<Output = ProcessOutput> + Send>>;
38
39/// Result of awaiting a command: `Ok` on success (or with `nothrow`), `Err`
40/// with the same [`ProcessOutput`] otherwise.
41pub type ZxResult = Result<ProcessOutput, ProcessOutput>;
42
43/// A prepared command (zx `ProcessPromise`). Nothing runs until it is
44/// awaited, [`run`](Self::run), [`run_sync`](Self::run_sync) or
45/// [`spawn`](Self::spawn) is called.
46#[derive(Debug, Clone)]
47pub struct ProcessPromise {
48    opts: Options,
49    cmd: String,
50    build_error: Option<ZxError>,
51    source: Option<Box<(ProcessPromise, PipeFrom)>>,
52}
53
54impl ProcessPromise {
55    /// A command with the given options.
56    pub fn new(opts: Options, cmd: impl Into<String>) -> Self {
57        Self {
58            opts,
59            cmd: cmd.into(),
60            build_error: None,
61            source: None,
62        }
63    }
64
65    /// A command that fails immediately with `err` when run.
66    pub fn failed(opts: Options, err: ZxError) -> Self {
67        Self {
68            opts,
69            cmd: String::new(),
70            build_error: Some(err),
71            source: None,
72        }
73    }
74
75    /// The command without prefix and postfix.
76    pub fn cmd(&self) -> &str {
77        &self.cmd
78    }
79
80    /// The command as passed to the shell: `prefix + cmd + postfix`.
81    pub fn full_cmd(&self) -> String {
82        format!("{}{}{}", self.opts.prefix, self.cmd, self.opts.postfix)
83    }
84
85    /// The options of this command.
86    pub fn options(&self) -> &Options {
87        &self.opts
88    }
89
90    /// Resolve with the output even when the command fails.
91    pub fn nothrow(mut self) -> Self {
92        self.opts.nothrow = true;
93        self
94    }
95
96    /// Suppress logging for this command.
97    pub fn quiet(mut self) -> Self {
98        self.opts.quiet = true;
99        self
100    }
101
102    /// Log the command and its output to stderr.
103    pub fn verbose(mut self) -> Self {
104        self.opts.verbose = true;
105        self
106    }
107
108    /// Kill the command after `timeout` (with the timeout signal).
109    pub fn timeout(mut self, timeout: Duration) -> Self {
110        self.opts.timeout = Some(timeout);
111        self
112    }
113
114    /// Kill the command after `timeout` with `signal`.
115    pub fn timeout_with(mut self, timeout: Duration, signal: impl Into<String>) -> Self {
116        self.opts.timeout = Some(timeout);
117        self.opts.timeout_signal = signal.into();
118        self
119    }
120
121    /// Default signal for [`RunningProcess::kill`].
122    pub fn kill_signal(mut self, signal: impl Into<String>) -> Self {
123        self.opts.kill_signal = signal.into();
124        self
125    }
126
127    /// Write `data` to the command's stdin.
128    pub fn input(mut self, data: impl Into<Vec<u8>>) -> Self {
129        self.opts.input = Some(data.into());
130        self
131    }
132
133    /// Feed this command's stdout into `dest`'s stdin; returns `dest`.
134    ///
135    /// When this command fails (and is not `nothrow`) the destination is
136    /// stopped and inherits its exit code, so the pipeline fails as a whole.
137    pub fn pipe(self, dest: ProcessPromise) -> ProcessPromise {
138        self.pipe_from(PipeFrom::Stdout, dest)
139    }
140
141    /// Feed this command's stderr into `dest`.
142    pub fn pipe_stderr(self, dest: ProcessPromise) -> ProcessPromise {
143        self.pipe_from(PipeFrom::Stderr, dest)
144    }
145
146    /// Feed this command's interleaved stdout and stderr into `dest`.
147    pub fn pipe_stdall(self, dest: ProcessPromise) -> ProcessPromise {
148        self.pipe_from(PipeFrom::Stdall, dest)
149    }
150
151    /// Feed the selected stream into `dest`. If `dest` already has a source,
152    /// this command is attached at the head of its pipeline.
153    pub fn pipe_from(self, from: PipeFrom, mut dest: ProcessPromise) -> ProcessPromise {
154        match dest.source.take() {
155            None => dest.source = Some(Box::new((self, from))),
156            Some(link) => {
157                let (head, head_from) = *link;
158                dest.source = Some(Box::new((self.pipe_from(from, head), head_from)));
159            }
160        }
161        dest
162    }
163
164    /// Run the command writing its stdout into the file at `path`.
165    pub async fn pipe_to_file(self, path: impl AsRef<Path>) -> ZxResult {
166        self.pipe_to_file_from(PipeFrom::Stdout, path).await
167    }
168
169    /// Run the command writing the selected stream into the file at `path`.
170    pub async fn pipe_to_file_from(self, from: PipeFrom, path: impl AsRef<Path>) -> ZxResult {
171        let nothrow = self.opts.nothrow;
172        let label = self.cmd.clone();
173        let mut file = match tokio::fs::File::create(path.as_ref()).await {
174            Ok(file) => file,
175            Err(err) => return Err(ProcessOutput::from_error(ErrorInfo::from_io(&err), label)),
176        };
177        let (tx, mut rx) = mpsc::unbounded_channel::<Vec<u8>>();
178        let writer = async move {
179            let mut result = Ok(());
180            while let Some(chunk) = rx.recv().await {
181                if result.is_ok() {
182                    result = file.write_all(&chunk).await;
183                }
184            }
185            result.and(file.flush().await)
186        };
187        let tap = Tap { from, tx };
188        let (out, written) =
189            tokio::join!(execute_chain(self, None, Some(tap), Control::new()), writer);
190        if let Err(err) = written {
191            return Err(ProcessOutput::from_error(ErrorInfo::from_io(&err), label));
192        }
193        settle(out, nothrow)
194    }
195
196    /// Run the command to completion.
197    pub async fn run(self) -> ZxResult {
198        let nothrow = self.opts.nothrow;
199        settle(
200            execute_chain(self, None, None, Control::new()).await,
201            nothrow,
202        )
203    }
204
205    /// Run the command synchronously (on a helper thread with its own
206    /// runtime, so it is safe to call from inside or outside of tokio).
207    pub fn run_sync(self) -> ZxResult {
208        let worker =
209            std::thread::spawn(move || {
210                match tokio::runtime::Builder::new_current_thread()
211                    .enable_all()
212                    .build()
213                {
214                    Ok(rt) => rt.block_on(self.run()),
215                    Err(err) => Err(ProcessOutput::from_error(ErrorInfo::from_io(&err), "")),
216                }
217            });
218        worker.join().unwrap_or_else(|_| {
219            Err(ProcessOutput::from_error(
220                ErrorInfo::new("zx runner thread panicked"),
221                "",
222            ))
223        })
224    }
225
226    /// Start the command in the background (requires a tokio runtime).
227    pub fn spawn(self) -> RunningProcess {
228        let ctl = Control::new();
229        let nothrow = self.opts.nothrow;
230        let kill_signal = self.opts.kill_signal.clone();
231        let handle = tokio::spawn(execute_chain(self, None, None, ctl.clone()));
232        RunningProcess {
233            handle,
234            ctl,
235            nothrow,
236            kill_signal,
237        }
238    }
239}
240
241impl IntoFuture for ProcessPromise {
242    type Output = ZxResult;
243    type IntoFuture = Pin<Box<dyn Future<Output = ZxResult> + Send>>;
244
245    fn into_future(self) -> Self::IntoFuture {
246        Box::pin(self.run())
247    }
248}
249
250fn settle(out: ProcessOutput, nothrow: bool) -> ZxResult {
251    if out.ok() || nothrow {
252        Ok(out)
253    } else {
254        Err(out)
255    }
256}
257
258/// A command started with [`ProcessPromise::spawn`].
259#[derive(Debug)]
260pub struct RunningProcess {
261    handle: tokio::task::JoinHandle<ProcessOutput>,
262    ctl: Control,
263    nothrow: bool,
264    kill_signal: String,
265}
266
267impl RunningProcess {
268    /// Pid of the (last) shell process, once it has been spawned.
269    pub async fn pid(&self) -> Option<u32> {
270        let mut rx = self.ctl.pid.subscribe();
271        loop {
272            if let Some(pid) = *rx.borrow_and_update() {
273                return Some(pid);
274            }
275            if self.is_finished() || rx.changed().await.is_err() {
276                return *rx.borrow();
277            }
278        }
279    }
280
281    /// `true` once the command has exited.
282    pub fn is_finished(&self) -> bool {
283        self.ctl.finished.load(Ordering::SeqCst)
284    }
285
286    /// Send `signal` (default: the `kill_signal` option) to the process tree.
287    pub fn kill(&self, signal: Option<&str>) -> Result<(), ZxError> {
288        if self.is_finished() {
289            return Err(ZxError::new("Too late to kill the process."));
290        }
291        let signal = signal.unwrap_or(&self.kill_signal).to_string();
292        #[cfg(unix)]
293        super::kill::parse_signal(&signal)?;
294        self.ctl.kill.send_replace(Some(signal));
295        Ok(())
296    }
297
298    /// Wait for the command to finish.
299    pub async fn wait(self) -> ZxResult {
300        match self.handle.await {
301            Ok(out) => settle(out, self.nothrow),
302            Err(err) => Err(ProcessOutput::from_error(
303                ErrorInfo::new(err.to_string()),
304                "",
305            )),
306        }
307    }
308}
309
310#[derive(Debug, Clone)]
311struct Control {
312    kill: Arc<watch::Sender<Option<String>>>,
313    pid: Arc<watch::Sender<Option<u32>>>,
314    finished: Arc<AtomicBool>,
315}
316
317impl Control {
318    fn new() -> Self {
319        Self {
320            kill: Arc::new(watch::Sender::new(None)),
321            pid: Arc::new(watch::Sender::new(None)),
322            finished: Arc::new(AtomicBool::new(false)),
323        }
324    }
325
326    fn finish(&self) {
327        self.finished.store(true, Ordering::SeqCst);
328        self.pid.send_modify(|_| {});
329    }
330}
331
332struct Tap {
333    from: PipeFrom,
334    tx: mpsc::UnboundedSender<Vec<u8>>,
335}
336
337/// Polls (5ms each) a failed pipe source waits for the destination to finish
338/// on its own before killing it.
339const PIPE_BREAK_GRACE_STEPS: usize = 20;
340
341fn execute_chain(
342    mut p: ProcessPromise,
343    feed: Option<mpsc::UnboundedReceiver<Vec<u8>>>,
344    tap: Option<Tap>,
345    ctl: Control,
346) -> BoxedRun {
347    Box::pin(async move {
348        let Some(link) = p.source.take() else {
349            return execute_one(p, feed, tap, ctl).await;
350        };
351        let (src, from) = *link;
352        let (tx, rx) = mpsc::unbounded_channel();
353        let src_nothrow = src.opts.nothrow;
354        let dest_ctl = ctl.clone();
355        let dest_signal = p.opts.kill_signal.clone();
356        let src_run = async move {
357            let out = execute_chain(src, feed, Some(Tap { from, tx }), Control::new()).await;
358            let broken = !out.ok() && !src_nothrow;
359            if broken {
360                // Give the destination a moment to drain what the source
361                // already produced (its stdin is closed now) before breaking.
362                for _ in 0..PIPE_BREAK_GRACE_STEPS {
363                    if dest_ctl.finished.load(Ordering::SeqCst) {
364                        break;
365                    }
366                    tokio::time::sleep(Duration::from_millis(5)).await;
367                }
368            }
369            if broken && !dest_ctl.finished.load(Ordering::SeqCst) {
370                let signal = out.signal.clone().unwrap_or(dest_signal);
371                dest_ctl.kill.send_replace(Some(signal));
372            }
373            (out, broken)
374        };
375        let ((src_out, broken), mut out) =
376            tokio::join!(src_run, execute_one(p, Some(rx), tap, ctl));
377        if broken {
378            out.exit_code = src_out.exit_code;
379            if src_out.signal.is_some() {
380                out.signal = src_out.signal;
381            }
382            if out.error.is_none() {
383                out.error = src_out.error;
384            }
385        }
386        out
387    })
388}
389
390async fn execute_one(
391    p: ProcessPromise,
392    feed: Option<mpsc::UnboundedReceiver<Vec<u8>>>,
393    tap: Option<Tap>,
394    ctl: Control,
395) -> ProcessOutput {
396    let out = spawn_and_wait(p, feed, tap, &ctl).await;
397    ctl.finish();
398    out
399}
400
401#[derive(Default)]
402struct Store {
403    stdout: Vec<u8>,
404    stderr: Vec<u8>,
405    stdall: Vec<u8>,
406}
407
408fn log_bytes(data: &[u8]) {
409    eprint!("{}", String::from_utf8_lossy(data));
410}
411
412fn configure_command(
413    opts: &Options,
414    shell: &str,
415    full_cmd: &str,
416    cwd: &Path,
417) -> tokio::process::Command {
418    let mut command = tokio::process::Command::new(shell);
419    command
420        .arg("-c")
421        .arg(full_cmd)
422        .current_dir(cwd)
423        .stdin(Stdio::piped())
424        .stdout(Stdio::piped())
425        .stderr(Stdio::piped());
426    if let Some(env) = &opts.env {
427        command.env_clear().envs(env);
428    }
429    crate::local_bin::apply_prefer_local(&mut command, opts.env.as_ref(), cwd, &opts.prefer_local);
430    #[cfg(unix)]
431    command.process_group(0);
432    command
433}
434
435async fn spawn_and_wait(
436    p: ProcessPromise,
437    feed: Option<mpsc::UnboundedReceiver<Vec<u8>>>,
438    tap: Option<Tap>,
439    ctl: &Control,
440) -> ProcessOutput {
441    let ProcessPromise {
442        opts,
443        cmd,
444        build_error,
445        ..
446    } = p;
447    if let Some(err) = build_error {
448        return ProcessOutput::from_error(ErrorInfo::new(err.message()), cmd);
449    }
450    let Some(shell) = opts.shell.clone() else {
451        let msg = format!("No shell is available: {DOCS_URL}/shell");
452        return ProcessOutput::from_error(ErrorInfo::new(msg), cmd);
453    };
454    let cwd = opts.effective_cwd();
455    if !cwd.is_dir() {
456        let msg = format!("The working directory '{}' does not exist.", cwd.display());
457        return ProcessOutput::from_error(ErrorInfo::new(msg), cmd);
458    }
459    let full_cmd = format!("{}{}{}", opts.prefix, cmd, opts.postfix);
460    let mut command = configure_command(&opts, &shell, &full_cmd, &cwd);
461    if opts.verbose && !opts.quiet {
462        super::log::log(&super::log::LogEntry::Cmd { cmd: cmd.clone() }, true);
463    }
464    let started = Instant::now();
465    let mut child = match command.spawn() {
466        Ok(child) => child,
467        Err(err) => {
468            return ProcessOutput::from_error(ErrorInfo::from_io(&err), cmd)
469                .with_duration(started.elapsed())
470        }
471    };
472    let pid = child.id();
473    ctl.pid.send_replace(pid);
474
475    let store = Mutex::new(Store::default());
476    let stdin = child.stdin.take();
477    let input = opts.input.clone();
478    let stdin_fut = async move {
479        let Some(mut writer) = stdin else { return };
480        if let Some(data) = input {
481            let _ = writer.write_all(&data).await;
482        } else if let Some(mut feed) = feed {
483            while let Some(chunk) = feed.recv().await {
484                if writer.write_all(&chunk).await.is_err() {
485                    break;
486                }
487            }
488        }
489        let _ = writer.shutdown().await;
490    };
491    let log_stdout = opts.verbose && !opts.quiet && tap.is_none();
492    let log_stderr = !opts.quiet;
493    let read_out = read_stream(child.stdout.take(), false, &store, tap.as_ref(), log_stdout);
494    let read_err = read_stream(child.stderr.take(), true, &store, tap.as_ref(), log_stderr);
495    let status = {
496        let wait = wait_child(&mut child, pid, &opts, ctl);
497        let main = async { tokio::join!(wait, read_out, read_err).0 };
498        tokio::pin!(main);
499        tokio::pin!(stdin_fut);
500        let mut stdin_done = false;
501        loop {
502            tokio::select! {
503                status = &mut main => break status,
504                _ = &mut stdin_fut, if !stdin_done => stdin_done = true,
505            }
506        }
507    };
508    drop(tap);
509    let duration = started.elapsed();
510    let store = store.into_inner().unwrap_or_else(|e| e.into_inner());
511    let mut out = ProcessOutput::new(
512        None,
513        None,
514        String::from_utf8_lossy(&store.stdout),
515        String::from_utf8_lossy(&store.stderr),
516        String::from_utf8_lossy(&store.stdall),
517    )
518    .with_from(cmd.clone())
519    .with_duration(duration);
520    match status {
521        Ok(status) => {
522            out.exit_code = status.code();
523            #[cfg(unix)]
524            {
525                use std::os::unix::process::ExitStatusExt;
526                out.signal = status.signal().and_then(signal_name);
527            }
528        }
529        Err(err) => out.error = Some(ErrorInfo::from_io(&err)),
530    }
531    #[cfg(not(unix))]
532    let _ = signal_name;
533    out
534}
535
536async fn wait_child(
537    child: &mut tokio::process::Child,
538    pid: Option<u32>,
539    opts: &Options,
540    ctl: &Control,
541) -> std::io::Result<std::process::ExitStatus> {
542    let deadline = opts.timeout.map(|t| tokio::time::Instant::now() + t);
543    let mut kill_rx = ctl.kill.subscribe();
544    let mut pending = kill_rx.borrow_and_update().clone();
545    let mut timed_out = false;
546    let mut kill_open = true;
547    loop {
548        if let (Some(signal), Some(pid)) = (pending.take(), pid) {
549            let _ = kill_tree(pid, &signal);
550        }
551        let sleep = async {
552            match deadline {
553                Some(at) => tokio::time::sleep_until(at).await,
554                None => std::future::pending().await,
555            }
556        };
557        tokio::select! {
558            status = child.wait() => return status,
559            _ = sleep, if !timed_out => {
560                timed_out = true;
561                pending = Some(opts.timeout_signal.clone());
562            }
563            changed = kill_rx.changed(), if kill_open => match changed {
564                Ok(()) => pending = kill_rx.borrow_and_update().clone(),
565                Err(_) => kill_open = false,
566            },
567        }
568    }
569}
570
571async fn read_stream<R: AsyncRead + Unpin>(
572    reader: Option<R>,
573    is_stderr: bool,
574    store: &Mutex<Store>,
575    tap: Option<&Tap>,
576    log: bool,
577) {
578    let Some(mut reader) = reader else { return };
579    let mut buf = vec![0u8; 8192];
580    loop {
581        let n = match reader.read(&mut buf).await {
582            Ok(0) | Err(_) => break,
583            Ok(n) => n,
584        };
585        let chunk = &buf[..n];
586        {
587            let mut store = store.lock().unwrap_or_else(|e| e.into_inner());
588            if is_stderr {
589                store.stderr.extend_from_slice(chunk);
590            } else {
591                store.stdout.extend_from_slice(chunk);
592            }
593            store.stdall.extend_from_slice(chunk);
594        }
595        if let Some(tap) = tap {
596            let wanted = match tap.from {
597                PipeFrom::Stdout => !is_stderr,
598                PipeFrom::Stderr => is_stderr,
599                PipeFrom::Stdall => true,
600            };
601            if wanted {
602                let _ = tap.tx.send(chunk.to_vec());
603            }
604        }
605        if log {
606            log_bytes(chunk);
607        }
608    }
609}
command_stream/zx/process.rs

command_stream/zx/
process.rs
