pidwait: optimize & cross-platform

src/uu/pidwait/Cargo.toml

@@ -17,6 +17,12 @@ clap = { workspace = true }
 regex = { workspace = true }
 uu_pgrep = { path = "../pgrep" }
 
+[target.'cfg(unix)'.dependencies]
+rustix = { version = "1", default-features = false, features = ["event", "process", "std"] }
+
+[target.'cfg(windows)'.dependencies]
+windows-sys = { workspace = true, features = ["Win32_Foundation", "Win32_System_Threading"] }
+
 [lib]
 path = "src/pidwait.rs"
 

src/uu/pidwait/src/pidwait.rs

@@ -4,15 +4,16 @@
 // file that was distributed with this source code.
 
 use clap::{arg, crate_version, Command};
+use std::time::Duration;
+use uu_pgrep::process::ProcessInformation;
 use uu_pgrep::process_matcher;
 use uucore::{error::UResult, format_usage, help_about, help_usage};
-use wait::wait;
-
-mod wait;
 
 const ABOUT: &str = help_about!("pidwait.md");
 const USAGE: &str = help_usage!("pidwait.md");
 
+mod platform;
+
 #[uucore::main]
 pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     let matches = uu_app().try_get_matches_from(args)?;
@@ -42,11 +43,23 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
         }
     }
 
-    wait(&proc_infos);
+    // It should be fine to reserve a `timeout` parameter for future use.
+    wait(&proc_infos, None)?;
 
     Ok(())
 }
 
+pub(crate) fn wait(
+    procs: &[ProcessInformation],
+    timeout: Option<Duration>,
+) -> Result<Option<()>, std::io::Error> {
+    if !procs.is_empty() {
+        platform::wait(procs, timeout)
+    } else {
+        Ok(None)
+    }
+}
+
 pub fn uu_app() -> Command {
     Command::new(env!("CARGO_PKG_NAME"))
         .version(crate_version!())

src/uu/pidwait/src/platform/bsd.rs

@@ -0,0 +1,51 @@
+// This file is part of the uutils procps package.
+//
+// For the full copyright and license information, please view the LICENSE
+// file that was distributed with this source code.
+
+// Reference: pidwait-any crate.
+// Thanks to @oxalica's implementation.
+
+// FIXME: Test this implementation
+
+use rustix::event::kqueue::{kevent, kqueue, Event, EventFilter, EventFlags, ProcessEvents};
+use rustix::process::Pid;
+use std::io::{Error, ErrorKind, Result};
+use std::mem::MaybeUninit;
+use std::time::Duration;
+use uu_pgrep::process::ProcessInformation;
+
+pub fn wait(procs: &[ProcessInformation], timeout: Option<Duration>) -> Result<Option<()>> {
+    let mut events = Vec::with_capacity(procs.len());
+    let kqueue = kqueue()?;
+    for proc in procs {
+        let pid = Pid::from_raw(proc.pid as i32).ok_or_else(|| {
+            Error::new(
+                ErrorKind::InvalidInput,
+                format!("Invalid PID: {}", proc.pid),
+            )
+        })?;
+        let event = Event::new(
+            EventFilter::Proc {
+                pid,
+                flags: ProcessEvents::EXIT,
+            },
+            EventFlags::ADD,
+            std::ptr::null_mut(),
+        );
+        events.push(event);
+    }
+    let ret = unsafe { kevent::<_, &mut [Event; 0]>(&kqueue, &events, &mut [], None)? };
+    debug_assert_eq!(ret, 0);
+    let mut buf = [MaybeUninit::uninit()];
+    let (events, _rest_buf) = unsafe { kevent(&kqueue, &[], &mut buf, timeout)? };
+    if events.is_empty() {
+        return Ok(None);
+    };
+    debug_assert!(matches!(
+        events[0].filter(),
+        EventFilter::Proc { flags, .. }
+        if flags.contains(ProcessEvents::EXIT)
+    ));
+    Ok(Some(()))
+}

src/uu/pidwait/src/platform/linux.rs

@@ -0,0 +1,44 @@
+// This file is part of the uutils procps package.
+//
+// For the full copyright and license information, please view the LICENSE
+// file that was distributed with this source code.
+
+// Reference: pidwait-any crate.
+// Thanks to @oxalica's implementation.
+
+use std::io::{Error, ErrorKind};
+use std::os::fd::OwnedFd;
+
+use rustix::event::{poll, PollFd, PollFlags};
+use rustix::io::Errno;
+use rustix::process::{pidfd_open, Pid, PidfdFlags};
+use std::io::Result;
+use std::time::Duration;
+use uu_pgrep::process::ProcessInformation;
+
+pub fn wait(procs: &[ProcessInformation], timeout: Option<Duration>) -> Result<Option<()>> {
+    let mut pidfds: Vec<OwnedFd> = Vec::with_capacity(procs.len());
+    for proc in procs {
+        let pid = Pid::from_raw(proc.pid as i32).ok_or_else(|| {
+            Error::new(
+                ErrorKind::InvalidInput,
+                format!("Invalid PID: {}", proc.pid),
+            )
+        })?;
+        let pidfd = pidfd_open(pid, PidfdFlags::empty())?;
+        pidfds.push(pidfd);
+    }
+    let timespec = timeout
+        .map(|t| t.try_into().map_err(|_| Errno::INVAL))
+        .transpose()?;
+    let mut fds: Vec<PollFd> = Vec::with_capacity(pidfds.len());
+    for pidfd in &pidfds {
+        fds.push(PollFd::new(pidfd, PollFlags::IN));
+    }
+    let ret = poll(&mut fds, timespec.as_ref())?;
+    if ret == 0 {
+        return Ok(None);
+    }
+    debug_assert!(fds[0].revents().contains(PollFlags::IN));
+    Ok(Some(()))
+}

src/uu/pidwait/src/platform/mod.rs

@@ -0,0 +1,27 @@
+#[cfg(target_os = "linux")]
+pub use self::linux::wait;
+
+#[cfg(windows)]
+pub use self::windows::wait;
+
+#[cfg(any(
+    target_os = "freebsd",
+    target_os = "macos",
+    target_os = "netbsd",
+    target_os = "openbsd",
+))]
+pub use self::bsd::wait;
+
+#[cfg(target_os = "linux")]
+mod linux;
+
+#[cfg(windows)]
+mod windows;
+
+#[cfg(any(
+    target_os = "freebsd",
+    target_os = "macos",
+    target_os = "netbsd",
+    target_os = "openbsd",
+))]
+mod bsd;

src/uu/pidwait/src/platform/windows.rs

@@ -0,0 +1,68 @@
+// This file is part of the uutils procps package.
+//
+// For the full copyright and license information, please view the LICENSE
+// file that was distributed with this source code.
+
+// Reference: pidwait-any crate.
+// Thanks to @oxalica's implementation.
+
+use std::time::Duration;
+use uu_pgrep::process::ProcessInformation;
+
+use std::ffi::c_void;
+use std::io::{Error, Result};
+use std::ptr::NonNull;
+
+use windows_sys::Win32::Foundation::{CloseHandle, WAIT_OBJECT_0, WAIT_TIMEOUT};
+use windows_sys::Win32::System::Threading::{
+    OpenProcess, WaitForMultipleObjects, INFINITE, PROCESS_SYNCHRONIZE,
+};
+
+struct HandleWrapper(NonNull<c_void>);
+unsafe impl Send for HandleWrapper {}
+impl Drop for HandleWrapper {
+    fn drop(&mut self) {
+        unsafe {
+            CloseHandle(self.0.as_ptr());
+        };
+    }
+}
+
+pub fn wait(procs: &[ProcessInformation], timeout: Option<Duration>) -> Result<Option<()>> {
+    let hprocess = unsafe {
+        let mut result = Vec::with_capacity(procs.len());
+        for proc in procs {
+            let handle = OpenProcess(PROCESS_SYNCHRONIZE, 0, proc.pid as u32);
+            result.push(HandleWrapper(
+                NonNull::new(handle).ok_or_else(Error::last_os_error)?,
+            ));
+        }
+        result
+    };
+    const _: [(); 1] = [(); (INFINITE == u32::MAX) as usize];
+    let timeout = match timeout {
+        Some(timeout) => timeout
+            .as_millis()
+            .try_into()
+            .unwrap_or(INFINITE - 1)
+            .min(INFINITE - 1),
+        None => INFINITE,
+    };
+    let ret = unsafe {
+        WaitForMultipleObjects(
+            hprocess.len() as u32,
+            hprocess
+                .into_iter()
+                .map(|proc| proc.0.as_ptr())
+                .collect::<Vec<_>>()
+                .as_ptr(),
+            1,
+            timeout,
+        )
+    };
+    match ret {
+        WAIT_OBJECT_0 => Ok(Some(())),
+        WAIT_TIMEOUT => Ok(None),
+        _ => Err(Error::last_os_error()),
+    }
+}