      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    271
    272
    273
    274
    275
    276
    277
    278
    279
    280
    281
    282
    283
    284
    285
    286
    287
    288
    289
    290
    291
    292
    293
    294
    295
    296
    297
    298
    299
    300
    301
    302
    303
    304
    305
    306
    307
    308
    309
    310
    311
    312
    313

    use super::batch_semaphore as ll; // low level implementation
    use super::{AcquireError, TryAcquireError};
    use std::sync::Arc;

    /// Counting semaphore performing asynchronous permit acquisition.
    ///
    /// A semaphore maintains a set of permits. Permits are used to synchronize
    /// access to a shared resource. A semaphore differs from a mutex in that it
    /// can allow more than one concurrent caller to access the shared resource at a
    /// time.
    ///
    /// When `acquire` is called and the semaphore has remaining permits, the
    /// function immediately returns a permit. However, if no remaining permits are
    /// available, `acquire` (asynchronously) waits until an outstanding permit is
    /// dropped. At this point, the freed permit is assigned to the caller.
    ///
    /// To use the `Semaphore` in a poll function, you can use the [`PollSemaphore`]
    /// utility.
    ///
    /// [`PollSemaphore`]: https://docs.rs/tokio-util/0.6/tokio_util/sync/struct.PollSemaphore.html
    #[derive(Debug)]
    pub struct Semaphore {
        /// The low level semaphore
        ll_sem: ll::Semaphore,
    }

    /// A permit from the semaphore.
    ///
    /// This type is created by the [`acquire`] method.
    ///
    /// [`acquire`]: crate::sync::Semaphore::acquire()
    #[must_use]
    #[derive(Debug)]
    pub struct SemaphorePermit<'a> {
        sem: &'a Semaphore,
        permits: u32,
    }

    /// An owned permit from the semaphore.
    ///
    /// This type is created by the [`acquire_owned`] method.
    ///
    /// [`acquire_owned`]: crate::sync::Semaphore::acquire_owned()
    #[must_use]
    #[derive(Debug)]
    pub struct OwnedSemaphorePermit {
        sem: Arc<Semaphore>,
        permits: u32,
    }

    #[test]
    #[cfg(not(loom))]
    fn bounds() {
        fn check_unpin<T: Unpin>() {}
        // This has to take a value, since the async fn's return type is unnameable.
        fn check_send_sync_val<T: Send + Sync>(_t: T) {}
        fn check_send_sync<T: Send + Sync>() {}
        check_unpin::<Semaphore>();
        check_unpin::<SemaphorePermit<'_>>();
        check_send_sync::<Semaphore>();

        let semaphore = Semaphore::new(0);
        check_send_sync_val(semaphore.acquire());
    }

    impl Semaphore {
        /// Creates a new semaphore with the initial number of permits.
        pub fn new(permits: usize) -> Self {
            Self {
                ll_sem: ll::Semaphore::new(permits),
            }
        }

        /// Creates a new semaphore with the initial number of permits.
        #[cfg(all(feature = "parking_lot", not(all(loom, test))))]
        #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))]
        pub const fn const_new(permits: usize) -> Self {
            Self {
                ll_sem: ll::Semaphore::const_new(permits),
            }
        }

        /// Returns the current number of available permits.
        pub fn available_permits(&self) -> usize {
            self.ll_sem.available_permits()
        }

        /// Adds `n` new permits to the semaphore.
        ///
        /// The maximum number of permits is `usize::MAX >> 3`, and this function will panic if the limit is exceeded.
        pub fn add_permits(&self, n: usize) {
            self.ll_sem.release(n);
        }

        /// Acquires a permit from the semaphore.
        ///
        /// If the semaphore has been closed, this returns an [`AcquireError`].
        /// Otherwise, this returns a [`SemaphorePermit`] representing the
        /// acquired permit.
        ///
        /// [`AcquireError`]: crate::sync::AcquireError
        /// [`SemaphorePermit`]: crate::sync::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.
        ///
        /// If the semaphore has been closed, this returns an [`AcquireError`].
        /// Otherwise, this returns a [`SemaphorePermit`] representing the
        /// acquired permits.
        ///
        /// [`AcquireError`]: crate::sync::AcquireError
        /// [`SemaphorePermit`]: crate::sync::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.
        ///
        /// If the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
        /// and a [`TryAcquireError::NoPermits`] if there are no permits left. Otherwise,
        /// this returns a [`SemaphorePermit`] representing the acquired permits.
        ///
        /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
        /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
        /// [`SemaphorePermit`]: crate::sync::SemaphorePermit
        pub fn try_acquire(&self) -> Result<SemaphorePermit<'_>, TryAcquireError> {
            match self.ll_sem.try_acquire(1) {
                Ok(_) => Ok(SemaphorePermit {
                    sem: self,
                    permits: 1,
                }),
                Err(e) => Err(e),
            }
        }

        /// Tries to acquire `n` permits from the semaphore.
        ///
        /// If the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
        /// and a [`TryAcquireError::NoPermits`] if there are no permits left. Otherwise,
        /// this returns a [`SemaphorePermit`] representing the acquired permits.
        ///
        /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
        /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
        /// [`SemaphorePermit`]: crate::sync::SemaphorePermit
        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(e) => Err(e),
            }
        }

        /// Acquires a permit from the semaphore.
        ///
        /// 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 a [`OwnedSemaphorePermit`] representing the
        /// acquired permit.
        ///
        /// [`Arc`]: std::sync::Arc
        /// [`AcquireError`]: crate::sync::AcquireError
        /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
        pub async fn acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, AcquireError> {
            self.ll_sem.acquire(1).await?;
            Ok(OwnedSemaphorePermit {
                sem: self,
                permits: 1,
            })
        }

        /// Acquires `n` permits from the semaphore.
        ///
        /// 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 a [`OwnedSemaphorePermit`] representing the
        /// acquired permit.
        ///
        /// [`Arc`]: std::sync::Arc
        /// [`AcquireError`]: crate::sync::AcquireError
        /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
        pub async fn acquire_many_owned(
            self: Arc<Self>,
            n: u32,
        ) -> Result<OwnedSemaphorePermit, AcquireError> {
            self.ll_sem.acquire(n).await?;
            Ok(OwnedSemaphorePermit {
                sem: self,
                permits: n,
            })
        }

        /// Tries to acquire a permit from the semaphore.
        ///
        /// The semaphore must be wrapped in an [`Arc`] to call this method. If
        /// the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
        /// and a [`TryAcquireError::NoPermits`] if there are no permits left.
        /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the
        /// acquired permit.
        ///
        /// [`Arc`]: std::sync::Arc
        /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
        /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
        /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
        pub fn try_acquire_owned(self: Arc<Self>) -> Result<OwnedSemaphorePermit, TryAcquireError> {
            match self.ll_sem.try_acquire(1) {
                Ok(_) => Ok(OwnedSemaphorePermit {
                    sem: self,
                    permits: 1,
                }),
                Err(e) => Err(e),
            }
        }

        /// Tries to acquire `n` permits from the semaphore.
        ///
        /// The semaphore must be wrapped in an [`Arc`] to call this method. If
        /// the semaphore has been closed, this returns a [`TryAcquireError::Closed`]
        /// and a [`TryAcquireError::NoPermits`] if there are no permits left.
        /// Otherwise, this returns a [`OwnedSemaphorePermit`] representing the
        /// acquired permit.
        ///
        /// [`Arc`]: std::sync::Arc
        /// [`TryAcquireError::Closed`]: crate::sync::TryAcquireError::Closed
        /// [`TryAcquireError::NoPermits`]: crate::sync::TryAcquireError::NoPermits
        /// [`OwnedSemaphorePermit`]: crate::sync::OwnedSemaphorePermit
        pub fn try_acquire_many_owned(
            self: Arc<Self>,
            n: u32,
        ) -> Result<OwnedSemaphorePermit, TryAcquireError> {
            match self.ll_sem.try_acquire(n) {
                Ok(_) => Ok(OwnedSemaphorePermit {
                    sem: self,
                    permits: n,
                }),
                Err(e) => Err(e),
            }
        }

        /// Closes the semaphore.
        ///
        /// This prevents the semaphore from issuing new permits and notifies all pending waiters.
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::sync::Semaphore;
        /// use std::sync::Arc;
        /// use tokio::sync::TryAcquireError;
        ///
        /// #[tokio::main]
        /// async fn main() {
        ///     let semaphore = Arc::new(Semaphore::new(1));
        ///     let semaphore2 = semaphore.clone();
        ///
        ///     tokio::spawn(async move {
        ///         let permit = semaphore.acquire_many(2).await;
        ///         assert!(permit.is_err());
        ///         println!("waiter received error");
        ///     });
        ///
        ///     println!("closing semaphore");
        ///     semaphore2.close();
        ///
        ///     // Cannot obtain more permits
        ///     assert_eq!(semaphore2.try_acquire().err(), Some(TryAcquireError::Closed))
        /// }
        /// ```
        pub fn close(&self) {
            self.ll_sem.close();
        }
    }

    impl<'a> SemaphorePermit<'a> {
        /// Forgets the permit **without** releasing it back to the semaphore.
        /// This can be used to reduce the amount of permits available from a
        /// semaphore.
        pub fn forget(mut self) {
            self.permits = 0;
        }
    }

    impl OwnedSemaphorePermit {
        /// Forgets the permit **without** releasing it back to the semaphore.
        /// This can be used to reduce the amount of permits available from a
        /// semaphore.
        pub fn forget(mut self) {
            self.permits = 0;
        }
    }

    impl<'a> Drop for SemaphorePermit<'_> {
        fn drop(&mut self) {
            self.sem.add_permits(self.permits as usize);
        }
    }

    impl Drop for OwnedSemaphorePermit {
        fn drop(&mut self) {
            self.sem.add_permits(self.permits as usize);
        }
    }