sync: move CancellationToken tests (#2477)

In preparation of work on `CancellationToken` internals, the tests are
moved into `tests/` and are updated to not depend on internals.

ci/azure-check-features.yml

@@ -30,3 +30,9 @@ jobs:
   #   tracking-issue: https://github.com/rust-lang/cargo/issues/5133
   - script: cargo hack check --all --each-feature -Z avoid-dev-deps
     displayName: cargo hack check --all --each-feature
+
+  # Try with unstable feature flags
+  - script: cargo hack check --all --each-feature -Z avoid-dev-deps
+    displayName: cargo hack check --all --each-feature
+    env:
+      RUSTFLAGS: '--cfg tokio_unstable'

ci/azure-loom.yml

@@ -21,7 +21,7 @@ jobs:
     parameters:
       rust_version: ${{ parameters.rust }}
 
-  - script: RUSTFLAGS="--cfg loom" cargo test --lib --release --features "full" -- --nocapture $(scope)
+  - script: RUSTFLAGS="--cfg loom --cfg tokio_unstable" cargo test --lib --release --features "full" -- --nocapture $(scope)
     env:
       LOOM_MAX_PREEMPTIONS: 2
       CI: 'True'

ci/azure-test-stable.yml

@@ -21,6 +21,23 @@ jobs:
 
   - template: azure-is-release.yml
 
+  # Run `tokio` with only `full`
+  - script: cargo test --features full
+    env:
+      RUST_BACKTRACE: 1
+      CI: 'True'
+    displayName: tokio - cargo test --features full
+    workingDirectory: $(Build.SourcesDirectory)/tokio
+
+  # Run `tokio` with "unstable" cfg flag
+  - script: cargo test --all-features
+    env:
+      RUSTFLAGS: '--cfg tokio_unstable'
+      RUST_BACKTRACE: 1
+      CI: 'True'
+    displayName: tokio - cargo test --features full
+    workingDirectory: $(Build.SourcesDirectory)/tokio
+
   - ${{ each crate in parameters.crates }}:
     # Run with all crate features
     - script: cargo test --all-features

tokio/src/macros/cfg.rs

@@ -353,3 +353,13 @@ macro_rules! cfg_uds {
         )*
     }
 }
+
+macro_rules! cfg_unstable {
+    ($($item:item)*) => {
+        $(
+            #[cfg(tokio_unstable)]
+            #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
+            $item
+        )*
+    }
+}

tokio/src/sync/cancellation_token.rs

@@ -1,17 +1,15 @@
 //! An asynchronously awaitable `CancellationToken`.
 //! The token allows to signal a cancellation request to one or more tasks.
 
-use crate::{
+use crate::loom::sync::atomic::AtomicUsize;
-    loom::sync::{atomic::AtomicUsize, Mutex},
+use crate::loom::sync::Mutex;
-    util::intrusive_double_linked_list::{LinkedList, ListNode},
+use crate::util::intrusive_double_linked_list::{LinkedList, ListNode};
-};
+
-use core::{
+use core::future::Future;
-    future::Future,
+use core::pin::Pin;
-    pin::Pin,
+use core::ptr::NonNull;
-    ptr::NonNull,
+use core::sync::atomic::Ordering;
-    sync::atomic::Ordering,
+use core::task::{Context, Poll, Waker};
-    task::{Context, Poll, Waker},
-};
 
 /// A token which can be used to signal a cancellation request to one or more
 /// tasks.
@@ -62,6 +60,25 @@ pub struct CancellationToken {
 unsafe impl Send for CancellationToken {}
 unsafe impl Sync for CancellationToken {}
 
+/// 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 CancellationToken =====
+
 impl core::fmt::Debug for CancellationToken {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_struct("CancellationToken")
@@ -114,8 +131,7 @@ impl Drop for CancellationToken {
 
 impl CancellationToken {
     /// Creates a new CancellationToken in the non-cancelled state.
-    #[allow(dead_code)]
+    pub fn new() -> CancellationToken {
-    pub(crate) fn new() -> CancellationToken {
         let state = Box::new(CancellationTokenState::new(
             None,
             StateSnapshot {
@@ -177,8 +193,7 @@ impl CancellationToken {
     ///     assert_eq!(5, join_handle.await.unwrap());
     /// }
     /// ```
-    #[allow(dead_code)]
+    pub fn child_token(&self) -> CancellationToken {
-    pub(crate) fn child_token(&self) -> CancellationToken {
         let inner = self.state();
 
         // Increment the refcount of this token. It will be referenced by the
@@ -232,31 +247,11 @@ impl CancellationToken {
         }
     }
 
-    /// 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 fn cancel(&self) {
-    pub(crate) fn cancel(&self) {
         self.state().cancel();
     }
 
@@ -295,22 +290,7 @@ impl CancellationToken {
     }
 }
 
-/// A Future that is resolved once the corresponding [`CancellationToken`]
+// ===== impl WaitForCancellationFuture =====
-/// 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 {
@@ -879,308 +859,3 @@ impl CancellationTokenState {
         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.cancel();
-        assert_eq!(wake_counter, 1);
-        assert_eq!(true, token.is_cancelled());
-
-        assert_eq!(
-            Poll::Ready(()),
-            wait_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(
-            Poll::Ready(()),
-            wait_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-    }
-
-    #[test]
-    fn cancel_child_token_through_parent() {
-        let (waker, wake_counter) = new_count_waker();
-        let token = CancellationToken::new();
-
-        let child_token = token.child_token();
-        assert_eq!(
-            StateSnapshot {
-                refcount: 1,
-                has_parent_ref: true,
-                cancel_state: CancellationState::NotCancelled,
-            },
-            child_token.state().snapshot()
-        );
-
-        let child_fut = child_token.cancelled();
-        pin!(child_fut);
-        let parent_fut = token.cancelled();
-        pin!(parent_fut);
-
-        assert_eq!(
-            Poll::Pending,
-            child_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(
-            Poll::Pending,
-            parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(wake_counter, 0);
-
-        token.cancel();
-        assert_eq!(wake_counter, 2);
-        assert_eq!(true, token.is_cancelled());
-        assert_eq!(true, child_token.is_cancelled());
-        assert_eq!(
-            StateSnapshot {
-                refcount: 1,
-                has_parent_ref: false,
-                cancel_state: CancellationState::Cancelled,
-            },
-            child_token.state().snapshot()
-        );
-
-        assert_eq!(
-            Poll::Ready(()),
-            child_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(
-            Poll::Ready(()),
-            parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-    }
-
-    #[test]
-    fn cancel_child_token_without_parent() {
-        let (waker, wake_counter) = new_count_waker();
-        let token = CancellationToken::new();
-
-        let child_token_1 = token.child_token();
-
-        let child_fut = child_token_1.cancelled();
-        pin!(child_fut);
-        let parent_fut = token.cancelled();
-        pin!(parent_fut);
-
-        assert_eq!(
-            Poll::Pending,
-            child_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(
-            Poll::Pending,
-            parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(wake_counter, 0);
-
-        child_token_1.cancel();
-        assert_eq!(wake_counter, 1);
-        assert_eq!(false, token.is_cancelled());
-        assert_eq!(true, child_token_1.is_cancelled());
-        assert_eq!(
-            StateSnapshot {
-                refcount: 1,
-                has_parent_ref: true,
-                cancel_state: CancellationState::Cancelled,
-            },
-            child_token_1.state().snapshot()
-        );
-
-        assert_eq!(
-            Poll::Ready(()),
-            child_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(
-            Poll::Pending,
-            parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-
-        let child_token_2 = token.child_token();
-        let child_fut_2 = child_token_2.cancelled();
-        pin!(child_fut_2);
-
-        assert_eq!(
-            Poll::Pending,
-            child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(
-            Poll::Pending,
-            parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-
-        token.cancel();
-        assert_eq!(wake_counter, 3);
-        assert_eq!(true, token.is_cancelled());
-        assert_eq!(true, child_token_2.is_cancelled());
-
-        assert_eq!(
-            Poll::Ready(()),
-            child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-        assert_eq!(
-            Poll::Ready(()),
-            parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
-        );
-    }
-
-    #[test]
-    fn create_child_token_after_parent_was_cancelled() {
-        for drop_child_first in [true, false].iter().cloned() {
-            let (waker, wake_counter) = new_count_waker();
-            let token = CancellationToken::new();
-            token.cancel();
-
-            let child_token = token.child_token();
-            assert_eq!(
-                StateSnapshot {
-                    refcount: 1,
-                    has_parent_ref: false,
-                    cancel_state: CancellationState::Cancelled,
-                },
-                child_token.state().snapshot()
-            );
-
-            {
-                let child_fut = child_token.cancelled();
-                pin!(child_fut);
-                let parent_fut = token.cancelled();
-                pin!(parent_fut);
-
-                assert_eq!(
-                    Poll::Ready(()),
-                    child_fut.as_mut().poll(&mut Context::from_waker(&waker))
-                );
-                assert_eq!(
-                    Poll::Ready(()),
-                    parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
-                );
-                assert_eq!(wake_counter, 0);
-
-                drop(child_fut);
-                drop(parent_fut);
-            }
-
-            if drop_child_first {
-                drop(child_token);
-                drop(token);
-            } else {
-                drop(token);
-                drop(child_token);
-            }
-        }
-    }
-
-    #[test]
-    fn drop_multiple_child_tokens() {
-        for drop_first_child_first in &[true, false] {
-            let token = CancellationToken::new();
-            let mut child_tokens = [None, None, None];
-            for i in 0..child_tokens.len() {
-                child_tokens[i] = Some(token.child_token());
-            }
-            assert_eq!(child_tokens.len(), token.child_tokens());
-
-            assert_eq!(
-                StateSnapshot {
-                    refcount: child_tokens.len() + 1,
-                    has_parent_ref: false,
-                    cancel_state: CancellationState::NotCancelled,
-                },
-                token.state().snapshot()
-            );
-            assert_eq!(
-                StateSnapshot {
-                    refcount: 1,
-                    has_parent_ref: true,
-                    cancel_state: CancellationState::NotCancelled,
-                },
-                child_tokens[0].as_ref().unwrap().state().snapshot()
-            );
-
-            let mut remaining_childs = child_tokens.len();
-            for i in 0..child_tokens.len() {
-                if *drop_first_child_first {
-                    child_tokens[i] = None;
-                } else {
-                    child_tokens[child_tokens.len() - 1 - i] = None;
-                }
-                remaining_childs -= 1;
-                assert_eq!(
-                    StateSnapshot {
-                        refcount: remaining_childs + 1,
-                        has_parent_ref: false,
-                        cancel_state: CancellationState::NotCancelled,
-                    },
-                    token.state().snapshot()
-                );
-                assert_eq!(remaining_childs, token.child_tokens());
-            }
-
-            drop(token);
-        }
-    }
-
-    #[test]
-    fn drop_parent_before_child_tokens() {
-        let token = CancellationToken::new();
-        let child1 = token.child_token();
-        let child2 = token.child_token();
-
-        drop(token);
-        assert_eq!(
-            StateSnapshot {
-                refcount: 1,
-                has_parent_ref: true,
-                cancel_state: CancellationState::NotCancelled,
-            },
-            child1.state().snapshot()
-        );
-
-        drop(child1);
-        drop(child2);
-    }
-
-    #[test]
-    fn parent_refcoutn_only_decreases_if_all_child_clones_are_released() {
-        let token = CancellationToken::new();
-        let child1 = token.child_token();
-        let child2 = child1.clone();
-
-        assert_eq!(2, token.state().snapshot().refcount);
-
-        drop(child1);
-        assert_eq!(2, token.state().snapshot().refcount);
-        drop(child2);
-        assert_eq!(1, token.state().snapshot().refcount);
-    }
-
-    #[test]
-    fn cancellation_future_is_send() {
-        let token = CancellationToken::new();
-        let fut = token.cancelled();
-
-        fn with_send<T: Send>(_: T) {}
-
-        with_send(fut);
-    }
-}

tokio/src/sync/mod.rs

@@ -425,10 +425,10 @@ cfg_sync! {
 
     pub mod broadcast;
 
-    #[cfg(tokio_unstable)]
+    cfg_unstable! {
-    mod cancellation_token;
+        mod cancellation_token;
-    #[cfg(tokio_unstable)]
+        pub use cancellation_token::{CancellationToken, WaitForCancellationFuture};
-    pub use cancellation_token::{CancellationToken, WaitForCancellationFuture};
+    }
 
     pub mod mpsc;
 

tokio/src/sync/tests/loom_cancellation_token.rs

@@ -1,4 +1,4 @@
-use crate::scope::CancellationToken;
+use crate::sync::CancellationToken;
 
 use loom::{future::block_on, thread};
 use tokio_test::assert_ok;

tokio/tests/async_send_sync.rs

@@ -259,3 +259,6 @@ async_assert_fn!(tokio::time::timeout_at(Instant, BoxFutureSync<()>): Send & Syn
 async_assert_fn!(tokio::time::timeout_at(Instant, BoxFutureSend<()>): Send & !Sync);
 async_assert_fn!(tokio::time::timeout_at(Instant, BoxFuture<()>): !Send & !Sync);
 async_assert_fn!(tokio::time::Interval::tick(_): Send & Sync);
+
+#[cfg(tokio_unstable)]
+assert_value!(tokio::sync::CancellationToken: Send & Sync);

tokio/tests/sync_cancellation_token.rs

@@ -0,0 +1,220 @@
+#![cfg(tokio_unstable)]
+
+use tokio::pin;
+use tokio::sync::CancellationToken;
+
+use core::future::Future;
+use core::task::{Context, Poll};
+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.cancel();
+    assert_eq!(wake_counter, 1);
+    assert_eq!(true, token.is_cancelled());
+
+    assert_eq!(
+        Poll::Ready(()),
+        wait_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(
+        Poll::Ready(()),
+        wait_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+}
+
+#[test]
+fn cancel_child_token_through_parent() {
+    let (waker, wake_counter) = new_count_waker();
+    let token = CancellationToken::new();
+
+    let child_token = token.child_token();
+    assert!(!child_token.is_cancelled());
+
+    let child_fut = child_token.cancelled();
+    pin!(child_fut);
+    let parent_fut = token.cancelled();
+    pin!(parent_fut);
+
+    assert_eq!(
+        Poll::Pending,
+        child_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(
+        Poll::Pending,
+        parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(wake_counter, 0);
+
+    token.cancel();
+    assert_eq!(wake_counter, 2);
+    assert_eq!(true, token.is_cancelled());
+    assert_eq!(true, child_token.is_cancelled());
+
+    assert_eq!(
+        Poll::Ready(()),
+        child_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(
+        Poll::Ready(()),
+        parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+}
+
+#[test]
+fn cancel_child_token_without_parent() {
+    let (waker, wake_counter) = new_count_waker();
+    let token = CancellationToken::new();
+
+    let child_token_1 = token.child_token();
+
+    let child_fut = child_token_1.cancelled();
+    pin!(child_fut);
+    let parent_fut = token.cancelled();
+    pin!(parent_fut);
+
+    assert_eq!(
+        Poll::Pending,
+        child_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(
+        Poll::Pending,
+        parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(wake_counter, 0);
+
+    child_token_1.cancel();
+    assert_eq!(wake_counter, 1);
+    assert_eq!(false, token.is_cancelled());
+    assert_eq!(true, child_token_1.is_cancelled());
+
+    assert_eq!(
+        Poll::Ready(()),
+        child_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(
+        Poll::Pending,
+        parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+
+    let child_token_2 = token.child_token();
+    let child_fut_2 = child_token_2.cancelled();
+    pin!(child_fut_2);
+
+    assert_eq!(
+        Poll::Pending,
+        child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(
+        Poll::Pending,
+        parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+
+    token.cancel();
+    assert_eq!(wake_counter, 3);
+    assert_eq!(true, token.is_cancelled());
+    assert_eq!(true, child_token_2.is_cancelled());
+
+    assert_eq!(
+        Poll::Ready(()),
+        child_fut_2.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+    assert_eq!(
+        Poll::Ready(()),
+        parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
+    );
+}
+
+#[test]
+fn create_child_token_after_parent_was_cancelled() {
+    for drop_child_first in [true, false].iter().cloned() {
+        let (waker, wake_counter) = new_count_waker();
+        let token = CancellationToken::new();
+        token.cancel();
+
+        let child_token = token.child_token();
+        assert!(child_token.is_cancelled());
+
+        {
+            let child_fut = child_token.cancelled();
+            pin!(child_fut);
+            let parent_fut = token.cancelled();
+            pin!(parent_fut);
+
+            assert_eq!(
+                Poll::Ready(()),
+                child_fut.as_mut().poll(&mut Context::from_waker(&waker))
+            );
+            assert_eq!(
+                Poll::Ready(()),
+                parent_fut.as_mut().poll(&mut Context::from_waker(&waker))
+            );
+            assert_eq!(wake_counter, 0);
+
+            drop(child_fut);
+            drop(parent_fut);
+        }
+
+        if drop_child_first {
+            drop(child_token);
+            drop(token);
+        } else {
+            drop(token);
+            drop(child_token);
+        }
+    }
+}
+
+#[test]
+fn drop_multiple_child_tokens() {
+    for drop_first_child_first in &[true, false] {
+        let token = CancellationToken::new();
+        let mut child_tokens = [None, None, None];
+        for i in 0..child_tokens.len() {
+            child_tokens[i] = Some(token.child_token());
+        }
+
+        assert!(!token.is_cancelled());
+        assert!(!child_tokens[0].as_ref().unwrap().is_cancelled());
+
+        for i in 0..child_tokens.len() {
+            if *drop_first_child_first {
+                child_tokens[i] = None;
+            } else {
+                child_tokens[child_tokens.len() - 1 - i] = None;
+            }
+            assert!(!token.is_cancelled());
+        }
+
+        drop(token);
+    }
+}
+
+#[test]
+fn drop_parent_before_child_tokens() {
+    let token = CancellationToken::new();
+    let child1 = token.child_token();
+    let child2 = token.child_token();
+
+    drop(token);
+    assert!(!child1.is_cancelled());
+
+    drop(child1);
+    drop(child2);
+}