#![allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
use crate::event::{AlertSeverity, Event, TaskEvent, WorkerEvent};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::collections::VecDeque;
use std::time::Duration;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct EventSnapshot {
pub timestamp: DateTime<Utc>,
pub pending: u64,
pub active: u64,
pub succeeded: u64,
pub failed: u64,
pub retried: u64,
pub revoked: u64,
pub queue_depth: u64,
pub worker_count: u64,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub queue_depths: HashMap<String, u64>,
}
impl EventSnapshot {
#[inline]
#[must_use]
pub const fn completed(&self) -> u64 {
self.succeeded + self.failed
}
#[inline]
#[must_use]
pub fn failure_rate(&self) -> f64 {
let total = self.completed();
if total == 0 {
0.0
} else {
self.failed as f64 / total as f64
}
}
}
#[derive(Debug, Clone, Default)]
pub struct EventSnapshotBuilder {
timestamp: Option<DateTime<Utc>>,
pending: u64,
active: u64,
succeeded: u64,
failed: u64,
retried: u64,
revoked: u64,
queue_depth: u64,
queue_depths: HashMap<String, u64>,
workers: std::collections::HashSet<String>,
worker_count_override: Option<u64>,
}
impl EventSnapshotBuilder {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn timestamp(mut self, ts: DateTime<Utc>) -> Self {
self.timestamp = Some(ts);
self
}
#[must_use]
pub fn pending(mut self, n: u64) -> Self {
self.pending = n;
self
}
#[must_use]
pub fn active(mut self, n: u64) -> Self {
self.active = n;
self
}
#[must_use]
pub fn succeeded(mut self, n: u64) -> Self {
self.succeeded = n;
self
}
#[must_use]
pub fn failed(mut self, n: u64) -> Self {
self.failed = n;
self
}
#[must_use]
pub fn retried(mut self, n: u64) -> Self {
self.retried = n;
self
}
#[must_use]
pub fn revoked(mut self, n: u64) -> Self {
self.revoked = n;
self
}
#[must_use]
pub fn worker_count(mut self, n: u64) -> Self {
self.worker_count_override = Some(n);
self
}
#[must_use]
pub fn queue_depth(mut self, queue: impl Into<String>, depth: u64) -> Self {
let q = queue.into();
let prev = self.queue_depths.insert(q, depth).unwrap_or(0);
self.queue_depth = self.queue_depth + depth - prev;
self
}
#[must_use]
pub fn total_queue_depth(mut self, depth: u64) -> Self {
self.queue_depth = depth;
self
}
pub fn observe(&mut self, event: &Event) -> &mut Self {
match event {
Event::Task(task) => match task {
TaskEvent::Sent { .. } => self.pending += 1,
TaskEvent::Started { .. } => self.active += 1,
TaskEvent::Succeeded { .. } => self.succeeded += 1,
TaskEvent::Failed { .. } => self.failed += 1,
TaskEvent::Retried { .. } => self.retried += 1,
TaskEvent::Revoked { .. } => self.revoked += 1,
TaskEvent::Received { .. } | TaskEvent::Rejected { .. } => {}
},
Event::Worker(worker) => match worker {
WorkerEvent::Online { hostname, .. } | WorkerEvent::Heartbeat { hostname, .. } => {
self.workers.insert(hostname.clone());
}
WorkerEvent::Offline { hostname, .. } => {
self.workers.remove(hostname);
}
},
}
self
}
pub fn observe_all<'a, I>(&mut self, events: I) -> &mut Self
where
I: IntoIterator<Item = &'a Event>,
{
for event in events {
self.observe(event);
}
self
}
#[must_use]
pub fn build(self) -> EventSnapshot {
let worker_count = self
.worker_count_override
.unwrap_or(self.workers.len() as u64);
EventSnapshot {
timestamp: self.timestamp.unwrap_or_else(Utc::now),
pending: self.pending,
active: self.active,
succeeded: self.succeeded,
failed: self.failed,
retried: self.retried,
revoked: self.revoked,
queue_depth: self.queue_depth,
worker_count,
queue_depths: self.queue_depths,
}
}
}
#[derive(Debug, Clone)]
pub enum AlertRuleKind {
FailureRate {
trip_threshold: f64,
clear_threshold: f64,
window: Duration,
min_samples: u64,
},
QueueDepth {
trip_depth: u64,
clear_depth: u64,
queue: Option<String>,
},
NoHeartbeat {
timeout: Duration,
},
}
#[derive(Debug, Clone)]
pub struct AlertRule {
pub name: String,
pub severity: AlertSeverity,
pub kind: AlertRuleKind,
pub cooldown: Duration,
}
impl AlertRule {
#[must_use]
pub fn failure_rate(
name: impl Into<String>,
severity: AlertSeverity,
trip_threshold: f64,
) -> Self {
Self {
name: name.into(),
severity,
kind: AlertRuleKind::FailureRate {
trip_threshold,
clear_threshold: trip_threshold,
window: Duration::from_secs(60),
min_samples: 1,
},
cooldown: Duration::from_secs(60),
}
}
#[must_use]
pub fn queue_depth(name: impl Into<String>, severity: AlertSeverity, trip_depth: u64) -> Self {
Self {
name: name.into(),
severity,
kind: AlertRuleKind::QueueDepth {
trip_depth,
clear_depth: trip_depth,
queue: None,
},
cooldown: Duration::from_secs(60),
}
}
#[must_use]
pub fn no_heartbeat(
name: impl Into<String>,
severity: AlertSeverity,
timeout: Duration,
) -> Self {
Self {
name: name.into(),
severity,
kind: AlertRuleKind::NoHeartbeat { timeout },
cooldown: Duration::from_secs(60),
}
}
#[must_use]
pub fn with_cooldown(mut self, cooldown: Duration) -> Self {
self.cooldown = cooldown;
self
}
#[must_use]
pub fn with_clear_threshold(mut self, clear: f64) -> Self {
if let AlertRuleKind::FailureRate {
clear_threshold, ..
} = &mut self.kind
{
*clear_threshold = clear;
}
self
}
#[must_use]
pub fn with_window(mut self, window: Duration) -> Self {
if let AlertRuleKind::FailureRate { window: w, .. } = &mut self.kind {
*w = window;
}
self
}
#[must_use]
pub fn with_min_samples(mut self, min_samples: u64) -> Self {
if let AlertRuleKind::FailureRate { min_samples: m, .. } = &mut self.kind {
*m = min_samples;
}
self
}
#[must_use]
pub fn with_clear_depth(mut self, clear: u64) -> Self {
if let AlertRuleKind::QueueDepth { clear_depth, .. } = &mut self.kind {
*clear_depth = clear;
}
self
}
#[must_use]
pub fn for_queue(mut self, queue: impl Into<String>) -> Self {
if let AlertRuleKind::QueueDepth { queue: q, .. } = &mut self.kind {
*q = Some(queue.into());
}
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AlertStatus {
Firing,
Resolved,
}
#[derive(Debug, Clone)]
pub struct RuleAlert {
pub rule_name: String,
pub severity: AlertSeverity,
pub status: AlertStatus,
pub message: String,
pub observed_value: f64,
pub timestamp: DateTime<Utc>,
}
impl RuleAlert {
#[inline]
#[must_use]
pub fn is_firing(&self) -> bool {
self.status == AlertStatus::Firing
}
}
#[derive(Debug)]
struct RuleState {
firing: bool,
last_fired: Option<DateTime<Utc>>,
completions: VecDeque<(DateTime<Utc>, bool)>,
}
impl RuleState {
fn new() -> Self {
Self {
firing: false,
last_fired: None,
completions: VecDeque::new(),
}
}
}
#[derive(Debug)]
pub struct AlertEvaluator {
rules: Vec<AlertRule>,
states: Vec<RuleState>,
last_heartbeat: Option<DateTime<Utc>>,
}
impl Default for AlertEvaluator {
fn default() -> Self {
Self::new()
}
}
impl AlertEvaluator {
#[must_use]
pub fn new() -> Self {
Self {
rules: Vec::new(),
states: Vec::new(),
last_heartbeat: None,
}
}
#[must_use]
pub fn with_rule(mut self, rule: AlertRule) -> Self {
self.add_rule(rule);
self
}
pub fn add_rule(&mut self, rule: AlertRule) {
self.rules.push(rule);
self.states.push(RuleState::new());
}
#[inline]
#[must_use]
pub fn rule_count(&self) -> usize {
self.rules.len()
}
#[must_use]
pub fn is_firing(&self, name: &str) -> bool {
self.rules
.iter()
.position(|r| r.name == name)
.and_then(|i| self.states.get(i))
.is_some_and(|s| s.firing)
}
pub fn observe_event(&mut self, event: &Event) -> Vec<RuleAlert> {
let now = event.timestamp();
if matches!(
event,
Event::Worker(WorkerEvent::Heartbeat { .. } | WorkerEvent::Online { .. })
) {
self.last_heartbeat = Some(now);
}
let completion = match event {
Event::Task(TaskEvent::Succeeded { .. }) => Some(false),
Event::Task(TaskEvent::Failed { .. }) => Some(true),
_ => None,
};
let mut alerts = Vec::new();
for idx in 0..self.rules.len() {
let kind = self.rules[idx].kind.clone();
match kind {
AlertRuleKind::FailureRate {
trip_threshold,
clear_threshold,
window,
min_samples,
} => {
let rule = self.rules[idx].clone();
let state = &mut self.states[idx];
if let Some(is_failure) = completion {
state.completions.push_back((now, is_failure));
}
Self::prune_completions(&mut state.completions, now, window);
let total = state.completions.len() as u64;
let failed = state.completions.iter().filter(|&&(_, f)| f).count() as u64;
let rate = if total == 0 {
0.0
} else {
failed as f64 / total as f64
};
let can_fire = total >= min_samples;
if let Some(alert) = Self::evaluate_threshold(
&rule,
state,
now,
rate,
trip_threshold,
clear_threshold,
can_fire,
true,
format!(
"failure rate {:.0}% ({failed}/{total}) over window",
rate * 100.0
),
) {
alerts.push(alert);
}
}
AlertRuleKind::NoHeartbeat { timeout } => {
let rule = self.rules[idx].clone();
if let Some(alert) = self.no_heartbeat_alert(idx, now, timeout, rule) {
alerts.push(alert);
}
}
AlertRuleKind::QueueDepth { .. } => {
}
}
}
alerts
}
pub fn observe_snapshot(&mut self, snapshot: &EventSnapshot) -> Vec<RuleAlert> {
let now = snapshot.timestamp;
if snapshot.worker_count > 0 {
self.last_heartbeat = Some(now);
}
let mut alerts = Vec::new();
for idx in 0..self.rules.len() {
let kind = self.rules[idx].kind.clone();
match kind {
AlertRuleKind::QueueDepth {
trip_depth,
clear_depth,
queue,
} => {
let depth = queue.as_ref().map_or(snapshot.queue_depth, |q| {
snapshot.queue_depths.get(q).copied().unwrap_or(0)
});
let rule = self.rules[idx].clone();
let state = &mut self.states[idx];
if let Some(alert) = Self::evaluate_threshold(
&rule,
state,
now,
depth as f64,
trip_depth as f64,
clear_depth as f64,
true,
false,
format!("queue depth {depth}"),
) {
alerts.push(alert);
}
}
AlertRuleKind::NoHeartbeat { timeout } => {
let rule = self.rules[idx].clone();
if let Some(alert) = self.no_heartbeat_alert(idx, now, timeout, rule) {
alerts.push(alert);
}
}
AlertRuleKind::FailureRate { .. } => {
}
}
}
alerts
}
pub fn tick(&mut self, now: DateTime<Utc>) -> Vec<RuleAlert> {
let mut alerts = Vec::new();
for idx in 0..self.rules.len() {
if let AlertRuleKind::NoHeartbeat { timeout } = self.rules[idx].kind {
let rule = self.rules[idx].clone();
if let Some(alert) = self.no_heartbeat_alert(idx, now, timeout, rule) {
alerts.push(alert);
}
}
}
alerts
}
fn no_heartbeat_alert(
&mut self,
idx: usize,
now: DateTime<Utc>,
timeout: Duration,
rule: AlertRule,
) -> Option<RuleAlert> {
let secs_silent = match self.last_heartbeat {
Some(last) => (now - last).num_milliseconds().max(0) as f64 / 1000.0,
None => f64::INFINITY,
};
let timeout_secs = timeout.as_secs_f64();
let tripped = secs_silent > timeout_secs;
let state = &mut self.states[idx];
let display = if secs_silent.is_finite() {
format!("no heartbeat for {secs_silent:.1}s")
} else {
"no heartbeat observed yet".to_string()
};
Self::transition(&rule, state, now, tripped, secs_silent, display)
}
#[allow(clippy::too_many_arguments)]
fn evaluate_threshold(
rule: &AlertRule,
state: &mut RuleState,
now: DateTime<Utc>,
value: f64,
trip: f64,
clear: f64,
can_fire: bool,
_trip_when_ge: bool,
display: String,
) -> Option<RuleAlert> {
let tripped = if state.firing {
value >= clear
} else {
can_fire && value >= trip
};
Self::transition(rule, state, now, tripped, value, display)
}
fn transition(
rule: &AlertRule,
state: &mut RuleState,
now: DateTime<Utc>,
tripped: bool,
observed_value: f64,
display: String,
) -> Option<RuleAlert> {
if tripped {
let was_firing = state.firing;
state.firing = true;
let cooldown_ok = match state.last_fired {
None => true,
Some(last) => {
let elapsed = (now - last).num_milliseconds().max(0) as u128;
elapsed >= rule.cooldown.as_millis()
}
};
if !was_firing || cooldown_ok {
state.last_fired = Some(now);
return Some(RuleAlert {
rule_name: rule.name.clone(),
severity: rule.severity,
status: AlertStatus::Firing,
message: format!("{}: {display}", rule.name),
observed_value,
timestamp: now,
});
}
None
} else {
if state.firing {
state.firing = false;
state.last_fired = None;
return Some(RuleAlert {
rule_name: rule.name.clone(),
severity: rule.severity,
status: AlertStatus::Resolved,
message: format!("{} resolved: {display}", rule.name),
observed_value,
timestamp: now,
});
}
None
}
}
fn prune_completions(
completions: &mut VecDeque<(DateTime<Utc>, bool)>,
now: DateTime<Utc>,
window: Duration,
) {
let cutoff = now - chrono::Duration::milliseconds(window.as_millis() as i64);
while let Some(&(ts, _)) = completions.front() {
if ts < cutoff {
completions.pop_front();
} else {
break;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::event::{TaskEventBuilder, WorkerEventBuilder};
use uuid::Uuid;
fn failed_event(at: DateTime<Utc>) -> Event {
Event::Task(TaskEvent::Failed {
task_id: Uuid::new_v4(),
task_name: "t".to_string(),
hostname: "w1".to_string(),
timestamp: at,
exception: "Boom".to_string(),
traceback: None,
})
}
fn succeeded_event(at: DateTime<Utc>) -> Event {
Event::Task(TaskEvent::Succeeded {
task_id: Uuid::new_v4(),
task_name: "t".to_string(),
hostname: "w1".to_string(),
timestamp: at,
runtime: 0.1,
result: None,
})
}
fn heartbeat_event(at: DateTime<Utc>) -> Event {
Event::Worker(WorkerEvent::Heartbeat {
hostname: "w1".to_string(),
timestamp: at,
active: 1,
processed: 1,
loadavg: None,
freq: 1.0,
})
}
#[test]
fn test_snapshot_builder_from_events() {
let now = Utc::now();
let events = vec![
TaskEventBuilder::new(Uuid::new_v4(), "t").sent("celery"),
succeeded_event(now),
failed_event(now),
WorkerEventBuilder::new("w1").online(),
WorkerEventBuilder::new("w2").online(),
];
let mut b = EventSnapshotBuilder::new();
b.observe_all(&events);
let snap = b.queue_depth("celery", 7).build();
assert_eq!(snap.pending, 1);
assert_eq!(snap.succeeded, 1);
assert_eq!(snap.failed, 1);
assert_eq!(snap.worker_count, 2);
assert_eq!(snap.queue_depth, 7);
assert!((snap.failure_rate() - 0.5).abs() < 1e-9);
}
#[test]
fn test_snapshot_serialization_roundtrip() {
let snap = EventSnapshotBuilder::new()
.pending(3)
.active(2)
.succeeded(10)
.failed(2)
.worker_count(4)
.queue_depth("q", 5)
.build();
let json = serde_json::to_string(&snap).unwrap();
let parsed: EventSnapshot = serde_json::from_str(&json).unwrap();
assert_eq!(snap, parsed);
}
#[test]
fn test_failure_rate_trips_and_recovers() {
let rule = AlertRule::failure_rate("fail", AlertSeverity::Error, 0.5)
.with_clear_threshold(0.2)
.with_min_samples(4)
.with_window(Duration::from_secs(3600))
.with_cooldown(Duration::from_secs(0));
let mut eval = AlertEvaluator::new().with_rule(rule);
let base = Utc::now();
let mut fired = false;
for i in 0..4 {
let at = base + chrono::Duration::seconds(i);
let alerts = eval.observe_event(&failed_event(at));
if alerts.iter().any(|a| a.is_firing()) {
fired = true;
}
}
assert!(fired, "rule should fire on high failure rate");
assert!(eval.is_firing("fail"));
let mut resolved = false;
for i in 4..40 {
let at = base + chrono::Duration::seconds(i);
let alerts = eval.observe_event(&succeeded_event(at));
if alerts.iter().any(|a| a.status == AlertStatus::Resolved) {
resolved = true;
}
}
assert!(resolved, "rule should resolve once failure rate drops");
assert!(!eval.is_firing("fail"));
}
#[test]
fn test_failure_rate_respects_min_samples() {
let rule = AlertRule::failure_rate("fail", AlertSeverity::Warning, 0.5)
.with_min_samples(10)
.with_cooldown(Duration::from_secs(0));
let mut eval = AlertEvaluator::new().with_rule(rule);
let base = Utc::now();
for i in 0..3 {
let at = base + chrono::Duration::seconds(i);
let alerts = eval.observe_event(&failed_event(at));
assert!(alerts.is_empty());
}
assert!(!eval.is_firing("fail"));
}
#[test]
fn test_queue_depth_hysteresis() {
let rule = AlertRule::queue_depth("depth", AlertSeverity::Warning, 100)
.with_clear_depth(50)
.with_cooldown(Duration::from_secs(0));
let mut eval = AlertEvaluator::new().with_rule(rule);
let now = Utc::now();
let low = EventSnapshotBuilder::new()
.timestamp(now)
.total_queue_depth(40)
.build();
assert!(eval.observe_snapshot(&low).is_empty());
let high = EventSnapshotBuilder::new()
.timestamp(now)
.total_queue_depth(150)
.build();
let alerts = eval.observe_snapshot(&high);
assert!(alerts.iter().any(|a| a.is_firing()));
assert!(eval.is_firing("depth"));
let mid = EventSnapshotBuilder::new()
.timestamp(now)
.total_queue_depth(60)
.build();
let _ = eval.observe_snapshot(&mid);
assert!(eval.is_firing("depth"), "60 > clear threshold 50");
let cleared = EventSnapshotBuilder::new()
.timestamp(now)
.total_queue_depth(40)
.build();
let alerts = eval.observe_snapshot(&cleared);
assert!(alerts.iter().any(|a| a.status == AlertStatus::Resolved));
assert!(!eval.is_firing("depth"));
}
#[test]
fn test_queue_depth_specific_queue() {
let rule = AlertRule::queue_depth("q-depth", AlertSeverity::Error, 10).for_queue("urgent");
let mut eval = AlertEvaluator::new().with_rule(rule);
let now = Utc::now();
let snap = EventSnapshotBuilder::new()
.timestamp(now)
.queue_depth("urgent", 20)
.queue_depth("bulk", 5)
.build();
let alerts = eval.observe_snapshot(&snap);
assert!(alerts.iter().any(|a| a.is_firing()));
}
#[test]
fn test_no_heartbeat_trips_via_tick() {
let rule = AlertRule::no_heartbeat("hb", AlertSeverity::Critical, Duration::from_secs(30));
let mut eval = AlertEvaluator::new().with_rule(rule);
let base = Utc::now();
let alerts = eval.observe_event(&heartbeat_event(base));
assert!(alerts.is_empty());
let _ = eval.tick(base + chrono::Duration::seconds(10));
assert!(!eval.is_firing("hb"));
let alerts = eval.tick(base + chrono::Duration::seconds(45));
assert!(alerts.iter().any(|a| a.is_firing()));
assert!(eval.is_firing("hb"));
let alerts = eval.observe_event(&heartbeat_event(base + chrono::Duration::seconds(46)));
assert!(alerts.iter().any(|a| a.status == AlertStatus::Resolved));
assert!(!eval.is_firing("hb"));
}
#[test]
fn test_cooldown_suppresses_repeat_firing() {
let rule = AlertRule::queue_depth("depth", AlertSeverity::Warning, 10)
.with_cooldown(Duration::from_secs(300));
let mut eval = AlertEvaluator::new().with_rule(rule);
let now = Utc::now();
let high = EventSnapshotBuilder::new()
.timestamp(now)
.total_queue_depth(50)
.build();
let first = eval.observe_snapshot(&high);
assert_eq!(first.iter().filter(|a| a.is_firing()).count(), 1);
let again = EventSnapshotBuilder::new()
.timestamp(now + chrono::Duration::seconds(5))
.total_queue_depth(50)
.build();
let second = eval.observe_snapshot(&again);
assert!(second.iter().all(|a| !a.is_firing()));
let later = EventSnapshotBuilder::new()
.timestamp(now + chrono::Duration::seconds(400))
.total_queue_depth(50)
.build();
let third = eval.observe_snapshot(&later);
assert_eq!(third.iter().filter(|a| a.is_firing()).count(), 1);
}
#[test]
fn test_multiple_rules_evaluated() {
let mut eval = AlertEvaluator::new()
.with_rule(AlertRule::queue_depth("d", AlertSeverity::Warning, 10))
.with_rule(AlertRule::no_heartbeat(
"hb",
AlertSeverity::Critical,
Duration::from_secs(5),
));
assert_eq!(eval.rule_count(), 2);
let now = Utc::now();
let snap = EventSnapshotBuilder::new()
.timestamp(now)
.total_queue_depth(99)
.worker_count(0)
.build();
let alerts = eval.observe_snapshot(&snap);
assert!(alerts.iter().any(|a| a.rule_name == "d" && a.is_firing()));
assert!(alerts.iter().any(|a| a.rule_name == "hb" && a.is_firing()));
}
#[test]
fn test_failure_rate_window_pruning() {
let rule = AlertRule::failure_rate("fail", AlertSeverity::Error, 0.5)
.with_min_samples(2)
.with_window(Duration::from_secs(10))
.with_cooldown(Duration::from_secs(0));
let mut eval = AlertEvaluator::new().with_rule(rule);
let base = Utc::now();
eval.observe_event(&failed_event(base));
eval.observe_event(&failed_event(base + chrono::Duration::seconds(1)));
assert!(eval.is_firing("fail"));
let mut resolved = false;
for i in 0..3 {
let at = base + chrono::Duration::seconds(100 + i);
let alerts = eval.observe_event(&succeeded_event(at));
if alerts.iter().any(|a| a.status == AlertStatus::Resolved) {
resolved = true;
}
}
assert!(resolved);
assert!(!eval.is_firing("fail"));
}
}