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

Code before:
    {
        RunUntilCancelledFuture {
            cancellation_token: self,
            future: fut,
        }
        .await
    }
}

// ===== impl RunUntilCancelledFuture =====

impl<'a, F: Future> core::fmt::Debug for RunUntilCancelledFuture<'a, F> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("RunUntilCancelledFuture").finish()
    }
}

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 mut this = self.project();
        if let Poll::Ready(res) = this.future.as_mut().poll(cx) {
            Poll::Ready(Some(res))
        } else if this.cancellation_token.is_cancelled() {
            Poll::Ready(None)
        } else {
            Poll::Pending
        }
    }
}

// ===== impl WaitForCancellationFuture =====

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

impl<'a> Future for WaitForCancellationFuture<'a> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let mut this = self.project();
        loop {
            if this.cancellation_token.is_cancelled() {
                return Poll::Ready(());
            }

            // No wakeups can be lost here because there is always a call to

Review comment (suggestion):
You probably want `poll` to get called when the token is cancelled, but that won't happen with this implementation. To achieve that, you should store a `WaitForCancellationFuture` inside `RunUntilCancelledFuture` and poll it instead of using `is_cancelled`.

Code after:
            }
        }

        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
                }
            }
        }

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

// ===== impl WaitForCancellationFuture =====

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

impl<'a> Future for WaitForCancellationFuture<'a> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let mut this = self.project();
        loop {
            if this.cancellation_token.is_cancelled() {
                return Poll::Ready(());
            }

            // No wakeups can be lost here because there is always a call to
            // `is_cancelled` between the creation of the future and the call to
            // `poll`, and the code that sets the cancelled flag does so before
            // waking the `Notified`.
            if this.future.as_mut().poll(cx).is_pending() {
                return Poll::Pending;
            }