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, - 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, 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, 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.