From a52f5071bf5ecf31c44e8aba5d8611400c50eb71 Mon Sep 17 00:00:00 2001
From: Darius Carrier <carrierdarius@gmail.com>
Date: Tue, 10 Nov 2020 09:39:30 -0800
Subject: sync: add acquire_many and try_acquire_many to `Sempahore` (#3067)

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

(limited to 'tokio')

diff --git a/tokio/src/sync/semaphore.rs b/tokio/src/sync/semaphore.rs
index 43dd9768..2acccfa2 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.
-- 
cgit v1.2.3