use std::sync::{Arc, Mutex};
pub struct JobPool {
workers: Vec<std::thread::JoinHandle<()>>,
sender: Option<crossbeam_channel::Sender<Job>>,
worker_count: usize,
}
type Job = Box<dyn FnOnce() + Send + 'static>;
impl JobPool {
pub fn new(num_threads: usize) -> Self {
let num_threads = num_threads.max(1);
let (sender, receiver) = crossbeam_channel::unbounded::<Job>();
let mut workers = Vec::with_capacity(num_threads);
for i in 0..num_threads {
let rx = receiver.clone();
let handle = std::thread::Builder::new()
.name(format!("kiran-worker-{i}"))
.spawn(move || {
while let Ok(job) = rx.recv() {
job();
}
})
.expect("failed to spawn worker thread");
workers.push(handle);
}
Self {
workers,
sender: Some(sender),
worker_count: num_threads,
}
}
pub fn auto() -> Self {
let cores = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(2);
Self::new((cores - 1).max(1))
}
pub fn submit(&self, job: impl FnOnce() + Send + 'static) {
if let Some(ref sender) = self.sender {
let _ = sender.send(Box::new(job));
}
}
pub fn scope<F>(&self, jobs: Vec<F>)
where
F: FnOnce() + Send + 'static,
{
if jobs.is_empty() {
return;
}
let remaining = Arc::new(Mutex::new(jobs.len()));
let condvar = Arc::new(std::sync::Condvar::new());
for job in jobs {
let remaining = Arc::clone(&remaining);
let condvar = Arc::clone(&condvar);
self.submit(move || {
job();
let mut count = remaining.lock().unwrap();
*count -= 1;
if *count == 0 {
condvar.notify_one();
}
});
}
let mut count = remaining.lock().unwrap();
while *count > 0 {
count = condvar.wait(count).unwrap();
}
}
#[must_use]
#[inline]
pub fn worker_count(&self) -> usize {
self.worker_count
}
}
impl Drop for JobPool {
fn drop(&mut self) {
self.sender.take();
for worker in self.workers.drain(..) {
let _ = worker.join();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicU32, Ordering};
#[test]
fn pool_creates_workers() {
let pool = JobPool::new(4);
assert_eq!(pool.worker_count(), 4);
}
#[test]
fn pool_auto() {
let pool = JobPool::auto();
assert!(pool.worker_count() >= 1);
}
#[test]
fn pool_submit_runs() {
let pool = JobPool::new(2);
let counter = Arc::new(AtomicU32::new(0));
let c = counter.clone();
pool.submit(move || {
c.fetch_add(1, Ordering::SeqCst);
});
std::thread::sleep(std::time::Duration::from_millis(50));
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
#[test]
fn pool_scope_waits() {
let pool = JobPool::new(4);
let counter = Arc::new(AtomicU32::new(0));
let jobs: Vec<_> = (0..10)
.map(|_| {
let c = counter.clone();
move || {
c.fetch_add(1, Ordering::SeqCst);
}
})
.collect();
pool.scope(jobs);
assert_eq!(counter.load(Ordering::SeqCst), 10);
}
#[test]
fn pool_scope_empty() {
let pool = JobPool::new(2);
pool.scope::<Box<dyn FnOnce() + Send>>(vec![]); }
#[test]
fn pool_scope_single_job() {
let pool = JobPool::new(2);
let counter = Arc::new(AtomicU32::new(0));
let c = counter.clone();
pool.scope(vec![move || {
c.fetch_add(42, Ordering::SeqCst);
}]);
assert_eq!(counter.load(Ordering::SeqCst), 42);
}
#[test]
fn pool_many_jobs() {
let pool = JobPool::new(4);
let counter = Arc::new(AtomicU32::new(0));
let jobs: Vec<_> = (0..1000)
.map(|_| {
let c = counter.clone();
move || {
c.fetch_add(1, Ordering::SeqCst);
}
})
.collect();
pool.scope(jobs);
assert_eq!(counter.load(Ordering::SeqCst), 1000);
}
#[test]
fn pool_drop_cleans_up() {
let pool = JobPool::new(2);
drop(pool); }
#[test]
fn pool_as_world_resource() {
let mut world = crate::World::new();
world.insert_resource(JobPool::new(2));
let pool = world.get_resource::<JobPool>().unwrap();
assert_eq!(pool.worker_count(), 2);
}
}