From 41ef0a585c64abf2facbc1997344df09482ce368 Mon Sep 17 00:00:00 2001
From: Jorge Prendes <jorge.prendes@gmail.com>
Date: Fri, 22 Nov 2024 22:39:20 +0000
Subject: [PATCH] Use async pid fd instead of blocking waitid to wait for a
 child process to exit

Signed-off-by: Jorge Prendes <jorge.prendes@gmail.com>
---
 .../src/sys/unix/container/instance.rs        | 21 +++++------
 .../containerd-shim-wasm/src/sys/unix/mod.rs  |  2 ++
 .../src/sys/unix/pid_fd.rs                    | 35 +++++++++++++++++++
 3 files changed, 48 insertions(+), 10 deletions(-)
 create mode 100644 crates/containerd-shim-wasm/src/sys/unix/pid_fd.rs

diff --git a/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs b/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs
index ca9b31d58..239697212 100644
--- a/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs
+++ b/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
diff --git a/crates/containerd-shim-wasm/src/sys/unix/mod.rs b/crates/containerd-shim-wasm/src/sys/unix/mod.rs
index f59351f57..2b0adebfb 100644
--- a/crates/containerd-shim-wasm/src/sys/unix/mod.rs
+++ b/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;
diff --git a/crates/containerd-shim-wasm/src/sys/unix/pid_fd.rs b/crates/containerd-shim-wasm/src/sys/unix/pid_fd.rs
new file mode 100644
index 000000000..8d04f73a0
--- /dev/null
+++ b/crates/containerd-shim-wasm/src/sys/unix/pid_fd.rs
@@ -0,0 +1,35 @@
+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 {
+            match waitid(Id::PIDFd(fd.as_fd()), WaitPidFlag::WEXITED | WaitPidFlag::WNOHANG)? {
+                WaitStatus::StillAlive => {
+                    let _ = fd.readable().await?;
+                }
+                status => {
+                    return Ok(status);
+                }
+            }
+        }
+    }
+}