diff --git a/src/actor/mod.rs b/src/actor/mod.rs index b2ee98d..fee2a68 100644 --- a/src/actor/mod.rs +++ b/src/actor/mod.rs @@ -193,7 +193,7 @@ mod tests { #[test] fn test_actor_system() { let thread_pool = Arc::new(ThreadPool::new(4)); - let system = ActorSystem::new(thread_pool); + let system = ActorSystem::new(thread_pool.clone()); let count = Arc::new(AtomicUsize::new(0)); let actor = CounterActor { @@ -212,6 +212,13 @@ mod tests { assert_eq!(count.load(Ordering::SeqCst), 6); + // Stop the actor handle.stop().unwrap(); + + // Give actor time to process the stop message + std::thread::sleep(Duration::from_millis(100)); + + // Drop the thread pool explicitly to wait for all threads + drop(thread_pool); } } diff --git a/src/utils/thread_pool.rs b/src/utils/thread_pool.rs index 47ac429..0fbf13c 100644 --- a/src/utils/thread_pool.rs +++ b/src/utils/thread_pool.rs @@ -4,9 +4,14 @@ use std::thread; type Job = Box; +enum Message { + NewJob(Job), + Terminate, +} + pub struct ThreadPool { workers: Vec, - sender: mpsc::Sender, + sender: Option>, active_threads: Arc, } @@ -30,7 +35,7 @@ impl ThreadPool { ThreadPool { workers, - sender, + sender: Some(sender), active_threads, } } @@ -40,7 +45,9 @@ impl ThreadPool { F: FnOnce() + Send + 'static, { let job = Box::new(f); - self.sender.send(job).unwrap(); + if let Some(ref sender) = self.sender { + sender.send(Message::NewJob(job)).unwrap(); + } } pub fn thread_count(&self) -> usize { @@ -54,6 +61,14 @@ impl ThreadPool { impl Drop for ThreadPool { fn drop(&mut self) { + // Send termination message to all workers + if let Some(sender) = self.sender.take() { + for _ in &self.workers { + sender.send(Message::Terminate).unwrap(); + } + } + + // Wait for all workers to finish for worker in &mut self.workers { if let Some(thread) = worker.thread.take() { thread.join().unwrap(); @@ -70,14 +85,14 @@ struct Worker { impl Worker { fn new( id: usize, - receiver: Arc>>, + receiver: Arc>>, _active_threads: Arc, _total_threads: usize, ) -> Worker { let thread = thread::spawn(move || loop { - let job = receiver.lock().unwrap().recv(); - match job { - Ok(job) => { + let message = receiver.lock().unwrap().recv(); + match message { + Ok(Message::NewJob(job)) => { // let active = // active_threads.fetch_add(1, Ordering::SeqCst) + 1; // trace!( @@ -98,7 +113,12 @@ impl Worker { // total_threads // ); } + Ok(Message::Terminate) => { + // Received termination signal + break; + } Err(_) => { + // Channel disconnected break; } }