finish rust book

Author: functora

rust/hello/src/lib.rs

@@ -5,7 +5,7 @@ use std::{
 
 pub struct ThreadPool {
     workers: Vec<Worker>,
-    sender: mpsc::Sender<Job>,
+    sender: Option<mpsc::Sender<Job>>,
 }
 
 impl ThreadPool {
@@ -17,14 +17,27 @@ impl ThreadPool {
         for id in 0..size {
             workers.push(Worker::new(id, Arc::clone(&receiver)));
         }
-        ThreadPool { workers, sender }
+        ThreadPool {
+            workers,
+            sender: Some(sender),
+        }
     }
     pub fn execute<F>(&self, f: F)
     where
         F: FnOnce() + Send + 'static,
     {
         let job = Box::new(f);
-        self.sender.send(job).unwrap();
+        self.sender.as_ref().unwrap().send(job).unwrap();
+    }
+}
+
+impl Drop for ThreadPool {
+    fn drop(&mut self) {
+        drop(self.sender.take());
+        for worker in &mut self.workers.drain(..) {
+            println!("Shutting down the worker {}", worker.id);
+            worker.thread.join().unwrap();
+        }
     }
 }
 
@@ -37,9 +50,17 @@ impl Worker {
     fn new(id: usize, receiver: Arc<Mutex<mpsc::Receiver<Job>>>) -> Worker {
         let thread = thread::spawn(move || {
             loop {
-                let job = receiver.lock().unwrap().recv().unwrap();
-                println!("Worker {id} got a job! Executing!");
-                job();
+                let msg = receiver.lock().unwrap().recv();
+                match msg {
+                    Ok(job) => {
+                        println!("Worker {id} got a job! Executing!");
+                        job();
+                    }
+                    Err(_) => {
+                        println!("Worker {id} disconnected! Shutting down!");
+                        break;
+                    }
+                };
             }
         });
         Worker { id, thread }

rust/hello/src/main.rs

@@ -9,7 +9,7 @@ use std::{
 fn main() {
     let listener = TcpListener::bind("127.0.0.1:7878").unwrap();
     let pool = ThreadPool::new(4);
-    for stream in listener.incoming() {
+    for stream in listener.incoming().take(2) {
         let stream = stream.unwrap();
         pool.execute(|| handle_connection(stream));
     }