2 files changed, 157 insertions, 0 deletions
diff --git a/tokio/src/sync/tests/loom_cancellation_token.rs b/tokio/src/sync/tests/loom_cancellation_token.rs
new file mode 100644
index 00000000..aac84f80
--- /dev/null
+++ b/tokio/src/sync/tests/loom_cancellation_token.rs
@@ -0,0 +1,155 @@
+use crate::scope::CancellationToken;
+
+use loom::{future::block_on, thread};
+use tokio_test::assert_ok;
+
+#[test]
+fn cancel_token() {
+    loom::model(|| {
+        let token = CancellationToken::new();
+        let token1 = token.clone();
+
+        let th1 = thread::spawn(move || {
+            block_on(async {
+                token1.cancelled().await;
+            });
+        });
+
+        let th2 = thread::spawn(move || {
+            token.cancel();
+        });
+
+        assert_ok!(th1.join());
+        assert_ok!(th2.join());
+    });
+}
+
+#[test]
+fn cancel_with_child() {
+    loom::model(|| {
+        let token = CancellationToken::new();
+        let token1 = token.clone();
+        let token2 = token.clone();
+        let child_token = token.child_token();
+
+        let th1 = thread::spawn(move || {
+            block_on(async {
+                token1.cancelled().await;
+            });
+        });
+
+        let th2 = thread::spawn(move || {
+            token2.cancel();
+        });
+
+        let th3 = thread::spawn(move || {
+            block_on(async {
+                child_token.cancelled().await;
+            });
+        });
+
+        assert_ok!(th1.join());
+        assert_ok!(th2.join());
+        assert_ok!(th3.join());
+    });
+}
+
+#[test]
+fn drop_token_no_child() {
+    loom::model(|| {
+        let token = CancellationToken::new();
+        let token1 = token.clone();
+        let token2 = token.clone();
+
+        let th1 = thread::spawn(move || {
+            drop(token1);
+        });
+
+        let th2 = thread::spawn(move || {
+            drop(token2);
+        });
+
+        let th3 = thread::spawn(move || {
+            drop(token);
+        });
+
+        assert_ok!(th1.join());
+        assert_ok!(th2.join());
+        assert_ok!(th3.join());
+    });
+}
+
+#[test]
+fn drop_token_with_childs() {
+    loom::model(|| {
+        let token1 = CancellationToken::new();
+        let child_token1 = token1.child_token();
+        let child_token2 = token1.child_token();
+
+        let th1 = thread::spawn(move || {
+            drop(token1);
+        });
+
+        let th2 = thread::spawn(move || {
+            drop(child_token1);
+        });
+
+        let th3 = thread::spawn(move || {
+            drop(child_token2);
+        });
+
+        assert_ok!(th1.join());
+        assert_ok!(th2.join());
+        assert_ok!(th3.join());
+    });
+}
+
+#[test]
+fn drop_and_cancel_token() {
+    loom::model(|| {
+        let token1 = CancellationToken::new();
+        let token2 = token1.clone();
+        let child_token = token1.child_token();
+
+        let th1 = thread::spawn(move || {
+            drop(token1);
+        });
+
+        let th2 = thread::spawn(move || {
+            token2.cancel();
+        });
+
+        let th3 = thread::spawn(move || {
+            drop(child_token);
+        });
+
+        assert_ok!(th1.join());
+        assert_ok!(th2.join());
+        assert_ok!(th3.join());
+    });
+}
+
+#[test]
+fn cancel_parent_and_child() {
+    loom::model(|| {
+        let token1 = CancellationToken::new();
+        let token2 = token1.clone();
+        let child_token = token1.child_token();
+
+        let th1 = thread::spawn(move || {
+            drop(token1);
+        });
+
+        let th2 = thread::spawn(move || {
+            token2.cancel();
+        });
+
+        let th3 = thread::spawn(move || {
+            child_token.cancel();
+        });
+
+        assert_ok!(th1.join());
+        assert_ok!(th2.join());
+        assert_ok!(th3.join());
+    });
+}
diff --git a/tokio/src/sync/tests/mod.rs b/tokio/src/sync/tests/mod.rs
index d571754c..6ba8c1f9 100644
--- a/tokio/src/sync/tests/mod.rs
+++ b/tokio/src/sync/tests/mod.rs
@@ -7,6 +7,8 @@ cfg_not_loom! {
 cfg_loom! {
     mod loom_atomic_waker;
     mod loom_broadcast;
+    #[cfg(tokio_unstable)]
+    mod loom_cancellation_token;
     mod loom_list;
     mod loom_mpsc;
     mod loom_notify;