pidwait: replace busy-loop with pidfd_open and poll on Linux

src/uu/pidwait/Cargo.toml

@@ -14,12 +14,21 @@ version.workspace = true
 workspace = true
 
 [dependencies]
-nix = { workspace = true }
 uucore = { workspace = true, features = ["entries"] }
 clap = { workspace = true }
 regex = { workspace = true }
+
 uu_pgrep = { path = "../pgrep" }
 
+[target.'cfg(unix)'.dependencies]
+rustix = { workspace = true, features = ["event", "std"] }
+
+[target.'cfg(windows)'.dependencies]
+windows-sys = { workspace = true, features = [
+	"Win32_Foundation",
+	"Win32_System_Threading",
+] }
+
 [lib]
 path = "src/pidwait.rs"
 

src/uu/pidwait/src/wait.rs

@@ -5,42 +5,78 @@
 
 use uu_pgrep::process::ProcessInformation;
 
-// Dirty, but it works.
-// TODO: Use better implementation instead
 #[cfg(target_os = "linux")]
 pub(crate) fn wait(procs: &[ProcessInformation]) {
-    use std::{thread::sleep, time::Duration};
+    use std::os::fd::OwnedFd;
 
-    let mut list = procs.to_vec();
+    use rustix::event::{poll, PollFd, PollFlags};
+    use rustix::process::{pidfd_open, Pid, PidfdFlags};
 
-    loop {
-        for proc in &list.clone() {
-            // Check is running
-            if !is_running(proc.pid) {
-                list.retain(|it| it.pid != proc.pid);
+    let mut pidfds: Vec<OwnedFd> = procs
+        .iter()
+        .filter_map(|proc| {
+            let pid = Pid::from_raw(proc.pid as i32)?;
+            pidfd_open(pid, PidfdFlags::empty()).ok()
+        })
+        .collect();
+
+    while !pidfds.is_empty() {
+        let to_remove = {
+            let mut fds: Vec<PollFd> = pidfds
+                .iter()
+                .map(|fd| PollFd::new(fd, PollFlags::IN))
+                .collect();
+
+            if poll(&mut fds, None).is_err() {
+                break;
             }
-        }
 
-        if list.is_empty() {
-            return;
-        }
+            fds.iter()
+                .enumerate()
+                .filter(|(_, pfd)| pfd.revents().contains(PollFlags::IN))
+                .map(|(i, _)| i)
+                .rev()
+                .collect::<Vec<_>>()
+        };
 
-        sleep(Duration::from_millis(50));
-    }
-}
-#[cfg(target_os = "linux")]
-fn is_running(pid: usize) -> bool {
-    use uu_pgrep::process::RunState;
-
-    match ProcessInformation::from_pid(pid) {
-        Ok(mut proc) => proc
-            .run_state()
-            .map(|it| it != RunState::Stopped)
-            .unwrap_or(false),
-        Err(_) => false,
+        for i in to_remove {
+            pidfds.remove(i);
+        }
     }
 }
 
-// Just for passing compile on other system.
 #[cfg(not(target_os = "linux"))]
 pub(crate) fn wait(_procs: &[ProcessInformation]) {}
+
+#[cfg(test)]
+mod tests {
+
+    #[cfg(target_os = "linux")]
+    #[test]
+    fn test_wait_single_process() {
+        use super::*;
+        use std::process::Command;
+        use std::time::Instant;
+
+        // NOTE: Manually tested with sleep 0.5, sleep 1, and sleep 2. Using 1s here to keep total
+        // test time reasonable; 2s would also pass.
+        let mut child = Command::new("sleep").arg("1").spawn().unwrap();
+        let pid = child.id() as usize;
+
+        let info = ProcessInformation::from_pid(pid).unwrap();
+        let start = Instant::now();
+        wait(&[info]);
+        let elapsed = start.elapsed();
+
+        assert!(
+            elapsed >= std::time::Duration::from_millis(900),
+            "wait returned too early: {elapsed:?}"
+        );
+        assert!(
+            elapsed < std::time::Duration::from_secs(3),
+            "wait took too long: {elapsed:?}"
+        );
+
+        let _ = child.wait();
+    }
+}