tokio 1.54.1

//! Thread pool for blocking operations

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};
use crate::runtime::{AutoBox, Builder, Callback, Handle};
use crate::util::metric_atomics::MetricAtomicUsize;
use crate::util::trace::{blocking_task, SpawnMeta};

use std::collections::{HashMap, VecDeque};
use std::fmt;
use std::io;
use std::sync::atomic::Ordering;
use std::time::Duration;

pub(crate) struct BlockingPool {
    spawner: Spawner,
    shutdown_rx: shutdown::Receiver,
}

#[derive(Clone)]
pub(crate) struct Spawner {
    inner: Arc<Inner>,
}

#[derive(Default)]
pub(crate) struct SpawnerMetrics {
    num_threads: MetricAtomicUsize,
    num_idle_threads: MetricAtomicUsize,
    queue_depth: MetricAtomicUsize,
}

impl SpawnerMetrics {
    pub(super) fn num_threads(&self) -> usize {
        self.num_threads.load(Ordering::Relaxed)
    }

    pub(super) fn num_idle_threads(&self) -> usize {
        self.num_idle_threads.load(Ordering::Relaxed)
    }

    cfg_unstable_metrics! {
        fn queue_depth(&self) -> usize {
            self.queue_depth.load(Ordering::Relaxed)
        }
    }

    fn inc_num_threads(&self) {
        self.num_threads.increment();
    }

    pub(super) fn dec_num_threads(&self) {
        self.num_threads.decrement();
    }

    pub(super) fn inc_num_idle_threads(&self) {
        self.num_idle_threads.increment();
    }

    pub(super) fn dec_num_idle_threads(&self) -> usize {
        self.num_idle_threads.decrement()
    }

    pub(super) fn inc_queue_depth(&self) {
        self.queue_depth.increment();
    }

    pub(super) fn dec_queue_depth(&self) {
        self.queue_depth.decrement();
    }
}

struct Inner {
    /// Queue + notification implementation.
    inner_impl: InnerImpl,

    /// Spawned threads use this name.
    thread_name: ThreadNameFn,

    /// Spawned thread stack size.
    stack_size: Option<usize>,

    /// Call after a thread starts.
    after_start: Option<Callback>,

    /// Call before a thread stops.
    before_stop: Option<Callback>,

    // Maximum number of threads.
    thread_cap: usize,

    // Number of runtime scheduler workers counted in `num_threads`.
    scheduler_threads: usize,

    // Customizable wait timeout.
    keep_alive: Duration,

    // Metrics about the pool.
    metrics: SpawnerMetrics,
}

/// Per-variant queue + notification + lock topology.
enum InnerImpl {
    Locked(LockedImpl),
    Sharded(ShardedImpl),
}

/// Single-mutex + condvar implementation.
struct LockedImpl {
    mutex: Mutex<LockedInner>,
    condvar: Condvar,
}

struct LockedInner {
    queue: VecDeque<Task>,
    num_notify: u32,
    /// Thread-management state. Split apart so that it's re-usable in a
    /// sharded queue implementation.
    thread_mgmt_state: ThreadManagementState,
}

/// 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.
pub(super) type ShutdownHandles = (
    Option<thread::JoinHandle<()>>,
    HashMap<usize, thread::JoinHandle<()>>,
);

/// Thread-management state used by every `InnerImpl` variant.
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.
    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.
    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).
    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 {
    task: task::UnownedTask<BlockingSchedule>,
    mandatory: Mandatory,
}

#[derive(PartialEq, Eq)]
pub(crate) enum Mandatory {
    #[cfg_attr(any(not(feature = "fs"), target_os = "emscripten"), allow(dead_code))]
    Mandatory,
    NonMandatory,
}

pub(crate) enum SpawnError {
    /// Pool is shutting down and the task was not scheduled
    ShuttingDown,
    /// There are no worker threads available to take the task
    /// and the OS failed to spawn a new one
    NoThreads(io::Error),
}

impl From<SpawnError> for io::Error {
    fn from(e: SpawnError) -> Self {
        match e {
            SpawnError::ShuttingDown => io::Error::other("blocking pool shutting down"),
            SpawnError::NoThreads(e) => e,
        }
    }
}

impl Task {
    pub(crate) fn new(task: task::UnownedTask<BlockingSchedule>, mandatory: Mandatory) -> Task {
        Task { task, mandatory }
    }

    pub(super) fn shutdown(self) {
        self.task.shutdown();
    }

    pub(super) fn run(self) {
        self.task.run();
    }

    pub(super) fn shutdown_or_run_if_mandatory(self) {
        match self.mandatory {
            Mandatory::NonMandatory => self.task.shutdown(),
            Mandatory::Mandatory => self.task.run(),
        }
    }
}

const KEEP_ALIVE: Duration = Duration::from_secs(10);

/// Runs the provided function on an executor dedicated to blocking operations.
/// Tasks will be scheduled as non-mandatory, meaning they may not get executed
/// in case of runtime shutdown.
#[track_caller]
#[cfg_attr(target_os = "wasi", allow(dead_code))]
pub(crate) fn spawn_blocking<F, R>(func: F) -> JoinHandle<R>
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    let rt = Handle::current();
    rt.spawn_blocking(func)
}

cfg_io_blocking! {
    #[cfg_attr(any(
        all(loom, not(test)), // the function is covered by loom tests
        test,
        target_os = "emscripten", // fs and io-std use the inline shim
    ), allow(dead_code))]
    /// Runs the provided function on an executor dedicated to blocking
    /// operations. Tasks will be scheduled as mandatory, meaning they are
    /// guaranteed to run unless a shutdown is already taking place. In case a
    /// shutdown is already taking place, `None` will be returned.
    pub(crate) fn spawn_mandatory_blocking<F, R>(func: F) -> Option<JoinHandle<R>>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let rt = Handle::current();
        rt.inner.blocking_spawner().spawn_mandatory_blocking(&rt, func)
    }
}

// ===== impl BlockingPool =====

impl BlockingPool {
    pub(crate) fn new(
        builder: &Builder,
        thread_cap: usize,
        scheduler_threads: usize,
    ) -> 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,
                    thread_name: builder.thread_name.clone(),
                    stack_size: builder.thread_stack_size,
                    after_start: builder.after_start.clone(),
                    before_stop: builder.before_stop.clone(),
                    thread_cap,
                    scheduler_threads,
                    keep_alive,
                    metrics: SpawnerMetrics::default(),
                }),
            },
            shutdown_rx,
        }
    }

    pub(crate) fn spawner(&self) -> &Spawner {
        &self.spawner
    }

    pub(crate) fn shutdown(&mut self, timeout: Option<Duration>) {
        // 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 inner = &self.spawner.inner;
        let (last_exited_thread, workers) = match inner.inner_impl.begin_shutdown(&inner.metrics) {
            Some(x) => x,
            None => return,
        };

        if self.shutdown_rx.wait(timeout) {
            let _ = last_exited_thread.map(thread::JoinHandle::join);

            // Loom requires that execution be deterministic, so sort by thread ID before joining.
            // (HashMaps use a randomly-seeded hash function, so the order is nondeterministic)
            #[cfg(loom)]
            let workers: Vec<(usize, thread::JoinHandle<()>)> = {
                let mut workers: Vec<_> = workers.into_iter().collect();
                workers.sort_by_key(|(id, _)| *id);
                workers
            };

            for (_id, handle) in workers {
                let _ = handle.join();
            }
        }
    }
}

impl Drop for BlockingPool {
    fn drop(&mut self) {
        self.shutdown(None);
    }
}

impl fmt::Debug for BlockingPool {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("BlockingPool").finish()
    }
}

// ===== impl Spawner =====

impl Spawner {
    #[track_caller]
    pub(crate) fn spawn_blocking<F, R>(&self, rt: &Handle, func: F) -> JoinHandle<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let fn_size = std::mem::size_of::<F>();
        let (join_handle, spawn_result) = if AutoBox::<F>::SHOULD_BOX {
            self.spawn_blocking_inner(
                Box::new(func),
                Mandatory::NonMandatory,
                SpawnMeta::new_unnamed(fn_size),
                rt,
            )
        } else {
            self.spawn_blocking_inner(
                func,
                Mandatory::NonMandatory,
                SpawnMeta::new_unnamed(fn_size),
                rt,
            )
        };

        match spawn_result {
            Ok(()) => join_handle,
            // Compat: do not panic here, return the join_handle even though it will never resolve
            Err(SpawnError::ShuttingDown) => join_handle,
            Err(SpawnError::NoThreads(e)) => {
                panic!("OS can't spawn worker thread: {e}")
            }
        }
    }

    cfg_io_blocking! {
        #[track_caller]
        #[cfg_attr(any(
            all(loom, not(test)), // the function is covered by loom tests
            test,
            target_os = "emscripten", // fs and io-std use the inline shim
        ), allow(dead_code))]
        pub(crate) fn spawn_mandatory_blocking<F, R>(&self, rt: &Handle, func: F) -> Option<JoinHandle<R>>
        where
            F: FnOnce() -> R + Send + 'static,
            R: Send + 'static,
        {
            let fn_size = std::mem::size_of::<F>();
            let (join_handle, spawn_result) = if AutoBox::<F>::SHOULD_BOX {
                self.spawn_blocking_inner(
                    Box::new(func),
                    Mandatory::Mandatory,
                    SpawnMeta::new_unnamed(fn_size),
                    rt,
                )
            } else {
                self.spawn_blocking_inner(
                    func,
                    Mandatory::Mandatory,
                    SpawnMeta::new_unnamed(fn_size),
                    rt,
                )
            };

            if spawn_result.is_ok() {
                Some(join_handle)
            } else {
                None
            }
        }
    }

    #[track_caller]
    pub(crate) fn spawn_blocking_inner<F, R>(
        &self,
        func: F,
        is_mandatory: Mandatory,
        spawn_meta: SpawnMeta<'_>,
        rt: &Handle,
    ) -> (JoinHandle<R>, Result<(), SpawnError>)
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let id = task::Id::next();
        let fut =
            blocking_task::<F, BlockingTask<F>>(BlockingTask::new(func), spawn_meta, id.as_u64());

        let (task, handle) = task::unowned(
            fut,
            BlockingSchedule::new(rt),
            id,
            task::SpawnLocation::capture(),
        );

        let spawned = self.spawn_task(Task::new(task, is_mandatory), rt);
        (handle, spawned)
    }

    pub(crate) fn num_blocking_threads(&self) -> usize {
        self.inner
            .metrics
            .num_threads()
            .saturating_sub(self.inner.scheduler_threads)
    }

    fn spawn_task(&self, task: Task, rt: &Handle) -> Result<(), SpawnError> {
        // The `on_no_idle` closure runs under the same lock as the queue
        // push, exactly like the pre-refactor code that called
        // `self.spawn_thread` directly while holding `Mutex<ThreadManagementState>`.
        self.inner
            .inner_impl
            .spawn_task(task, &self.inner.metrics, |thread_mgmt_state| {
                // No threads are able to process the task.

                if self.inner.metrics.num_threads() == self.inner.thread_cap {
                    // At max number of threads
                    return Ok(());
                }

                assert!(thread_mgmt_state.shutdown_tx.is_some());
                let shutdown_tx = thread_mgmt_state.shutdown_tx.clone();

                if let Some(shutdown_tx) = shutdown_tx {
                    let id = thread_mgmt_state.worker_thread_index;

                    match self.spawn_thread(shutdown_tx, rt, id) {
                        Ok(handle) => {
                            self.inner.metrics.inc_num_threads();
                            thread_mgmt_state.worker_thread_index += 1;
                            thread_mgmt_state.worker_threads.insert(id, handle);
                        }
                        Err(ref e)
                            if is_temporary_os_thread_error(e)
                                && self.num_blocking_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));
                        }
                    }
                }

                Ok(())
            })
    }

    fn spawn_thread(
        &self,
        shutdown_tx: shutdown::Sender,
        rt: &Handle,
        id: usize,
    ) -> io::Result<thread::JoinHandle<()>> {
        let mut builder = thread::Builder::new().name((self.inner.thread_name)());

        if let Some(stack_size) = self.inner.stack_size {
            builder = builder.stack_size(stack_size);
        }

        let rt = rt.clone();

        builder.spawn(move || {
            // Only the reference should be moved into the closure
            let _enter = rt.enter();
            rt.inner.blocking_spawner().inner.run(id);
            drop(shutdown_tx);
        })
    }
}

cfg_unstable_metrics! {
    impl Spawner {
        pub(crate) fn num_idle_threads(&self) -> usize {
            self.inner.metrics.num_idle_threads()
        }

        pub(crate) fn queue_depth(&self) -> usize {
            self.inner.metrics.queue_depth()
        }
    }
}

// Each method on `InnerImpl` dispatches to the matching method on the
// concrete variant. The variant methods own the entire critical section
// for their operation, which makes it self-evident that the `Locked`
// variant's behavior is identical to the pre-refactor code and gives a
// future `Concurrent` variant a symmetric slot to fill.
impl InnerImpl {
    fn spawn_task<F>(
        &self,
        task: Task,
        metrics: &SpawnerMetrics,
        on_no_idle: F,
    ) -> Result<(), SpawnError>
    where
        F: FnOnce(&mut ThreadManagementState) -> Result<(), SpawnError>,
    {
        match self {
            InnerImpl::Locked(l) => l.spawn_task(task, metrics, on_no_idle),
            InnerImpl::Sharded(s) => s.spawn_task(task, metrics, on_no_idle),
        }
    }

    fn run_worker(
        &self,
        metrics: &SpawnerMetrics,
        keep_alive: Duration,
        worker_thread_id: usize,
    ) -> 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, metrics: &SpawnerMetrics) -> Option<ShutdownHandles> {
        match self {
            InnerImpl::Locked(l) => l.begin_shutdown(),
            InnerImpl::Sharded(s) => s.begin_shutdown(metrics),
        }
    }
}

// This is the original, single-lock implementation of the `spawn_blocking`
// queue. Method scope was principally designed around ensuring that when the
// 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).
    fn spawn_task<F>(
        &self,
        task: Task,
        metrics: &SpawnerMetrics,
        on_no_idle: F,
    ) -> Result<(), SpawnError>
    where
        F: FnOnce(&mut ThreadManagementState) -> Result<(), SpawnError>,
    {
        let mut locked = self.mutex.lock();

        if locked.thread_mgmt_state.shutdown {
            drop(locked);

            // Shutdown the task: it's fine to shutdown this task
            // (even if mandatory) because it was scheduled after the
            // shutdown of the runtime began.
            task.task.shutdown();
            return Err(SpawnError::ShuttingDown);
        }

        locked.queue.push_back(task);
        metrics.inc_queue_depth();

        if metrics.num_idle_threads() == 0 {
            on_no_idle(&mut locked.thread_mgmt_state)?;
        } 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.
            metrics.dec_num_idle_threads();
            locked.num_notify += 1;
            self.condvar.notify_one();
        }

        Ok(())
    }

    /// Run a worker thread's main loop.
    fn run_worker(
        &self,
        metrics: &SpawnerMetrics,
        keep_alive: Duration,
        worker_thread_id: usize,
    ) -> Option<thread::JoinHandle<()>> {
        let mut locked = self.mutex.lock();
        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) = locked.queue.pop_front() {
                metrics.dec_queue_depth();
                drop(locked);
                task.run();

                locked = self.mutex.lock();
            }

            // IDLE
            metrics.inc_num_idle_threads();
            // mark this thread as currently counted in `num_idle_threads`.
            is_counted_idle = true;

            while !locked.thread_mgmt_state.shutdown {
                let lock_result = self.condvar.wait_timeout(locked, keep_alive).unwrap();

                locked = lock_result.0;
                let timeout_result = lock_result.1;

                if locked.num_notify != 0 {
                    // We have received a legitimate wakeup,
                    // acknowledge it by decrementing the counter
                    // and transition to the BUSY state.
                    locked.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 !locked.thread_mgmt_state.shutdown && timeout_result.timed_out() {
                    join_on_thread = locked.thread_mgmt_state.worker_timed_out(worker_thread_id);

                    break 'main;
                }

                // Spurious wakeup detected, go back to sleep.
            }

            if locked.thread_mgmt_state.shutdown {
                // Drain the queue
                while let Some(task) = locked.queue.pop_front() {
                    metrics.dec_queue_depth();
                    drop(locked);

                    task.shutdown_or_run_if_mandatory();

                    locked = self.mutex.lock();
                }

                break;
            }
        }

        // Thread exit
        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 = metrics.dec_num_idle_threads();
            assert_ne!(
                prev_idle, 0,
                "`num_idle_threads` underflowed on thread exit"
            );
        }

        if locked.thread_mgmt_state.shutdown && metrics.num_threads() == 0 {
            self.condvar.notify_one();
        }

        drop(locked);

        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.
    fn begin_shutdown(&self) -> Option<ShutdownHandles> {
        let mut locked = self.mutex.lock();
        let handles = locked.thread_mgmt_state.begin_shutdown()?;
        self.condvar.notify_all();
        Some(handles)
    }
}

// Tells whether the error when spawning a thread is temporary.
#[inline]
fn is_temporary_os_thread_error(error: &io::Error) -> bool {
    matches!(error.kind(), io::ErrorKind::WouldBlock)
}

impl Inner {
    fn run(&self, worker_thread_id: usize) {
        if let Some(f) = &self.after_start {
            f();
        }

        let join_on_thread =
            self.inner_impl
                .run_worker(&self.metrics, self.keep_alive, worker_thread_id);

        if let Some(f) = &self.before_stop {
            f();
        }

        if let Some(handle) = join_on_thread {
            let _ = handle.join();
        }
    }
}

impl fmt::Debug for Spawner {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("blocking::Spawner").finish()
    }
}