sync: wake mpsc receiver when a queued reserve[_many] returns permits (#8260)

tokio/src/sync/mpsc/bounded.rs

@@ -1275,8 +1275,37 @@ impl<T> Sender<T> {
         if n > self.max_capacity() {
             return Err(SendError(()));
         }
-        match self.chan.semaphore().semaphore.acquire(n).await {
-            Ok(()) => Ok(()),
+
+        // If the reservation fails or is cancelled while it holds some but not
+        // all of the permits it asked for, the inner `Acquire` returns those
+        // permits to the semaphore on drop, without the receiver notification
+        // that `Permit` and `PermitIterator` do on release. Wake the receiver in
+        // that case so a `recv()` blocked on the channel becoming idle can
+        // observe a close. On success the caller receives the permits and takes
+        // over that job, so the guard is forgotten below.
+        struct WakeReceiverOnDrop<'a, T> {
+            chan: &'a chan::Tx<T, Semaphore>,
+        }
+
+        impl<T> Drop for WakeReceiverOnDrop<'_, T> {
+            fn drop(&mut self) {
+                use chan::Semaphore;
+
+                let semaphore = self.chan.semaphore();
+                if semaphore.is_closed() && semaphore.is_idle() {
+                    self.chan.wake_rx();
+                }
+            }
+        }
+
+        let guard = WakeReceiverOnDrop { chan: &self.chan };
+        let result = self.chan.semaphore().semaphore.acquire(n).await;
+
+        match result {
+            Ok(()) => {
+                std::mem::forget(guard);
+                Ok(())
+            }
             Err(_) => Err(SendError(())),
         }
     }

tokio/tests/sync_mpsc.rs

@@ -519,6 +519,110 @@ async fn recv_close_gets_none_reserved() {
     assert!(rx.recv().await.is_none());
 }
 
+#[test]
+fn failed_reserve_many_wakes_closed_receiver() {
+    let (tx, mut rx) = mpsc::channel::<()>(2);
+
+    // Hold one of the two slots so the reservation can only take the other one
+    // and has to queue for the second.
+    let permit = tx.try_reserve().unwrap();
+
+    let mut reserve = tokio_test::task::spawn(tx.reserve_many(2));
+    assert_pending!(reserve.poll());
+
+    rx.close();
+
+    let mut recv = tokio_test::task::spawn(rx.recv());
+    assert_pending!(recv.poll());
+
+    // One slot is still outstanding, so the channel is not idle and the
+    // receiver must not be woken yet.
+    drop(permit);
+    assert!(!recv.is_woken());
+
+    // The reservation fails on the closed channel and returns the slot it had
+    // already acquired, leaving the channel idle. The receiver must be woken.
+    assert_ready_err!(reserve.poll());
+    assert!(recv.is_woken());
+    assert_ready!(recv.poll());
+}
+
+#[test]
+fn cancelled_reserve_many_wakes_closed_receiver() {
+    let (tx, mut rx) = mpsc::channel::<()>(2);
+
+    let permit = tx.try_reserve().unwrap();
+
+    let mut reserve = tokio_test::task::spawn(tx.reserve_many(2));
+    assert_pending!(reserve.poll());
+
+    rx.close();
+
+    let mut recv = tokio_test::task::spawn(rx.recv());
+    assert_pending!(recv.poll());
+
+    drop(permit);
+    assert!(!recv.is_woken());
+
+    // Cancelling the reservation returns the slot it had already acquired,
+    // leaving the channel idle. The receiver must be woken.
+    drop(reserve);
+    assert!(recv.is_woken());
+    assert_ready!(recv.poll());
+}
+
+#[test]
+fn failed_reserve_wakes_closed_receiver() {
+    let (tx, mut rx) = mpsc::channel::<()>(1);
+
+    // Hold the only slot so the reservation below has to queue.
+    let permit = tx.try_reserve().unwrap();
+
+    let mut reserve = tokio_test::task::spawn(tx.reserve());
+    assert_pending!(reserve.poll());
+
+    // Returning the slot assigns it to the queued reservation, which is woken
+    // but not yet polled, so the channel stays non-idle.
+    drop(permit);
+    assert!(reserve.is_woken());
+
+    rx.close();
+
+    let mut recv = tokio_test::task::spawn(rx.recv());
+    assert_pending!(recv.poll());
+
+    // The reservation fails on the closed channel despite holding the
+    // assigned slot; returning it leaves the channel idle. The receiver must
+    // be woken.
+    assert_ready_err!(reserve.poll());
+    assert!(recv.is_woken());
+    assert_ready!(recv.poll());
+}
+
+#[test]
+fn cancelled_reserve_wakes_closed_receiver() {
+    let (tx, mut rx) = mpsc::channel::<()>(1);
+
+    let permit = tx.try_reserve().unwrap();
+
+    let mut reserve = tokio_test::task::spawn(tx.reserve());
+    assert_pending!(reserve.poll());
+
+    drop(permit);
+    assert!(reserve.is_woken());
+
+    rx.close();
+
+    let mut recv = tokio_test::task::spawn(rx.recv());
+    assert_pending!(recv.poll());
+
+    // Cancelling the reservation returns the slot it was assigned, leaving
+    // the channel idle. The receiver must be woken.
+    drop(reserve);
+    assert!(recv.is_woken());
+    assert_ready!(recv.poll());
+}
+
 #[maybe_tokio_test]
 async fn tx_close_gets_none() {
     let (_, mut rx) = mpsc::channel::<i32>(10);