rt(alt): fix a number of concurrency bugs (#5907)

Expands loom coverage and fixes a number of bugs.

Closes #5888

tokio/Cargo.toml

@@ -157,7 +157,7 @@ wasm-bindgen-test = "0.3.0"
 mio-aio = { version = "0.7.0", features = ["tokio"] }
 
 [target.'cfg(loom)'.dev-dependencies]
-loom = { version = "0.6", features = ["futures", "checkpoint"] }
+loom = { version = "0.7", features = ["futures", "checkpoint"] }
 
 [package.metadata.docs.rs]
 all-features = true

tokio/src/loom/mocked.rs

@@ -15,6 +15,7 @@ pub(crate) mod sync {
         }
 
         #[inline]
+        #[track_caller]
         pub(crate) fn lock(&self) -> MutexGuard<'_, T> {
             self.0.lock().unwrap()
         }

tokio/src/runtime/builder.rs

@@ -93,6 +93,8 @@ pub struct Builder {
     /// How many ticks before yielding to the driver for timer and I/O events?
     pub(super) event_interval: u32,
 
+    pub(super) local_queue_capacity: usize,
+
     /// When true, the multi-threade scheduler LIFO slot should not be used.
     ///
     /// This option should only be exposed as unstable.
@@ -297,6 +299,12 @@ impl Builder {
             global_queue_interval: None,
             event_interval,
 
+            #[cfg(not(loom))]
+            local_queue_capacity: 256,
+
+            #[cfg(loom)]
+            local_queue_capacity: 4,
+
             seed_generator: RngSeedGenerator::new(RngSeed::new()),
 
             #[cfg(tokio_unstable)]
@@ -1046,6 +1054,14 @@ impl Builder {
         }
     }
 
+    cfg_loom! {
+        pub(crate) fn local_queue_capacity(&mut self, value: usize) -> &mut Self {
+            assert!(value.is_power_of_two());
+            self.local_queue_capacity = value;
+            self
+        }
+    }
+
     fn build_current_thread_runtime(&mut self) -> io::Result<Runtime> {
         use crate::runtime::scheduler::{self, CurrentThread};
         use crate::runtime::{runtime::Scheduler, Config};
@@ -1074,6 +1090,7 @@ impl Builder {
                 after_unpark: self.after_unpark.clone(),
                 global_queue_interval: self.global_queue_interval,
                 event_interval: self.event_interval,
+                local_queue_capacity: self.local_queue_capacity,
                 #[cfg(tokio_unstable)]
                 unhandled_panic: self.unhandled_panic.clone(),
                 disable_lifo_slot: self.disable_lifo_slot,
@@ -1224,6 +1241,7 @@ cfg_rt_multi_thread! {
                     after_unpark: self.after_unpark.clone(),
                     global_queue_interval: self.global_queue_interval,
                     event_interval: self.event_interval,
+                    local_queue_capacity: self.local_queue_capacity,
                     #[cfg(tokio_unstable)]
                     unhandled_panic: self.unhandled_panic.clone(),
                     disable_lifo_slot: self.disable_lifo_slot,
@@ -1271,6 +1289,7 @@ cfg_rt_multi_thread! {
                         after_unpark: self.after_unpark.clone(),
                         global_queue_interval: self.global_queue_interval,
                         event_interval: self.event_interval,
+                        local_queue_capacity: self.local_queue_capacity,
                         #[cfg(tokio_unstable)]
                         unhandled_panic: self.unhandled_panic.clone(),
                         disable_lifo_slot: self.disable_lifo_slot,

tokio/src/runtime/config.rs

@@ -1,4 +1,7 @@
-#![cfg_attr(any(not(feature = "full"), target_family = "wasm"), allow(dead_code))]
+#![cfg_attr(
+    any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"),
+    allow(dead_code)
+)]
 use crate::runtime::Callback;
 use crate::util::RngSeedGenerator;
 
@@ -9,6 +12,9 @@ pub(crate) struct Config {
     /// How many ticks before yielding to the driver for timer and I/O events?
     pub(crate) event_interval: u32,
 
+    /// How big to make each worker's local queue
+    pub(crate) local_queue_capacity: usize,
+
     /// Callback for a worker parking itself
     pub(crate) before_park: Option<Callback>,
 

tokio/src/runtime/driver.rs

@@ -2,7 +2,10 @@
 
 // Eventually, this file will see significant refactoring / cleanup. For now, we
 // don't need to worry much about dead code with certain feature permutations.
-#![cfg_attr(not(feature = "full"), allow(dead_code))]
+#![cfg_attr(
+    any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"),
+    allow(dead_code)
+)]
 
 use crate::runtime::park::{ParkThread, UnparkThread};
 
@@ -58,6 +61,10 @@ impl Driver {
         ))
     }
 
+    pub(crate) fn is_enabled(&self) -> bool {
+        self.inner.is_enabled()
+    }
+
     pub(crate) fn park(&mut self, handle: &Handle) {
         self.inner.park(handle)
     }
@@ -154,6 +161,13 @@ cfg_io_driver! {
     }
 
     impl IoStack {
+        pub(crate) fn is_enabled(&self) -> bool {
+            match self {
+                IoStack::Enabled(..) => true,
+                IoStack::Disabled(..) => false,
+            }
+        }
+
         pub(crate) fn park(&mut self, handle: &Handle) {
             match self {
                 IoStack::Enabled(v) => v.park(handle),
@@ -217,6 +231,11 @@ cfg_not_io_driver! {
         pub(crate) fn shutdown(&mut self, _handle: &Handle) {
             self.0.shutdown();
         }
+
+        /// This is not a "real" driver, so it is not considered enabled.
+        pub(crate) fn is_enabled(&self) -> bool {
+            false
+        }
     }
 }
 
@@ -298,6 +317,13 @@ cfg_time! {
     }
 
     impl TimeDriver {
+        pub(crate) fn is_enabled(&self) -> bool {
+            match self {
+                TimeDriver::Enabled { .. } => true,
+                TimeDriver::Disabled(inner) => inner.is_enabled(),
+            }
+        }
+
         pub(crate) fn park(&mut self, handle: &Handle) {
             match self {
                 TimeDriver::Enabled { driver, .. } => driver.park(handle),

tokio/src/runtime/scheduler/inject/shared.rs

@@ -109,6 +109,8 @@ impl<T: 'static> Shared<T> {
     pub(crate) unsafe fn pop_n<'a>(&'a self, synced: &'a mut Synced, n: usize) -> Pop<'a, T> {
         use std::cmp;
 
+        debug_assert!(n > 0);
+
         // safety: All updates to the len atomic are guarded by the mutex. As
         // such, a non-atomic load followed by a store is safe.
         let len = self.len.unsync_load();

tokio/src/runtime/scheduler/inject/synced.rs

@@ -1,3 +1,8 @@
+#![cfg_attr(
+    any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"),
+    allow(dead_code)
+)]
+
 use crate::runtime::task;
 
 pub(crate) struct Synced {
@@ -29,4 +34,8 @@ impl Synced {
         // safety: a `Notified` is pushed into the queue and now it is popped!
         Some(unsafe { task::Notified::from_raw(task) })
     }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        self.head.is_none()
+    }
 }

tokio/src/runtime/scheduler/multi_thread_alt/idle.rs

@@ -60,7 +60,12 @@ impl Idle {
         (idle, synced)
     }
 
+    pub(super) fn needs_searching(&self) -> bool {
+        self.needs_searching.load(Acquire)
+    }
+
     pub(super) fn num_idle(&self, synced: &Synced) -> usize {
+        #[cfg(not(loom))]
         debug_assert_eq!(synced.available_cores.len(), self.num_idle.load(Acquire));
         synced.available_cores.len()
     }
@@ -131,13 +136,7 @@ impl Idle {
         }
 
         // We need to establish a stronger barrier than with `notify_local`
-        if self
-            .num_searching
-            .compare_exchange(0, 1, AcqRel, Acquire)
-            .is_err()
-        {
-            return;
-        }
+        self.num_searching.fetch_add(1, AcqRel);
 
         self.notify_synced(synced, shared);
     }
@@ -158,6 +157,7 @@ impl Idle {
                 synced.assigned_cores[worker] = Some(core);
 
                 let num_idle = synced.idle.available_cores.len();
+                #[cfg(not(loom))]
                 debug_assert_eq!(num_idle, self.num_idle.load(Acquire) - 1);
 
                 // Update the number of sleeping workers
@@ -221,6 +221,7 @@ impl Idle {
             let num_idle = synced.idle.available_cores.len();
             self.num_idle.store(num_idle, Release);
         } else {
+            #[cfg(not(loom))]
             debug_assert_eq!(
                 synced.idle.available_cores.len(),
                 self.num_idle.load(Acquire)
@@ -260,11 +261,11 @@ impl Idle {
         // The core should not be searching at this point
         debug_assert!(!core.is_searching);
 
-        // Check that this isn't the final worker to go idle *and*
-        // `needs_searching` is set.
-        debug_assert!(!self.needs_searching.load(Acquire) || num_active_workers(&synced.idle) > 1);
+        // Check that there are no pending tasks in the global queue
+        debug_assert!(synced.inject.is_empty());
 
         let num_idle = synced.idle.available_cores.len();
+        #[cfg(not(loom))]
         debug_assert_eq!(num_idle, self.num_idle.load(Acquire));
 
         self.idle_map.set(core.index);
@@ -314,7 +315,7 @@ impl Idle {
         }
     }
 
-    fn transition_worker_to_searching(&self, core: &mut Core) {
+    pub(super) fn transition_worker_to_searching(&self, core: &mut Core) {
         core.is_searching = true;
         self.num_searching.fetch_add(1, AcqRel);
         self.needs_searching.store(false, Release);
@@ -324,10 +325,7 @@ impl Idle {
     ///
     /// Returns `true` if this is the final searching worker. The caller
     /// **must** notify a new worker.
-    pub(super) fn transition_worker_from_searching(&self, core: &mut Core) -> bool {
-        debug_assert!(core.is_searching);
-        core.is_searching = false;
-
+    pub(super) fn transition_worker_from_searching(&self) -> bool {
         let prev = self.num_searching.fetch_sub(1, AcqRel);
         debug_assert!(prev > 0);