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

Code before:
    #[must_use = "futures do nothing unless polled"]
    pub struct WaitForCancellationFuture<'a> {
        cancellation_token: &'a CancellationToken,
        #[pin]
        future: tokio::sync::futures::Notified<'a>,
    }
}

pin_project! {
    /// A Future that is resolved once the corresponding [`CancellationToken`]
    /// is cancelled.
    ///
    /// This is the counterpart to [`WaitForCancellationFuture`] that takes
    /// [`CancellationToken`] by value instead of using a reference.
    #[must_use = "futures do nothing unless polled"]
    pub struct WaitForCancellationFutureOwned {
        // `Notified` references to the cancellation_token internally,
        // but camouflages the relationship with `'static`.
        // This can lead to Undefined Behavior, breaking the pointer aliasing rules
        // if the reference is dangling.
        // 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.

Review comment (bug):
Hmm, I think this could be better phrased like this:
```text
This field internally has a reference to the cancellation token, but camouflages
the relationship with `'static`. To avoid undefined behavior, we must ensure
that the reference is only used while the cancellation token is still alive. To
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.
```

Code after:
    /// is cancelled.
    #[must_use = "futures do nothing unless polled"]
    pub struct WaitForCancellationFuture<'a> {
        cancellation_token: &'a CancellationToken,
        #[pin]
        future: tokio::sync::futures::Notified<'a>,
    }
}

pin_project! {
    /// A Future that is resolved once the corresponding [`CancellationToken`]
    /// is cancelled.
    ///
    /// This is the counterpart to [`WaitForCancellationFuture`] that takes
    /// [`CancellationToken`] by value instead of using a reference.
    #[must_use = "futures do nothing unless polled"]
    pub struct WaitForCancellationFutureOwned {
        // This field internally has a reference to the cancellation token, but camouflages
        // the relationship with `'static`. To avoid Undefined Behavior, we must ensure
        // that the reference is only used while the cancellation token is still alive. To
        // 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.