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, AtomicUsize, Ordering};
use std::sync::mpsc::Sender;
use std::sync::{Arc, Mutex, mpsc};
use std::thread;
use std::time::Duration;

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

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)),
                max_jobs: 50_000,
                job_count: Arc::new(AtomicUsize::new(0)),
            }
        } 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));
            let job_count = Arc::new(AtomicUsize::new(0));

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

            Self {
                workers,
                sender: Some(sender),
                num_threads: pool_size,
                kill_signal,
                max_jobs: 50_000,
                job_count,
            }
        }
    }

    /// Sets the maximum number of jobs that can be queued at once.
    ///
    /// When the queue is full, `execute()` will block with exponential backoff
    /// (1ms → 50ms ceiling) until a slot becomes available.
    ///
    /// # Arguments
    /// * `limit` - Maximum number of jobs (must be > 0)
    ///
    /// # Panics
    /// Panics if `limit` is 0.
    ///
    /// # Example
    /// ```
    /// use jlizard_simple_threadpool::threadpool::ThreadPool;
    ///
    /// let pool = ThreadPool::new(4).max_jobs(100_000);
    /// ```
    pub fn max_jobs(mut self, limit: usize) -> Self {
        if limit == 0 {
            panic!("max_jobs must be > 0");
        }
        self.max_jobs = limit;
        self
    }

    /// 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
    ///
    /// # Backpressure
    /// When the job queue reaches `max_jobs` capacity, this method blocks with exponential
    /// backoff (starting at 1ms, doubling up to 50ms ceiling) until a slot becomes available.
    ///
    /// # Errors
    /// Returns an error if the worker threads have been dropped or the channel is closed.
    pub fn execute<F>(&self, f: F) -> Result<(), Box<dyn Error>>
    where
        F: FnOnce() + Send + 'static,
    {
        if self.is_single_threaded() {
            f();
            Ok(())
        } else {
            let mut backoff_ms = 1u64;
            const MAX_BACKOFF_MS: u64 = 50;

            loop {
                let current = self.job_count.load(Ordering::Relaxed);

                if current < self.max_jobs {
                    self.job_count.fetch_add(1, Ordering::Relaxed);

                    match self.sender.as_ref().unwrap().send(Box::new(f)) {
                        Ok(_) => return Ok(()),
                        Err(e) => {
                            self.job_count.fetch_sub(1, Ordering::Relaxed);
                            return Err(e.into());
                        }
                    }
                } else {
                    thread::sleep(Duration::from_millis(backoff_ms));
                    backoff_ms = (backoff_ms * 2).min(MAX_BACKOFF_MS);
                }
            }
        }
    }

    /// 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);
    }

    // ===== Backpressure & Queue Limit Tests =====

    #[test]
    fn test_builder_pattern_api() {
        // Test method chaining
        let pool = ThreadPool::new(4).max_jobs(1000);
        assert_eq!(pool.num_threads, 4);
        assert_eq!(pool.max_jobs, 1000);
        assert!(!pool.is_single_threaded());

        // Test default still works
        let pool_default = ThreadPool::new(2);
        assert_eq!(pool_default.max_jobs, 50_000);
    }

    #[test]
    fn test_backward_compatibility() {
        // Existing API should work unchanged
        let pool = ThreadPool::new(3);
        let counter = Arc::new(Mutex::new(0));
        let counter_clone = Arc::clone(&counter);

        pool.execute(move || {
            *counter_clone.lock().unwrap() += 1;
        })
        .expect("Failed to execute");

        drop(pool);
        assert_eq!(*counter.lock().unwrap(), 1);
    }

    #[test]
    fn test_single_threaded_ignores_max_jobs() {
        // Single-threaded mode should bypass queue limits entirely
        let pool = ThreadPool::new(1).max_jobs(5);
        assert!(pool.is_single_threaded());

        let counter = Arc::new(Mutex::new(0));

        // Execute 100 jobs - should all complete immediately without blocking
        for _ in 0..100 {
            let counter_clone = Arc::clone(&counter);
            pool.execute(move || {
                *counter_clone.lock().unwrap() += 1;
            })
            .expect("Failed to execute");
        }

        // All jobs execute synchronously
        assert_eq!(*counter.lock().unwrap(), 100);
    }

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

        let pool = ThreadPool::new(2).max_jobs(100);
        let completed = Arc::new(Mutex::new(0));

        // Submit 10 jobs that take 50ms each
        for _ in 0..10 {
            let completed_clone = Arc::clone(&completed);
            pool.execute(move || {
                std::thread::sleep(Duration::from_millis(50));
                *completed_clone.lock().unwrap() += 1;
            })
            .expect("Failed to execute");
        }

        // Job count should be > 0 while jobs are running
        std::thread::sleep(Duration::from_millis(20));
        let count_while_running = pool.job_count.load(Ordering::Relaxed);
        assert!(count_while_running > 0);

        // Wait for all jobs to complete
        drop(pool);

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

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

        // Pool with 2 workers, max 5 jobs
        let pool = Arc::new(ThreadPool::new(2).max_jobs(5));
        let executed = Arc::new(Mutex::new(Vec::new()));

        // Submit 5 slow jobs to fill the queue
        for i in 0..5 {
            let executed_clone = Arc::clone(&executed);
            pool.execute(move || {
                std::thread::sleep(Duration::from_millis(100));
                executed_clone.lock().unwrap().push(i);
            })
            .expect("Failed to execute");
        }

        // Give jobs time to start processing
        std::thread::sleep(Duration::from_millis(20));

        // Queue should now be full (5 jobs queued/running)
        let count = pool.job_count.load(Ordering::Relaxed);
        assert_eq!(count, 5);

        // Submitting 6th job should block until a slot opens
        let pool_clone = Arc::clone(&pool);
        let executed_clone = Arc::clone(&executed);

        std::thread::spawn(move || {
            pool_clone
                .execute(move || {
                    executed_clone.lock().unwrap().push(99);
                })
                .expect("Failed to execute");
        });

        // Wait for all jobs to complete
        std::thread::sleep(Duration::from_millis(300));
        drop(pool);

        // All 6 jobs should have completed
        let final_executed = executed.lock().unwrap();
        assert_eq!(final_executed.len(), 6);
        assert!(final_executed.contains(&99));
    }

    #[test]
    fn test_exponential_backoff_timing() {

        // Create pool with 1 worker and max 2 jobs to easily saturate
        let pool = Arc::new(ThreadPool::new(1).max_jobs(2));

        // Fill queue with 2 slow jobs
        for _ in 0..2 {
            pool.execute(|| {
                std::thread::sleep(Duration::from_millis(200));
            })
            .expect("Failed to execute");
        }

        // Give first job time to start
        std::thread::sleep(Duration::from_millis(20));

        // Next execute() should block with backoff
        let pool_clone = Arc::clone(&pool);

        std::thread::spawn(move || {
            pool_clone
                .execute(|| {
                    // Quick job
                })
                .expect("Failed to execute");
        });

        // Wait for completion
        std::thread::sleep(Duration::from_millis(500));
        drop(pool);

        // Test passes if we don't hang indefinitely
    }

    #[test]
    fn test_queue_recovers_after_full() {

        let pool = ThreadPool::new(2).max_jobs(10);
        let completed = Arc::new(Mutex::new(0));

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

        // Wait for some jobs to complete
        std::thread::sleep(Duration::from_millis(150));

        // Should be able to submit more jobs now
        for _ in 0..5 {
            let completed_clone = Arc::clone(&completed);
            pool.execute(move || {
                *completed_clone.lock().unwrap() += 1;
            })
            .expect("Failed to execute");
        }

        drop(pool);

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

    #[test]
    fn test_concurrent_submissions_near_limit() {

        let pool = Arc::new(ThreadPool::new(4).max_jobs(20));

        // Spawn multiple threads submitting jobs concurrently
        let mut handles = vec![];
        for _ in 0..5 {
            let pool_clone = Arc::clone(&pool);

            let handle = std::thread::spawn(move || {
                for _ in 0..10 {
                    pool_clone
                        .execute(move || {
                            std::thread::sleep(Duration::from_millis(10));
                        })
                        .expect("Failed to execute");
                }
            });
            handles.push(handle);
        }

        // Wait for all submission threads
        for handle in handles {
            handle.join().unwrap();
        }

        drop(pool);

        // All 50 jobs should complete without races
        // (Can't check completed count easily due to closure move semantics,
        // but test passes if no panics/deadlocks occur)
    }

    #[test]
    #[should_panic(expected = "max_jobs must be > 0")]
    fn test_max_jobs_zero_panics() {
        let _pool = ThreadPool::new(4).max_jobs(0);
    }
}