From 08f1b67fcb9564483f2e6e5ab635d5636f2a0b80 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Mon, 12 Apr 2021 09:23:03 +0200
Subject: [PATCH] sync: add `mpsc::Sender::capacity` (#3690)

Simply exposes the number of available permits of the semaphore.

This makes some kinds of bookkeeping easier without having to manually keep counts using atomics.

Fixes #2642
---
 tokio/src/sync/mpsc/bounded.rs | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/tokio/src/sync/mpsc/bounded.rs b/tokio/src/sync/mpsc/bounded.rs
index 5d17c1e199b..1f670bfa743 100644
--- a/tokio/src/sync/mpsc/bounded.rs
+++ b/tokio/src/sync/mpsc/bounded.rs
@@ -699,6 +699,39 @@ impl<T> Sender<T> {
     pub fn same_channel(&self, other: &Self) -> bool {
         self.chan.same_channel(&other.chan)
     }
+
+    /// Returns the current capacity of the channel.
+    ///
+    /// The capacity goes down when sending a value by calling [`send`] or by reserving capacity
+    /// with [`reserve`]. The capacity goes up when values are received by the [`Receiver`].
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, mut rx) = mpsc::channel::<()>(5);
+    ///
+    ///     assert_eq!(tx.capacity(), 5);
+    ///
+    ///     // Making a reservation drops the capacity by one.
+    ///     let permit = tx.reserve().await.unwrap();
+    ///     assert_eq!(tx.capacity(), 4);
+    ///
+    ///     // Sending and receiving a value increases the caapcity by one.
+    ///     permit.send(());
+    ///     rx.recv().await.unwrap();
+    ///     assert_eq!(tx.capacity(), 5);
+    /// }
+    /// ```
+    ///
+    /// [`send`]: Sender::send
+    /// [`reserve`]: Sender::reserve
+    pub fn capacity(&self) -> usize {
+        self.chan.semaphore().0.available_permits()
+    }
 }
 
 impl<T> Clone for Sender<T> {