Repository: Thomasdezeeuw/a10
Pull request #103 (2024-03-30 13:58:22)

Add process:wait(_on)
Wraps the waitid(2) system call/io_uring operation.

This adds a new process module to hold the functions and related types.


Patch:
diff --git a/Makefile b/Makefile
index 77c34ea..f67e142 100644
--- a/Makefile
+++ b/Makefile
@@ -63,6 +63,7 @@ clippy:
 		--allow clippy::use-self \
 		\
 		--allow clippy::cast-possible-truncation \
+		--allow clippy::cast-possible-wrap \
 
 doc:
 	cargo doc
diff --git a/src/cancel.rs b/src/cancel.rs
index bf33690..020f803 100644
--- a/src/cancel.rs
+++ b/src/cancel.rs
@@ -164,7 +164,7 @@ impl<'fd> Future for CancelOp<'fd> {
             Poll::Ready(result) => {
                 self.state = OpState::Done;
                 match result {
-                    Ok((_, res)) => Poll::Ready(Ok(debug_assert!(res == 0))),
+                    Ok((_, _)) => Poll::Ready(Ok(())),
                     Err(err) => Poll::Ready(Err(err)),
                 }
             }
diff --git a/src/lib.rs b/src/lib.rs
index 5033e12..5454c2d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -152,6 +152,7 @@ pub mod msg;
 pub mod net;
 mod op;
 pub mod poll;
+pub mod process;
 pub mod signals;
 
 // TODO: replace this with definitions from the `libc` crate once available.
@@ -866,7 +867,7 @@ impl SubmissionQueue {
     where
         F: FnOnce(&mut Submission),
     {
-        log::trace!(op_index = op_index.0; "canceling operation as receiver (Future) was dropped before completion");
+        log::trace!(op_index = op_index.0; "canceling operation");
         if let Some(operation) = self.shared.queued_ops.get(op_index.0) {
             let mut operation = operation.lock().unwrap();
             if let Some(op) = &mut *operation {
diff --git a/src/op.rs b/src/op.rs
index 20e07ca..ba0555d 100644
--- a/src/op.rs
+++ b/src/op.rs
@@ -666,6 +666,23 @@ impl Submission {
         };
     }
 
+    #[allow(clippy::cast_possible_wrap)]
+    pub(crate) unsafe fn waitid(
+        &mut self,
+        id: libc::id_t,
+        id_type: libc::idtype_t,
+        options: libc::c_int,
+        info: *mut libc::signalfd_siginfo,
+    ) {
+        self.inner.opcode = libc::IORING_OP_WAITID as u8;
+        self.inner.fd = id as _;
+        self.inner.__bindgen_anon_1 = libc::io_uring_sqe__bindgen_ty_1 { addr2: info as _ };
+        self.inner.len = id_type;
+        self.inner.__bindgen_anon_5 = libc::io_uring_sqe__bindgen_ty_5 {
+            file_index: options as _,
+        };
+    }
+
     pub(crate) unsafe fn wake(&mut self, ring_fd: RawFd) {
         self.msg(ring_fd, u64::MAX, 0, 0);
         self.no_completion_event();
@@ -889,6 +906,15 @@ impl fmt::Debug for Submission {
                     .field("msg1", &self.inner.len)
                     .field("msg2", unsafe { &self.inner.__bindgen_anon_1.off });
             }
+            libc::IORING_OP_WAITID => {
+                f.field("opcode", &"IORING_OP_WAITID")
+                    .field("id", &self.inner.fd)
+                    .field("id_type", &self.inner.len)
+                    .field("options", unsafe {
+                        &self.inner.__bindgen_anon_5.file_index
+                    })
+                    .field("info", unsafe { &self.inner.__bindgen_anon_1.addr2 });
+            }
             _ => {
                 // NOTE: we can't access the unions safely without know what
                 // fields to read.
diff --git a/src/process.rs b/src/process.rs
new file mode 100644
index 0000000..029aa7b
--- /dev/null
+++ b/src/process.rs
@@ -0,0 +1,128 @@
+//! Process handling.
+
+use std::future::Future;
+use std::io;
+use std::mem::MaybeUninit;
+use std::pin::Pin;
+use std::process::Child;
+use std::task::{self, Poll};
+
+use crate::cancel::{Cancel, CancelOp, CancelResult};
+use crate::op::{poll_state, OpState};
+use crate::SubmissionQueue;
+
+/// Wait on the child `process`.
+///
+/// See [`wait`].
+pub fn wait_on(sq: SubmissionQueue, process: &Child, options: libc::c_int) -> WaitId {
+    wait(sq, WaitOn::Process(process.id()), options)
+}
+
+/// Obtain status information on termination, stop, and/or continue events in
+/// one of the caller's child processes.
+///
+/// Also see [`wait_on`] to wait on a [`Child`] process.
+#[doc(alias = "waitid")]
+pub fn wait(sq: SubmissionQueue, wait: WaitOn, options: libc::c_int) -> WaitId {
+    WaitId {
+        sq,
+        state: OpState::NotStarted((wait, options)),
+        info: Some(Box::new(MaybeUninit::uninit())),
+    }
+}
+
+/// Defines on what process (or processes) to wait.
+#[doc(alias = "idtype")]
+#[doc(alias = "idtype_t")]
+#[derive(Copy, Clone, Debug)]
+pub enum WaitOn {
+    /// Wait for the child process.
+    #[doc(alias = "P_PID")]
+    Process(libc::id_t),
+    /// Wait for any child process in the process group with ID.
+    #[doc(alias = "P_PGID")]
+    Group(libc::id_t),
+    /// Wait for all childeren.
+    #[doc(alias = "P_ALL")]
+    All,
+}
+
+/// [`Future`] behind [`wait`].
+#[derive(Debug)]
+#[must_use = "`Future`s do nothing unless polled"]
+pub struct WaitId {
+    sq: SubmissionQueue,
+    /// Buffer to write into, needs to stay in memory so the kernel can
+    /// access it safely.
+    info: Option<Box<MaybeUninit<libc::signalfd_siginfo>>>,
+    state: OpState<(WaitOn, libc::c_int)>,
+}
+
+impl Future for WaitId {
+    type Output = io::Result<Box<libc::siginfo_t>>;
+
+    fn poll(mut self: Pin<&mut Self>, ctx: &mut task::Context<'_>) -> Poll<Self::Output> {
+        let op_index = poll_state!(
+            WaitId,
+            self.state,
+            self.sq,
+            ctx,
+            |submission, (wait, options)| unsafe {
+                let (id_type, pid) = match wait {
+                    WaitOn::Process(pid) => (libc::P_PID, pid),
+                    WaitOn::Group(pid) => (libc::P_PGID, pid),
+                    WaitOn::All => (libc::P_ALL, 0), // NOTE: id is ignored.
+                };
+                let info = self.info.as_ref().unwrap().as_ptr().cast_mut();
+                submission.waitid(pid, id_type, options, info);
+            }
+        );
+
+        match self.sq.poll_op(ctx, op_index) {
+            Poll::Ready(result) => {
+                self.state = OpState::Done;
+                match result {
+                    Ok((_, _)) => Poll::Ready(Ok(unsafe {
+                        Box::from_raw(Box::into_raw(self.info.take().unwrap()).cast())
+                    })),
+                    Err(err) => Poll::Ready(Err(err)),
+                }
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+impl Cancel for WaitId {
+    fn try_cancel(&mut self) -> CancelResult {
+        self.state.try_cancel(&self.sq)
+    }
+
+    fn cancel(&mut self) -> CancelOp {
+        self.state.cancel(&self.sq)
+    }
+}
+
+impl Drop for WaitId {
+    fn drop(&mut self) {
+        if let Some(info) = self.info.take() {
+            match self.state {
+                OpState::Running(op_index) => {
+                    // Only drop the signal `info` field once we know the
+                    // operation has finished, otherwise the kernel might write
+                    // into memory we have deallocated.
+                    let result = self.sq.cancel_op(op_index, info, |submission| unsafe {
+                        submission.cancel_op(op_index);
+                        // We'll get a canceled completion event if we succeeded, which
+                        // is sufficient to cleanup the operation.
+                        submission.no_completion_event();
+                    });
+                    if let Err(err) = result {
+                        log::error!("dropped a10::WaitId before canceling it, attempt to cancel failed: {err}");
+                    }
+                }
+                OpState::NotStarted((_, _)) | OpState::Done => drop(info),
+            }
+        }
+    }
+}


Test patch:
diff --git a/tests/ring.rs b/tests/ring.rs
index 984337e..da3e953 100644
--- a/tests/ring.rs
+++ b/tests/ring.rs
@@ -9,6 +9,7 @@ use std::io::{self, Read, Write};
 use std::mem::take;
 use std::os::fd::{AsFd, FromRawFd, RawFd};
 use std::pin::{pin, Pin};
+use std::process::Command;
 use std::sync::{Arc, Mutex};
 use std::task::{self, Poll, Wake};
 use std::thread;
@@ -18,7 +19,7 @@ use a10::cancel::Cancel;
 use a10::fs::OpenOptions;
 use a10::msg::{MsgListener, MsgToken, SendMsg};
 use a10::poll::{MultishotPoll, OneshotPoll};
-use a10::{mem, AsyncFd, Config, Ring, SubmissionQueue};
+use a10::{mem, process, AsyncFd, Config, Ring, SubmissionQueue};
 
 mod util;
 use util::{
@@ -370,3 +371,63 @@ fn madvise() {
         ))
         .expect("failed madvise");
 }
+
+#[test]
+fn process_wait_on() {
+    require_kernel!(6, 7);
+    let sq = test_queue();
+    let waker = Waker::new();
+
+    is_send::<process::WaitId>();
+    is_sync::<process::WaitId>();
+
+    let process = Command::new("true").spawn().unwrap();
+    let pid = process.id();
+
+    let info = waker
+        .block_on(process::wait_on(sq, &process, libc::WEXITED))
+        .expect("failed wait");
+
+    assert_eq!(info.si_signo, libc::SIGCHLD);
+    assert_eq!(info.si_code, libc::CLD_EXITED);
+    // SAFETY: these fields are set if `si_signo == SIGCHLD`.
+    assert_eq!(unsafe { info.si_pid() }, pid as i32);
+    assert_eq!(unsafe { info.si_status() }, libc::EXIT_SUCCESS);
+    assert!(unsafe { info.si_utime() } >= 1);
+    assert!(unsafe { info.si_stime() } >= 1);
+}
+
+#[test]
+fn process_wait_on_cancel() {
+    require_kernel!(6, 7);
+    let sq = test_queue();
+    let waker = Waker::new();
+
+    let mut process = Command::new("sleep").arg("1000").spawn().unwrap();
+
+    let mut future = process::wait_on(sq, &process, libc::WEXITED);
+    let result = poll_nop(Pin::new(&mut future));
+    if !result.is_pending() {
+        panic!("unexpected result, expected it to return Poll::Pending");
+    }
+
+    waker.block_on(future.cancel()).unwrap();
+
+    process.kill().unwrap();
+    process.wait().unwrap();
+}
+
+#[test]
+fn process_wait_on_drop_before_complete() {
+    require_kernel!(6, 7);
+    let sq = test_queue();
+
+    let process = Command::new("sleep").arg("1000").spawn().unwrap();
+
+    let mut future = process::wait_on(sq, &process, libc::WEXITED);
+    let result = poll_nop(Pin::new(&mut future));
+    if !result.is_pending() {
+        panic!("unexpected result, expected it to return Poll::Pending");
+    }
+    drop(future);
+}