use crate::{
Result,
batch::BatchId,
error::HammerworkError,
job::Job,
priority::PriorityWeights,
queue::{DatabaseQueue, JobOutcome, JobQueue, RecordedOutcome},
rate_limit::{RateLimit, RateLimiter, ThrottleConfig},
retry::RetryStrategy,
stats::{JobEvent, JobEventType, StatisticsCollector},
};
#[cfg(feature = "metrics")]
use crate::metrics::PrometheusMetricsCollector;
#[cfg(feature = "alerting")]
use crate::alerting::{AlertManager, AlertingConfig};
#[cfg(feature = "webhooks")]
use crate::events::{EventManager, JobError, JobLifecycleEvent, JobLifecycleEventType};
use chrono::{DateTime, Utc};
use sqlx::Database;
use std::{
future::Future,
panic::AssertUnwindSafe,
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Duration,
};
use tokio::{
sync::{mpsc, watch},
task::JoinSet,
time::sleep,
};
use tracing::{debug, error, info, warn};
pub const DEFAULT_SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(30);
pub const DEFAULT_LEASE_DURATION: Duration = crate::queue::DEFAULT_LEASE_DURATION;
pub const DEFAULT_MONITORING_INTERVAL: Duration = Duration::from_secs(30);
#[cfg(any(feature = "metrics", feature = "alerting"))]
const MIN_MONITORING_INTERVAL: Duration = Duration::from_millis(10);
pub const DEFAULT_STALE_JOB_REAPER_INTERVAL: Duration = Duration::from_secs(60);
pub const MAX_RETRY_DELAY: Duration = Duration::from_secs(365 * 24 * 60 * 60);
const MAX_ERROR_BACKOFF: Duration = Duration::from_secs(60);
const MIN_ERROR_BACKOFF: Duration = Duration::from_millis(100);
pub const MIN_POLL_INTERVAL: Duration = Duration::from_millis(10);
const MIN_HEARTBEAT_INTERVAL: Duration = Duration::from_millis(100);
const WORKER_RESTART_DELAY: Duration = Duration::from_secs(1);
type HandlerFuture = Pin<Box<dyn Future<Output = Result<JobResult>> + Send>>;
struct CatchUnwind<F> {
inner: F,
}
impl<F: Future + Unpin> Future for CatchUnwind<F> {
type Output = std::thread::Result<F::Output>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let inner = &mut self.inner;
match std::panic::catch_unwind(AssertUnwindSafe(|| Pin::new(&mut *inner).poll(cx))) {
Ok(Poll::Pending) => Poll::Pending,
Ok(Poll::Ready(output)) => Poll::Ready(Ok(output)),
Err(payload) => Poll::Ready(Err(payload)),
}
}
}
fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
if let Some(message) = payload.downcast_ref::<&str>() {
(*message).to_string()
} else if let Some(message) = payload.downcast_ref::<String>() {
message.clone()
} else {
"non-string panic payload".to_string()
}
}
fn handler_panic_error(payload: &(dyn std::any::Any + Send)) -> HammerworkError {
HammerworkError::Worker {
message: format!("Job handler panicked: {}", panic_message(payload)),
}
}
fn retry_at_from_delay(now: DateTime<Utc>, retry_delay: Duration) -> DateTime<Utc> {
let delay = chrono::Duration::from_std(retry_delay.min(MAX_RETRY_DELAY))
.unwrap_or(chrono::Duration::MAX);
now.checked_add_signed(delay)
.unwrap_or(DateTime::<Utc>::MAX_UTC)
}
fn effective_attempt_limit(job_max_attempts: i32, worker_cap: Option<i32>) -> i32 {
match worker_cap {
Some(cap) => job_max_attempts.min(cap),
None => job_max_attempts,
}
}
#[cfg(any(feature = "metrics", feature = "alerting"))]
struct AbortOnDrop(tokio::task::JoinHandle<()>);
#[cfg(any(feature = "metrics", feature = "alerting"))]
impl Drop for AbortOnDrop {
fn drop(&mut self) {
self.0.abort();
}
}
enum Acquired {
Job(Box<Job>),
Idle,
Shutdown,
}
fn shutdown_requested(shutdown_rx: &mut mpsc::Receiver<()>) -> bool {
!matches!(
shutdown_rx.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
)
}
fn error_backoff(base: Duration, consecutive_errors: u32) -> Duration {
let base = base.max(MIN_ERROR_BACKOFF);
let exponent = consecutive_errors.saturating_sub(1).min(16);
base.checked_mul(1u32 << exponent)
.unwrap_or(MAX_ERROR_BACKOFF)
.min(MAX_ERROR_BACKOFF.max(base))
}
#[derive(Debug, Clone)]
pub struct JobHookEvent {
pub job: Job,
pub timestamp: DateTime<Utc>,
pub duration: Option<Duration>,
pub error: Option<String>,
}
pub type JobHookHandler = Arc<dyn Fn(JobHookEvent) + Send + Sync>;
#[derive(Clone, Default)]
pub struct JobEventHooks {
pub on_job_start: Option<JobHookHandler>,
pub on_job_complete: Option<JobHookHandler>,
pub on_job_fail: Option<JobHookHandler>,
pub on_job_timeout: Option<JobHookHandler>,
pub on_job_retry: Option<JobHookHandler>,
}
impl JobEventHooks {
pub fn new() -> Self {
Self::default()
}
pub fn on_start<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.on_job_start = Some(Arc::new(handler));
self
}
pub fn on_complete<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.on_job_complete = Some(Arc::new(handler));
self
}
pub fn on_fail<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.on_job_fail = Some(Arc::new(handler));
self
}
pub fn on_timeout<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.on_job_timeout = Some(Arc::new(handler));
self
}
pub fn on_retry<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.on_job_retry = Some(Arc::new(handler));
self
}
pub(crate) fn fire_job_start(&self, job: Job) {
if let Some(handler) = &self.on_job_start {
let event = JobHookEvent {
job,
timestamp: Utc::now(),
duration: None,
error: None,
};
handler(event);
}
}
pub(crate) fn fire_job_complete(&self, job: Job, duration: Duration) {
if let Some(handler) = &self.on_job_complete {
let event = JobHookEvent {
job,
timestamp: Utc::now(),
duration: Some(duration),
error: None,
};
handler(event);
}
}
pub(crate) fn fire_job_fail(&self, job: Job, error: String) {
if let Some(handler) = &self.on_job_fail {
let event = JobHookEvent {
job,
timestamp: Utc::now(),
duration: None,
error: Some(error),
};
handler(event);
}
}
pub(crate) fn fire_job_timeout(&self, job: Job, duration: Duration) {
if let Some(handler) = &self.on_job_timeout {
let event = JobHookEvent {
job,
timestamp: Utc::now(),
duration: Some(duration),
error: Some("Job timed out".to_string()),
};
handler(event);
}
}
pub(crate) fn fire_job_retry(&self, job: Job, error: String) {
if let Some(handler) = &self.on_job_retry {
let event = JobHookEvent {
job,
timestamp: Utc::now(),
duration: None,
error: Some(error),
};
handler(event);
}
}
}
#[derive(Debug, Clone)]
pub struct AutoscaleConfig {
pub enabled: bool,
pub min_workers: usize,
pub max_workers: usize,
pub scale_up_threshold: usize,
pub scale_down_threshold: usize,
pub cooldown_period: Duration,
pub scale_step: usize,
pub evaluation_window: Duration,
}
impl Default for AutoscaleConfig {
fn default() -> Self {
Self {
enabled: true,
min_workers: 1,
max_workers: 10,
scale_up_threshold: 5,
scale_down_threshold: 2,
cooldown_period: Duration::from_secs(60),
scale_step: 1,
evaluation_window: Duration::from_secs(30),
}
}
}
impl AutoscaleConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_enabled(mut self, enabled: bool) -> Self {
self.enabled = enabled;
self
}
pub fn with_min_workers(mut self, min_workers: usize) -> Self {
self.min_workers = min_workers.max(1); self
}
pub fn with_max_workers(mut self, max_workers: usize) -> Self {
self.max_workers = max_workers.max(self.min_workers);
self
}
pub fn with_scale_up_threshold(mut self, threshold: usize) -> Self {
self.scale_up_threshold = threshold.max(1);
self
}
pub fn with_scale_down_threshold(mut self, threshold: usize) -> Self {
self.scale_down_threshold = threshold;
self
}
pub fn with_cooldown_period(mut self, period: Duration) -> Self {
self.cooldown_period = period;
self
}
pub fn with_scale_step(mut self, step: usize) -> Self {
self.scale_step = step.max(1);
self
}
pub fn with_evaluation_window(mut self, window: Duration) -> Self {
self.evaluation_window = window;
self
}
pub fn conservative() -> Self {
Self {
enabled: true,
min_workers: 2,
max_workers: 5,
scale_up_threshold: 10,
scale_down_threshold: 1,
cooldown_period: Duration::from_secs(300), scale_step: 1,
evaluation_window: Duration::from_secs(60),
}
}
pub fn aggressive() -> Self {
Self {
enabled: true,
min_workers: 1,
max_workers: 20,
scale_up_threshold: 3,
scale_down_threshold: 1,
cooldown_period: Duration::from_secs(30),
scale_step: 2,
evaluation_window: Duration::from_secs(15),
}
}
pub fn disabled() -> Self {
Self {
enabled: false,
..Self::default()
}
}
}
#[derive(Debug, Clone, Default)]
pub struct AutoscaleMetrics {
pub active_workers: usize,
pub avg_queue_depth: f64,
pub current_queue_depth: u64,
pub jobs_per_second: f64,
pub worker_utilization: f64,
pub time_since_last_scale: Duration,
pub last_scale_time: Option<chrono::DateTime<Utc>>,
}
#[derive(Debug, Clone, Copy)]
enum ScalingDecision {
ScaleUp,
ScaleDown,
}
type QueueDepthHistory = Arc<std::sync::RwLock<Vec<(chrono::DateTime<Utc>, u64)>>>;
struct Scaling<DB: Database> {
template: Worker<DB>,
desired: watch::Receiver<usize>,
initial: usize,
}
async fn wait_for_change(rx: &mut Option<watch::Receiver<usize>>) -> Option<usize> {
match rx {
Some(rx) => match rx.changed().await {
Ok(()) => Some(*rx.borrow_and_update()),
Err(_) => None,
},
None => std::future::pending().await,
}
}
#[derive(Debug, Clone, Default)]
pub struct BatchProcessingStats {
pub jobs_processed: u64,
pub jobs_completed: u64,
pub jobs_failed: u64,
pub batches_completed: u64,
pub batches_successful: u64,
pub total_processing_time_ms: u64,
pub average_processing_time_ms: f64,
pub last_processed_job: Option<DateTime<Utc>>,
}
impl BatchProcessingStats {
pub fn success_rate(&self) -> f64 {
if self.jobs_processed == 0 {
0.0
} else {
self.jobs_completed as f64 / self.jobs_processed as f64
}
}
pub fn batch_success_rate(&self) -> f64 {
if self.batches_completed == 0 {
0.0
} else {
self.batches_successful as f64 / self.batches_completed as f64
}
}
pub fn update_average_processing_time(&mut self) {
if self.jobs_completed > 0 {
self.average_processing_time_ms =
self.total_processing_time_ms as f64 / self.jobs_completed as f64;
}
}
}
pub type JobHandler = Arc<
dyn Fn(Job) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send>>
+ Send
+ Sync,
>;
pub type JobHandlerWithResult = Arc<
dyn Fn(Job) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<JobResult>> + Send>>
+ Send
+ Sync,
>;
#[derive(Debug, Clone)]
pub struct JobResult {
pub data: Option<serde_json::Value>,
}
impl JobResult {
pub fn success() -> Self {
Self { data: None }
}
pub fn with_data(data: serde_json::Value) -> Self {
Self { data: Some(data) }
}
}
impl Default for JobResult {
fn default() -> Self {
Self::success()
}
}
#[derive(Clone)]
pub enum JobHandlerType {
Legacy(JobHandler),
WithResult(JobHandlerWithResult),
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub struct Worker<DB: Database> {
queue: Arc<JobQueue<DB>>,
queue_name: String,
handler: JobHandlerType,
poll_interval: Duration,
max_retries: Option<i32>,
retry_delay: Duration,
default_retry_strategy: Option<RetryStrategy>,
default_timeout: Option<Duration>,
priority_weights: Option<PriorityWeights>,
stats_collector: Option<Arc<dyn StatisticsCollector>>,
rate_limiter: Option<RateLimiter>,
throttle_config: Option<ThrottleConfig>,
concurrency_limit: Option<Arc<tokio::sync::Semaphore>>,
#[cfg(feature = "metrics")]
metrics_collector: Option<Arc<PrometheusMetricsCollector>>,
#[cfg(feature = "alerting")]
alert_manager: Option<Arc<AlertManager>>,
last_job_time: Arc<std::sync::RwLock<DateTime<Utc>>>,
batch_processing_enabled: bool,
batch_stats: Arc<std::sync::RwLock<BatchProcessingStats>>,
event_hooks: JobEventHooks,
spawn_manager: Option<Arc<crate::spawn::SpawnManager<DB>>>,
#[cfg(feature = "webhooks")]
event_manager: Option<Arc<crate::events::EventManager>>,
shutdown_grace_period: Duration,
lease_duration: Duration,
monitoring_interval: Option<Duration>,
}
impl<DB: Database + Send + Sync + 'static> Clone for Worker<DB>
where
JobQueue<DB>: DatabaseQueue<Database = DB> + Send + Sync,
{
fn clone(&self) -> Self {
Self {
queue: Arc::clone(&self.queue),
queue_name: self.queue_name.clone(),
handler: self.handler.clone(),
poll_interval: self.poll_interval,
max_retries: self.max_retries,
retry_delay: self.retry_delay,
default_retry_strategy: self.default_retry_strategy.clone(),
default_timeout: self.default_timeout,
priority_weights: self.priority_weights.clone(),
stats_collector: self.stats_collector.clone(),
rate_limiter: self.rate_limiter.clone(),
throttle_config: self.throttle_config.clone(),
concurrency_limit: self.concurrency_limit.clone(),
#[cfg(feature = "metrics")]
metrics_collector: self.metrics_collector.clone(),
#[cfg(feature = "alerting")]
alert_manager: self.alert_manager.clone(),
last_job_time: Arc::new(std::sync::RwLock::new(Utc::now())),
batch_processing_enabled: self.batch_processing_enabled,
batch_stats: Arc::new(std::sync::RwLock::new(BatchProcessingStats::default())),
event_hooks: self.event_hooks.clone(),
spawn_manager: self.spawn_manager.clone(),
#[cfg(feature = "webhooks")]
event_manager: self.event_manager.clone(),
shutdown_grace_period: self.shutdown_grace_period,
lease_duration: self.lease_duration,
monitoring_interval: self.monitoring_interval,
}
}
}
impl<DB: Database + Send + Sync + 'static> Worker<DB>
where
JobQueue<DB>: DatabaseQueue<Database = DB> + Send + Sync,
{
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn new(queue: Arc<JobQueue<DB>>, queue_name: String, handler: JobHandler) -> Self {
Self {
queue,
queue_name,
handler: JobHandlerType::Legacy(handler),
poll_interval: Duration::from_secs(1),
max_retries: None,
retry_delay: Duration::from_secs(30),
default_retry_strategy: None,
default_timeout: None,
priority_weights: None,
stats_collector: None,
rate_limiter: None,
throttle_config: None,
concurrency_limit: None,
#[cfg(feature = "metrics")]
metrics_collector: None,
#[cfg(feature = "alerting")]
alert_manager: None,
last_job_time: Arc::new(std::sync::RwLock::new(Utc::now())),
batch_processing_enabled: false,
batch_stats: Arc::new(std::sync::RwLock::new(BatchProcessingStats::default())),
event_hooks: JobEventHooks::default(),
spawn_manager: None,
#[cfg(feature = "webhooks")]
event_manager: None,
shutdown_grace_period: DEFAULT_SHUTDOWN_GRACE_PERIOD,
lease_duration: DEFAULT_LEASE_DURATION,
monitoring_interval: None,
}
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn new_with_result_handler(
queue: Arc<JobQueue<DB>>,
queue_name: String,
handler: JobHandlerWithResult,
) -> Self {
Self {
queue,
queue_name,
handler: JobHandlerType::WithResult(handler),
poll_interval: Duration::from_secs(1),
max_retries: None,
retry_delay: Duration::from_secs(30),
default_retry_strategy: None,
default_timeout: None,
priority_weights: None,
stats_collector: None,
rate_limiter: None,
throttle_config: None,
concurrency_limit: None,
#[cfg(feature = "metrics")]
metrics_collector: None,
#[cfg(feature = "alerting")]
alert_manager: None,
last_job_time: Arc::new(std::sync::RwLock::new(Utc::now())),
batch_processing_enabled: false,
batch_stats: Arc::new(std::sync::RwLock::new(BatchProcessingStats::default())),
event_hooks: JobEventHooks::default(),
spawn_manager: None,
#[cfg(feature = "webhooks")]
event_manager: None,
shutdown_grace_period: DEFAULT_SHUTDOWN_GRACE_PERIOD,
lease_duration: DEFAULT_LEASE_DURATION,
monitoring_interval: None,
}
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_stats_collector(mut self, stats_collector: Arc<dyn StatisticsCollector>) -> Self {
self.stats_collector = Some(stats_collector);
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_poll_interval(mut self, interval: Duration) -> Self {
self.poll_interval = interval;
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_max_retries(mut self, max_retries: i32) -> Self {
self.max_retries = Some(max_retries);
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_retry_delay(mut self, delay: Duration) -> Self {
self.retry_delay = delay;
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_default_retry_strategy(mut self, strategy: RetryStrategy) -> Self {
self.default_retry_strategy = Some(strategy);
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_default_timeout(mut self, timeout: Duration) -> Self {
self.default_timeout = Some(timeout);
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_shutdown_grace_period(mut self, grace_period: Duration) -> Self {
self.shutdown_grace_period = grace_period;
self
}
pub fn with_lease_duration(mut self, lease: Duration) -> Self {
self.lease_duration = lease;
self
}
pub fn with_monitoring_interval(mut self, interval: Duration) -> Self {
self.monitoring_interval = Some(interval);
self
}
pub fn lease_duration(&self) -> Duration {
self.lease_duration
}
pub fn with_config(self, config: &crate::config::WorkerConfig) -> Self {
self.with_poll_interval(config.polling_interval)
.with_default_timeout(config.job_timeout)
.with_priority_weights(config.priority_weights.clone())
.with_default_retry_strategy(config.retry_strategy.clone())
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_hammerwork_config(mut self, config: &crate::HammerworkConfig) -> Self {
self = self.with_config(&config.worker);
if let Some(throttle) = config.rate_limiting.throttle_for(&self.queue_name) {
self = self.with_throttle_config(throttle);
}
#[cfg(feature = "alerting")]
{
self = self.with_alerting_config(config.alerting.clone());
}
self
}
pub fn with_priority_weights(mut self, weights: PriorityWeights) -> Self {
self.priority_weights = Some(weights);
self
}
pub fn with_strict_priority(mut self) -> Self {
self.priority_weights = Some(PriorityWeights::strict());
self
}
pub fn with_weighted_priority(mut self) -> Self {
self.priority_weights = Some(PriorityWeights::new());
self
}
pub fn with_rate_limit(mut self, rate_limit: RateLimit) -> Self {
self.rate_limiter = Some(RateLimiter::new(rate_limit));
self
}
pub fn with_throttle_config(mut self, throttle_config: ThrottleConfig) -> Self {
if !throttle_config.enabled {
return self;
}
if let Some(rate_limit) = throttle_config.to_rate_limit() {
self.rate_limiter = Some(RateLimiter::new(rate_limit));
}
self.concurrency_limit = throttle_config.max_concurrent.map(|max| {
let permits = usize::try_from(max)
.unwrap_or(usize::MAX)
.clamp(1, tokio::sync::Semaphore::MAX_PERMITS);
Arc::new(tokio::sync::Semaphore::new(permits))
});
self.throttle_config = Some(throttle_config);
self
}
#[cfg(feature = "metrics")]
pub fn with_metrics_collector(
mut self,
metrics_collector: Arc<PrometheusMetricsCollector>,
) -> Self {
self.metrics_collector = Some(metrics_collector);
self
}
#[cfg(feature = "alerting")]
pub fn with_alerting_config(mut self, alerting_config: AlertingConfig) -> Self {
self.alert_manager = Some(Arc::new(AlertManager::new(alerting_config)));
self
}
#[cfg(feature = "alerting")]
pub fn with_alert_manager(mut self, alert_manager: Arc<AlertManager>) -> Self {
self.alert_manager = Some(alert_manager);
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
#[cfg(feature = "webhooks")]
pub fn with_event_manager(mut self, event_manager: Arc<EventManager>) -> Self {
self.event_manager = Some(event_manager);
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_batch_processing_enabled(mut self, enabled: bool) -> Self {
self.batch_processing_enabled = enabled;
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn with_event_hooks(mut self, hooks: JobEventHooks) -> Self {
self.event_hooks = hooks;
self
}
pub fn with_spawn_manager(
mut self,
spawn_manager: Arc<crate::spawn::SpawnManager<DB>>,
) -> Self {
self.spawn_manager = Some(spawn_manager);
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn on_job_start<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.event_hooks.on_job_start = Some(Arc::new(handler));
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn on_job_complete<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.event_hooks.on_job_complete = Some(Arc::new(handler));
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn on_job_fail<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.event_hooks.on_job_fail = Some(Arc::new(handler));
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn on_job_timeout<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.event_hooks.on_job_timeout = Some(Arc::new(handler));
self
}
#[cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#[cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
pub fn on_job_retry<F>(mut self, handler: F) -> Self
where
F: Fn(JobHookEvent) + Send + Sync + 'static,
{
self.event_hooks.on_job_retry = Some(Arc::new(handler));
self
}
pub async fn run(&self, mut shutdown_rx: mpsc::Receiver<()>) -> Result<()> {
info!("Worker started for queue: {}", self.queue_name);
#[cfg(any(feature = "metrics", feature = "alerting"))]
let _monitoring_task = AbortOnDrop(self.start_monitoring_task());
let mut consecutive_errors: u32 = 0;
loop {
if shutdown_requested(&mut shutdown_rx) {
info!("Worker shutting down for queue: {}", self.queue_name);
break;
}
let _permit = match &self.concurrency_limit {
Some(limit) => tokio::select! {
biased;
_ = shutdown_rx.recv() => {
info!("Worker shutting down for queue: {}", self.queue_name);
break;
}
permit = Arc::clone(limit).acquire_owned() => permit.ok(),
},
None => None,
};
let acquired = self.acquire_job(&mut shutdown_rx).await;
let (outcome, shutting_down) = match acquired {
Ok(Acquired::Shutdown) => {
info!("Worker shutting down for queue: {}", self.queue_name);
break;
}
Ok(Acquired::Job(job)) => self.run_job_until_done(*job, &mut shutdown_rx).await,
Ok(Acquired::Idle) => (Ok(()), false),
Err(e) => {
error!("Error dequeuing job from queue {}: {}", self.queue_name, e);
(Err(e), false)
}
};
match outcome {
Ok(()) => consecutive_errors = 0,
Err(e) => {
consecutive_errors = consecutive_errors.saturating_add(1);
if !shutting_down {
let backoff = error_backoff(self.error_backoff_base(), consecutive_errors);
warn!(
"Worker for queue {} hit an error ({} in a row), backing off for {:?}: {}",
self.queue_name, consecutive_errors, backoff, e
);
tokio::select! {
biased;
_ = shutdown_rx.recv() => {
info!("Worker shutting down for queue: {}", self.queue_name);
break;
}
_ = sleep(backoff) => {}
}
}
}
}
if shutting_down {
info!("Worker shutting down for queue: {}", self.queue_name);
break;
}
}
Ok(())
}
async fn run_job_until_done(
&self,
job: Job,
shutdown_rx: &mut mpsc::Receiver<()>,
) -> (Result<()>, bool) {
let job_id = job.id;
let processing = self.process_job(job);
tokio::pin!(processing);
let outcome = tokio::select! {
biased;
outcome = &mut processing => return (self.log_process_outcome(job_id, outcome), false),
_ = shutdown_rx.recv() => {
info!(
"Shutdown requested while job {} is running; waiting up to {:?} for it to finish",
job_id, self.shutdown_grace_period
);
tokio::time::timeout(self.shutdown_grace_period, &mut processing).await
}
};
match outcome {
Ok(outcome) => (self.log_process_outcome(job_id, outcome), true),
Err(_) => {
warn!(
"Job {} did not finish within the {:?} shutdown grace period and was \
cancelled; it stays Running until its lease expires and the stale job \
reaper reclaims it",
job_id, self.shutdown_grace_period
);
(Ok(()), true)
}
}
}
fn log_process_outcome(&self, job_id: crate::job::JobId, outcome: Result<()>) -> Result<()> {
if let Err(ref e) = outcome {
error!(
"Failed to record the outcome of job {} on queue {}: {}",
job_id, self.queue_name, e
);
}
outcome
}
fn error_backoff_base(&self) -> Duration {
self.throttle_config
.as_ref()
.and_then(|throttle| throttle.backoff_on_error)
.unwrap_or(self.poll_interval)
}
async fn acquire_job(&self, shutdown_rx: &mut mpsc::Receiver<()>) -> Result<Acquired> {
let rate_token = match &self.rate_limiter {
Some(rate_limiter) => {
let permit = tokio::select! {
biased;
_ = shutdown_rx.recv() => return Ok(Acquired::Shutdown),
permit = rate_limiter.acquire() => permit,
};
if let Err(e) = permit {
warn!("Rate limiter error: {}", e);
return Ok(self.idle_wait(shutdown_rx).await);
}
Some(rate_limiter)
}
None => None,
};
let job = self.claim_job().await;
if !matches!(job, Ok(Some(_)))
&& let Some(rate_limiter) = rate_token
{
rate_limiter.refund();
}
match job? {
Some(job) => {
debug!(
"Processing job: {} with priority: {:?}",
job.id, job.priority
);
Ok(Acquired::Job(Box::new(job)))
}
None => Ok(self.idle_wait(shutdown_rx).await),
}
}
async fn claim_job(&self) -> Result<Option<Job>> {
match self.queue.is_queue_paused(&self.queue_name).await {
Ok(true) => {
debug!(
"Queue '{}' is paused, skipping job dequeue",
self.queue_name
);
return Ok(None);
}
Ok(false) => {}
Err(e) => {
warn!("Failed to check queue pause status: {}", e);
}
}
self.queue
.dequeue_leased(
&self.queue_name,
self.priority_weights.as_ref(),
self.lease_duration,
)
.await
}
async fn idle_wait(&self, shutdown_rx: &mut mpsc::Receiver<()>) -> Acquired {
tokio::select! {
biased;
_ = shutdown_rx.recv() => Acquired::Shutdown,
_ = sleep(self.poll_interval.max(MIN_POLL_INTERVAL)) => Acquired::Idle,
}
}
fn attempt_limit(&self, job: &Job) -> i32 {
effective_attempt_limit(job.max_attempts, self.max_retries)
}
fn retry_at_for(&self, job: &Job) -> DateTime<Utc> {
let next_attempt = u32::try_from(job.attempts.saturating_add(1)).unwrap_or(u32::MAX);
let retry_delay = if let Some(ref job_strategy) = job.retry_strategy {
job_strategy.calculate_delay(next_attempt)
} else if let Some(ref default_strategy) = self.default_retry_strategy {
default_strategy.calculate_delay(next_attempt)
} else {
self.retry_delay
};
retry_at_from_delay(Utc::now(), retry_delay)
}
async fn record_outcome(
&self,
job: &Job,
outcome: JobOutcome,
) -> Result<Option<RecordedOutcome>> {
self.record(job, outcome, None).await
}
async fn record(
&self,
job: &Job,
outcome: JobOutcome,
children: Option<Vec<Job>>,
) -> Result<Option<RecordedOutcome>> {
let queue = Arc::clone(&self.queue);
let finished_job = job.clone();
let recorded = tokio::spawn(async move {
match children {
Some(children) => {
queue
.complete_job_run_with_children(&finished_job, children)
.await
}
None => queue.finish_job_run(&finished_job, outcome).await,
}
})
.await
.map_err(|e| HammerworkError::Worker {
message: format!("recording the outcome of job {} failed: {e}", job.id),
})??;
match recorded {
Some(recorded) => {
if let Some(next_run_at) = recorded.next_run_at {
info!(
"Rescheduled recurring job {} for its next run at {}",
job.id, next_run_at
);
}
if !recorded.unblocked.is_empty() {
debug!(
"Job {} completed; {} dependent job(s) are now runnable",
job.id,
recorded.unblocked.len()
);
}
if !recorded.cancelled.is_empty() {
warn!(
"Job {} failed; {} dependent, workflow or batch job(s) were failed with it",
job.id,
recorded.cancelled.len()
);
}
Ok(Some(recorded))
}
None => {
warn!(
"Job {} is no longer running attempt {} (it was reclaimed as stale, \
changed by an operator or is running again elsewhere); discarding the \
outcome of this run",
job.id, job.attempts
);
Ok(None)
}
}
}
async fn process_job(&self, job: Job) -> Result<()> {
let job_id = job.id;
let batch_id = job.batch_id;
let start_time = Utc::now();
#[cfg(feature = "tracing")]
let _span = crate::tracing::create_job_span(&job, "job.process");
if self.batch_processing_enabled && batch_id.is_some() {
self.update_batch_stats(|stats| {
stats.jobs_processed += 1;
stats.last_processed_job = Some(start_time);
});
}
self.event_hooks.fire_job_start(job.clone());
self.record_event(JobEvent {
job_id,
queue_name: job.queue_name.clone(),
event_type: JobEventType::Started,
priority: job.priority,
processing_time_ms: None,
error_message: None,
timestamp: start_time,
})
.await;
let attempts_left = job.attempts < self.attempt_limit(&job);
let handler_job = match self.queue.decrypt_job(job.clone()).await {
Ok(handler_job) => handler_job,
Err(e) => return self.handle_failure(&job, e, attempts_left).await,
};
let spawn_parent = self.spawn_manager.is_some().then(|| handler_job.clone());
let timeout_duration = job
.timeout
.filter(|timeout| !timeout.is_zero())
.or(self.default_timeout.filter(|timeout| !timeout.is_zero()));
let handler_result = match timeout_duration {
Some(timeout) => tokio::time::timeout(timeout, self.execute_handler(handler_job))
.await
.map_err(|_| timeout),
None => Ok(self.execute_handler(handler_job).await),
};
match handler_result {
Ok(Ok(job_result)) => {
self.handle_success(&job, spawn_parent, job_result, start_time, attempts_left)
.await
}
Ok(Err(e)) => self.handle_failure(&job, e, attempts_left).await,
Err(timeout) => self.handle_timeout(&job, timeout, attempts_left).await,
}
}
async fn prepare_spawn(
&self,
job: &Job,
spawn_parent: Option<&Job>,
) -> Result<Option<(Vec<Job>, crate::spawn::SpawnConfig)>> {
let Some(spawn_manager) = &self.spawn_manager else {
return Ok(None);
};
let parent = spawn_parent.unwrap_or(job);
let Some(config) = crate::spawn::spawn_config_from_payload(&parent.payload)? else {
return Ok(None);
};
match spawn_manager
.prepare_spawn(parent, &config, self.queue.clone())
.await?
{
Some(children) => Ok(Some((children, config))),
None => {
debug!(
"No spawn handler registered for job type: {}",
job.queue_name
);
Ok(None)
}
}
}
async fn handle_success(
&self,
job: &Job,
spawn_parent: Option<Job>,
job_result: JobResult,
start_time: DateTime<Utc>,
attempts_left: bool,
) -> Result<()> {
let job_id = job.id;
debug!("Job {} completed successfully", job_id);
let processing_time_ms = (Utc::now() - start_time).num_milliseconds().max(0) as u64;
if let Some(result_data) = job_result.data {
let expires_at = job
.result_config
.ttl
.map(|ttl| crate::queue::saturating_add_to(Utc::now(), ttl));
match job.result_config.storage {
crate::job::ResultStorage::Database => {
if let Err(e) = self
.queue
.store_job_result(job_id, result_data, expires_at)
.await
{
error!("Failed to store result for job {}: {}", job_id, e);
} else {
debug!("Stored result for job {}", job_id);
}
}
crate::job::ResultStorage::Memory => {
self.queue
.memory_results
.store(job_id, result_data, expires_at);
debug!("Stored result for job {} in memory", job_id);
}
crate::job::ResultStorage::None => {}
}
}
let spawn = match self.prepare_spawn(job, spawn_parent.as_ref()).await {
Ok(spawn) => spawn,
Err(e) => {
warn!("Failed to spawn child jobs for job {}: {}", job_id, e);
return self.handle_failure(job, e, attempts_left).await;
}
};
let (children, spawn_config) = match spawn {
Some((children, config)) => (Some(children), Some(config)),
None => (None, None),
};
let recorded = match self.record(job, JobOutcome::Completed, children).await {
Ok(Some(recorded)) => recorded,
Ok(None) => return Ok(()),
Err(e) if spawn_config.is_some() => {
warn!("Failed to enqueue the child jobs of job {}: {}", job_id, e);
return self.handle_failure(job, e, attempts_left).await;
}
Err(e) => return Err(e),
};
if let (Some(config), Some(spawn_manager)) = (spawn_config, &self.spawn_manager) {
let spawn_result = crate::spawn::SpawnResult {
parent_job_id: job_id,
spawned_jobs: recorded.spawned.clone(),
spawned_at: Utc::now(),
spawn_operation_id: config.operation_id,
};
info!(
"Job {} spawned {} child jobs: {:?}",
job_id,
spawn_result.spawned_jobs.len(),
spawn_result.spawned_jobs
);
let parent = spawn_parent.as_ref().unwrap_or(job);
if let Err(e) = spawn_manager
.notify_spawn_complete(parent, &spawn_result)
.await
{
warn!("Spawn completion callback for job {} failed: {}", job_id, e);
}
}
if self.batch_processing_enabled
&& let Some(batch_id) = job.batch_id
{
self.update_batch_stats(|stats| {
stats.jobs_completed += 1;
stats.total_processing_time_ms += processing_time_ms;
stats.update_average_processing_time();
});
if let Err(e) = self.check_and_update_batch_status(batch_id).await {
warn!(
"Failed to update batch status for batch {}: {}",
batch_id, e
);
}
}
#[cfg(feature = "tracing")]
{
let span = tracing::Span::current();
span.record("success", true);
span.record("processing_time_ms", processing_time_ms);
}
let processing_duration = Duration::from_millis(processing_time_ms);
self.event_hooks
.fire_job_complete(job.clone(), processing_duration);
self.record_event(JobEvent {
job_id,
queue_name: job.queue_name.clone(),
event_type: JobEventType::Completed,
priority: job.priority,
processing_time_ms: Some(processing_time_ms),
error_message: None,
timestamp: Utc::now(),
})
.await;
Ok(())
}
async fn handle_failure(
&self,
job: &Job,
error: HammerworkError,
attempts_left: bool,
) -> Result<()> {
let job_id = job.id;
error!("Job {} failed: {}", job_id, error);
let error_message = error.to_string();
#[cfg(feature = "tracing")]
{
let span = tracing::Span::current();
span.record("error", true);
span.record("error.type", "job_failure");
span.record("error.message", &error_message);
span.record("job.will_retry", attempts_left);
}
let outcome = if attempts_left {
JobOutcome::Retry {
retry_at: self.retry_at_for(job),
error: error_message.clone(),
timed_out: false,
}
} else {
warn!(
"Job {} failed on its last attempt ({} of {}), marking it dead",
job_id,
job.attempts,
self.attempt_limit(job)
);
JobOutcome::Dead {
error: error_message.clone(),
}
};
let Some(recorded) = self.record_outcome(job, outcome).await? else {
return Ok(());
};
self.after_failed_run(job, &error_message, attempts_left)
.await;
if attempts_left {
info!(
"Retrying job {} (attempt {} of {})",
job_id,
job.attempts + 1,
self.attempt_limit(job)
);
self.event_hooks
.fire_job_retry(job.clone(), error_message.clone());
self.record_event(JobEvent {
job_id,
queue_name: job.queue_name.clone(),
event_type: JobEventType::Retried,
priority: job.priority,
processing_time_ms: None,
error_message: Some(error_message),
timestamp: Utc::now(),
})
.await;
} else {
debug!("Job {} run ended with status {:?}", job_id, recorded.status);
self.event_hooks
.fire_job_fail(job.clone(), error_message.clone());
self.record_event(JobEvent {
job_id,
queue_name: job.queue_name.clone(),
event_type: JobEventType::Dead,
priority: job.priority,
processing_time_ms: None,
error_message: Some(error_message),
timestamp: Utc::now(),
})
.await;
}
Ok(())
}
async fn handle_timeout(
&self,
job: &Job,
timeout: Duration,
attempts_left: bool,
) -> Result<()> {
let job_id = job.id;
let error_message = format!("Job timed out after {:?}", timeout);
warn!("Job {} timed out after {:?}", job_id, timeout);
let outcome = if attempts_left {
JobOutcome::Retry {
retry_at: self.retry_at_for(job),
error: error_message.clone(),
timed_out: true,
}
} else {
JobOutcome::TimedOut {
error: error_message.clone(),
}
};
if self.record_outcome(job, outcome).await?.is_none() {
return Ok(());
}
#[cfg(feature = "tracing")]
{
let span = tracing::Span::current();
span.record("error", true);
span.record("error.type", "timeout");
span.record("error.message", &error_message);
span.record("job.will_retry", attempts_left);
}
self.after_failed_run(job, &error_message, attempts_left)
.await;
self.event_hooks.fire_job_timeout(job.clone(), timeout);
self.record_event(JobEvent {
job_id,
queue_name: job.queue_name.clone(),
event_type: JobEventType::TimedOut,
priority: job.priority,
processing_time_ms: Some(timeout.as_millis() as u64),
error_message: Some(error_message.clone()),
timestamp: Utc::now(),
})
.await;
if attempts_left {
info!(
"Retrying timed-out job {} (attempt {} of {})",
job_id,
job.attempts + 1,
self.attempt_limit(job)
);
self.event_hooks
.fire_job_retry(job.clone(), error_message.clone());
self.record_event(JobEvent {
job_id,
queue_name: job.queue_name.clone(),
event_type: JobEventType::Retried,
priority: job.priority,
processing_time_ms: None,
error_message: Some(error_message),
timestamp: Utc::now(),
})
.await;
}
Ok(())
}
async fn after_failed_run(&self, job: &Job, error_message: &str, attempts_left: bool) {
if !self.batch_processing_enabled {
return;
}
let Some(batch_id) = job.batch_id else {
return;
};
self.update_batch_stats(|stats| {
stats.jobs_failed += 1;
});
if attempts_left {
return;
}
if let Err(e) = self
.handle_batch_job_failure(batch_id, job.id, error_message)
.await
{
warn!(
"Failed to handle batch job failure for batch {}: {}",
batch_id, e
);
}
if let Err(e) = self.check_and_update_batch_status(batch_id).await {
warn!(
"Failed to update batch status for batch {}: {}",
batch_id, e
);
}
}
fn start_handler(&self, job: Job) -> HandlerFuture {
match &self.handler {
JobHandlerType::Legacy(handler) => {
let future = handler(job);
Box::pin(async move { future.await.map(|_| JobResult::success()) })
}
JobHandlerType::WithResult(handler) => handler(job),
}
}
async fn execute_handler(&self, job: Job) -> Result<JobResult> {
let job_id = job.id;
let run = Job {
payload: serde_json::Value::Null,
..job.clone()
};
let handler_future =
match std::panic::catch_unwind(AssertUnwindSafe(|| self.start_handler(job))) {
Ok(future) => future,
Err(payload) => {
error!("Handler for job {} panicked", job_id);
return Err(handler_panic_error(payload.as_ref()));
}
};
let guarded = CatchUnwind {
inner: handler_future,
};
tokio::select! {
biased;
outcome = guarded => outcome.unwrap_or_else(|payload| {
error!("Handler for job {} panicked", job_id);
Err(handler_panic_error(payload.as_ref()))
}),
() = self.maintain_lease(&run) => Err(HammerworkError::Worker {
message: "job lease heartbeat stopped unexpectedly".to_string(),
}),
}
}
async fn maintain_lease(&self, run: &Job) {
let job_id = run.id;
if self.lease_duration.is_zero() {
return std::future::pending().await;
}
let interval = (self.lease_duration / 3).max(MIN_HEARTBEAT_INTERVAL);
loop {
sleep(interval).await;
match self.queue.heartbeat_job(run, self.lease_duration).await {
Ok(true) => debug!("Renewed lease on job {}", job_id),
Ok(false) => {
warn!(
"Lost the lease on job {}: it is no longer Running this run (it may \
have been reclaimed as stale); the handler keeps running",
job_id
);
return std::future::pending().await;
}
Err(e) => warn!("Failed to renew the lease on job {}: {}", job_id, e),
}
}
}
fn update_batch_stats<F>(&self, updater: F)
where
F: FnOnce(&mut BatchProcessingStats),
{
if let Ok(mut stats) = self.batch_stats.write() {
updater(&mut stats);
}
}
#[cfg(feature = "webhooks")]
fn convert_to_lifecycle_event(&self, event: &JobEvent) -> JobLifecycleEvent {
use std::collections::HashMap;
let event_type = match event.event_type {
JobEventType::Started => JobLifecycleEventType::Started,
JobEventType::Completed => JobLifecycleEventType::Completed,
JobEventType::Failed => JobLifecycleEventType::Failed,
JobEventType::Retried => JobLifecycleEventType::Retried,
JobEventType::Dead => JobLifecycleEventType::Dead,
JobEventType::TimedOut => JobLifecycleEventType::TimedOut,
};
let error = event.error_message.as_ref().map(|message| JobError {
message: message.clone(),
error_type: match event.event_type {
JobEventType::TimedOut => Some("timeout".to_string()),
JobEventType::Failed => Some("processing_error".to_string()),
JobEventType::Dead => Some("max_retries_exceeded".to_string()),
_ => None,
},
details: None,
retry_attempt: None,
});
let mut metadata = HashMap::new();
metadata.insert("worker_queue".to_string(), self.queue_name.clone());
if let Some(processing_time) = event.processing_time_ms {
metadata.insert(
"processing_time_ms".to_string(),
processing_time.to_string(),
);
}
JobLifecycleEvent {
event_id: uuid::Uuid::new_v4(),
job_id: event.job_id,
queue_name: event.queue_name.clone(),
event_type,
priority: event.priority,
timestamp: event.timestamp,
processing_time_ms: event.processing_time_ms,
error,
payload: None, metadata,
}
}
async fn check_and_update_batch_status(&self, batch_id: BatchId) -> Result<()> {
let batch_result = self.queue.get_batch_status(batch_id).await?;
if batch_result.pending_jobs == 0 {
let completion_rate = batch_result.success_rate();
if completion_rate >= 0.95 {
info!(
"Batch {} completed successfully with {:.1}% success rate",
batch_id,
completion_rate * 100.0
);
} else {
warn!(
"Batch {} completed with {:.1}% success rate ({} failures)",
batch_id,
completion_rate * 100.0,
batch_result.failed_jobs
);
}
self.update_batch_stats(|stats| {
stats.batches_completed += 1;
if completion_rate >= 0.95 {
stats.batches_successful += 1;
}
});
}
Ok(())
}
async fn handle_batch_job_failure(
&self,
batch_id: BatchId,
job_id: uuid::Uuid,
error_message: &str,
) -> Result<()> {
let batch_result = self.queue.get_batch_status(batch_id).await?;
warn!(
"Job {} in batch {} failed: {}. Batch status: {}/{} jobs remaining",
job_id, batch_id, error_message, batch_result.pending_jobs, batch_result.total_jobs
);
Ok(())
}
pub fn get_batch_stats(&self) -> BatchProcessingStats {
if let Ok(stats) = self.batch_stats.read() {
stats.clone()
} else {
BatchProcessingStats::default()
}
}
async fn record_event(&self, event: JobEvent) {
if let Some(stats_collector) = &self.stats_collector
&& let Err(e) = stats_collector.record_event(event.clone()).await
{
warn!("Failed to record statistics event: {}", e);
}
#[cfg(feature = "metrics")]
if let Some(metrics_collector) = &self.metrics_collector
&& let Err(e) = metrics_collector.record_job_event(&event).await
{
warn!("Failed to record metrics event: {}", e);
}
#[cfg(feature = "webhooks")]
if let Some(event_manager) = &self.event_manager {
let lifecycle_event = self.convert_to_lifecycle_event(&event);
if let Err(e) = event_manager.publish_event(lifecycle_event).await {
warn!("Failed to publish lifecycle event: {}", e);
}
}
if matches!(
event.event_type,
JobEventType::Completed
| JobEventType::Failed
| JobEventType::Dead
| JobEventType::TimedOut
) && let Ok(mut last_time) = self.last_job_time.write()
{
*last_time = event.timestamp;
}
}
#[cfg(any(feature = "metrics", feature = "alerting"))]
fn start_monitoring_task(&self) -> tokio::task::JoinHandle<()> {
let queue_name = self.queue_name.clone();
let queue = Arc::clone(&self.queue);
#[cfg(feature = "alerting")]
let last_job_time = Arc::clone(&self.last_job_time);
#[cfg(feature = "metrics")]
let metrics_collector = self.metrics_collector.clone();
#[cfg(feature = "alerting")]
let alert_manager = self.alert_manager.clone();
#[cfg(feature = "alerting")]
let stats_collector = self.stats_collector.clone();
#[cfg(feature = "metrics")]
let collector_interval = metrics_collector.as_ref().map(|m| m.update_interval());
#[cfg(not(feature = "metrics"))]
let collector_interval: Option<Duration> = None;
let monitor_every = self
.monitoring_interval
.or(collector_interval)
.unwrap_or(DEFAULT_MONITORING_INTERVAL);
tokio::spawn(async move {
let mut interval = tokio::time::interval(monitor_every.max(MIN_MONITORING_INTERVAL));
loop {
interval.tick().await;
#[cfg(feature = "metrics")]
let wants_depth = metrics_collector.is_some();
#[cfg(not(feature = "metrics"))]
let wants_depth = false;
#[cfg(feature = "alerting")]
let wants_depth = wants_depth || alert_manager.is_some();
let queue_depth = if wants_depth {
match queue.get_queue_depth(&queue_name).await {
Ok(depth) => Some(depth),
Err(e) => {
warn!("Failed to get queue depth for monitoring: {}", e);
None
}
}
} else {
None
};
#[cfg(feature = "metrics")]
if let (Some(metrics_collector), Some(queue_depth)) =
(&metrics_collector, queue_depth)
&& let Err(e) = metrics_collector
.update_queue_depth(&queue_name, queue_depth)
.await
{
warn!("Failed to update queue depth metrics: {}", e);
}
#[cfg(feature = "alerting")]
if let (Some(alert_manager), Some(queue_depth)) = (&alert_manager, queue_depth)
&& let Err(e) = alert_manager
.check_queue_depth(&queue_name, queue_depth)
.await
{
warn!("Failed to check queue depth alerts: {}", e);
}
#[cfg(feature = "alerting")]
if let Some(alert_manager) = &alert_manager {
let last_time_value = {
if let Ok(last_time) = last_job_time.read() {
Some(*last_time)
} else {
None
}
};
if let Some(last_time_value) = last_time_value
&& let Err(e) = alert_manager
.check_worker_starvation(&queue_name, last_time_value)
.await
{
warn!("Failed to check worker starvation: {}", e);
}
}
#[cfg(feature = "alerting")]
if let (Some(alert_manager), Some(stats_collector)) =
(&alert_manager, &stats_collector)
{
match stats_collector
.get_queue_statistics(&queue_name, Duration::from_secs(300))
.await
{
Ok(stats) => {
if let Err(e) =
alert_manager.check_thresholds(&queue_name, &stats).await
{
warn!("Failed to check alert thresholds: {}", e);
}
}
Err(e) => {
warn!("Failed to get queue statistics for alerting: {}", e);
}
}
}
}
})
}
}
pub struct WorkerPool<DB: Database> {
workers: Vec<Worker<DB>>,
shutdown_signal: Option<watch::Sender<bool>>,
supervisor_done: Option<watch::Receiver<bool>>,
reaper_interval: Option<Duration>,
reaper_older_than: Option<Duration>,
encrypted_job_purge_interval: Option<Duration>,
archival: Option<PoolArchival>,
stats_collector: Option<Arc<dyn StatisticsCollector>>,
worker_template: Option<Worker<DB>>,
autoscale_config: AutoscaleConfig,
autoscale_metrics: Arc<std::sync::RwLock<AutoscaleMetrics>>,
queue_depth_history: QueueDepthHistory,
autoscale_task: Option<tokio::task::JoinHandle<()>>,
}
impl<DB: Database + Send + Sync + 'static> WorkerPool<DB>
where
JobQueue<DB>: DatabaseQueue<Database = DB> + Send + Sync,
{
pub fn new() -> Self {
Self {
workers: Vec::new(),
shutdown_signal: None,
supervisor_done: None,
reaper_interval: Some(DEFAULT_STALE_JOB_REAPER_INTERVAL),
reaper_older_than: None,
encrypted_job_purge_interval: None,
archival: None,
stats_collector: None,
worker_template: None,
autoscale_config: AutoscaleConfig::disabled(),
autoscale_metrics: Arc::new(std::sync::RwLock::new(AutoscaleMetrics::default())),
queue_depth_history: Arc::new(std::sync::RwLock::new(Vec::new())),
autoscale_task: None,
}
}
pub fn from_config(worker: Worker<DB>, config: &crate::config::WorkerConfig) -> Self {
let mut pool = Self::new()
.with_autoscaling(config.autoscale_config())
.with_worker_template(worker.clone());
for _ in 1..config.pool_size {
pool.add_worker(worker.clone());
}
pool.add_worker(worker);
pool
}
pub fn from_hammerwork_config(
worker: Worker<DB>,
config: &crate::config::HammerworkConfig,
) -> Result<Self> {
config.encryption.validate()?;
config.worker.validate()?;
config.archive.validate()?;
let mut pool = Self::from_config(worker, &config.worker);
pool.encrypted_job_purge_interval = config.encryption.purge_interval();
if config.archive.enabled {
pool = pool.with_archival(
config.archive.archival_policy(),
config.archive.archival_config(),
config.archive.check_interval,
);
}
Ok(pool)
}
pub fn with_stats_collector(mut self, stats_collector: Arc<dyn StatisticsCollector>) -> Self {
self.stats_collector = Some(stats_collector);
self
}
pub fn with_autoscaling(mut self, config: AutoscaleConfig) -> Self {
self.autoscale_config = config;
self
}
pub fn without_autoscaling(mut self) -> Self {
self.autoscale_config = AutoscaleConfig::disabled();
self
}
pub fn with_stale_job_reaper(mut self, interval: Duration, older_than: Duration) -> Self {
self.reaper_interval = Some(interval);
self.reaper_older_than = Some(older_than);
self
}
pub fn without_stale_job_reaper(mut self) -> Self {
self.reaper_interval = None;
self
}
pub fn with_encrypted_job_purge(mut self, interval: Duration) -> Self {
self.encrypted_job_purge_interval = Some(interval);
self
}
pub fn without_encrypted_job_purge(mut self) -> Self {
self.encrypted_job_purge_interval = None;
self
}
pub fn with_archival(
mut self,
policy: crate::archive::ArchivalPolicy,
config: crate::archive::ArchivalConfig,
interval: Duration,
) -> Self {
self.archival = Some(PoolArchival {
policy,
config,
interval,
});
self
}
pub fn without_archival(mut self) -> Self {
self.archival = None;
self
}
pub fn with_worker_template(mut self, worker: Worker<DB>) -> Self {
self.worker_template = Some(worker);
self
}
pub fn add_worker(&mut self, mut worker: Worker<DB>) {
if let Some(stats_collector) = &self.stats_collector {
worker.stats_collector = Some(Arc::clone(stats_collector));
}
if self.worker_template.is_none()
&& self.autoscale_config.enabled
&& self.workers.is_empty()
{
self.worker_template = Some(worker.clone());
}
self.workers.push(worker);
}
pub async fn start(&mut self) -> Result<()> {
info!("Starting worker pool with {} workers", self.workers.len());
let (signal_tx, signal_rx) = watch::channel(false);
let (done_tx, done_rx) = watch::channel(false);
self.shutdown_signal = Some(signal_tx);
self.supervisor_done = Some(done_rx);
let workers = std::mem::take(&mut self.workers);
let maintenance: Vec<tokio::task::JoinHandle<()>> = [
self.start_stale_job_reaper(&workers, signal_rx.clone()),
self.start_encrypted_job_purge(&workers, signal_rx.clone()),
self.start_archival(&workers, signal_rx.clone()),
]
.into_iter()
.flatten()
.collect();
let scaling = if self.autoscale_config.enabled {
self.start_autoscaling_task(&workers)
} else {
None
};
if let Ok(mut metrics) = self.autoscale_metrics.write() {
metrics.active_workers = match &scaling {
Some(scaling) => scaling.initial,
None => workers.len(),
};
}
let supervisor = tokio::spawn(Self::supervise(
workers,
signal_rx,
done_tx,
maintenance,
scaling,
));
supervisor.await.map_err(|e| HammerworkError::Worker {
message: format!("Worker pool supervisor failed: {}", e),
})?
}
fn spawn_worker(
set: &mut JoinSet<(usize, std::thread::Result<Result<()>>)>,
index: usize,
worker: Worker<DB>,
delay: Option<Duration>,
) -> mpsc::Sender<()> {
let (shutdown_tx, shutdown_rx) = mpsc::channel(1);
set.spawn(async move {
if let Some(delay) = delay {
sleep(delay).await;
}
let run = CatchUnwind {
inner: Box::pin(async move { worker.run(shutdown_rx).await }),
};
(index, run.await)
});
shutdown_tx
}
async fn supervise(
workers: Vec<Worker<DB>>,
mut signal_rx: watch::Receiver<bool>,
done_tx: watch::Sender<bool>,
maintenance: Vec<tokio::task::JoinHandle<()>>,
scaling: Option<Scaling<DB>>,
) -> Result<()> {
let mut templates: Vec<Worker<DB>> = workers.to_vec();
let mut set = JoinSet::new();
let mut senders: Vec<mpsc::Sender<()>> = workers
.into_iter()
.enumerate()
.map(|(index, worker)| Self::spawn_worker(&mut set, index, worker, None))
.collect();
let mut retired: Vec<bool> = vec![false; senders.len()];
let mut finished: Vec<bool> = vec![false; senders.len()];
let (template, mut desired_rx) = match scaling {
Some(scaling) => (Some(scaling.template), Some(scaling.desired)),
None => (None, None),
};
let mut shutting_down = *signal_rx.borrow();
if shutting_down {
for tx in &senders {
let _ = tx.try_send(());
}
}
loop {
tokio::select! {
_ = signal_rx.changed(), if !shutting_down => {
info!("Worker pool shutting down; waiting for in-flight jobs");
shutting_down = true;
for tx in &senders {
let _ = tx.try_send(());
}
}
changed = wait_for_change(&mut desired_rx), if !shutting_down => {
let (Some(desired), Some(template)) = (changed, template.as_ref()) else {
desired_rx = None;
continue;
};
let active: Vec<usize> = (0..senders.len())
.filter(|i| !retired[*i] && templates[*i].queue_name == template.queue_name)
.collect();
if desired > active.len() {
for _ in active.len()..desired {
let worker = template.clone();
let reusable = (0..senders.len()).find(|i| retired[*i] && finished[*i]);
let index = reusable.unwrap_or(senders.len());
let tx = Self::spawn_worker(&mut set, index, worker.clone(), None);
if index == senders.len() {
templates.push(worker);
senders.push(tx);
retired.push(false);
finished.push(false);
} else {
templates[index] = worker;
senders[index] = tx;
retired[index] = false;
finished[index] = false;
}
}
info!(
"Autoscaling: started workers for queue {}, now {}",
template.queue_name, desired
);
} else {
for index in active.iter().rev().take(active.len() - desired) {
retired[*index] = true;
let _ = senders[*index].try_send(());
}
if desired < active.len() {
info!(
"Autoscaling: retiring workers of queue {}, now {}",
template.queue_name, desired
);
}
}
}
joined = set.join_next() => {
let Some(joined) = joined else { break };
let index = match joined {
Ok((index, Ok(Ok(())))) => {
if shutting_down {
continue;
}
if retired[index] {
finished[index] = true;
continue;
}
warn!("Worker {} stopped unexpectedly", index);
index
}
Ok((index, Ok(Err(e)))) => {
error!("Worker {} failed: {}", index, e);
index
}
Ok((index, Err(payload))) => {
error!(
"Worker {} panicked: {}",
index,
panic_message(payload.as_ref())
);
index
}
Err(e) => {
error!("Worker task ended abnormally: {}", e);
continue;
}
};
if retired[index] {
finished[index] = true;
} else if !shutting_down {
info!("Restarting worker {} in {:?}", index, WORKER_RESTART_DELAY);
senders[index] = Self::spawn_worker(
&mut set,
index,
templates[index].clone(),
Some(WORKER_RESTART_DELAY),
);
}
}
}
}
for task in maintenance {
task.abort();
}
info!("All workers stopped");
done_tx.send_replace(true);
Ok(())
}
fn start_stale_job_reaper(
&self,
workers: &[Worker<DB>],
signal_rx: watch::Receiver<bool>,
) -> Option<tokio::task::JoinHandle<()>> {
let interval = self.reaper_interval?;
let queue = Arc::clone(&workers.first()?.queue);
let older_than = self.reaper_older_than.unwrap_or_else(|| {
workers
.iter()
.map(|worker| worker.lease_duration)
.max()
.filter(|lease| !lease.is_zero())
.unwrap_or(DEFAULT_LEASE_DURATION)
});
Some(spawn_periodic(interval, signal_rx, move || {
let queue = Arc::clone(&queue);
async move {
match queue.requeue_stale_jobs(older_than).await {
Ok(recovery) if !recovery.is_empty() => info!(
"Stale job reaper requeued {} and marked {} dead",
recovery.requeued.len(),
recovery.dead.len()
),
Ok(_) => {}
Err(e) => warn!("Stale job reaper failed: {}", e),
}
}
}))
}
fn start_encrypted_job_purge(
&self,
workers: &[Worker<DB>],
signal_rx: watch::Receiver<bool>,
) -> Option<tokio::task::JoinHandle<()>> {
let interval = self.encrypted_job_purge_interval?;
let queue = Arc::clone(&workers.first()?.queue);
Some(spawn_periodic(interval, signal_rx, move || {
let queue = Arc::clone(&queue);
async move {
match queue.purge_expired_encrypted_jobs().await {
Ok(purge) if purge.total() > 0 => info!(
"Encrypted job purge deleted {} jobs and {} archived jobs",
purge.jobs, purge.archived_jobs
),
Ok(_) => {}
Err(e) => warn!("Encrypted job purge failed: {}", e),
}
}
}))
}
fn start_archival(
&self,
workers: &[Worker<DB>],
signal_rx: watch::Receiver<bool>,
) -> Option<tokio::task::JoinHandle<()>> {
let archival = self.archival.clone()?;
let queue = Arc::clone(&workers.first()?.queue);
if !archival.policy.enabled {
info!("Automatic archival is configured with a disabled policy; not starting it");
return None;
}
let archiver = Arc::new(
crate::archive::JobArchiver::new(queue.pool.clone())
.with_default_policy(archival.policy)
.with_config(archival.config),
);
let stop_rx = signal_rx.clone();
Some(spawn_periodic(archival.interval, signal_rx, move || {
let (queue, archiver, stop_rx) =
(Arc::clone(&queue), Arc::clone(&archiver), stop_rx.clone());
async move {
let pass = archiver
.run_scheduled_pass(queue.as_ref(), Some("worker_pool"), || *stop_rx.borrow())
.await;
match pass {
Ok(pass) if pass.jobs_archived > 0 || pass.jobs_purged > 0 => info!(
"Archival archived {} jobs and purged {} archived jobs{}",
pass.jobs_archived,
pass.jobs_purged,
if pass.stopped_early {
" (stopped early for shutdown)"
} else {
""
}
),
Ok(_) => debug!("Archival found nothing to archive or purge"),
Err(e) => warn!("Archival failed: {}", e),
}
}
}))
}
fn start_autoscaling_task(&mut self, workers: &[Worker<DB>]) -> Option<Scaling<DB>> {
let Some(worker_template) = self.worker_template.as_ref().or(workers.first()) else {
warn!("Cannot start autoscaling: no worker template available");
return None;
};
let mut template = worker_template.clone();
let initial_workers = workers
.iter()
.filter(|worker| worker.queue_name == template.queue_name)
.count();
if let Some(stats_collector) = &self.stats_collector {
template.stats_collector = Some(Arc::clone(stats_collector));
}
let queue = Arc::clone(&template.queue);
let queue_name = template.queue_name.clone();
let config = self.autoscale_config.clone();
let metrics = Arc::clone(&self.autoscale_metrics);
let history = Arc::clone(&self.queue_depth_history);
let (desired_tx, desired_rx) = watch::channel(initial_workers);
let task = tokio::spawn(async move {
Self::autoscaling_loop(queue, queue_name, config, metrics, history, desired_tx).await;
});
if let Some(previous) = self.autoscale_task.replace(task) {
previous.abort();
}
info!(
"Autoscaling task started for queue: {}",
template.queue_name
);
Some(Scaling {
template,
desired: desired_rx,
initial: initial_workers,
})
}
async fn autoscaling_loop(
queue: Arc<JobQueue<DB>>,
queue_name: String,
config: AutoscaleConfig,
metrics: Arc<std::sync::RwLock<AutoscaleMetrics>>,
history: QueueDepthHistory,
desired: watch::Sender<usize>,
) {
let mut interval =
tokio::time::interval((config.evaluation_window / 2).max(Duration::from_millis(1)));
loop {
interval.tick().await;
match Self::evaluate_scaling_decision(&queue, &queue_name, &config, &metrics, &history)
.await
{
Ok(Some(count)) => {
if desired.send(count).is_err() {
return;
}
}
Ok(None) => {}
Err(e) => warn!("Autoscaling evaluation error: {}", e),
}
}
}
async fn evaluate_scaling_decision(
queue: &Arc<JobQueue<DB>>,
queue_name: &str,
config: &AutoscaleConfig,
metrics: &Arc<std::sync::RwLock<AutoscaleMetrics>>,
history: &QueueDepthHistory,
) -> Result<Option<usize>> {
let current_depth = queue.get_queue_depth(queue_name).await?;
Ok(apply_queue_depth(
current_depth,
Utc::now(),
config,
metrics,
history,
))
}
pub fn get_autoscale_metrics(&self) -> AutoscaleMetrics {
if let Ok(metrics) = self.autoscale_metrics.read() {
metrics.clone()
} else {
AutoscaleMetrics::default()
}
}
pub async fn shutdown(&self) -> Result<()> {
info!("Shutting down worker pool");
if let Some(task) = &self.autoscale_task {
task.abort();
info!("Autoscaling task stopped");
}
if let Some(signal) = &self.shutdown_signal {
signal.send_replace(true);
}
if let Some(done) = &self.supervisor_done {
let mut done = done.clone();
let _ = done.wait_for(|finished| *finished).await;
}
Ok(())
}
pub fn stats_collector(&self) -> Option<Arc<dyn StatisticsCollector>> {
self.stats_collector.clone()
}
}
impl<DB: Database + Send + Sync + 'static> Default for WorkerPool<DB>
where
JobQueue<DB>: DatabaseQueue<Database = DB> + Send + Sync,
{
fn default() -> Self {
Self::new()
}
}
fn apply_queue_depth(
current_depth: u64,
now: DateTime<Utc>,
config: &AutoscaleConfig,
metrics: &Arc<std::sync::RwLock<AutoscaleMetrics>>,
history: &QueueDepthHistory,
) -> Option<usize> {
let avg_depth = match history.write() {
Ok(mut hist) => {
hist.push((now, current_depth));
let cutoff = now
- chrono::Duration::from_std(config.evaluation_window)
.unwrap_or(chrono::Duration::seconds(30));
hist.retain(|(timestamp, _)| *timestamp > cutoff);
if hist.is_empty() {
current_depth as f64
} else {
hist.iter().map(|(_, depth)| *depth as f64).sum::<f64>() / hist.len() as f64
}
}
Err(_) => current_depth as f64,
};
let mut m = metrics.write().ok()?;
m.current_queue_depth = current_depth;
m.avg_queue_depth = avg_depth;
let time_since_last = m
.last_scale_time
.map(|t| now - t)
.and_then(|d| d.to_std().ok())
.unwrap_or(config.cooldown_period);
m.time_since_last_scale = time_since_last;
if time_since_last < config.cooldown_period {
return None;
}
let depth_per_worker = if m.active_workers > 0 {
avg_depth / m.active_workers as f64
} else {
avg_depth
};
let decision = if depth_per_worker > config.scale_up_threshold as f64
&& m.active_workers < config.max_workers
{
ScalingDecision::ScaleUp
} else if depth_per_worker < config.scale_down_threshold as f64
&& m.active_workers > config.min_workers
{
ScalingDecision::ScaleDown
} else {
return None;
};
let new_count = match decision {
ScalingDecision::ScaleUp => (m.active_workers + config.scale_step).min(config.max_workers),
ScalingDecision::ScaleDown => {
(m.active_workers.saturating_sub(config.scale_step)).max(config.min_workers)
}
};
info!(
"Autoscaling: {:?} from {} to {} workers (avg queue depth: {:.1})",
decision, m.active_workers, new_count, m.avg_queue_depth
);
m.active_workers = new_count;
m.last_scale_time = Some(now);
Some(new_count)
}
#[derive(Debug, Clone)]
struct PoolArchival {
policy: crate::archive::ArchivalPolicy,
config: crate::archive::ArchivalConfig,
interval: Duration,
}
fn spawn_periodic<F, Fut>(
interval: Duration,
mut signal_rx: watch::Receiver<bool>,
run: F,
) -> tokio::task::JoinHandle<()>
where
F: Fn() -> Fut + Send + 'static,
Fut: std::future::Future<Output = ()> + Send + 'static,
{
tokio::spawn(async move {
let mut ticker = tokio::time::interval(interval.max(Duration::from_millis(1)));
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
tokio::select! {
_ = signal_rx.changed() => break,
_ = ticker.tick() => {}
}
if *signal_rx.borrow() {
break;
}
if let Err(e) = tokio::spawn(run()).await
&& e.is_panic()
{
error!("Periodic maintenance task panicked: {}", e);
}
}
})
}
impl<DB: Database> Drop for WorkerPool<DB> {
fn drop(&mut self) {
if let Some(task) = &self.autoscale_task {
task.abort();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[test]
fn test_error_handling() {
let error = HammerworkError::Worker {
message: "Test error".to_string(),
};
assert_eq!(error.to_string(), "Worker error: Test error");
}
#[test]
fn test_job_timeout_detection_logic() {
use crate::job::{Job, JobStatus};
use serde_json::json;
use std::time::Duration;
let mut job = Job::new("timeout_test".to_string(), json!({"data": "test"}))
.with_timeout(Duration::from_millis(100));
assert!(!job.should_timeout());
job.started_at = Some(chrono::Utc::now() - chrono::Duration::milliseconds(50));
job.status = JobStatus::Running;
assert!(!job.should_timeout());
job.started_at = Some(chrono::Utc::now() - chrono::Duration::milliseconds(200));
assert!(job.should_timeout());
let mut job_no_timeout = Job::new("no_timeout".to_string(), json!({"data": "test"}));
job_no_timeout.started_at = Some(chrono::Utc::now() - chrono::Duration::hours(1));
assert!(!job_no_timeout.should_timeout());
}
#[tokio::test]
async fn test_timeout_statistics_integration() {
use crate::stats::{InMemoryStatsCollector, JobEvent, JobEventType};
use std::sync::Arc;
let stats_collector = Arc::new(InMemoryStatsCollector::new_default());
let timeout_event = JobEvent {
job_id: uuid::Uuid::new_v4(),
queue_name: "timeout_queue".to_string(),
event_type: JobEventType::TimedOut,
priority: crate::priority::JobPriority::Normal,
processing_time_ms: Some(5000), error_message: Some("Job timed out after 5s".to_string()),
timestamp: chrono::Utc::now(),
};
stats_collector.record_event(timeout_event).await.unwrap();
let stats = stats_collector
.get_queue_statistics("timeout_queue", Duration::from_secs(60))
.await
.unwrap();
assert_eq!(stats.total_processed, 1);
assert_eq!(stats.timed_out, 1);
assert_eq!(stats.error_rate, 1.0); }
#[test]
fn test_worker_rate_limit_configuration() {
use crate::rate_limit::RateLimit;
let rate_limit = RateLimit::per_second(10).with_burst_limit(20);
assert_eq!(rate_limit.rate, 10);
assert_eq!(rate_limit.burst_limit, 20);
assert_eq!(rate_limit.per, Duration::from_secs(1));
let per_minute = RateLimit::per_minute(60);
assert_eq!(per_minute.rate, 60);
assert_eq!(per_minute.per, Duration::from_secs(60));
let per_hour = RateLimit::per_hour(3600);
assert_eq!(per_hour.rate, 3600);
assert_eq!(per_hour.per, Duration::from_secs(3600));
}
#[test]
fn test_throttle_config_configuration() {
use crate::rate_limit::ThrottleConfig;
let throttle_config = ThrottleConfig::new()
.max_concurrent(5)
.rate_per_minute(100)
.backoff_on_error(Duration::from_secs(30))
.enabled(true);
assert_eq!(throttle_config.max_concurrent, Some(5));
assert_eq!(throttle_config.rate_per_minute, Some(100));
assert_eq!(
throttle_config.backoff_on_error,
Some(Duration::from_secs(30))
);
assert!(throttle_config.enabled);
let rate_limit = throttle_config.to_rate_limit().unwrap();
assert_eq!(rate_limit.rate, 100);
assert_eq!(rate_limit.per, Duration::from_secs(60));
}
#[tokio::test]
async fn test_rate_limiter_integration() {
use crate::rate_limit::{RateLimit, RateLimiter};
let rate_limit = RateLimit::per_second(5); let rate_limiter = RateLimiter::new(rate_limit);
assert!(rate_limiter.try_acquire());
assert!(rate_limiter.try_acquire());
assert!(rate_limiter.try_acquire());
assert!(rate_limiter.try_acquire());
assert!(rate_limiter.try_acquire());
assert!(!rate_limiter.try_acquire());
let start = std::time::Instant::now();
rate_limiter.acquire().await.unwrap();
let elapsed = start.elapsed();
assert!(elapsed < Duration::from_millis(500)); }
#[test]
fn test_worker_backoff_configuration() {
use crate::rate_limit::ThrottleConfig;
let throttle_config = ThrottleConfig::new().backoff_on_error(Duration::from_secs(60));
assert_eq!(
throttle_config.backoff_on_error,
Some(Duration::from_secs(60))
);
let poll_interval = Duration::from_secs(1);
let backoff_duration = throttle_config.backoff_on_error.unwrap_or(poll_interval);
assert_eq!(backoff_duration, Duration::from_secs(60));
let no_backoff_config = ThrottleConfig::new();
let backoff_duration = no_backoff_config.backoff_on_error.unwrap_or(poll_interval);
assert_eq!(backoff_duration, poll_interval);
}
#[tokio::test]
async fn test_rate_limiter_token_availability() {
use crate::rate_limit::{RateLimit, RateLimiter};
let rate_limit = RateLimit::per_second(10); let rate_limiter = RateLimiter::new(rate_limit);
let initial_tokens = rate_limiter.available_tokens();
assert_eq!(initial_tokens, 10.0);
assert!(rate_limiter.try_acquire());
assert!(rate_limiter.try_acquire());
let remaining_tokens = rate_limiter.available_tokens();
assert!((8.0..8.5).contains(&remaining_tokens), "{remaining_tokens}");
}
#[test]
fn test_rate_limit_edge_cases() {
use crate::rate_limit::RateLimit;
let low_rate = RateLimit::per_hour(1);
assert_eq!(low_rate.rate, 1);
assert_eq!(low_rate.per, Duration::from_secs(3600));
let high_rate = RateLimit::per_second(1000);
assert_eq!(high_rate.rate, 1000);
assert_eq!(high_rate.burst_limit, 1000);
let custom_burst = RateLimit::per_second(10).with_burst_limit(50);
assert_eq!(custom_burst.burst_limit, 50);
}
#[test]
fn test_throttle_config_defaults() {
use crate::rate_limit::ThrottleConfig;
let default_config = ThrottleConfig::default();
assert!(default_config.enabled);
assert!(default_config.max_concurrent.is_none());
assert!(default_config.rate_per_minute.is_none());
assert!(default_config.backoff_on_error.is_none());
let new_config = ThrottleConfig::new();
assert_eq!(new_config.enabled, default_config.enabled);
assert_eq!(new_config.max_concurrent, default_config.max_concurrent);
}
#[test]
fn test_autoscale_config_defaults() {
let config = AutoscaleConfig::default();
assert!(config.enabled);
assert_eq!(config.min_workers, 1);
assert_eq!(config.max_workers, 10);
assert_eq!(config.scale_up_threshold, 5);
assert_eq!(config.scale_down_threshold, 2);
assert_eq!(config.cooldown_period, Duration::from_secs(60));
assert_eq!(config.scale_step, 1);
assert_eq!(config.evaluation_window, Duration::from_secs(30));
}
#[test]
fn test_autoscale_config_builder() {
let config = AutoscaleConfig::new()
.with_min_workers(2)
.with_max_workers(20)
.with_scale_up_threshold(8)
.with_scale_down_threshold(1)
.with_cooldown_period(Duration::from_secs(120))
.with_scale_step(2)
.with_evaluation_window(Duration::from_secs(45));
assert_eq!(config.min_workers, 2);
assert_eq!(config.max_workers, 20);
assert_eq!(config.scale_up_threshold, 8);
assert_eq!(config.scale_down_threshold, 1);
assert_eq!(config.cooldown_period, Duration::from_secs(120));
assert_eq!(config.scale_step, 2);
assert_eq!(config.evaluation_window, Duration::from_secs(45));
}
#[test]
fn test_autoscale_config_presets() {
let conservative = AutoscaleConfig::conservative();
assert_eq!(conservative.min_workers, 2);
assert_eq!(conservative.max_workers, 5);
assert_eq!(conservative.scale_up_threshold, 10);
assert_eq!(conservative.cooldown_period, Duration::from_secs(300));
let aggressive = AutoscaleConfig::aggressive();
assert_eq!(aggressive.min_workers, 1);
assert_eq!(aggressive.max_workers, 20);
assert_eq!(aggressive.scale_up_threshold, 3);
assert_eq!(aggressive.cooldown_period, Duration::from_secs(30));
let disabled = AutoscaleConfig::disabled();
assert!(!disabled.enabled);
}
#[test]
fn test_autoscale_config_validation() {
let config = AutoscaleConfig::new().with_min_workers(0);
assert_eq!(config.min_workers, 1);
let config = AutoscaleConfig::new()
.with_min_workers(5)
.with_max_workers(3);
assert_eq!(config.max_workers, 5);
let config = AutoscaleConfig::new().with_scale_up_threshold(0);
assert_eq!(config.scale_up_threshold, 1);
let config = AutoscaleConfig::new().with_scale_step(0);
assert_eq!(config.scale_step, 1);
}
#[test]
fn test_autoscale_metrics_default() {
let metrics = AutoscaleMetrics::default();
assert_eq!(metrics.active_workers, 0);
assert_eq!(metrics.avg_queue_depth, 0.0);
assert_eq!(metrics.current_queue_depth, 0);
assert_eq!(metrics.jobs_per_second, 0.0);
assert_eq!(metrics.worker_utilization, 0.0);
assert_eq!(metrics.time_since_last_scale, Duration::from_secs(0));
assert!(metrics.last_scale_time.is_none());
}
#[test]
fn test_job_event_hooks_default() {
let hooks = JobEventHooks::default();
assert!(hooks.on_job_start.is_none());
assert!(hooks.on_job_complete.is_none());
assert!(hooks.on_job_fail.is_none());
assert!(hooks.on_job_timeout.is_none());
assert!(hooks.on_job_retry.is_none());
}
#[test]
fn test_job_event_hooks_new() {
let hooks = JobEventHooks::new();
assert!(hooks.on_job_start.is_none());
assert!(hooks.on_job_complete.is_none());
assert!(hooks.on_job_fail.is_none());
assert!(hooks.on_job_timeout.is_none());
assert!(hooks.on_job_retry.is_none());
}
#[test]
fn test_job_event_hooks_builder() {
use std::sync::{Arc, Mutex};
let events = Arc::new(Mutex::new(Vec::new()));
let events_start = Arc::clone(&events);
let events_complete = Arc::clone(&events);
let events_fail = Arc::clone(&events);
let events_timeout = Arc::clone(&events);
let events_retry = Arc::clone(&events);
let hooks = JobEventHooks::new()
.on_start(move |event: JobHookEvent| {
events_start
.lock()
.unwrap()
.push(format!("start:{}", event.job.id));
})
.on_complete(move |event: JobHookEvent| {
events_complete
.lock()
.unwrap()
.push(format!("complete:{}", event.job.id));
})
.on_fail(move |event: JobHookEvent| {
events_fail
.lock()
.unwrap()
.push(format!("fail:{}", event.job.id));
})
.on_timeout(move |event: JobHookEvent| {
events_timeout
.lock()
.unwrap()
.push(format!("timeout:{}", event.job.id));
})
.on_retry(move |event: JobHookEvent| {
events_retry
.lock()
.unwrap()
.push(format!("retry:{}", event.job.id));
});
assert!(hooks.on_job_start.is_some());
assert!(hooks.on_job_complete.is_some());
assert!(hooks.on_job_fail.is_some());
assert!(hooks.on_job_timeout.is_some());
assert!(hooks.on_job_retry.is_some());
}
#[test]
fn test_job_hook_event_creation() {
use crate::Job;
use serde_json::json;
use std::time::Duration;
let job = Job::new("test_queue".to_string(), json!({"test": "data"}))
.with_trace_id("trace-123")
.with_correlation_id("corr-456");
let event = JobHookEvent {
job: job.clone(),
timestamp: Utc::now(),
duration: Some(Duration::from_millis(500)),
error: Some("Test error".to_string()),
};
assert_eq!(event.job.id, job.id);
assert_eq!(event.job.queue_name, "test_queue");
assert_eq!(event.job.trace_id, Some("trace-123".to_string()));
assert_eq!(event.job.correlation_id, Some("corr-456".to_string()));
assert_eq!(event.duration, Some(Duration::from_millis(500)));
assert_eq!(event.error, Some("Test error".to_string()));
}
#[test]
fn test_job_event_hooks_fire_methods() {
use crate::Job;
use serde_json::json;
use std::sync::{Arc, Mutex};
use std::time::Duration;
let events = Arc::new(Mutex::new(Vec::new()));
let job = Job::new("test_queue".to_string(), json!({"test": "data"}));
{
let events_clone = Arc::clone(&events);
let hooks = JobEventHooks::new().on_start(move |event: JobHookEvent| {
events_clone
.lock()
.unwrap()
.push(format!("start:{}", event.job.queue_name));
});
hooks.fire_job_start(job.clone());
let captured_events = events.lock().unwrap();
assert_eq!(captured_events.len(), 1);
assert_eq!(captured_events[0], "start:test_queue");
}
events.lock().unwrap().clear();
{
let events_clone = Arc::clone(&events);
let hooks = JobEventHooks::new().on_complete(move |event: JobHookEvent| {
events_clone.lock().unwrap().push(format!(
"complete:{}:{}ms",
event.job.queue_name,
event.duration.unwrap_or_default().as_millis()
));
});
hooks.fire_job_complete(job.clone(), Duration::from_millis(150));
let captured_events = events.lock().unwrap();
assert_eq!(captured_events.len(), 1);
assert_eq!(captured_events[0], "complete:test_queue:150ms");
}
events.lock().unwrap().clear();
{
let events_clone = Arc::clone(&events);
let hooks = JobEventHooks::new().on_fail(move |event: JobHookEvent| {
events_clone.lock().unwrap().push(format!(
"fail:{}:{}",
event.job.queue_name,
event.error.unwrap_or_default()
));
});
hooks.fire_job_fail(job.clone(), "Connection timeout".to_string());
let captured_events = events.lock().unwrap();
assert_eq!(captured_events.len(), 1);
assert_eq!(captured_events[0], "fail:test_queue:Connection timeout");
}
events.lock().unwrap().clear();
{
let events_clone = Arc::clone(&events);
let hooks = JobEventHooks::new().on_timeout(move |event: JobHookEvent| {
events_clone.lock().unwrap().push(format!(
"timeout:{}:{}ms",
event.job.queue_name,
event.duration.unwrap_or_default().as_millis()
));
});
hooks.fire_job_timeout(job.clone(), Duration::from_secs(30));
let captured_events = events.lock().unwrap();
assert_eq!(captured_events.len(), 1);
assert_eq!(captured_events[0], "timeout:test_queue:30000ms");
}
events.lock().unwrap().clear();
{
let events_clone = Arc::clone(&events);
let hooks = JobEventHooks::new().on_retry(move |event: JobHookEvent| {
events_clone.lock().unwrap().push(format!(
"retry:{}:{}",
event.job.queue_name,
event.error.unwrap_or_default()
));
});
hooks.fire_job_retry(job.clone(), "API rate limit exceeded".to_string());
let captured_events = events.lock().unwrap();
assert_eq!(captured_events.len(), 1);
assert_eq!(
captured_events[0],
"retry:test_queue:API rate limit exceeded"
);
}
}
#[test]
fn test_job_event_hooks_clone() {
use std::sync::{Arc, Mutex};
let events = Arc::new(Mutex::new(Vec::new()));
let events_clone = Arc::clone(&events);
let hooks = JobEventHooks::new().on_start(move |event: JobHookEvent| {
events_clone
.lock()
.unwrap()
.push(format!("cloned:{}", event.job.id));
});
let hooks_clone = hooks.clone();
let job = crate::Job::new("test".to_string(), serde_json::json!({}));
hooks.fire_job_start(job.clone());
hooks_clone.fire_job_start(job);
let captured_events = events.lock().unwrap();
assert_eq!(captured_events.len(), 2);
assert!(captured_events[0].starts_with("cloned:"));
assert!(captured_events[1].starts_with("cloned:"));
}
#[test]
fn test_retry_at_from_delay_never_panics_and_clamps() {
let now = Utc::now();
assert_eq!(
retry_at_from_delay(now, Duration::from_secs(30)),
now + chrono::Duration::seconds(30)
);
let max = now + chrono::Duration::from_std(MAX_RETRY_DELAY).unwrap();
assert_eq!(retry_at_from_delay(now, Duration::MAX), max);
assert_eq!(retry_at_from_delay(now, MAX_RETRY_DELAY * 2), max);
assert_eq!(
retry_at_from_delay(DateTime::<Utc>::MAX_UTC, Duration::from_secs(1)),
DateTime::<Utc>::MAX_UTC
);
}
#[test]
fn test_error_backoff_grows_is_capped_and_never_zero() {
assert_eq!(error_backoff(Duration::ZERO, 1), MIN_ERROR_BACKOFF);
assert_eq!(
error_backoff(Duration::from_secs(1), 1),
Duration::from_secs(1)
);
assert_eq!(
error_backoff(Duration::from_secs(1), 3),
Duration::from_secs(4)
);
assert_eq!(
error_backoff(Duration::from_secs(1), u32::MAX),
MAX_ERROR_BACKOFF
);
assert_eq!(
error_backoff(Duration::from_secs(120), 5),
Duration::from_secs(120)
);
}
#[tokio::test]
async fn test_catch_unwind_turns_panics_into_errors() {
let ok: HandlerFuture = Box::pin(async { Ok(JobResult::success()) });
assert!(CatchUnwind { inner: ok }.await.unwrap().is_ok());
let panicking: HandlerFuture = Box::pin(async {
tokio::task::yield_now().await;
panic!("boom")
});
let payload = CatchUnwind { inner: panicking }.await.unwrap_err();
let error = handler_panic_error(payload.as_ref());
assert_eq!(
error.to_string(),
"Worker error: Job handler panicked: boom"
);
let formatted: HandlerFuture = Box::pin(async { panic!("code {}", 42) });
let payload = CatchUnwind { inner: formatted }.await.unwrap_err();
assert_eq!(panic_message(payload.as_ref()), "code 42");
}
#[tokio::test]
async fn spawn_periodic_runs_until_shutdown() {
let runs = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let (signal_tx, signal_rx) = watch::channel(false);
let counter = Arc::clone(&runs);
let task = spawn_periodic(Duration::from_millis(10), signal_rx, move || {
let counter = Arc::clone(&counter);
async move {
counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
});
tokio::time::sleep(Duration::from_millis(100)).await;
signal_tx.send_replace(true);
tokio::time::timeout(Duration::from_secs(5), task)
.await
.expect("the loop stops on shutdown")
.unwrap();
let seen = runs.load(std::sync::atomic::Ordering::SeqCst);
assert!(seen >= 2, "ran {seen} times");
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(runs.load(std::sync::atomic::Ordering::SeqCst), seen);
}
#[tokio::test]
async fn spawn_periodic_never_cancels_a_run() {
use std::sync::atomic::{AtomicBool, Ordering};
let started = Arc::new(tokio::sync::Notify::new());
let finished = Arc::new(AtomicBool::new(false));
let (signal_tx, signal_rx) = watch::channel(false);
let (run_started, run_finished) = (Arc::clone(&started), Arc::clone(&finished));
let task = spawn_periodic(Duration::from_secs(3600), signal_rx, move || {
let (started, finished) = (Arc::clone(&run_started), Arc::clone(&run_finished));
async move {
started.notify_one();
tokio::time::sleep(Duration::from_millis(200)).await;
finished.store(true, Ordering::SeqCst);
}
});
started.notified().await;
signal_tx.send_replace(true);
task.abort();
assert!(!finished.load(Ordering::SeqCst));
tokio::time::sleep(Duration::from_millis(400)).await;
assert!(
finished.load(Ordering::SeqCst),
"the in-flight run completed after the loop was stopped"
);
}
#[cfg(any(feature = "postgres", feature = "mysql"))]
#[test]
fn encrypted_job_purge_is_opt_in_and_configurable() {
#[cfg(feature = "postgres")]
type Db = sqlx::Postgres;
#[cfg(all(feature = "mysql", not(feature = "postgres")))]
type Db = sqlx::MySql;
let pool = WorkerPool::<Db>::new();
assert_eq!(pool.encrypted_job_purge_interval, None);
let pool = pool.with_encrypted_job_purge(Duration::from_secs(30));
assert_eq!(
pool.encrypted_job_purge_interval,
Some(Duration::from_secs(30))
);
assert_eq!(
pool.without_encrypted_job_purge()
.encrypted_job_purge_interval,
None
);
}
#[cfg(any(feature = "postgres", feature = "mysql"))]
#[test]
fn archival_is_opt_in() {
#[cfg(feature = "postgres")]
type Db = sqlx::Postgres;
#[cfg(all(feature = "mysql", not(feature = "postgres")))]
type Db = sqlx::MySql;
let pool = WorkerPool::<Db>::new();
assert!(pool.archival.is_none());
let policy = crate::archive::ArchivalPolicy::new().with_batch_size(7);
let pool = pool.with_archival(
policy,
crate::archive::ArchivalConfig::new().with_compression_level(2),
Duration::from_secs(90),
);
let archival = pool.archival.as_ref().unwrap();
assert_eq!(archival.interval, Duration::from_secs(90));
assert_eq!(archival.policy.batch_size, 7);
assert_eq!(archival.config.compression_level, 2);
assert!(pool.without_archival().archival.is_none());
}
#[test]
fn panic_message_describes_every_payload_kind() {
assert_eq!(panic_message(&"static str"), "static str");
assert_eq!(panic_message(&"owned".to_string()), "owned");
assert_eq!(panic_message(&42_u32), "non-string panic payload");
assert_eq!(
handler_panic_error(&"boom").to_string(),
"Worker error: Job handler panicked: boom"
);
}
#[test]
fn shutdown_requested_on_signal_or_dropped_sender() {
let (tx, mut rx) = mpsc::channel(1);
assert!(!shutdown_requested(&mut rx));
tx.try_send(()).unwrap();
assert!(shutdown_requested(&mut rx));
drop(tx);
assert!(
shutdown_requested(&mut rx),
"a dropped sender means shutdown"
);
}
fn autoscale_state(
active_workers: usize,
) -> (Arc<std::sync::RwLock<AutoscaleMetrics>>, QueueDepthHistory) {
(
Arc::new(std::sync::RwLock::new(AutoscaleMetrics {
active_workers,
..Default::default()
})),
Arc::new(std::sync::RwLock::new(Vec::new())),
)
}
fn fast_autoscale() -> AutoscaleConfig {
AutoscaleConfig::new()
.with_min_workers(1)
.with_max_workers(3)
.with_scale_up_threshold(4)
.with_scale_down_threshold(1)
.with_cooldown_period(Duration::ZERO)
.with_evaluation_window(Duration::from_secs(30))
}
#[test]
fn autoscaling_scales_up_within_max_workers() {
let config = fast_autoscale().with_scale_step(5);
let (metrics, history) = autoscale_state(1);
let decision = apply_queue_depth(10, Utc::now(), &config, &metrics, &history);
assert_eq!(decision, Some(3));
let m = metrics.read().unwrap().clone();
assert_eq!(m.active_workers, 3);
assert_eq!(m.current_queue_depth, 10);
assert_eq!(m.avg_queue_depth, 10.0);
assert!(m.last_scale_time.is_some());
let decision = apply_queue_depth(1000, Utc::now(), &config, &metrics, &history);
assert_eq!(decision, None);
assert_eq!(metrics.read().unwrap().active_workers, 3);
}
#[test]
fn autoscaling_scales_down_to_min_workers() {
let config = fast_autoscale().with_scale_step(2);
let (metrics, history) = autoscale_state(3);
let decision = apply_queue_depth(0, Utc::now(), &config, &metrics, &history);
assert_eq!(decision, Some(1));
let decision = apply_queue_depth(0, Utc::now(), &config, &metrics, &history);
assert_eq!(decision, None, "never below min_workers");
assert_eq!(metrics.read().unwrap().active_workers, 1);
}
#[test]
fn autoscaling_holds_between_thresholds_and_during_cooldown() {
let (metrics, history) = autoscale_state(2);
let decision = apply_queue_depth(6, Utc::now(), &fast_autoscale(), &metrics, &history);
assert_eq!(decision, None);
let config = fast_autoscale().with_cooldown_period(Duration::from_secs(3600));
let (metrics, history) = autoscale_state(1);
let now = Utc::now();
assert_eq!(
apply_queue_depth(100, now, &config, &metrics, &history),
Some(2),
"the first decision is not in a cooldown"
);
assert_eq!(
apply_queue_depth(
100,
now + chrono::Duration::seconds(1),
&config,
&metrics,
&history
),
None,
"a second decision within the cooldown period is suppressed"
);
let m = metrics.read().unwrap().clone();
assert_eq!(m.active_workers, 2);
assert_eq!(m.time_since_last_scale, Duration::from_secs(1));
}
#[test]
fn autoscaling_averages_depth_over_the_evaluation_window() {
let config = fast_autoscale()
.with_scale_up_threshold(100)
.with_scale_down_threshold(0)
.with_evaluation_window(Duration::from_secs(10));
let (metrics, history) = autoscale_state(1);
let start = Utc::now();
for (offset, depth) in [(0, 30), (4, 10), (8, 20)] {
apply_queue_depth(
depth,
start + chrono::Duration::seconds(offset),
&config,
&metrics,
&history,
);
}
assert_eq!(metrics.read().unwrap().avg_queue_depth, 20.0);
apply_queue_depth(
40,
start + chrono::Duration::seconds(12),
&config,
&metrics,
&history,
);
assert_eq!(history.read().unwrap().len(), 3);
let m = metrics.read().unwrap().clone();
assert_eq!(m.avg_queue_depth, (10.0 + 20.0 + 40.0) / 3.0);
assert_eq!(m.current_queue_depth, 40);
}
#[cfg(feature = "postgres")]
fn lazy_worker(queue_name: &str) -> Worker<sqlx::Postgres> {
let pool = sqlx::PgPool::connect_lazy("postgres://localhost/hammerwork_unused").unwrap();
let handler: JobHandler = Arc::new(|_job| Box::pin(async { Ok(()) }));
Worker::new(
Arc::new(JobQueue::new(pool)),
queue_name.to_string(),
handler,
)
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn with_config_applies_the_worker_section() {
let config = crate::config::WorkerConfig {
polling_interval: Duration::from_millis(250),
job_timeout: Duration::from_secs(7),
priority_weights: PriorityWeights::strict(),
retry_strategy: RetryStrategy::fixed(Duration::from_secs(42)),
..Default::default()
};
let worker = lazy_worker("cfg").with_config(&config);
assert_eq!(worker.poll_interval, Duration::from_millis(250));
assert_eq!(worker.default_timeout, Some(Duration::from_secs(7)));
assert!(worker.priority_weights.as_ref().unwrap().is_strict());
assert_eq!(
worker.default_retry_strategy,
Some(RetryStrategy::fixed(Duration::from_secs(42)))
);
let worker = lazy_worker("cfg").with_weighted_priority();
assert!(!worker.priority_weights.as_ref().unwrap().is_strict());
let worker = worker.with_strict_priority();
assert!(worker.priority_weights.as_ref().unwrap().is_strict());
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn with_hammerwork_config_applies_the_queue_throttle() {
let mut config = crate::HammerworkConfig::new();
config.worker.polling_interval = Duration::from_millis(100);
config.rate_limiting.enabled = true;
config.rate_limiting.default_throttle = ThrottleConfig::new().rate_per_minute(600);
config.rate_limiting.queue_throttles.insert(
"emails".to_string(),
ThrottleConfig::new()
.max_concurrent(2)
.rate_per_minute(60)
.backoff_on_error(Duration::from_secs(5)),
);
let emails = lazy_worker("emails").with_hammerwork_config(&config);
assert_eq!(emails.poll_interval, Duration::from_millis(100));
assert_eq!(
emails.throttle_config.as_ref().unwrap().max_concurrent,
Some(2)
);
assert_eq!(
emails
.concurrency_limit
.as_ref()
.unwrap()
.available_permits(),
2
);
assert!(emails.rate_limiter.is_some());
assert_eq!(emails.error_backoff_base(), Duration::from_secs(5));
let other = lazy_worker("other").with_hammerwork_config(&config);
assert_eq!(
other.throttle_config.as_ref().unwrap().rate_per_minute,
Some(600)
);
assert!(other.concurrency_limit.is_none());
assert_eq!(
other.error_backoff_base(),
Duration::from_millis(100),
"without backoff_on_error the poll interval is the backoff base"
);
config.rate_limiting.enabled = false;
let unthrottled = lazy_worker("emails").with_hammerwork_config(&config);
assert!(unthrottled.throttle_config.is_none());
assert!(unthrottled.rate_limiter.is_none());
}
#[cfg(all(feature = "postgres", feature = "alerting"))]
#[tokio::test]
async fn with_hammerwork_config_applies_alerting() {
let mut config = crate::HammerworkConfig::new();
config.alerting = AlertingConfig::new().alert_on_queue_depth(17);
let worker = lazy_worker("alerts").with_hammerwork_config(&config);
let manager = worker.alert_manager.as_ref().expect("an alert manager");
assert_eq!(manager.config().queue_depth_threshold, Some(17));
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn disabled_throttle_is_ignored_and_max_concurrent_is_at_least_one() {
let mut disabled = ThrottleConfig::new()
.max_concurrent(1)
.rate_per_minute(1)
.backoff_on_error(Duration::from_secs(9));
disabled.enabled = false;
let worker = lazy_worker("t").with_throttle_config(disabled);
assert!(worker.throttle_config.is_none());
assert!(worker.rate_limiter.is_none());
assert!(worker.concurrency_limit.is_none());
assert_eq!(worker.error_backoff_base(), worker.poll_interval);
let worker = lazy_worker("t").with_throttle_config(ThrottleConfig::new().max_concurrent(0));
assert_eq!(
worker
.concurrency_limit
.as_ref()
.unwrap()
.available_permits(),
1,
"max_concurrent = 0 would block the worker forever"
);
let clone = worker.clone();
assert!(Arc::ptr_eq(
worker.concurrency_limit.as_ref().unwrap(),
clone.concurrency_limit.as_ref().unwrap()
));
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn pool_from_config_uses_pool_size_and_autoscaling() {
let config = crate::config::WorkerConfig {
pool_size: 3,
autoscaling_enabled: true,
min_workers: 2,
max_workers: 6,
..Default::default()
};
let pool = WorkerPool::from_config(lazy_worker("p").with_config(&config), &config);
assert_eq!(pool.workers.len(), 3);
assert!(pool.worker_template.is_some());
assert!(pool.autoscale_config.enabled);
assert_eq!(pool.autoscale_config.min_workers, 2);
assert_eq!(pool.autoscale_config.max_workers, 6);
let config = crate::config::WorkerConfig {
pool_size: 0,
..Default::default()
};
let pool = WorkerPool::from_config(lazy_worker("p"), &config);
assert_eq!(pool.workers.len(), 1, "a pool always has a worker");
assert!(!pool.autoscale_config.enabled);
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn pool_from_hammerwork_config_schedules_purge_and_validates_encryption() {
let mut config = crate::HammerworkConfig::new();
config.worker.pool_size = 2;
config.encryption.purge_interval_secs = Some(90);
let pool = WorkerPool::from_hammerwork_config(lazy_worker("p"), &config).unwrap();
assert_eq!(pool.workers.len(), 2);
assert_eq!(
pool.encrypted_job_purge_interval,
Some(Duration::from_secs(90))
);
config.encryption.purge_interval_secs = Some(0);
assert!(
WorkerPool::from_hammerwork_config(lazy_worker("p"), &config).is_err(),
"an invalid encryption section is rejected"
);
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn pool_stats_collector_is_applied_to_workers() {
use crate::stats::InMemoryStatsCollector;
let stats: Arc<dyn StatisticsCollector> = Arc::new(InMemoryStatsCollector::new_default());
let mut pool = WorkerPool::new().with_stats_collector(Arc::clone(&stats));
pool.add_worker(lazy_worker("s"));
assert!(Arc::ptr_eq(
pool.workers[0].stats_collector.as_ref().unwrap(),
&stats
));
assert!(Arc::ptr_eq(
pool.stats_collector().as_ref().unwrap(),
&stats
));
let mut pool = WorkerPool::new().with_autoscaling(AutoscaleConfig::new());
assert!(pool.worker_template.is_none());
pool.add_worker(lazy_worker("s"));
assert!(pool.worker_template.is_some());
let mut pool = WorkerPool::new().without_autoscaling();
pool.add_worker(lazy_worker("s"));
assert!(pool.worker_template.is_none());
let mut pool = WorkerPool::new();
assert!(!pool.autoscale_config.enabled);
pool.add_worker(lazy_worker("s"));
assert!(pool.worker_template.is_none());
}
#[cfg(all(feature = "postgres", feature = "webhooks"))]
#[tokio::test]
async fn lifecycle_events_carry_the_job_event_details() {
let worker = lazy_worker("events");
let job_id = uuid::Uuid::new_v4();
let event = |event_type, error: Option<&str>| JobEvent {
job_id,
queue_name: "events".to_string(),
event_type,
priority: crate::priority::JobPriority::High,
processing_time_ms: Some(12),
error_message: error.map(str::to_string),
timestamp: Utc::now(),
};
let completed = worker.convert_to_lifecycle_event(&event(JobEventType::Completed, None));
assert_eq!(completed.event_type, JobLifecycleEventType::Completed);
assert_eq!(completed.job_id, job_id);
assert_eq!(completed.priority, crate::priority::JobPriority::High);
assert!(completed.error.is_none());
assert_eq!(completed.metadata["worker_queue"], "events");
assert_eq!(completed.metadata["processing_time_ms"], "12");
for (event_type, lifecycle, error_type) in [
(
JobEventType::TimedOut,
JobLifecycleEventType::TimedOut,
Some("timeout"),
),
(
JobEventType::Failed,
JobLifecycleEventType::Failed,
Some("processing_error"),
),
(
JobEventType::Dead,
JobLifecycleEventType::Dead,
Some("max_retries_exceeded"),
),
(JobEventType::Retried, JobLifecycleEventType::Retried, None),
(JobEventType::Started, JobLifecycleEventType::Started, None),
] {
let converted = worker.convert_to_lifecycle_event(&event(event_type, Some("bad")));
assert_eq!(converted.event_type, lifecycle);
let error = converted.error.expect("the error message is kept");
assert_eq!(error.message, "bad");
assert_eq!(error.error_type.as_deref(), error_type);
}
}
#[cfg(all(feature = "postgres", feature = "metrics"))]
#[tokio::test]
async fn monitoring_task_stops_when_run_is_dropped() {
let metrics = Arc::new(
PrometheusMetricsCollector::new(crate::metrics::MetricsConfig::default()).unwrap(),
);
let worker = lazy_worker("monitor")
.with_metrics_collector(Arc::clone(&metrics))
.with_poll_interval(Duration::from_millis(10));
assert_eq!(Arc::strong_count(&metrics), 2);
let (_shutdown_tx, shutdown_rx) = mpsc::channel(1);
let run = tokio::spawn(async move { worker.run(shutdown_rx).await });
let deadline = std::time::Instant::now() + Duration::from_secs(5);
while Arc::strong_count(&metrics) < 3 && std::time::Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(5)).await;
}
assert_eq!(Arc::strong_count(&metrics), 3);
run.abort();
let _ = run.await;
let deadline = std::time::Instant::now() + Duration::from_secs(5);
while Arc::strong_count(&metrics) > 1 && std::time::Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(5)).await;
}
assert_eq!(
Arc::strong_count(&metrics),
1,
"the monitoring task outlived the worker"
);
}
#[cfg(feature = "postgres")]
#[tokio::test]
async fn zero_poll_interval_does_not_hot_loop() {
let worker = lazy_worker("idle").with_poll_interval(Duration::ZERO);
let (_shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
let started = std::time::Instant::now();
for _ in 0..5 {
assert!(matches!(
worker.idle_wait(&mut shutdown_rx).await,
Acquired::Idle
));
}
assert!(started.elapsed() >= MIN_POLL_INTERVAL * 5);
}
}