can wait for threads & processes with a timeout now

3 files changed, 137 insertions, 0 deletions

diff --git a/src/wait/child.rs b/src/wait/child.rs
new file mode 100644
index 0000000..29a7d70
--- /dev/null
+++ b/src/wait/child.rs
@@ -0,0 +1,81 @@
+//! based on https://www.man7.org/linux/man-pages/man2/pidfd_open.2.html
+#![cfg(target_os = "linux")]
+
+use std::{
+    io,
+    mem::ManuallyDrop,
+    ops::{Deref, DerefMut},
+    os::fd::{BorrowedFd, RawFd},
+    process::{Child, ExitStatus},
+    ptr,
+};
+
+use libc::{SYS_pidfd_open, syscall};
+use nix::poll::{PollFd, PollFlags};
+
+pub struct ChildWaiter {
+    fd: RawFd,
+    child: Child,
+}
+
+#[derive(Debug)]
+struct PidFdOpenError;
+
+impl std::fmt::Display for PidFdOpenError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "pid_fd_open")
+    }
+}
+
+impl std::error::Error for PidFdOpenError {}
+
+impl ChildWaiter {
+    pub fn new(child: Child) -> io::Result<Self> {
+        let fd = unsafe { syscall(SYS_pidfd_open, child.id(), 0) };
+        if fd < 0 {
+            Err(io::Error::new(io::ErrorKind::Other, PidFdOpenError))
+        } else {
+            let fd = fd as RawFd;
+            Ok(Self { child, fd })
+        }
+    }
+
+    pub fn wait(&mut self, timeout_ms: u16) -> io::Result<Option<ExitStatus>> {
+        let mut poll_fds = [PollFd::new(
+            unsafe { BorrowedFd::borrow_raw(self.fd) },
+            PollFlags::POLLIN,
+        )];
+        let _ = nix::poll::poll(&mut poll_fds, timeout_ms);
+        self.child.try_wait()
+    }
+
+    pub fn into_inner(self) -> Child {
+        unsafe {
+            libc::close(self.fd);
+        }
+        let this = ManuallyDrop::new(self);
+        unsafe { ptr::read(&this.child) }
+    }
+}
+
+impl Deref for ChildWaiter {
+    type Target = Child;
+
+    fn deref(&self) -> &Self::Target {
+        &self.child
+    }
+}
+
+impl DerefMut for ChildWaiter {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.child
+    }
+}
+
+impl Drop for ChildWaiter {
+    fn drop(&mut self) {
+        unsafe {
+            libc::close(self.fd);
+        }
+    }
+}
diff --git a/src/wait/mod.rs b/src/wait/mod.rs
new file mode 100644
index 0000000..63a083b
--- /dev/null
+++ b/src/wait/mod.rs
@@ -0,0 +1,5 @@
+mod child;
+mod thread;
+
+pub use child::*;
+pub use thread::*;
diff --git a/src/wait/thread.rs b/src/wait/thread.rs
new file mode 100644
index 0000000..0eadc3c
--- /dev/null
+++ b/src/wait/thread.rs
@@ -0,0 +1,51 @@
+use std::{
+    sync::mpsc::{Receiver, channel},
+    thread::JoinHandle, time::Duration,
+};
+
+use crate::defer;
+
+pub struct ThreadWaiter<T> {
+    handle: JoinHandle<T>,
+    chan: Receiver<()>,
+    done: bool,
+}
+
+pub fn spawn<T, F>(fun: F) -> ThreadWaiter<T>
+where
+    T: Send + 'static,
+    F: Send + 'static,
+    F: FnOnce() -> T,
+{
+    let (tx, rx) = channel();
+
+    let handle = std::thread::spawn(move || {
+        defer! {
+            let _ = tx.send(());
+        };
+        fun()
+    });
+
+    ThreadWaiter {
+        handle,
+        chan: rx,
+        done: false,
+    }
+}
+
+impl<T> ThreadWaiter<T> {
+    pub fn try_join(&mut self, timeout_ms: u16) -> bool {
+        if self.done {
+            return true;
+        }
+
+        if let Ok(()) = self.chan.recv_timeout(Duration::from_millis(timeout_ms as _)) {
+            self.done = true;
+        }
+
+        self.done
+    }
+    pub fn into_inner(self) -> JoinHandle<T> {
+        self.handle
+    }
+}