use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
#[derive(Debug, Clone)]
pub struct BreakerConfig {
pub failure_threshold: u32,
pub reset_after_secs: u64,
}
impl Default for BreakerConfig {
fn default() -> Self {
Self {
failure_threshold: 5,
reset_after_secs: 60,
}
}
}
#[derive(Debug)]
struct ProviderState {
failures: AtomicU32,
first_failure_at: AtomicU64,
opened_at: AtomicU64,
}
impl ProviderState {
fn new() -> Self {
Self {
failures: AtomicU32::new(0),
first_failure_at: AtomicU64::new(0),
opened_at: AtomicU64::new(0),
}
}
}
pub struct ProviderHealthRegistry {
states: RwLock<HashMap<String, ProviderState>>,
config: BreakerConfig,
}
impl ProviderHealthRegistry {
pub fn new(config: BreakerConfig) -> Self {
Self {
states: RwLock::new(HashMap::new()),
config,
}
}
pub fn is_healthy(&self, provider: &str) -> bool {
let states = self.states.read();
let state = match states.get(provider) {
Some(s) => s,
None => return true, };
let opened = state.opened_at.load(Ordering::SeqCst);
if opened == 0 {
return true; }
let elapsed = now_ms().saturating_sub(opened);
elapsed >= self.config.reset_after_secs * 1000
}
pub fn record_success(&self, provider: &str) {
let mut states = self.states.write();
let state = states
.entry(provider.to_string())
.or_insert_with(ProviderState::new);
state.failures.store(0, Ordering::SeqCst);
state.first_failure_at.store(0, Ordering::SeqCst);
state.opened_at.store(0, Ordering::SeqCst);
}
pub fn record_failure(&self, provider: &str) {
let mut states = self.states.write();
let state = states
.entry(provider.to_string())
.or_insert_with(ProviderState::new);
let n = state.failures.fetch_add(1, Ordering::SeqCst) + 1;
let _ = state.first_failure_at.compare_exchange(
0,
now_ms(),
Ordering::SeqCst,
Ordering::SeqCst,
);
if n >= self.config.failure_threshold {
state.opened_at.store(now_ms(), Ordering::SeqCst);
state.first_failure_at.store(0, Ordering::SeqCst);
}
}
}
fn now_ms() -> u64 {
use std::time::SystemTime;
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn healthy_by_default() {
let reg = ProviderHealthRegistry::new(BreakerConfig::default());
assert!(reg.is_healthy("unknown-provider"));
}
#[test]
fn opens_after_threshold_failures() {
let reg = ProviderHealthRegistry::new(BreakerConfig {
failure_threshold: 3,
reset_after_secs: 60,
});
reg.record_failure("p");
reg.record_failure("p");
assert!(reg.is_healthy("p"));
reg.record_failure("p");
assert!(!reg.is_healthy("p"));
}
#[test]
fn success_resets_failures() {
let reg = ProviderHealthRegistry::new(BreakerConfig {
failure_threshold: 3,
reset_after_secs: 3600, });
reg.record_failure("p");
reg.record_failure("p");
reg.record_success("p");
reg.record_failure("p");
assert!(reg.is_healthy("p"));
}
}