use std::collections::VecDeque;
use std::io;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::mpsc::{Receiver, Sender, SyncSender, channel, sync_channel};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use crate::runtime::log::{ErrlogSevEnum, errlog_sev_printf};
use crate::runtime::task::{InheritedRuntime, StackSizeClass, ThreadPriority, enter_ioc_thread};
#[derive(Clone, Copy, Debug)]
pub struct WorkerRole {
pub suffix: &'static str,
pub stack: StackSizeClass,
pub priority: ThreadPriority,
}
enum Assignment {
Joinable {
body: Box<dyn FnOnce() + Send + 'static>,
ambient: InheritedRuntime,
done: SyncSender<thread::Result<()>>,
},
Detached {
body: Box<dyn FnOnce() + Send + 'static>,
ambient: InheritedRuntime,
label: String,
},
Stop,
}
struct SetState {
leased: bool,
running: usize,
parked: bool,
}
impl SetState {
fn became_free(&mut self) -> bool {
if !self.leased && self.running == 0 && !self.parked {
self.parked = true;
true
} else {
false
}
}
}
struct SetHandle {
senders: Vec<Sender<Assignment>>,
state: Mutex<SetState>,
}
struct Registry {
idle: VecDeque<Arc<SetHandle>>,
all: Vec<Arc<SetHandle>>,
created: usize,
joins: Vec<JoinHandle<()>>,
stopping: bool,
}
struct PoolInner {
roster: Box<[WorkerRole]>,
name_prefix: &'static str,
capacity: usize,
reg: Mutex<Registry>,
}
impl PoolInner {
fn lock(&self) -> std::sync::MutexGuard<'_, Registry> {
self.reg.lock().unwrap_or_else(|e| e.into_inner())
}
}
fn lock_set(set: &SetHandle) -> std::sync::MutexGuard<'_, SetState> {
set.state.lock().unwrap_or_else(|e| e.into_inner())
}
fn free_if_idle(inner: &Arc<PoolInner>, set: &Arc<SetHandle>, freed: bool) {
if !freed {
return;
}
let mut reg = inner.lock();
if reg.stopping {
return;
}
reg.idle.push_back(set.clone());
}
pub struct SetLease {
inner: Arc<PoolInner>,
set: Arc<SetHandle>,
}
impl Drop for SetLease {
fn drop(&mut self) {
let freed = {
let mut st = lock_set(&self.set);
st.leased = false;
st.became_free()
};
free_if_idle(&self.inner, &self.set, freed);
}
}
pub struct Worker {
set: Arc<SetHandle>,
tx: Sender<Assignment>,
}
impl Worker {
fn charge(&self) {
lock_set(&self.set).running += 1;
}
pub fn run<F>(self, body: F) -> Job
where
F: FnOnce() + Send + 'static,
{
let (done, done_rx) = sync_channel(1);
self.charge();
let _ = self.tx.send(Assignment::Joinable {
body: Box::new(body),
ambient: InheritedRuntime::capture(),
done,
});
Job { done: done_rx }
}
pub fn run_detached<F>(self, label: String, body: F)
where
F: FnOnce() + Send + 'static,
{
self.charge();
let _ = self.tx.send(Assignment::Detached {
body: Box::new(body),
ambient: InheritedRuntime::capture(),
label,
});
}
}
pub struct Job {
done: Receiver<thread::Result<()>>,
}
impl Job {
pub fn join(self) -> thread::Result<()> {
self.done.recv().unwrap_or(Ok(()))
}
}
fn announce_panic(label: &str) {
errlog_sev_printf(
ErrlogSevEnum::Major,
&format!(
"{label}: the connection thread panicked; this connection is being \
torn down. Other connections are unaffected."
),
);
}
fn worker_loop(inner: Arc<PoolInner>, set: Arc<SetHandle>, rx: Receiver<Assignment>) {
while let Ok(assignment) = rx.recv() {
match assignment {
Assignment::Stop => break,
Assignment::Joinable {
body,
ambient,
done,
} => {
let outcome = ambient.run(|| catch_unwind(AssertUnwindSafe(body)));
let _ = done.send(outcome);
finish_job(&inner, &set);
}
Assignment::Detached {
body,
ambient,
label,
} => {
let outcome = ambient.run(|| catch_unwind(AssertUnwindSafe(body)));
if outcome.is_err() {
announce_panic(&label);
}
finish_job(&inner, &set);
}
}
}
}
fn finish_job(inner: &Arc<PoolInner>, set: &Arc<SetHandle>) {
let freed = {
let mut st = lock_set(set);
st.running -= 1;
st.became_free()
};
free_if_idle(inner, set, freed);
}
pub struct WorkerPool<const N: usize> {
inner: Arc<PoolInner>,
}
impl<const N: usize> WorkerPool<N> {
pub fn new(name_prefix: &'static str, roster: [WorkerRole; N], capacity: usize) -> Self {
Self {
inner: Arc::new(PoolInner {
roster: Box::new(roster),
name_prefix,
capacity,
reg: Mutex::new(Registry {
idle: VecDeque::new(),
all: Vec::new(),
created: 0,
joins: Vec::new(),
stopping: false,
}),
}),
}
}
pub fn acquire(&self) -> io::Result<(SetLease, [Worker; N])> {
enum Decision {
Reuse(Arc<SetHandle>),
Grow(usize),
Full,
}
let decision = {
let mut reg = self.inner.lock();
if reg.stopping {
return Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"worker pool is shutting down",
));
}
if let Some(set) = reg.idle.pop_front() {
Decision::Reuse(set)
} else if reg.created < self.inner.capacity {
let index = reg.created;
reg.created += 1;
Decision::Grow(index)
} else {
Decision::Full
}
};
let set = match decision {
Decision::Full => {
return Err(io::Error::new(
io::ErrorKind::WouldBlock,
"worker pool at capacity",
));
}
Decision::Reuse(set) => set,
Decision::Grow(index) => match self.spawn_set(index) {
Ok((set, joins)) => {
let mut reg = self.inner.lock();
reg.joins.extend(joins);
reg.all.push(set.clone());
set
}
Err(e) => {
self.inner.lock().created -= 1;
return Err(e);
}
},
};
{
let mut st = lock_set(&set);
st.leased = true;
st.parked = false;
}
let workers: Vec<Worker> = (0..N)
.map(|slot| Worker {
set: set.clone(),
tx: set.senders[slot].clone(),
})
.collect();
let workers: [Worker; N] = workers
.try_into()
.unwrap_or_else(|_| unreachable!("N workers for an N-role set"));
let lease = SetLease {
inner: self.inner.clone(),
set,
};
Ok((lease, workers))
}
fn spawn_set(&self, index: usize) -> io::Result<(Arc<SetHandle>, Vec<JoinHandle<()>>)> {
let mut senders = Vec::with_capacity(N);
let mut receivers = Vec::with_capacity(N);
for _ in 0..N {
let (tx, rx) = channel::<Assignment>();
senders.push(tx);
receivers.push(rx);
}
let set = Arc::new(SetHandle {
senders,
state: Mutex::new(SetState {
leased: false,
running: 0,
parked: false,
}),
});
let mut joins: Vec<JoinHandle<()>> = Vec::with_capacity(N);
for (slot, rx) in receivers.into_iter().enumerate() {
let role = self.inner.roster[slot];
let name = format!("{}-{} {index}", self.inner.name_prefix, role.suffix);
let inner = self.inner.clone();
let set_for_worker = set.clone();
let spawned = thread::Builder::new()
.name(name)
.stack_size(role.stack.bytes())
.spawn(move || {
let _ = enter_ioc_thread(role.priority);
worker_loop(inner, set_for_worker, rx);
});
match spawned {
Ok(handle) => joins.push(handle),
Err(e) => {
for tx in &set.senders {
let _ = tx.send(Assignment::Stop);
}
for handle in joins {
let _ = handle.join();
}
return Err(e);
}
}
}
Ok((set, joins))
}
pub fn worker_count(&self) -> usize {
self.inner.lock().created * N
}
pub fn set_usage(&self) -> (usize, usize, usize) {
let reg = self.inner.lock();
let busy = reg.created - reg.idle.len();
(busy, reg.created, self.inner.capacity)
}
}
impl<const N: usize> Drop for WorkerPool<N> {
fn drop(&mut self) {
let (senders, joins) = {
let mut reg = self.inner.lock();
reg.stopping = true;
reg.idle.clear();
let joins = std::mem::take(&mut reg.joins);
let mut senders: Vec<Sender<Assignment>> = Vec::new();
for set in ®.all {
senders.extend(set.senders.iter().cloned());
}
(senders, joins)
};
for tx in senders {
let _ = tx.send(Assignment::Stop);
}
for handle in joins {
let _ = handle.join();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
fn roster2() -> [WorkerRole; 2] {
[
WorkerRole {
suffix: "reader",
stack: StackSizeClass::Small,
priority: ThreadPriority::Low,
},
WorkerRole {
suffix: "writer",
stack: StackSizeClass::Small,
priority: ThreadPriority::Low,
},
]
}
#[test]
fn sequential_borrows_reuse_one_set() {
let pool: WorkerPool<2> = WorkerPool::new("test-pool", roster2(), 4);
const BORROWS: usize = 8;
for i in 0..BORROWS {
let (lease, [reader, writer]) = pool.acquire().expect("borrow");
let ran = Arc::new(AtomicUsize::new(0));
let r = ran.clone();
let jr = reader.run(move || {
r.fetch_add(1, Ordering::SeqCst);
});
let w = ran.clone();
let jw = writer.run(move || {
w.fetch_add(1, Ordering::SeqCst);
});
assert!(jr.join().is_ok());
assert!(jw.join().is_ok());
drop(lease);
let deadline = Instant::now() + Duration::from_secs(5);
while pool.set_usage().0 != 0 {
assert!(Instant::now() < deadline, "set never returned to idle");
thread::yield_now();
}
assert_eq!(ran.load(Ordering::SeqCst), 2);
assert_eq!(
pool.worker_count(),
2,
"borrow {i} created new threads instead of reusing the idle set"
);
}
assert_eq!(
pool.worker_count(),
2,
"{BORROWS} sequential borrows must have created exactly one set"
);
}
#[test]
fn a_set_is_not_reidled_while_a_job_runs() {
let pool: WorkerPool<2> = WorkerPool::new("test-hold", roster2(), 4);
let (lease, [reader, writer]) = pool.acquire().expect("borrow");
let gate = Arc::new((Mutex::new(false), std::sync::Condvar::new()));
let g = gate.clone();
let blocking = reader.run(move || {
let (m, cv) = &*g;
let mut open = m.lock().unwrap();
while !*open {
open = cv.wait(open).unwrap();
}
});
let quick = writer.run(|| {});
assert!(quick.join().is_ok());
drop(lease);
assert_eq!(
pool.set_usage().0,
1,
"a running job must keep its set busy"
);
{
let (m, cv) = &*gate;
*m.lock().unwrap() = true;
cv.notify_all();
}
assert!(blocking.join().is_ok());
let deadline = Instant::now() + Duration::from_secs(5);
while pool.set_usage().0 != 0 {
assert!(
Instant::now() < deadline,
"set never returned after its last job"
);
thread::yield_now();
}
assert_eq!(pool.worker_count(), 2);
}
#[test]
fn acquire_refuses_at_capacity_without_creating_a_thread() {
let pool: WorkerPool<2> = WorkerPool::new("test-cap", roster2(), 1);
let (lease, _workers) = pool.acquire().expect("first borrow");
let before = pool.worker_count();
let refused = pool.acquire();
assert_eq!(
refused.err().map(|e| e.kind()),
Some(io::ErrorKind::WouldBlock),
"a full pool must refuse with EAGAIN, not queue or grow"
);
assert_eq!(
pool.worker_count(),
before,
"a refusal must create no thread"
);
drop(lease);
}
#[test]
fn a_panicked_job_returns_its_set_and_the_worker_survives() {
let pool: WorkerPool<2> = WorkerPool::new("test-panic", roster2(), 2);
let (lease, [reader, writer]) = pool.acquire().expect("borrow");
let boom = reader.run(|| panic!("job blew up"));
let ok = writer.run(|| {});
assert!(boom.join().is_err(), "the panic must reach the joiner");
assert!(ok.join().is_ok());
drop(lease);
let deadline = Instant::now() + Duration::from_secs(5);
while pool.set_usage().0 != 0 {
assert!(Instant::now() < deadline, "panicked set never returned");
thread::yield_now();
}
let created_before = pool.worker_count();
let (lease2, [r2, w2]) = pool.acquire().expect("borrow after panic");
assert!(r2.run(|| {}).join().is_ok());
assert!(w2.run(|| {}).join().is_ok());
drop(lease2);
assert_eq!(
pool.worker_count(),
created_before,
"a lost worker is never recreated, and a survivor needs no replacement"
);
}
#[test]
fn a_detached_job_returns_its_set() {
let pool: WorkerPool<2> = WorkerPool::new("test-detach", roster2(), 2);
let (lease, [reader, writer]) = pool.acquire().expect("borrow");
let ran = Arc::new(AtomicUsize::new(0));
let r = ran.clone();
reader.run_detached("conn".into(), move || {
r.fetch_add(1, Ordering::SeqCst);
});
let done = writer.run(|| {});
assert!(done.join().is_ok());
drop(lease);
let deadline = Instant::now() + Duration::from_secs(5);
while pool.set_usage().0 != 0 {
assert!(Instant::now() < deadline, "detached set never returned");
thread::yield_now();
}
assert_eq!(ran.load(Ordering::SeqCst), 1);
}
#[test]
fn dropping_the_pool_joins_its_workers() {
let pool: WorkerPool<2> = WorkerPool::new("test-drop", roster2(), 2);
let (lease, [reader, writer]) = pool.acquire().expect("borrow");
assert!(reader.run(|| {}).join().is_ok());
assert!(writer.run(|| {}).join().is_ok());
drop(lease);
let deadline = Instant::now() + Duration::from_secs(5);
while pool.set_usage().0 != 0 {
assert!(Instant::now() < deadline, "set never returned before drop");
thread::yield_now();
}
drop(pool);
}
}