diff --git a/tokio/src/sync/semaphore.rs b/tokio/src/sync/semaphore.rs index 8b23aa396..b4c2e3a92 100644 --- a/tokio/src/sync/semaphore.rs +++ b/tokio/src/sync/semaphore.rs @@ -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, 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, 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) -> Result { + 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, + n: u32, + ) -> Result { + 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. diff --git a/tokio/tests/sync_semaphore.rs b/tokio/tests/sync_semaphore.rs index d976f1e68..f204f85c4 100644 --- a/tokio/tests/sync_semaphore.rs +++ b/tokio/tests/sync_semaphore.rs @@ -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); +}