jobs-vibe-coded

jobs.rs (5919B)

---

 
 use std::sync::{mpsc, Arc, Mutex};
 use std::thread;
 
 // A worker that stores no T data directly.
 struct Worker {
     id: usize,
     handle: Option<thread::JoinHandle<()>>,
 }
 
 impl Worker {
     /// Generic function that spawns a thread capable of processing `Message<T>`.
     /// This does the same job as `Worker<T>::new` did, but doesn't make `Worker` itself generic.
     fn spawn_worker_for<T: Send + 'static>(
         id: usize,
         receiver: Arc<Mutex<mpsc::Receiver<Message<T>>>>,
     ) -> Self {
         let handle = thread::spawn(move || loop {
             match receiver.lock().expect("lock poisoned").recv() {
                 Ok(Message::NewJob(data, job)) => {
                     println!("Worker {id} got a job; executing on data...");
                     (job)(data);
                 }
                 Ok(Message::Terminate) => {
                     println!("Worker {id} told to terminate.");
                     break;
                 }
                 Err(_) => {
                     println!("Worker {id} failed to receive message; terminating.");
                     break;
                 }
             }
         });
 
         Worker {
             id,
             handle: Some(handle),
         }
     }
 }
 
 // Our job and message definitions
 type Job<T> = Box<dyn FnOnce(T) + Send + 'static>;
 
 pub enum Message<T> {
     NewJob(T, Job<T>),
     Terminate,
 }
 
 // The thread pool is still generic over T
 pub struct ThreadPool<T> {
     workers: Vec<Worker>,
     sender: mpsc::Sender<Message<T>>,
 }
 
 impl<T: Send + 'static> ThreadPool<T> {
     pub fn new(size: usize) -> Self {
         assert!(size > 0);
         let (sender, receiver) = mpsc::channel::<Message<T>>();
         let receiver = Arc::new(Mutex::new(receiver));
 
         // Create plain `Worker` structs using a generic spawn function
         let mut workers = Vec::with_capacity(size);
         for id in 0..size {
             workers.push(Worker::spawn_worker_for(id, Arc::clone(&receiver)));
         }
 
         Self { workers, sender }
     }
 
     pub fn execute<F>(&self, data: T, f: F)
     where
         F: FnOnce(T) + Send + 'static,
     {
         let job = Box::new(f) as Job<T>;
         self.sender
             .send(Message::NewJob(data, job))
             .expect("Failed to send job to worker.");
     }
 }
 
 impl<T> Drop for ThreadPool<T> {
     fn drop(&mut self) {
         // Ask each worker to terminate
         for _ in &self.workers {
             let _ = self.sender.send(Message::Terminate);
         }
 
         // Join them
         for worker in &mut self.workers {
             if let Some(handle) = worker.handle.take() {
                 handle.join().expect("Worker thread panicked.");
             }
         }
     }
 }
 
 // Example usage
 fn main() {
     let pool = ThreadPool::new(4);
 
     for i in 0..8 {
         pool.execute(i, move |num| {
             println!("Data {num} processed in a worker thread.");
         });
     }
 }
 
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use std::sync::{mpsc, Arc, Mutex};
     use std::thread;
     use std::time::Duration;
 
     /// Test basic task execution and verify that we can gather all results.
     #[test]
     fn test_execute_tasks() {
         let pool = ThreadPool::new(4);
         let (tx, rx) = mpsc::channel();
 
         // Send 8 jobs to the thread pool. Each job sends back its data via tx.
         for i in 0..8 {
             let tx_clone = tx.clone();
             pool.execute(i, move |num| {
                 // Simulate some work
                 thread::sleep(Duration::from_millis(10));
                 tx_clone.send(num).expect("Failed to send result");
             });
         }
 
         // Gather results
         let mut results = Vec::new();
         for _ in 0..8 {
             results.push(rx.recv().expect("Failed to receive result"));
         }
         results.sort();
         assert_eq!(results, (0..8).collect::<Vec<_>>());
     }
 
     /// Test that the thread pool does not allow a size of 0.
     /// This test should panic due to the assert!(size > 0) in `ThreadPool::new`.
     #[test]
     #[should_panic(expected = "size > 0")]
     fn test_zero_size_pool_should_panic() {
         let _ = ThreadPool::<i32>::new(0);
     }
 
     /// Test dropping the pool after submitting jobs.
     /// If everything is correct, the workers will gracefully terminate.
     #[test]
     fn test_dropping_pool() {
         let pool = ThreadPool::new(2);
         let (tx, rx) = mpsc::channel();
 
         for i in 0..4 {
             let tx_clone = tx.clone();
             pool.execute(i, move |num| {
                 // Simulate work
                 thread::sleep(Duration::from_millis(5));
                 tx_clone.send(num).unwrap();
             });
         }
 
         // Scope so pool goes out of scope and triggers drop after sending messages
         drop(pool);
 
         // We should still be able to receive all 4 messages
         let mut results = vec![];
         for _ in 0..4 {
             results.push(rx.recv().expect("Failed to receive result"));
         }
         results.sort();
         assert_eq!(results, vec![0, 1, 2, 3]);
     }
 
     /// A test that ensures multiple data types can be handled.
     /// In reality, we just need T: Send + 'static, but let's be explicit.
     #[test]
     fn test_execute_string_tasks() {
         let pool = ThreadPool::new(2);
         let (tx, rx) = mpsc::channel();
 
         let strings = vec!["alpha", "beta", "gamma", "delta"];
         for s in strings {
             let tx_clone = tx.clone();
             pool.execute(s.to_string(), move |st: String| {
                 // Just pass the string back
                 tx_clone.send(st).unwrap();
             });
         }
 
         // Collect results
         let mut results = Vec::new();
         for _ in 0..4 {
             let r = rx.recv().unwrap();
             results.push(r);
         }
         results.sort();
         assert_eq!(results, vec!["alpha", "beta", "delta", "gamma"]);
     }
 }