use crate::{ io::{interest::Interest, PollEvented}, process::{ imp::{orphan::Wait, OrphanQueue}, kill::Kill, }, util::error::RUNTIME_SHUTTING_DOWN_ERROR, }; use libc::{syscall, SYS_pidfd_open, ENOSYS, PIDFD_NONBLOCK}; use mio::{event::Source, unix::SourceFd}; use std::{ fs::File, future::Future, io, marker::Unpin, ops::Deref, os::unix::io::{AsRawFd, FromRawFd, RawFd}, pin::Pin, process::ExitStatus, sync::atomic::{AtomicBool, Ordering::Relaxed}, task::{ready, Context, Poll}, }; #[derive(Debug)] struct Pidfd { fd: File, } impl Pidfd { fn open(pid: u32) -> Option { // Store false (0) to reduce executable size static NO_PIDFD_SUPPORT: AtomicBool = AtomicBool::new(false); if NO_PIDFD_SUPPORT.load(Relaxed) { return None; } // Safety: The following function calls invovkes syscall pidfd_open, // which takes two parameter: pidfd_open(fd: c_int, flag: c_int) let fd = unsafe { syscall(SYS_pidfd_open, pid, PIDFD_NONBLOCK) }; if fd == -1 { let errno = io::Error::last_os_error().raw_os_error().unwrap(); if errno == ENOSYS { NO_PIDFD_SUPPORT.store(true, Relaxed) } None } else { // Safety: pidfd_open returns -1 on error or a valid fd with ownership. Some(Pidfd { fd: unsafe { File::from_raw_fd(fd as i32) }, }) } } } impl AsRawFd for Pidfd { fn as_raw_fd(&self) -> RawFd { self.fd.as_raw_fd() } } impl Source for Pidfd { fn register( &mut self, registry: &mio::Registry, token: mio::Token, interest: mio::Interest, ) -> io::Result<()> { SourceFd(&self.as_raw_fd()).register(registry, token, interest) } fn reregister( &mut self, registry: &mio::Registry, token: mio::Token, interest: mio::Interest, ) -> io::Result<()> { SourceFd(&self.as_raw_fd()).reregister(registry, token, interest) } fn deregister(&mut self, registry: &mio::Registry) -> io::Result<()> { SourceFd(&self.as_raw_fd()).deregister(registry) } } #[derive(Debug)] struct PidfdReaperInner where W: Unpin, { inner: W, pidfd: PollEvented, } #[allow(deprecated)] fn is_rt_shutdown_err(err: &io::Error) -> bool { if let Some(inner) = err.get_ref() { // Using `Error::description()` is more efficient than `format!("{inner}")`, // so we use it here even if it is deprecated. err.kind() == io::ErrorKind::Other && inner.source().is_none() && inner.description() == RUNTIME_SHUTTING_DOWN_ERROR } else { false } } impl Future for PidfdReaperInner where W: Wait + Unpin, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = Pin::into_inner(self); match ready!(this.pidfd.poll_read_ready(cx)) { Err(err) if is_rt_shutdown_err(&err) => { this.pidfd.reregister(Interest::READABLE)?; ready!(this.pidfd.poll_read_ready(cx))? } res => res?, } Poll::Ready(Ok(this .inner .try_wait()? .expect("pidfd is ready to read, the process should have exited"))) } } #[derive(Debug)] pub(crate) struct PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { inner: Option>, orphan_queue: Q, } impl Deref for PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { type Target = W; fn deref(&self) -> &Self::Target { &self.inner.as_ref().expect("inner has gone away").inner } } impl PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { pub(crate) fn new(inner: W, orphan_queue: Q) -> Result, W)> { if let Some(pidfd) = Pidfd::open(inner.id()) { match PollEvented::new_with_interest(pidfd, Interest::READABLE) { Ok(pidfd) => Ok(Self { inner: Some(PidfdReaperInner { pidfd, inner }), orphan_queue, }), Err(io_error) => Err((Some(io_error), inner)), } } else { Err((None, inner)) } } pub(crate) fn inner_mut(&mut self) -> &mut W { &mut self.inner.as_mut().expect("inner has gone away").inner } } impl Future for PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { type Output = io::Result; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { Pin::new( Pin::into_inner(self) .inner .as_mut() .expect("inner has gone away"), ) .poll(cx) } } impl Kill for PidfdReaper where W: Wait + Unpin + Kill, Q: OrphanQueue + Unpin, { fn kill(&mut self) -> io::Result<()> { self.inner_mut().kill() } } impl Drop for PidfdReaper where W: Wait + Unpin, Q: OrphanQueue + Unpin, { fn drop(&mut self) { let mut orphan = self.inner.take().expect("inner has gone away").inner; if let Ok(Some(_)) = orphan.try_wait() { return; } self.orphan_queue.push_orphan(orphan); } } #[cfg(all(test, not(loom), not(miri)))] mod test { use super::*; use crate::{ process::unix::orphan::test::MockQueue, runtime::{Builder as RuntimeBuilder, Runtime}, }; use std::process::{Command, Output}; fn create_runtime() -> Runtime { RuntimeBuilder::new_current_thread() .enable_io() .build() .unwrap() } fn run_test(fut: impl Future) { create_runtime().block_on(fut) } fn is_pidfd_available() -> bool { let Output { stdout, status, .. } = Command::new("uname").arg("-r").output().unwrap(); assert!(status.success()); let stdout = String::from_utf8_lossy(&stdout); let mut kernel_version_iter = match stdout.split_once('-') { Some((version, _)) => version, _ => &stdout, } .split('.'); let major: u32 = kernel_version_iter.next().unwrap().parse().unwrap(); let minor: u32 = kernel_version_iter.next().unwrap().parse().unwrap(); major >= 6 || (major == 5 && minor >= 10) } #[test] fn test_pidfd_reaper_poll() { if !is_pidfd_available() { eprintln!("pidfd is not available on this linux kernel, skip this test"); return; } let queue = MockQueue::new(); run_test(async { let child = Command::new("true").spawn().unwrap(); let pidfd_reaper = PidfdReaper::new(child, &queue).unwrap(); let exit_status = pidfd_reaper.await.unwrap(); assert!(exit_status.success()); }); assert!(queue.all_enqueued.borrow().is_empty()); } #[test] fn test_pidfd_reaper_kill() { if !is_pidfd_available() { eprintln!("pidfd is not available on this linux kernel, skip this test"); return; } let queue = MockQueue::new(); run_test(async { let child = Command::new("sleep").arg("1800").spawn().unwrap(); let mut pidfd_reaper = PidfdReaper::new(child, &queue).unwrap(); pidfd_reaper.kill().unwrap(); let exit_status = pidfd_reaper.await.unwrap(); assert!(!exit_status.success()); }); assert!(queue.all_enqueued.borrow().is_empty()); } #[test] fn test_pidfd_reaper_drop() { if !is_pidfd_available() { eprintln!("pidfd is not available on this linux kernel, skip this test"); return; } let queue = MockQueue::new(); let mut child = Command::new("sleep").arg("1800").spawn().unwrap(); run_test(async { let _pidfd_reaper = PidfdReaper::new(&mut child, &queue).unwrap(); }); assert_eq!(queue.all_enqueued.borrow().len(), 1); child.kill().unwrap(); child.wait().unwrap(); } }