Skip to content

Commit

Permalink
sync: add acquire_many and try_acquire_many to Sempahore (#3067)
Browse files Browse the repository at this point in the history
Fixes: #1550
  • Loading branch information
dcarrier committed Nov 10, 2020
1 parent f1f8c3c commit a52f507
Showing 1 changed file with 22 additions and 2 deletions.
24 changes: 22 additions & 2 deletions tokio/src/sync/semaphore.rs
Expand Up @@ -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.
Expand All @@ -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.
Expand Down Expand Up @@ -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) {
Expand All @@ -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.
Expand Down

0 comments on commit a52f507

Please sign in to comment.