Repository: Thomasdezeeuw/a10
Pull request #76 (2023-07-16 13:19:37)

Add user space messaging
This add three new functions to SubmissionQueue:
 * receive_msg: creates a MsgListener to listen for messages.
 * try_send_msg: send a message to the MsgListener.
 * send_msg: same as try_send_msg, but retries if the submission queue is full (returns a Future).

It also adds a new module called msg, with the following types:
 * MsgListener: listener returned by SubmissionQueue::receive_msg.
 * SendMsg: future behind SubmissionQueue::send_msg.
 * MsgToken: token used to send a message to a specific listener.


Patch:
diff --git a/src/lib.rs b/src/lib.rs
index 1082316..92eae73 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -148,6 +148,7 @@ pub mod extract;
 pub mod fs;
 pub mod io;
 pub mod mem;
+pub mod msg;
 pub mod net;
 mod op;
 pub mod poll;
@@ -162,6 +163,7 @@ use config::munmap;
 pub use config::Config;
 #[doc(no_inline)]
 pub use extract::Extract;
+use msg::{MsgListener, MsgToken, SendMsg};
 use op::{QueuedOperation, Submission};
 use poll::{MultishotPoll, OneshotPoll};
 
@@ -465,6 +467,56 @@ impl SubmissionQueue {
         });
     }
 
+    /// Setup a listener for user space messages.
+    ///
+    /// The returned [`MsgListener`] iterator will return all messages send
+    /// using [`SubmissionQueue::try_send_msg`] and
+    /// [`SubmissionQueue::send_msg`] using the returned `MsgToken`.
+    ///
+    /// # Notes
+    ///
+    /// This will return an error if too many operations are already queued,
+    /// this is usually resolved by calling [`Ring::poll`].
+    ///
+    /// The returned `MsgToken` has an implicitly lifetime linked to
+    /// `MsgListener`. If `MsgListener` is dropped the `MsgToken` will
+    /// become invalid.
+    ///
+    /// Due to the limitations mentioned above it's advised to consider the
+    /// usefulness of the type severly limited. The returned `MsgListener`
+    /// iterator should live for the entire lifetime of the `Ring`, to ensure we
+    /// don't use `MsgToken` after it became invalid. Furthermore to ensure
+    /// the creation of it succeeds it should be done early in the lifetime of
+    /// `Ring`.
+    pub fn msg_listener(self) -> io::Result<(MsgListener, MsgToken)> {
+        MsgListener::new(self)
+    }
+
+    /// Try to send a message to iterator listening for message using `MsgToken`.
+    ///
+    /// This will use the io_uring submission queue to share `data` with the
+    /// receiving end. This means that it will wake up the thread if it's
+    /// currently [polling].
+    ///
+    /// This will fail if the submission queue is currently full. See
+    /// [`SubmissionQueue::send_msg`] for a version that tries again when the
+    /// submission queue is full.
+    ///
+    /// See [`SubmissionQueue::msg_listener`] for examples.
+    ///
+    /// [polling]: Ring::poll
+    pub fn try_send_msg(&self, token: MsgToken, data: u32) -> io::Result<()> {
+        self.add_no_result(|submission| unsafe {
+            submission.msg(self.shared.ring_fd.as_raw_fd(), (token.0).0 as u64, data, 0);
+        })?;
+        Ok(())
+    }
+
+    /// Send a message to iterator listening for message using `MsgToken`.
+    pub const fn send_msg<'a>(&'a self, token: MsgToken, data: u32) -> SendMsg<'a> {
+        SendMsg::new(self, token, data)
+    }
+
     /// Wait for an event specified in `mask` on the file descriptor `fd`.
     ///
     /// Ths is similar to calling `poll(2)` the file descriptor.
@@ -550,6 +602,34 @@ impl SubmissionQueue {
         }
     }
 
+    /// Queue a new operation without making a submission.
+    fn queue_multishot(&self) -> Result<OpIndex, QueueFull> {
+        self._queue(QueuedOperation::new_multishot)
+    }
+
+    /// See [`queue_multishot`].
+    fn _queue<O>(&self, new_op: O) -> Result<OpIndex, QueueFull>
+    where
+        O: FnOnce() -> QueuedOperation,
+    {
+        // Get an index to the queued operation queue.
+        let shared = &*self.shared;
+        let Some(op_index) = shared.op_indices.next_available() else {
+            return Err(QueueFull(()));
+        };
+
+        let queued_op = new_op();
+        // SAFETY: the `AtomicBitMap` always returns valid indices for
+        // `op_queue` (it's the whole point of it).
+        let old_queued_op = replace(
+            &mut *shared.queued_ops[op_index].lock().unwrap(),
+            Some(queued_op),
+        );
+        debug_assert!(old_queued_op.is_none());
+
+        Ok(OpIndex(op_index))
+    }
+
     /// Same as [`SubmissionQueue::add`], but ignores the result.
     #[allow(clippy::mutex_integer)] // For `array_index`, need to the lock for more.
     fn add_no_result<F>(&self, submit: F) -> Result<(), QueueFull>
diff --git a/src/msg.rs b/src/msg.rs
new file mode 100644
index 0000000..536568f
--- /dev/null
+++ b/src/msg.rs
@@ -0,0 +1,95 @@
+//! User space messages.
+//!
+//! To setup a [`MsgListener`] use [`SubmissionQueue::msg_listener`]. It returns
+//! the listener as well as a [`MsgToken`], which can be used in
+//! [`SubmissionQueue::try_send_msg`] and [`SubmissionQueue::send_msg`] to send
+//! a message to the created `MsgListener`.
+
+use std::future::Future;
+use std::io;
+use std::pin::Pin;
+use std::task::{self, Poll};
+
+use crate::{OpIndex, SubmissionQueue};
+
+/// Token used to the messages.
+///
+/// See [`SubmissionQueue::msg_listener`].
+#[derive(Copy, Clone, Debug)]
+#[allow(clippy::module_name_repetitions)]
+pub struct MsgToken(pub(crate) OpIndex);
+
+/// [`AsyncIterator`] behind [`SubmissionQueue::msg_listener`].
+///
+/// [`AsyncIterator`]: std::async_iter::AsyncIterator
+#[derive(Debug)]
+#[must_use = "`Future`s do nothing unless polled"]
+#[allow(clippy::module_name_repetitions)]
+pub struct MsgListener {
+    sq: SubmissionQueue,
+    op_index: OpIndex,
+}
+
+impl MsgListener {
+    /// Create a new `MsgListener`.
+    pub(crate) fn new(sq: SubmissionQueue) -> io::Result<(MsgListener, MsgToken)> {
+        let op_index = sq.queue_multishot()?;
+        Ok((MsgListener { sq, op_index }, MsgToken(op_index)))
+    }
+
+    /// This is the same as the `AsyncIterator::poll_next` function, but then
+    /// available on stable Rust.
+    pub fn poll_next(self: Pin<&mut Self>, ctx: &mut task::Context<'_>) -> Poll<Option<u32>> {
+        log::trace!(op_index = self.op_index.0; "polling multishot messages");
+        if let Some(operation) = self.sq.shared.queued_ops.get(self.op_index.0) {
+            let mut operation = operation.lock().unwrap();
+            if let Some(op) = &mut *operation {
+                return match op.poll_msg(ctx) {
+                    Poll::Ready(data) => Poll::Ready(Some(data.1)),
+                    Poll::Pending => Poll::Pending,
+                };
+            }
+        }
+        panic!("a10::MsgListener called incorrectly");
+    }
+}
+
+#[cfg(feature = "nightly")]
+impl std::async_iter::AsyncIterator for MsgListener {
+    type Item = u32;
+
+    fn poll_next(self: Pin<&mut Self>, ctx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
+        self.poll_next(ctx)
+    }
+}
+
+/// [`Future`] behind [`SubmissionQueue::send_msg`].
+#[derive(Debug)]
+#[must_use = "`Future`s do nothing unless polled"]
+#[allow(clippy::module_name_repetitions)]
+pub struct SendMsg<'a> {
+    sq: &'a SubmissionQueue,
+    token: MsgToken,
+    data: u32,
+}
+
+impl<'a> SendMsg<'a> {
+    /// Create a new `SendMsg`.
+    pub(crate) const fn new(sq: &'a SubmissionQueue, token: MsgToken, data: u32) -> SendMsg {
+        SendMsg { sq, token, data }
+    }
+}
+
+impl<'a> Future for SendMsg<'a> {
+    type Output = io::Result<()>;
+
+    fn poll(self: Pin<&mut Self>, ctx: &mut task::Context<'_>) -> Poll<Self::Output> {
+        match self.sq.try_send_msg(self.token, self.data) {
+            Ok(()) => Poll::Ready(Ok(())),
+            Err(_) => {
+                self.sq.wait_for_submission(ctx.waker().clone());
+                Poll::Pending
+            }
+        }
+    }
+}
diff --git a/src/op.rs b/src/op.rs
index 374f74f..5e4c65e 100644
--- a/src/op.rs
+++ b/src/op.rs
@@ -124,6 +124,31 @@ impl QueuedOperation {
         Poll::Pending
     }
 
+    /// Poll the operation for a message.
+    ///
+    /// Returns the `flags` and the `result` (always positive).
+    pub(crate) fn poll_msg(&mut self, ctx: &mut task::Context<'_>) -> Poll<(u16, u32)> {
+        match &mut self.kind {
+            QueuedOperationKind::Multishot { results } => {
+                if !results.is_empty() {
+                    let completion = results.remove(0);
+                    let data = u32::from_ne_bytes(completion.result.to_ne_bytes());
+                    return Poll::Ready((completion.flags, data));
+                }
+            }
+            QueuedOperationKind::Single { .. } => {
+                panic!("QueuedOperation::poll_msg called incorrectly")
+            }
+        }
+
+        // Still in progress.
+        let waker = ctx.waker();
+        if !matches!(&self.waker, Some(w) if w.will_wake(waker)) {
+            self.waker = Some(waker.clone());
+        }
+        Poll::Pending
+    }
+
     /// Returns true if the operation is done.
     pub(crate) const fn is_done(&self) -> bool {
         self.done


Test patch:
diff --git a/tests/ring.rs b/tests/ring.rs
index 520bfa9..fb0bd0c 100644
--- a/tests/ring.rs
+++ b/tests/ring.rs
@@ -16,11 +16,15 @@ use std::time::{Duration, Instant};
 
 use a10::cancel::Cancel;
 use a10::fs::OpenOptions;
+use a10::msg::{MsgListener, MsgToken, SendMsg};
 use a10::poll::OneshotPoll;
 use a10::{mem, AsyncFd, Config, Ring, SubmissionQueue};
 
 mod util;
-use util::{defer, expect_io_errno, init, is_send, is_sync, next, poll_nop, test_queue, Waker};
+use util::{
+    defer, expect_io_errno, init, is_send, is_sync, next, poll_nop, require_kernel, test_queue,
+    Waker,
+};
 
 const DATA: &[u8] = b"Hello, World!";
 
@@ -185,6 +189,38 @@ fn wake_ring_after_ring_dropped() {
     sq.wake();
 }
 
+#[test]
+fn message_sending() {
+    require_kernel!(5, 18);
+
+    const DATA1: u32 = 123;
+    const DATA2: u32 = u32::MAX;
+
+    let sq = test_queue();
+    let waker = Waker::new();
+
+    is_send::<MsgListener>();
+    is_sync::<MsgListener>();
+    is_send::<SendMsg>();
+    is_sync::<SendMsg>();
+    is_send::<MsgToken>();
+    is_sync::<MsgToken>();
+
+    let (msg_listener, msg_token) = sq.clone().msg_listener().unwrap();
+    let mut msg_listener = pin!(msg_listener);
+
+    assert!(poll_nop(Pin::new(&mut next(msg_listener.as_mut()))).is_pending());
+
+    // Send some messages.
+    sq.try_send_msg(msg_token, DATA1).unwrap();
+    waker.block_on(pin!(sq.send_msg(msg_token, DATA2))).unwrap();
+
+    assert_eq!(waker.block_on(next(msg_listener.as_mut())), Some(DATA1));
+    assert_eq!(waker.block_on(next(msg_listener.as_mut())), Some(DATA2));
+
+    assert!(poll_nop(Pin::new(&mut next(msg_listener.as_mut()))).is_pending());
+}
+
 #[test]
 fn oneshot_poll() {
     let sq = test_queue();
diff --git a/tests/util/mod.rs b/tests/util/mod.rs
index b2f95bc..ce00489 100644
--- a/tests/util/mod.rs
+++ b/tests/util/mod.rs
@@ -289,7 +289,7 @@ macro_rules! op_async_iter {
     ($name: ty => $item: ty) => {
         #[cfg(not(feature = "nightly"))]
         impl AsyncIterator for $name {
-            type Item = io::Result<$item>;
+            type Item = $item;
 
             fn poll_next(
                 self: Pin<&mut Self>,
@@ -301,9 +301,10 @@ macro_rules! op_async_iter {
     };
 }
 
-op_async_iter!(a10::net::MultishotAccept<'_> => AsyncFd);
-op_async_iter!(a10::net::MultishotRecv<'_> => a10::io::ReadBuf);
-op_async_iter!(a10::poll::MultishotPoll<'_> => a10::poll::PollEvent);
+op_async_iter!(a10::msg::MsgListener => u32);
+op_async_iter!(a10::net::MultishotAccept<'_> => io::Result<AsyncFd>);
+op_async_iter!(a10::net::MultishotRecv<'_> => io::Result<a10::io::ReadBuf>);
+op_async_iter!(a10::poll::MultishotPoll<'_> => io::Result<a10::poll::PollEvent>);
 
 /// Return a [`Future`] that return the next item in the `iter` or `None`.
 pub(crate) fn next<I: AsyncIterator>(iter: I) -> Next<I> {