Use async pid fd instead of blocking waitid to wait for a child process to exit

crates/containerd-shim-wasm/src/sys/unix/container/instance.rs

@@ -10,9 +10,7 @@ use libcontainer::container::builder::ContainerBuilder;
 use libcontainer::container::Container;
 use libcontainer::signal::Signal;
 use libcontainer::syscall::syscall::SyscallType;
-use nix::errno::Errno;
-use nix::sys::wait::{waitid, Id as WaitID, WaitPidFlag, WaitStatus};
-use nix::unistd::Pid;
+use nix::sys::wait::WaitStatus;
 use oci_spec::image::Platform;
 use zygote::{WireError, Zygote};
 
@@ -24,6 +22,7 @@ use crate::sandbox::{
     containerd, Error as SandboxError, Instance as SandboxInstance, InstanceConfig,
 };
 use crate::sys::container::executor::Executor;
+use crate::sys::pid_fd::PidFd;
 use crate::sys::stdio::open;
 
 const DEFAULT_CONTAINER_ROOT_DIR: &str = "/run/containerd";
@@ -107,20 +106,22 @@ impl<E: Engine + Default> SandboxInstance for Instance<E> {
         let mut container = self.container.lock().expect("Poisoned mutex");
         let pid = container.pid().context("failed to get pid")?.as_raw();
 
+        // Use a pidfd FD so that we can wait for the process to exit asynchronously.
+        let pidfd = PidFd::new(pid)?;
+
         container.start()?;
 
         let exit_code = self.exit_code.clone();
         thread::spawn(move || {
             // move the exit code guard into this thread
             let _guard = guard;
 
-            let status = match waitid(WaitID::Pid(Pid::from_raw(pid)), WaitPidFlag::WEXITED) {
+            let status = match pidfd.wait().block_on() {
                 Ok(WaitStatus::Exited(_, status)) => status,
-                Ok(WaitStatus::Signaled(_, sig, _)) => sig as i32,
-                Ok(_) => 0,
-                Err(Errno::ECHILD) => {
-                    log::info!("no child process");
-                    0
+                Ok(WaitStatus::Signaled(_, sig, _)) => 128 + sig as i32,
+                Ok(res) => {
+                    log::error!("waitpid unexpected result: {res:?}");
+                    137
                 }
                 Err(e) => {
                     log::error!("waitpid failed: {e}");
@@ -130,7 +131,7 @@ impl<E: Engine + Default> SandboxInstance for Instance<E> {
             let _ = exit_code.set((status, Utc::now()));
         });
 
-        Ok(pid as u32)
+        Ok(pid as _)
     }
 
     /// Send a signal to the instance

crates/containerd-shim-wasm/src/sys/unix/mod.rs

@@ -1,3 +1,5 @@
 pub mod container;
 pub mod metrics;
 pub mod stdio;
+
+mod pid_fd;

crates/containerd-shim-wasm/src/sys/unix/pid_fd.rs

@@ -0,0 +1,58 @@
+use std::os::fd::{AsFd, FromRawFd as _, OwnedFd, RawFd};
+
+use containerd_shim::monitor::{monitor_subscribe, wait_pid, Subscription, Topic};
+use libc::pid_t;
+use nix::errno::Errno;
+use nix::sys::wait::{waitid, Id, WaitPidFlag, WaitStatus};
+use nix::unistd::Pid;
+use tokio::io::unix::AsyncFd;
+
+pub(super) struct PidFd {
+    fd: OwnedFd,
+    pid: pid_t,
+    subs: Subscription,
+}
+
+impl PidFd {
+    pub(super) fn new(pid: impl Into<pid_t>) -> anyhow::Result<Self> {
+        use libc::{syscall, SYS_pidfd_open, PIDFD_NONBLOCK};
+        let pid = pid.into();
+        let pidfd = unsafe { syscall(SYS_pidfd_open, pid, PIDFD_NONBLOCK) };
+        if pidfd == -1 {
+            return Err(std::io::Error::last_os_error().into());
+        }
+        let fd = unsafe { OwnedFd::from_raw_fd(pidfd as RawFd) };
+        let subs = monitor_subscribe(Topic::Pid)?;
+        Ok(Self { fd, pid, subs })
+    }
+
+    pub(super) async fn wait(self) -> std::io::Result<WaitStatus> {
+        let fd = AsyncFd::new(self.fd)?;
+        loop {
+            // Check with non-blocking waitid before awaiting on fd.
+            // On some platforms, the readiness detecting mechanism relies on
+            // edge-triggered notifications.
+            // This means that we could miss a notification if the process exits
+            // before we create the AsyncFd.
+            // See https://docs.rs/tokio/latest/tokio/io/unix/struct.AsyncFd.html
+            match waitid(
+                Id::PIDFd(fd.as_fd()),
+                WaitPidFlag::WEXITED | WaitPidFlag::WNOHANG,
+            ) {
+                Ok(WaitStatus::StillAlive) => {
+                    let _ = fd.readable().await?;
+                }
+                Ok(status) => {
+                    return Ok(status);
+                }
+                Err(Errno::ECHILD) => {
+                    // The process has already been reaped by the containerd-shim reaper.
+                    // Get the status from there.
+                    let status = wait_pid(self.pid, self.subs);
+                    return Ok(WaitStatus::Exited(Pid::from_raw(self.pid), status));
+                }
+                Err(err) => return Err(err.into()),
+            }
+        }
+    }
+}