#![allow(clippy::disallowed_types)]
use std::{
mem,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
thread,
};
use crossbeam_channel::{Receiver, Sender, TryRecvError, select, unbounded};
use crate::{pool::actor_pool::Runnable, sync::mutex::Mutex};
pub(crate) enum TaskItem {
Job(Box<dyn FnOnce() + Send + 'static>),
Actor(Arc<dyn Runnable>),
}
pub(crate) struct TaskPool {
tx: Sender<TaskItem>,
shutdown: Arc<AtomicBool>,
shutdown_tx: Mutex<Option<Sender<()>>>,
joins: Mutex<Vec<thread::JoinHandle<()>>>,
threads: usize,
}
impl TaskPool {
pub(crate) fn new(threads: usize, name_prefix: &'static str) -> Self {
let (tx, rx) = unbounded::<TaskItem>();
let (shutdown_tx, shutdown_rx) = unbounded::<()>();
let shutdown = Arc::new(AtomicBool::new(false));
let joins = (0..threads)
.map(|i| {
let rx = rx.clone();
let shutdown_rx = shutdown_rx.clone();
let shutdown = Arc::clone(&shutdown);
thread::Builder::new()
.name(format!("{name_prefix}-{i}"))
.spawn(move || task_loop(rx, shutdown_rx, shutdown))
.unwrap_or_else(|_| panic!("failed to spawn {name_prefix} worker thread"))
})
.collect();
Self {
tx,
shutdown,
shutdown_tx: Mutex::new(Some(shutdown_tx)),
joins: Mutex::new(joins),
threads,
}
}
pub(crate) fn spawn(&self, job: impl FnOnce() + Send + 'static) {
let _ = self.tx.send(TaskItem::Job(Box::new(job)));
}
pub(crate) fn injector(&self) -> Sender<TaskItem> {
self.tx.clone()
}
pub(crate) fn thread_count(&self) -> usize {
self.threads
}
pub(crate) fn shutdown(&self) {
if self.shutdown.swap(true, Ordering::AcqRel) {
return;
}
drop(self.shutdown_tx.lock().take());
let joins = mem::take(&mut *self.joins.lock());
let current = thread::current().id();
for handle in joins {
if handle.thread().id() != current {
let _ = handle.join();
}
}
}
}
fn task_loop(rx: Receiver<TaskItem>, shutdown_rx: Receiver<()>, shutdown: Arc<AtomicBool>) {
loop {
match rx.try_recv() {
Ok(item) => run_guarded(item),
Err(TryRecvError::Empty) => {
if shutdown.load(Ordering::Acquire) {
return;
}
select! {
recv(rx) -> item => match item {
Ok(item) => run_guarded(item),
Err(_) => return,
},
recv(shutdown_rx) -> _ => return,
}
}
Err(TryRecvError::Disconnected) => return,
}
}
}
fn run_guarded(item: TaskItem) {
match item {
TaskItem::Job(job) => job(),
TaskItem::Actor(actor) => actor.run(),
}
}