use std::sync::OnceLock;
use rayon::prelude::*;
pub struct JobPool {
pool: rayon::ThreadPool,
}
impl JobPool {
fn build() -> JobPool {
Self::with_threads(
CONFIGURED_THREADS
.get()
.copied()
.unwrap_or_else(default_threads),
)
}
fn with_threads(threads: usize) -> JobPool {
let threads = threads.max(1);
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.thread_name(|i| format!("cn-job-{i}"))
.build()
.expect("failed to build job thread pool");
tracing::info!("JobPool: {threads} worker thread(s)");
JobPool { pool }
}
pub fn thread_count(&self) -> usize {
self.pool.current_num_threads()
}
pub fn parallel_for<T, F>(&self, items: &mut [T], f: F)
where
T: Send,
F: Fn(&mut T) + Send + Sync,
{
if items.len() < 2 {
items.iter_mut().for_each(f);
return;
}
self.pool.install(|| items.par_iter_mut().for_each(f));
}
pub fn install<R, F>(&self, f: F) -> R
where
F: FnOnce() -> R + Send,
R: Send,
{
self.pool.install(f)
}
}
static CONFIGURED_THREADS: OnceLock<usize> = OnceLock::new();
fn default_threads() -> usize {
std::thread::available_parallelism()
.map(|n| n.get().saturating_sub(1).max(1))
.unwrap_or(1)
}
pub fn configure(threads: usize) {
let _ = CONFIGURED_THREADS.set(threads.max(1));
}
pub fn pool() -> &'static JobPool {
static POOL: OnceLock<JobPool> = OnceLock::new();
POOL.get_or_init(JobPool::build)
}
pub fn serial_pool() -> &'static JobPool {
static POOL: OnceLock<JobPool> = OnceLock::new();
POOL.get_or_init(|| JobPool::with_threads(1))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pool_is_a_singleton() {
assert!(std::ptr::eq(pool(), pool()));
}
#[test]
fn with_threads_sets_the_worker_count() {
assert_eq!(JobPool::with_threads(3).thread_count(), 3);
assert_eq!(JobPool::with_threads(0).thread_count(), 1);
}
#[test]
fn default_threads_is_at_least_one() {
assert!(default_threads() >= 1);
}
#[test]
fn parallel_for_visits_every_item() {
let mut data: Vec<u32> = (0..10_000).collect();
pool().parallel_for(&mut data, |x| *x += 1);
assert!(data.iter().enumerate().all(|(i, &x)| x == i as u32 + 1));
}
#[test]
fn parallel_for_handles_empty_and_single() {
let mut empty: Vec<u32> = Vec::new();
pool().parallel_for(&mut empty, |x| *x += 1);
assert!(empty.is_empty());
let mut single = vec![41u32];
pool().parallel_for(&mut single, |x| *x += 1);
assert_eq!(single, vec![42]);
}
}