runtime: add an opt-in sharded spawn_blocking queue (#8337)

Re-lands the sharded queue from #7757 (reverted due to #8056), disabled by
default. Opt in via the unstable `Builder::enable_sharded_blocking_queue` or
the `TOKIO_UNSTABLE_SHARDED_BLOCKING_QUEUE` environment variable.

.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

.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

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

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;
 }

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<thread::JoinHandle<()>>,
     HashMap<usize, thread::JoinHandle<()>>,
 );
 
 /// Thread-management state used by every `InnerImpl` variant.
-struct ThreadManagementState {
-    shutdown: bool,
-    shutdown_tx: Option<shutdown::Sender>,
+pub(super) struct ThreadManagementState {
+    pub(super) shutdown: bool,
+    pub(super) shutdown_tx: Option<shutdown::Sender>,
     /// 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
     /// <https://github.com/tokio-rs/tokio/commit/646fbae76535e397ef79dbcaacb945d4c829f666>
     /// for more information.
-    last_exiting_thread: Option<thread::JoinHandle<()>>,
+    pub(super) last_exiting_thread: Option<thread::JoinHandle<()>>,
     /// This holds the `JoinHandles` for all running threads; on shutdown, the thread
     /// calling shutdown handles joining on these.
-    worker_threads: HashMap<usize, thread::JoinHandle<()>>,
+    pub(super) worker_threads: HashMap<usize, thread::JoinHandle<()>>,
     /// 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<ShutdownHandles> {
+        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<thread::JoinHandle<()>> {
+        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<thread::JoinHandle<()>> {
         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<ShutdownHandles> {
+    fn begin_shutdown(&self, metrics: &SpawnerMetrics) -> Option<ShutdownHandles> {
         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<ShutdownHandles> {
         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)
     }
 }
 

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<Task>,
+    /// 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<Shard>; 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<ShardedCoord>,
+    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<Task> {
+        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<F>(
+        &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<thread::JoinHandle<()>> {
+        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<ShutdownHandles> {
+        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)
+    }
+}

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<dyn Fn() -> 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

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);

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
+/// <https://github.com/tokio-rs/tokio/issues/8056>: 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)

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/