diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c5a023538..974a2050f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -315,6 +315,30 @@ jobs: # the unstable cfg to RustDoc RUSTDOCFLAGS: --cfg tokio_unstable + # Run the test suite with the sharded `spawn_blocking` queue. + test-sharded-blocking-queue: + name: test tokio full with sharded blocking queue + needs: basics + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - name: Install Rust ${{ env.rust_stable }} + uses: dtolnay/rust-toolchain@stable + with: + toolchain: ${{ env.rust_stable }} + + - name: Install cargo-nextest + uses: taiki-e/install-action@v2 + with: + tool: cargo-nextest + + - uses: Swatinem/rust-cache@v2 + - name: test tokio full with sharded blocking queue + run: cargo nextest run --features full + working-directory: tokio + env: + TOKIO_UNSTABLE_SHARDED_BLOCKING_QUEUE: "1" + test-unstable-taskdump: name: test tokio full --unstable --taskdump needs: basics diff --git a/.github/workflows/loom.yml b/.github/workflows/loom.yml index 904958d80..c2fc66ff0 100644 --- a/.github/workflows/loom.yml +++ b/.github/workflows/loom.yml @@ -28,6 +28,13 @@ jobs: # base_ref is null when it's not a pull request if: github.repository_owner == 'tokio-rs' && (contains(github.event.pull_request.labels.*.name, 'R-loom-blocking') || (github.base_ref == null)) runs-on: ubuntu-latest + strategy: + matrix: + # Run the blocking pool loom tests against both `spawn_blocking` + # queue implementations. + include: + - sharded_blocking_queue: "0" + - sharded_blocking_queue: "1" steps: - uses: actions/checkout@v7 - name: Install Rust ${{ env.rust_stable }} @@ -38,6 +45,8 @@ jobs: - name: run tests run: cargo test --lib --release --features full -- --nocapture loom_blocking working-directory: tokio + env: + TOKIO_UNSTABLE_SHARDED_BLOCKING_QUEUE: ${{ matrix.sharded_blocking_queue }} loom-sync: name: loom tokio::sync diff --git a/spellcheck.dic b/spellcheck.dic index aeb1e9db4..9d4479334 100644 --- a/spellcheck.dic +++ b/spellcheck.dic @@ -1,4 +1,4 @@ -324 +327 & + < @@ -141,6 +141,7 @@ hashsets HdrHistogram ICMP ie +iff Illumos impl implementers @@ -229,6 +230,7 @@ reregistering resize resized RMW +RNG runtime runtime's runtimes @@ -244,6 +246,7 @@ signalling SmallCrush Solaris spawner +spawners Splitter spmc spsc diff --git a/tokio/src/runtime/blocking/mod.rs b/tokio/src/runtime/blocking/mod.rs index c42924be7..b0e42264d 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; + 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 9fd5ffe21..5507bf136 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -3,6 +3,7 @@ use crate::loom::sync::{Arc, Condvar, Mutex}; use crate::loom::thread; use crate::runtime::blocking::schedule::BlockingSchedule; +use crate::runtime::blocking::sharded::ShardedImpl; use crate::runtime::blocking::{shutdown, BlockingTask}; use crate::runtime::builder::ThreadNameFn; use crate::runtime::task::{self, JoinHandle}; @@ -34,11 +35,11 @@ pub(crate) struct SpawnerMetrics { } impl SpawnerMetrics { - fn num_threads(&self) -> usize { + pub(super) fn num_threads(&self) -> usize { self.num_threads.load(Ordering::Relaxed) } - fn num_idle_threads(&self) -> usize { + pub(super) fn num_idle_threads(&self) -> usize { self.num_idle_threads.load(Ordering::Relaxed) } @@ -52,23 +53,23 @@ impl SpawnerMetrics { self.num_threads.increment(); } - fn dec_num_threads(&self) { + pub(super) fn dec_num_threads(&self) { self.num_threads.decrement(); } - fn inc_num_idle_threads(&self) { + pub(super) fn inc_num_idle_threads(&self) { self.num_idle_threads.increment(); } - fn dec_num_idle_threads(&self) -> usize { + pub(super) fn dec_num_idle_threads(&self) -> usize { self.num_idle_threads.decrement() } - fn inc_queue_depth(&self) { + pub(super) fn inc_queue_depth(&self) { self.queue_depth.increment(); } - fn dec_queue_depth(&self) { + pub(super) fn dec_queue_depth(&self) { self.queue_depth.decrement(); } } @@ -102,6 +103,7 @@ struct Inner { /// Per-variant queue + notification + lock topology. enum InnerImpl { Locked(LockedImpl), + Sharded(ShardedImpl), } /// Single-mutex + condvar implementation. @@ -121,27 +123,56 @@ struct LockedInner { /// State handed back from `InnerImpl::begin_shutdown` to the caller so it /// can join the worker threads after the wait completes: an optional /// previously-timed-out worker, plus the map of currently-running workers. -type ShutdownHandles = ( +pub(super) type ShutdownHandles = ( Option>, HashMap>, ); /// Thread-management state used by every `InnerImpl` variant. -struct ThreadManagementState { - shutdown: bool, - shutdown_tx: Option, +pub(super) struct ThreadManagementState { + pub(super) shutdown: bool, + pub(super) shutdown_tx: Option, /// Prior to shutdown, we clean up `JoinHandles` by having each timed-out /// thread join on the previous timed-out thread. This is not strictly /// necessary but helps avoid Valgrind false positives, see /// /// for more information. - last_exiting_thread: Option>, + pub(super) last_exiting_thread: Option>, /// This holds the `JoinHandles` for all running threads; on shutdown, the thread /// calling shutdown handles joining on these. - worker_threads: HashMap>, + pub(super) worker_threads: HashMap>, /// This is a counter used to iterate `worker_threads` in a consistent order (for loom's /// benefit). - worker_thread_index: usize, + pub(super) worker_thread_index: usize, +} + +impl ThreadManagementState { + /// Flag the pool as shutting down and hand back the worker `JoinHandle`s + /// to join. Returns `None` if shutdown has already begun. The caller is + /// responsible for waking all waiting workers. + pub(super) fn begin_shutdown(&mut self) -> Option { + if self.shutdown { + return None; + } + self.shutdown = true; + self.shutdown_tx = None; + + let last_exited_thread = std::mem::take(&mut self.last_exiting_thread); + let workers = std::mem::take(&mut self.worker_threads); + Some((last_exited_thread, workers)) + } + + /// Bookkeeping for a worker exiting on its keep-alive timeout: leaves + /// its own handle for the next timed-out worker (or shutdown) to join, + /// and returns the previous timed-out thread's handle, which the caller + /// must join after releasing the lock. + pub(super) fn worker_timed_out( + &mut self, + worker_thread_id: usize, + ) -> Option> { + let my_handle = self.worker_threads.remove(&worker_thread_id); + std::mem::replace(&mut self.last_exiting_thread, my_handle) + } } pub(crate) struct Task { @@ -180,11 +211,15 @@ impl Task { Task { task, mandatory } } - fn run(self) { + pub(super) fn shutdown(self) { + self.task.shutdown(); + } + + pub(super) fn run(self) { self.task.run(); } - fn shutdown_or_run_if_mandatory(self) { + pub(super) fn shutdown_or_run_if_mandatory(self) { match self.mandatory { Mandatory::NonMandatory => self.task.shutdown(), Mandatory::Mandatory => self.task.run(), @@ -234,23 +269,24 @@ impl BlockingPool { let (shutdown_tx, shutdown_rx) = shutdown::channel(); let keep_alive = builder.keep_alive.unwrap_or(KEEP_ALIVE); + let thread_mgmt_state = ThreadManagementState { + shutdown: false, + shutdown_tx: Some(shutdown_tx), + last_exiting_thread: None, + worker_threads: HashMap::new(), + worker_thread_index: 0, + }; + + let inner_impl = if builder.sharded_blocking_queue { + InnerImpl::Sharded(ShardedImpl::new(thread_mgmt_state)) + } else { + InnerImpl::Locked(LockedImpl::new(thread_mgmt_state)) + }; + BlockingPool { spawner: Spawner { inner: Arc::new(Inner { - inner_impl: InnerImpl::Locked(LockedImpl { - mutex: Mutex::new(LockedInner { - queue: VecDeque::new(), - num_notify: 0, - thread_mgmt_state: ThreadManagementState { - shutdown: false, - shutdown_tx: Some(shutdown_tx), - last_exiting_thread: None, - worker_threads: HashMap::new(), - worker_thread_index: 0, - }, - }), - condvar: Condvar::new(), - }), + inner_impl, thread_name: builder.thread_name.clone(), stack_size: builder.thread_stack_size, after_start: builder.after_start.clone(), @@ -272,7 +308,8 @@ impl BlockingPool { // The function can be called multiple times. First, by explicitly // calling `shutdown` then by the drop handler calling `shutdown`. This // prevents shutting down twice. - let (last_exited_thread, workers) = match self.spawner.inner.inner_impl.begin_shutdown() { + let inner = &self.spawner.inner; + let (last_exited_thread, workers) = match inner.inner_impl.begin_shutdown(&inner.metrics) { Some(x) => x, None => return, }; @@ -509,6 +546,7 @@ impl InnerImpl { { match self { InnerImpl::Locked(l) => l.spawn_task(task, metrics, on_no_idle), + InnerImpl::Sharded(s) => s.spawn_task(task, metrics, on_no_idle), } } @@ -520,12 +558,14 @@ impl InnerImpl { ) -> Option> { match self { InnerImpl::Locked(l) => l.run_worker(metrics, keep_alive, worker_thread_id), + InnerImpl::Sharded(s) => s.run_worker(metrics, keep_alive, worker_thread_id), } } - fn begin_shutdown(&self) -> Option { + fn begin_shutdown(&self, metrics: &SpawnerMetrics) -> Option { match self { InnerImpl::Locked(l) => l.begin_shutdown(), + InnerImpl::Sharded(s) => s.begin_shutdown(metrics), } } } @@ -535,6 +575,17 @@ impl InnerImpl { // code was refactored to enable adding a sharded queue implementation, this // was self-evidently behaviorally identical to the original implementation. impl LockedImpl { + fn new(thread_mgmt_state: ThreadManagementState) -> LockedImpl { + LockedImpl { + mutex: Mutex::new(LockedInner { + queue: VecDeque::new(), + num_notify: 0, + thread_mgmt_state, + }), + condvar: Condvar::new(), + } + } + /// Push a task and either notify an idle worker or invoke /// `on_no_idle` (which is responsible for spawning a new worker if /// possible). @@ -625,19 +676,7 @@ impl LockedImpl { // entering the shutdown phase, we want to perform // the cleanup logic. if !locked.thread_mgmt_state.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 = locked - .thread_mgmt_state - .worker_threads - .remove(&worker_thread_id); - join_on_thread = std::mem::replace( - &mut locked.thread_mgmt_state.last_exiting_thread, - my_handle, - ); + join_on_thread = locked.thread_mgmt_state.worker_timed_out(worker_thread_id); break 'main; } @@ -689,16 +728,9 @@ impl LockedImpl { /// `JoinHandle`s for the caller to join. fn begin_shutdown(&self) -> Option { let mut locked = self.mutex.lock(); - if locked.thread_mgmt_state.shutdown { - return None; - } - locked.thread_mgmt_state.shutdown = true; - locked.thread_mgmt_state.shutdown_tx = None; + let handles = locked.thread_mgmt_state.begin_shutdown()?; self.condvar.notify_all(); - - let last_exited_thread = std::mem::take(&mut locked.thread_mgmt_state.last_exiting_thread); - let workers = std::mem::take(&mut locked.thread_mgmt_state.worker_threads); - Some((last_exited_thread, workers)) + Some(handles) } } diff --git a/tokio/src/runtime/blocking/sharded.rs b/tokio/src/runtime/blocking/sharded.rs new file mode 100644 index 000000000..895acf9bb --- /dev/null +++ b/tokio/src/runtime/blocking/sharded.rs @@ -0,0 +1,353 @@ +//! A sharded queue for the blocking pool's tasks. +//! +//! Tasks live in `NUM_SHARDS` queues, each with its own mutex. Spawners push +//! to a shard chosen via the thread-local RNG; workers pop by scanning the +//! shards, starting from one derived from their worker id. A mask tracks which +//! shards have tasks so scans rarely lock empty shards. +//! +//! Worker lifecycle (thread spawning, parking/waking, timeouts, shutdown) is +//! coordinated by the `coord` mutex + `condvar`, using the same claim +//! protocol as the default single-mutex queue. Unlike that queue, `coord` is +//! held only for the claim-or-spawn decision and for a worker's transition +//! to idle — never around queue operations or task execution — so spawners +//! and workers contend mostly on `1/NUM_SHARDS` of a shard lock each. + +use crate::loom::sync::atomic::{AtomicBool, AtomicUsize}; +use crate::loom::sync::{Condvar, Mutex}; +use crate::loom::thread; + +use std::collections::VecDeque; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use super::pool::{ShutdownHandles, SpawnError, SpawnerMetrics, Task, ThreadManagementState}; + +/// Number of shards. Must be a power of two. +/// +/// Under loom, use a small shard count to keep the state space tractable +/// while still exercising the cross-shard scanning in `pop`. +#[cfg(not(loom))] +const NUM_SHARDS: usize = 16; +#[cfg(loom)] +const NUM_SHARDS: usize = 2; + +struct Shard { + queue: VecDeque, + /// Set (under the shard's lock) when the shard is drained for shutdown; + /// pushes to a sealed shard are rejected. This is what guarantees that a + /// spawner racing with shutdown cannot leave a task behind: a push + /// either loses (rejected, task is shut down) or wins, in which case the + /// sealer that later drains this shard collects the task. + sealed: bool, +} + +pub(super) struct ShardedImpl { + shards: [Mutex; NUM_SHARDS], + /// One bit per shard, set when that shard's queue is non-empty. Only + /// updated while holding that shard's lock, so it is exact at every + /// lock release; unlocked loads may be stale (see `pop`). + non_empty_mask: AtomicUsize, + coord: Mutex, + condvar: Condvar, + /// Mirror of `ThreadManagementState::shutdown`, to let spawners reject + /// tasks without taking `coord`. + is_shutdown: AtomicBool, + /// Round-robin push counter, used to pick a shard when the thread-local + /// RNG is unavailable (loom requires each execution to be deterministic). + #[cfg(loom)] + push_index: AtomicUsize, +} + +/// State protected by `ShardedImpl::coord`. This mirrors `LockedInner`, +/// except that the queue itself lives in the shards. +struct ShardedCoord { + /// Pending worker wakeups. A spawner claims an idle worker by + /// decrementing `num_idle_threads` and incrementing this; a woken worker + /// acknowledges by decrementing it. Distinguishes real wakeups from + /// spurious ones. + num_notify: u32, + thread_mgmt_state: ThreadManagementState, +} + +impl ShardedImpl { + pub(super) fn new(thread_mgmt_state: ThreadManagementState) -> ShardedImpl { + ShardedImpl { + shards: std::array::from_fn(|_| { + Mutex::new(Shard { + queue: VecDeque::new(), + sealed: false, + }) + }), + non_empty_mask: AtomicUsize::new(0), + coord: Mutex::new(ShardedCoord { + num_notify: 0, + thread_mgmt_state, + }), + condvar: Condvar::new(), + is_shutdown: AtomicBool::new(false), + #[cfg(loom)] + push_index: AtomicUsize::new(0), + } + } + + /// Pick a shard to push to. Use the thread-local RNG so that concurrent + /// spawners spread across the shards. + #[cfg(not(loom))] + fn push_shard_index(&self) -> usize { + crate::runtime::context::thread_rng_n(NUM_SHARDS as u32) as usize + } + + /// Under loom the RNG would make each execution take a different path, + /// breaking loom's requirement that executions be deterministic, so use + /// round-robin selection instead. + #[cfg(loom)] + fn push_shard_index(&self) -> usize { + self.push_index.fetch_add(1, Ordering::Relaxed) % NUM_SHARDS + } + + /// Push a task onto one of the shards, or hand it back if the chosen + /// shard has been sealed for shutdown. + /// + /// The queue-depth metric is incremented under the shard lock so that + /// the pop that consumes this task (whose decrement is ordered after + /// this lock's release) can never transiently wrap the counter. + fn push(&self, task: Task, metrics: &SpawnerMetrics) -> Result<(), Task> { + let index = self.push_shard_index(); + let mut shard = self.shards[index].lock(); + if shard.sealed { + return Err(task); + } + shard.queue.push_back(task); + metrics.inc_queue_depth(); + if shard.queue.len() == 1 { + self.non_empty_mask.fetch_or(1 << index, Ordering::Relaxed); + } + Ok(()) + } + + /// Pop a task, checking the worker's preferred shard first. + fn pop(&self, preferred_shard: usize) -> Option { + let mask = self.non_empty_mask.load(Ordering::Relaxed); + if mask == 0 { + return None; + } + + let start = preferred_shard % NUM_SHARDS; + for i in 0..NUM_SHARDS { + let index = (start + i) % NUM_SHARDS; + if mask & (1 << index) == 0 { + continue; + } + + let mut shard = self.shards[index].lock(); + match shard.queue.pop_front() { + Some(task) => { + if shard.queue.is_empty() { + self.non_empty_mask + .fetch_and(!(1 << index), Ordering::Relaxed); + } + return Some(task); + } + None => { + // The shard was emptied (and its bit cleared) after the + // mask was loaded; move on to the next candidate. + } + } + } + + None + } + + /// Drain every shard, sealing each so that later pushes are rejected, + /// and run-or-cancel the collected tasks. Called by workers during + /// shutdown (and by `begin_shutdown` when there are no workers left to + /// do it). Sealing is idempotent, so concurrent callers are fine. + fn drain_and_seal(&self, metrics: &SpawnerMetrics, preferred_shard: usize) { + let start = preferred_shard % NUM_SHARDS; + for i in 0..NUM_SHARDS { + let index = (start + i) % NUM_SHARDS; + let tasks = { + let mut shard = self.shards[index].lock(); + shard.sealed = true; + self.non_empty_mask + .fetch_and(!(1 << index), Ordering::Relaxed); + std::mem::take(&mut shard.queue) + }; + for task in tasks { + metrics.dec_queue_depth(); + task.shutdown_or_run_if_mandatory(); + } + } + } + + /// Push a task and either notify an idle worker or invoke + /// `on_no_idle` (which is responsible for spawning a new worker if + /// possible). + pub(super) fn spawn_task( + &self, + task: Task, + metrics: &SpawnerMetrics, + on_no_idle: F, + ) -> Result<(), SpawnError> + where + F: FnOnce(&mut ThreadManagementState) -> Result<(), SpawnError>, + { + if self.is_shutdown.load(Ordering::Acquire) { + // It's fine to shutdown this task (even if mandatory): it was + // scheduled after the shutdown of the runtime began. + task.shutdown(); + return Err(SpawnError::ShuttingDown); + } + + // Push before taking `coord`, so spawners don't serialize on a + // pool-wide lock held across the queue operation. + if let Err(task) = self.push(task, metrics) { + // The shard was already drained and sealed for shutdown: reject + // the task, exactly as if the shutdown check above had caught it. + task.shutdown(); + return Err(SpawnError::ShuttingDown); + } + + let mut coord = self.coord.lock(); + + if coord.thread_mgmt_state.shutdown { + // Shutdown raced with our push, but the push beat the seal, so + // whichever worker (or `begin_shutdown`) seals that shard is + // guaranteed to collect the task and run it (if mandatory) or + // shut it down. Nothing to do here. + return Ok(()); + } + + if metrics.num_idle_threads() == 0 { + on_no_idle(&mut coord.thread_mgmt_state)?; + } else { + // Claim an idle worker (see `num_notify`). Signal after + // releasing `coord` so the woken worker doesn't immediately + // block on it; the counter increment, made under `coord`, is + // what guarantees the wakeup cannot be lost. + metrics.dec_num_idle_threads(); + coord.num_notify += 1; + drop(coord); + self.condvar.notify_one(); + } + + Ok(()) + } + + /// Run a worker thread's main loop. + pub(super) fn run_worker( + &self, + metrics: &SpawnerMetrics, + keep_alive: Duration, + worker_thread_id: usize, + ) -> Option> { + let mut join_on_thread = None; + let mut coord; + + 'main: loop { + // BUSY: run tasks without holding `coord`, so that spawners and + // other workers are not blocked on this worker. + while let Some(task) = self.pop(worker_thread_id) { + metrics.dec_queue_depth(); + task.run(); + } + + coord = self.coord.lock(); + + // Re-check the shards under `coord` before going idle: a task + // may have been pushed after the scan above, its spawner seeing + // this worker as busy and so neither notifying nor spawning. + // `coord` orders this re-check against every claim-or-spawn + // decision: a spawner that decided first pushed (and set the + // mask bit) before its `coord` critical section, so the task is + // visible here; one that decides later sees this worker counted + // idle and claims it. + if let Some(task) = self.pop(worker_thread_id) { + metrics.dec_queue_depth(); + drop(coord); + task.run(); + continue 'main; + } + + // IDLE + metrics.inc_num_idle_threads(); + + while !coord.thread_mgmt_state.shutdown { + let lock_result = self.condvar.wait_timeout(coord, keep_alive).unwrap(); + + coord = lock_result.0; + let timeout_result = lock_result.1; + + if coord.num_notify != 0 { + // A legitimate wakeup; the spawner already decremented + // `num_idle_threads` on this worker's behalf. + coord.num_notify -= 1; + drop(coord); + continue 'main; + } + + // Even if the condvar "timed out", if the pool is + // entering the shutdown phase, we want to perform + // the cleanup logic. + if !coord.thread_mgmt_state.shutdown && timeout_result.timed_out() { + join_on_thread = coord.thread_mgmt_state.worker_timed_out(worker_thread_id); + + break 'main; + } + + // Spurious wakeup detected, go back to sleep. + } + + // The pool is shutting down: drain and seal the shards, so that + // a spawner racing with shutdown either gets its task collected + // here or gets its push rejected (see `spawn_task`). + drop(coord); + self.drain_and_seal(metrics, worker_thread_id); + + coord = self.coord.lock(); + break 'main; + } + + // Thread exit + metrics.dec_num_threads(); + + // Unlike `LockedImpl`, this worker is always counted in + // `num_idle_threads` here: both loop exits (timeout and shutdown) + // are reached after the IDLE transition without a spawner having + // claimed this worker. + let prev_idle = metrics.dec_num_idle_threads(); + assert_ne!( + prev_idle, 0, + "`num_idle_threads` underflowed on thread exit" + ); + + if coord.thread_mgmt_state.shutdown && metrics.num_threads() == 0 { + self.condvar.notify_one(); + } + + drop(coord); + + join_on_thread + } + + /// Begin pool shutdown: set the shutdown flag, drop the shutdown + /// sender, wake all waiting workers, and hand back the worker + /// `JoinHandle`s for the caller to join. + pub(super) fn begin_shutdown(&self, metrics: &SpawnerMetrics) -> Option { + let mut coord = self.coord.lock(); + let handles = coord.thread_mgmt_state.begin_shutdown()?; + self.is_shutdown.store(true, Ordering::Release); + self.condvar.notify_all(); + + // Every live worker seals the shards on its way out, but if there + // are no workers (and none can be spawned now that `shutdown` is + // set), seal here so a racing spawner's push can't be stranded. + let no_workers = metrics.num_threads() == 0; + drop(coord); + if no_workers { + self.drain_and_seal(metrics, 0); + } + + Some(handles) + } +} diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index 2966a91b4..1143a5b83 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -155,6 +155,9 @@ pub struct Builder { /// Whether or not to enable eager hand-off for the I/O and time drivers (in /// `tokio_unstable`). enable_eager_driver_handoff: bool, + + /// When true, the blocking pool uses the sharded queue implementation. + pub(super) sharded_blocking_queue: bool, } cfg_unstable! { @@ -243,6 +246,16 @@ cfg_unstable! { pub(crate) type ThreadNameFn = std::sync::Arc String + Send + Sync + 'static>; +/// The default for the `sharded_blocking_queue` option: enabled iff the +/// `TOKIO_UNSTABLE_SHARDED_BLOCKING_QUEUE` environment variable is set to a +/// value other than `0`. +fn sharded_blocking_queue_default() -> bool { + match std::env::var_os("TOKIO_UNSTABLE_SHARDED_BLOCKING_QUEUE") { + Some(value) => !value.is_empty() && value != "0", + None => false, + } +} + #[derive(Clone, Copy)] pub(crate) enum Kind { CurrentThread, @@ -356,6 +369,8 @@ impl Builder { // Eager driver handoff is disabled by default. enable_eager_driver_handoff: false, + + sharded_blocking_queue: sharded_blocking_queue_default(), } } @@ -470,6 +485,47 @@ impl Builder { self } + /// Enables the sharded `spawn_blocking` queue, which is disabled by + /// default. + /// + /// By default, the blocking pool's task queue is protected by a single + /// mutex, which can become a point of contention when many threads spawn + /// blocking tasks concurrently. When this option is enabled, tasks are + /// instead distributed across several independently-locked queue shards. + /// + /// The sharded queue can also be enabled by setting the + /// `TOKIO_UNSTABLE_SHARDED_BLOCKING_QUEUE` environment variable to any + /// value other than `0`. + /// + /// [Click here to share your experience with the sharded queue](https://github.com/tokio-rs/tokio/issues/8067) + /// + /// **Note**: This is an [unstable API][unstable]. The sharded + /// `spawn_blocking` queue is an experimental feature that may be removed + /// or become the default behavior in 1.x releases. See + /// [the documentation on unstable features][unstable] for details. + /// + /// # Examples + /// + /// ``` + /// # #[cfg(not(target_family = "wasm"))] + /// # { + /// use tokio::runtime; + /// + /// let rt = runtime::Builder::new_multi_thread() + /// .enable_sharded_blocking_queue() + /// .build() + /// .unwrap(); + /// # } + /// ``` + /// + /// [unstable]: crate#unstable-features + #[cfg(tokio_unstable)] + #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] + pub fn enable_sharded_blocking_queue(&mut self) -> &mut Self { + self.sharded_blocking_queue = true; + self + } + /// Sets the number of worker threads the `Runtime` will use. /// /// This can be any number above 0 though it is advised to keep this value diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index d9863fe85..22c3e3063 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -127,7 +127,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/runtime/tests/loom_blocking.rs b/tokio/src/runtime/tests/loom_blocking.rs index ce1f26441..cf1021ed5 100644 --- a/tokio/src/runtime/tests/loom_blocking.rs +++ b/tokio/src/runtime/tests/loom_blocking.rs @@ -132,6 +132,56 @@ fn spawn_blocking_then_shutdown() { }); } +/// Regression-style test for the class of bug behind +/// : a `spawn_blocking` while +/// the pool is at its thread cap must not be stranded when it races with the +/// only worker transitioning between busy and idle. +#[test] +fn spawn_blocking_at_thread_cap_runs() { + loom::model(|| { + let rt = crate::runtime::Builder::new_current_thread() + .max_blocking_threads(1) + .thread_keep_alive(Duration::from_secs(7200)) // don't let the thread exit on its own + .build() + .unwrap(); + let rt_hdl = rt.handle().clone(); + + // Spawn a worker thread and wait for its task to finish, so the + // worker is somewhere between running a task and parking idle. + let jh0 = rt_hdl.spawn_blocking(|| {}); + loom::future::block_on(jh0).unwrap(); + + // The pool is now at its thread cap, so this task can only run if + // the existing worker picks it up; if the spawn is lost, this + // deadlocks. + let jh1 = rt_hdl.spawn_blocking(|| {}); + loom::future::block_on(jh1).unwrap(); + + drop(rt); + }); +} + +/// A `spawn_blocking` racing runtime shutdown must never strand the task: +/// its `JoinHandle` must resolve (the task ran or was cancelled) no matter +/// how the spawn interleaves with shutdown's drain. +#[test] +fn spawn_blocking_racing_shutdown_resolves() { + loom::model(|| { + let rt = crate::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let handle = rt.handle().clone(); + + let spawner = loom::thread::spawn(move || handle.spawn_blocking(|| {})); + + drop(rt); + + let jh = spawner.join().unwrap(); + // This deadlocks (which loom detects) if the task was lost. + let _ = loom::future::block_on(jh); + }); +} + fn mk_runtime(num_threads: usize) -> Runtime { runtime::Builder::new_multi_thread() .worker_threads(num_threads) diff --git a/tokio/src/util/rand.rs b/tokio/src/util/rand.rs index efd235154..3e71e330b 100644 --- a/tokio/src/util/rand.rs +++ b/tokio/src/util/rand.rs @@ -67,11 +67,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/