fix counter bug

src/bin/main.rs

@@ -10,24 +10,26 @@ use std::thread;
 use std::time::Duration;
 
 fn main() {
-    let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
+    let addr = "127.0.0.1:7878";
-    let counter = Arc::new(Mutex::new(0usize));
+    let listener = TcpListener::bind(&addr).unwrap();
+    println!("Listening on http://{}", &addr);
 
+    let counter = Arc::new(Mutex::new(0usize));
     let tasks = ThreadPool::new(100);
 
-    for stream in listener.incoming() {
+    for stream in listener.incoming()/*.take(2)*/ {
-        let stream = stream.unwrap();
+        println!("accepted new connection");
-        let counter = Arc::clone(&counter);
 
-        tasks.run(move || {
+        if let Ok(stream) = stream {
-            {
+            let counter = Arc::clone(&counter);
-                // in a subscope to release borrow automatically
+
-                let mut count = counter.lock().unwrap();
+            tasks.run(move || {
-                *count += 1;
+                handle_connection(counter, stream);
-            }
+            });
-            handle_connection(counter, stream);
+        }
-        });
     }
+
+    println!("Shutting down...");
 }
 
 fn handle_connection(counter: Arc<Mutex<usize>>, mut stream: TcpStream) {
@@ -37,7 +39,8 @@ fn handle_connection(counter: Arc<Mutex<usize>>, mut stream: TcpStream) {
 
     {
         // in a subscope to unlock automatically
-        let count = counter.lock().unwrap();
+        let mut count = counter.lock().unwrap();
+        *count += 1;
         println!("Request: {}\n{}", count, String::from_utf8_lossy(&buffer[..]));
     }
 

src/lib.rs

@@ -16,29 +16,65 @@ type Job = Box<dyn FnBox + Send + 'static>;
 
 struct Worker {
     id: usize,
-    thread: thread::JoinHandle<()>,
+    thread: Option<thread::JoinHandle<()>>,
+}
+
+enum Message {
+  NewJob(Job),
+  Terminate,
 }
 
 impl Worker {
-    pub fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>) -> Worker {
+    pub fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Message>>>) -> Worker {
         Worker {
             id,
-            thread: thread::spawn(move || {
+            thread: Some(thread::spawn(move || {
                 loop {
-                    let job = receiver.lock().expect("unrecoverable poisened panicked thread state").recv().unwrap();
+                    if let Ok(job) = receiver.lock().expect("unrecoverable poisened panicked thread state").recv() {
-
+                        match job {
-                    println!("such busy, so worker #{}!", id);
+                            Message::NewJob(job) => {
-
+                                println!("such busy, so worker #{}!", id);
-                    job.call_box();
+                                job.call_box();
+                            }
+                            Message::Terminate => {
+                                println!("such sad, so dead #{}!", id);
+                                break;
+                            }
+                        }
+                    } else {
+                        // TODO respawn? (unless in terminate mode)
+                        // or why would the recv fail? already terminated?
+                        // thread panic not handled?
+                        break;
+                    }
                 }
-            }),
+            })),
         }
     }
 }
 
 pub struct ThreadPool {
     workers: Vec<Worker>,
-    sender: mpsc::Sender<Job>,
+    sender: mpsc::Sender<Message>,
+}
+
+impl Drop for ThreadPool {
+    fn drop(&mut self) {
+        // Each worker will only receive Terminate once
+        for _ in &mut self.workers {
+            self.sender.send(Message::Terminate).unwrap();
+        }
+
+        // Each worker will join once complete
+        for worker in &mut self.workers {
+            println!("Shutting down worker {} ...", worker.id);
+
+            if let Some(thread) = worker.thread.take() {
+                thread.join().unwrap();
+                println!("Killed {}!", worker.id);
+            }
+        }
+    }
 }
 
 impl ThreadPool {
@@ -64,6 +100,6 @@ impl ThreadPool {
     pub fn run<F: FnOnce() + Send + 'static>(&self, task: F)
     {
         let task = Box::new(task);
-        self.sender.send(task).unwrap();
+        self.sender.send(Message::NewJob(task)).unwrap();
     }
 }