poolio 0.2.3

//! poolio is a thread pool implementation using only channels for concurrency.
//!
//! ## Design
//!
//! A poolio thread pool is essentially made up of a 'supervisor' thread and a specified number of 'worker' threads.
//! A worker's only purpose is executing jobs (in the form of closures), while the supervisor is responsible for everything else - most importantly, assigning jobs to workers that it receives from outside the pool via the public API.
//! To this end, the thread pool is set up so that the supervisor can communicate with each worker separately and concurrently.
//! This ensures that each worker remains equally busy.
//! A single supervisor-worker communication cycle is roughly as follows:
//! 1. The worker tells the supervisor its current status.
//! 2. The supervisor decides what to tell the worker to do based on the current order-message from outside the pool and the worker's status.
//! 3. The supervisor tells the worker what to do.
//! 4. The worker attempts to perform the task assigned by the supervisor.
//! 5. The worker tells the supervisor its current status.
//!
//! The following graphic illustrates the aforementioned communication model between a supervisor thread S and a worker thread W:
//!
//! <pre>
//!    W
//!    _
//!    .
//!    .
//!    send-status
//!    .   O
//!    .     O
//!    .       O                 send-message
//!    .         O                   O
//!    .           O               O
//!    recv         recv         O
//!   * .  O       O  . .      O
//!  .   .   O   O   .   .   O
//! .     e    O    m     recv . . | S
//!  .   .   O   O   .   *
//!   . .  O       O  . .
//!    send-status  send-message
//!
//! X | . . * : arrow starting at | and ending at * representing the control-flow of thread X
//! O O O O O : channel
//! e : execute job
//! m : manage workers
//! </pre>
//!
//! ## Usage
//!
//! To use a poolio [`ThreadPool`], you simply set one up using the [`ThreadPool::new`] method and task the pool to run jobs using the [`ThreadPool::execute`] method.
//!
//! # Examples
//!
//! Setting up a pool to make a server multi-threaded:
//!
//! ```
//! fn handle(req: usize) {
//!     println!("Handled!")
//! }
//!
//! let server_requests = [1, 2, 3, 4, 5, 6, 7, 8, 9];
//!
//! let pool = poolio::ThreadPool::new(3, poolio::PanicSwitch::Kill).unwrap();
//!
//! for req in server_requests {
//!     pool.execute(move || {
//!         handle(req);
//!     });
//! }
//! ```

mod thread {
    //! This module is a wrapper for parts of the [`std::thread`] module to handle ownership issues when joining threads embedded in a larger data structure.
    //! It allows you to spawn threads that return a handle, which you can join normally even if the handle is part of a larger data structure.

    use std::thread;

    /// Wraps [`std::thread::JoinHandle<T>`] to allow for "stealing" the handle for joining.
    pub type JoinHandle = Option<thread::JoinHandle<()>>;

    /// Wraps [`std::thread::spawn`] in an [`Option::Some`].
    #[inline]
    pub fn spawn<F>(f: F) -> JoinHandle
    where
        F: FnOnce() + Send + 'static,
    {
        Some(thread::spawn(f))
    }

    /// Takes the thread handle from the call site to pass it to [`std::thread::JoinHandle<T>::join`].
    /// - `thread` is a reference to the handle this function intends to take.
    ///
    /// # Panics
    ///
    /// Panics if the `thread` is `None` or if joining the thread fails (which occurs if the thread panicked).
    pub fn join(thread: &mut JoinHandle) {
        let thread = thread.take();

        match thread {
            Some(thread) => {
                if let Err(e) = thread.join() {
                    panic!("{:?}", e);
                }
            }
            None => panic!("Cannot join: no thread has been provided."),
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;

        #[test]
        fn test_spawn() {
            assert!(spawn(|| {}).is_some());
        }

        #[test]
        fn test_join() {
            let mut thread = spawn(|| {});
            join(&mut thread);
            assert!(thread.is_none());
        }

        #[test]
        #[should_panic]
        fn test_join_panic_some() {
            join(&mut spawn(|| panic!("Oh no!")));
        }

        #[test]
        #[should_panic]
        fn test_join_panic_none() {
            join(&mut None);
        }
    }
}

use thread::JoinHandle;

use std::fmt;
use std::panic::UnwindSafe;

use crossbeam::channel::unbounded as channel;
use crossbeam::channel::Sender;

/// The type of jobs the [`ThreadPool`] can run.
type Job = Box<dyn FnOnce() + UnwindSafe + Send + 'static>;

/// Messages containing orders for the [`ThreadPool`].
enum Message {
    /// A message ordering the pool to execute a job.
    NewJob(Job),
    /// A message ordering the pool to finish its remaining jobs and then shut down.
    Terminate,
}

impl fmt::Display for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Self::NewJob(_) => write!(f, "[NewJob]"),
            Self::Terminate => write!(f, "[Terminate]"),
        }
    }
}

/// Configuration for how the [`ThreadPool`] handles panics in jobs.
pub enum PanicSwitch {
    /// The pool finishes parallel running jobs and then kills the entire process if a job panics.
    Kill,
    /// The pool ignores panicked jobs and simply respawns the affected threads.
    Respawn,
}

/// Abstracts thread pools.
pub struct ThreadPool {
    /// Interface to the pool-controlling thread.
    supervisor: Supervisor,
}

impl ThreadPool {
    /// Sets up a new pool.
    /// - `size` is the (non-zero) number of worker threads in the pool.
    /// - `mode` is the setting for the panic switch.
    ///
    /// # Errors
    ///
    /// Returns an error if `size` is 0 (as a pool without worker threads is invalid).
    ///
    /// # Examples
    ///
    /// Setting up a pool with three worker threads in kill-mode:
    ///
    /// ```
    /// let pool = poolio::ThreadPool::new(3, poolio::PanicSwitch::Kill).unwrap();
    /// ```
    pub fn new<'a>(size: usize, mode: PanicSwitch) -> Result<Self, &'a str> {
        if size == 0 {
            return Err("Setting up a pool with no workers is not allowed.");
        }

        let pool = Self {
            supervisor: Supervisor::new(size, mode),
        };
        Ok(pool)
    }

    /// Runs a job in `self`.
    /// - `f` is the job to be run, provided as a closure.
    ///
    /// # Panics
    ///
    /// Panics if the pool is unreachable.
    ///
    /// # Notes
    ///
    /// If `f` panics, the behavior is determined by the [`PanicSwitch`] setting of `self`.
    ///
    /// # Examples
    ///
    /// Setting up a pool and printing two strings concurrently:
    ///
    /// ```
    /// let pool = poolio::ThreadPool::new(2, poolio::PanicSwitch::Kill).unwrap();
    /// pool.execute(|| println!{"house"});
    /// pool.execute(|| println!{"cat"});
    /// ```
    pub fn execute<F>(&self, f: F)
    where
        F: FnOnce() + UnwindSafe + Send + 'static,
    {
        let job = Box::new(f);

        self.send(Message::NewJob(job));
    }

    /// Attempts to shut down `self` gracefully.
    ///
    /// # Panics
    ///
    /// Panics if:
    /// 1. The pool is unreachable.
    /// 2. Joining the threads causes a panic.
    ///
    /// # Notes
    ///
    /// Graceful shutdown ensures all remaining jobs are finished (except for panics in [`PanicSwitch::Kill`] mode).
    fn terminate(&mut self) {
        self.send(Message::Terminate);

        thread::join(&mut self.supervisor.thread);
    }

    /// Wraps sending a [`Message`] to the pool.
    ///
    /// # Panics
    ///
    /// Panics if the receiver has already been deallocated.
    fn send(&self, msg: Message) {
        let panic_message = format!("Ordering {} failed. Pool is unreachable.", msg);

        self.supervisor.orders_s.send(msg).expect(&panic_message);
    }
}

impl Drop for ThreadPool {
    /// Attempts to shut down `self` gracefully.
    ///
    /// # Panics
    ///
    /// Panics if:
    /// 1. The pool is unreachable.
    /// 2. Joining the threads causes a panic.
    ///
    /// Note: A panic during a drop will abort the entire process.
    ///
    /// # Notes
    ///
    /// Graceful shutdown ensures all remaining jobs are finished (except for panics in [`PanicSwitch::Kill`] mode).
    fn drop(&mut self) {
        self.terminate();
    }
}

/// A numeric type used to identify workers.
type StaffNumber = usize;

/// States a worker can be in when not busy.
enum Status {
    Idle(StaffNumber),
    Panic(StaffNumber),
}

impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Self::Idle(_) => write!(f, "[idle]"),
            Self::Panic(_) => write!(f, "[panic]"),
        }
    }
}

/// Abstracts supervisors.
struct Supervisor {
    /// Channel for sending orders.
    orders_s: Sender<Message>,
    /// Handle to join the supervisor thread.
    thread: JoinHandle,
}

impl Supervisor {
    /// Sets up a supervisor.
    /// - `number_of_workers` is the number of workers to employ.
    /// - `mode` configures behavior when workers report panicked jobs.
    fn new(mut number_of_workers: usize, mode: PanicSwitch) -> Self {
        // This channel is used by the pool to contact the supervisor.
        let (orders_s, orders_r) = channel();

        let thread = thread::spawn(move || {
            // This channel is used by the workers to contact the supervisor.
            let (statuses_s, statuses_r) = channel();

            // Construct `number_of_workers` worker threads.
            let mut workers = Vec::with_capacity(number_of_workers);
            for id in 0..number_of_workers {
                workers.push(Worker::new(id, statuses_s.clone()));
            }

            // Track how many jobs have panicked.
            let mut panicked_jobs = 0;

            // Keep running to distribute jobs among idle workers.
            'distribute_jobs: while let Message::NewJob(job) = orders_r.recv().unwrap() {
                'query_status: loop {
                    match statuses_r.recv().unwrap() {
                        Status::Idle(id) => {
                            workers[id]
                                .instructions_s
                                .send(Message::NewJob(job))
                                .unwrap();
                            break 'query_status;
                        }
                        Status::Panic(id) => {
                            thread::join(&mut workers[id].thread);
                            match mode {
                                PanicSwitch::Kill => {
                                    panicked_jobs += 1;
                                    number_of_workers -= 1;
                                    break 'distribute_jobs;
                                }
                                PanicSwitch::Respawn => {
                                    workers[id] = Worker::new(id, statuses_s.clone());
                                }
                            }
                        }
                    }
                }
            }

            // Destruct all remaining worker threads.
            while number_of_workers != 0 {
                match statuses_r.recv().unwrap() {
                    Status::Idle(id) => {
                        workers[id].instructions_s.send(Message::Terminate).unwrap();
                        thread::join(&mut workers[id].thread);
                    }
                    Status::Panic(id) => {
                        thread::join(&mut workers[id].thread);
                        if matches!(mode, PanicSwitch::Kill) {
                            panicked_jobs += 1;
                        }
                    }
                }
                number_of_workers -= 1;
            }

            if panicked_jobs > 0 {
                eprintln!("Aborting process: {} panicked jobs.", panicked_jobs);
                std::process::abort();
            }

            // Ensure that `orders_r` lives as long as the thread to prevent reachability errors.
            drop(orders_r);
        });

        Self { orders_s, thread }
    }
}

/// Abstracts workers.
struct Worker {
    /// Channel for sending instructions.
    instructions_s: Sender<Message>,
    /// Handle to join the worker thread.
    thread: JoinHandle,
}

impl Worker {
    /// Sets up a new worker.
    /// - `id` is the worker's staff number.
    /// - `statuses_s` is where the worker reports its current status.
    fn new(id: StaffNumber, statuses_s: Sender<Status>) -> Self {
        // This channel is used by the supervisor to contact this worker.
        let (instructions_s, instructions_r) = channel();

        let thread = thread::spawn(move || {
            // Report for duty.
            statuses_s.send(Status::Idle(id)).unwrap();

            // Keep running to execute jobs.
            loop {
                let message = instructions_r.recv().unwrap();

                match message {
                    Message::NewJob(job) => match std::panic::catch_unwind(job) {
                        Ok(()) => {
                            statuses_s.send(Status::Idle(id)).unwrap();
                        }
                        Err(_) => {
                            statuses_s.send(Status::Panic(id)).unwrap();
                            break;
                        }
                    },
                    Message::Terminate => break,
                }
            }
        });

        Self {
            instructions_s,
            thread,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    use std::sync::Arc;

    // settings
    const SIZE: usize = 2; //= 6; && = 12; && = 36;
    const MODE: PanicSwitch = PanicSwitch::Respawn; //= PanicSwitch::Kill;
    const ID: StaffNumber = 0;

    #[test]
    fn test_threadpool_new_ok() {
        let pool = ThreadPool::new(SIZE, MODE);
        assert!(pool.is_ok());
    }

    #[test]
    fn test_threadpool_new_err() {
        let pool = ThreadPool::new(0, MODE);
        assert!(pool.is_err());
    }

    #[test]
    fn test_threadpool_execute() {
        const N: usize = 5;

        let pool = ThreadPool::new(SIZE, MODE).unwrap();

        let counter = Arc::new(AtomicUsize::new(0));

        let count_to = |n: usize| {
            for _ in 0..n {
                let counter = Arc::clone(&counter);
                pool.execute(move || {
                    counter.fetch_add(1, Ordering::SeqCst);
                });
            }
        };

        for _ in 0..N {
            count_to(SIZE);
            if matches!(MODE, PanicSwitch::Respawn) {
                pool.execute(|| panic!("Oh no!"));
            }
        }

        drop(pool);

        assert_eq!(N * SIZE, counter.load(Ordering::SeqCst));
    }

    #[test]
    fn test_worker_thread_newjob() {
        let (statuses_s, statuses_r) = channel();
        let mut worker = Worker::new(ID, statuses_s);

        assert!(matches!(statuses_r.recv().unwrap(), Status::Idle(ID)));

        let flag = Arc::new(AtomicBool::new(false));
        let flag_ref = Arc::clone(&flag);
        let job = Box::new(move || {
            flag_ref.store(true, Ordering::SeqCst);
        });
        worker.instructions_s.send(Message::NewJob(job)).unwrap();
        assert!(matches!(statuses_r.recv().unwrap(), Status::Idle(ID)));
        assert!(flag.load(Ordering::SeqCst));

        let job = Box::new(|| panic!("Oh no!"));
        worker.instructions_s.send(Message::NewJob(job)).unwrap();
        assert!(matches!(statuses_r.recv().unwrap(), Status::Panic(ID)));

        thread::join(&mut worker.thread);
    }

    #[test]
    fn test_worker_thread_terminate() {
        let (statuses_s, statuses_r) = channel();
        let mut worker = Worker::new(ID, statuses_s);

        assert!(matches!(statuses_r.recv().unwrap(), Status::Idle(ID)));

        worker.instructions_s.send(Message::Terminate).unwrap();

        thread::join(&mut worker.thread);
    }
}