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

Code before:
                //
                // `self` is pinned, so self.future is also pinned.
                .map(|fut| unsafe { Pin::new_unchecked(fut) })
                .map(|fut| fut.poll(cx).is_pending())
                .unwrap_or_default()
            {
                return Poll::Pending;
            }

            let notified = this.cancellation_token.inner.notified();

            // Safety:
            //  - We transmute notified into Notified<'static>
            //    since there's no 'static
            //  - self is pinned, so it's safe to have self-reference
            //    data structure.
            //  - this.cancellation_token contains an Arc to TreeNode,
            //    which is guaranteed to have stable dereference.
            //    So even if `Notified` implements `Unpin` and `self`
            //    get moved, it would still be valid.
            this.future = Some(unsafe { mem::transmute(notified) });
        }
    }
}

impl WaitForCancellationFutureOwned {
    fn do_drop<'a>(&'a mut self) {
        if let Some(future) = self.future.take() {
            // Safety:
            //
            // The future itself refererences cancellation_token, so its
            // lifetime must be at least as long as 'a.
            let future: tokio::sync::futures::Notified<'a> = unsafe { mem::transmute(future) };
            drop(future);
        }
    }

    // Use explicit lifetime here just to be clear what this function is doing.
    #[allow(clippy::needless_lifetimes)]
    fn get_future_mut<'a>(&'a mut self) -> Option<&mut tokio::sync::futures::Notified<'a>> {
        self.future
            .as_mut()
            // Safety:
            //
            // The future itself references cancellation_token, so its
            // lifetime must be at least as long as 'a.
            .map(|fut| unsafe { mem::transmute(fut) })
    }
}

impl Drop for WaitForCancellationFutureOwned {

Review comment (suggestion):
This moves the `Notified` future, but it is pinned so it may not be moved.

Code after:
    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`.
            if this.future.as_mut().poll(cx).is_pending() {
                return Poll::Pending;
            }

            // # Safety
            //
            // cancellation_token is dropped after future due to the field ordering.
            this.future
                .set(unsafe { Self::new_future(this.cancellation_token) });
        }
    }
}