use std::sync::{
Arc, Barrier,
atomic::{AtomicBool, AtomicUsize, Ordering},
mpsc,
};
use std::time::Duration;
use timerwheel::Error;
use timerwheel::executor::{BoxTask, Executor, ExecutorMetrics, Pool, RejectedTask};
use timerwheel::policy::RejectPolicy;
struct CountingExecutorMetricSink {
observations: Arc<AtomicUsize>,
}
impl timerwheel::executor::MetricSink for CountingExecutorMetricSink {
fn observe_metrics(&self, metrics: ExecutorMetrics) {
if metrics.submitted_tasks > 0 {
self.observations.fetch_add(1, Ordering::SeqCst);
}
}
}
struct RejectingExecutor;
impl Executor for RejectingExecutor {
fn try_execute(&self, task: BoxTask) -> std::result::Result<(), RejectedTask> {
Err(RejectedTask::new(Error::Saturated, task))
}
fn shutdown(&self) -> timerwheel::Result<()> {
Ok(())
}
fn metrics(&self) -> ExecutorMetrics {
ExecutorMetrics::default()
}
}
#[test]
fn pool_executes_tasks() {
let pool = Pool::builder()
.workers(2)
.queue_capacity(8)
.build()
.expect("pool builds");
let (tx, rx) = mpsc::channel();
pool.execute(move || {
tx.send(7usize).expect("send succeeds");
})
.expect("execute succeeds");
assert_eq!(
rx.recv_timeout(Duration::from_secs(1)).expect("task runs"),
7
);
assert!(pool.shutdown().is_ok());
}
#[test]
fn pool_try_execute_rejects_when_queue_is_full() {
let pool = Pool::builder()
.workers(1)
.queue_capacity(1)
.reject_policy(RejectPolicy::Reject)
.build()
.expect("pool builds");
let started = Arc::new(Barrier::new(2));
let release = Arc::new(Barrier::new(2));
let started_task = Arc::clone(&started);
let release_task = Arc::clone(&release);
pool.execute(move || {
started_task.wait();
release_task.wait();
})
.expect("first task starts");
started.wait();
pool.try_execute(|| {}).expect("queued task succeeds");
assert_eq!(
pool.try_execute(|| {}).expect_err("full queue rejects"),
Error::Saturated
);
release.wait();
assert!(pool.shutdown().is_ok());
}
#[test]
fn rejected_task_returns_original_task() {
let ran = Arc::new(AtomicBool::new(false));
let marker = Arc::clone(&ran);
let task: BoxTask = Box::new(move || marker.store(true, Ordering::SeqCst));
let rejected = match RejectingExecutor.try_execute(task) {
Ok(()) => panic!("task should be rejected"),
Err(rejected) => rejected,
};
assert_eq!(rejected.error(), &Error::Saturated);
rejected.into_task().run();
assert!(ran.load(Ordering::SeqCst));
}
#[test]
fn executor_can_be_used_as_trait_object() {
let executor: Arc<dyn Executor> = Arc::new(RejectingExecutor);
let task: BoxTask = Box::new(|| {});
let rejected = match executor.try_execute(task) {
Ok(()) => panic!("task should be rejected"),
Err(rejected) => rejected,
};
assert_eq!(rejected.error(), &Error::Saturated);
}
#[test]
fn pool_executes_trait_object_task() {
let pool = Pool::builder()
.workers(1)
.queue_capacity(8)
.build()
.expect("pool builds");
let (tx, rx) = mpsc::channel();
let task: BoxTask = Box::new(move || tx.send(()).expect("send succeeds"));
assert!(Executor::try_execute(&pool, task).is_ok());
rx.recv_timeout(Duration::from_secs(1))
.expect("trait task should run");
assert!(pool.shutdown().is_ok());
}
#[test]
fn pool_try_execute_returns_task_when_full() {
let pool = Pool::builder()
.workers(1)
.queue_capacity(1)
.build()
.expect("pool builds");
let (release_tx, release_rx) = mpsc::channel();
let (started_tx, started_rx) = mpsc::channel();
pool.execute(move || {
started_tx.send(()).expect("send succeeds");
release_rx.recv().expect("release received");
})
.expect("blocking task accepted");
started_rx
.recv_timeout(Duration::from_secs(1))
.expect("blocking task starts");
pool.execute(|| {}).expect("queued task accepted");
let calls = Arc::new(AtomicUsize::new(0));
let calls_for_task = Arc::clone(&calls);
let task: BoxTask = Box::new(move || {
calls_for_task.fetch_add(1, Ordering::SeqCst);
});
let rejected = match Executor::try_execute(&pool, task) {
Ok(()) => panic!("queue should be full"),
Err(rejected) => rejected,
};
rejected.into_task().run();
assert_eq!(calls.load(Ordering::SeqCst), 1);
release_tx.send(()).expect("release succeeds");
assert!(pool.shutdown().is_ok());
}
#[test]
fn pool_shutdown_is_idempotent_and_closes_submission() {
let pool = Pool::builder().build().expect("pool builds");
assert!(pool.shutdown().is_ok());
assert!(pool.shutdown().is_ok());
assert_eq!(
pool.execute(|| {})
.expect_err("execute after shutdown rejects"),
Error::Closed
);
}
#[test]
fn pool_recovers_from_task_panic() {
let pool = Pool::builder()
.workers(1)
.queue_capacity(8)
.build()
.expect("pool builds");
let (tx, rx) = mpsc::channel();
pool.execute(|| panic!("expected panic"))
.expect("panic task accepted");
pool.execute(move || {
tx.send("after panic").expect("send succeeds");
})
.expect("second task accepted");
assert_eq!(
rx.recv_timeout(Duration::from_secs(1))
.expect("worker should continue"),
"after panic"
);
assert!(pool.shutdown().is_ok());
assert_eq!(pool.metrics().panicked_tasks, 1);
}
#[test]
fn pool_reports_metrics_to_metric_sink() {
let observations = Arc::new(AtomicUsize::new(0));
let pool = Pool::builder()
.workers(1)
.queue_capacity(8)
.metric_sink(CountingExecutorMetricSink {
observations: Arc::clone(&observations),
})
.build()
.expect("pool builds");
let (tx, rx) = mpsc::channel();
pool.execute(move || {
tx.send(()).expect("send succeeds");
})
.expect("execute succeeds");
rx.recv_timeout(Duration::from_secs(1))
.expect("task should run");
assert!(pool.shutdown().is_ok());
assert!(observations.load(Ordering::SeqCst) > 0);
}