tokio/process/unix/
pidfd_reaper.rs

1use crate::{
2    io::{interest::Interest, PollEvented},
3    process::{
4        imp::{orphan::Wait, OrphanQueue},
5        kill::Kill,
6    },
7    util::error::RUNTIME_SHUTTING_DOWN_ERROR,
8};
9
10use libc::{syscall, SYS_pidfd_open, ENOSYS, PIDFD_NONBLOCK};
11use mio::{event::Source, unix::SourceFd};
12use std::{
13    fs::File,
14    future::Future,
15    io,
16    marker::Unpin,
17    ops::Deref,
18    os::unix::io::{AsRawFd, FromRawFd, RawFd},
19    pin::Pin,
20    process::ExitStatus,
21    sync::atomic::{AtomicBool, Ordering::Relaxed},
22    task::{Context, Poll},
23};
24
25#[derive(Debug)]
26struct Pidfd {
27    fd: File,
28}
29
30impl Pidfd {
31    fn open(pid: u32) -> Option<Pidfd> {
32        // Store false (0) to reduce executable size
33        static NO_PIDFD_SUPPORT: AtomicBool = AtomicBool::new(false);
34
35        if NO_PIDFD_SUPPORT.load(Relaxed) {
36            return None;
37        }
38
39        // Safety: The following function calls invovkes syscall pidfd_open,
40        // which takes two parameter: pidfd_open(fd: c_int, flag: c_int)
41        let fd = unsafe { syscall(SYS_pidfd_open, pid, PIDFD_NONBLOCK) };
42        if fd == -1 {
43            let errno = io::Error::last_os_error().raw_os_error().unwrap();
44
45            if errno == ENOSYS {
46                NO_PIDFD_SUPPORT.store(true, Relaxed)
47            }
48
49            None
50        } else {
51            // Safety: pidfd_open returns -1 on error or a valid fd with ownership.
52            Some(Pidfd {
53                fd: unsafe { File::from_raw_fd(fd as i32) },
54            })
55        }
56    }
57}
58
59impl AsRawFd for Pidfd {
60    fn as_raw_fd(&self) -> RawFd {
61        self.fd.as_raw_fd()
62    }
63}
64
65impl Source for Pidfd {
66    fn register(
67        &mut self,
68        registry: &mio::Registry,
69        token: mio::Token,
70        interest: mio::Interest,
71    ) -> io::Result<()> {
72        SourceFd(&self.as_raw_fd()).register(registry, token, interest)
73    }
74
75    fn reregister(
76        &mut self,
77        registry: &mio::Registry,
78        token: mio::Token,
79        interest: mio::Interest,
80    ) -> io::Result<()> {
81        SourceFd(&self.as_raw_fd()).reregister(registry, token, interest)
82    }
83
84    fn deregister(&mut self, registry: &mio::Registry) -> io::Result<()> {
85        SourceFd(&self.as_raw_fd()).deregister(registry)
86    }
87}
88
89#[derive(Debug)]
90struct PidfdReaperInner<W>
91where
92    W: Unpin,
93{
94    inner: W,
95    pidfd: PollEvented<Pidfd>,
96}
97
98#[allow(deprecated)]
99fn is_rt_shutdown_err(err: &io::Error) -> bool {
100    if let Some(inner) = err.get_ref() {
101        // Using `Error::description()` is more efficient than `format!("{inner}")`,
102        // so we use it here even if it is deprecated.
103        err.kind() == io::ErrorKind::Other
104            && inner.source().is_none()
105            && inner.description() == RUNTIME_SHUTTING_DOWN_ERROR
106    } else {
107        false
108    }
109}
110
111impl<W> Future for PidfdReaperInner<W>
112where
113    W: Wait + Unpin,
114{
115    type Output = io::Result<ExitStatus>;
116
117    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
118        let this = Pin::into_inner(self);
119
120        match this.pidfd.registration().poll_read_ready(cx) {
121            Poll::Ready(Ok(evt)) => {
122                if let Some(exit_code) = this.inner.try_wait()? {
123                    return Poll::Ready(Ok(exit_code));
124                }
125                this.pidfd.registration().clear_readiness(evt);
126            }
127            Poll::Ready(Err(err)) if is_rt_shutdown_err(&err) => {}
128            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
129            Poll::Pending => return Poll::Pending,
130        };
131
132        this.pidfd.reregister(Interest::READABLE)?;
133        cx.waker().wake_by_ref();
134        Poll::Pending
135    }
136}
137
138#[derive(Debug)]
139pub(crate) struct PidfdReaper<W, Q>
140where
141    W: Wait + Unpin,
142    Q: OrphanQueue<W> + Unpin,
143{
144    inner: Option<PidfdReaperInner<W>>,
145    orphan_queue: Q,
146}
147
148impl<W, Q> Deref for PidfdReaper<W, Q>
149where
150    W: Wait + Unpin,
151    Q: OrphanQueue<W> + Unpin,
152{
153    type Target = W;
154
155    fn deref(&self) -> &Self::Target {
156        &self.inner.as_ref().expect("inner has gone away").inner
157    }
158}
159
160impl<W, Q> PidfdReaper<W, Q>
161where
162    W: Wait + Unpin,
163    Q: OrphanQueue<W> + Unpin,
164{
165    pub(crate) fn new(inner: W, orphan_queue: Q) -> Result<Self, (Option<io::Error>, W)> {
166        if let Some(pidfd) = Pidfd::open(inner.id()) {
167            match PollEvented::new_with_interest(pidfd, Interest::READABLE) {
168                Ok(pidfd) => Ok(Self {
169                    inner: Some(PidfdReaperInner { pidfd, inner }),
170                    orphan_queue,
171                }),
172                Err(io_error) => Err((Some(io_error), inner)),
173            }
174        } else {
175            Err((None, inner))
176        }
177    }
178
179    pub(crate) fn inner_mut(&mut self) -> &mut W {
180        &mut self.inner.as_mut().expect("inner has gone away").inner
181    }
182}
183
184impl<W, Q> Future for PidfdReaper<W, Q>
185where
186    W: Wait + Unpin,
187    Q: OrphanQueue<W> + Unpin,
188{
189    type Output = io::Result<ExitStatus>;
190
191    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
192        Pin::new(
193            Pin::into_inner(self)
194                .inner
195                .as_mut()
196                .expect("inner has gone away"),
197        )
198        .poll(cx)
199    }
200}
201
202impl<W, Q> Kill for PidfdReaper<W, Q>
203where
204    W: Wait + Unpin + Kill,
205    Q: OrphanQueue<W> + Unpin,
206{
207    fn kill(&mut self) -> io::Result<()> {
208        self.inner_mut().kill()
209    }
210}
211
212impl<W, Q> Drop for PidfdReaper<W, Q>
213where
214    W: Wait + Unpin,
215    Q: OrphanQueue<W> + Unpin,
216{
217    fn drop(&mut self) {
218        let mut orphan = self.inner.take().expect("inner has gone away").inner;
219        if let Ok(Some(_)) = orphan.try_wait() {
220            return;
221        }
222
223        self.orphan_queue.push_orphan(orphan);
224    }
225}
226
227#[cfg(all(test, not(loom), not(miri)))]
228mod test {
229    use super::*;
230    use crate::{
231        process::unix::orphan::test::MockQueue,
232        runtime::{Builder as RuntimeBuilder, Runtime},
233    };
234    use std::process::{Command, Output};
235
236    fn create_runtime() -> Runtime {
237        RuntimeBuilder::new_current_thread()
238            .enable_io()
239            .build()
240            .unwrap()
241    }
242
243    fn run_test(fut: impl Future<Output = ()>) {
244        create_runtime().block_on(fut)
245    }
246
247    fn is_pidfd_available() -> bool {
248        let Output { stdout, status, .. } = Command::new("uname").arg("-r").output().unwrap();
249        assert!(status.success());
250        let stdout = String::from_utf8_lossy(&stdout);
251
252        let mut kernel_version_iter = match stdout.split_once('-') {
253            Some((version, _)) => version,
254            _ => &stdout,
255        }
256        .split('.');
257
258        let major: u32 = kernel_version_iter.next().unwrap().parse().unwrap();
259        let minor: u32 = kernel_version_iter.next().unwrap().parse().unwrap();
260
261        major >= 6 || (major == 5 && minor >= 10)
262    }
263
264    #[test]
265    fn test_pidfd_reaper_poll() {
266        if !is_pidfd_available() {
267            eprintln!("pidfd is not available on this linux kernel, skip this test");
268            return;
269        }
270
271        let queue = MockQueue::new();
272
273        run_test(async {
274            let child = Command::new("true").spawn().unwrap();
275            let pidfd_reaper = PidfdReaper::new(child, &queue).unwrap();
276
277            let exit_status = pidfd_reaper.await.unwrap();
278            assert!(exit_status.success());
279        });
280
281        assert!(queue.all_enqueued.borrow().is_empty());
282    }
283
284    #[test]
285    fn test_pidfd_reaper_kill() {
286        if !is_pidfd_available() {
287            eprintln!("pidfd is not available on this linux kernel, skip this test");
288            return;
289        }
290
291        let queue = MockQueue::new();
292
293        run_test(async {
294            let child = Command::new("sleep").arg("1800").spawn().unwrap();
295            let mut pidfd_reaper = PidfdReaper::new(child, &queue).unwrap();
296
297            pidfd_reaper.kill().unwrap();
298
299            let exit_status = pidfd_reaper.await.unwrap();
300            assert!(!exit_status.success());
301        });
302
303        assert!(queue.all_enqueued.borrow().is_empty());
304    }
305
306    #[test]
307    fn test_pidfd_reaper_drop() {
308        if !is_pidfd_available() {
309            eprintln!("pidfd is not available on this linux kernel, skip this test");
310            return;
311        }
312
313        let queue = MockQueue::new();
314
315        let mut child = Command::new("sleep").arg("1800").spawn().unwrap();
316
317        run_test(async {
318            let _pidfd_reaper = PidfdReaper::new(&mut child, &queue).unwrap();
319        });
320
321        assert_eq!(queue.all_enqueued.borrow().len(), 1);
322
323        child.kill().unwrap();
324        child.wait().unwrap();
325    }
326}
tokio/process/unix/pidfd_reaper.rs

tokio/process/unix/
pidfd_reaper.rs
