486 lines
17 KiB
Rust
486 lines
17 KiB
Rust
use std::marker::PhantomData;
|
|
use std::sync::Arc;
|
|
use std::sync::mpsc::{self, Receiver, RecvTimeoutError, SyncSender, TrySendError};
|
|
use std::thread::{self, JoinHandle, Scope};
|
|
use std::time::Duration;
|
|
|
|
/// Executor shared by every worker of a pool. The `'static` bound required to
|
|
/// move the executor into detached threads is expressed on the pool types
|
|
/// instead of this trait so scoped pools can borrow their environment.
|
|
pub trait ObjectTaskExecutor<T, R>: Send + Sync {
|
|
fn execute(&self, worker_index: usize, task: T) -> R;
|
|
}
|
|
|
|
enum ObjectWorkerMessage<T> {
|
|
Task(T),
|
|
Shutdown,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum ObjectWorkerSubmitError<T> {
|
|
QueueFull { worker_index: usize, task: T },
|
|
Disconnected { worker_index: usize, task: T },
|
|
}
|
|
|
|
pub struct ObjectWorkerPool<T, R, E>
|
|
where
|
|
T: Send + 'static,
|
|
R: Send + 'static,
|
|
E: ObjectTaskExecutor<T, R> + 'static,
|
|
{
|
|
task_txs: Vec<SyncSender<ObjectWorkerMessage<T>>>,
|
|
result_rx: Receiver<R>,
|
|
workers: Vec<JoinHandle<()>>,
|
|
next_worker_idx: usize,
|
|
_executor: Arc<E>,
|
|
}
|
|
|
|
impl<T, R, E> ObjectWorkerPool<T, R, E>
|
|
where
|
|
T: Send + 'static,
|
|
R: Send + 'static,
|
|
E: ObjectTaskExecutor<T, R> + 'static,
|
|
{
|
|
pub fn new(worker_count: usize, queue_capacity: usize, executor: E) -> Result<Self, String> {
|
|
if worker_count == 0 {
|
|
return Err("ObjectWorkerPool requires at least one worker".to_string());
|
|
}
|
|
if queue_capacity == 0 {
|
|
return Err("ObjectWorkerPool requires queue_capacity > 0".to_string());
|
|
}
|
|
|
|
let executor = Arc::new(executor);
|
|
let (result_tx, result_rx) = mpsc::channel::<R>();
|
|
let mut task_txs = Vec::with_capacity(worker_count);
|
|
let mut workers = Vec::with_capacity(worker_count);
|
|
|
|
for worker_index in 0..worker_count {
|
|
let (task_tx, task_rx) = mpsc::sync_channel::<ObjectWorkerMessage<T>>(queue_capacity);
|
|
let result_tx = result_tx.clone();
|
|
let executor = Arc::clone(&executor);
|
|
let handle = thread::Builder::new()
|
|
.name(format!("object-validation-worker-{worker_index}"))
|
|
.spawn(move || object_worker_loop(worker_index, task_rx, result_tx, executor))
|
|
.map_err(|e| format!("spawn object worker failed: {e}"))?;
|
|
task_txs.push(task_tx);
|
|
workers.push(handle);
|
|
}
|
|
|
|
Ok(Self {
|
|
task_txs,
|
|
result_rx,
|
|
workers,
|
|
next_worker_idx: 0,
|
|
_executor: executor,
|
|
})
|
|
}
|
|
|
|
pub fn worker_count(&self) -> usize {
|
|
self.task_txs.len()
|
|
}
|
|
|
|
pub fn next_worker_index(&self) -> usize {
|
|
self.next_worker_idx
|
|
}
|
|
|
|
pub fn try_submit_round_robin(&mut self, task: T) -> Result<usize, ObjectWorkerSubmitError<T>> {
|
|
let worker_index = self.next_worker_idx % self.task_txs.len();
|
|
match self.task_txs[worker_index].try_send(ObjectWorkerMessage::Task(task)) {
|
|
Ok(()) => {
|
|
self.next_worker_idx = (worker_index + 1) % self.task_txs.len();
|
|
Ok(worker_index)
|
|
}
|
|
Err(TrySendError::Full(ObjectWorkerMessage::Task(task))) => {
|
|
Err(ObjectWorkerSubmitError::QueueFull { worker_index, task })
|
|
}
|
|
Err(TrySendError::Disconnected(ObjectWorkerMessage::Task(task))) => {
|
|
Err(ObjectWorkerSubmitError::Disconnected { worker_index, task })
|
|
}
|
|
Err(TrySendError::Full(ObjectWorkerMessage::Shutdown))
|
|
| Err(TrySendError::Disconnected(ObjectWorkerMessage::Shutdown)) => {
|
|
unreachable!("shutdown is never submitted via try_submit_round_robin")
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn recv_result_timeout(&self, timeout: Duration) -> Result<Option<R>, String> {
|
|
match self.result_rx.recv_timeout(timeout) {
|
|
Ok(result) => Ok(Some(result)),
|
|
Err(RecvTimeoutError::Timeout) => Ok(None),
|
|
Err(RecvTimeoutError::Disconnected) => {
|
|
Err("object worker result channel disconnected".to_string())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn shutdown(mut self) -> Result<(), String> {
|
|
self.shutdown_inner()
|
|
}
|
|
|
|
fn shutdown_inner(&mut self) -> Result<(), String> {
|
|
if self.workers.is_empty() {
|
|
return Ok(());
|
|
}
|
|
for tx in &self.task_txs {
|
|
tx.send(ObjectWorkerMessage::Shutdown)
|
|
.map_err(|e| format!("send shutdown to object worker failed: {e}"))?;
|
|
}
|
|
let mut first_err = None;
|
|
for handle in self.workers.drain(..) {
|
|
if let Err(e) = handle.join() {
|
|
if first_err.is_none() {
|
|
first_err = Some(format!("join object worker failed: {e:?}"));
|
|
}
|
|
}
|
|
}
|
|
if let Some(err) = first_err {
|
|
return Err(err);
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl<T, R, E> Drop for ObjectWorkerPool<T, R, E>
|
|
where
|
|
T: Send + 'static,
|
|
R: Send + 'static,
|
|
E: ObjectTaskExecutor<T, R> + 'static,
|
|
{
|
|
fn drop(&mut self) {
|
|
let _ = self.shutdown_inner();
|
|
}
|
|
}
|
|
|
|
fn object_worker_loop<T, R, E>(
|
|
worker_index: usize,
|
|
task_rx: Receiver<ObjectWorkerMessage<T>>,
|
|
result_tx: mpsc::Sender<R>,
|
|
executor: Arc<E>,
|
|
) where
|
|
T: Send,
|
|
R: Send,
|
|
E: ObjectTaskExecutor<T, R>,
|
|
{
|
|
loop {
|
|
match task_rx.recv() {
|
|
Ok(ObjectWorkerMessage::Task(task)) => {
|
|
let result = executor.execute(worker_index, task);
|
|
if result_tx.send(result).is_err() {
|
|
break;
|
|
}
|
|
}
|
|
Ok(ObjectWorkerMessage::Shutdown) | Err(_) => break,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Scoped variant of `ObjectWorkerPool`: workers are spawned on a
|
|
/// `std::thread::Scope`, so tasks, results and the executor may borrow their
|
|
/// environment (`'env`) instead of being `'static`. The pool never sends
|
|
/// `Shutdown`; workers exit when every task sender is dropped, which happens
|
|
/// when the pool itself is dropped ahead of the scope join.
|
|
pub struct ScopedObjectWorkerPool<'scope, 'env, T, R, E>
|
|
where
|
|
T: Send + 'env,
|
|
R: Send + 'env,
|
|
E: ObjectTaskExecutor<T, R> + 'env,
|
|
{
|
|
task_txs: Vec<SyncSender<ObjectWorkerMessage<T>>>,
|
|
result_rx: Receiver<R>,
|
|
next_worker_idx: usize,
|
|
_executor: Arc<E>,
|
|
// Join handles are intentionally not stored: dropping a `ScopedJoinHandle`
|
|
// detaches the worker and the enclosing scope joins it on exit, after the
|
|
// dropped task senders have made every worker return.
|
|
_marker: PhantomData<(&'scope (), &'env ())>,
|
|
}
|
|
|
|
impl<'scope, 'env, T, R, E> ScopedObjectWorkerPool<'scope, 'env, T, R, E>
|
|
where
|
|
T: Send + 'env,
|
|
R: Send + 'env,
|
|
E: ObjectTaskExecutor<T, R> + 'env,
|
|
{
|
|
pub fn new(
|
|
scope: &'scope Scope<'scope, 'env>,
|
|
worker_count: usize,
|
|
queue_capacity: usize,
|
|
executor: E,
|
|
) -> Result<Self, String> {
|
|
if worker_count == 0 {
|
|
return Err("ScopedObjectWorkerPool requires at least one worker".to_string());
|
|
}
|
|
if queue_capacity == 0 {
|
|
return Err("ScopedObjectWorkerPool requires queue_capacity > 0".to_string());
|
|
}
|
|
|
|
let executor = Arc::new(executor);
|
|
let (result_tx, result_rx) = mpsc::channel::<R>();
|
|
let mut task_txs = Vec::with_capacity(worker_count);
|
|
|
|
for worker_index in 0..worker_count {
|
|
let (task_tx, task_rx) = mpsc::sync_channel::<ObjectWorkerMessage<T>>(queue_capacity);
|
|
let result_tx = result_tx.clone();
|
|
let executor = Arc::clone(&executor);
|
|
thread::Builder::new()
|
|
.name(format!("object-validation-worker-{worker_index}"))
|
|
.spawn_scoped(scope, move || {
|
|
object_worker_loop(worker_index, task_rx, result_tx, executor)
|
|
})
|
|
.map_err(|e| format!("spawn scoped object worker failed: {e}"))?;
|
|
task_txs.push(task_tx);
|
|
}
|
|
|
|
Ok(Self {
|
|
task_txs,
|
|
result_rx,
|
|
next_worker_idx: 0,
|
|
_executor: executor,
|
|
_marker: PhantomData,
|
|
})
|
|
}
|
|
|
|
pub fn worker_count(&self) -> usize {
|
|
self.task_txs.len()
|
|
}
|
|
|
|
pub fn try_submit_round_robin(&mut self, task: T) -> Result<usize, ObjectWorkerSubmitError<T>> {
|
|
let worker_index = self.next_worker_idx % self.task_txs.len();
|
|
match self.task_txs[worker_index].try_send(ObjectWorkerMessage::Task(task)) {
|
|
Ok(()) => {
|
|
self.next_worker_idx = (worker_index + 1) % self.task_txs.len();
|
|
Ok(worker_index)
|
|
}
|
|
Err(TrySendError::Full(ObjectWorkerMessage::Task(task))) => {
|
|
Err(ObjectWorkerSubmitError::QueueFull { worker_index, task })
|
|
}
|
|
Err(TrySendError::Disconnected(ObjectWorkerMessage::Task(task))) => {
|
|
Err(ObjectWorkerSubmitError::Disconnected { worker_index, task })
|
|
}
|
|
Err(TrySendError::Full(ObjectWorkerMessage::Shutdown))
|
|
| Err(TrySendError::Disconnected(ObjectWorkerMessage::Shutdown)) => {
|
|
unreachable!("shutdown is never submitted via try_submit_round_robin")
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn recv_result_timeout(&self, timeout: Duration) -> Result<Option<R>, String> {
|
|
match self.result_rx.recv_timeout(timeout) {
|
|
Ok(result) => Ok(Some(result)),
|
|
Err(RecvTimeoutError::Timeout) => Ok(None),
|
|
Err(RecvTimeoutError::Disconnected) => {
|
|
Err("scoped object worker result channel disconnected".to_string())
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<'scope, 'env, T, R, E> Drop for ScopedObjectWorkerPool<'scope, 'env, T, R, E>
|
|
where
|
|
T: Send + 'env,
|
|
R: Send + 'env,
|
|
E: ObjectTaskExecutor<T, R> + 'env,
|
|
{
|
|
fn drop(&mut self) {
|
|
// Close every worker input queue so blocked `recv` calls return and
|
|
// the scoped workers exit before the enclosing scope joins them.
|
|
self.task_txs.clear();
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{ObjectTaskExecutor, ObjectWorkerPool, ObjectWorkerSubmitError};
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::{Arc, Barrier};
|
|
use std::time::Duration;
|
|
|
|
#[derive(Clone)]
|
|
struct EchoExecutor;
|
|
|
|
impl ObjectTaskExecutor<u32, (usize, u32)> for EchoExecutor {
|
|
fn execute(&self, worker_index: usize, task: u32) -> (usize, u32) {
|
|
(worker_index, task)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn object_worker_pool_rejects_invalid_config_and_shutdowns_explicitly() {
|
|
let err = match ObjectWorkerPool::new(0, 1, EchoExecutor) {
|
|
Ok(_) => panic!("zero workers should be rejected"),
|
|
Err(err) => err,
|
|
};
|
|
assert!(err.contains("at least one worker"));
|
|
let err = match ObjectWorkerPool::new(1, 0, EchoExecutor) {
|
|
Ok(_) => panic!("zero queue should be rejected"),
|
|
Err(err) => err,
|
|
};
|
|
assert!(err.contains("queue_capacity > 0"));
|
|
|
|
let pool = ObjectWorkerPool::new(2, 1, EchoExecutor).expect("pool");
|
|
assert_eq!(pool.worker_count(), 2);
|
|
assert_eq!(pool.next_worker_index(), 0);
|
|
pool.shutdown().expect("shutdown");
|
|
}
|
|
|
|
#[test]
|
|
fn object_worker_pool_round_robin_submits_to_worker_queues() {
|
|
let mut pool = ObjectWorkerPool::new(3, 4, EchoExecutor).expect("pool");
|
|
assert_eq!(pool.try_submit_round_robin(10).expect("submit 10"), 0);
|
|
assert_eq!(pool.try_submit_round_robin(11).expect("submit 11"), 1);
|
|
assert_eq!(pool.try_submit_round_robin(12).expect("submit 12"), 2);
|
|
assert_eq!(pool.try_submit_round_robin(13).expect("submit 13"), 0);
|
|
|
|
let mut results = Vec::new();
|
|
for _ in 0..4 {
|
|
results.push(
|
|
pool.recv_result_timeout(Duration::from_secs(1))
|
|
.expect("result channel")
|
|
.expect("result"),
|
|
);
|
|
}
|
|
results.sort_by_key(|(_, task)| *task);
|
|
assert_eq!(results, vec![(0, 10), (1, 11), (2, 12), (0, 13)]);
|
|
}
|
|
|
|
struct BlockingExecutor {
|
|
barrier: Arc<Barrier>,
|
|
started: Arc<AtomicBool>,
|
|
}
|
|
|
|
impl ObjectTaskExecutor<u32, u32> for BlockingExecutor {
|
|
fn execute(&self, _worker_index: usize, task: u32) -> u32 {
|
|
self.started.store(true, Ordering::SeqCst);
|
|
self.barrier.wait();
|
|
task
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn object_worker_pool_reports_full_worker_queue_without_advancing_round_robin() {
|
|
let barrier = Arc::new(Barrier::new(2));
|
|
let started = Arc::new(AtomicBool::new(false));
|
|
let mut pool = ObjectWorkerPool::new(
|
|
1,
|
|
1,
|
|
BlockingExecutor {
|
|
barrier: Arc::clone(&barrier),
|
|
started: Arc::clone(&started),
|
|
},
|
|
)
|
|
.expect("pool");
|
|
|
|
assert_eq!(pool.try_submit_round_robin(1).expect("first task"), 0);
|
|
let deadline = std::time::Instant::now() + Duration::from_secs(1);
|
|
while !started.load(Ordering::SeqCst) {
|
|
assert!(
|
|
std::time::Instant::now() < deadline,
|
|
"worker did not start first task"
|
|
);
|
|
std::thread::sleep(Duration::from_millis(1));
|
|
}
|
|
assert_eq!(pool.try_submit_round_robin(2).expect("queued task"), 0);
|
|
match pool.try_submit_round_robin(3) {
|
|
Err(ObjectWorkerSubmitError::QueueFull { worker_index, task }) => {
|
|
assert_eq!(worker_index, 0);
|
|
assert_eq!(task, 3);
|
|
}
|
|
other => panic!("expected queue full, got {other:?}"),
|
|
}
|
|
assert_eq!(pool.next_worker_index(), 0);
|
|
|
|
barrier.wait();
|
|
assert_eq!(
|
|
pool.recv_result_timeout(Duration::from_secs(1))
|
|
.expect("result channel"),
|
|
Some(1)
|
|
);
|
|
barrier.wait();
|
|
assert_eq!(
|
|
pool.recv_result_timeout(Duration::from_secs(1))
|
|
.expect("result channel"),
|
|
Some(2)
|
|
);
|
|
}
|
|
|
|
struct BorrowingEchoExecutor<'a> {
|
|
base: &'a u32,
|
|
}
|
|
|
|
impl<'a> ObjectTaskExecutor<u32, u32> for BorrowingEchoExecutor<'a> {
|
|
fn execute(&self, _worker_index: usize, task: u32) -> u32 {
|
|
task + *self.base
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn scoped_object_worker_pool_borrows_environment_and_processes_tasks() {
|
|
let base = 100u32;
|
|
std::thread::scope(|scope| {
|
|
let mut pool = super::ScopedObjectWorkerPool::new(
|
|
scope,
|
|
2,
|
|
2,
|
|
BorrowingEchoExecutor { base: &base },
|
|
)
|
|
.expect("scoped pool");
|
|
assert_eq!(pool.worker_count(), 2);
|
|
pool.try_submit_round_robin(1).expect("submit 1");
|
|
pool.try_submit_round_robin(2).expect("submit 2");
|
|
let mut results = Vec::new();
|
|
for _ in 0..2 {
|
|
results.push(
|
|
pool.recv_result_timeout(Duration::from_secs(1))
|
|
.expect("result channel")
|
|
.expect("result"),
|
|
);
|
|
}
|
|
results.sort();
|
|
assert_eq!(results, vec![101, 102]);
|
|
// Dropping the pool inside the scope closes the task queues; the
|
|
// workers exit on their own and the scope join below must not hang.
|
|
drop(pool);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn scoped_object_worker_pool_reports_full_queue() {
|
|
std::thread::scope(|scope| {
|
|
let barrier = Arc::new(Barrier::new(2));
|
|
let started = Arc::new(AtomicBool::new(false));
|
|
let mut pool = super::ScopedObjectWorkerPool::new(
|
|
scope,
|
|
1,
|
|
1,
|
|
BlockingExecutor {
|
|
barrier: Arc::clone(&barrier),
|
|
started: Arc::clone(&started),
|
|
},
|
|
)
|
|
.expect("scoped pool");
|
|
pool.try_submit_round_robin(1).expect("first task");
|
|
let deadline = std::time::Instant::now() + Duration::from_secs(1);
|
|
while !started.load(Ordering::SeqCst) {
|
|
assert!(
|
|
std::time::Instant::now() < deadline,
|
|
"scoped worker did not start first task"
|
|
);
|
|
std::thread::sleep(Duration::from_millis(1));
|
|
}
|
|
pool.try_submit_round_robin(2).expect("queued task");
|
|
match pool.try_submit_round_robin(3) {
|
|
Err(ObjectWorkerSubmitError::QueueFull { worker_index, task }) => {
|
|
assert_eq!(worker_index, 0);
|
|
assert_eq!(task, 3);
|
|
}
|
|
other => panic!("expected queue full, got {other:?}"),
|
|
}
|
|
// Dropping the pool closes the task queues; releasing the barrier
|
|
// afterwards lets the blocked worker finish task 1, observe the
|
|
// closed result channel, and exit before the scope join.
|
|
drop(pool);
|
|
barrier.wait();
|
|
});
|
|
}
|
|
}
|