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    pub async fn wait(&mut self) -> Result<ProcessWaitResult, error::Error> {
50        #[allow(unused_mut, reason = "only mutated on some platforms")]
51        let mut sigtstp = sys::signal::tstp_signal_listener()?;
52        #[allow(unused_mut, reason = "only mutated on some platforms")]
53        let mut sigchld = sys::signal::chld_signal_listener()?;
54
55        // The child may have stopped before there were listeners to hear about
56        // it (e.g., straight after `fg` resumed it).
57        if let Some(pid) = self.pid
58            && sys::signal::poll_for_stopped_child(pid)?
59        {
60            return Ok(ProcessWaitResult::Stopped);
61        }
62
63        #[allow(clippy::ignored_unit_patterns)]
64        loop {
65            tokio::select! {
66                output = &mut self.exec_future => {
67                    break Ok(ProcessWaitResult::Completed(output?))
68                },
69                // Only this child's stop counts: another wait may be under way at the same time
70                // (e.g., a background job's), for a child of its own, and the terminal's stop
71                // may reach the shell without stopping this child.
72                _ = sigtstp.recv() => {
73                    if self.poll_for_stop()? {
74                        break Ok(ProcessWaitResult::Stopped);
75                    }
76                },
77                _ = sigchld.recv() => {
78                    if self.poll_for_stop()? {
79                        break Ok(ProcessWaitResult::Stopped);
80                    }
81                },
82                _ = sys::signal::await_ctrl_c() => {
83                    // SIGINT got thrown. Handle it and continue looping. The child should
84                    // have received it as well, and either handled it or ended up getting
85                    // terminated (in which case we'll see the child exit).
86                },
87            }
88        }
89    }
90
91    /// Returns whether the process has stopped, consuming its stop status.
92    fn poll_for_stop(&self) -> Result<bool, error::Error> {
93        match self.pid {
94            Some(pid) => sys::signal::poll_for_stopped_children(pid),
95            None => Ok(false),
96        }
97    }
98
99    pub(crate) fn poll(&mut self) -> Option<Result<std::process::Output, error::Error>> {
100        let checkable_future = &mut self.exec_future;
101        checkable_future
102            .now_or_never()
103            .map(|result| result.map_err(Into::into))
104    }
105}
106
107/// Represents the result of waiting for an executing process.
108pub enum ProcessWaitResult {
109    /// The process completed.
110    Completed(std::process::Output),
111    /// The process stopped and has not yet completed.
112    Stopped,
113}