Repository: tokio-rs/tokio
File: tokio-util/src/sync/poll_semaphore.rs (Rust)

Code before:
    ///
    /// 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>> {
        self.poll_acquire_many(1, cx)
    }

    /// Poll to acquire many permits 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 the permits become 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_many(
        &mut self,
        permits: u32,
        cx: &mut Context<'_>,
    ) -> Poll<Option<OwnedSemaphorePermit>> {
        let permit_future = match self.permit_fut.as_mut() {
            Some((prev_permits, fut)) if *prev_permits == permits => fut,
            Some((old_permits, fut_box)) => {
                // We're requesting a different number of permits, so replace the future
                // and record the new amount.
                let fut = Arc::clone(&self.semaphore).acquire_many_owned(permits);
                fut_box.set(fut);
                *old_permits = permits;
                fut_box
            }
            None => {
                // avoid allocations completely if we can grab a permit immediately
                match Arc::clone(&self.semaphore).try_acquire_many_owned(permits) {
                    Ok(permit) => return Poll::Ready(Some(permit)),
                    Err(TryAcquireError::Closed) => return Poll::Ready(None),
                    Err(TryAcquireError::NoPermits) => {}
                }

                let next_fut = Arc::clone(&self.semaphore).acquire_many_owned(permits);
                &mut self
                    .permit_fut
                    .get_or_insert((permits, ReusableBoxFuture::new(next_fut)))
                    .1

Review comment (suggestion):
The context is always the first argument. See e.g. the IO traits for examples of this.
```suggestion
    pub fn poll_acquire_many(
        &mut self,
        cx: &mut Context<'_>,
        permits: u32,
```

Code after:
    ///
    /// 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>> {
        self.poll_acquire_many(cx, 1)
    }

    /// Poll to acquire many permits 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 the permits become 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_many(
        &mut self,
        cx: &mut Context<'_>,
        permits: u32,
    ) -> Poll<Option<OwnedSemaphorePermit>> {
        let permit_future = match self.permit_fut.as_mut() {
            Some((prev_permits, fut)) if *prev_permits == permits => fut,
            Some((old_permits, fut_box)) => {
                // We're requesting a different number of permits, so replace the future
                // and record the new amount.
                let fut = Arc::clone(&self.semaphore).acquire_many_owned(permits);
                fut_box.set(fut);
                *old_permits = permits;
                fut_box
            }
            None => {
                // avoid allocations completely if we can grab a permit immediately
                match Arc::clone(&self.semaphore).try_acquire_many_owned(permits) {
                    Ok(permit) => return Poll::Ready(Some(permit)),
                    Err(TryAcquireError::Closed) => return Poll::Ready(None),
                    Err(TryAcquireError::NoPermits) => {}
                }

                let next_fut = Arc::clone(&self.semaphore).acquire_many_owned(permits);
                &mut self
                    .permit_fut
                    .get_or_insert((permits, ReusableBoxFuture::new(next_fut)))
                    .1