brutils 0.1.51

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

/// Simple thread pool implementation
pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: mpsc::Sender<Message>
}

impl ThreadPool {
    /// Creates new thread pool with specified number of worker threads
    pub fn new(size: usize) -> Self {
        assert!(size > 0);
        let mut workers = Vec::with_capacity(size);
        let (sender, receiver) = mpsc::channel();
        let receiver = Arc::new(Mutex::new(receiver));

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

        ThreadPool { workers, sender }
    }

    /// Schedules new job to be executed
    pub fn execute<T: FnOnce() + Send + 'static>(&self, f: T) {
        let job = Box::new(f);
        self.sender.send(Message::NewJob(job)).expect("Failed to send job to worker threads");
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        for _ in &mut self.workers {
            self.sender.send(Message::Terminate).expect("Failed to send message to terminate worker thread");
        }

        for worker in &mut self.workers {
            if let Some(thread) = worker.thread.take() {
                thread.join().expect("Failed to join worker thread");
            }
        }
    }
}

struct Worker {
    thread: Option<thread::JoinHandle<()>>
}

impl Worker {
    fn new(receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Self {
        let thread = thread::spawn(move || {
            loop {
                let message = receiver.lock()
                    .expect("Failed to lock mutex while receiving job on worker thread")
                    .recv()
                    .expect("Failed to receive job on worker thread");
                
                match message {
                    Message::NewJob(job) => job.call_box(),
                    Message::Terminate => break
                }
            }
        });

        Worker { thread: Some(thread) }
    }
}

type Job = Box<FnBox + Send + 'static>;

trait FnBox {
    fn call_box(self: Box<Self>);
}

impl<T: FnOnce()> FnBox for T {
    fn call_box(self: Box<T>) {
        (*self)();
    }
}

enum Message {
    NewJob(Job),
    Terminate
}