From bf8c77bea1e6742d1fbff0a2790e4ac5064f6933 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Wed, 7 Apr 2021 11:40:03 -0700
Subject: [PATCH] util: add `PollSemaphore::{add_permits, available_permits}`
 (#3683)

## Motivation

The `tokio::sync::Semaphore` type provides a
[`Semaphore::available_permits` method][1] which returns the current
number of permits available on the semaphore.
`tokio_util::sync::PollSemaphore` [does not expose such a method][2]. It
is possible to use `PollSemaphore::into_inner` or
`PollSemaphore::clone_inner` to unwrap the inner semaphore, but this may
require cloning/dropping the semaphore's `Arc` which shouldn't be
necessary to access the permits.

## Solution

This commit adds `PollSemaphore::available_permits`. It also adds
`PollSemaphore::add_permits` and an `AsRef<Semaphore>` impl while
we're here.

[1]: https://docs.rs/tokio/1.4.0/tokio/sync/struct.Semaphore.html#method.available_permits
[2]: https://docs.rs/tokio-util/0.6.5/tokio_util/sync/struct.PollSemaphore.html#implementations

Closes #3682
---
 tokio-util/src/sync/poll_semaphore.rs | 29 +++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tokio-util/src/sync/poll_semaphore.rs b/tokio-util/src/sync/poll_semaphore.rs
index a0a531c91..fcf943e9b 100644
--- a/tokio-util/src/sync/poll_semaphore.rs
+++ b/tokio-util/src/sync/poll_semaphore.rs
@@ -82,6 +82,29 @@ impl PollSemaphore {
             }
         }
     }
+
+    /// Returns the current number of available permits.
+    ///
+    /// This is equivalent to the [`Semaphore::available_permits`] method on the
+    /// `tokio::sync::Semaphore` type.
+    ///
+    /// [`Semaphore::available_permits`]: tokio::sync::Semaphore::available_permits
+    pub fn available_permits(&self) -> usize {
+        self.semaphore.available_permits()
+    }
+
+    /// Adds `n` new permits to the semaphore.
+    ///
+    /// The maximum number of permits is `usize::MAX >> 3`, and this function
+    /// will panic if the limit is exceeded.
+    ///
+    /// This is equivalent to the [`Semaphore::add_permits`] method on the
+    /// `tokio::sync::Semaphore` type.
+    ///
+    /// [`Semaphore::add_permits`]: tokio::sync::Semaphore::add_permits
+    pub fn add_permits(&self, n: usize) {
+        self.semaphore.add_permits(n);
+    }
 }
 
 impl Stream for PollSemaphore {
@@ -105,3 +128,9 @@ impl fmt::Debug for PollSemaphore {
             .finish()
     }
 }
+
+impl AsRef<Semaphore> for PollSemaphore {
+    fn as_ref(&self) -> &Semaphore {
+        &*self.semaphore
+    }
+}