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: 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 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

(limited to 'tokio/src/sync/mpsc')

diff --git a/tokio/src/sync/mpsc/chan.rs b/tokio/src/sync/mpsc/chan.rs
index 148ee3ad..0a53cda2 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);
-- 
cgit v1.2.3