FIX: Thread leak in pool (#9)

Co-authored-by: Artem Vasiliev <[email protected]>

Author: YaZasnyal

src/lib.rs

@@ -160,16 +160,18 @@ impl Threadpool {
 		if let Some(stack_size) = data.stack_size {
 			builder = builder.stack_size(stack_size);
 		}
+		// Increase the thread count counter
+		data.thread_count.fetch_add(1, Ordering::SeqCst);
+		// Create a new sentry watcher
+		let sentry = Sentry::new(coreid, Arc::downgrade(&data));
+		// Clone receiver
+		let receiver = data.receiver.clone();
 		// Spawn a new worker thread
 		let _ = builder.spawn(move || {
-			// Create a new sentry watcher
-			let sentry = Sentry::new(coreid, &data);
-			// Increase the thread count counter
-			data.thread_count.fetch_add(1, Ordering::SeqCst);
 			// Loop continuously, processing any jobs
 			loop {
 				// Pull a message from the job channel
-				let job = match data.receiver.recv_blocking() {
+				let job = match receiver.recv_blocking() {
 					// We received a job to process
 					Ok(job) => job,
 					// This threadpool was dropped

src/sentry.rs

@@ -1,16 +1,16 @@
 use crate::Data;
 use crate::Threadpool;
-use std::sync::Arc;
+use std::sync::Weak;
 
-pub(crate) struct Sentry<'a> {
+pub(crate) struct Sentry {
 	active: bool,
 	coreid: Option<usize>,
-	data: &'a Arc<Data>,
+	data: Weak<Data>,
 }
 
-impl<'a> Sentry<'a> {
+impl Sentry {
 	/// Create a new sentry tracker
-	pub fn new(coreid: Option<usize>, data: &'a Arc<Data>) -> Sentry<'a> {
+	pub fn new(coreid: Option<usize>, data: Weak<Data>) -> Sentry {
 		Sentry {
 			data,
 			coreid,
@@ -23,15 +23,19 @@ impl<'a> Sentry<'a> {
 	}
 }
 
-impl Drop for Sentry<'_> {
+impl Drop for Sentry {
 	fn drop(&mut self) {
+		let Some(data) = self.data.upgrade() else {
+			return;
+		};
+		
 		// If this sentry was still active,
 		// then the task panicked without
 		// properly cancelling the sentry,
 		// so we should start a new thread.
 		if self.active {
 			// Spawn another new thread
-			Threadpool::spin_up(self.coreid, self.data.clone());
+			Threadpool::spin_up(self.coreid, data.clone());
 		}
 	}
 }