Implement cooperative multitasking for UnixDatagram

tokio/src/runtime/coop.rs

@@ -196,6 +196,17 @@ cfg_coop! {
         }).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained()))))
     }
 
+    #[inline]
+    pub(crate) fn try_decrement() {
+        let _ = context::budget(|cell| {
+            let mut budget = cell.get();
+
+            budget.decrement();
+
+            cell.set(budget);
+         });
+    }
+
     cfg_rt! {
         cfg_metrics! {
             #[inline(always)]

tokio/src/runtime/io/registration.rs

@@ -223,7 +223,11 @@ impl Registration {
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                     self.clear_readiness(event);
                 }
-                x => return x,
+                x => {
+                    crate::runtime::coop::try_decrement();
+
+                    return x;
+                }
             }
         }
     }

tokio/src/runtime/io/scheduled_io.rs

@@ -159,6 +159,8 @@ struct Readiness<'a> {
 
     /// Entry in the waiter `LinkedList`.
     waiter: UnsafeCell<Waiter>,
+
+    is_waiter_registered: bool,
 }
 
 enum State {
@@ -429,6 +431,7 @@ impl ScheduledIo {
                 interest,
                 _p: PhantomPinned,
             }),
+            is_waiter_registered: false,
         }
     }
 }
@@ -458,11 +461,46 @@ impl Future for Readiness<'_> {
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         use std::sync::atomic::Ordering::SeqCst;
 
-        let (scheduled_io, state, waiter) = unsafe {
+        let (scheduled_io, state, waiter, is_waiter_registered) = unsafe {
             let me = self.get_unchecked_mut();
-            (&me.scheduled_io, &mut me.state, &me.waiter)
+            (
+                &me.scheduled_io,
+                &mut me.state,
+                &me.waiter,
+                &mut me.is_waiter_registered,
+            )
         };
 
+        if !crate::runtime::coop::has_budget_remaining() {
+            // Wasn't ready, take the lock (and check again while locked).
+            let mut waiters = scheduled_io.waiters.lock();
+
+            let w = unsafe { &mut *waiter.get() };
+
+            if *is_waiter_registered {
+                // Update the waker, if necessary.
+                if !w.waker.as_ref().unwrap().will_wake(cx.waker()) {
+                    w.waker = Some(cx.waker().clone());
+                }
+            } else {
+                // Safety: called while locked
+                w.waker = Some(cx.waker().clone());
+
+                // Insert the waiter into the linked list
+                //
+                // safety: pointers from `UnsafeCell` are never null.
+                waiters
+                    .list
+                    .push_front(unsafe { NonNull::new_unchecked(waiter.get()) });
+
+                *is_waiter_registered = true;
+            }
+
+            drop(waiters);
+
+            return Poll::Pending;
+        }
+
         loop {
             match *state {
                 State::Init => {
@@ -512,17 +550,22 @@ impl Future for Readiness<'_> {
 
                     // Not ready even after locked, insert into list...
 
-                    // Safety: called while locked
-                    unsafe {
-                        (*waiter.get()).waker = Some(cx.waker().clone());
+                    if !*is_waiter_registered {
+                        // Safety: called while locked
+                        unsafe {
+                            (*waiter.get()).waker = Some(cx.waker().clone());
+                        }
+
+                        // Insert the waiter into the linked list
+                        //
+                        // safety: pointers from `UnsafeCell` are never null.
+                        waiters
+                            .list
+                            .push_front(unsafe { NonNull::new_unchecked(waiter.get()) });
+
+                        *is_waiter_registered = true;
                     }
 
-                    // Insert the waiter into the linked list
-                    //
-                    // safety: pointers from `UnsafeCell` are never null.
-                    waiters
-                        .list
-                        .push_front(unsafe { NonNull::new_unchecked(waiter.get()) });
                     *state = State::Waiting;
                 }
                 State::Waiting => {

tokio/tests/uds_datagram.rs

@@ -411,3 +411,47 @@ async fn poll_ready() -> io::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "current_thread")]
+async fn coop_uds() -> io::Result<()> {
+    use std::sync::atomic::{AtomicU32, Ordering};
+    use std::time::{Duration, Instant};
+
+    const HELLO: &[u8] = b"hello world";
+    const DURATION: Duration = Duration::from_secs(1);
+
+    let dir = tempfile::tempdir().unwrap();
+    let server_path = dir.path().join("server.sock");
+
+    let client = std::os::unix::net::UnixDatagram::unbound().unwrap();
+    let server = UnixDatagram::bind(&server_path).unwrap();
+
+    let counter = Arc::new(AtomicU32::new(0));
+
+    let counter_jh = tokio::spawn({
+        let counter = counter.clone();
+
+        async move {
+            loop {
+                tokio::time::sleep(Duration::from_millis(250)).await;
+                counter.fetch_add(1, Ordering::Relaxed);
+            }
+        }
+    });
+
+    let mut buf = [0; HELLO.len()];
+    let start = Instant::now();
+    while Instant::now().duration_since(start) < DURATION {
+        let _ = client.send_to(HELLO, &server_path);
+        let _ = server.recv(&mut buf[..]).await.unwrap();
+    }
+
+    counter_jh.abort();
+    let _ = counter_jh.await;
+
+    let expected = ((DURATION.as_secs() * 4) as f64 * 0.5) as u32;
+    let counter = counter.load(Ordering::Relaxed);
+    assert!(counter >= expected);
+
+    Ok(())
+}