Repository: tokio-rs/tokio (language: rust)

Issue:
`watch::Receiver::changed` does use cooperative scheduling
**Version**

tokio 1.40.0

**Platform**

Linux hostname 6.10.8-arch1-1 #1 SMP PREEMPT_DYNAMIC Wed, 04 Sep 2024 15:16:37 +0000 x86_64 GNU/Linux

**Description**

I stumbled upon a surprising behavior that caused a program to never end, because dropping
the runtime never ended.

The behavior can be better explained with this reproducer:

```rust
fn main() {
    let rt = tokio::runtime::Runtime::new().unwrap();

    // With an mpsc, there is no blocking on runtime drop
    // With a watch, the runtime drop can block infinitely.

    // let (sender, mut receiver) = tokio::sync::mpsc::channel::<()>(1);
    let (sender, mut receiver) = tokio::sync::watch::channel(false);

    // This task will block the runtime during drop
    rt.spawn(async move {
        loop {
            // let _r = receiver.recv().await;
            let _r = receiver.changed().await;
        }
    });
    drop(sender);

    println!("drop runtime");
    drop(rt);
    println!("runtime has been dropped");
}
```

Running this, "drop runtime" is printed but "runtime has been dropped" is never printed and the program never ends. This is because the task spin loops on the call to `receiver.changed().await` which returns an error.

It would be expected that this `.await` would yield and thus allow the runtime to drop the task, but
that does not appear to be the case. This is afaict surprising behavior from my understanding on how shutting down the runtime is documented, especially given that for other types of queues (for example mpsc), the recv() call does yield even if the sender has been dropped for example.

Of course the code isn't ideal here and the bug can clearly be avoided, but I thought it would be interesting to report. Either this is a bug and the `changed().await` should yield, or maybe the documentation on this function should make it clear that in the case where the sender has been dropped, this call does not yield?



Discussion / hints:
It's entirely normal that async operations don't yield if they can complete immediately. However, we do have a mechanism to avoid this kind of bug called [cooperative scheduling](https://docs.rs/tokio/latest/tokio/task/index.html#cooperative-scheduling), which we should apply to the watch channel.

Fix patch:
diff --git a/tokio/src/runtime/coop.rs b/tokio/src/runtime/coop.rs
index aaca8b6baa2..c01e5e3de8b 100644
--- a/tokio/src/runtime/coop.rs
+++ b/tokio/src/runtime/coop.rs
@@ -135,8 +135,11 @@ cfg_rt! {
 }
 
 cfg_coop! {
+    use pin_project_lite::pin_project;
     use std::cell::Cell;
-    use std::task::{Context, Poll};
+    use std::future::Future;
+    use std::pin::Pin;
+    use std::task::{ready, Context, Poll};
 
     #[must_use]
     pub(crate) struct RestoreOnPending(Cell<Budget>);
@@ -240,6 +243,44 @@ cfg_coop! {
             self.0.is_none()
         }
     }
+
+    pin_project! {
+        /// Future wrapper to ensure cooperative scheduling.
+        ///
+        /// When being polled `poll_proceed` is called before the inner future is polled to check
+        /// if the inner future has exceeded its budget. If the inner future resolves, this will
+        /// automatically call `RestoreOnPending::made_progress` before resolving this future with
+        /// the result of the inner one. If polling the inner future is pending, polling this future
+        /// type will also return a `Poll::Pending`.
+        #[must_use = "futures do nothing unless polled"]
+        pub(crate) struct Coop<F: Future> {
+            #[pin]
+            pub(crate) fut: F,
+        }
+    }
+
+    impl<F: Future> Future for Coop<F> {
+        type Output = F::Output;
+
+        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            let coop = ready!(poll_proceed(cx));
+            let me = self.project();
+            if let Poll::Ready(ret) = me.fut.poll(cx) {
+                coop.made_progress();
+                Poll::Ready(ret)
+            } else {
+                Poll::Pending
+            }
+        }
+    }
+
+    /// Run a future with a budget constraint for cooperative scheduling.
+    /// If the future exceeds its budget while being polled, control is yielded back to the
+    /// runtime.
+    #[inline]
+    pub(crate) fn cooperative<F: Future>(fut: F) -> Coop<F> {
+        Coop { fut }
+    }
 }
 
 #[cfg(all(test, not(loom)))]
diff --git a/tokio/src/sync/watch.rs b/tokio/src/sync/watch.rs
index 490b9e4df88..af72e30dd32 100644
--- a/tokio/src/sync/watch.rs
+++ b/tokio/src/sync/watch.rs
@@ -111,6 +111,7 @@
 //! [`Sender::closed`]: crate::sync::watch::Sender::closed
 //! [`Sender::subscribe()`]: crate::sync::watch::Sender::subscribe
 
+use crate::runtime::coop::cooperative;
 use crate::sync::notify::Notify;
 
 use crate::loom::sync::atomic::AtomicUsize;
@@ -743,7 +744,7 @@ impl<T> Receiver<T> {
     /// }
     /// ```
     pub async fn changed(&mut self) -> Result<(), error::RecvError> {
-        changed_impl(&self.shared, &mut self.version).await
+        cooperative(changed_impl(&self.shared, &mut self.version)).await
     }
 
     /// Waits for a value that satisfies the provided condition.
@@ -807,6 +808,13 @@ impl<T> Receiver<T> {
     /// }
     /// ```
     pub async fn wait_for(
+        &mut self,
+        f: impl FnMut(&T) -> bool,
+    ) -> Result<Ref<'_, T>, error::RecvError> {
+        cooperative(self.wait_for_inner(f)).await
+    }
+
+    async fn wait_for_inner(
         &mut self,
         mut f: impl FnMut(&T) -> bool,
     ) -> Result<Ref<'_, T>, error::RecvError> {
@@ -1224,19 +1232,22 @@ impl<T> Sender<T> {
     /// }
     /// ```
     pub async fn closed(&self) {
-        crate::trace::async_trace_leaf().await;
+        cooperative(async {
+            crate::trace::async_trace_leaf().await;
 
-        while self.receiver_count() > 0 {
-            let notified = self.shared.notify_tx.notified();
+            while self.receiver_count() > 0 {
+                let notified = self.shared.notify_tx.notified();
 
-            if self.receiver_count() == 0 {
-                return;
-            }
+                if self.receiver_count() == 0 {
+                    return;
+                }
 
-            notified.await;
-            // The channel could have been reopened in the meantime by calling
-            // `subscribe`, so we loop again.
-        }
+                notified.await;
+                // The channel could have been reopened in the meantime by calling
+                // `subscribe`, so we loop again.
+            }
+        })
+        .await;
     }
 
     /// Creates a new [`Receiver`] connected to this `Sender`.


Test patch:
diff --git a/tokio/tests/sync_watch.rs b/tokio/tests/sync_watch.rs
index 17f0c81087a..4418f88e57b 100644
--- a/tokio/tests/sync_watch.rs
+++ b/tokio/tests/sync_watch.rs
@@ -368,3 +368,85 @@ async fn receiver_is_notified_when_last_sender_is_dropped() {
 
     assert!(t.is_woken());
 }
+
+#[tokio::test]
+async fn receiver_changed_is_cooperative() {
+    let (tx, mut rx) = watch::channel(());
+
+    drop(tx);
+
+    tokio::select! {
+        biased;
+        _ = async {
+            loop {
+                assert!(rx.changed().await.is_err());
+            }
+        } => {},
+        _ = tokio::task::yield_now() => {},
+    }
+}
+
+#[tokio::test]
+async fn receiver_changed_is_cooperative_ok() {
+    let (tx, mut rx) = watch::channel(());
+
+    tokio::select! {
+        biased;
+        _ = async {
+            loop {
+                assert!(tx.send(()).is_ok());
+                assert!(rx.changed().await.is_ok());
+            }
+        } => {},
+        _ = tokio::task::yield_now() => {},
+    }
+}
+
+#[tokio::test]
+async fn receiver_wait_for_is_cooperative() {
+    let (tx, mut rx) = watch::channel(0);
+
+    drop(tx);
+
+    tokio::select! {
+        biased;
+        _ = async {
+            loop {
+                assert!(rx.wait_for(|val| *val == 1).await.is_err());
+            }
+        } => {},
+        _ = tokio::task::yield_now() => {},
+    }
+}
+
+#[tokio::test]
+async fn receiver_wait_for_is_cooperative_ok() {
+    let (tx, mut rx) = watch::channel(0);
+
+    tokio::select! {
+        biased;
+        _ = async {
+            loop {
+                assert!(tx.send(1).is_ok());
+                assert!(rx.wait_for(|val| *val == 1).await.is_ok());
+            }
+        } => {},
+        _ = tokio::task::yield_now() => {},
+    }
+}
+
+#[tokio::test]
+async fn sender_closed_is_cooperative() {
+    let (tx, rx) = watch::channel(());
+
+    drop(rx);
+
+    tokio::select! {
+        _ = async {
+            loop {
+                tx.closed().await;
+            }
+        } => {},
+        _ = tokio::task::yield_now() => {},
+    }
+}