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

Code before:
        // This is fixed here by the use of `MaybeDanglingFuture`, which prevent
        // the compiler to generate the UB.
        // Since `future` is the first field, it is dropped before the
        // cancellation_token field. This ensures that the reference inside the
        // `Notified` remains valid, hance the safety requirements of
        // `MaybeDanglingFuture` are satisfied.
        #[pin]
        future: MaybeDanglingFuture<tokio::sync::futures::Notified<'static>>,
        cancellation_token: CancellationToken,
    }
}

/// A wrapper type that avoids implicitly asserting the contents are currently valid.
/// It's safe to contain a dangling reference, as long as it is not used.
///
/// # Safety
///
/// The inner type must be dropped before the data it references to.
///
// Reference
// See <https://users.rust-lang.org/t/unsafe-code-review-semi-owning-weak-rwlock-t-guard/95706>
//
// TODO: replace this with an official solution once RFC #3336 or similar is available.
// <https://github.com/rust-lang/rfcs/pull/3336>
#[repr(transparent)]
struct MaybeDanglingFuture<T>(core::mem::MaybeUninit<T>);

impl<T> Drop for MaybeDanglingFuture<T> {
    fn drop(&mut self) {
        // Safety: `0` is always initialized.
        // It may be dangling, but it's the caller responsibility to ensure the data is valid.
        unsafe { core::ptr::drop_in_place(self.0.as_mut_ptr()) };
    }
}

impl<T> MaybeDanglingFuture<T> {
    fn new(inner: T) -> Self {
        Self(core::mem::MaybeUninit::new(inner))
    }
}

impl<F: Future> Future for MaybeDanglingFuture<F> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // Safety: `0` is always initialized.
        // It may be dangling, but it's the caller responsibility to ensure the data is valid.
        let fut = unsafe { self.map_unchecked_mut(|this| this.0.assume_init_mut()) };
        fut.poll(cx)
    }
}

Review comment (refactor):
I would move this into another file and just call it `MaybeDangling`. (But I wouldn't remove the future impl.)

Code after:
        // do that, we ensure that the future is the first field, so that it is dropped
        // before the cancellation token.
        //
        // We use `MaybeDanglingFuture` here because without it, the compiler could assert
        // the reference inside `future` to be valid even after the destructor of that
        // field runs. (Specifically, when the `WaitForCancellationFutureOwned` is passed
        // as an argument to a function, the reference can be asserted to be valid for the
        // rest of that function.) To avoid that, we use `MaybeDangling` which tells the
        // compiler that the reference stored inside it might not be valid.
        //
        // See <https://users.rust-lang.org/t/unsafe-code-review-semi-owning-weak-rwlock-t-guard/95706>
        // for more info.
        #[pin]
        future: MaybeDangling<tokio::sync::futures::Notified<'static>>,
        cancellation_token: CancellationToken,
    }
}

// ===== impl CancellationToken =====

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

impl Clone for CancellationToken {
    /// Creates a clone of the `CancellationToken` which will get cancelled
    /// whenever the current token gets cancelled, and vice versa.
    fn clone(&self) -> Self {
        tree_node::increase_handle_refcount(&self.inner);
        CancellationToken {
            inner: self.inner.clone(),
        }
    }
}

impl Drop for CancellationToken {
    fn drop(&mut self) {
        tree_node::decrease_handle_refcount(&self.inner);
    }
}

impl Default for CancellationToken {
    fn default() -> CancellationToken {
        CancellationToken::new()
    }
}