shellcaster 2.0.1

use std::sync::{mpsc, Arc, Mutex};
use std::thread;

// Much of the threadpool implementation here was taken directly from
// the Rust Book: https://doc.rust-lang.org/book/ch20-02-multithreaded.html
// and https://doc.rust-lang.org/book/ch20-03-graceful-shutdown-and-cleanup.html

/// Manages a threadpool of a given size, sending jobs to workers as
/// necessary. Implements Drop trait to allow threads to complete
/// their current jobs before being stopped.
pub struct Threadpool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<JobMessage>,
}

impl Threadpool {
    /// Creates a new Threadpool of a given size.
    pub fn new(n_threads: usize) -> Threadpool {
        let (sender, receiver) = mpsc::channel();
        let receiver_lock = Arc::new(Mutex::new(receiver));

        let mut workers = Vec::with_capacity(n_threads);

        for _ in 0..n_threads {
            workers.push(Worker::new(Arc::clone(&receiver_lock)));
        }

        return Threadpool {
            workers: workers,
            sender: sender,
        };
    }

    /// Adds a new job to the threadpool, passing closure to first
    /// available worker.
    pub fn execute<F>(&self, func: F)
    where F: FnOnce() + Send + 'static {
        let job = Box::new(func);
        self.sender
            .send(JobMessage::NewJob(job))
            .expect("Thread messaging error");
    }
}

impl Drop for Threadpool {
    /// Upon going out of scope, Threadpool sends terminate message to
    /// all workers but allows them to complete current jobs.
    fn drop(&mut self) {
        for _ in &self.workers {
            self.sender
                .send(JobMessage::Terminate)
                .expect("Thread messaging error");
        }

        for worker in &mut self.workers {
            if let Some(thread) = worker.thread.take() {
                // joins to ensure threads finish job before stopping
                thread.join().expect("Error dropping threads");
            }
        }
    }
}

type Job = Box<dyn FnOnce() + Send + 'static>;

/// Messages used by Threadpool to communicate with Workers.
enum JobMessage {
    NewJob(Job),
    Terminate,
}

/// Used by Threadpool to complete jobs. Each Worker manages a single
/// thread.
struct Worker {
    thread: Option<thread::JoinHandle<()>>,
}

impl Worker {
    /// Creates a new Worker, which waits for Jobs to be passed by the
    /// Threadpool.
    fn new(receiver: Arc<Mutex<mpsc::Receiver<JobMessage>>>) -> Worker {
        let thread = thread::spawn(move || loop {
            let message = receiver
                .lock()
                .expect("Threadpool error")
                .recv()
                .expect("Thread messaging error");

            match message {
                JobMessage::NewJob(job) => job(),
                JobMessage::Terminate => break,
            }
        });

        return Worker {
            thread: Some(thread),
        };
    }
}