xuko-core 0.6.1

//! Thread pool

use std::{collections::VecDeque, thread};

/// A task for a [`ThreadPool`]
pub trait Worker: Send + 'static {
    /// Execute the task associated with the [`Worker`]
    fn exec(&mut self);
}

/// [`Worker`] is implemented for any compatible function or closure
impl<T: Fn() + Send + Sync + 'static> Worker for T {
    fn exec(&mut self) {
        self()
    }
}

/// A basic [Thread Pool](https://en.wikipedia.org/wiki/Thread_pool)
///
/// # Examples
///
/// ```
/// use std::thread;
/// use std::time::Duration;
/// use xuko_core::thread_pool::ThreadPool;
///
/// let mut pool = ThreadPool::new(8);
/// for i in 0..pool.capacity() {
///     let v = (i + 1) as u64;
///     pool.enqueue(move || {
///         thread::sleep(Duration::from_secs(v));
///         println!("Task finished! {v}")
///     });
/// }
/// // blocks until all threads in the pool finish
/// pool.run_all();
/// ```
pub struct ThreadPool {
    pool: Vec<thread::JoinHandle<()>>,
    queue: VecDeque<Box<dyn Worker>>,
    pool_size: usize,
}

impl ThreadPool {
    /// Create a new [`ThreadPool`].
    pub fn new(size: usize) -> Self {
        Self {
            pool: Vec::with_capacity(size),
            queue: VecDeque::new(),
            pool_size: size,
        }
    }

    /// Enqueue a [`Worker`] to the [`ThreadPool`].
    pub fn enqueue(&mut self, worker: impl Worker) {
        self.queue.push_back(Box::new(worker));
    }

    /// Removes threads from the pool once they're finished and spawn new ones from the queue.
    ///
    /// Returns `true` if the pool is empty even after emptying the queue.
    pub fn update_threads(&mut self) -> bool {
        for i in self
            .pool
            .iter()
            .enumerate()
            .filter(|(_, t)| t.is_finished())
            .map(|(i, _)| i)
            .collect::<Vec<_>>()
        {
            self.pool.swap_remove(i);
        }

        for _ in 0..self.queue.len() {
            #[expect(
                clippy::collapsible_if,
                reason = "We don't want to take ownership of the worker until we know we can append a worker"
            )]
            if self.pool.len() < self.pool_size {
                if let Some(mut w) = self.queue.pop_back() {
                    self.pool.push(thread::spawn(move || {
                        w.exec();
                    }));
                }
            }
        }

        self.pool.is_empty()
    }

    /// Gets the capacity of the [`ThreadPool`].
    pub const fn capacity(&self) -> usize {
        self.pool_size
    }

    /// Returns the amount of enqueued [`Worker`]s.
    pub fn queued_count(&self) -> usize {
        self.queue.len()
    }

    /// Returns the amount of running threads.
    pub const fn active_count(&self) -> usize {
        self.pool.len()
    }

    /// Returns `true` when both the queue and pool are empty.
    pub fn is_finished(&self) -> bool {
        self.queued_count() == 0 && self.active_count() == 0
    }

    /// Blocks until all the active threads to finish and the queue is empty
    pub fn run_all(&mut self) {
        while !self.update_threads() {}
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::*;

    #[test]
    fn thread_pool() {
        let mut pool = ThreadPool::new(8);
        for i in 0..pool.capacity() {
            let v = (i + 1) as u64;
            pool.enqueue(move || {
                thread::sleep(Duration::from_secs(v));
                println!("Task finished! {v}")
            });
        }

        // waits for all the threads to finish
        while !pool.update_threads() {}
    }
}