nuclei/sync/cancellation/
wait.rs

1//! The futures of a token: the borrowed and the owned wait for cancellation,
2//! and the race of a future against one.
3
4#![allow(unsafe_code)]
5
6use core::fmt;
7use core::future::Future;
8use core::mem::MaybeUninit;
9use core::pin::Pin;
10use core::task::{Context, Poll};
11
12use super::CancellationToken;
13use crate::sync::event::{Event, EventWait};
14
15/// `fut`, or `None` if the token is cancelled first; `fut` wins a tie.
16pub(super) struct RunUntilCancelled<'a, F: Future> {
17    pub(super) cancellation: WaitForCancellationFuture<'a>,
18    pub(super) future: F,
19}
20
21impl<F: Future> Future for RunUntilCancelled<'_, F> {
22    type Output = Option<F::Output>;
23
24    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
25        // SAFETY: `future` is pinned structurally: it is never moved out of
26        // the struct, the struct has no `Drop`, and it is `Unpin` only when
27        // `F` is. `cancellation` is `Unpin`.
28        let this = unsafe { self.get_unchecked_mut() };
29        // SAFETY: as above.
30        let future = unsafe { Pin::new_unchecked(&mut this.future) };
31        if let Poll::Ready(output) = future.poll(cx) {
32            Poll::Ready(Some(output))
33        } else if Pin::new(&mut this.cancellation).poll(cx).is_ready() {
34            Poll::Ready(None)
35        } else {
36            Poll::Pending
37        }
38    }
39}
40
41/// A future that resolves once the corresponding [`CancellationToken`] is
42/// cancelled.
43///
44/// Created by [`CancellationToken::cancelled`].
45#[must_use = "futures do nothing unless polled"]
46pub struct WaitForCancellationFuture<'a> {
47    pub(super) wait: EventWait<'a>,
48}
49
50impl fmt::Debug for WaitForCancellationFuture<'_> {
51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52        f.debug_struct("WaitForCancellationFuture").finish()
53    }
54}
55
56impl Future for WaitForCancellationFuture<'_> {
57    type Output = ();
58
59    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
60        Pin::new(&mut self.get_mut().wait).poll(cx)
61    }
62}
63
64/// Tells the compiler its contents may dangle: a reference inside is not
65/// assumed valid for as long as a function argument holding it is in scope,
66/// which the owned future needs because it drops the node its reference
67/// points into while it is itself being dropped.
68#[repr(transparent)]
69struct MaybeDangling<T>(MaybeUninit<T>);
70
71impl<T> MaybeDangling<T> {
72    fn new(inner: T) -> Self {
73        Self(MaybeUninit::new(inner))
74    }
75
76    fn get_mut(&mut self) -> &mut T {
77        // SAFETY: the cell is initialised by `new` and dropped only by `Drop`.
78        unsafe { self.0.assume_init_mut() }
79    }
80}
81
82impl<T> Drop for MaybeDangling<T> {
83    fn drop(&mut self) {
84        // SAFETY: the cell is initialised and dropped only here.
85        unsafe { self.0.assume_init_drop() };
86    }
87}
88
89/// A future that resolves once the corresponding [`CancellationToken`] is
90/// cancelled, and owns the token.
91///
92/// Created by [`CancellationToken::cancelled_owned`].
93#[must_use = "futures do nothing unless polled"]
94pub struct WaitForCancellationFutureOwned {
95    // The wait borrows the node `token` keeps alive, with the borrow's
96    // lifetime erased to `'static`. It is declared first so it drops first:
97    // it must never be used or dropped after `token` is gone.
98    wait: MaybeDangling<EventWait<'static>>,
99    _token: CancellationToken,
100}
101
102impl WaitForCancellationFutureOwned {
103    pub(super) fn new(token: CancellationToken) -> Self {
104        let event: *const Event = &token.node.cancelled;
105        // SAFETY: the event lives inside the heap allocation `token` owns, so
106        // it does not move when the future does and lives as long as `token`,
107        // which is a field of the future that drops after `wait`.
108        let wait = unsafe { (*event).wait() };
109        Self {
110            wait: MaybeDangling::new(wait),
111            _token: token,
112        }
113    }
114}
115
116impl fmt::Debug for WaitForCancellationFutureOwned {
117    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118        f.debug_struct("WaitForCancellationFutureOwned").finish()
119    }
120}
121
122impl Future for WaitForCancellationFutureOwned {
123    type Output = ();
124
125    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
126        Pin::new(self.get_mut().wait.get_mut()).poll(cx)
127    }
128}
nuclei/sync/cancellation/wait.rs

nuclei/sync/cancellation/
wait.rs
