kithara_platform/common/cancel/
wait.rs

1use std::{
2    fmt,
3    future::Future,
4    pin::Pin,
5    sync::Arc,
6    task::{Context, Poll},
7};
8
9use super::node::{Node, Slot};
10
11/// Future returned by [`CancelToken::cancelled`](super::CancelToken::cancelled).
12/// Resolves once the token's subtree is cancelled. `Unpin`; cancel-safe (drop
13/// unregisters its slot).
14pub struct Cancelled<'a> {
15    node: &'a Arc<Node>,
16    slot: Option<u64>,
17    done: bool,
18}
19
20impl<'a> Cancelled<'a> {
21    pub(super) fn new(node: &'a Arc<Node>) -> Self {
22        Self {
23            node,
24            slot: None,
25            done: false,
26        }
27    }
28}
29
30impl Future for Cancelled<'_> {
31    type Output = ();
32
33    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
34        let me = self.get_mut();
35        poll_cancelled(
36            &mut Fields {
37                node: me.node,
38                slot: &mut me.slot,
39                done: &mut me.done,
40            },
41            cx,
42        )
43    }
44}
45
46impl Drop for Cancelled<'_> {
47    fn drop(&mut self) {
48        if let Some(id) = self.slot.take() {
49            self.node.unregister(id);
50        }
51    }
52}
53
54impl fmt::Debug for Cancelled<'_> {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        f.debug_struct("Cancelled").finish_non_exhaustive()
57    }
58}
59
60struct Fields<'a> {
61    node: &'a Arc<Node>,
62    slot: &'a mut Option<u64>,
63    done: &'a mut bool,
64}
65
66fn poll_cancelled(f: &mut Fields<'_>, cx: &mut Context<'_>) -> Poll<()> {
67    if *f.done {
68        return Poll::Ready(());
69    }
70    if f.node.is_cancelled() {
71        if let Some(id) = f.slot.take() {
72            f.node.unregister(id);
73        }
74        *f.done = true;
75        return Poll::Ready(());
76    }
77    if let Some(id) = *f.slot {
78        f.node.refresh_task(id, cx.waker());
79        return Poll::Pending;
80    }
81    // Register, then handle the race: cancel may have fired between the
82    // is_cancelled() above and the registration. register() returns None if the
83    // node already fired (born/late) — resolve at once.
84    if let Some(id) = f.node.register(Slot::Task(cx.waker().clone())) {
85        *f.slot = Some(id);
86        Poll::Pending
87    } else {
88        *f.done = true;
89        Poll::Ready(())
90    }
91}
kithara_platform/common/cancel/wait.rs

kithara_platform/common/cancel/
wait.rs
