mpsc: ensure try_reserve error is consistent with try_send (#4119)

tokio-rs · Sep 19, 2021 · e9f6fae · e9f6fae
1 parent f1b8967
commit e9f6fae
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/tokio/src/sync/mpsc/bounded.rs b/tokio/src/sync/mpsc/bounded.rs
@@ -851,7 +851,8 @@ impl<T> Sender<T> {
     pub fn try_reserve(&self) -> Result<Permit<'_, T>, TrySendError<()>> {
         match self.chan.semaphore().0.try_acquire(1) {
             Ok(_) => {}
-            Err(_) => return Err(TrySendError::Full(())),
+            Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(())),
+            Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(())),
         }
 
         Ok(Permit { chan: &self.chan })
@@ -915,7 +916,8 @@ impl<T> Sender<T> {
     pub fn try_reserve_owned(self) -> Result<OwnedPermit<T>, TrySendError<Self>> {
         match self.chan.semaphore().0.try_acquire(1) {
             Ok(_) => {}
-            Err(_) => return Err(TrySendError::Full(self)),
+            Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(self)),
+            Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(self)),
         }
 
         Ok(OwnedPermit {

diff --git a/tokio/tests/sync_mpsc.rs b/tokio/tests/sync_mpsc.rs
@@ -410,13 +410,15 @@ fn dropping_rx_closes_channel_for_try() {
 
     drop(rx);
 
-    {
-        let err = assert_err!(tx.try_send(msg.clone()));
-        match err {
-            TrySendError::Closed(..) => {}
-            _ => panic!(),
-        }
-    }
+    assert!(matches!(
+        tx.try_send(msg.clone()),
+        Err(TrySendError::Closed(_))
+    ));
+    assert!(matches!(tx.try_reserve(), Err(TrySendError::Closed(_))));
+    assert!(matches!(
+        tx.try_reserve_owned(),
+        Err(TrySendError::Closed(_))
+    ));
 
     assert_eq!(1, Arc::strong_count(&msg));
 }