Rewrite tokio-util::sync::CancellationToken

[Finomnis]

Fixes leak from #4639.

Motivation

CancellationToken uses a lot of unsafe code and seems to be error-prone.

Solution

@Darksonn suggested a rewrite.

After a discussion on discord this is my attempt to combine and polish our ideas.

A thorough review will be necessary.

State

• Architectural discussion on discord
• Rewrite attempt
• Close TODOs in code
• Clippy
• More tests? (Related: Enable ASAN tests in CI #4645)
• Thorough review of @Darksonn

[Darksonn]

Overall this looks good.

[Ten0]

I realized that there is an alternate way to implement the future. I think this implementation is slightly simpler, and it has the advantage that it's impossible to first get a true from token.is_cancelled() followed by getting a Pending from polling the future.

pin_project! {
    /// A Future that is resolved once the corresponding [`CancellationToken`]
    /// is cancelled.
    #[must_use = "futures do nothing unless polled"]
    pub struct WaitForCancellationFuture<'a> {
        cancellation_token: &'a CancellationToken,
        #[pin]
        future: tokio::sync::futures::Notified<'a>,
    }
}

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(());
            }

            if this.future.as_mut().poll(cx).is_pending() {
                return Poll::Pending;
            }

            this.future.set(this.cancellation_token.inner.waker.notified());
        }
    }
}

impl CancellationToken {
    pub fn cancelled(&self) -> WaitForCancellationFuture<'_> {
        WaitForCancellationFuture {
            cancellation_token: self,
            future: self.inner.waker.notified(),
        }
    }
}

// and just delete tree_node::get_future

I'm struggling to understand the point of this.future.set(this.cancellation_token.inner.waker.notified());.
I'm not sure why this can't be written as:

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

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

        this.future.as_mut().poll(cx)
    }
}

As far as I can tell this would preserve the following properties:

• We can't get a is_cancelled then have poll return Pending
• We can't get a poll return Ready then have is_cancelled = false because it's set before
• 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 the Notified.

Would anyone be able to provide a more in-depth explanation? Thanks in advance.

[NobodyXu]

@Ten0 Because once future returns Poll::Ready it usually cannot be polled again, otherwise it might panic.

Previous implementation always return Poll::Ready(()) if the token is akready cancelled and changing this might be a breaking change.

[Ten0]

Wouldn't property 2 "We can't get a poll return Ready then have is_cancelled = false because it's set before" guarantee that it would still return ready immediately again without re polling the inner future?

[NobodyXu]

Wouldn't property 2 "We can't get a poll return Ready then have is_cancelled = false because it's set before" guarantee that it would still return ready immediately again without re polling the inner future?

You are right, you can indeed rewrite it and remove the loop.

[Finomnis]

I remember having had this discussion already, and there were reasons. @Darksonn might have to pitch in, she explained it to me back then

[Finomnis]

I think the reason was that Notified is susceptible to spurious wakeups. So getting a Ready from it doesn't guarantee that the CancellationToken is actually cancelled, so we decided that the only positive return path of this function should be through an is_cancelled check. To make it easier to prove correctness.

[Finomnis]

Regarding the order: would you mind explaining how it would be possible to first get is_cancelled and then afterwards get Pending? I can't see this flow path right now...

[Ten0]

I think the reason was that Notified is susceptible to spurious wakeups.

I can't see anywhere that it's susceptible to spurious wake ups though (unless those are explicitly triggered by the tree implem).
That being said if indeed Notified is susceptible to such (explicit?) spurious wakeups and cant't be re-polled after it has been notified then I understand this.

Thanks! :)

would you mind explaining how it would be possible to first get is_cancelled and then afterwards get Pending? I can't see this flow path right now...

My original message was saying the alternate implementation seems to preserve the properties that this still can't happen (for the reasons in the comment in the code). I'm not sure what you're referring to.

[Finomnis]

Ah OK never mind then, I must have misunderstood your original comment then :)

To my understanding Notified is meant to signal that some information changed and had to be re-evaluated. It is not meant to actually propagate said information.

[Ten0]

To my understanding Notified is meant to signal that some information changed and had to be re-evaluated. It is not meant to actually propagate said information.

When said information is a boolean that can only go from false to true that might be enough. But I get the will for consistency.

[Darksonn]

The implementation of Notified can indeed result in spurious wakeups. I don't think the cases where it can happen actually matter for this particular use-case, but I prefer to not rely on that implementation detail.

[Ten0]

Alright, thanks! Might be worth mentioning those in the documentation then (did I miss them?).