sync: add blocking_acquire methods to Semaphore (#8269)

This commit is contained in:
Rachit2323
2026-08-11 09:51:50 +02:00
committed by GitHub
parent af93763009
commit 625954f365
2 changed files with 267 additions and 0 deletions
+169
View File
@@ -751,6 +751,94 @@ impl Semaphore {
}
}
/// Acquires a permit from the semaphore, blocking the current thread until
/// one is available.
///
/// If the semaphore has been closed, this returns an [`AcquireError`].
/// Otherwise, this returns a [`SemaphorePermit`] representing the
/// acquired permit.
///
/// This method is intended for use in synchronous code, such as inside
/// [`spawn_blocking`] or when the semaphore is shared between asynchronous
/// and synchronous code. It is the blocking equivalent of [`acquire`].
///
/// # Panics
///
/// This function panics if called within an asynchronous execution
/// context.
///
/// # Examples
///
/// ```
/// # #[cfg(not(target_family = "wasm"))]
/// # {
/// use std::sync::Arc;
/// use tokio::sync::Semaphore;
///
/// #[tokio::main(flavor = "current_thread")]
/// async fn main() {
/// let semaphore = Arc::new(Semaphore::new(2));
///
/// let semaphore2 = semaphore.clone();
/// let blocking_task = tokio::task::spawn_blocking(move || {
/// // Inside a blocking context we cannot use `.await`, so we use
/// // `blocking_acquire` instead.
/// let _permit = semaphore2.blocking_acquire().unwrap();
///
/// // ... perform blocking work while holding the permit ...
/// });
///
/// blocking_task.await.unwrap();
/// }
/// # }
/// ```
///
/// [`AcquireError`]: crate::sync::AcquireError
/// [`SemaphorePermit`]: crate::sync::SemaphorePermit
/// [`spawn_blocking`]: crate::task::spawn_blocking
/// [`acquire`]: Semaphore::acquire
#[track_caller]
#[cfg(feature = "sync")]
pub fn blocking_acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
crate::future::block_on(self.acquire())
}
/// Acquires `n` permits from the semaphore, blocking the current thread
/// until they are available.
///
/// If the semaphore has been closed, this returns an [`AcquireError`].
/// Otherwise, this returns a [`SemaphorePermit`] representing the
/// acquired permits.
///
/// This method is the blocking equivalent of [`acquire_many`].
///
/// # Panics
///
/// This function panics if called within an asynchronous execution
/// context.
///
/// # Examples
///
/// ```
/// use tokio::sync::Semaphore;
///
/// fn main() {
/// let semaphore = Semaphore::new(5);
///
/// let permit = semaphore.blocking_acquire_many(3).unwrap();
/// assert_eq!(semaphore.available_permits(), 2);
/// }
/// ```
///
/// [`AcquireError`]: crate::sync::AcquireError
/// [`SemaphorePermit`]: crate::sync::SemaphorePermit
/// [`acquire_many`]: Semaphore::acquire_many
#[track_caller]
#[cfg(feature = "sync")]
pub fn blocking_acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, AcquireError> {
crate::future::block_on(self.acquire_many(n))
}
/// Acquires a permit from the semaphore.
///
/// The semaphore must be wrapped in an [`Arc`] to call this method.
@@ -960,6 +1048,87 @@ impl Semaphore {
}
}
/// Acquires a permit from the semaphore, blocking the current thread until
/// one is available.
///
/// The semaphore must be wrapped in an [`Arc`] to call this method.
/// If the semaphore has been closed, this returns an [`AcquireError`].
/// Otherwise, this returns an [`OwnedSemaphorePermit`] representing the
/// acquired permit.
///
/// This method is the blocking equivalent of [`acquire_owned`].
///
/// # Panics
///
/// This function panics if called within an asynchronous execution
/// context.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
/// use tokio::sync::Semaphore;
///
/// fn main() {
/// let semaphore = Arc::new(Semaphore::new(2));
///
/// let permit = semaphore.clone().blocking_acquire_owned().unwrap();
/// assert_eq!(semaphore.available_permits(), 1);
/// }
/// ```
///
/// [`Arc`]: std::sync::Arc
/// [`AcquireError`]: crate::sync::AcquireError
/// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
/// [`acquire_owned`]: Semaphore::acquire_owned
#[track_caller]
#[cfg(feature = "sync")]
pub fn blocking_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, AcquireError> {
crate::future::block_on(self.acquire_owned())
}
/// Acquires `n` permits from the semaphore, blocking the current thread
/// until they are available.
///
/// The semaphore must be wrapped in an [`Arc`] to call this method.
/// If the semaphore has been closed, this returns an [`AcquireError`].
/// Otherwise, this returns an [`OwnedSemaphorePermit`] representing the
/// acquired permits.
///
/// This method is the blocking equivalent of [`acquire_many_owned`].
///
/// # Panics
///
/// This function panics if called within an asynchronous execution
/// context.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
/// use tokio::sync::Semaphore;
///
/// fn main() {
/// let semaphore = Arc::new(Semaphore::new(5));
///
/// let permit = semaphore.clone().blocking_acquire_many_owned(3).unwrap();
/// assert_eq!(semaphore.available_permits(), 2);
/// }
/// ```
///
/// [`Arc`]: std::sync::Arc
/// [`AcquireError`]: crate::sync::AcquireError
/// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
/// [`acquire_many_owned`]: Semaphore::acquire_many_owned
#[track_caller]
#[cfg(feature = "sync")]
pub fn blocking_acquire_many_owned(
self: Arc<Self>,
n: u32,
) -> Result<OwnedSemaphorePermit, AcquireError> {
crate::future::block_on(self.acquire_many_owned(n))
}
/// Closes the semaphore.
///
/// This prevents the semaphore from issuing new permits and notifies all pending waiters.
+98
View File
@@ -228,3 +228,101 @@ fn no_panic_at_maxpermits() {
let s = Semaphore::new(Semaphore::MAX_PERMITS - 1);
s.add_permits(1);
}
#[test]
fn blocking_acquire() {
let sem = Semaphore::new(1);
let permit = sem.blocking_acquire().unwrap();
assert_eq!(sem.available_permits(), 0);
drop(permit);
assert_eq!(sem.available_permits(), 1);
}
#[test]
#[cfg(not(target_family = "wasm"))] // spawns a thread, which wasm doesn't support
fn blocking_acquire_waits_for_permit() {
let sem = Arc::new(Semaphore::new(0));
let sem2 = sem.clone();
let handle = std::thread::spawn(move || {
// Blocks until a permit becomes available.
let _permit = sem2.blocking_acquire().unwrap();
});
sem.add_permits(1);
handle.join().unwrap();
}
#[test]
fn blocking_acquire_many() {
let sem = Semaphore::new(5);
let permit = sem.blocking_acquire_many(3).unwrap();
assert_eq!(sem.available_permits(), 2);
drop(permit);
assert_eq!(sem.available_permits(), 5);
}
#[test]
fn blocking_acquire_owned() {
let sem = Arc::new(Semaphore::new(1));
let permit = sem.clone().blocking_acquire_owned().unwrap();
assert_eq!(sem.available_permits(), 0);
drop(permit);
assert_eq!(sem.available_permits(), 1);
}
#[test]
fn blocking_acquire_many_owned() {
let sem = Arc::new(Semaphore::new(5));
let permit = sem.clone().blocking_acquire_many_owned(3).unwrap();
assert_eq!(sem.available_permits(), 2);
drop(permit);
assert_eq!(sem.available_permits(), 5);
}
#[test]
fn blocking_acquire_closed() {
let sem = Arc::new(Semaphore::new(1));
sem.close();
assert!(sem.blocking_acquire().is_err());
assert!(sem.blocking_acquire_many(2).is_err());
assert!(sem.clone().blocking_acquire_owned().is_err());
assert!(sem.clone().blocking_acquire_many_owned(2).is_err());
}
#[tokio::test]
#[cfg(feature = "full")]
#[should_panic = "Cannot block the current thread from within a runtime"]
async fn blocking_acquire_in_async_context() {
let sem = Semaphore::new(1);
// Calling a blocking method from an async context must panic.
let _permit = sem.blocking_acquire();
}
#[tokio::test]
#[cfg(feature = "full")]
#[should_panic = "Cannot block the current thread from within a runtime"]
async fn blocking_acquire_many_in_async_context() {
let sem = Semaphore::new(1);
// Calling a blocking method from an async context must panic.
let _permit = sem.blocking_acquire_many(1);
}
#[tokio::test]
#[cfg(feature = "full")]
#[should_panic = "Cannot block the current thread from within a runtime"]
async fn blocking_acquire_owned_in_async_context() {
let sem = Arc::new(Semaphore::new(1));
// Calling a blocking method from an async context must panic.
let _permit = sem.blocking_acquire_owned();
}
#[tokio::test]
#[cfg(feature = "full")]
#[should_panic = "Cannot block the current thread from within a runtime"]
async fn blocking_acquire_many_owned_in_async_context() {
let sem = Arc::new(Semaphore::new(1));
// Calling a blocking method from an async context must panic.
let _permit = sem.blocking_acquire_many_owned(1);
}