1use std::num::NonZeroUsize;
7use std::sync::atomic::AtomicUsize;
8use std::sync::atomic::Ordering::Relaxed;
9
10static NUM_THREADS: AtomicUsize = AtomicUsize::new(0);
11
12#[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
28pub fn set_thread_count(count: NonZeroUsize) {
32 NUM_THREADS.store(count.get(), Relaxed);
33}
34
35pub 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}