sync: add acquire_many and try_acquire_many to `Sempahore` (#3067)

Fixes: #1550

1 files changed, 22 insertions, 2 deletions

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.