tests: simplify loom tests (#3995)

tokio/src/runtime/tests/loom_pool.rs

@@ -11,7 +11,7 @@ use crate::{spawn, task};
 use tokio_test::assert_ok;
 
 use loom::sync::atomic::{AtomicBool, AtomicUsize};
-use loom::sync::{Arc, Mutex};
+use loom::sync::Arc;
 
 use pin_project_lite::pin_project;
 use std::future::Future;
@@ -19,6 +19,57 @@ use std::pin::Pin;
 use std::sync::atomic::Ordering::{Relaxed, SeqCst};
 use std::task::{Context, Poll};
 
+mod atomic_take {
+    use loom::sync::atomic::AtomicBool;
+    use std::mem::MaybeUninit;
+    use std::sync::atomic::Ordering::SeqCst;
+
+    pub(super) struct AtomicTake<T> {
+        inner: MaybeUninit<T>,
+        taken: AtomicBool,
+    }
+
+    impl<T> AtomicTake<T> {
+        pub(super) fn new(value: T) -> Self {
+            Self {
+                inner: MaybeUninit::new(value),
+                taken: AtomicBool::new(false),
+            }
+        }
+
+        pub(super) fn take(&self) -> Option<T> {
+            // safety: Only one thread will see the boolean change from false
+            // to true, so that thread is able to take the value.
+            match self.taken.fetch_or(true, SeqCst) {
+                false => unsafe { Some(std::ptr::read(self.inner.as_ptr())) },
+                true => None,
+            }
+        }
+    }
+
+    impl<T> Drop for AtomicTake<T> {
+        fn drop(&mut self) {
+            drop(self.take());
+        }
+    }
+}
+
+#[derive(Clone)]
+struct AtomicOneshot<T> {
+    value: std::sync::Arc<atomic_take::AtomicTake<oneshot::Sender<T>>>,
+}
+impl<T> AtomicOneshot<T> {
+    fn new(sender: oneshot::Sender<T>) -> Self {
+        Self {
+            value: std::sync::Arc::new(atomic_take::AtomicTake::new(sender)),
+        }
+    }
+
+    fn assert_send(&self, value: T) {
+        self.value.take().unwrap().send(value);
+    }
+}
+
 /// Tests are divided into groups to make the runs faster on CI.
 mod group_a {
     use super::*;
@@ -52,7 +103,7 @@ mod group_a {
             let c1 = Arc::new(AtomicUsize::new(0));
 
             let (tx, rx) = oneshot::channel();
-            let tx1 = Arc::new(Mutex::new(Some(tx)));
+            let tx1 = AtomicOneshot::new(tx);
 
             // Spawn a task
             let c2 = c1.clone();
@@ -60,7 +111,7 @@ mod group_a {
             pool.spawn(track(async move {
                 spawn(track(async move {
                     if 1 == c1.fetch_add(1, Relaxed) {
-                        tx1.lock().unwrap().take().unwrap().send(());
+                        tx1.assert_send(());
                     }
                 }));
             }));
@@ -69,7 +120,7 @@ mod group_a {
             pool.spawn(track(async move {
                 spawn(track(async move {
                     if 1 == c2.fetch_add(1, Relaxed) {
-                        tx2.lock().unwrap().take().unwrap().send(());
+                        tx2.assert_send(());
                     }
                 }));
             }));
@@ -119,7 +170,7 @@ mod group_b {
 
             let (block_tx, block_rx) = oneshot::channel();
             let (done_tx, done_rx) = oneshot::channel();
-            let done_tx = Arc::new(Mutex::new(Some(done_tx)));
+            let done_tx = AtomicOneshot::new(done_tx);
 
             pool.spawn(track(async move {
                 crate::task::block_in_place(move || {
@@ -136,7 +187,7 @@ mod group_b {
 
                 pool.spawn(track(async move {
                     if NUM == cnt.fetch_add(1, Relaxed) + 1 {
-                        done_tx.lock().unwrap().take().unwrap().send(());
+                        done_tx.assert_send(());
                     }
                 }));
             }
@@ -266,17 +317,17 @@ mod group_d {
             let c1 = Arc::new(AtomicUsize::new(0));
 
             let (done_tx, done_rx) = oneshot::channel();
-            let done_tx1 = Arc::new(Mutex::new(Some(done_tx)));
+            let done_tx1 = AtomicOneshot::new(done_tx);
+            let done_tx2 = done_tx1.clone();
 
             // Spawn a task
             let c2 = c1.clone();
-            let done_tx2 = done_tx1.clone();
             pool.spawn(track(async move {
                 gated().await;
                 gated().await;
 
                 if 1 == c1.fetch_add(1, Relaxed) {
-                    done_tx1.lock().unwrap().take().unwrap().send(());
+                    done_tx1.assert_send(());
                 }
             }));
 
@@ -286,7 +337,7 @@ mod group_d {
                 gated().await;
 
                 if 1 == c2.fetch_add(1, Relaxed) {
-                    done_tx2.lock().unwrap().take().unwrap().send(());
+                    done_tx2.assert_send(());
                 }
             }));