diff --git a/spellcheck.dic b/spellcheck.dic index e3cbfe077..0cf8ae850 100644 --- a/spellcheck.dic +++ b/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 diff --git a/tokio/src/runtime/blocking/mod.rs b/tokio/src/runtime/blocking/mod.rs index 0eaf33fc3..c42924be7 100644 --- a/tokio/src/runtime/blocking/mod.rs +++ b/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; } diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 38948038e..dae98bc94 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -1,9 +1,8 @@ //! 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}; @@ -11,7 +10,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; +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, + /// 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, + num_notify: u32, shutdown: bool, shutdown_tx: Option, /// 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(); + // No threads are able to process the task. - // 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 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(); - // 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; + if let Some(shutdown_tx) = shutdown_tx { + 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(); + + shared = self.shared.lock(); } - // Check for shutdown before going idle - if self.queue.is_shutdown() { - break; - } - - // IDLE: Wait for new tasks (spurious wakeups handled internally) + // IDLE self.metrics.inc_num_idle_threads(); + // mark this thread as currently counted in `num_idle_threads`. + is_counted_idle = true; - 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(); + while !shared.shutdown { + let lock_result = self.condvar.wait_timeout(shared, self.keep_alive).unwrap(); - // 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); - } + 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; + } + + // 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(); } diff --git a/tokio/src/runtime/blocking/sharded_queue.rs b/tokio/src/runtime/blocking/sharded_queue.rs deleted file mode 100644 index 2fcc925c4..000000000 --- a/tokio/src/runtime/blocking/sharded_queue.rs +++ /dev/null @@ -1,238 +0,0 @@ -//! 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 42e17bc66..3bf991940 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -121,7 +121,7 @@ tokio_thread_local! { } } -#[cfg(any(feature = "macros", feature = "rt"))] +#[cfg(any(feature = "macros", all(feature = "sync", 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 52f0962c3..67c45693c 100644 --- a/tokio/src/util/rand.rs +++ b/tokio/src/util/rand.rs @@ -68,7 +68,11 @@ impl FastRand { } } - #[cfg(any(feature = "macros", feature = "sync", feature = "rt"))] + #[cfg(any( + feature = "macros", + feature = "rt-multi-thread", + all(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/