From 93f8cb8df2bb0dceb7921556165a8ed8efed9151 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Mon, 21 Sep 2020 14:29:22 -0700
Subject: [PATCH] sync: fix missing notification during mpsc close (#2854)

When the mpsc channel receiver closes the channel, receiving should
return `None` once all in-progress sends have completed. When a sender
reserves capacity, this prevents the receiver from fully shutting down.
Previously, when the sender, after reserving capacity, dropped without
sending a message, the receiver was not notified. This results in
blocking the shutdown process until all sender handles drop.

This patch adds a receiver notification when the channel is both closed
and all outstanding sends have completed.
---
 tokio/src/sync/mpsc/chan.rs | 22 +++++++++++++++++-----
 tokio/tests/sync_mpsc.rs    | 22 ++++++++++++++++++++++
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/tokio/src/sync/mpsc/chan.rs b/tokio/src/sync/mpsc/chan.rs
index 148ee3ad766..0a53cda2038 100644
--- a/tokio/src/sync/mpsc/chan.rs
+++ b/tokio/src/sync/mpsc/chan.rs
@@ -75,7 +75,11 @@ pub(crate) trait Semaphore {
 
     /// The permit is dropped without a value being sent. In this case, the
     /// permit must be returned to the semaphore.
-    fn drop_permit(&self, permit: &mut Self::Permit);
+    ///
+    /// # Return
+    ///
+    /// Returns true if the permit was acquired.
+    fn drop_permit(&self, permit: &mut Self::Permit) -> bool;
 
     fn is_idle(&self) -> bool;
 
@@ -192,7 +196,7 @@ where
 
     pub(crate) fn disarm(&mut self) {
         // TODO: should this error if not acquired?
-        self.inner.semaphore.drop_permit(&mut self.permit)
+        self.inner.semaphore.drop_permit(&mut self.permit);
     }
 
     /// Send a message and notify the receiver.
@@ -234,7 +238,11 @@ where
     S: Semaphore,
 {
     fn drop(&mut self) {
-        self.inner.semaphore.drop_permit(&mut self.permit);
+        let notify = self.inner.semaphore.drop_permit(&mut self.permit);
+
+        if notify && self.inner.semaphore.is_idle() {
+            self.inner.rx_waker.wake();
+        }
 
         if self.inner.tx_count.fetch_sub(1, AcqRel) != 1 {
             return;
@@ -424,8 +432,10 @@ impl Semaphore for (crate::sync::semaphore_ll::Semaphore, usize) {
         Permit::new()
     }
 
-    fn drop_permit(&self, permit: &mut Permit) {
+    fn drop_permit(&self, permit: &mut Permit) -> bool {
+        let ret = permit.is_acquired();
         permit.release(1, &self.0);
+        ret
     }
 
     fn add_permit(&self) {
@@ -477,7 +487,9 @@ impl Semaphore for AtomicUsize {
 
     fn new_permit() {}
 
-    fn drop_permit(&self, _permit: &mut ()) {}
+    fn drop_permit(&self, _permit: &mut ()) -> bool {
+        false
+    }
 
     fn add_permit(&self) {
         let prev = self.fetch_sub(2, Release);
diff --git a/tokio/tests/sync_mpsc.rs b/tokio/tests/sync_mpsc.rs
index f571a71caf8..919bddbfb18 100644
--- a/tokio/tests/sync_mpsc.rs
+++ b/tokio/tests/sync_mpsc.rs
@@ -534,3 +534,25 @@ async fn blocking_send_async() {
     let (mut tx, _rx) = mpsc::channel::<()>(1);
     let _ = tx.blocking_send(());
 }
+
+#[test]
+fn ready_close_cancel_bounded() {
+    use futures::future::poll_fn;
+
+    let (mut tx, mut rx) = mpsc::channel::<()>(100);
+    let _tx2 = tx.clone();
+
+    {
+        let mut ready = task::spawn(async { poll_fn(|cx| tx.poll_ready(cx)).await });
+        assert_ready_ok!(ready.poll());
+    }
+
+    rx.close();
+
+    let mut recv = task::spawn(async { rx.recv().await });
+    assert_pending!(recv.poll());
+
+    drop(tx);
+
+    assert!(recv.is_woken());
+}