safina_threadpool/

lib.rs

//! # ARCHIVED ARCHIVED ARCHIVED
//! This crate is archived and will not be updated.
//!
//! The code is now at
//! [`safina::threadpool`](https://docs.rs/safina/latest/safina/threadpool/index.html) in the
//! [`safina`](https://crates.io/crates/safina) crate.
//!
//! ----
//!
//! # safina-threadpool
//!
//! A threadpool.
//!
//! You can use it alone or with [`safina`](https://crates.io/crates/safina),
//! a safe async runtime.
//!
//! # Features
//! - Add a closure or `FnOnce` to the pool and one of the threads will execute it
//! - Automatically restarts panicked threads
//! - Retries after failing to spawn a thread
//! - Drop the `ThreadPool` struct to stop all idle threads.
//! - Destroy the pool and wait for all threads to stop
//! - `forbid(unsafe_code)`
//! - Depends only on `std`
//! - 100% test coverage
//!
//! # Limitations
//! - Not optimized
//!
//! # Examples
//! ```rust
//! # type ProcessResult = ();
//! # fn process_data(data: (), sender: std::sync::mpsc::Sender<ProcessResult>) -> ProcessResult {
//! #    sender.send(()).unwrap();
//! # }
//! # fn f() {
//! # let data_source = vec![(),()];
//! let pool =
//!     safina_threadpool::ThreadPool::new("worker", 2).unwrap();
//! let receiver = {
//!     let (sender, receiver) =
//!         std::sync::mpsc::channel();
//!     for data in data_source {
//!         let sender_clone = sender.clone();
//!         pool.schedule(
//!             move || process_data(data, sender_clone));
//!     }
//!     receiver
//! };
//! let results: Vec<ProcessResult> =
//!     receiver.iter().collect();
//! // ...
//! # }
//! ```
//!
//! # Alternatives
//! - [`blocking`](https://crates.io/crates/blocking)
//!   - Popular
//!   - A little `unsafe` code
//!   - [blocking/issues/24: Recover from thread spawn failure](https://github.com/smol-rs/blocking/issues/24)
//! - [`threadpool`](https://crates.io/crates/threadpool)
//!   - Popular
//!   - Well maintained
//!   - Dependencies have `unsafe` code
//!   - [rust-threadpool/issues/97: Feature: provide a way to shutdown the thread pool](https://github.com/rust-threadpool/rust-threadpool/issues/97)
//!   - Panics when failing to spawn a thread.
//! - [`scoped_threadpool`](https://crates.io/crates/scoped_threadpool)
//!   - Popular
//!   - Contains `unsafe` code
//!   - Unmaintained
//!   - Does not restart workers on panic.
//! - [`scheduled-thread-pool`](https://crates.io/crates/scheduled-thread-pool)
//!   - Used by a popular connection pool library
//!   - Dependencies have `unsafe` code
//!   - Schedule jobs to run immediately, periodically, or after a specified delay.
//! - [`workerpool`](https://crates.io/crates/workerpool)
//!   - Dependencies have `unsafe` code
//! - [`threads_pool`](https://crates.io/crates/threads_pool)
//!   - Full of `unsafe`
//! - [`thread-pool`](https://crates.io/crates/thread-pool)
//!   - Old
//!   - Dependencies have `unsafe` code
//! - [`tasque`](https://crates.io/crates/tasque)
//!   - Dependencies have `unsafe` code
//! - [`fast-threadpool`](https://crates.io/crates/fast-threadpool)
//!   - Dependencies have `unsafe` code
//! - [`blocking-permit`](https://crates.io/crates/blocking-permit)
//!   - Full of `unsafe`
//! - [`rayon-core`](https://crates.io/crates/rayon-core)
//!   - Full of `unsafe`
//!
//! # Changelog
//! <details>
//! <summary>Changelog</summary>
//!
//! - v0.2.4 - Update docs.
//! - v0.2.3 - Implement `From<NewThreadPoolError>` and `From<TryScheduleError>` for `std::io::Error`.
//! - v0.2.2 - Add `ThreadPool::join` and `ThreadPool::try_join`.
//! - v0.2.1 - Improve test coverage.
//! - v0.2.0
//!   - `ThreadPool::new` to return `Result`.
//!   - `ThreadPool::try_schedule` to return an error when it fails to restart panicked threads.
//!   - `ThreadPool::schedule` to handle failure starting replacement threads.
//! - v0.1.4 - Stop threads on drop.
//! - v0.1.3 - Support stable Rust!  Needs 1.51+.
//! - v0.1.2 - Add another example
//! - v0.1.1 - Simplified internals and improved documentation.
//! - v0.1.0 - First release
//!
//! </details>
//!
//! # TO DO
//! - Make `join` and `try_join` work with `Arc<ThreadPool>`.
//! - Log a warning when all threads panicked.
//! - Update test coverage.
//! - Add a public `respawn_threads` function.
//! - Add a stress test
//! - Add a benchmark.  See benchmarks in <https://crates.io/crates/executors>
//! - Add a way for a job to schedule another job on the same thread, with stealing.
#![forbid(unsafe_code)]

mod atomic_counter;

use atomic_counter::AtomicCounter;
use core::fmt::{Debug, Display, Formatter};
use core::time::Duration;
use std::error::Error;
use std::io::ErrorKind;
use std::sync::mpsc::{Receiver, RecvTimeoutError, SyncSender, TrySendError};
use std::sync::{Arc, Mutex};
use std::time::Instant;

#[cfg(feature = "testing")]
#[doc(hidden)]
pub static INTERNAL_MAX_THREADS: core::sync::atomic::AtomicUsize =
    core::sync::atomic::AtomicUsize::new(usize::MAX);

fn sleep_ms(ms: u64) {
    std::thread::sleep(Duration::from_millis(ms));
}

fn err_eq(a: &std::io::Error, b: &std::io::Error) -> bool {
    a.kind() == b.kind() && format!("{}", a) == format!("{}", b)
}

#[derive(Debug)]
pub enum StartThreadsError {
    /// The pool has no threads and `std::thread::Builder::spawn` returned the included error.
    NoThreads(std::io::Error),
    /// The pool has at least one thread and `std::thread::Builder::spawn` returned the included error.
    Respawn(std::io::Error),
}
impl Display for StartThreadsError {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        match self {
            StartThreadsError::NoThreads(e) => write!(
                f,
                "ThreadPool workers all panicked, failed starting replacement threads: {}",
                e
            ),
            StartThreadsError::Respawn(e) => {
                write!(
                    f,
                    "ThreadPool failed starting threads to replace panicked threads: {}",
                    e
                )
            }
        }
    }
}
impl Error for StartThreadsError {}
impl PartialEq for StartThreadsError {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (StartThreadsError::NoThreads(a), StartThreadsError::NoThreads(b))
            | (StartThreadsError::Respawn(a), StartThreadsError::Respawn(b)) => err_eq(a, b),
            _ => false,
        }
    }
}
impl Eq for StartThreadsError {}

#[derive(Debug)]
pub enum NewThreadPoolError {
    Parameter(String),
    /// `std::thread::Builder::spawn` returned the included error.
    Spawn(std::io::Error),
}
impl Display for NewThreadPoolError {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        match self {
            NewThreadPoolError::Parameter(s) => write!(f, "{}", s),
            NewThreadPoolError::Spawn(e) => {
                write!(f, "ThreadPool failed starting threads: {}", e)
            }
        }
    }
}
impl Error for NewThreadPoolError {}
impl PartialEq for NewThreadPoolError {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (NewThreadPoolError::Parameter(a), NewThreadPoolError::Parameter(b)) => a == b,
            (NewThreadPoolError::Spawn(a), NewThreadPoolError::Spawn(b)) => err_eq(a, b),
            _ => false,
        }
    }
}
impl Eq for NewThreadPoolError {}
impl From<StartThreadsError> for NewThreadPoolError {
    fn from(err: StartThreadsError) -> Self {
        match err {
            StartThreadsError::NoThreads(e) | StartThreadsError::Respawn(e) => {
                NewThreadPoolError::Spawn(e)
            }
        }
    }
}
impl From<NewThreadPoolError> for std::io::Error {
    fn from(new_thread_pool_error: NewThreadPoolError) -> Self {
        match new_thread_pool_error {
            NewThreadPoolError::Parameter(s) => std::io::Error::new(ErrorKind::InvalidInput, s),
            NewThreadPoolError::Spawn(s) => {
                std::io::Error::new(ErrorKind::Other, format!("failed to start threads: {}", s))
            }
        }
    }
}

#[derive(Debug)]
pub enum TryScheduleError {
    QueueFull,
    /// The pool has no threads and `std::thread::Builder::spawn` returned the included error.
    NoThreads(std::io::Error),
    /// The pool has at least one thread and `std::thread::Builder::spawn` returned the included error.
    Respawn(std::io::Error),
}
impl Display for TryScheduleError {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        match self {
            TryScheduleError::QueueFull => write!(f, "ThreadPool queue is full"),
            TryScheduleError::NoThreads(e) => write!(
                f,
                "ThreadPool workers all panicked, failed starting replacement threads: {}",
                e
            ),
            TryScheduleError::Respawn(e) => {
                write!(
                    f,
                    "ThreadPool failed starting threads to replace panicked threads: {}",
                    e
                )
            }
        }
    }
}
impl Error for TryScheduleError {}
impl PartialEq for TryScheduleError {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (TryScheduleError::QueueFull, TryScheduleError::QueueFull) => true,
            (TryScheduleError::NoThreads(a), TryScheduleError::NoThreads(b))
            | (TryScheduleError::Respawn(a), TryScheduleError::Respawn(b)) => err_eq(a, b),
            _ => false,
        }
    }
}
impl Eq for TryScheduleError {}
impl From<StartThreadsError> for TryScheduleError {
    fn from(err: StartThreadsError) -> Self {
        match err {
            StartThreadsError::NoThreads(e) => TryScheduleError::NoThreads(e),
            StartThreadsError::Respawn(e) => TryScheduleError::Respawn(e),
        }
    }
}
impl From<TryScheduleError> for std::io::Error {
    fn from(try_schedule_error: TryScheduleError) -> Self {
        match try_schedule_error {
            TryScheduleError::QueueFull => {
                std::io::Error::new(ErrorKind::WouldBlock, "TryScheduleError::QueueFull")
            }
            TryScheduleError::NoThreads(e) => std::io::Error::new(
                e.kind(),
                format!(
                    "ThreadPool workers all panicked, failed starting replacement threads: {}",
                    e
                ),
            ),
            TryScheduleError::Respawn(e) => std::io::Error::new(
                e.kind(),
                format!(
                    "ThreadPool failed starting threads to replace panicked threads: {}",
                    e
                ),
            ),
        }
    }
}

struct Inner {
    name: &'static str,
    next_name_num: AtomicCounter,
    size: usize,
    receiver: Mutex<Receiver<Box<dyn FnOnce() + Send>>>,
}

impl Inner {
    fn num_live_threads(self: &Arc<Self>) -> usize {
        Arc::strong_count(self) - 1
    }

    fn work(self: &Arc<Self>) {
        loop {
            let recv_result = self
                .receiver
                .lock()
                .unwrap()
                .recv_timeout(Duration::from_millis(500));
            match recv_result {
                Ok(f) => {
                    let _ignored = self.start_threads();
                    f();
                }
                Err(RecvTimeoutError::Timeout) => {}
                // ThreadPool was dropped.
                Err(RecvTimeoutError::Disconnected) => return,
            };
            let _ignored = self.start_threads();
        }
    }

    #[allow(clippy::unused_self)]
    #[allow(unused_variables)]
    fn spawn_thread(
        &self,
        num_live_threads: usize,
        name: String,
        f: impl FnOnce() + Send + 'static,
    ) -> Result<(), std::io::Error> {
        // I found no way to make std::thread fail reliably on both macOS & Linux.  So we simulate it with this.
        #[cfg(feature = "testing")]
        if num_live_threads >= INTERNAL_MAX_THREADS.load(std::sync::atomic::Ordering::Acquire) {
            return Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                "err1".to_string(),
            ));
        }

        std::thread::Builder::new().name(name).spawn(f)?;
        Ok(())
    }

    fn start_thread(self: &Arc<Self>) -> Result<(), StartThreadsError> {
        let self_clone = self.clone();
        let num_live_threads = self.num_live_threads() - 1;
        if num_live_threads < self.size {
            self.spawn_thread(
                num_live_threads,
                format!("{}{}", self.name, self.next_name_num.next()),
                move || self_clone.work(),
            )
            .map_err(|e| {
                if num_live_threads == 0 {
                    StartThreadsError::NoThreads(e)
                } else {
                    StartThreadsError::Respawn(e)
                }
            })?;
        }
        Ok(())
    }

    fn start_threads(self: &Arc<Self>) -> Result<(), StartThreadsError> {
        while self.num_live_threads() < self.size {
            self.start_thread()?;
        }
        Ok(())
    }
}

/// A collection of threads and a queue for jobs (`FnOnce` structs) they execute.
///
/// Threads stop when they execute a job that panics.
/// If one thread survives, it will recreate all the threads.
/// The next call to [`schedule`](#method.schedule) or [`try_schedule`](#method.try_schedule)
/// also recreates threads.
///
/// If your threadpool load is bursty and you want to automatically recover
/// from an all-threads-panicked state, you could use
/// [`safina_timer`](https://crates.io/crates/safina-timer) to periodically call
/// [`schedule`](#method.schedule) or [`try_schedule`](#method.try_schedule).
///
/// After drop, threads stop as they become idle.
///
/// # Example
/// ```rust
/// # type ProcessResult = ();
/// # fn process_data(data: (), sender: std::sync::mpsc::Sender<ProcessResult>) -> ProcessResult {
/// #    sender.send(()).unwrap();
/// # }
/// # fn f() {
/// # let data_source = vec![(),()];
/// let pool =
///     safina_threadpool::ThreadPool::new("worker", 2).unwrap();
/// let receiver = {
///     let (sender, receiver) =
///         std::sync::mpsc::channel();
///     for data in data_source {
///         let sender_clone = sender.clone();
///         pool.schedule(
///             move || process_data(data, sender_clone));
///     }
///     receiver
/// };
/// let results: Vec<ProcessResult> =
///     receiver.iter().collect();
/// // ...
/// # }
/// ```
///
/// ```rust
/// # use core::time::Duration;
/// # use std::sync::Arc;
/// let pool =
///     Arc::new(
///         safina_threadpool::ThreadPool::new("worker", 2).unwrap());
/// let executor = safina_executor::Executor::default();
/// safina_timer::start_timer_thread();
/// let pool_clone = pool.clone();
/// executor.spawn(async move {
///     loop {
///         safina_timer::sleep_for(Duration::from_millis(500)).await;
///         pool_clone.schedule(|| {});
///     }
/// });
/// # assert_eq!(2, pool.num_live_threads());
/// # for _ in 0..2 {
/// #     pool.schedule(|| {
/// #         std::thread::sleep(Duration::from_millis(100));
/// #         panic!("ignore this panic")
/// #     });
/// # }
/// # std::thread::sleep(Duration::from_millis(200));
/// # assert_eq!(0, pool.num_live_threads());
/// # std::thread::sleep(Duration::from_millis(500));
/// # assert_eq!(2, pool.num_live_threads());
/// ```
pub struct ThreadPool {
    inner: Arc<Inner>,
    sender: SyncSender<Box<dyn FnOnce() + Send>>,
}
impl ThreadPool {
    /// Creates a new thread pool containing `size` threads.
    /// The threads all start immediately.
    ///
    /// Threads are named with `name` with a number.
    /// For example, `ThreadPool::new("worker", 2)`
    /// creates threads named "worker-1" and "worker-2".
    /// If one of those threads panics, the pool creates "worker-3".
    ///
    /// After the `ThreadPool` struct drops, the threads continue processing
    /// jobs and stop when the queue is empty.
    ///
    /// # Errors
    /// Returns an error when `name` is empty, `size` is zero, or it fails to start the threads.
    pub fn new(name: &'static str, size: usize) -> Result<Self, NewThreadPoolError> {
        if name.is_empty() {
            return Err(NewThreadPoolError::Parameter(
                "ThreadPool::new called with empty name".to_string(),
            ));
        }
        if size < 1 {
            return Err(NewThreadPoolError::Parameter(format!(
                "ThreadPool::new called with invalid size value: {:?}",
                size
            )));
        }
        // Use a channel with bounded size.
        // If the channel was unbounded, the process could OOM when throughput goes down.
        let (sender, receiver) = std::sync::mpsc::sync_channel(size * 200);
        let pool = ThreadPool {
            inner: Arc::new(Inner {
                name,
                next_name_num: AtomicCounter::new(),
                size,
                receiver: Mutex::new(receiver),
            }),
            sender,
        };
        pool.inner.start_threads()?;
        Ok(pool)
    }

    /// Returns the number of threads in the pool.
    #[must_use]
    pub fn size(&self) -> usize {
        self.inner.size
    }

    /// Returns the number of threads currently alive.
    #[must_use]
    pub fn num_live_threads(&self) -> usize {
        self.inner.num_live_threads()
    }

    #[cfg(feature = "testing")]
    #[doc(hidden)]
    #[must_use]
    pub fn num_live_threads_fn(&self) -> Box<dyn Fn() -> usize> {
        let inner_clone = self.inner.clone();
        Box::new(move || inner_clone.num_live_threads())
    }

    /// Adds a job to the queue.  The next idle thread will execute it.
    /// Jobs are started in FIFO order.
    ///
    /// Blocks when the queue is full or no threads are running.
    /// See [`try_schedule`](#method.try_schedule).
    ///
    /// Recreates any threads that panicked.
    /// Retries on failure to start a new thread.
    ///
    /// Puts `f` in a [`Box`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html) before
    /// adding it to the queue.
    #[allow(clippy::missing_panics_doc)]
    pub fn schedule<F: FnOnce() + Send + 'static>(&self, f: F) {
        let mut opt_box_f: Option<Box<dyn FnOnce() + Send + 'static>> = Some(Box::new(f));
        loop {
            match self.inner.start_threads() {
                Ok(()) | Err(StartThreadsError::Respawn(_)) => {
                    // At least one thread is running.
                }
                Err(StartThreadsError::NoThreads(_)) => {
                    sleep_ms(10);
                    continue;
                }
            }
            opt_box_f = match self.sender.try_send(opt_box_f.take().unwrap()) {
                Ok(()) => return,
                Err(TrySendError::Disconnected(_)) => unreachable!(),
                Err(TrySendError::Full(box_f)) => Some(box_f),
            };
            sleep_ms(10);
        }
    }

    /// Adds a job to the queue and then starts threads to replace any panicked threads.
    /// The next idle thread will execute the job.
    /// Starts jobs in FIFO order.
    ///
    /// Puts `f` in a [`Box`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html) before
    /// adding it to the queue.
    ///
    /// # Errors
    /// Returns an error when the queue is full or it fails to start a thread.
    /// If the return value is not `TryScheduleError::QueueFull` then it added the job to the queue.
    #[allow(clippy::missing_panics_doc)]
    pub fn try_schedule(&self, f: impl FnOnce() + Send + 'static) -> Result<(), TryScheduleError> {
        match self.sender.try_send(Box::new(f)) {
            Ok(_) => {}
            Err(TrySendError::Disconnected(_)) => unreachable!(),
            Err(TrySendError::Full(_)) => return Err(TryScheduleError::QueueFull),
        };
        self.inner.start_threads().map_err(std::convert::Into::into)
    }

    /// Consumes the thread pool and waits for all threads to stop.
    pub fn join(self) {
        let inner = self.inner.clone();
        drop(self);
        while inner.num_live_threads() > 0 {
            sleep_ms(10);
        }
    }

    /// Consumes the thread pool and waits for all threads to stop.
    ///
    /// # Errors
    /// Returns an error if the threads do not stop within the timeout duration.
    pub fn try_join(self, timeout: Duration) -> Result<(), String> {
        let inner = self.inner.clone();
        drop(self);
        let deadline = Instant::now() + timeout;
        loop {
            if inner.num_live_threads() < 1 {
                return Ok(());
            }
            if deadline < Instant::now() {
                return Err("timed out waiting for ThreadPool workers to stop".to_string());
            }
            sleep_ms(10);
        }
    }
}
impl Debug for ThreadPool {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), core::fmt::Error> {
        write!(
            f,
            "ThreadPool{{{:?},size={:?}}}",
            self.inner.name, self.inner.size
        )
    }
}