rt: revert #7757 to fix regression in spawn_blocking (#8057)

This reverts commit 1604bc3.

Unfortunately, this commit introduced a regression that causes programs
using `spawn_blocking` to hang (see #8056). To fix the regression, we
need to undo this change and publish a v1.52.1 release as soon as
possible.

In the future, we may wish to bring back a sharded queue for
`spawn_blocking` tasks, either based on the implementation added in
#7757 or a new one. However, since this is a substantial change to the
runtime internals, I think such a change should probably be done as an
unstable, opt-in `tokio::runtime::Builder` setting initially, so that we
don't regress existing users. I had hoped we could do this now, but
unfortunately, the sharded queue implementation from #7757 is kind of
tightly coupled with the rest of the `spawn_blocking` machinery and
cannot be easily swapped out --- and the hang still occurs with
`NUM_SHARDS` set to 1, so there isn't an easy way to turn it on and off.
Therefore, in the interest of getting a fix out ASAP, this is just a
simple revert.

Fixes #8056

Author: hawkw

spellcheck.dic

@@ -1,4 +1,4 @@
-318
+316
 &
 +
 <
@@ -32,7 +32,6 @@
 8MB
 ABI
 accessors
-adaptively
 adaptor
 adaptors
 Adaptors
@@ -225,7 +224,6 @@ reregistering
 resize
 resized
 RMW
-RNG
 runtime
 runtime's
 runtimes

tokio/src/runtime/blocking/mod.rs

@@ -6,8 +6,6 @@
 mod pool;
 pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner};
 
-mod sharded_queue;
-
 cfg_fs! {
     pub(crate) use pool::spawn_mandatory_blocking;
 }

tokio/src/runtime/blocking/pool.rs

@@ -1,17 +1,16 @@
 //! Thread pool for blocking operations
 
-use crate::loom::sync::{Arc, Mutex};
+use crate::loom::sync::{Arc, Condvar, Mutex};
 use crate::loom::thread;
 use crate::runtime::blocking::schedule::BlockingSchedule;
-use crate::runtime::blocking::sharded_queue::{ShardedQueue, WaitResult};
 use crate::runtime::blocking::{shutdown, BlockingTask};
 use crate::runtime::builder::ThreadNameFn;
 use crate::runtime::task::{self, JoinHandle};
 use crate::runtime::{Builder, Callback, Handle, BOX_FUTURE_THRESHOLD};
 use crate::util::metric_atomics::MetricAtomicUsize;
 use crate::util::trace::{blocking_task, SpawnMeta};
 
-use std::collections::HashMap;
+use std::collections::{HashMap, VecDeque};
 use std::fmt;
 use std::io;
 use std::sync::atomic::Ordering;
@@ -75,12 +74,12 @@ impl SpawnerMetrics {
 }
 
 struct Inner {
-    /// Sharded queue for task distribution.
-    queue: ShardedQueue,
-
-    /// State shared between worker threads (thread management only).
+    /// State shared between worker threads.
     shared: Mutex<Shared>,
 
+    /// Pool threads wait on this.
+    condvar: Condvar,
+
     /// Spawned threads use this name.
     thread_name: ThreadNameFn,
 
@@ -104,6 +103,8 @@ struct Inner {
 }
 
 struct Shared {
+    queue: VecDeque<Task>,
+    num_notify: u32,
     shutdown: bool,
     shutdown_tx: Option<shutdown::Sender>,
     /// Prior to shutdown, we clean up `JoinHandles` by having each timed-out
@@ -213,14 +214,16 @@ impl BlockingPool {
         BlockingPool {
             spawner: Spawner {
                 inner: Arc::new(Inner {
-                    queue: ShardedQueue::new(),
                     shared: Mutex::new(Shared {
+                        queue: VecDeque::new(),
+                        num_notify: 0,
                         shutdown: false,
                         shutdown_tx: Some(shutdown_tx),
                         last_exiting_thread: None,
                         worker_threads: HashMap::new(),
                         worker_thread_index: 0,
                     }),
+                    condvar: Condvar::new(),
                     thread_name: builder.thread_name.clone(),
                     stack_size: builder.thread_stack_size,
                     after_start: builder.after_start.clone(),
@@ -250,7 +253,7 @@ impl BlockingPool {
 
         shared.shutdown = true;
         shared.shutdown_tx = None;
-        self.spawner.inner.queue.shutdown();
+        self.spawner.inner.condvar.notify_all();
 
         let last_exited_thread = std::mem::take(&mut shared.last_exiting_thread);
         let workers = std::mem::take(&mut shared.worker_threads);
@@ -388,8 +391,9 @@ impl Spawner {
     }
 
     fn spawn_task(&self, task: Task, rt: &Handle) -> Result<(), SpawnError> {
-        // Check shutdown without holding the lock
-        if self.inner.queue.is_shutdown() {
+        let mut shared = self.inner.shared.lock();
+
+        if shared.shutdown {
             // Shutdown the task: it's fine to shutdown this task (even if
             // mandatory) because it was scheduled after the shutdown of the
             // runtime began.
@@ -399,64 +403,52 @@ impl Spawner {
             return Err(SpawnError::ShuttingDown);
         }
 
-        // Push to the sharded queue
-        self.inner.queue.push(task);
+        shared.queue.push_back(task);
         self.inner.metrics.inc_queue_depth();
 
-        // Check if we need to spawn a new thread or notify an idle one
         if self.inner.metrics.num_idle_threads() == 0 {
-            // No idle threads - might need to spawn one
-            if self.inner.metrics.num_threads() < self.inner.thread_cap {
-                // Try to spawn a new thread
-                let mut shared = self.inner.shared.lock();
-
-                // Double-check conditions after acquiring the lock
-                if shared.shutdown {
-                    // Shutdown raced with our push. The task is in the
-                    // sharded queue but workers may have already exited.
-                    // Drain it here so mandatory tasks still run.
-                    drop(shared);
-                    while let Some(task) = self.inner.queue.pop(0) {
-                        self.inner.metrics.dec_queue_depth();
-                        task.shutdown_or_run_if_mandatory();
-                    }
-                    return Ok(());
-                }
+            // No threads are able to process the task.
+
+            if self.inner.metrics.num_threads() == self.inner.thread_cap {
+                // At max number of threads
+            } else {
+                assert!(shared.shutdown_tx.is_some());
+                let shutdown_tx = shared.shutdown_tx.clone();
+
+                if let Some(shutdown_tx) = shutdown_tx {
+                    let id = shared.worker_thread_index;
 
-                // Re-check thread count (another thread might have spawned one)
-                if self.inner.metrics.num_threads() < self.inner.thread_cap {
-                    if let Some(shutdown_tx) = shared.shutdown_tx.clone() {
-                        let id = shared.worker_thread_index;
-
-                        match self.spawn_thread(shutdown_tx, rt, id) {
-                            Ok(handle) => {
-                                self.inner.metrics.inc_num_threads();
-                                shared.worker_thread_index += 1;
-                                shared.worker_threads.insert(id, handle);
-                            }
-                            Err(ref e)
-                                if is_temporary_os_thread_error(e)
-                                    && self.inner.metrics.num_threads() > 0 =>
-                            {
-                                // OS temporarily failed to spawn a new thread.
-                                // The task will be picked up eventually by a currently
-                                // busy thread.
-                            }
-                            Err(e) => {
-                                // The OS refused to spawn the thread and there is no thread
-                                // to pick up the task that has just been pushed to the queue.
-                                return Err(SpawnError::NoThreads(e));
-                            }
+                    match self.spawn_thread(shutdown_tx, rt, id) {
+                        Ok(handle) => {
+                            self.inner.metrics.inc_num_threads();
+                            shared.worker_thread_index += 1;
+                            shared.worker_threads.insert(id, handle);
+                        }
+                        Err(ref e)
+                            if is_temporary_os_thread_error(e)
+                                && self.inner.metrics.num_threads() > 0 =>
+                        {
+                            // OS temporarily failed to spawn a new thread.
+                            // The task will be picked up eventually by a currently
+                            // busy thread.
+                        }
+                        Err(e) => {
+                            // The OS refused to spawn the thread and there is no thread
+                            // to pick up the task that has just been pushed to the queue.
+                            return Err(SpawnError::NoThreads(e));
                         }
                     }
                 }
-            } else {
-                // At max threads, notify anyway in case threads are waiting
-                self.inner.queue.notify_one();
             }
         } else {
-            // There are idle threads waiting, notify one
-            self.inner.queue.notify_one();
+            // Notify an idle worker thread. The notification counter
+            // is used to count the needed amount of notifications
+            // exactly. Thread libraries may generate spurious
+            // wakeups, this counter is used to keep us in a
+            // consistent state.
+            self.inner.metrics.dec_num_idle_threads();
+            shared.num_notify += 1;
+            self.inner.condvar.notify_one();
         }
 
         Ok(())
@@ -513,62 +505,94 @@ impl Inner {
             f();
         }
 
-        // Use worker_thread_id as the preferred shard
-        let preferred_shard = worker_thread_id;
+        let mut shared = self.shared.lock();
         let mut join_on_thread = None;
+        // is this thread currently counted in `num_idle_threads`?
+        let mut is_counted_idle;
 
         'main: loop {
-            // BUSY: Process tasks from the queue
-            while let Some(task) = self.queue.pop(preferred_shard) {
+            // BUSY
+            while let Some(task) = shared.queue.pop_front() {
                 self.metrics.dec_queue_depth();
+                drop(shared);
                 task.run();
-            }
 
-            // Check for shutdown before going idle
-            if self.queue.is_shutdown() {
-                break;
+                shared = self.shared.lock();
             }
 
-            // IDLE: Wait for new tasks (spurious wakeups handled internally)
+            // IDLE
             self.metrics.inc_num_idle_threads();
-
-            match self.queue.wait_for_task(preferred_shard, self.keep_alive) {
-                WaitResult::Task(task) => {
-                    self.metrics.dec_num_idle_threads();
-                    self.metrics.dec_queue_depth();
-                    task.run();
-                }
-                WaitResult::Shutdown => {
-                    self.metrics.dec_num_idle_threads();
-                    break 'main;
+            // mark this thread as currently counted in `num_idle_threads`.
+            is_counted_idle = true;
+
+            while !shared.shutdown {
+                let lock_result = self.condvar.wait_timeout(shared, self.keep_alive).unwrap();
+
+                shared = lock_result.0;
+                let timeout_result = lock_result.1;
+
+                if shared.num_notify != 0 {
+                    // We have received a legitimate wakeup,
+                    // acknowledge it by decrementing the counter
+                    // and transition to the BUSY state.
+                    shared.num_notify -= 1;
+                    // since this is a legitimate wakeup,
+                    // the `Spawner::spawn_task` has already decremented `num_idle_threads`.
+                    is_counted_idle = false;
+                    break;
                 }
-                WaitResult::Timeout => {
-                    self.metrics.dec_num_idle_threads();
-
-                    // Clean up thread handle
-                    let mut shared = self.shared.lock();
-                    if !shared.shutdown {
-                        let my_handle = shared.worker_threads.remove(&worker_thread_id);
-                        join_on_thread =
-                            std::mem::replace(&mut shared.last_exiting_thread, my_handle);
-                    }
+
+                // Even if the condvar "timed out", if the pool is entering the
+                // shutdown phase, we want to perform the cleanup logic.
+                if !shared.shutdown && timeout_result.timed_out() {
+                    // We'll join the prior timed-out thread's JoinHandle after dropping the lock.
+                    // This isn't done when shutting down, because the thread calling shutdown will
+                    // handle joining everything.
+                    let my_handle = shared.worker_threads.remove(&worker_thread_id);
+                    join_on_thread = std::mem::replace(&mut shared.last_exiting_thread, my_handle);
 
                     break 'main;
                 }
+
+                // Spurious wakeup detected, go back to sleep.
             }
-        }
 
-        // Drain remaining tasks if shutting down
-        if self.queue.is_shutdown() {
-            while let Some(task) = self.queue.pop(preferred_shard) {
-                self.metrics.dec_queue_depth();
-                task.shutdown_or_run_if_mandatory();
+            if shared.shutdown {
+                // Drain the queue
+                while let Some(task) = shared.queue.pop_front() {
+                    self.metrics.dec_queue_depth();
+                    drop(shared);
+
+                    task.shutdown_or_run_if_mandatory();
+
+                    shared = self.shared.lock();
+                }
+
+                break;
             }
         }
 
         // Thread exit
         self.metrics.dec_num_threads();
 
+        // Is this thread currently counted in `num_idle_threads`?
+        if is_counted_idle {
+            // `num_idle_threads` should now be tracked exactly, panic
+            // with a descriptive message if it is not the
+            // case.
+            let prev_idle = self.metrics.dec_num_idle_threads();
+            assert_ne!(
+                prev_idle, 0,
+                "`num_idle_threads` underflowed on thread exit"
+            );
+        }
+
+        if shared.shutdown && self.metrics.num_threads() == 0 {
+            self.condvar.notify_one();
+        }
+
+        drop(shared);
+
         if let Some(f) = &self.before_stop {
             f();
         }