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

Code before:
    pub async fn run_until_cancelled<F>(&self, fut: F) -> Option<F::Output>
    where
        F: Future,
    {
        pin_project! {
            /// A Future that is resolved once the corresponding [`CancellationToken`]
            /// is cancelled or a given Future gets resolved. It is biased towards the
            /// Future completion.
            #[must_use = "futures do nothing unless polled"]
            struct RunUntilCancelledFuture<'a, F: Future> {
                #[pin]
                cancellation: WaitForCancellationFuture<'a>,
                #[pin]
                future: F,
            }
        }

        impl<'a, F: Future> Future for RunUntilCancelledFuture<'a, F> {
            type Output = Option<F::Output>;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                let this = self.project();
                if this.cancellation.poll(cx).is_ready() {
                    Poll::Ready(None)
                } else if let Poll::Ready(res) = this.future.poll(cx) {
                    Poll::Ready(Some(res))
                } else {
                    Poll::Pending
                }
            }
        }

        RunUntilCancelledFuture {
            cancellation: self.cancelled(),
            future: fut,
        }
        .await
    }

    /// Runs a future to completion and returns its result wrapped inside of an `Option`
    /// unless the [`CancellationToken`] is cancelled. In that case the function returns
    /// `None` and the future gets dropped.
    ///
    /// The function takes self by value and returns a future that owns the token.
    ///
    /// # Cancellation safety
    ///
    /// This method is only cancel safe if `fut` is cancel safe.
    pub async fn run_until_cancelled_owned<F>(self, fut: F) -> Option<F::Output>
    where
        F: Future,

Review comment (performance):
## For reviewers

This affects the performance of the happy path as `this.cancellation.poll(cx)` internally locks the waiter list, this is especially visible for the use case mentioned in this PR.

```rust
loop {
  let Some(item) = shutdown.run_until_cancelled(stream.next()).await else {
     info!("Exit due to shutdown");
     break;
  };
  ...
}
```

However, reliability over performance.

Code after:
    /// For all subsequent polls, in case of concurrent completion and
    /// cancellation, this is biased towards the future completion.
    ///
    /// # Cancellation safety
    ///
    /// This method is only cancel safe if `fut` is cancel safe.
    pub async fn run_until_cancelled<F>(&self, fut: F) -> Option<F::Output>
    where
        F: Future,
    {
        pin_project! {
            /// A Future that is resolved once the corresponding [`CancellationToken`]
            /// is cancelled or a given Future gets resolved. It is biased towards the
            /// Future completion.
            #[must_use = "futures do nothing unless polled"]
            struct RunUntilCancelledFuture<'a, F: Future> {
                #[pin]
                cancellation: WaitForCancellationFuture<'a>,
                #[pin]
                future: F,
            }
        }

        impl<'a, F: Future> Future for RunUntilCancelledFuture<'a, F> {
            type Output = Option<F::Output>;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                let this = self.project();
                if let Poll::Ready(res) = this.future.poll(cx) {
                    Poll::Ready(Some(res))
                } else if this.cancellation.poll(cx).is_ready() {
                    Poll::Ready(None)
                } else {
                    Poll::Pending
                }
            }
        }

        if self.is_cancelled() {
            None
        } else {
            RunUntilCancelledFuture {
                cancellation: self.cancelled(),
                future: fut,
            }
            .await
        }
    }

    /// Runs a future to completion and returns its result wrapped inside of an `Option`
    /// unless the [`CancellationToken`] is cancelled. In that case the function returns