sync: add CancellationToken (#2263)

As a first step towards structured concurrency, this change adds a
CancellationToken for graceful cancellation of tasks.

The task can be awaited by an arbitrary amount of tasks due to the usage
of an intrusive list.

The token can be cloned. In addition to this child tokens can be derived.
When the parent token gets cancelled, all child tokens will also get
cancelled.

7 files changed, 2139 insertions, 0 deletions

diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml
index 063d208e..63bb7ac0 100644
--- a/tokio/Cargo.toml
+++ b/tokio/Cargo.toml
@@ -123,6 +123,7 @@ optional = true
 [dev-dependencies]
 tokio-test = { version = "0.2.0", path = "../tokio-test" }
 futures = { version = "0.3.0", features = ["async-await"] }
+futures-test = "0.3.0"
 proptest = "0.9.4"
 tempfile = "3.1.0"
 
diff --git a/tokio/src/sync/cancellation_token.rs b/tokio/src/sync/cancellation_token.rs
new file mode 100644
index 00000000..66c302ca
--- /dev/null
+++ b/tokio/src/sync/cancellation_token.rs
@@ -0,0 +1,1186 @@
+//! An asynchronously awaitable `CancellationToken`.
+//! The token allows to signal a cancellation request to one or more tasks.
+
+use crate::{
+    loom::sync::{atomic::AtomicUsize, Mutex},
+    util::intrusive_double_linked_list::{LinkedList, ListNode},
+};
+use core::{
+    future::Future,
+    pin::Pin,
+    ptr::NonNull,
+    sync::atomic::Ordering,
+    task::{Context, Poll, Waker},
+};
+
+/// A token which can be used to signal a cancellation request to one or more
+/// tasks.
+///
+/// Tasks can call [`CancellationToken::cancelled()`] in order to
+/// obtain a Future which will be resolved when cancellation is requested.
+///
+/// Cancellation can be requested through the [`CancellationToken::cancel`] method.
+///
+/// # Examples
+///
+/// ```ignore
+/// use tokio::select;
+/// use tokio::scope::CancellationToken;
+///
+/// #[tokio::main]
+/// async fn main() {
+///     let token = CancellationToken::new();
+///     let cloned_token = token.clone();
+///
+///     let join_handle = tokio::spawn(async move {
+///         // Wait for either cancellation or a very long time
+///         select! {
+///             _ = cloned_token.cancelled() => {
+///                 // The token was cancelled
+///                 5
+///             }
+///             _ = tokio::time::delay_for(std::time::Duration::from_secs(9999)) => {
+///                 99
+///             }
+///         }
+///     });
+///
+///     tokio::spawn(async move {
+///         tokio::time::delay_for(std::time::Duration::from_millis(10)).await;
+///         token.cancel();
+///     });
+///
+///     assert_eq!(5, join_handle.await.unwrap());
+/// }
+/// ```
+pub struct CancellationToken {
+    inner: NonNull<CancellationTokenState>,
+}
+
+// Safety: The CancellationToken is thread-safe and can be moved between threads,
+// since all methods are internally synchronized.
+unsafe impl Send for CancellationToken {}
+unsafe impl Sync for 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 {
+    fn clone(&self) -> Self {
+        // Safety: The state inside a `CancellationToken` is always valid, since
+        // is reference counted
+        let inner = self.state();
+
+        // Tokens are cloned by increasing their refcount
+        let current_state = inner.snapshot();
+        inner.increment_refcount(current_state);
+
+        CancellationToken { inner: self.inner }
+    }
+}
+
+impl Drop for CancellationToken {
+    fn drop(&mut self) {
+        let token_state_pointer = self.inner;
+
+        // Safety: The state inside a `CancellationToken` is always valid, since
+        // is reference counted
+        let inner = unsafe { &mut *self.inner.as_ptr() };
+
+        let mut current_state = inner.snapshot();
+
+        // We need to safe the parent, since the state might be released by the
+        // next call
+        let parent = inner.parent;
+
+        // Drop our own refcount
+        current_state = inner.decrement_refcount(current_state);
+
+        // If this was the last reference, unregister from the parent
+        if current_state.refcount == 0 {
+            if let Some(mut parent) = parent {
+                // Safety: Since we still retain a reference on the parent, it must be valid.
+                let parent = unsafe { parent.as_mut() };
+                parent.unregister_child(token_state_pointer, current_state);
+            }
+        }
+    }
+}
+
+impl CancellationToken {
+    /// Creates a new CancellationToken in the non-cancelled state.
+    #[allow(dead_code)]
+    pub(crate) fn new() -> CancellationToken {
+        let state = Box::new(CancellationTokenState::new(
+            None,
+            StateSnapshot {
+                cancel_state: CancellationState::NotCancelled,
+                has_parent_ref: false,
+                refcount: 1,
+            },
+        ));
+
+        // Safety: We just created the Box. The pointer is guaranteed to be
+        // not null
+        CancellationToken {
+            inner: unsafe { NonNull::new_unchecked(Box::into_raw(state)) },
+        }
+    }
+
+    /// Returns a reference to the utilized `CancellationTokenState`.
+    fn state(&self) -> &CancellationTokenState {
+        // Safety: The state inside a `CancellationToken` is always valid, since
+        // is reference counted
+        unsafe { &*self.inner.as_ptr() }
+    }
+
+    /// Creates a `CancellationToken` which will get cancelled whenever the
+    /// current token gets cancelled.
+    ///
+    /// If the current token is already cancelled, the child token will get
+    /// returned in cancelled state.
+    ///
+    /// # Examples
+    ///
+    /// ```ignore
+    /// use tokio::select;
+    /// use tokio::scope::CancellationToken;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let token = CancellationToken::new();
+    ///     let child_token = token.child_token();
+    ///
+    ///     let join_handle = tokio::spawn(async move {
+    ///         // Wait for either cancellation or a very long time
+    ///         select! {
+    ///             _ = child_token.cancelled() => {
+    ///                 // The token was cancelled
+    ///                 5
+    ///             }
+    ///             _ = tokio::time::delay_for(std::time::Duration::from_secs(9999)) => {
+    ///                 99
+    ///             }
+    ///         }
+    ///     });
+    ///
+    ///     tokio::spawn(async move {
+    ///         tokio::time::delay_for(std::time::Duration::from_millis(10)).await;
+    ///         token.cancel();
+    ///     });
+    ///
+    ///     assert_eq!(5, join_handle.await.unwrap());
+    /// }
+    /// ```
+    #[allow(dead_code)]
+    pub(crate) fn child_token(&self) -> CancellationToken {
+        let inner = self.state();
+
+        // Increment the refcount of this token. It will be referenced by the
+        // child, independent of whether the child is immediately cancelled or
+        // not.
+        let _current_state = inner.increment_refcount(inner.snapshot());
+
+        let mut unpacked_child_state = StateSnapshot {
+            has_parent_ref: true,
+            refcount: 1,
+            cancel_state: CancellationState::NotCancelled,
+        };
+        let mut child_token_state = Box::new(CancellationTokenState::new(
+            Some(self.inner),
+            unpacked_child_state,
+        ));
+
+        {
+            let mut guard = inner.synchronized.lock().unwrap();
+            if guard.is_cancelled {
+                // This task was already cancelled. In this case we should not
+                // insert the child into the list, since it would never get removed
+                // from the list.
+                (*child_token_state.synchronized.lock().unwrap()).is_cancelled = true;
+                unpacked_child_state.cancel_state = CancellationState::Cancelled;
+                // Since it's not in the list, the parent doesn't need to retain
+                // a reference to it.
+                unpacked_child_state.has_parent_ref = false;
+                child_token_state
+                    .state
+                    .store(unpacked_child_state.pack(), Ordering::SeqCst);
+            } else {
+                if let Some(mut first_child) = guard.first_child {
+                    child_token_state.from_parent.next_peer = Some(first_child);
+                    // Safety: We manipulate other child task inside the Mutex
+                    // and retain a parent reference on it. The child token can't
+                    // get invalidated while the Mutex is held.
+                    unsafe {
+                        first_child.as_mut().from_parent.prev_peer =
+                            Some((&mut *child_token_state).into())
+                    };
+                }
+                guard.first_child = Some((&mut *child_token_state).into());
+            }
+        };
+
+        let child_token_ptr = Box::into_raw(child_token_state);
+        // Safety: We just created the pointer from a `Box`
+        CancellationToken {
+            inner: unsafe { NonNull::new_unchecked(child_token_ptr) },
+        }
+    }
+
+    /// Returns the number of child tokens
+    #[cfg(all(test, not(loom)))]
+    fn child_tokens(&self) -> usize {
+        let mut result = 0;
+        let inner = self.state();
+
+        let guard = inner.synchronized.lock().unwrap();
+        let mut child = guard.first_child;
+        while let Some(mut c) = child {
+            result += 1;
+            // Safety: The child state is accessed from within a Mutex. Since
+            // the child needs to take the Mutex to unregister itself before
+            // getting destroyed, it is guaranteed to be alive.
+            child = unsafe { c.as_mut().from_parent.next_peer };
+        }
+
+        result
+    }
+
+    /// Cancel the [`CancellationToken`] and all child tokens which had been
+    /// derived from it.
+    ///
+    /// This will wake up all tasks which are waiting for cancellation.
+    #[allow(dead_code)]
+    pub(crate) fn cancel(&self) {
+        self.state().cancel();
+    }
+
+    /// Returns `true` if the `CancellationToken` had been cancelled
+    pub fn is_cancelled(&self) -> bool {
+        self.state().is_cancelled()
+    }
+
+    /// Returns a `Future` that gets fulfilled when cancellation is requested.
+    pub fn cancelled(&self) -> WaitForCancellationFuture<'_> {
+        WaitForCancellationFuture {
+            cancellation_token: Some(self),
+            wait_node: ListNode::new(WaitQueueEntry::new()),
+            is_registered: false,
+        }
+    }
+
+    unsafe fn register(
+        &self,
+        wait_node: &mut ListNode<WaitQueueEntry>,
+        cx: &mut Context<'_>,
+    ) -> Poll<()> {
+        self.state().register(wait_node, cx)
+    }
+
+    fn check_for_cancellation(
+        &self,
+        wait_node: &mut ListNode<WaitQueueEntry>,
+        cx: &mut Context<'_>,
+    ) -> Poll<()> {
+        self.state().check_for_cancellation(wait_node, cx)
+    }
+
+    fn unregister(&self, wait_node: &mut ListNode<WaitQueueEntry>) {
+        self.state().unregister(wait_node)
+    }
+}
+
+/// A Future that is resolved once the corresponding [`CancellationToken`]
+/// was cancelled
+#[must_use = "futures do nothing unless polled"]
+pub struct WaitForCancellationFuture<'a> {
+    /// The CancellationToken that is associated with this WaitForCancellationFuture
+    cancellation_token: Option<&'a CancellationToken>,
+    /// Node for waiting at the cancellation_token
+    wait_node: ListNode<WaitQueueEntry>,
+    /// Whether this future was registered at the token yet as a waiter
+    is_registered: bool,
+}
+
+// Safety: Futures can be sent between threads as long as the underlying
+// cancellation_token is thread-safe (Sync),
+// which allows to poll/register/unregister from a different thread.
+unsafe impl<'a> Send for WaitForCancellationFuture<'a> {}
+
+impl<'a> core::fmt::Debug for WaitForCancellationFuture<'a> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("WaitForCancellationFuture").finish()
+    }
+}
+
+impl<'a> Future for WaitForCancellationFuture<'a> {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        // Safety: We do not move anything out of `WaitForCancellationFuture`
+        let mut_self: &mut WaitForCancellationFuture<'_> = unsafe { Pin::get_unchecked_mut(self) };
+
+        let cancellation_token = mut_self
+            .cancellation_token
+            .expect("polled WaitForCancellationFuture after completion");
+
+        let poll_res = if !mut_self.is_registered {
+            // Safety: The `ListNode` is pinned through the Future,
+            // and we will unregister it in `WaitForCancellationFuture::drop`
+            // before the Future is dropped and the memory reference is invalidated.
+            unsafe { cancellation_token.register(&mut mut_self.wait_node, cx) }
+        } else {
+            cancellation_token.check_for_cancellation(&mut mut_self.wait_node, cx)
+        };
+
+        if let Poll::Ready(()) = poll_res {
+            // The cancellation_token was signalled
+            mut_self.cancellation_token = None;
+            // A signalled Token means the Waker won't be enqueued anymore
+            mut_self.is_registered = false;
+            mut_self.wait_node.task = None;
+        } else {
+            // This `Future` and its stored `Waker` stay registered at the
+            // `CancellationToken`
+            mut_self.is_registered = true;
+        }
+
+        poll_res
+    }
+}
+
+impl<'a> Drop for WaitForCancellationFuture<'a> {
+    fn drop(&mut self) {
+        // If this WaitForCancellationFuture has been polled and it was added to the
+        // wait queue at the cancellation_token, it must be removed before dropping.
+        // Otherwise the cancellation_token would access invalid memory.
+        if let Some(token) = self.cancellation_token {
+            if self.is_registered {
+                token.unregister(&mut self.wait_node);
+            }
+        }
+    }
+}
+
+/// Tracks how the future had interacted with the [`CancellationToken`]
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+enum PollState {
+    /// The task has never interacted with the [`CancellationToken`].
+    New,
+    /// The task was added to the wait queue at the [`CancellationToken`].
+    Waiting,
+    /// The task has been polled to completion.
+    Done,
+}
+
+/// Tracks the WaitForCancellationFuture waiting state.
+/// Access to this struct is synchronized through the mutex in the CancellationToken.
+struct WaitQueueEntry {
+    /// The task handle of the waiting task
+    task: Option<Waker>,
+    // Current polling state. This state is only updated inside the Mutex of
+    // the CancellationToken.
+    state: PollState,
+}
+
+impl WaitQueueEntry {
+    /// Creates a new WaitQueueEntry
+    fn new() -> WaitQueueEntry {
+        WaitQueueEntry {
+            task: None,
+            state: PollState::New,
+        }
+    }
+}
+
+struct SynchronizedState {
+    waiters: LinkedList<WaitQueueEntry>,
+    first_child: Option<NonNull<CancellationTokenState>>,
+    is_cancelled: bool,
+}
+
+impl SynchronizedState {
+    fn new() -> Self {
+        Self {
+            waiters: LinkedList::new(),
+            first_child: None,
+            is_cancelled: false,
+        }
+    }
+}
+
+/// Information embedded in child tokens which is synchronized through the Mutex
+/// in their parent.
+struct SynchronizedThroughParent {
+    next_peer: Option<NonNull<CancellationTokenState>>,
+    prev_peer: Option<NonNull<CancellationTokenState>>,
+}
+
+/// Possible states of a `CancellationToken`
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+enum CancellationState {
+    NotCancelled = 0,
+    Cancelling = 1,
+    Cancelled = 2,
+}
+
+impl CancellationState {
+    fn pack(self) -> usize {
+        self as usize
+    }
+
+    fn unpack(value: usize) -> Self {
+        match value {
+            0 => CancellationState::NotCancelled,
+            1 => CancellationState::Cancelling,
+            2 => CancellationState::Cancelled,
+            _ => unreachable!("Invalid value"),
+        }
+    }
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+struct StateSnapshot {
+    /// The amount of references to this particular CancellationToken.
+    /// `CancellationToken` structs hold these references to a `CancellationTokenState`.
+    /// Also the state is referenced by the state of each child.
+    refcount: usize,
+    /// Whether the state is still referenced by it's parent and can therefore
+    /// not be freed.
+    has_parent_ref: bool,
+    /// Whether the token is cancelled
+    cancel_state: CancellationState,
+}
+
+impl StateSnapshot {
+    /// Packs the snapshot into a `usize`
+    fn pack(self) -> usize {
+        self.refcount << 3 | if self.has_parent_ref { 4 } else { 0 } | self.cancel_state.pack()
+    }
+
+    /// Unpacks the snapshot from a `usize`
+    fn unpack(value: usize) -> Self {
+        let refcount = value >> 3;
+        let has_parent_ref = value & 4 != 0;
+        let cancel_state = CancellationState::unpack(value & 0x03);
+
+        StateSnapshot {
+            refcount,
+            has_parent_ref,
+            cancel_state,
+        }
+    }
+
+    /// Whether this `CancellationTokenState` is still referenced by any
+    /// `CancellationToken`.
+    fn has_refs(&self) -> bool {
+        self.refcount != 0 || self.has_parent_ref
+    }
+}
+
+/// The maximum permitted amount of references to a CancellationToken. This
+/// is derived from the intent to never use more than 32bit in the `Snapshot`.
+const MAX_REFS: u32 = (std::u32::MAX - 7) >> 3;
+
+/// Internal state of the `CancellationToken` pair above
+struct CancellationTokenState {
+    state: AtomicUsize,
+    parent: Option<NonNull<CancellationTokenState>>,
+    from_parent: SynchronizedThroughParent,
+    synchronized: Mutex<SynchronizedState>,
+}
+
+impl CancellationTokenState {
+    fn new(
+        parent: Option<NonNull<CancellationTokenState>>,
+        state: StateSnapshot,
+    ) -> CancellationTokenState {
+        CancellationTokenState {
+            parent,
+            from_parent: SynchronizedThroughParent {
+                prev_peer: None,
+                next_peer: None,
+            },
+            state: AtomicUsize::new(state.pack()),
+            synchronized: Mutex::new(SynchronizedState::new()),
+        }
+    }
+
+    /// Returns a snapshot of the current atomic state of the token
+    fn snapshot(&self) -> StateSnapshot {
+        StateSnapshot::unpack(self.state.load(Ordering::SeqCst))
+    }
+
+    fn atomic_update_state<F>(&self, mut current_state: StateSnapshot, func: F) -> StateSnapshot
+    where
+        F: Fn(StateSnapshot) -> StateSnapshot,
+    {
+        let mut current_packed_state = current_state.pack();
+        loop {
+            let next_state = func(current_state);
+            match self.state.compare_exchange(
+                current_packed_state,
+                next_state.pack(),
+                Ordering::SeqCst,
+                Ordering::SeqCst,
+            ) {
+                Ok(_) => {
+                    return next_state;
+                }
+                Err(actual) => {
+                    current_packed_state = actual;
+                    current_state = StateSnapshot::unpack(actual);
+                }
+            }
+        }
+    }
+
+    fn increment_refcount(&self, current_state: StateSnapshot) -> StateSnapshot {
+        self.atomic_update_state(current_state, |mut state: StateSnapshot| {
+            if state.refcount >= MAX_REFS as usize {
+                eprintln!("[ERROR] Maximum reference count for CancellationToken was exceeded");
+                std::process::abort();
+            }
+            state.refcount += 1;
+            state
+        })
+    }
+
+    fn decrement_refcount(&self, current_state: StateSnapshot) -> StateSnapshot {
+        let current_state = self.atomic_update_state(current_state, |mut state: StateSnapshot| {
+            state.refcount -= 1;
+            state
+        });
+
+        // Drop the State if it is not referenced anymore
+        if !current_state.has_refs() {
+            // Safety: `CancellationTokenState` is always stored in refcounted
+            // Boxes
+            let _ = unsafe { Box::from_raw(self as *const Self as *mut Self) };
+        }
+
+        current_state
+    }
+
+    fn remove_parent_ref(&self, current_state: StateSnapshot) -> StateSnapshot {
+        let current_state = self.atomic_update_state(current_state, |mut state: StateSnapshot| {
+            state.has_parent_ref = false;
+            state
+        });
+
+        // Drop the State if it is not referenced anymore
+        if !current_state.has_refs() {
+            // Safety: `CancellationTokenState` is always stored in refcounted
+            // Boxes
+            let _ = unsafe { Box::from_raw(self as *const Self as *mut Self) };
+        }
+
+        current_state
+    }
+
+    /// Unregisters a child from the parent token.
+    /// The child tokens state is not exactly known at this point in time.
+    /// If the parent token is cancelled, the child token gets removed from the
+    /// parents list, and might therefore already have been freed. If the parent
+    /// token is not cancelled, the child token is still valid.
+    fn unregister_child(
+        &mut self,
+        mut child_state: NonNull<CancellationTokenState>,
+        current_child_state: StateSnapshot,
+    ) {
+        let removed_child = {
+            // Remove the child toke from the parents linked list
+            let mut guard = self.synchronized.lock().unwrap();
+            if !guard.is_cancelled {
+                // Safety: Since the token was not cancelled, the child must
+                // still be in the list and valid.
+                let mut child_state = unsafe { child_state.as_mut() };
+                debug_assert!(child_state.snapshot().has_parent_ref);
+
+                if guard.first_child == Some(child_state.into()) {
+                    guard.first_child = child_state.from_parent.next_peer;
+                }
+                // Safety: If peers wouldn't be valid anymore, they would try
+                // to remove themselves from the list. This would require locking
+                // the Mutex that we currently own.
+                unsafe {
+                    if let Some(mut prev_peer) = child_state.from_parent.prev_peer {
+                        prev_peer.as_mut().from_parent.next_peer =
+                            child_state.from_parent.next_peer;
+                    }
+                    if let Some(mut next_peer) = child_state.from_parent.next_peer {
+                        next_peer.as_mut().from_parent.prev_peer =
+                            child_state.from_parent.prev_peer;
+                    }
+                }
+                child_state.from_parent.prev_peer = None;
+                child_state.from_parent.next_peer = None;
+
+                // The child is no longer referenced by the parent, since we were able
+                // to remove its reference from the parents list.
+                true
+            } else {
+                // Do not touch the linked list anymore. If the parent is cancelled
+                // it will move all childs outside of the Mutex and manipulate
+                // the pointers there. Manipulating the pointers here too could
+                // lead to races. Therefore leave them just as as and let the
+                // parent deal with it. The parent will make sure to retain a
+                // reference to this state as long as it manipulates the list
+                // pointers. Therefore the pointers are not dangling.
+                false
+            }
+        };
+
+        if removed_child {
+            // If the token removed itself from the parents list, it can reset
+            // the the parent ref status. If it is isn't able to do so, because the
+            // parent removed it from the list, there is no need to do this.
+            // The parent ref acts as as another reference count. Therefore
+            // removing this reference can free the object.
+            // Safety: The token was in the list. This means the parent wasn't
+            // cancelled before, and the token must still be alive.
+            unsafe { child_state.as_mut().remove_parent_ref(current_child_state) };
+        }
+
+        // Decrement the refcount on the parent and free it if necessary
+        self.decrement_refcount(self.snapshot());
+    }
+
+    fn cancel(&self) {
+        // Move the state of the CancellationToken from `NotCancelled` to `Cancelling`
+        let mut current_state = self.snapshot();
+
+        let state_after_cancellation = loop {
+            if current_state.cancel_state != CancellationState::NotCancelled {
+                // Another task already initiated the cancellation
+                return;
+            }
+
+            let mut next_state = current_state;
+            next_state.cancel_state = CancellationState::Cancelling;
+            match self.state.compare_exchange(
+                current_state.pack(),
+                next_state.pack(),
+                Ordering::SeqCst,
+                Ordering::SeqCst,
+            ) {
+                Ok(_) => break next_state,
+                Err(actual) => current_state = StateSnapshot::unpack(actual),
+            }
+        };
+
+        // This task cancelled the token
+
+        // Take the task list out of the Token
+        // We do not want to cancel child token inside this lock. If one of the
+        // child tasks would have additional child tokens, we would recursively
+        // take locks.
+
+        // Doing this action has an impact if the child token is dropped concurrently:
+        // It will try to deregister itself from the parent task, but can not find
+        // itself in the task list anymore. Therefore it needs to assume the parent
+        // has extracted the list and will process it. It may not modify the list.
+        // This is OK from a memory safety perspective, since the parent still
+        // retains a reference to the child task until it finished iterating over
+        // it.
+
+        let mut first_child = {
+            let mut guard = self.synchronized.lock().unwrap();
+            // Save the cancellation also inside the Mutex
+            // This allows child tokens which want to detach themselves to detect
+            // that this is no longer required since the parent cleared the list.
+            guard.is_cancelled = true;
+
+            // Wakeup all waiters
+            // This happens inside the lock to make cancellation reliable
+            // If we would access waiters outside of the lock, the pointers
+            // may no longer be valid.
+            // Typically this shouldn't be an issue, since waking a task should
+            // only move it from the blocked into the ready state and not have
+            // further side effects.
+
+            // Use a reverse iterator, so that the oldest waiter gets
+            // scheduled first
+            guard.waiters.reverse_drain(|waiter| {
+                // We are not allowed to move the `Waker` out of the list node.
+                // The `Future` relies on the fact that the old `Waker` stays there
+                // as long as the `Future` has not completed in order to perform
+                // the `will_wake()` check.
+                // Therefore `wake_by_ref` is used instead of `wake()`
+                if let Some(handle) = &mut waiter.task {
+                    handle.wake_by_ref();
+                }
+                // Mark the waiter to have been removed from the list.
+                waiter.state = PollState::Done;
+            });
+
+            guard.first_child.take()
+        };
+
+        while let Some(mut child) = first_child {
+            // Safety: We know this is a valid pointer since it is in our child pointer
+            // list. It can't have been freed in between, since we retain a a reference
+            // to each child.
+            let mut_child = unsafe { child.as_mut() };
+
+            // Get the next child and clean up list pointers
+            first_child = mut_child.from_parent.next_peer;
+            mut_child.from_parent.prev_peer = None;
+            mut_child.from_parent.next_peer = None;
+
+            // Cancel the child task
+            mut_child.cancel();
+
+            // Drop the parent reference. This `CancellationToken` is not interested
+            // in interacting with the child anymore.
+            // This is ONLY allowed once we promised not to touch the state anymore
+            // after this interaction.
+            mut_child.remove_parent_ref(mut_child.snapshot());
+        }
+
+        // The cancellation has completed
+        // At this point in time tasks which registered a wait node can be sure
+        // that this wait node already had been dequeued from the list without
+        // needing to inspect the list.
+        self.atomic_update_state(state_after_cancellation, |mut state| {
+            state.cancel_state = CancellationState::Cancelled;
+            state
+        });
+    }
+
+    /// Returns `true` if the `CancellationToken` had been cancelled
+    fn is_cancelled(&self) -> bool {
+        let current_state = self.snapshot();
+        current_state.cancel_state != CancellationState::NotCancelled
+    }
+
+    /// Registers a waiting task at the `CancellationToken`.
+    /// Safety: This method is only safe as long as the waiting waiting task
+    /// will properly unregister the wait node before it gets moved.
+    unsafe fn register(
+        &self,
+        wait_node: &mut ListNode<WaitQueueEntry>,
+        cx: &mut Context<'_>,
+    ) -> Poll<()> {
+        debug_assert_eq!(PollState::New, wait_node.state);
+        let current_state = self.snapshot();
+
+        // Perform an optimistic cancellation check before. This is not strictly
+        // necessary since we also check for cancellation in the Mutex, but
+        // reduces the necessary work to be performed for tasks which already
+        // had been cancelled.
+        if current_state.cancel_state != CancellationState::NotCancelled {
+            return Poll::Ready(());
+        }
+
+        // So far the token is not cancelled. However it could be cancelld before
+        // we get the chance to store the `Waker`. Therfore we need to check
+        // for cancellation again inside the mutex.
+        let mut guard = self.synchronized.lock().unwrap();
+        if guard.is_cancelled {
+            // Cancellation was signalled
+            wait_node.state = PollState::Done;
+            Poll::Ready(())
+        } else {
+            // Added the task to the wait queue
+            wait_node.task = Some(cx.waker().clone());
+            wait_node.state = PollState::Waiting;
+            guard.waiters.add_front(wait_node);
+            Poll::Pending
+        }
+    }
+
+    fn check_for_cancellation(
+        &self,
+        wait_node: &mut ListNode<WaitQueueEntry>,
+        cx: &mut Context<'_>,
+    ) -> Poll<()> {
+        debug_assert!(
+            wait_node.task.is_some(),
+            "Method can only be called after task had been registered"
+        );
+
+        let current_state = self.snapshot();
+
+        if current_state.cancel_state != CancellationState::NotCancelled {
+            // If the cancellation had been fully completed we know that our `Waker`
+            // is no longer registered at the `CancellationToken`.
+            // Otherwise the cancel call may or may not yet have iterated
+            // through the waiters list and removed the wait nodes.
+            // If it hasn't yet, we need to remove it. Otherwise an attempt to
+            // reuse the `wait_node´ might get freed due to the `WaitForCancellationFuture`
+            // getting dropped before the cancellation had interacted with it.
+            if current_state.cancel_state != CancellationState::Cancelled {
+                self.unregister(wait_node);
+            }
+            Poll::Ready(())
+        } else {
+            // Check if we need to swap the `Waker`. This will make the check more
+            // expensive, since the `Waker` is synchronized through the Mutex.
+            // If we don't need to perform a `Waker` update, an atomic check for
+            // cancellation is sufficient.
+            let need_waker_update = wait_node
+                .task
+                .as_ref()
+                .map(|waker| waker.will_wake(cx.waker()))
+                .unwrap_or(true);
+
+            if need_waker_update {
+                let guard = self.synchronized.lock().unwrap();
+                if guard.is_cancelled {
+                    // Cancellation was signalled. Since this cancellation signal
+                    // is set inside the Mutex, the old waiter must already have
+                    // been removed from the waiting list
+                    debug_assert_eq!(PollState::Done, wait_node.state);
+                    wait_node.task = None;
+                    Poll::Ready(())
+                } else {
+                    // The WaitForCancellationFuture is already in the queue.
+                    // The CancellationToken can't have been cancelled,
+                    // since this would change the is_cancelled flag inside the mutex.
+                    // Therefore we just have to update the Waker. A follow-up
+                    // cancellation will always use the new waker.
+                    wait_node.task = Some(cx.waker().clone());
+                    Poll::Pending
+                }
+            } else {
+                // Do nothing. If the token gets cancelled, this task will get
+                // woken again and can fetch the cancellation.
+                Poll::Pending
+            }
+        }
+    }
+
+    fn unregister(&self, wait_node: &mut ListNode<WaitQueueEntry>) {
+        debug_assert!(
+            wait_node.task.is_some(),
+            "waiter can not be active without task"
+        );
+
+        let mut guard = self.synchronized.lock().unwrap();
+        // WaitForCancellationFuture only needs to get removed if it has been added to
+        // the wait queue of the CancellationToken.
+        // This has happened in the PollState::Waiting case.
+        if let PollState::Waiting = wait_node.state {
+            // Safety: Due to the state, we know that the node must be part
+            // of the waiter list
+            if !unsafe { guard.waiters.remove(wait_node) } {
+                // Panic if the address isn't found. This can only happen if the contract was
+                // violated, e.g. the WaitQueueEntry got moved after the initial poll.
+                panic!("Future could not be removed from wait queue");
+            }
+            wait_node.state = PollState::Done;
+        }
+        wait_node.task = None;
+    }
+}
+
+#[cfg(all(test, not(loom)))]
+mod tests {
+    use super::*;
+    use crate::pin;
+    use futures_test::task::new_count_waker;
+
+    #[test]
+    fn cancel_token() {
+        let (waker, wake_counter) = new_count_waker();
+        let token = CancellationToken::new();
+        assert_eq!(false, token.is_cancelled());
+
+        let wait_fut = token.cancelled();
+        pin!(wait_fut);
+
+        assert_eq!(
+            Poll::Pending,
+            wait_fut.as_mut().poll(&mut Context::from_waker(&waker))
+        );
+        assert_eq!(wake_counter, 0);
+
+        let wait_fut_2 = token.cancelled();
+        pin!(wait_fut_2);
+
+        token.can