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)
19///
20/// # Examples
21///
22/// ```
23/// use std::thread;
24/// use std::time::Duration;
25/// use xuko_core::thread_pool::ThreadPool;
26///
27/// let mut pool = ThreadPool::new(8);
28/// for i in 0..pool.capacity() {
29///     let v = (i + 1) as u64;
30///     pool.enqueue(move || {
31///         thread::sleep(Duration::from_secs(v));
32///         println!("Task finished! {v}")
33///     });
34/// }
35/// // blocks until all threads in the pool finish
36/// pool.run_all();
37/// ```
38pub struct ThreadPool {
39    pool: Vec<thread::JoinHandle<()>>,
40    queue: VecDeque<Box<dyn Worker>>,
41    pool_size: usize,
42}
43
44impl ThreadPool {
45    /// Create a new [`ThreadPool`].
46    pub fn new(size: usize) -> Self {
47        Self {
48            pool: Vec::with_capacity(size),
49            queue: VecDeque::new(),
50            pool_size: size,
51        }
52    }
53
54    /// Enqueue a [`Worker`] to the [`ThreadPool`].
55    pub fn enqueue(&mut self, worker: impl Worker) {
56        self.queue.push_back(Box::new(worker));
57    }
58
59    /// Removes threads from the pool once they're finished and spawn new ones from the queue.
60    ///
61    /// Returns `true` if the pool is empty even after emptying the queue.
62    pub fn update_threads(&mut self) -> bool {
63        for i in self
64            .pool
65            .iter()
66            .enumerate()
67            .filter(|(_, t)| t.is_finished())
68            .map(|(i, _)| i)
69            .collect::<Vec<_>>()
70        {
71            self.pool.swap_remove(i);
72        }
73
74        for _ in 0..self.queue.len() {
75            #[expect(
76                clippy::collapsible_if,
77                reason = "We don't want to take ownership of the worker until we know we can append a worker"
78            )]
79            if self.pool.len() < self.pool_size {
80                if let Some(mut w) = self.queue.pop_back() {
81                    self.pool.push(thread::spawn(move || {
82                        w.exec();
83                    }));
84                }
85            }
86        }
87
88        self.pool.is_empty()
89    }
90
91    /// Gets the capacity of the [`ThreadPool`].
92    pub const fn capacity(&self) -> usize {
93        self.pool_size
94    }
95
96    /// Returns the amount of enqueued [`Worker`]s.
97    pub fn queued_count(&self) -> usize {
98        self.queue.len()
99    }
100
101    /// Returns the amount of running threads.
102    pub const fn active_count(&self) -> usize {
103        self.pool.len()
104    }
105
106    /// Returns `true` when both the queue and pool are empty.
107    pub fn is_finished(&self) -> bool {
108        self.queued_count() == 0 && self.active_count() == 0
109    }
110
111    /// Blocks until all the active threads to finish and the queue is empty
112    pub fn run_all(&mut self) {
113        while !self.update_threads() {}
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use std::time::Duration;
120
121    use super::*;
122
123    #[test]
124    fn thread_pool() {
125        let mut pool = ThreadPool::new(8);
126        for i in 0..pool.capacity() {
127            let v = (i + 1) as u64;
128            pool.enqueue(move || {
129                thread::sleep(Duration::from_secs(v));
130                println!("Task finished! {v}")
131            });
132        }
133
134        // waits for all the threads to finish
135        while !pool.update_threads() {}
136    }
137}
xuko_core/thread_pool.rs

xuko_core/
thread_pool.rs
