chore: use the loom mutex wrapper everywhere (#3958)

tokio/src/loom/std/atomic_u64.rs

@@ -15,8 +15,8 @@ mod imp {
 
 #[cfg(any(target_arch = "arm", target_arch = "mips", target_arch = "powerpc"))]
 mod imp {
+    use crate::loom::sync::Mutex;
     use std::sync::atomic::Ordering;
-    use std::sync::Mutex;
 
     #[derive(Debug)]
     pub(crate) struct AtomicU64 {
@@ -31,15 +31,15 @@ mod imp {
         }
 
         pub(crate) fn load(&self, _: Ordering) -> u64 {
-            *self.inner.lock().unwrap()
+            *self.inner.lock()
         }
 
         pub(crate) fn store(&self, val: u64, _: Ordering) {
-            *self.inner.lock().unwrap() = val;
+            *self.inner.lock() = val;
         }
 
         pub(crate) fn fetch_or(&self, val: u64, _: Ordering) -> u64 {
-            let mut lock = self.inner.lock().unwrap();
+            let mut lock = self.inner.lock();
             let prev = *lock;
             *lock = prev | val;
             prev
@@ -52,7 +52,7 @@ mod imp {
             _success: Ordering,
             _failure: Ordering,
         ) -> Result<u64, u64> {
-            let mut lock = self.inner.lock().unwrap();
+            let mut lock = self.inner.lock();
 
             if *lock == current {
                 *lock = new;

tokio/src/runtime/shell.rs

@@ -1,132 +0,0 @@
-#![allow(clippy::redundant_clone)]
-
-use crate::future::poll_fn;
-use crate::park::{Park, Unpark};
-use crate::runtime::driver::Driver;
-use crate::sync::Notify;
-use crate::util::{waker_ref, Wake};
-
-use std::sync::{Arc, Mutex};
-use std::task::Context;
-use std::task::Poll::{Pending, Ready};
-use std::{future::Future, sync::PoisonError};
-
-#[derive(Debug)]
-pub(super) struct Shell {
-    driver: Mutex<Option<Driver>>,
-
-    notify: Notify,
-
-    /// TODO: don't store this
-    unpark: Arc<Handle>,
-}
-
-#[derive(Debug)]
-struct Handle(<Driver as Park>::Unpark);
-
-impl Shell {
-    pub(super) fn new(driver: Driver) -> Shell {
-        let unpark = Arc::new(Handle(driver.unpark()));
-
-        Shell {
-            driver: Mutex::new(Some(driver)),
-            notify: Notify::new(),
-            unpark,
-        }
-    }
-
-    pub(super) fn block_on<F>(&self, f: F) -> F::Output
-    where
-        F: Future,
-    {
-        let mut enter = crate::runtime::enter(true);
-
-        pin!(f);
-
-        loop {
-            if let Some(driver) = &mut self.take_driver() {
-                return driver.block_on(f);
-            } else {
-                let notified = self.notify.notified();
-                pin!(notified);
-
-                if let Some(out) = enter
-                    .block_on(poll_fn(|cx| {
-                        if notified.as_mut().poll(cx).is_ready() {
-                            return Ready(None);
-                        }
-
-                        if let Ready(out) = f.as_mut().poll(cx) {
-                            return Ready(Some(out));
-                        }
-
-                        Pending
-                    }))
-                    .expect("Failed to `Enter::block_on`")
-                {
-                    return out;
-                }
-            }
-        }
-    }
-
-    fn take_driver(&self) -> Option<DriverGuard<'_>> {
-        let mut lock = self.driver.lock().unwrap();
-        let driver = lock.take()?;
-
-        Some(DriverGuard {
-            inner: Some(driver),
-            shell: &self,
-        })
-    }
-}
-
-impl Wake for Handle {
-    /// Wake by value
-    fn wake(self: Arc<Self>) {
-        Wake::wake_by_ref(&self);
-    }
-
-    /// Wake by reference
-    fn wake_by_ref(arc_self: &Arc<Self>) {
-        arc_self.0.unpark();
-    }
-}
-
-struct DriverGuard<'a> {
-    inner: Option<Driver>,
-    shell: &'a Shell,
-}
-
-impl DriverGuard<'_> {
-    fn block_on<F: Future>(&mut self, f: F) -> F::Output {
-        let driver = self.inner.as_mut().unwrap();
-
-        pin!(f);
-
-        let waker = waker_ref(&self.shell.unpark);
-        let mut cx = Context::from_waker(&waker);
-
-        loop {
-            if let Ready(v) = crate::coop::budget(|| f.as_mut().poll(&mut cx)) {
-                return v;
-            }
-
-            driver.park().unwrap();
-        }
-    }
-}
-
-impl Drop for DriverGuard<'_> {
-    fn drop(&mut self) {
-        if let Some(inner) = self.inner.take() {
-            self.shell
-                .driver
-                .lock()
-                .unwrap_or_else(PoisonError::into_inner)
-                .replace(inner);
-
-            self.shell.notify.notify_one();
-        }
-    }
-}

tokio/src/runtime/tests/loom_oneshot.rs

@@ -1,7 +1,6 @@
+use crate::loom::sync::{Arc, Mutex};
 use loom::sync::Notify;
 
-use std::sync::{Arc, Mutex};
-
 pub(crate) fn channel<T>() -> (Sender<T>, Receiver<T>) {
     let inner = Arc::new(Inner {
         notify: Notify::new(),
@@ -31,7 +30,7 @@ struct Inner<T> {
 
 impl<T> Sender<T> {
     pub(crate) fn send(self, value: T) {
-        *self.inner.value.lock().unwrap() = Some(value);
+        *self.inner.value.lock() = Some(value);
         self.inner.notify.notify();
     }
 }
@@ -39,7 +38,7 @@ impl<T> Sender<T> {
 impl<T> Receiver<T> {
     pub(crate) fn recv(self) -> T {
         loop {
-            if let Some(v) = self.inner.value.lock().unwrap().take() {
+            if let Some(v) = self.inner.value.lock().take() {
                 return v;
             }
 

tokio/src/sync/barrier.rs

@@ -1,7 +1,6 @@
+use crate::loom::sync::Mutex;
 use crate::sync::watch;
 
-use std::sync::Mutex;
-
 /// A barrier enables multiple tasks to synchronize the beginning of some computation.
 ///
 /// ```
@@ -94,7 +93,7 @@ impl Barrier {
         // NOTE: the extra scope here is so that the compiler doesn't think `state` is held across
         // a yield point, and thus marks the returned future as !Send.
         let generation = {
-            let mut state = self.state.lock().unwrap();
+            let mut state = self.state.lock();
             let generation = state.generation;
             state.arrived += 1;
             if state.arrived == self.n {

tokio/src/task/local.rs

@@ -1,4 +1,5 @@
 //! Runs `!Send` futures on the current thread.
+use crate::loom::sync::{Arc, Mutex};
 use crate::runtime::task::{self, JoinHandle, Task};
 use crate::sync::AtomicWaker;
 use crate::util::linked_list::{Link, LinkedList};
@@ -9,7 +10,6 @@ use std::fmt;
 use std::future::Future;
 use std::marker::PhantomData;
 use std::pin::Pin;
-use std::sync::{Arc, Mutex};
 use std::task::Poll;
 
 use pin_project_lite::pin_project;
@@ -538,7 +538,6 @@ impl LocalSet {
                 .shared
                 .queue
                 .lock()
-                .unwrap()
                 .pop_front()
                 .or_else(|| self.context.tasks.borrow_mut().queue.pop_front())
         } else {
@@ -547,7 +546,7 @@ impl LocalSet {
                 .borrow_mut()
                 .queue
                 .pop_front()
-                .or_else(|| self.context.shared.queue.lock().unwrap().pop_front())
+                .or_else(|| self.context.shared.queue.lock().pop_front())
         }
     }
 
@@ -611,7 +610,7 @@ impl Drop for LocalSet {
                 task.shutdown();
             }
 
-            for task in self.context.shared.queue.lock().unwrap().drain(..) {
+            for task in self.context.shared.queue.lock().drain(..) {
                 task.shutdown();
             }
 
@@ -661,7 +660,7 @@ impl Shared {
                 cx.tasks.borrow_mut().queue.push_back(task);
             }
             _ => {
-                self.queue.lock().unwrap().push_back(task);
+                self.queue.lock().push_back(task);
                 self.waker.wake();
             }
         });

tokio/src/time/clock.rs

@@ -29,7 +29,7 @@ cfg_not_test_util! {
 
 cfg_test_util! {
     use crate::time::{Duration, Instant};
-    use std::sync::{Arc, Mutex};
+    use crate::loom::sync::{Arc, Mutex};
 
     cfg_rt! {
         fn clock() -> Option<Clock> {
@@ -102,7 +102,7 @@ cfg_test_util! {
     /// runtime.
     pub fn resume() {
         let clock = clock().expect("time cannot be frozen from outside the Tokio runtime");
-        let mut inner = clock.inner.lock().unwrap();
+        let mut inner = clock.inner.lock();
 
         if inner.unfrozen.is_some() {
             panic!("time is not frozen");
@@ -164,7 +164,7 @@ cfg_test_util! {
         }
 
         pub(crate) fn pause(&self) {
-            let mut inner = self.inner.lock().unwrap();
+            let mut inner = self.inner.lock();
 
             if !inner.enable_pausing {
                 drop(inner); // avoid poisoning the lock
@@ -178,12 +178,12 @@ cfg_test_util! {
         }
 
         pub(crate) fn is_paused(&self) -> bool {
-            let inner = self.inner.lock().unwrap();
+            let inner = self.inner.lock();
             inner.unfrozen.is_none()
         }
 
         pub(crate) fn advance(&self, duration: Duration) {
-            let mut inner = self.inner.lock().unwrap();
+            let mut inner = self.inner.lock();
 
             if inner.unfrozen.is_some() {
                 panic!("time is not frozen");
@@ -193,7 +193,7 @@ cfg_test_util! {
         }
 
         pub(crate) fn now(&self) -> Instant {
-            let inner = self.inner.lock().unwrap();
+            let inner = self.inner.lock();
 
             let mut ret = inner.base;