Repository: Thomasdezeeuw/a10
Pull request #79 (2023-07-19 20:02:26)

Add ReceiveSignals
Type that combines Signals and signals::Receive to not have to deal with lifetime of the fd.


Patch:
diff --git a/src/op.rs b/src/op.rs
index 5e4c65e..b451b94 100644
--- a/src/op.rs
+++ b/src/op.rs
@@ -92,6 +92,7 @@ impl QueuedOperation {
     /// For multishot operations: if this returns `Poll::Pending` the caller
     /// should check `is_done` to determine if the previous result was the last
     /// one.
+    #[allow(clippy::needless_pass_by_ref_mut)] // Match `Future` API.
     pub(crate) fn poll(&mut self, ctx: &mut task::Context<'_>) -> Poll<io::Result<(u16, i32)>> {
         match &mut self.kind {
             QueuedOperationKind::Single { result } => {
@@ -127,6 +128,7 @@ impl QueuedOperation {
     /// Poll the operation for a message.
     ///
     /// Returns the `flags` and the `result` (always positive).
+    #[allow(clippy::needless_pass_by_ref_mut)] // Match `Future` API.
     pub(crate) fn poll_msg(&mut self, ctx: &mut task::Context<'_>) -> Poll<(u16, u32)> {
         match &mut self.kind {
             QueuedOperationKind::Multishot { results } => {
@@ -269,6 +271,7 @@ impl Submission {
     }
 
     /// Create a read vectored submission starting at `offset`.
+    #[allow(clippy::needless_pass_by_ref_mut)] // Needed for `bufs`.
     pub(crate) unsafe fn read_vectored_at(
         &mut self,
         fd: RawFd,
diff --git a/src/signals.rs b/src/signals.rs
index 0bd6ffd..10d483c 100644
--- a/src/signals.rs
+++ b/src/signals.rs
@@ -3,13 +3,14 @@
 //! See the [`Signals`] documentation.
 
 use std::mem::{size_of, MaybeUninit};
+use std::task::{self, Poll};
 use std::{fmt, io, ptr};
 
 use log::{error, trace};
 
 use crate::libc::{self, syscall};
-use crate::op::{op_future, NO_OFFSET};
-use crate::{AsyncFd, SubmissionQueue};
+use crate::op::{op_future, OpState, NO_OFFSET};
+use crate::{AsyncFd, QueueFull, SubmissionQueue};
 
 /// Notification of process signals.
 ///
@@ -109,6 +110,20 @@ impl Signals {
         let info = Box::new(MaybeUninit::uninit());
         Receive::new(&self.fd, info, ())
     }
+
+    /// Receive multiple signals.
+    ///
+    /// This is an combined, owned version of `Signals` and `Receive` (the
+    /// future behind `Signals::receive`). This is useful if you don't want to
+    /// deal with the `'fd` lifetime.
+    pub fn receive_signals(self) -> ReceiveSignals {
+        ReceiveSignals {
+            signals: self,
+            // TODO: replace with `Box::new_uninit` once `new_uninit` is stable.
+            info: Box::new(MaybeUninit::uninit()),
+            state: OpState::NotStarted(()),
+        }
+    }
 }
 
 /// Create a `sigset_t` from `signals`.
@@ -207,3 +222,71 @@ fn sigprocmask(how: libc::c_int, set: &libc::sigset_t) -> io::Result<()> {
     libc::syscall!(pthread_sigmask(how, set, ptr::null_mut()))?;
     Ok(())
 }
+
+/// Receive multiple signals.
+#[derive(Debug)]
+#[must_use = "`Future`s do nothing unless polled"]
+#[allow(clippy::module_name_repetitions)]
+pub struct ReceiveSignals {
+    signals: Signals,
+    info: Box<MaybeUninit<libc::signalfd_siginfo>>,
+    state: OpState<()>,
+}
+
+impl ReceiveSignals {
+    /// Poll the next signal.
+    pub fn poll_signal<'a>(
+        &'a mut self,
+        ctx: &mut task::Context<'_>,
+    ) -> Poll<Option<io::Result<&'a libc::signalfd_siginfo>>> {
+        let ReceiveSignals {
+            signals,
+            info,
+            state,
+        } = self;
+        let op_index = match state {
+            OpState::Running(op_index) => *op_index,
+            OpState::NotStarted(()) => {
+                let result = signals.fd.sq.add(|submission| unsafe {
+                    submission.read_at(
+                        signals.fd.fd(),
+                        (**info).as_mut_ptr().cast(),
+                        size_of::<libc::signalfd_siginfo>() as u32,
+                        NO_OFFSET,
+                    );
+                });
+                match result {
+                    Ok(op_index) => {
+                        *state = OpState::Running(op_index);
+                        op_index
+                    }
+                    Err(QueueFull(())) => {
+                        signals.fd.sq.wait_for_submission(ctx.waker().clone());
+                        return Poll::Pending;
+                    }
+                }
+            }
+            OpState::Done => return Poll::Ready(None),
+        };
+
+        match signals.fd.sq.poll_op(ctx, op_index) {
+            Poll::Ready(Ok((_, n))) => {
+                // Reset the state so that we start reading another signal in
+                // the next call.
+                *state = OpState::NotStarted(());
+                #[allow(clippy::cast_sign_loss)] // Negative values are mapped to errors.
+                {
+                    debug_assert_eq!(n as usize, size_of::<libc::signalfd_siginfo>());
+                }
+                // SAFETY: the kernel initialised the info allocation for us as
+                // part of the read call.
+                Poll::Ready(Some(Ok(unsafe { MaybeUninit::assume_init_ref(&*info) })))
+            }
+            Poll::Ready(Err(err)) => {
+                *state = OpState::Done; // Consider the error as fatal.
+                Poll::Ready(Some(Err(err)))
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}


Test patch:
diff --git a/tests/signals.rs b/tests/signals.rs
index 882c92a..08fb6ca 100644
--- a/tests/signals.rs
+++ b/tests/signals.rs
@@ -5,13 +5,13 @@ use std::pin::Pin;
 use std::sync::{Arc, Barrier};
 use std::task::Poll;
 use std::time::Instant;
-use std::{env, fmt, io, panic, process, ptr, thread};
+use std::{env, fmt, io, panic, process, ptr, task, thread};
 
 use a10::signals::{self, Signals};
 use a10::Ring;
 
 mod util;
-use util::{is_send, is_sync, poll_nop, syscall};
+use util::{is_send, is_sync, poll_nop, syscall, NOP_WAKER};
 
 const SIGNALS: [libc::c_int; 31] = [
     libc::SIGHUP,
@@ -87,17 +87,20 @@ const SIGNAL_NAMES: [&str; SIGNALS.len()] = [
 
 fn main() {
     let start = Instant::now();
-    println!("\nrunning {} tests", (2 * SIGNALS.len()) + 1);
+    println!("\nrunning {} tests", (3 * SIGNALS.len()) + 1);
 
     is_send::<Signals>();
     is_sync::<Signals>();
     is_send::<signals::Receive>();
     is_sync::<signals::Receive>();
+    is_send::<signals::ReceiveSignals>();
+    is_sync::<signals::ReceiveSignals>();
 
     let quiet = env::args().any(|arg| matches!(&*arg, "-q" | "--quiet"));
     let mut harness = TestHarness::setup(quiet);
     harness.test_single_threaded();
     harness.test_multi_threaded();
+    harness.test_receive_signals(); // MUST be last as it takes `Signals`.
 
     let mut passed = harness.passed;
     let mut failed = harness.failed;
@@ -111,7 +114,7 @@ fn main() {
 
 struct TestHarness {
     ring: Ring,
-    signals: Signals,
+    signals: Option<Signals>,
     passed: usize,
     failed: usize,
     quiet: bool,
@@ -121,10 +124,9 @@ impl TestHarness {
     fn setup(quiet: bool) -> TestHarness {
         let ring = Ring::new(2).unwrap();
         let sq = ring.submission_queue().clone();
-        let signals = Signals::from_signals(sq, SIGNALS).unwrap();
         TestHarness {
             ring,
-            signals,
+            signals: Some(Signals::from_signals(sq, SIGNALS).unwrap()),
             passed: 0,
             failed: 0,
             quiet,
@@ -133,11 +135,12 @@ impl TestHarness {
 
     fn test_single_threaded(&mut self) {
         let pid = process::id();
+        let signals = self.signals.as_ref().unwrap();
         for (signal, name) in SIGNALS.into_iter().zip(SIGNAL_NAMES) {
             print_test_start(self.quiet, format_args!("single_threaded ({name})"));
             let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
                 send_signal(pid, signal).unwrap();
-                receive_signal(&mut self.ring, &self.signals, signal);
+                receive_signal(&mut self.ring, &signals, signal);
             }));
             if res.is_ok() {
                 print_test_ok(self.quiet);
@@ -167,10 +170,11 @@ impl TestHarness {
             }
         });
 
+        let signals = self.signals.as_ref().unwrap();
         for (signal, name) in SIGNALS.into_iter().zip(SIGNAL_NAMES) {
             print_test_start(self.quiet, format_args!("multi_threaded ({name})"));
             let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
-                receive_signal(&mut self.ring, &self.signals, signal);
+                receive_signal(&mut self.ring, signals, signal);
             }));
             if res.is_ok() {
                 print_test_ok(self.quiet);
@@ -184,6 +188,35 @@ impl TestHarness {
 
         handle.join().unwrap();
     }
+
+    fn test_receive_signals(&mut self) {
+        let pid = process::id();
+        let mut receive_signal = self.signals.take().unwrap().receive_signals();
+        let task_waker = unsafe { task::Waker::from_raw(NOP_WAKER) };
+        let mut task_ctx = task::Context::from_waker(&task_waker);
+        for (signal, name) in SIGNALS.into_iter().zip(SIGNAL_NAMES) {
+            print_test_start(self.quiet, format_args!("single_threaded ({name})"));
+            let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                send_signal(pid, signal).unwrap();
+
+                // Check if the signals can be received.
+                let signal_info = loop {
+                    match receive_signal.poll_signal(&mut task_ctx) {
+                        Poll::Ready(result) => break result.unwrap().unwrap(),
+                        Poll::Pending => self.ring.poll(None).unwrap(),
+                    }
+                };
+                assert_eq!(signal_info.ssi_signo as libc::c_int, signal);
+            }));
+            if res.is_ok() {
+                print_test_ok(self.quiet);
+                self.passed += 1
+            } else {
+                print_test_failed(self.quiet);
+                self.failed += 1
+            };
+        }
+    }
 }
 
 fn test_cleanup(quiet: bool, passed: &mut usize, failed: &mut usize) {
diff --git a/tests/util/mod.rs b/tests/util/mod.rs
index ce00489..ddd7a0b 100644
--- a/tests/util/mod.rs
+++ b/tests/util/mod.rs
@@ -201,18 +201,22 @@ impl Waker {
     }
 }
 
+const NOP_WAKER_VTABLE: task::RawWakerVTable = task::RawWakerVTable::new(
+    |_| task::RawWaker::new(ptr::null(), &NOP_WAKER_VTABLE), // clone.
+    |_| {},                                                  // wake.
+    |_| {},                                                  // wake_by_ref.
+    |_| {},                                                  // drop.
+);
+
+/// No-op waker used by [`poll_nop`].
+pub(crate) const NOP_WAKER: task::RawWaker = task::RawWaker::new(ptr::null(), &NOP_WAKER_VTABLE);
+
 /// Poll the `future` once with a no-op waker.
 pub(crate) fn poll_nop<Fut>(future: Pin<&mut Fut>) -> Poll<Fut::Output>
 where
     Fut: Future,
 {
-    const NOP_WAKER: task::RawWakerVTable = task::RawWakerVTable::new(
-        |_| task::RawWaker::new(ptr::null(), &NOP_WAKER), // clone.
-        |_| {},                                           // wake.
-        |_| {},                                           // wake_by_ref.
-        |_| {},                                           // drop.
-    );
-    let task_waker = unsafe { task::Waker::from_raw(task::RawWaker::new(ptr::null(), &NOP_WAKER)) };
+    let task_waker = unsafe { task::Waker::from_raw(NOP_WAKER) };
     let mut task_ctx = task::Context::from_waker(&task_waker);
     Future::poll(future, &mut task_ctx)
 }