utils/
multithreading.rs

1//! Multithreading helpers.
2//!
3//! The main purpose of this module is to allow the number of worker threads used by each puzzle
4//! solution to be controlled by a CLI argument.
5
6use std::num::NonZeroUsize;
7use std::sync::atomic::AtomicUsize;
8use std::sync::atomic::Ordering::Relaxed;
9
10static NUM_THREADS: AtomicUsize = AtomicUsize::new(0);
11
12/// Get the number of worker threads to use.
13///
14/// Defaults to [`std::thread::available_parallelism`] unless set by [`set_thread_count`].
15#[must_use]
16pub fn get_thread_count() -> NonZeroUsize {
17    if let Some(threads) = NonZeroUsize::new(NUM_THREADS.load(Relaxed)) {
18        threads
19    } else {
20        let default = std::thread::available_parallelism().unwrap_or(NonZeroUsize::new(8).unwrap());
21        match NUM_THREADS.compare_exchange(0, default.get(), Relaxed, Relaxed) {
22            Ok(_) => default,
23            Err(not_zero) => NonZeroUsize::new(not_zero).unwrap(),
24        }
25    }
26}
27
28/// Set the number of worker threads to use.
29///
30/// This will affect any future call to [`get_thread_count`].
31pub fn set_thread_count(count: NonZeroUsize) {
32    NUM_THREADS.store(count.get(), Relaxed);
33}
34
35/// Run a worker function concurrently using a pool of worker threads.
36///
37/// This is a wrapper around [`std::thread::scope`] for spawning a pool of identical worker threads.
38///
39/// The number of workers is controlled by [`get_thread_count`].
40pub fn worker_pool(worker: impl Fn() + Copy + Send) {
41    let threads = get_thread_count().get();
42    if threads == 1 {
43        worker();
44    } else {
45        std::thread::scope(|scope| {
46            for _ in 0..threads {
47                scope.spawn(worker);
48            }
49        });
50    }
51}