move thread in worker + impl manager stop

src/main.rs

@@ -5,14 +5,12 @@ mod model;
 mod worker;
 
 use std::sync::{mpsc::channel, Arc, Mutex};
-use std::thread;
-use std::time;
 
 use chrono::prelude::*;
 use database::init_database_pool;
 use message::{Message, Subject};
 use model::{Job, JobAction};
-use worker::{CheckHandler, DeployHandler, Manager, WorkerStatus};
+use worker::{CheckHandler, DeployHandler, Manager};
 
 fn main() {
     let now: DateTime<Utc> = Utc::now();
@@ -31,15 +29,18 @@ fn main() {
     checker.launch_workers::<Job, CheckHandler>(5, Arc::new(check_handler));
     checker.subscribe(receiver);
 
+    // test message handling and subscription
     for i in 0..500 {
         let message = Message::new(i, Subject::Action(JobAction::Deploy));
         deployer.put_message(message);
     }
 
-    let wait = time::Duration::from_millis(100);
-    while !deployer.healthcheck(WorkerStatus::Stopped) {
-        thread::sleep(wait);
-    }
+    // deployer and checker (order matters)
+    deployer.stop();
+    assert_eq!(0, deployer.len());
+
+    checker.stop();
+    assert_eq!(0, checker.len());
 
     let elapsed = Utc::now() - now;
     println!(

src/worker/worker.rs

@@ -1,10 +1,10 @@
 use std::collections::VecDeque;
 use std::sync::{mpsc::Receiver, Arc, Condvar, Mutex};
-use std::thread;
+use std::{thread, time};
 
 use super::handler::Handler;
 use crate::error::HandlerError;
-use crate::message::Message;
+use crate::message::{Message, Subject};
 use crate::model::*;
 
 // Queue is a simple queue data structure.
@@ -25,7 +25,6 @@ impl<T> Queue<T> {
 }
 
 #[derive(Debug, Copy, Clone, PartialEq)]
-#[allow(dead_code)]
 pub enum WorkerStatus {
     Pending,
     Running,
@@ -33,59 +32,35 @@ pub enum WorkerStatus {
     Stopped,
 }
 
-// Worker holds the computation in thread with a status.
 struct Worker {
     id: u32,
     manager: String,
-    status: WorkerStatus,
+    queue: Arc<Queue<Message>>,
+    status: Arc<Mutex<WorkerStatus>>,
 }
 
 impl Worker {
-    fn new(id: u32, manager: String) -> Self {
+    fn new(id: u32, manager: String, queue: Arc<Queue<Message>>) -> Self {
         Self {
             id: id,
             manager: manager,
-            status: WorkerStatus::Pending,
+            queue: queue,
+            status: Arc::new(Mutex::new(WorkerStatus::Pending)),
         }
     }
 
-    fn set_status(&mut self, status: WorkerStatus) {
-        self.status = status;
-    }
-
-    #[allow(dead_code)]
     fn get_status(&self) -> WorkerStatus {
-        self.status
-    }
+        let guard = self.status.lock().unwrap();
+        println!(
+            "[worker({} - {})] status: {:?}",
+            self.manager, self.id, *guard
+        );
+        *guard
     }
 
-// Manager is a pool of workers and holds a queue containing `Job` to run.
-pub struct Manager {
-    name: String,
-    workers: Vec<Arc<Mutex<Worker>>>,
-    queue: Arc<Queue<Message>>,
-}
-
-impl Manager {
-    pub fn new(name: &str) -> Self {
-        Self {
-            name: name.to_string(),
-            workers: vec![],
-            queue: Arc::new(Queue::new()),
-        }
-    }
-
-    pub fn launch_workers<T: Runner, U: Handler>(
-        &mut self,
-        nb_workers: u32,
-        shared_handler: Arc<U>,
-    ) {
-        for i in 0..nb_workers {
-            let shared_worker = Arc::new(Mutex::new(Worker::new(i, self.name.clone())));
-            self.workers.push(shared_worker.clone());
-
-            let worker = shared_worker.clone();
+    fn launch<T: Runner, U: Handler>(&self, shared_handler: Arc<U>) {
         let queue = self.queue.clone();
+        let status = self.status.clone();
         let handler = shared_handler.clone();
 
         thread::spawn(move || loop {
@@ -99,6 +74,11 @@ impl Manager {
                 }
             };
 
+            if message.get_subject() == Subject::StopManager {
+                Worker::set_status(&status, WorkerStatus::Stopped);
+                break;
+            }
+
             drop(guard);
 
             let mut runner = match T::try_from(message) {
@@ -109,19 +89,63 @@ impl Manager {
                 }
             };
 
-                Manager::set_worker_status(&worker, WorkerStatus::Running);
+            Worker::set_status(&status, WorkerStatus::Running);
             if let Err(e) = handler.handle(runner.run()) {
                 match e {
                     HandlerError::Unknown => {
-                            Manager::set_worker_status(&worker, WorkerStatus::Failed);
+                        Worker::set_status(&status, WorkerStatus::Failed);
                         break;
                     }
                     _ => (),
                 }
             }
-                Manager::set_worker_status(&worker, WorkerStatus::Pending);
+            Worker::set_status(&status, WorkerStatus::Pending);
         });
     }
+
+    fn set_status(old: &Arc<Mutex<WorkerStatus>>, new: WorkerStatus) {
+        let mut guard = old.lock().unwrap();
+        *guard = new;
+    }
+}
+
+#[derive(PartialEq)]
+pub enum ManagerStatus {
+    Up,
+    Down,
+    Stopping,
+}
+
+// Manager is a pool of workers and holds a queue containing `Job` to run.
+pub struct Manager {
+    name: String,
+    status: ManagerStatus,
+    workers: Vec<Worker>,
+    queue: Arc<Queue<Message>>,
+}
+
+impl Manager {
+    pub fn new(name: &str) -> Self {
+        Self {
+            name: name.to_string(),
+            workers: vec![],
+            status: ManagerStatus::Down,
+            queue: Arc::new(Queue::new()),
+        }
+    }
+
+    pub fn launch_workers<T: Runner, U: Handler>(
+        &mut self,
+        nb_workers: u32,
+        shared_handler: Arc<U>,
+    ) {
+        for i in 0..nb_workers {
+            let worker = Worker::new(i, self.name.clone(), self.queue.clone());
+            worker.launch::<T, U>(shared_handler.clone());
+            self.workers.push(worker);
+        }
+
+        self.status = ManagerStatus::Up;
     }
 
     // subscribe subscribes to a `Receiver` channel.
@@ -144,29 +168,42 @@ impl Manager {
         });
     }
 
+    pub fn len(&self) -> usize {
+        self.queue.content.lock().unwrap().len()
+    }
+
     pub fn put_message(&self, message: Message) {
+        if self.status != ManagerStatus::Up {
+            return;
+        }
+
         let mut q = self.queue.content.lock().unwrap();
         q.push_back(message);
 
         self.queue.not_empty.notify_one();
     }
 
+    pub fn stop(&mut self) {
+        for _ in &self.workers {
+            self.put_message(Message::stop());
+        }
+
+        self.status = ManagerStatus::Stopping;
+
+        let wait = time::Duration::from_millis(100);
+        while !self.healthcheck(WorkerStatus::Stopped) {
+            thread::sleep(wait);
+        }
+
+        self.status = ManagerStatus::Down;
+    }
+
     pub fn healthcheck(&self, target: WorkerStatus) -> bool {
         for w in &self.workers {
-            if w.lock().unwrap().get_status() != target {
+            if w.get_status() != target {
                 return false;
             }
         }
         true
     }
-
-    fn set_worker_status(worker: &Arc<Mutex<Worker>>, status: WorkerStatus) {
-        let mut guard = worker.lock().unwrap();
-        guard.set_status(status);
-
-        println!(
-            "[worker({} - {})] status: {:?}",
-            guard.manager, guard.id, guard.status
-        );
-    }
 }