rt: track loom changes + tweak queue (#2315)

Loom is having a big refresh to improve performance and tighten up the
concurrency model. This diff tracks those changes.

Included in the changes is the removal of `CausalCell` deferred checks.
This is due to it technically being undefined behavior in the C++11
memory model. To address this, the work-stealing queue is updated to
avoid needing this behavior. This is done by limiting the queue to have
one concurrent stealer.

6 files changed, 104 insertions, 46 deletions

diff --git a/tokio/src/loom/std/atomic_ptr.rs b/tokio/src/loom/std/atomic_ptr.rs
new file mode 100644
index 00000000..eb8e4755
--- /dev/null
+++ b/tokio/src/loom/std/atomic_ptr.rs
@@ -0,0 +1,32 @@
+use std::fmt;
+use std::ops::Deref;
+
+/// `AtomicPtr` providing an additional `load_unsync` function.
+pub(crate) struct AtomicPtr<T> {
+    inner: std::sync::atomic::AtomicPtr<T>,
+}
+
+impl<T> AtomicPtr<T> {
+    pub(crate) fn new(ptr: *mut T) -> AtomicPtr<T> {
+        let inner = std::sync::atomic::AtomicPtr::new(ptr);
+        AtomicPtr { inner }
+    }
+
+    pub(crate) fn with_mut<R>(&mut self, f: impl FnOnce(&mut *mut T) -> R) -> R {
+        f(self.inner.get_mut())
+    }
+}
+
+impl<T> Deref for AtomicPtr<T> {
+    type Target = std::sync::atomic::AtomicPtr<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<T> fmt::Debug for AtomicPtr<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.deref().fmt(fmt)
+    }
+}
diff --git a/tokio/src/loom/std/atomic_u8.rs b/tokio/src/loom/std/atomic_u8.rs
new file mode 100644
index 00000000..9f394a79
--- /dev/null
+++ b/tokio/src/loom/std/atomic_u8.rs
@@ -0,0 +1,44 @@
+use std::cell::UnsafeCell;
+use std::fmt;
+use std::ops::Deref;
+
+/// `AtomicU8` providing an additional `load_unsync` function.
+pub(crate) struct AtomicU8 {
+    inner: UnsafeCell<std::sync::atomic::AtomicU8>,
+}
+
+unsafe impl Send for AtomicU8 {}
+unsafe impl Sync for AtomicU8 {}
+
+impl AtomicU8 {
+    pub(crate) fn new(val: u8) -> AtomicU8 {
+        let inner = UnsafeCell::new(std::sync::atomic::AtomicU8::new(val));
+        AtomicU8 { inner }
+    }
+
+    /// Performs an unsynchronized load.
+    ///
+    /// # Safety
+    ///
+    /// All mutations must have happened before the unsynchronized load.
+    /// Additionally, there must be no concurrent mutations.
+    pub(crate) unsafe fn unsync_load(&self) -> u8 {
+        *(*self.inner.get()).get_mut()
+    }
+}
+
+impl Deref for AtomicU8 {
+    type Target = std::sync::atomic::AtomicU8;
+
+    fn deref(&self) -> &Self::Target {
+        // safety: it is always safe to access `&self` fns on the inner value as
+        // we never perform unsafe mutations.
+        unsafe { &*self.inner.get() }
+    }
+}
+
+impl fmt::Debug for AtomicU8 {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.deref().fmt(fmt)
+    }
+}
diff --git a/tokio/src/loom/std/atomic_usize.rs b/tokio/src/loom/std/atomic_usize.rs
index 78644b05..0fe998f1 100644
--- a/tokio/src/loom/std/atomic_usize.rs
+++ b/tokio/src/loom/std/atomic_usize.rs
@@ -25,6 +25,11 @@ impl AtomicUsize {
     pub(crate) unsafe fn unsync_load(&self) -> usize {
         *(*self.inner.get()).get_mut()
     }
+
+    pub(crate) fn with_mut<R>(&mut self, f: impl FnOnce(&mut usize) -> R) -> R {
+        // safety: we have mutable access
+        f(unsafe { (*self.inner.get()).get_mut() })
+    }
 }
 
 impl ops::Deref for AtomicUsize {
diff --git a/tokio/src/loom/std/causal_cell.rs b/tokio/src/loom/std/causal_cell.rs
deleted file mode 100644
index 8300437a..00000000
--- a/tokio/src/loom/std/causal_cell.rs
+++ /dev/null
@@ -1,40 +0,0 @@
-use std::cell::UnsafeCell;
-
-#[derive(Debug)]
-pub(crate) struct CausalCell<T>(UnsafeCell<T>);
-
-#[derive(Default)]
-pub(crate) struct CausalCheck(());
-
-impl<T> CausalCell<T> {
-    pub(crate) fn new(data: T) -> CausalCell<T> {
-        CausalCell(UnsafeCell::new(data))
-    }
-
-    pub(crate) fn with<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce(*const T) -> R,
-    {
-        f(self.0.get())
-    }
-
-    pub(crate) fn with_deferred<F, R>(&self, f: F) -> (R, CausalCheck)
-    where
-        F: FnOnce(*const T) -> R,
-    {
-        (f(self.0.get()), CausalCheck::default())
-    }
-
-    pub(crate) fn with_mut<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce(*mut T) -> R,
-    {
-        f(self.0.get())
-    }
-}
-
-impl CausalCheck {
-    pub(crate) fn check(self) {}
-
-    pub(crate) fn join(&mut self, _other: CausalCheck) {}
-}
diff --git a/tokio/src/loom/std/mod.rs b/tokio/src/loom/std/mod.rs
index a56d778a..6d7bcee1 100644
--- a/tokio/src/loom/std/mod.rs
+++ b/tokio/src/loom/std/mod.rs
@@ -1,12 +1,13 @@
 #![cfg_attr(any(not(feature = "full"), loom), allow(unused_imports, dead_code))]
 
-mod atomic_u32;
+mod atomic_ptr;
 mod atomic_u64;
+mod atomic_u8;
 mod atomic_usize;
-mod causal_cell;
+mod unsafe_cell;
 
 pub(crate) mod cell {
-    pub(crate) use super::causal_cell::{CausalCell, CausalCheck};
+    pub(crate) use super::unsafe_cell::UnsafeCell;
 }
 
 #[cfg(any(feature = "sync", feature = "io-driver"))]
@@ -58,12 +59,12 @@ pub(crate) mod sync {
     pub(crate) use std::sync::{Condvar, Mutex, MutexGuard, WaitTimeoutResult};
 
     pub(crate) mod atomic {
-        pub(crate) use crate::loom::std::atomic_u32::AtomicU32;
+        pub(crate) use crate::loom::std::atomic_ptr::AtomicPtr;
         pub(crate) use crate::loom::std::atomic_u64::AtomicU64;
+        pub(crate) use crate::loom::std::atomic_u8::AtomicU8;
         pub(crate) use crate::loom::std::atomic_usize::AtomicUsize;
 
-        pub(crate) use std::sync::atomic::AtomicU8;
-        pub(crate) use std::sync::atomic::{fence, AtomicPtr};
+        pub(crate) use std::sync::atomic::AtomicU16;
         pub(crate) use std::sync::atomic::{spin_loop_hint, AtomicBool};
     }
 }
diff --git a/tokio/src/loom/std/unsafe_cell.rs b/tokio/src/loom/std/unsafe_cell.rs
new file mode 100644
index 00000000..f2b03d8d
--- /dev/null
+++ b/tokio/src/loom/std/unsafe_cell.rs
@@ -0,0 +1,16 @@
+#[derive(Debug)]
+pub(crate) struct UnsafeCell<T>(std::cell::UnsafeCell<T>);
+
+impl<T> UnsafeCell<T> {
+    pub(crate) fn new(data: T) -> UnsafeCell<T> {
+        UnsafeCell(std::cell::UnsafeCell::new(data))
+    }
+
+    pub(crate) fn with<R>(&self, f: impl FnOnce(*const T) -> R) -> R {
+        f(self.0.get())
+    }
+
+    pub(crate) fn with_mut<R>(&self, f: impl FnOnce(*mut T) -> R) -> R {
+        f(self.0.get())
+    }
+}