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

Code before:
    }
}

// ===== impl WaitForCancellationFutureOwned =====

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

impl WaitForCancellationFutureOwned {
    fn new(cancellation_token: CancellationToken) -> Self {
        let future = cancellation_token.inner.notified();

        // Safety:
        //  - We transmute notified into Notified<'static>
        //    since there's no 'self
        //    Since future is declared before cancellation_token,
        //    it will be dropped before cancellation_token so it
        //    will always holds a valid reference.
        //  - this.cancellation_token contains an Arc to TreeNode,
        //    which is guaranteed to have stable dereference.
        //    So it's ok to move it here as long as `Notified::enabled`
        //    or `Notified::poll` is not called.
        let future: tokio::sync::futures::Notified<'static> = unsafe { mem::transmute(future) };

        WaitForCancellationFutureOwned {
            future,
            cancellation_token,
        }
    }
}

impl Future for WaitForCancellationFutureOwned {
    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;

Review comment (refactor):
You've duplicated this. Can you put it into a helper method?

Additionally, I have a slight preference for obtaining the future like this:
```Rust
let inner_ptr = Arc::as_ptr(&cancellation_token.inner);
let future = (*inner_ptr).notified();
```
Why? Because the documentation for `Arc::as_ptr` says the following:

> The pointer is valid for as long as there are strong counts in the Arc.

So there is no ambiguity about whether our references remains valid if the Arc is moved or anything like that. Additionally, this would let you avoid the transmute.

Code after:

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

impl WaitForCancellationFutureOwned {
    fn new(cancellation_token: CancellationToken) -> Self {
        WaitForCancellationFutureOwned {
            // cancellation_token holds a heap allocation and is guaranteed to have a
            // stable deref, thus it would be ok to move the cancellation_token while
            // the future holds a reference to it.
            //
            // # Safety
            //
            // cancellation_token is dropped after future due to the field ordering.
            future: unsafe { Self::new_future(&cancellation_token) },
            cancellation_token,
        }
    }

    /// # Safety
    /// The returned future must be destroyed before the cancellation token is
    /// destroyed.
    unsafe fn new_future(
        cancellation_token: &CancellationToken,
    ) -> tokio::sync::futures::Notified<'static> {
        let inner_ptr = Arc::as_ptr(&cancellation_token.inner);
        // SAFETY: The `Arc::as_ptr` method guarantees that `inner_ptr` remains
        // valid until the strong count of the Arc drops to zero, and the caller
        // guarantees that they will drop the future before that happens.
        (*inner_ptr).notified()
    }
}

impl Future for WaitForCancellationFutureOwned {
    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`.