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

@@ -9,9 +9,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 crate::container::Engine;
@@ -22,6 +20,7 @@ use crate::sandbox::{
     containerd, Error as SandboxError, Instance as SandboxInstance, InstanceConfig, Stdio,
 };
 use crate::sys::container::executor::Executor;
+use crate::sys::pid_fd::PidFd;
 
 static DEFAULT_CONTAINER_ROOT_DIR: &str = "/run/containerd";
 
@@ -80,7 +79,10 @@ impl<E: Engine> SandboxInstance for Instance<E> {
 
         let container_root = get_instance_root(&self.rootdir, &self.id)?;
         let mut container = Container::load(container_root)?;
-        let pid = container.pid().context("failed to get pid")?.as_raw();
+        let pid = container.pid().context("failed to get pid")?;
+
+        // Use a pidfd FD so that we can wait for the process to exit asynchronously.
+        let pidfd = PidFd::new(pid)?;
 
         container.start()?;
 
@@ -89,13 +91,12 @@ impl<E: Engine> SandboxInstance for Instance<E> {
             // 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(res) => {
+                    log::error!("waitpid unexpected result: {res:?}");
+                    137
                 }
                 Err(e) => {
                     log::error!("waitpid failed: {e}");
@@ -105,7 +106,7 @@ impl<E: Engine> SandboxInstance for Instance<E> {
             let _ = exit_code.set((status, Utc::now()));
         });
 
-        Ok(pid as u32)
+        Ok(pid.as_raw() 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,44 @@
+use std::os::fd::{AsFd, FromRawFd as _, OwnedFd, RawFd};
+
+use libc::pid_t;
+use nix::sys::wait::{waitid, Id, WaitPidFlag, WaitStatus};
+use tokio::io::unix::AsyncFd;
+
+pub(super) struct PidFd {
+    fd: OwnedFd,
+}
+
+impl PidFd {
+    pub(super) fn new(pid: impl Into<pid_t>) -> anyhow::Result<Self> {
+        use libc::{syscall, SYS_pidfd_open, PIDFD_NONBLOCK};
+        let pidfd = unsafe { syscall(SYS_pidfd_open, pid.into(), PIDFD_NONBLOCK) };
+        if pidfd == -1 {
+            return Err(std::io::Error::last_os_error().into());
+        }
+        let fd = unsafe { OwnedFd::from_raw_fd(pidfd as RawFd) };
+        Ok(Self { fd })
+    }
+
+    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,
+            )? {
+                WaitStatus::StillAlive => {
+                    let _ = fd.readable().await?;
+                }
+                status => {
+                    return Ok(status);
+                }
+            }
+        }
+    }
+}