multishot: implement detection of finished

When a multishot IO will no longer yield a value the IORING_CQE_F_MORE
flag will not be set. Detect that here and make all multishot IOs return
`Poll::Ready(None)` in that case to signal that the end of the stream
has been reached.

Author: hexagonal-sun

examples/udp.rs

@@ -24,11 +24,11 @@ fn main() {
     });
 
     Executor::spawn(async {
-        let mut buf = [0xadu8; 20];
+        let buf = [0xadu8; 20];
         let udpsock = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
         Timer::sleep(Duration::from_secs(1)).unwrap().await;
         let len = udpsock
-            .send_to(&mut buf, (Ipv4Addr::LOCALHOST, 9998))
+            .send_to(&buf, (Ipv4Addr::LOCALHOST, 9998))
             .await
             .unwrap();
 

examples/web_server.rs

@@ -134,12 +134,10 @@ async fn handle_connection(mut conn: TcpStream) -> anyhow::Result<()> {
 
     let file = if path == PathBuf::from("/") {
         ARGS.get().unwrap().webroot.join("index.html")
+    } else if let Ok(path) = path.strip_prefix("/") {
+        ARGS.get().unwrap().webroot.join(path)
     } else {
-        if let Ok(path) = path.strip_prefix("/") {
-            ARGS.get().unwrap().webroot.join(path)
-        } else {
-            return send_response_hdr(&mut conn, Response::NotFound, 0).await;
-        }
+        return send_response_hdr(&mut conn, Response::NotFound, 0).await;
     };
 
     send_file(conn, file).await

src/futures/tcp.rs

@@ -42,7 +42,7 @@ use io_uring::{opcode, types};
 use libc::{AF_INET, AF_INET6, SOCK_STREAM};
 use tokio_stream::Stream;
 
-use crate::reactor::{Reactor, ReactorIo};
+use crate::reactor::{MultishotReactorIo, Reactor, ReactorIo};
 
 use super::{
     read::{AsyncRead, AsyncReader},
@@ -56,7 +56,7 @@ use super::{
 /// accept connections on the specified address.
 pub struct TcpListener {
     inner: OwnedFd,
-    io: ReactorIo,
+    io: MultishotReactorIo,
 }
 
 fn mk_sock(addr: &SocketAddr) -> std::io::Result<OwnedFd> {
@@ -127,7 +127,7 @@ impl Stream for TcpListener {
                 )
             })
             .map(|x| {
-                Some({
+                x.map(|x| {
                     x.map(|fd| TcpStream {
                         inner: unsafe { OwnedFd::from_raw_fd(fd) },
                     })

src/reactor/mod.rs

@@ -1,9 +1,11 @@
-use std::{mem::transmute, task::Waker};
-use uring::{ReactorUring, UringIo};
+use std::task::Waker;
+
+use uring::{MultishotUringIo, OneshotUringIo, ReactorUring};
 
 mod uring;
 
-pub type ReactorIo = UringIo<'static, Waker>;
+pub type ReactorIo = OneshotUringIo<Waker>;
+pub type MultishotReactorIo = MultishotUringIo<Waker>;
 
 pub(crate) struct Reactor {}
 
@@ -13,11 +15,11 @@ thread_local! {
 
 impl Reactor {
     pub fn new_io() -> ReactorIo {
-        REACTOR.with(|r| unsafe { transmute(r.new_io()) })
+        REACTOR.with(|r| r.new_oneshot_io())
     }
 
-    pub fn new_multishot_io() -> ReactorIo {
-        REACTOR.with(|r| unsafe { transmute(r.new_multishot_io()) })
+    pub fn new_multishot_io() -> MultishotReactorIo {
+        REACTOR.with(|r| r.new_multishot_io())
     }
 
     pub fn react() {

src/reactor/uring.rs

@@ -1,29 +1,32 @@
-pub(crate) use io::UringIo;
-use io_uring::{squeue, CompletionQueue, IoUring};
+pub(crate) use io::{multishot::MultishotUringIo, oneshot::OneshotUringIo};
+use io_uring::{cqueue, squeue, CompletionQueue, IoUring};
 use result::RingResults;
 use slab::Slab;
-use std::cell::{RefCell, RefMut};
+use std::{
+    cell::{RefCell, RefMut},
+    rc::Rc,
+};
 
 mod io;
 mod result;
 
 pub struct ReactorUring<T: Clone> {
-    inner: RefCell<ReactorInner<T>>,
+    inner: Rc<RefCell<ReactorInner<T>>>,
 }
 
 impl<T: Clone> ReactorUring<T> {
     pub fn new() -> Self {
         Self {
-            inner: RefCell::new(ReactorInner::new()),
+            inner: Rc::new(RefCell::new(ReactorInner::new())),
         }
     }
 
-    pub fn new_io(&self) -> UringIo<'_, T> {
-        UringIo::new(&self.inner, IoKind::Oneshot)
+    pub fn new_oneshot_io(&self) -> OneshotUringIo<T> {
+        OneshotUringIo::new(self.inner.clone())
     }
 
-    pub fn new_multishot_io(&self) -> UringIo<'_, T> {
-        UringIo::new(&self.inner, IoKind::Multi)
+    pub fn new_multishot_io(&self) -> MultishotUringIo<T> {
+        MultishotUringIo::new(self.inner.clone())
     }
 
     pub fn react(&self) -> IoCompletionIter<'_, T> {
@@ -47,7 +50,7 @@ impl<T: Clone> ReactorUring<T> {
     }
 }
 
-struct ReactorInner<T> {
+pub(crate) struct ReactorInner<T> {
     uring: IoUring,
     pending: Slab<PendingIo<T>>,
     results: RingResults,
@@ -75,8 +78,11 @@ impl<T> ReactorInner<T> {
         }
     }
 
-    fn submit_io(&mut self, entry: squeue::Entry, obj: T, kind: IoKind) -> usize {
-        let result_slab_idx = self.results.get(kind).create_slot();
+    fn submit_io(&mut self, entry: squeue::Entry, obj: T, kind: IoKind) -> (u64, usize) {
+        let result_slab_idx = match kind {
+            IoKind::Oneshot => self.results.get_oneshot().create_slot(),
+            IoKind::Multi => self.results.get_multishot().create_slot(),
+        };
 
         let slot = self.pending.insert(PendingIo {
             assoc_obj: obj,
@@ -91,7 +97,7 @@ impl<T> ReactorInner<T> {
                 .unwrap();
         }
 
-        result_slab_idx
+        (slot as u64, result_slab_idx)
     }
 }
 
@@ -113,13 +119,21 @@ impl<T: Clone> Iterator for IoCompletionIter<'_, T> {
             .unwrap()
             .clone();
 
-        self.ring
-            .results
-            .get(pending_io.kind)
-            .set_result(entry.result(), pending_io.result_slab_idx);
-
-        if let IoKind::Oneshot = pending_io.kind {
-            self.ring.pending.remove(entry.user_data() as usize);
+        match pending_io.kind {
+            IoKind::Oneshot => {
+                self.ring
+                    .results
+                    .get_oneshot()
+                    .set_result(entry.result(), pending_io.result_slab_idx);
+                self.ring.pending.remove(entry.user_data() as usize);
+            }
+            IoKind::Multi => {
+                let results = self.ring.results.get_multishot();
+                results.push_result(entry.result(), pending_io.result_slab_idx);
+                if !cqueue::more(entry.flags()) {
+                    results.set_finished(pending_io.result_slab_idx);
+                }
+            }
         }
 
         Some(pending_io.assoc_obj)
@@ -189,7 +203,7 @@ mod tests {
         run_test(|a, b, uring| {
             let mut buf = [0];
 
-            let mut io = uring.new_io();
+            let mut io = uring.new_oneshot_io();
             let result = io.submit_or_get_result(|| {
                 (
                     opcode::Read::new(types::Fd(a.as_raw_fd()), buf.as_mut_ptr(), 1).build(),
@@ -224,7 +238,7 @@ mod tests {
         run_test(|a, b, uring| {
             let mut buf = [0];
 
-            let mut io = uring.new_io();
+            let mut io = uring.new_oneshot_io();
             assert!(matches!(
                 io.submit_or_get_result(|| {
                     (
@@ -255,7 +269,7 @@ mod tests {
         run_test(|a, b, uring| {
             let buf = [0];
 
-            let mut io = uring.new_io();
+            let mut io = uring.new_oneshot_io();
             let result = io.submit_or_get_result(|| {
                 (
                     opcode::Write::new(types::Fd(a.as_raw_fd()), buf.as_ptr(), buf.len() as _)
@@ -293,7 +307,7 @@ mod tests {
         run_test(|a, b, uring| {
             let mut buf = [0, 0];
 
-            let mut io1 = uring.new_io();
+            let mut io1 = uring.new_oneshot_io();
             assert!(matches!(
                 io1.submit_or_get_result(|| {
                     (
@@ -304,7 +318,7 @@ mod tests {
                 Poll::Pending
             ));
 
-            let mut io2 = uring.new_io();
+            let mut io2 = uring.new_oneshot_io();
             assert!(matches!(
                 io2.submit_or_get_result(|| {
                     (
@@ -344,7 +358,7 @@ mod tests {
         run_test(|a, b, uring| {
             let buf = [0xbe, 0xef];
 
-            let mut io1 = uring.new_io();
+            let mut io1 = uring.new_oneshot_io();
             assert!(matches!(
                 io1.submit_or_get_result(|| {
                     (
@@ -355,7 +369,7 @@ mod tests {
                 Poll::Pending
             ));
 
-            let mut io2 = uring.new_io();
+            let mut io2 = uring.new_oneshot_io();
             assert!(matches!(
                 io2.submit_or_get_result(|| {
                     (

src/reactor/uring/io.rs

@@ -0,0 +1,10 @@
+pub mod multishot;
+pub mod oneshot;
+
+fn reactor_value_to_result(v: i32) -> std::io::Result<i32> {
+    if v < 0 {
+        Err(std::io::Error::from_raw_os_error(v.abs()))
+    } else {
+        Ok(v)
+    }
+}