xuko_core/
thread_pool.rs

1//! Thread pool
2
3use std::{collections::VecDeque, thread};
4
5/// A task for a [`ThreadPool`]
6pub trait Worker: Send + 'static {
7    /// Execute the task associated with the [`Worker`]
8    fn exec(&mut self);
9}
10
11/// [`Worker`] is implemented for any compatible function or closure
12impl<T: Fn() + Send + Sync + 'static> Worker for T {
13    fn exec(&mut self) {
14        self()
15    }
16}
17
18/// A basic [Thread Pool](https://en.wikipedia.org/wiki/Thread_pool)
19pub struct ThreadPool {
20    pool: Vec<thread::JoinHandle<()>>,
21    queue: VecDeque<Box<dyn Worker>>,
22    pool_size: usize,
23}
24
25impl ThreadPool {
26    /// Create a new [`ThreadPool`].
27    pub fn new(size: usize) -> Self {
28        Self {
29            pool: Vec::with_capacity(size),
30            queue: VecDeque::new(),
31            pool_size: size,
32        }
33    }
34
35    /// Enqueue a [`Worker`] to the [`ThreadPool`].
36    pub fn enqueue(&mut self, worker: impl Worker) {
37        self.queue.push_back(Box::new(worker));
38    }
39
40    /// Removes threads from the pool once they're finished and spawn new ones from the queue.
41    ///
42    /// Returns `true` if the pool is empty even after emptying the queue.
43    pub fn update_threads(&mut self) -> bool {
44        for i in self
45            .pool
46            .iter()
47            .enumerate()
48            .filter(|(_, t)| t.is_finished())
49            .map(|(i, _)| i)
50            .collect::<Vec<_>>()
51        {
52            self.pool.swap_remove(i);
53        }
54
55        for _ in 0..self.queue.len() {
56            #[expect(
57                clippy::collapsible_if,
58                reason = "We don't want to take ownership of the worker until we know we can append a worker"
59            )]
60            if self.pool.len() < self.pool_size {
61                if let Some(mut w) = self.queue.pop_back() {
62                    self.pool.push(thread::spawn(move || {
63                        w.exec();
64                    }));
65                }
66            }
67        }
68
69        self.pool.is_empty()
70    }
71
72    /// Gets the capacity of the [`ThreadPool`].
73    pub const fn capacity(&self) -> usize {
74        self.pool_size
75    }
76
77    /// Returns the amount of enqueued [`Worker`]s.
78    pub fn queued_count(&self) -> usize {
79        self.queue.len()
80    }
81
82    /// Returns the amount of running threads.
83    pub const fn active_count(&self) -> usize {
84        self.pool.len()
85    }
86
87    /// Returns `true` when both the queue and pool are empty.
88    pub fn is_finished(&self) -> bool {
89        self.queued_count() == 0 && self.active_count() == 0
90    }
91
92    /// Wait for all the active threads to finish until the queue is empty
93    pub fn run_all(&mut self) {
94        while !self.update_threads() {}
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use std::time::Duration;
101
102    use super::*;
103
104    #[test]
105    fn thread_pool() {
106        let mut pool = ThreadPool::new(8);
107        for i in 0..pool.capacity() {
108            let v = (i + 1) as u64;
109            pool.enqueue(move || {
110                thread::sleep(Duration::from_secs(v));
111                println!("Task finished! {v}")
112            });
113        }
114
115        // waits for all the threads to finish
116        while !pool.update_threads() {}
117    }
118}
xuko_core/thread_pool.rs

xuko_core/
thread_pool.rs
