From 07c2be78af645dc83d64456bdd6ee01cf8037e45 Mon Sep 17 00:00:00 2001
From: Darius Carrier <carrierdarius@gmail.com>
Date: Wed, 28 Oct 2020 23:35:36 +0000
Subject: [PATCH] tokio-sync: add acquire_many and try_acquire_many methods to
 `Sempahore`

Fixes: #1550
---
 tokio/src/sync/semaphore.rs | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/tokio/src/sync/semaphore.rs b/tokio/src/sync/semaphore.rs
index 43dd97681ba..2acccfa2fc1 100644
--- a/tokio/src/sync/semaphore.rs
+++ b/tokio/src/sync/semaphore.rs
@@ -27,7 +27,7 @@ pub struct Semaphore {
 #[derive(Debug)]
 pub struct SemaphorePermit<'a> {
     sem: &'a Semaphore,
-    permits: u16,
+    permits: u32,
 }
 
 /// An owned permit from the semaphore.
@@ -39,7 +39,7 @@ pub struct SemaphorePermit<'a> {
 #[derive(Debug)]
 pub struct OwnedSemaphorePermit {
     sem: Arc<Semaphore>,
-    permits: u16,
+    permits: u32,
 }
 
 /// Error returned from the [`Semaphore::try_acquire`] function.
@@ -104,6 +104,15 @@ impl Semaphore {
         }
     }
 
+    /// Acquires `n` permits from the semaphore
+    pub async fn acquire_many(&self, n: u32) -> SemaphorePermit<'_> {
+        self.ll_sem.acquire(n).await.unwrap();
+        SemaphorePermit {
+            sem: &self,
+            permits: n,
+        }
+    }
+
     /// Tries to acquire a permit from the semaphore.
     pub fn try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError> {
         match self.ll_sem.try_acquire(1) {
@@ -115,6 +124,17 @@ impl Semaphore {
         }
     }
 
+    /// Tries to acquire `n` permits from the semaphore.
+    pub fn try_acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, TryAcquireError> {
+        match self.ll_sem.try_acquire(n) {
+            Ok(_) => Ok(SemaphorePermit {
+                sem: self,
+                permits: n,
+            }),
+            Err(_) => Err(TryAcquireError(())),
+        }
+    }
+
     /// Acquires permit from the semaphore.
     ///
     /// The semaphore must be wrapped in an [`Arc`] to call this method.