driver: improve cqueue dispatch

src/driver/mod.rs

@@ -77,20 +77,30 @@ impl Driver {
     }
 
     pub(crate) fn tick(&mut self) {
-        let mut cq = self.uring.completion();
-        cq.sync();
-
-        for cqe in cq {
-            if cqe.user_data() == u64::MAX {
-                // Result of the cancellation action. There isn't anything we
-                // need to do here. We must wait for the CQE for the operation
-                // that was canceled.
-                continue;
-            }
+        loop {
+            let mut cq = self.uring.completion();
+
+            cq.sync();
 
-            let index = cqe.user_data() as _;
+            // if the cqueue is full, we are experiencing overflow and need to keep looping
+            let should_continue = cq.is_full();
+
+            for cqe in cq {
+                if cqe.user_data() == u64::MAX {
+                    // Result of the cancellation action. There isn't anything we
+                    // need to do here. We must wait for the CQE for the operation
+                    // that was canceled.
+                    continue;
+                }
 
-            self.ops.complete(index, cqe.into());
+                let index = cqe.user_data() as _;
+
+                self.ops.complete(index, cqe.into());
+            }
+
+            if !should_continue {
+                break;
+            }
         }
     }
 

src/runtime/mod.rs

@@ -62,7 +62,22 @@ impl Runtime {
         let rt = tokio::runtime::Builder::new_current_thread()
             .on_thread_park(|| {
                 CONTEXT.with(|x| {
-                    let _ = x.with_driver_mut(|d| d.uring.submit());
+                    let _ = x.with_driver_mut(|d| {
+                        // optimization: we can potentially avoid an epoll_wait call if we try to
+                        // dispatch completions before we park on epoll
+                        d.tick();
+                        d.submit()
+                    });
+                });
+            })
+            .on_thread_unpark(|| {
+                CONTEXT.with(|x| {
+                    x.with_driver_mut(|d| {
+                        // Dispatch completions to wake tasks based on any completed ops.
+                        // this is an optimization to try and avoid the whole "give the io driver
+                        // a dedicated task" thing
+                        d.tick();
+                    });
                 });
             })
             .enable_all()
@@ -106,11 +121,10 @@ impl Runtime {
 
         self.local.spawn_local(drive);
 
-        self.rt
-            .block_on(self.local.run_until(crate::future::poll_fn(|cx| {
-                // assert!(drive.as_mut().poll(cx).is_pending());
-                future.as_mut().poll(cx)
-            })))
+        self.rt.block_on(
+            self.local
+                .run_until(crate::future::poll_fn(|cx| future.as_mut().poll(cx))),
+        )
     }
 }
 

tests/driver.rs

@@ -1,4 +1,5 @@
 use tempfile::NamedTempFile;
+use tokio::task::JoinSet;
 
 use tokio_uring::{buf::IoBuf, fs::File};
 
@@ -90,6 +91,28 @@ fn too_many_submissions() {
     });
 }
 
+#[test]
+fn completion_overflow() {
+    let spawn_cnt = 100;
+    let squeue_entries = 2;
+    let cqueue_entries = 2 * squeue_entries;
+
+    tokio_uring::builder()
+        .entries(squeue_entries)
+        .uring_builder(tokio_uring::uring_builder().setup_cqsize(cqueue_entries))
+        .start(async move {
+            let mut js = JoinSet::new();
+
+            for _ in 0..spawn_cnt {
+                js.spawn_local(tokio_uring::no_op());
+            }
+
+            while let Some(res) = js.join_next().await {
+                res.unwrap().unwrap();
+            }
+        });
+}
+
 fn tempfile() -> NamedTempFile {
     NamedTempFile::new().unwrap()
 }