1mod capture;
4use crate::worker::{CancelToken, Failure, FailureKind};
5pub use capture::{
6 capture, capture_with, stream_with, CaptureError, CapturePolicy, CommandOutput, StdinPolicy,
7 StreamError, StreamOutput, StreamPolicy,
8};
9use parking_lot::Mutex;
10use std::io;
11use std::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, ExitStatus};
12use std::sync::Arc;
13
14#[derive(Default)]
15struct Group {
16 pid: Option<u32>,
17 #[cfg(unix)]
20 cancelled: bool,
21 #[cfg(unix)]
22 lease: Option<std::os::unix::net::UnixStream>,
23}
24impl Group {
25 fn signal(&self) -> Result<(), Failure> {
26 #[cfg(unix)]
27 if let Some(lease) = &self.lease {
28 return match lease.shutdown(std::net::Shutdown::Write) {
29 Ok(()) => Ok(()),
30 Err(error) if error.kind() == io::ErrorKind::NotConnected => Ok(()),
31 Err(error) => Err(Failure::new(
32 FailureKind::Io,
33 format!("close process lease: {error}"),
34 )),
35 };
36 }
37 #[cfg(unix)]
38 if let Some(pid) = self.pid {
39 if unsafe { libc::kill(-(pid as libc::pid_t), libc::SIGKILL) } == -1 {
43 let error = io::Error::last_os_error();
44 if error.raw_os_error() != Some(libc::ESRCH) {
45 return Err(Failure::new(
46 FailureKind::Io,
47 format!("kill process group: {error}"),
48 ));
49 }
50 }
51 }
52 Ok(())
53 }
54 #[cfg(unix)]
55 fn cancel(&mut self) -> Result<(), Failure> {
56 self.cancelled = true;
57 self.signal()
58 }
59}
60
61pub struct OwnedProcess {
64 child: Child,
65 group: Arc<Mutex<Group>>,
66 token: CancelToken,
67 reaped: bool,
68}
69impl OwnedProcess {
70 pub fn spawn(command: &mut Command, token: &CancelToken) -> Result<Self, Failure> {
71 #[cfg(not(unix))]
72 {
73 let _ = (command, token);
75 return Err(Failure::new(
76 FailureKind::Unavailable,
77 "process-group supervision requires Unix",
78 ));
79 }
80 #[cfg(unix)]
81 Self::spawn_unix(command, token, None)
82 }
83
84 #[cfg(unix)]
87 pub fn spawn_leased(
88 command: &mut Command,
89 token: &CancelToken,
90 lease: std::os::unix::net::UnixStream,
91 ) -> Result<Self, Failure> {
92 Self::spawn_unix(command, token, Some(lease))
93 }
94
95 #[cfg(unix)]
96 fn spawn_unix(
97 command: &mut Command,
98 token: &CancelToken,
99 lease: Option<std::os::unix::net::UnixStream>,
100 ) -> Result<Self, Failure> {
101 use std::os::unix::process::CommandExt;
102 if lease.is_none() {
103 command.process_group(0);
104 }
105 let group = Arc::new(Mutex::new(Group {
106 lease,
107 ..Group::default()
108 }));
109 let callback = group.clone();
110 token.register_cancel_resource(move || callback.lock().cancel())?;
111 if token.is_cancelled() {
112 token.clear_cancel_resource();
113 return Err(Failure::new(
114 FailureKind::Unavailable,
115 "process cancelled before spawn",
116 ));
117 }
118 let child = match command.spawn() {
119 Ok(child) => child,
120 Err(error) => {
121 token.clear_cancel_resource();
122 return Err(Failure::new(FailureKind::Spawn, error.to_string()));
123 }
124 };
125 let process = Self {
126 child,
127 group,
128 token: token.clone(),
129 reaped: false,
130 };
131 {
132 let mut group = process.group.lock();
133 group.pid = Some(process.child.id());
134 if group.cancelled || token.is_cancelled() {
135 group.cancel()?;
136 }
137 }
138 Ok(process)
139 }
140 pub fn take_stdin(&mut self) -> Option<ChildStdin> {
141 self.child.stdin.take()
142 }
143 pub fn take_stdout(&mut self) -> Option<ChildStdout> {
144 self.child.stdout.take()
145 }
146 pub fn take_stderr(&mut self) -> Option<ChildStderr> {
147 self.child.stderr.take()
148 }
149 pub fn terminate(&mut self) -> Result<(), Failure> {
150 self.group.lock().signal()
151 }
152
153 pub fn has_exited(&self) -> Result<bool, Failure> {
156 child_has_exited(&self.child)
157 }
158 pub fn wait(&mut self) -> Result<ExitStatus, Failure> {
161 while !self.reaped && !self.has_exited()? {
164 std::thread::park_timeout(std::time::Duration::from_millis(20));
165 }
166 self.group.lock().pid = None;
167 loop {
168 match self.child.wait() {
169 Ok(status) => {
170 self.reaped = true;
171 self.token.clear_cancel_resource();
172 return Ok(status);
173 }
174 Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
175 Err(error) => return Err(Failure::new(FailureKind::Wait, error.to_string())),
176 }
177 }
178 }
179}
180impl Drop for OwnedProcess {
181 fn drop(&mut self) {
182 if !self.reaped {
183 let _ = self.terminate();
184 #[cfg(unix)]
187 if self.group.lock().lease.is_none() {
188 let _ = self.child.kill();
189 }
190 #[cfg(not(unix))]
191 let _ = self.child.kill();
192 let _ = self.wait();
193 }
194 self.token.clear_cancel_resource();
195 }
196}
197
198pub fn child_has_exited(child: &Child) -> Result<bool, Failure> {
201 #[cfg(not(unix))]
202 {
203 let _ = child;
206 return Err(Failure::new(
207 FailureKind::Unavailable,
208 "process supervision requires Unix",
209 ));
210 }
211 #[cfg(unix)]
212 {
213 let mut info: libc::siginfo_t = unsafe { std::mem::zeroed() };
215 loop {
216 let result = unsafe {
219 libc::waitid(
220 libc::P_PID,
221 child.id() as libc::id_t,
222 &mut info,
223 libc::WEXITED | libc::WNOHANG | libc::WNOWAIT,
224 )
225 };
226 if result == 0 {
227 return Ok(unsafe { info.si_pid() } != 0);
229 }
230 let error = io::Error::last_os_error();
231 if error.kind() != io::ErrorKind::Interrupted {
232 return Err(Failure::new(FailureKind::Wait, error.to_string()));
233 }
234 }
235 }
236}
237
238#[cfg(all(test, unix))]
239mod tests {
240 use super::*;
241 use crate::worker::{self, CancelReason, Outcome};
242 use std::io::Read;
243 use std::os::fd::OwnedFd;
244 use std::os::unix::net::UnixStream;
245 use std::process::Stdio;
246 use std::sync::mpsc;
247 use std::time::Duration;
248
249 #[test]
250 fn leased_cancellation_allows_helper_cleanup_before_reaping() {
251 let (ready_tx, ready_rx) = mpsc::channel();
252 let (done_tx, done_rx) = mpsc::channel();
253 let handle = worker::spawn_effect(
254 "leased-cleanup-test",
255 move |outcome| {
256 done_tx.send(outcome).unwrap();
257 },
258 move |token| {
259 let (lease, child_lease) = UnixStream::pair().unwrap();
260 let mut command = Command::new("/bin/sh");
261 command
262 .args(["-c", "printf READY; cat >/dev/null; printf CLEANED"])
263 .stdin(Stdio::from(OwnedFd::from(child_lease)))
264 .stdout(Stdio::piped())
265 .stderr(Stdio::null());
266 let mut process = OwnedProcess::spawn_leased(&mut command, &token, lease).unwrap();
267 let mut stdout = process.take_stdout().unwrap();
268 let mut ready = [0; 5];
269 stdout.read_exact(&mut ready).unwrap();
270 assert_eq!(&ready, b"READY");
271 ready_tx.send(()).unwrap();
272 let mut output = Vec::new();
273 stdout.read_to_end(&mut output).unwrap();
274 let status = process.wait().unwrap();
275 Outcome::Success((output, status.success()))
276 },
277 );
278 ready_rx.recv_timeout(Duration::from_secs(5)).unwrap();
279 handle.cancel(CancelReason::Shutdown);
280 let outcome = done_rx.recv_timeout(Duration::from_secs(5)).unwrap();
281 let Outcome::Success((output, successful)) = outcome else {
282 panic!("cleanup outcome lost: {outcome:?}");
283 };
284 assert!(successful);
285 assert_eq!(output, b"CLEANED");
286 }
287}