Skip to main content

shrub_core/
processes.rs

1//! Process management
2
3use futures::FutureExt;
4
5use crate::{error, sys};
6
7/// A waitable future that will yield the results of a child process's
8/// execution.
9pub(crate) type WaitableChildProcess = std::pin::Pin<
10    Box<dyn futures::Future<Output = Result<std::process::Output, std::io::Error>> + Send + Sync>,
11>;
12
13/// Tracks a child process being awaited.
14pub struct ChildProcess {
15    /// A waitable future that will yield the results of a child process's
16    /// execution.
17    exec_future: WaitableChildProcess,
18    /// If available, the process ID of the child.
19    pid: Option<sys::process::ProcessId>,
20    /// If available, the process group ID of the child.
21    pgid: Option<sys::process::ProcessId>,
22}
23
24impl ChildProcess {
25    /// Wraps a child process and its future.
26    pub fn new(
27        child: sys::process::Child,
28        pid: Option<sys::process::ProcessId>,
29        pgid: Option<sys::process::ProcessId>,
30    ) -> Self {
31        Self {
32            exec_future: Box::pin(child.wait_with_output()),
33            pid,
34            pgid,
35        }
36    }
37
38    /// Returns the process's ID.
39    pub const fn pid(&self) -> Option<sys::process::ProcessId> {
40        self.pid
41    }
42
43    /// Returns the process's group ID.
44    pub const fn pgid(&self) -> Option<sys::process::ProcessId> {
45        self.pgid
46    }
47
48    /// Waits for the process to exit.
49    ///
50    /// # Errors
51    ///
52    /// Returns [`error::ErrorKind::PlatformError`] or
53    /// [`error::ErrorKind::IoError`] if registering the signal listeners,
54    /// checking for a stopped child, or waiting on the child fails.
55    pub async fn wait(&mut self) -> Result<ProcessWaitResult, error::Error> {
56        #[allow(unused_mut, reason = "only mutated on some platforms")]
57        let mut sigtstp = sys::signal::tstp_signal_listener()?;
58        #[allow(unused_mut, reason = "only mutated on some platforms")]
59        let mut sigchld = sys::signal::chld_signal_listener()?;
60
61        // The child may have stopped before there were listeners to hear about
62        // it (e.g., straight after `fg` resumed it).
63        if let Some(pid) = self.pid
64            && sys::signal::poll_for_stopped_child(pid)?
65        {
66            return Ok(ProcessWaitResult::Stopped);
67        }
68
69        #[allow(clippy::ignored_unit_patterns)]
70        loop {
71            tokio::select! {
72                output = &mut self.exec_future => {
73                    break Ok(ProcessWaitResult::Completed(output?))
74                },
75                // Only this child's stop counts: another wait may be under way at the same time
76                // (e.g., a background job's), for a child of its own, and the terminal's stop
77                // may reach the shell without stopping this child.
78                _ = sigtstp.recv() => {
79                    if self.poll_for_stop()? {
80                        break Ok(ProcessWaitResult::Stopped);
81                    }
82                },
83                _ = sigchld.recv() => {
84                    if self.poll_for_stop()? {
85                        break Ok(ProcessWaitResult::Stopped);
86                    }
87                },
88                _ = sys::signal::await_ctrl_c() => {
89                    // SIGINT got thrown. Handle it and continue looping. The child should
90                    // have received it as well, and either handled it or ended up getting
91                    // terminated (in which case we'll see the child exit).
92                },
93            }
94        }
95    }
96
97    /// Returns whether the process has stopped, consuming its stop status.
98    fn poll_for_stop(&self) -> Result<bool, error::Error> {
99        match self.pid {
100            Some(pid) => sys::signal::poll_for_stopped_children(pid),
101            None => Ok(false),
102        }
103    }
104
105    pub(crate) fn poll(&mut self) -> Option<Result<std::process::Output, error::Error>> {
106        let checkable_future = &mut self.exec_future;
107        checkable_future
108            .now_or_never()
109            .map(|result| result.map_err(Into::into))
110    }
111}
112
113/// Represents the result of waiting for an executing process.
114pub enum ProcessWaitResult {
115    /// The process completed.
116    Completed(std::process::Output),
117    /// The process stopped and has not yet completed.
118    Stopped,
119}