qubit_thread_pool/dynamic/

thread_pool.rs

/*******************************************************************************
 *
 *    Copyright (c) 2025 - 2026 Haixing Hu.
 *
 *    SPDX-License-Identifier: Apache-2.0
 *
 *    Licensed under the Apache License, Version 2.0.
 *
 ******************************************************************************/
use std::{
    sync::Arc,
    time::Duration,
};

use qubit_function::{
    Callable,
    Runnable,
};

use qubit_executor::task::spi::TaskEndpointPair;
use qubit_executor::{
    TaskHandle,
    TrackedTask,
};

use super::thread_pool_builder::ThreadPoolBuilder;
use super::thread_pool_inner::ThreadPoolInner;
use crate::{
    ExecutorServiceBuilderError,
    PoolJob,
    ThreadPoolStats,
};
use qubit_executor::service::{
    ExecutorService,
    ExecutorServiceLifecycle,
    StopReport,
    SubmissionError,
};

/// OS thread pool implementing [`ExecutorService`].
///
/// `ThreadPool` accepts fallible tasks, stores accepted waiting tasks in
/// internal queues, and executes them on worker threads. The global waiting
/// queue is FIFO, but task start and completion order are not a strict FIFO API
/// guarantee. Workers are created lazily up to the configured core size, tasks
/// are queued after that, and the pool may grow up to the maximum size when a
/// bounded queue is full. Callable submissions return [`TaskHandle`], while
/// tracked submissions return [`TrackedTask`].
///
/// `shutdown` is graceful: already accepted queued tasks are allowed to run.
/// `stop` is abrupt: queued tasks that have not started are completed
/// with [`TaskExecutionError::Cancelled`](qubit_executor::TaskExecutionError::Cancelled).
///
pub struct ThreadPool {
    /// Shared pool state and worker coordination primitives.
    inner: Arc<ThreadPoolInner>,
}

impl ThreadPool {
    pub(super) fn from_inner(inner: Arc<ThreadPoolInner>) -> Self {
        Self { inner }
    }

    /// Creates a thread pool with equal core and maximum pool sizes.
    ///
    /// # Parameters
    ///
    /// * `pool_size` - Value applied as both core and maximum pool size.
    ///
    /// # Returns
    ///
    /// `Ok(ThreadPool)` if all workers are spawned successfully.
    ///
    /// # Errors
    ///
    /// Returns [`ExecutorServiceBuilderError`] if the resulting maximum pool size is
    /// zero or a worker thread cannot be spawned.
    #[inline]
    pub fn new(pool_size: usize) -> Result<Self, ExecutorServiceBuilderError> {
        Self::builder().pool_size(pool_size).build()
    }

    /// Creates a builder for configuring a thread pool.
    ///
    /// # Returns
    ///
    /// A builder with default core and maximum pool sizes and an unbounded queue.
    #[inline]
    pub fn builder() -> ThreadPoolBuilder {
        ThreadPoolBuilder::default()
    }

    /// Returns the number of queued tasks waiting for a worker.
    ///
    /// # Returns
    ///
    /// The number of accepted tasks that have not started yet.
    #[inline]
    pub fn queued_count(&self) -> usize {
        self.inner.queued_count()
    }

    /// Returns the number of tasks currently held by workers.
    ///
    /// # Returns
    ///
    /// The number of tasks that workers have taken from the queue and have not
    /// yet finished processing.
    #[inline]
    pub fn running_count(&self) -> usize {
        self.inner.running_count()
    }

    /// Returns how many worker threads are still running in this pool.
    ///
    /// # Returns
    ///
    /// The number of live worker loops still owned by this pool. This is a
    /// runtime count and is not required to match configured
    /// [`Self::core_pool_size`] or [`Self::maximum_pool_size`].
    #[inline]
    pub fn live_worker_count(&self) -> usize {
        self.inner.read_state(|state| state.live_workers)
    }

    /// Returns the configured core pool size.
    ///
    /// # Returns
    ///
    /// The number of workers kept for normal load before tasks are queued.
    #[inline]
    pub fn core_pool_size(&self) -> usize {
        self.inner.read_state(|state| state.core_pool_size)
    }

    /// Returns the configured maximum pool size.
    ///
    /// # Returns
    ///
    /// The maximum number of worker threads this pool may create.
    #[inline]
    pub fn maximum_pool_size(&self) -> usize {
        self.inner.read_state(|state| state.maximum_pool_size)
    }

    /// Returns a point-in-time snapshot of pool counters.
    ///
    /// # Returns
    ///
    /// A snapshot containing worker, queue, and task counters observed under
    /// the pool state lock.
    #[inline]
    pub fn stats(&self) -> ThreadPoolStats {
        self.inner.stats()
    }

    /// Submits a custom pool job.
    ///
    /// This low-level extension point is intended for higher-level services
    /// that need to pair pool execution with their own task registry or
    /// cancellation bookkeeping. Prefer the [`ExecutorService`] submission
    /// methods for ordinary runnable and callable work. Custom job callbacks
    /// run synchronously on the thread that reaches the corresponding lifecycle
    /// event and should stay short and non-blocking. Callback panics are
    /// contained; if an acceptance callback panics, the job is not queued,
    /// run, or cancelled.
    ///
    /// # Parameters
    ///
    /// * `job` - Custom job to execute or cancel later.
    ///
    /// # Returns
    ///
    /// `Ok(())` when the pool accepts the job.
    ///
    /// # Errors
    ///
    /// Returns [`SubmissionError::Shutdown`] after shutdown, returns
    /// [`SubmissionError::Saturated`] when the bounded pool cannot accept
    /// more work, or returns [`SubmissionError::WorkerSpawnFailed`] when a
    /// required worker cannot be created.
    #[inline]
    pub fn submit_job(&self, job: PoolJob) -> Result<(), SubmissionError> {
        self.inner.submit(job)
    }

    /// Blocks until all accepted work has completed.
    ///
    /// This is a join-style wait for quiescence: it does not request shutdown
    /// and does not wait for worker threads to exit. Concurrent submissions may
    /// extend the wait until those accepted jobs also drain.
    #[inline]
    pub fn join(&self) {
        self.inner.wait_until_idle();
    }

    /// Starts one core worker if the pool has fewer live workers than its
    /// configured core size.
    ///
    /// # Returns
    ///
    /// `Ok(true)` if a worker was started, or `Ok(false)` if no core worker
    /// was needed.
    ///
    /// # Errors
    ///
    /// Returns [`SubmissionError::Shutdown`] if the pool is shut down, or
    /// [`SubmissionError::WorkerSpawnFailed`] if worker creation fails.
    #[inline]
    pub fn prestart_core_thread(&self) -> Result<bool, SubmissionError> {
        self.inner.prestart_core_thread()
    }

    /// Starts all missing core workers.
    ///
    /// # Returns
    ///
    /// The number of workers started.
    ///
    /// # Errors
    ///
    /// Returns [`SubmissionError::Shutdown`] if the pool is shut down, or
    /// [`SubmissionError::WorkerSpawnFailed`] if worker creation fails.
    #[inline]
    pub fn prestart_all_core_threads(&self) -> Result<usize, SubmissionError> {
        self.inner.prestart_all_core_threads()
    }

    /// Updates the core pool size.
    ///
    /// Increasing the core size changes future admission and prestart limits,
    /// but it does not eagerly create workers or reschedule already queued work.
    /// Call [`Self::prestart_all_core_threads`] when eager creation is desired.
    /// Decreasing the core size lets excess idle workers retire according to
    /// the keep-alive policy.
    ///
    /// # Parameters
    ///
    /// * `core_pool_size` - New core pool size.
    ///
    /// # Returns
    ///
    /// `Ok(())` if the size is accepted.
    ///
    /// # Errors
    ///
    /// Returns [`ExecutorServiceBuilderError::CorePoolSizeExceedsMaximum`] when the
    /// new core size would exceed the current maximum size.
    pub fn set_core_pool_size(&self, core_pool_size: usize) -> Result<(), ExecutorServiceBuilderError> {
        self.inner.set_core_pool_size(core_pool_size)
    }

    /// Updates the maximum pool size.
    ///
    /// Excess workers are not interrupted. They retire after finishing current
    /// work or timing out while idle.
    ///
    /// # Parameters
    ///
    /// * `maximum_pool_size` - New maximum pool size.
    ///
    /// # Returns
    ///
    /// `Ok(())` if the size is accepted.
    ///
    /// # Errors
    ///
    /// Returns [`ExecutorServiceBuilderError::ZeroMaximumPoolSize`] when the maximum
    /// size is zero, or [`ExecutorServiceBuilderError::CorePoolSizeExceedsMaximum`]
    /// when it would be smaller than the current core size.
    pub fn set_maximum_pool_size(&self, maximum_pool_size: usize) -> Result<(), ExecutorServiceBuilderError> {
        self.inner.set_maximum_pool_size(maximum_pool_size)
    }

    /// Updates how long excess idle workers may wait before exiting.
    ///
    /// # Parameters
    ///
    /// * `keep_alive` - New idle timeout for workers above the core size.
    ///
    /// # Returns
    ///
    /// `Ok(())` if the timeout is accepted.
    ///
    /// # Errors
    ///
    /// Returns [`ExecutorServiceBuilderError::ZeroKeepAlive`] when `keep_alive` is
    /// zero.
    pub fn set_keep_alive(&self, keep_alive: Duration) -> Result<(), ExecutorServiceBuilderError> {
        self.inner.set_keep_alive(keep_alive)
    }

    /// Updates whether core workers may also retire after keep-alive timeout.
    ///
    /// # Parameters
    ///
    /// * `allow` - Whether core workers are subject to idle timeout.
    pub fn allow_core_thread_timeout(&self, allow: bool) {
        self.inner.allow_core_thread_timeout(allow);
    }
}

impl Drop for ThreadPool {
    /// Requests graceful shutdown when the pool value is dropped.
    fn drop(&mut self) {
        self.inner.shutdown();
    }
}

impl ExecutorService for ThreadPool {
    type ResultHandle<R, E>
        = TaskHandle<R, E>
    where
        R: Send + 'static,
        E: Send + 'static;

    type TrackedHandle<R, E>
        = TrackedTask<R, E>
    where
        R: Send + 'static,
        E: Send + 'static;

    /// Accepts a runnable and queues it for pool workers.
    fn submit<T, E>(&self, task: T) -> Result<(), SubmissionError>
    where
        T: Runnable<E> + Send + 'static,
        E: Send + 'static,
    {
        self.inner.submit(PoolJob::detached(task))
    }

    /// Accepts a callable and queues it for pool workers.
    ///
    /// # Parameters
    ///
    /// * `task` - Callable to execute on a pool worker.
    ///
    /// # Returns
    ///
    /// A [`TaskHandle`] for the accepted task.
    ///
    /// # Errors
    ///
    /// Returns [`SubmissionError::Shutdown`] after shutdown, returns
    /// [`SubmissionError::Saturated`] when the bounded pool cannot accept
    /// more work, or returns [`SubmissionError::WorkerSpawnFailed`] when a
    /// required worker cannot be created.
    fn submit_callable<C, R, E>(&self, task: C) -> Result<Self::ResultHandle<R, E>, SubmissionError>
    where
        C: Callable<R, E> + Send + 'static,
        R: Send + 'static,
        E: Send + 'static,
    {
        let (handle, completion) = TaskEndpointPair::new().into_parts();
        let job = PoolJob::from_task(task, completion);
        self.inner.submit(job)?;
        Ok(handle)
    }

    /// Accepts a callable and queues it with a tracked handle.
    fn submit_tracked_callable<C, R, E>(&self, task: C) -> Result<Self::TrackedHandle<R, E>, SubmissionError>
    where
        C: Callable<R, E> + Send + 'static,
        R: Send + 'static,
        E: Send + 'static,
    {
        let (handle, completion) = TaskEndpointPair::new().into_tracked_parts();
        let job = PoolJob::from_task(task, completion);
        self.inner.submit(job)?;
        Ok(handle)
    }

    /// Stops accepting new tasks after already queued work is drained.
    ///
    /// Queued and running tasks remain eligible to complete normally.
    #[inline]
    fn shutdown(&self) {
        self.inner.shutdown();
    }

    /// Stops accepting tasks and cancels queued tasks that have not started.
    ///
    /// # Returns
    ///
    /// A report containing the number of queued jobs cancelled and the number
    /// of jobs running at the time of the request.
    #[inline]
    fn stop(&self) -> StopReport {
        self.inner.stop()
    }

    /// Returns the current lifecycle state.
    #[inline]
    fn lifecycle(&self) -> ExecutorServiceLifecycle {
        self.inner.lifecycle()
    }

    /// Returns whether shutdown has been requested.
    #[inline]
    fn is_not_running(&self) -> bool {
        self.inner.is_not_running()
    }

    /// Returns whether shutdown was requested and all workers have exited.
    #[inline]
    fn is_terminated(&self) -> bool {
        self.inner.is_terminated()
    }

    /// Blocks until the pool has terminated.
    fn wait_termination(&self) {
        self.inner.wait_for_termination();
    }
}