jlizard_simple_threadpool/

threadpool.rs

//! Thread pool implementation for concurrent PDF generation.
//!
//! Supports single-threaded mode (`max_jobs=1`) for debugging and sequential processing,
//! or multi-threaded mode for parallel processing of multiple BOM files.
//!
//! When `max_jobs` is 0 or not set, the pool uses all available CPU cores for maximum parallelism.

use crate::common::Job;
use crate::worker::Worker;
use std::error::Error;

#[cfg(feature = "log")]
use log::debug;

use std::fmt::{Display, Formatter};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::Sender;
use std::sync::{Arc, Mutex, mpsc};
use std::thread;

pub struct ThreadPool {
    workers: Vec<Worker>,
    sender: Option<Sender<Job>>,
    num_threads: u8,
    kill_signal: Arc<AtomicBool>,
}

impl ThreadPool {
    /// Creates a new thread pool with the following behavior constraints:
    /// - pool_size is `0`: runs in multithreaded default mode using maximum parallelism
    /// - pool_size is `1`: runs in single-threaded mode (all jobs are run in the main thread)
    /// - pool_size is `1<N<=255`: runs in multithreaded mode with `N` jobs
    pub fn new(pool_size: u8) -> Self {
        if pool_size == 0 {
            Self::default()
        } else if pool_size == 1 {
            Self {
                workers: Vec::new(),
                sender: None,
                num_threads: pool_size,
                kill_signal: Arc::new(AtomicBool::new(false)),
            }
        } else {
            let (sender, receiver) = mpsc::channel::<Job>();

            let mut workers = Vec::with_capacity(pool_size as usize);

            let receiver = Arc::new(Mutex::new(receiver));
            let kill_signal = Arc::new(AtomicBool::new(false));

            for id in 1..=pool_size {
                workers.push(Worker::new(
                    id,
                    Arc::clone(&receiver),
                    Arc::clone(&kill_signal),
                ));
            }

            Self {
                workers,
                sender: Some(sender),
                num_threads: pool_size,
                kill_signal,
            }
        }
    }

    /// Executes a job on the thread pool.
    ///
    /// # Behavior
    /// - **Single-threaded mode** (`max_jobs=1`): Job executes synchronously in the calling thread
    /// - **Multi-threaded mode**: Job is queued and executed asynchronously by worker threads
    pub fn execute<F>(&self, f: F) -> Result<(), Box<dyn Error>>
    where
        F: FnOnce() + Send + 'static,
    {
        if self.is_single_threaded() {
            f();
            Ok(())
        } else {
            self.sender
                .as_ref()
                .unwrap()
                .send(Box::new(f))
                .map_err(|e| e.into())
        }
    }

    /// Returns `true` if running in single-threaded mode.
    ///
    /// Single-threaded mode is active when `max_jobs=1`, resulting in:
    /// - No worker threads spawned
    /// - No message passing channel created
    /// - All jobs executed synchronously in the main thread
    pub fn is_single_threaded(&self) -> bool {
        self.sender.is_none() && self.workers.is_empty()
    }

    /// Signals all worker threads to stop processing after completing their current jobs.
    ///
    /// This method sets the kill signal which workers check periodically.
    /// Workers will complete their current job before stopping.
    /// The pool's Drop implementation will wait for all workers to finish.
    pub fn signal_stop(&self) {
        self.kill_signal.store(true, Ordering::Relaxed);
    }

    /// Returns a clone of the kill signal Arc that can be passed into jobs.
    ///
    /// This allows jobs to signal the thread pool to stop from within the job itself.
    /// Useful for scenarios like finding a hash collision where one worker needs to
    /// stop all other workers.
    ///
    /// # Example
    /// ```no_run
    /// use jlizard_simple_threadpool::threadpool::ThreadPool;
    /// use std::sync::atomic::Ordering;
    ///
    /// let pool = ThreadPool::new(4);
    /// let kill_signal = pool.get_kill_signal();
    ///
    /// pool.execute(move || {
    ///     // Do some work...
    ///     if /* condition met */ true {
    ///         // Signal all workers to stop
    ///         kill_signal.store(true, Ordering::Relaxed);
    ///     }
    /// }).expect("Failed to execute");
    /// ```
    pub fn get_kill_signal(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.kill_signal)
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        // Drop the sender first which causes receivers to error out gracefully
        // This allows pending jobs in the queue to complete
        drop(self.sender.take());

        #[cfg(feature = "log")]
        {
            debug!("Waiting for workers to finish");
        }

        // Workers will exit naturally when:
        // 1. The channel is closed (sender dropped) and queue is empty, OR
        // 2. The kill_signal was set (via signal_stop())
        for worker in &mut self.workers {
            #[cfg(feature = "log")]
            {
                debug!("Shutting down worker {}", worker.id);
            }
            worker.thread.take().unwrap().join().unwrap();
        }

        #[cfg(feature = "log")]
        {
            debug!("All workers stopped");
        }
    }
}

impl Default for ThreadPool {
    fn default() -> Self {
        let max_threads = thread::available_parallelism().map(|e| e.get()).expect("Unable to find any threads to run with. Possible system-side restrictions or limitations");

        // saturate to u8::MAX if number of threads is larger than what u8 can hold
        ThreadPool::new(u8::try_from(max_threads).unwrap_or(u8::MAX))
    }
}

impl Display for ThreadPool {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        if self.is_single_threaded() {
            write!(
                f,
                "Concurrency Disabled: running all jobs sequentially in main thread. A user override forced this through an VEX2PDF_MAX_JOBS or the --max-jobs cli argument"
            )
        } else {
            write!(
                f,
                "Concurrency Enabled: running with {} jobs",
                self.num_threads
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Mutex};
    use std::time::Duration;

    #[test]
    fn test_threadpool_creation_modes() {
        // Test pool with size 0 (default - max parallelism)
        let pool_default = ThreadPool::new(0);
        assert!(pool_default.num_threads > 0);
        assert!(!pool_default.is_single_threaded());

        // Test pool with size 1 (single-threaded)
        let pool_single = ThreadPool::new(1);
        assert_eq!(pool_single.num_threads, 1);
        assert!(pool_single.is_single_threaded());
        assert!(pool_single.workers.is_empty());
        assert!(pool_single.sender.is_none());

        // Test pool with size 4 (multi-threaded)
        let pool_multi = ThreadPool::new(4);
        assert_eq!(pool_multi.num_threads, 4);
        assert!(!pool_multi.is_single_threaded());
        assert_eq!(pool_multi.workers.len(), 4);
        assert!(pool_multi.sender.is_some());
    }

    #[test]
    fn test_single_threaded_execution() {
        let pool = ThreadPool::new(1);
        let counter = Arc::new(Mutex::new(0));
        let counter_clone = Arc::clone(&counter);

        // Execute job synchronously
        pool.execute(move || {
            let mut num = counter_clone.lock().unwrap();
            *num += 1;
        })
        .expect("Failed to execute job");

        // In single-threaded mode, job executes immediately
        let value = *counter.lock().unwrap();
        assert_eq!(value, 1);
    }

    #[test]
    fn test_multi_threaded_execution() {
        let pool = ThreadPool::new(2);
        let results = Arc::new(Mutex::new(Vec::new()));

        // Execute multiple jobs
        for i in 0..5 {
            let results_clone = Arc::clone(&results);
            pool.execute(move || {
                std::thread::sleep(Duration::from_millis(10));
                results_clone.lock().unwrap().push(i);
            })
            .expect("Failed to execute job");
        }

        // Drop pool to wait for all jobs to complete
        drop(pool);

        // Verify all jobs completed
        let final_results = results.lock().unwrap();
        assert_eq!(final_results.len(), 5);
        // Results may be in any order due to concurrency
        for i in 0..5 {
            assert!(final_results.contains(&i));
        }
    }

    #[test]
    fn test_get_num_threads() {
        let pool1 = ThreadPool::new(1);
        assert_eq!(pool1.num_threads, 1);

        let pool4 = ThreadPool::new(4);
        assert_eq!(pool4.num_threads, 4);

        let pool_default = ThreadPool::default();
        assert!(pool_default.num_threads > 0);
    }

    #[test]
    fn test_is_single_threaded() {
        let pool_single = ThreadPool::new(1);
        assert!(pool_single.is_single_threaded());

        let pool_multi = ThreadPool::new(2);
        assert!(!pool_multi.is_single_threaded());

        let pool_default = ThreadPool::default();
        assert!(!pool_default.is_single_threaded());
    }

    #[test]
    fn test_pool_graceful_shutdown() {
        let pool = ThreadPool::new(3);
        let completed = Arc::new(Mutex::new(0));

        // Execute several jobs
        for _ in 0..10 {
            let completed_clone = Arc::clone(&completed);
            pool.execute(move || {
                std::thread::sleep(Duration::from_millis(20));
                *completed_clone.lock().unwrap() += 1;
            })
            .expect("Failed to execute job");
        }

        // Drop pool - should wait for all jobs to complete
        drop(pool);

        // All jobs should have completed
        assert_eq!(*completed.lock().unwrap(), 10);
    }

    #[test]
    fn test_signal_stop_method() {
        let pool = ThreadPool::new(4);
        let completed = Arc::new(Mutex::new(0));

        // Execute several quick jobs
        for _ in 0..5 {
            let completed_clone = Arc::clone(&completed);
            pool.execute(move || {
                std::thread::sleep(Duration::from_millis(10));
                *completed_clone.lock().unwrap() += 1;
            })
            .expect("Failed to execute job");
        }

        // Signal stop
        pool.signal_stop();

        // Drop pool to wait for shutdown
        drop(pool);

        // At least some jobs should have completed before stop
        let count = *completed.lock().unwrap();
        assert!(count >= 1 && count <= 5);
    }

    #[test]
    fn test_get_kill_signal() {
        let pool = ThreadPool::new(2);
        let kill_signal = pool.get_kill_signal();

        // Verify we can read the signal
        assert!(!kill_signal.load(std::sync::atomic::Ordering::Relaxed));

        // Signal stop through the cloned kill signal
        kill_signal.store(true, std::sync::atomic::Ordering::Relaxed);

        // Pool should also see it
        drop(pool);
    }

    #[test]
    fn test_job_signals_stop_to_other_workers() {
        use std::sync::atomic::Ordering;

        let pool = Arc::new(ThreadPool::new(4));
        let completed = Arc::new(Mutex::new(Vec::new()));
        let collision_found = Arc::new(AtomicBool::new(false));

        // Submit jobs where job 2 will "find a collision" early
        for i in 0..20 {
            let pool_clone = Arc::clone(&pool);
            let completed_clone = Arc::clone(&completed);
            let collision_found_clone = Arc::clone(&collision_found);
            let kill_signal = pool.get_kill_signal();

            pool.execute(move || {
                // Check if already stopped before starting work
                if kill_signal.load(Ordering::Relaxed) {
                    return;
                }

                std::thread::sleep(Duration::from_millis(10));

                // Job 2 simulates finding a collision (early job to ensure it runs)
                if i == 2 {
                    collision_found_clone.store(true, Ordering::Relaxed);
                    pool_clone.signal_stop();
                    completed_clone.lock().unwrap().push(i);
                } else {
                    // Only complete if collision not yet found
                    if !collision_found_clone.load(Ordering::Relaxed) {
                        completed_clone.lock().unwrap().push(i);
                    }
                }
            })
            .expect("Failed to execute job");
        }

        // Give jobs time to execute
        std::thread::sleep(Duration::from_millis(150));

        // Wait for pool to finish
        drop(pool);

        // Verify collision was found
        assert!(collision_found.load(Ordering::Relaxed));

        // Not all jobs should have completed (some were stopped)
        let completed_jobs = completed.lock().unwrap();
        assert!(completed_jobs.len() < 20);
        assert!(completed_jobs.contains(&2)); // The collision job completed
    }

    #[test]
    fn test_workers_complete_current_job_before_stopping() {
        use std::sync::atomic::Ordering;

        let pool = ThreadPool::new(2);
        let job_started = Arc::new(AtomicBool::new(false));
        let job_completed = Arc::new(AtomicBool::new(false));

        let job_started_clone = Arc::clone(&job_started);
        let job_completed_clone = Arc::clone(&job_completed);

        // Start a long job
        pool.execute(move || {
            job_started_clone.store(true, Ordering::Relaxed);
            std::thread::sleep(Duration::from_millis(100));
            job_completed_clone.store(true, Ordering::Relaxed);
        })
        .expect("Failed to execute job");

        // Wait for job to start
        std::thread::sleep(Duration::from_millis(50));
        assert!(job_started.load(Ordering::Relaxed));

        // Signal stop while job is running
        pool.signal_stop();

        // Drop pool to wait for completion
        drop(pool);

        // Job should have completed before worker stopped
        assert!(job_completed.load(Ordering::Relaxed));
    }

    #[test]
    fn test_no_new_jobs_after_signal_stop() {
        use std::sync::atomic::Ordering;

        let pool = ThreadPool::new(3);
        let executed = Arc::new(AtomicBool::new(false));
        let executed_clone = Arc::clone(&executed);

        // Signal stop immediately
        pool.signal_stop();

        // Try to execute a job
        pool.execute(move || {
            executed_clone.store(true, Ordering::Relaxed);
        })
        .expect("Failed to execute job");

        // Give some time for potential execution
        std::thread::sleep(Duration::from_millis(200));

        // Job might or might not execute depending on timing
        // This is expected behavior - jobs in queue may still execute
        // The important thing is workers stop checking for new jobs

        drop(pool);
        // Test passes if we don't hang
    }

    #[test]
    fn test_kill_signal_in_single_threaded_mode() {
        let pool = ThreadPool::new(1);
        assert!(pool.is_single_threaded());

        // Get kill signal - should exist even in single-threaded mode
        let kill_signal = pool.get_kill_signal();
        assert!(!kill_signal.load(std::sync::atomic::Ordering::Relaxed));

        // Signal stop
        pool.signal_stop();
        assert!(kill_signal.load(std::sync::atomic::Ordering::Relaxed));

        // Drop should work fine
        drop(pool);
    }
}