Expose ability to acquire multiple permits

tokio-rs#1550

tokio/src/sync/semaphore.rs

@@ -86,21 +86,33 @@ impl Semaphore {
         self.ll_sem.release(n);
     }
 
-    /// Acquires permit from the semaphore.
+    /// Acquires one permit from the semaphore.
+    #[inline]
     pub async fn acquire(&self) -> SemaphorePermit<'_> {
-        self.ll_sem.acquire(1).await.unwrap();
+        self.acquire_n(1).await
+    }
+
+    /// Acquires a number of permits from the semaphore.
+    pub async fn acquire_n(&self, num_permits: u16) -> SemaphorePermit<'_> {
+        self.ll_sem.acquire(num_permits).await.unwrap();
         SemaphorePermit {
             sem: &self,
-            permits: 1,
+            permits: num_permits,
         }
     }
 
     /// Tries to acquire a permit from the semaphore.
+    #[inline]
     pub fn try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError> {
-        match self.ll_sem.try_acquire(1) {
+        self.try_acquire_n(1)
+    }
+
+    /// Tries to acquire a number of permits from the semaphore.
+    pub fn try_acquire_n(&self, num_permits: u16) -> Result<SemaphorePermit<'_>, TryAcquireError> {
+        match self.ll_sem.try_acquire(num_permits) {
             Ok(_) => Ok(SemaphorePermit {
                 sem: self,
-                permits: 1,
+                permits: num_permits,
             }),
             Err(_) => Err(TryAcquireError(())),
         }
@@ -111,11 +123,21 @@ impl Semaphore {
     /// The semaphore must be wrapped in an [`Arc`] to call this method.
     ///
     /// [`Arc`]: std::sync::Arc
+    #[inline]
     pub async fn acquire_owned(self: Arc<Self>) -> OwnedSemaphorePermit {
-        self.ll_sem.acquire(1).await.unwrap();
+        self.acquire_owned_n(1).await
+    }
+
+    /// Acquires a number of permits from the semaphore.
+    ///
+    /// The semaphore must be wrapped in an [`Arc`] to call this method.
+    ///
+    /// [`Arc`]: std::sync::Arc
+    pub async fn acquire_owned_n(self: Arc<Self>, num_permits: u16) -> OwnedSemaphorePermit {
+        self.ll_sem.acquire(num_permits).await.unwrap();
         OwnedSemaphorePermit {
             sem: self.clone(),
-            permits: 1,
+            permits: num_permits,
         }
     }
 
@@ -124,11 +146,21 @@ impl Semaphore {
     /// The semaphore must be wrapped in an [`Arc`] to call this method.
     ///
     /// [`Arc`]: std::sync::Arc
+    #[inline]
     pub fn try_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, TryAcquireError> {
-        match self.ll_sem.try_acquire(1) {
+        self.try_acquire_owned_n(1)
+    }
+
+    /// Tries to acquire a number of permits from the semaphore.
+    ///
+    /// The semaphore must be wrapped in an [`Arc`] to call this method.
+    ///
+    /// [`Arc`]: std::sync::Arc
+    pub fn try_acquire_owned_n(self: Arc<Self>, num_permits: u16) -> Result<OwnedSemaphorePermit, TryAcquireError> {
+        match self.ll_sem.try_acquire(num_permits) {
             Ok(_) => Ok(OwnedSemaphorePermit {
                 sem: self.clone(),
-                permits: 1,
+                permits: num_permits,
             }),
             Err(_) => Err(TryAcquireError(())),
         }