rt: fix bug in work-stealing queue (#2387)

Fixes a couple bugs in the work-stealing queue introduced as
part of #2315. First, the cursor needs to be able to represent more
values than the size of the buffer. This is to be able to track if
`tail` is ahead of `head` or if they are identical. This bug resulted in
the "overflow" path being taken before the buffer was full.

The second bug can happen when a queue is being stolen from concurrently
with stealing into. In this case, it is possible for buffer slots to be
overwritten before they are released by the stealer. This is harder to
happen in practice due to the first bug preventing the queue from
filling up 100%, but could still happen. It triggered an assertion in
`steal_into`. This bug slipped through due to a bug in loom not
correctly catching the case. The loom bug is fixed as part of
tokio-rs/loom#119.

Fixes: #2382

4 files changed, 48 insertions, 21 deletions

diff --git a/tokio/src/loom/std/atomic_u16.rs b/tokio/src/loom/std/atomic_u16.rs
new file mode 100644
index 00000000..70390972
--- /dev/null
+++ b/tokio/src/loom/std/atomic_u16.rs
@@ -0,0 +1,44 @@
+use std::cell::UnsafeCell;
+use std::fmt;
+use std::ops::Deref;
+
+/// `AtomicU16` providing an additional `load_unsync` function.
+pub(crate) struct AtomicU16 {
+    inner: UnsafeCell<std::sync::atomic::AtomicU16>,
+}
+
+unsafe impl Send for AtomicU16 {}
+unsafe impl Sync for AtomicU16 {}
+
+impl AtomicU16 {
+    pub(crate) fn new(val: u16) -> AtomicU16 {
+        let inner = UnsafeCell::new(std::sync::atomic::AtomicU16::new(val));
+        AtomicU16 { 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) -> u16 {
+        *(*self.inner.get()).get_mut()
+    }
+}
+
+impl Deref for AtomicU16 {
+    type Target = std::sync::atomic::AtomicU16;
+
+    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 AtomicU16 {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.deref().fmt(fmt)
+    }
+}
diff --git a/tokio/src/loom/std/atomic_u32.rs b/tokio/src/loom/std/atomic_u32.rs
index c83cfa2b..6f786c51 100644
--- a/tokio/src/loom/std/atomic_u32.rs
+++ b/tokio/src/loom/std/atomic_u32.rs
@@ -15,16 +15,6 @@ impl AtomicU32 {
         let inner = UnsafeCell::new(std::sync::atomic::AtomicU32::new(val));
         AtomicU32 { 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) -> u32 {
-        *(*self.inner.get()).get_mut()
-    }
 }
 
 impl Deref for AtomicU32 {
diff --git a/tokio/src/loom/std/atomic_u8.rs b/tokio/src/loom/std/atomic_u8.rs
index 9f394a79..4fcd0df3 100644
--- a/tokio/src/loom/std/atomic_u8.rs
+++ b/tokio/src/loom/std/atomic_u8.rs
@@ -15,16 +15,6 @@ impl 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 {
diff --git a/tokio/src/loom/std/mod.rs b/tokio/src/loom/std/mod.rs
index 6d7bcee1..28644544 100644
--- a/tokio/src/loom/std/mod.rs
+++ b/tokio/src/loom/std/mod.rs
@@ -1,6 +1,8 @@
 #![cfg_attr(any(not(feature = "full"), loom), allow(unused_imports, dead_code))]
 
 mod atomic_ptr;
+mod atomic_u16;
+mod atomic_u32;
 mod atomic_u64;
 mod atomic_u8;
 mod atomic_usize;
@@ -60,11 +62,12 @@ pub(crate) mod sync {
 
     pub(crate) mod atomic {
         pub(crate) use crate::loom::std::atomic_ptr::AtomicPtr;
+        pub(crate) use crate::loom::std::atomic_u16::AtomicU16;
+        pub(crate) use crate::loom::std::atomic_u32::AtomicU32;
         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::AtomicU16;
         pub(crate) use std::sync::atomic::{spin_loop_hint, AtomicBool};
     }
 }