utils/multithreading/
impl_native.rs

1use std::num::NonZeroUsize;
2use std::sync::atomic::AtomicUsize;
3use std::sync::atomic::Ordering::Relaxed;
4
5static NUM_THREADS: AtomicUsize = AtomicUsize::new(0);
6
7/// Get the number of worker threads to use.
8///
9/// Defaults to [`std::thread::available_parallelism`] unless set by [`set_thread_count`].
10#[must_use]
11pub fn get_thread_count() -> NonZeroUsize {
12    if let Some(threads) = NonZeroUsize::new(NUM_THREADS.load(Relaxed)) {
13        threads
14    } else {
15        let default = std::thread::available_parallelism().unwrap_or(NonZeroUsize::new(8).unwrap());
16        match NUM_THREADS.compare_exchange(0, default.get(), Relaxed, Relaxed) {
17            Ok(_) => default,
18            Err(not_zero) => NonZeroUsize::new(not_zero).unwrap(),
19        }
20    }
21}
22
23/// Set the number of worker threads to use.
24///
25/// This will affect any future call to [`get_thread_count`].
26pub fn set_thread_count(count: NonZeroUsize) {
27    NUM_THREADS.store(count.get(), Relaxed);
28}
29
30/// Run a worker function concurrently using a pool of worker threads.
31///
32/// This is a wrapper around [`std::thread::scope`] for spawning a pool of identical worker threads.
33///
34/// The number of workers is controlled by [`get_thread_count`].
35pub fn worker_pool(worker: impl Fn() + Copy + Send) {
36    let threads = get_thread_count().get();
37    if threads == 1 {
38        worker();
39    } else {
40        #[cfg(feature = "scoped-tasks")]
41        {
42            super::scoped_tasks::scope(|scope| {
43                for _ in 1..threads {
44                    scope.force_spawn(worker);
45                }
46                worker();
47            });
48        }
49
50        #[cfg(not(feature = "scoped-tasks"))]
51        {
52            std::thread::scope(move |scope| {
53                for _ in 1..threads {
54                    scope.spawn(worker);
55                }
56                worker();
57            });
58        }
59    }
60}
utils/multithreading/impl_native.rs

utils/multithreading/
impl_native.rs
