Expose Semaphore::close

tokio/src/sync/batch_semaphore.rs

@@ -41,15 +41,21 @@ struct Waitlist {
     closed: bool,
 }
 
-/// Error returned by `Semaphore::try_acquire`.
+/// Error returned from the `Semaphore::try_acquire` function.
 #[derive(Debug)]
-pub(crate) enum TryAcquireError {
+pub enum TryAcquireError {
+    /// The semaphore has been closed and cannot issue new permits
     Closed,
+
+    /// The has no available permits.
     NoPermits,
 }
 /// Error returned by `Semaphore::acquire`.
+///
+/// An `acquire` operation can only fail if the semaphore has been
+/// closed.
 #[derive(Debug)]
-pub(crate) struct AcquireError(());
+pub struct AcquireError(());
 
 pub(crate) struct Acquire<'a> {
     node: Waiter,
@@ -164,8 +170,6 @@ impl Semaphore {
 
     /// Closes the semaphore. This prevents the semaphore from issuing new
     /// permits and notifies all pending waiters.
-    // This will be used once the bounded MPSC is updated to use the new
-    // semaphore implementation.
     pub(crate) fn close(&self) {
         let mut waiters = self.waiters.lock();
         // If the semaphore's permits counter has enough permits for an

tokio/src/sync/mod.rs

@@ -443,8 +443,10 @@ cfg_sync! {
     pub mod oneshot;
 
     pub(crate) mod batch_semaphore;
+    pub use batch_semaphore::{AcquireError, TryAcquireError};
+
     mod semaphore;
-    pub use semaphore::{Semaphore, SemaphorePermit, OwnedSemaphorePermit, TryAcquireError};
+    pub use semaphore::{Semaphore, SemaphorePermit, OwnedSemaphorePermit};
 
     mod rwlock;
     pub use rwlock::{RwLock, RwLockReadGuard, RwLockWriteGuard};

tokio/src/sync/semaphore.rs

@@ -1,4 +1,5 @@
 use super::batch_semaphore as ll; // low level implementation
+use super::{AcquireError, TryAcquireError};
 use std::sync::Arc;
 
 /// Counting semaphore performing asynchronous permit acquisition.
@@ -42,15 +43,6 @@ pub struct OwnedSemaphorePermit {
     permits: u32,
 }
 
-/// Error returned from the [`Semaphore::try_acquire`] function.
-///
-/// A `try_acquire` operation can only fail if the semaphore has no available
-/// permits.
-///
-/// [`Semaphore::try_acquire`]: Semaphore::try_acquire
-#[derive(Debug)]
-pub struct TryAcquireError(());
-
 #[test]
 #[cfg(not(loom))]
 fn bounds() {
@@ -96,21 +88,21 @@ impl Semaphore {
     }
 
     /// Acquires permit from the semaphore.
-    pub async fn acquire(&self) -> SemaphorePermit<'_> {
-        self.ll_sem.acquire(1).await.unwrap();
-        SemaphorePermit {
+    pub async fn acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
+        self.ll_sem.acquire(1).await?;
+        Ok(SemaphorePermit {
             sem: &self,
             permits: 1,
-        }
+        })
     }
 
     /// Acquires `n` permits from the semaphore
-    pub async fn acquire_many(&self, n: u32) -> SemaphorePermit<'_> {
-        self.ll_sem.acquire(n).await.unwrap();
-        SemaphorePermit {
+    pub async fn acquire_many(&self, n: u32) -> Result<SemaphorePermit<'_>, AcquireError> {
+        self.ll_sem.acquire(n).await?;
+        Ok(SemaphorePermit {
             sem: &self,
             permits: n,
-        }
+        })
     }
 
     /// Tries to acquire a permit from the semaphore.
@@ -120,7 +112,7 @@ impl Semaphore {
                 sem: self,
                 permits: 1,
             }),
-            Err(_) => Err(TryAcquireError(())),
+            Err(e) => Err(e),
         }
     }
 
@@ -131,7 +123,7 @@ impl Semaphore {
                 sem: self,
                 permits: n,
             }),
-            Err(_) => Err(TryAcquireError(())),
+            Err(e) => Err(e),
         }
     }
 
@@ -140,12 +132,12 @@ impl Semaphore {
     /// The semaphore must be wrapped in an [`Arc`] to call this method.
     ///
     /// [`Arc`]: std::sync::Arc
-    pub async fn acquire_owned(self: Arc<Self>) -> OwnedSemaphorePermit {
-        self.ll_sem.acquire(1).await.unwrap();
-        OwnedSemaphorePermit {
+    pub async fn acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, AcquireError> {
+        self.ll_sem.acquire(1).await?;
+        Ok(OwnedSemaphorePermit {
             sem: self,
             permits: 1,
-        }
+        })
     }
 
     /// Tries to acquire a permit from the semaphore.
@@ -159,9 +151,16 @@ impl Semaphore {
                 sem: self,
                 permits: 1,
             }),
-            Err(_) => Err(TryAcquireError(())),
+            Err(e) => Err(e),
         }
     }
+
+    /// Closes the semaphore.
+    ///
+    /// This prevents the semaphore from issuing new permits and notifies all pending waiters.
+    pub fn close(&self) {
+        self.ll_sem.close();
+    }
 }
 
 impl<'a> SemaphorePermit<'a> {