Added graceful termination. Book done for real this time XD

2 files changed, 43 insertions, 12 deletions

diff --git a/hello_server/src/bin/main.rs b/hello_server/src/bin/main.rs
index 522777a..bcf1e46 100644
--- a/hello_server/src/bin/main.rs
+++ b/hello_server/src/bin/main.rs
@@ -10,7 +10,7 @@ fn main() {
     let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
     let pool = ThreadPool::new(4).unwrap();
 
-    for stream in listener.incoming() {
+    for stream in listener.incoming().take(2) {
         let stream = stream.unwrap();
 
         pool.execute(|| {
@@ -18,6 +18,8 @@ fn main() {
         });
     }
 
+    println!("shutting down...");
+
 }
 
 fn handle_connection(mut stream: TcpStream) {
diff --git a/hello_server/src/lib.rs b/hello_server/src/lib.rs
index c8e6a3e..cd1f616 100644
--- a/hello_server/src/lib.rs
+++ b/hello_server/src/lib.rs
@@ -3,6 +3,11 @@ use std::sync::Mutex;
 use std::sync::mpsc;
 use std::thread;
 
+enum Message {
+    NewJob(Job),
+    Terminate,
+}
+
 trait FnBox {
     fn call_box(self: Box<Self>);
 }
@@ -17,7 +22,7 @@ type Job = Box<FnBox + Send + 'static>;
 
 pub struct ThreadPool {
     workers: Vec<Worker>,
-    sender: mpsc::Sender<Job>,
+    sender: mpsc::Sender<Message>,
 }
 
 impl ThreadPool {
@@ -55,34 +60,58 @@ impl ThreadPool {
             F: FnOnce() + Send + 'static 
     {
         let job = Box::new(f);
-        self.sender.send(job).unwrap();
+        self.sender.send(Message::NewJob(job)).unwrap();
     }
 
 }
 
+impl Drop for ThreadPool {
+    fn drop(&mut self) {
+        for _ in &mut self.workers {
+            self.sender.send(Message::Terminate).unwrap();
+        }
+
+        println!("shutting down all workers...");
+
+        for worker in &mut self.workers {
+            println!("shutting down worker {}", worker.id);
+
+            if let Some(thread) = worker.thread.take() {
+                thread.join().unwrap();
+            }
+        }
+    }
+}
+
 struct Worker {
     id: usize,
-    thread: thread::JoinHandle<()>,
+    thread: Option<thread::JoinHandle<()>>,
 }
 
 impl Worker {
-    pub fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>)
+    pub fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Message>>>)
         -> Result<Worker, &'static str>
     {
         let thread = thread::spawn(move || { 
             loop {
-                let job = receiver
-                    .lock().expect("lock poisoned!")
-                    .recv().unwrap();
-                println!("worker {} got job; executing...", id);
-
-                job.call_box();
+                let message = receiver.lock().unwrap().recv().unwrap();
+
+                match message {
+                    Message::NewJob(job) => {
+                        println!("worker {} got job, executing...", id);
+                        job.call_box();
+                    },
+                    Message::Terminate => {
+                        println!("worker {} was told to terminate...", id);
+                        break;
+                    },
+                }
             }
         });
 
         Ok(Worker {
             id,
-            thread,
+            thread: Some(thread),
         })
     }
 }