# poll_semaphore.rs - source
Source of the Rust file `/home/runner/.cargo/registry/src/github.com-1ecc6299db9ec823/tokio-util-0.6.5/src/sync/poll_semaphore.rs`.
Keywords: rust, rustlang, rust-lang

     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

    use futures_core::{ready, Stream};
    use std::fmt;
    use std::pin::Pin;
    use std::sync::Arc;
    use std::task::{Context, Poll};
    use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore};

    use super::ReusableBoxFuture;

    /// A wrapper around [`Semaphore`] that provides a `poll_acquire` method.
    ///
    /// [`Semaphore`]: tokio::sync::Semaphore
    pub struct PollSemaphore {
        semaphore: Arc<Semaphore>,
        permit_fut: ReusableBoxFuture<Result<OwnedSemaphorePermit, AcquireError>>,
    }

    impl PollSemaphore {
        /// Create a new `PollSemaphore`.
        pub fn new(semaphore: Arc<Semaphore>) -> Self {
            let fut = Arc::clone(&semaphore).acquire_owned();

            Self {
                semaphore,
                permit_fut: ReusableBoxFuture::new(fut),
            }
        }

        /// Closes the semaphore.
        pub fn close(&self) {
            self.semaphore.close()
        }

        /// Obtain a clone of the inner semaphore.
        pub fn clone_inner(&self) -> Arc<Semaphore> {
            self.semaphore.clone()
        }

        /// Get back the inner semaphore.
        pub fn into_inner(self) -> Arc<Semaphore> {
            self.semaphore
        }

        /// Poll to acquire a permit from the semaphore.
        ///
        /// This can return the following values:
        ///
        ///  - `Poll::Pending` if a permit is not currently available.
        ///  - `Poll::Ready(Some(permit))` if a permit was acquired.
        ///  - `Poll::Ready(None)` if the semaphore has been closed.
        ///
        /// When this method returns `Poll::Pending`, the current task is scheduled
        /// to receive a wakeup when a permit becomes available, or when the
        /// semaphore is closed. Note that on multiple calls to `poll_acquire`, only
        /// the `Waker` from the `Context` passed to the most recent call is
        /// scheduled to receive a wakeup.
        pub fn poll_acquire(&mut self, cx: &mut Context<'_>) -> Poll<Option<OwnedSemaphorePermit>> {
            let result = ready!(self.permit_fut.poll(cx));

            let next_fut = Arc::clone(&self.semaphore).acquire_owned();
            self.permit_fut.set(next_fut);

            match result {
                Ok(permit) => Poll::Ready(Some(permit)),
                Err(_closed) => Poll::Ready(None),
            }
        }
    }

    impl Stream for PollSemaphore {
        type Item = OwnedSemaphorePermit;

        fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<OwnedSemaphorePermit>> {
            Pin::into_inner(self).poll_acquire(cx)
        }
    }

    impl Clone for PollSemaphore {
        fn clone(&self) -> PollSemaphore {
            PollSemaphore::new(self.clone_inner())
        }
    }

    impl fmt::Debug for PollSemaphore {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("PollSemaphore")
                .field("semaphore", &self.semaphore)
                .finish()
        }
    }