From 1604bc335157be9a131e24cfde55cab3c90ebc92 Mon Sep 17 00:00:00 2001 From: Alex Gaynor Date: Fri, 10 Apr 2026 10:06:19 -0400 Subject: [PATCH] rt: improve `spawn_blocking` scalability with sharded queue (#7757) --- spellcheck.dic | 4 +- tokio/src/runtime/blocking/mod.rs | 2 + tokio/src/runtime/blocking/pool.rs | 208 ++++++++--------- tokio/src/runtime/blocking/sharded_queue.rs | 238 ++++++++++++++++++++ tokio/src/runtime/context.rs | 2 +- tokio/src/util/rand.rs | 6 +- 6 files changed, 337 insertions(+), 123 deletions(-) create mode 100644 tokio/src/runtime/blocking/sharded_queue.rs diff --git a/spellcheck.dic b/spellcheck.dic index 1169078d3..3a7e586c0 100644 --- a/spellcheck.dic +++ b/spellcheck.dic @@ -1,4 +1,4 @@ -314 +316 & + < @@ -32,6 +32,7 @@ 8MB ABI accessors +adaptively adaptor adaptors Adaptors @@ -223,6 +224,7 @@ reregistering resize resized RMW +RNG runtime runtime's runtimes diff --git a/tokio/src/runtime/blocking/mod.rs b/tokio/src/runtime/blocking/mod.rs index c42924be7..0eaf33fc3 100644 --- a/tokio/src/runtime/blocking/mod.rs +++ b/tokio/src/runtime/blocking/mod.rs @@ -6,6 +6,8 @@ mod pool; pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner}; +mod sharded_queue; + cfg_fs! { pub(crate) use pool::spawn_mandatory_blocking; } diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index dae98bc94..38948038e 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -1,8 +1,9 @@ //! Thread pool for blocking operations -use crate::loom::sync::{Arc, Condvar, Mutex}; +use crate::loom::sync::{Arc, 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}; @@ -10,7 +11,7 @@ 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, VecDeque}; +use std::collections::HashMap; use std::fmt; use std::io; use std::sync::atomic::Ordering; @@ -74,11 +75,11 @@ impl SpawnerMetrics { } struct Inner { - /// State shared between worker threads. - shared: Mutex, + /// Sharded queue for task distribution. + queue: ShardedQueue, - /// Pool threads wait on this. - condvar: Condvar, + /// State shared between worker threads (thread management only). + shared: Mutex, /// Spawned threads use this name. thread_name: ThreadNameFn, @@ -103,8 +104,6 @@ struct Inner { } struct Shared { - queue: VecDeque, - num_notify: u32, shutdown: bool, shutdown_tx: Option, /// Prior to shutdown, we clean up `JoinHandles` by having each timed-out @@ -214,16 +213,14 @@ 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(), @@ -253,7 +250,7 @@ impl BlockingPool { shared.shutdown = true; shared.shutdown_tx = None; - self.spawner.inner.condvar.notify_all(); + self.spawner.inner.queue.shutdown(); let last_exited_thread = std::mem::take(&mut shared.last_exiting_thread); let workers = std::mem::take(&mut shared.worker_threads); @@ -391,9 +388,8 @@ impl Spawner { } fn spawn_task(&self, task: Task, rt: &Handle) -> Result<(), SpawnError> { - let mut shared = self.inner.shared.lock(); - - if shared.shutdown { + // Check shutdown without holding the lock + if self.inner.queue.is_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. @@ -403,52 +399,64 @@ impl Spawner { return Err(SpawnError::ShuttingDown); } - shared.queue.push_back(task); + // Push to the sharded queue + self.inner.queue.push(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 threads are able to process the task. + // 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(); - 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(); + // 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(()); + } - 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 { - // 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(); + // There are idle threads waiting, notify one + self.inner.queue.notify_one(); } Ok(()) @@ -505,94 +513,62 @@ impl Inner { f(); } - let mut shared = self.shared.lock(); + // Use worker_thread_id as the preferred shard + let preferred_shard = worker_thread_id; let mut join_on_thread = None; - // is this thread currently counted in `num_idle_threads`? - let mut is_counted_idle; 'main: loop { - // BUSY - while let Some(task) = shared.queue.pop_front() { + // BUSY: Process tasks from the queue + while let Some(task) = self.queue.pop(preferred_shard) { self.metrics.dec_queue_depth(); - drop(shared); task.run(); - - shared = self.shared.lock(); } - // IDLE + // Check for shutdown before going idle + if self.queue.is_shutdown() { + break; + } + + // IDLE: Wait for new tasks (spurious wakeups handled internally) self.metrics.inc_num_idle_threads(); - // 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; + 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; + } + WaitResult::Timeout => { + self.metrics.dec_num_idle_threads(); - // 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); + // 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); + } break 'main; } - - // Spurious wakeup detected, go back to sleep. } + } - 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; + // 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(); } } // 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(); } diff --git a/tokio/src/runtime/blocking/sharded_queue.rs b/tokio/src/runtime/blocking/sharded_queue.rs new file mode 100644 index 000000000..2fcc925c4 --- /dev/null +++ b/tokio/src/runtime/blocking/sharded_queue.rs @@ -0,0 +1,238 @@ +//! A sharded concurrent queue for the blocking pool. +//! +//! This implementation distributes tasks across multiple shards to reduce +//! lock contention when many threads are spawning blocking tasks concurrently. +//! The push operations use per-shard locking, while notifications use a global +//! condvar for simplicity. +//! +//! For shard selection, we use the same approach as `sync::watch`: prefer +//! randomness when available to reduce contention, falling back to circular +//! access when the random number generator is not available. + +use crate::loom::sync::{Condvar, Mutex}; + +use std::collections::VecDeque; +use std::sync::atomic::AtomicBool; +#[cfg(loom)] +use std::sync::atomic::AtomicUsize; +#[cfg(loom)] +use std::sync::atomic::Ordering::Relaxed; +use std::sync::atomic::Ordering::{Acquire, Release}; +use std::time::Duration; + +use super::pool::Task; + +/// Number of shards. Must be a power of 2. +/// Under loom, use a single shard to keep the state space tractable — +/// the concurrency properties we need to verify (condvar signaling, +/// shutdown ordering) are independent of shard count. +#[cfg(not(loom))] +const NUM_SHARDS: usize = 16; +#[cfg(loom)] +const NUM_SHARDS: usize = 1; + +/// A single shard containing a queue protected by its own mutex. +struct Shard { + /// The task queue for this shard. + queue: Mutex>, +} + +impl Shard { + fn new() -> Self { + Shard { + queue: Mutex::new(VecDeque::new()), + } + } + + /// Push a task to this shard's queue. + fn push(&self, task: Task) { + let mut queue = self.queue.lock(); + + // Check if pushing would require reallocation (when len == capacity). + // If so, allocate outside the lock to avoid blocking readers. + while queue.len() == queue.capacity() { + let current_len = queue.len(); + // Use 2x growth factor, minimum 4 + let new_cap = current_len.saturating_mul(2).max(4); + + // Release lock before allocating + drop(queue); + + let mut new_queue = VecDeque::with_capacity(new_cap); + + queue = self.queue.lock(); + // If the queue is: + // a) Not full anymore => push to the current queue + // b) Full and our new queue is big enough => copy items to the new + // queue and push to it. + // c) Full and our new queue is too small => try again. + if queue.len() == queue.capacity() { + if new_queue.capacity() > queue.len() { + new_queue.extend(queue.drain(..)); + *queue = new_queue; + break; + } + } else { + break; + } + } + + queue.push_back(task); + } + + /// Try to pop a task from this shard's queue. + fn pop(&self) -> Option { + let mut queue = self.queue.lock(); + queue.pop_front() + } +} + +/// A sharded queue that distributes tasks across multiple shards. +pub(super) struct ShardedQueue { + /// The shards - each with its own mutex-protected queue. + shards: [Shard; NUM_SHARDS], + /// Atomic counter for round-robin task distribution. + /// Only used when randomness is not available (loom). + #[cfg(loom)] + push_index: AtomicUsize, + /// Global shutdown flag. + shutdown: AtomicBool, + /// Global condition variable for worker notifications. + /// We use a single condvar to avoid the complexity of per-shard waiting. + condvar: Condvar, + /// Mutex paired with the condvar. Protects the notification counter + /// (`num_notify`), which tracks how many tasks have been pushed and + /// need to be picked up by idle workers. + condvar_mutex: Mutex, +} + +impl ShardedQueue { + pub(super) fn new() -> Self { + ShardedQueue { + shards: std::array::from_fn(|_| Shard::new()), + #[cfg(loom)] + push_index: AtomicUsize::new(0), + shutdown: AtomicBool::new(false), + condvar: Condvar::new(), + condvar_mutex: Mutex::new(0), + } + } + + /// Select the next shard index for pushing a task -- when the RNG is + /// available. + #[cfg(not(loom))] + fn next_push_index(&self, num_shards: usize) -> usize { + crate::runtime::context::thread_rng_n(num_shards as u32) as usize + } + + /// Select the next shard index for pushing a task -- when the RNG is not + /// available (loom). + #[cfg(loom)] + fn next_push_index(&self, num_shards: usize) -> usize { + self.push_index.fetch_add(1, Relaxed) & (num_shards - 1) + } + + /// Push a task to the queue. + pub(super) fn push(&self, task: Task) { + let index = self.next_push_index(NUM_SHARDS); + self.shards[index].push(task); + } + + /// Notify one waiting worker that a task is available. + /// + /// Increments the notification counter under `condvar_mutex` and signals + /// the condvar. The counter acts as a persistent notification that cannot + /// be lost — even if no worker is currently waiting on the condvar, the + /// next worker to enter `wait_for_task` will see the counter and know + /// there is work to do. + pub(super) fn notify_one(&self) { + let mut guard = self.condvar_mutex.lock(); + *guard += 1; + drop(guard); + self.condvar.notify_one(); + } + + /// Try to pop a task, checking the preferred shard first, then others. + pub(super) fn pop(&self, preferred_shard: usize) -> Option { + // Check shards starting from preferred, wrapping around + let start = preferred_shard % NUM_SHARDS; + for i in 0..NUM_SHARDS { + let index = (start + i) % NUM_SHARDS; + if let Some(task) = self.shards[index].pop() { + return Some(task); + } + } + + None + } + + /// Set the shutdown flag and wake all workers. + pub(super) fn shutdown(&self) { + // Set the flag while holding condvar_mutex so that any worker + // currently inside `wait_for_task` (which also holds condvar_mutex + // while checking) is guaranteed to see the flag on its next check. + { + let _guard = self.condvar_mutex.lock(); + self.shutdown.store(true, Release); + } + self.condvar.notify_all(); + } + + /// Check if shutdown has been initiated. + pub(super) fn is_shutdown(&self) -> bool { + self.shutdown.load(Acquire) + } + + /// Wait for a task notification with timeout. Returns when a task has + /// been pushed (the caller should then `pop`), shutdown occurs, or the + /// wait times out. + /// + /// Uses a notification counter under `condvar_mutex` to prevent lost + /// wakeups — the same pattern as the original single-mutex blocking pool. + pub(super) fn wait_for_task(&self, preferred_shard: usize, timeout: Duration) -> WaitResult { + let mut guard = self.condvar_mutex.lock(); + + loop { + if self.is_shutdown() { + return WaitResult::Shutdown; + } + + if *guard > 0 { + // A notification is pending — a task was pushed. + *guard -= 1; + drop(guard); + // Pop outside the condvar_mutex to avoid holding two locks. + if let Some(task) = self.pop(preferred_shard) { + return WaitResult::Task(task); + } + // The task was already consumed in the caller's BUSY loop + // (race between push+notify and the worker's pop loop). + // Re-enter the wait. + guard = self.condvar_mutex.lock(); + continue; + } + + let (g, timeout_result) = self.condvar.wait_timeout(guard, timeout).unwrap(); + guard = g; + + if timeout_result.timed_out() && *guard == 0 { + // Double-check: shutdown may have raced with the timeout. + if self.is_shutdown() { + return WaitResult::Shutdown; + } + return WaitResult::Timeout; + } + // Woken by notify or spurious wakeup — loop back to check. + } + } +} + +/// Result of waiting for a task. +pub(super) enum WaitResult { + /// A task was found. + Task(Task), + /// The wait timed out. + Timeout, + /// Shutdown was initiated. + Shutdown, +} diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index 3bf991940..42e17bc66 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -121,7 +121,7 @@ tokio_thread_local! { } } -#[cfg(any(feature = "macros", all(feature = "sync", feature = "rt")))] +#[cfg(any(feature = "macros", feature = "rt"))] pub(crate) fn thread_rng_n(n: u32) -> u32 { CONTEXT.with(|ctx| { let mut rng = ctx.rng.get().unwrap_or_else(FastRand::new); diff --git a/tokio/src/util/rand.rs b/tokio/src/util/rand.rs index 67c45693c..52f0962c3 100644 --- a/tokio/src/util/rand.rs +++ b/tokio/src/util/rand.rs @@ -68,11 +68,7 @@ impl FastRand { } } - #[cfg(any( - feature = "macros", - feature = "rt-multi-thread", - all(feature = "sync", feature = "rt") - ))] + #[cfg(any(feature = "macros", feature = "sync", feature = "rt"))] pub(crate) fn fastrand_n(&mut self, n: u32) -> u32 { // This is similar to fastrand() % n, but faster. // See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/