use std::{
net::SocketAddr,
sync::Mutex,
time::{Duration, Instant},
};
use dashmap::DashMap;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CircuitState {
Closed,
Open,
HalfOpen,
}
pub enum CircuitCheck {
Allowed(CircuitToken),
Rejected,
}
pub struct CircuitToken {
generation: u64,
}
#[derive(Clone, Debug)]
pub struct CircuitBreakerConfig {
pub threshold: u32,
pub recovery_window: Duration,
pub half_open_timeout: Duration,
}
#[derive(Debug)]
pub struct CircuitBreaker {
inner: Mutex<CircuitInner>,
config: CircuitBreakerConfig,
}
#[derive(Debug)]
struct CircuitInner {
consecutive_failures: u32,
half_opened_at: Option<Instant>,
opened_at: Option<Instant>,
last_activity: Instant,
state: CircuitState,
generation: u64,
}
impl CircuitInner {
fn issue_token(&self) -> CircuitCheck {
CircuitCheck::Allowed(CircuitToken {
generation: self.generation,
})
}
fn transition_to_half_open(&mut self) -> CircuitCheck {
self.generation = self.generation.wrapping_add(1);
self.state = CircuitState::HalfOpen;
self.half_opened_at = Some(Instant::now());
self.issue_token()
}
fn try_open_to_half_open(&mut self, config: &CircuitBreakerConfig) -> CircuitCheck {
if self.opened_at.is_some_and(|t| t.elapsed() >= config.recovery_window) {
self.transition_to_half_open()
} else {
CircuitCheck::Rejected
}
}
fn try_reset_stale_probe(&mut self, config: &CircuitBreakerConfig) -> CircuitCheck {
if self
.half_opened_at
.is_some_and(|t| t.elapsed() >= config.half_open_timeout)
{
self.state = CircuitState::Open;
self.opened_at = Some(Instant::now());
self.half_opened_at = None;
self.try_open_to_half_open(config)
} else {
CircuitCheck::Rejected
}
}
}
impl CircuitBreaker {
pub fn new(config: CircuitBreakerConfig) -> Self {
Self {
inner: Mutex::new(CircuitInner {
consecutive_failures: 0,
half_opened_at: None,
opened_at: None,
last_activity: Instant::now(),
state: CircuitState::Closed,
generation: 0,
}),
config,
}
}
#[expect(clippy::expect_used, reason = "poisoned mutex is unrecoverable")]
pub fn precheck(&self) -> bool {
let inner = self.inner.lock().expect("circuit breaker lock poisoned");
match inner.state {
CircuitState::Closed => true,
CircuitState::Open => inner
.opened_at
.is_some_and(|t| t.elapsed() >= self.config.recovery_window),
CircuitState::HalfOpen => inner
.half_opened_at
.is_some_and(|t| t.elapsed() >= self.config.half_open_timeout),
}
}
#[expect(clippy::expect_used, reason = "poisoned mutex is unrecoverable")]
pub fn try_acquire(&self) -> CircuitCheck {
let mut inner = self.inner.lock().expect("circuit breaker lock poisoned");
match inner.state {
CircuitState::Closed => inner.issue_token(),
CircuitState::Open => inner.try_open_to_half_open(&self.config),
CircuitState::HalfOpen => inner.try_reset_stale_probe(&self.config),
}
}
#[expect(clippy::expect_used, reason = "poisoned mutex is unrecoverable")]
#[expect(clippy::needless_pass_by_value, reason = "consumed to prevent double-recording")]
pub fn record_success(&self, token: CircuitToken) {
let mut inner = self.inner.lock().expect("circuit breaker lock poisoned");
if token.generation != inner.generation {
return;
}
inner.last_activity = Instant::now();
match inner.state {
CircuitState::Closed => {
inner.consecutive_failures = 0;
},
CircuitState::HalfOpen => {
inner.state = CircuitState::Closed;
inner.consecutive_failures = 0;
inner.half_opened_at = None;
inner.opened_at = None;
},
CircuitState::Open => {},
}
}
#[expect(clippy::expect_used, reason = "poisoned mutex is unrecoverable")]
#[expect(clippy::needless_pass_by_value, reason = "consumed to prevent double-recording")]
pub fn record_failure(&self, token: CircuitToken) {
let mut inner = self.inner.lock().expect("circuit breaker lock poisoned");
if token.generation != inner.generation {
return;
}
inner.last_activity = Instant::now();
match inner.state {
CircuitState::Closed => {
inner.consecutive_failures = inner.consecutive_failures.saturating_add(1);
if inner.consecutive_failures >= self.config.threshold {
inner.state = CircuitState::Open;
inner.opened_at = Some(Instant::now());
}
},
CircuitState::HalfOpen => {
inner.state = CircuitState::Open;
inner.opened_at = Some(Instant::now());
inner.half_opened_at = None;
},
CircuitState::Open => {},
}
}
#[expect(clippy::expect_used, reason = "poisoned mutex is unrecoverable")]
fn is_idle(&self, idle_threshold: Duration) -> bool {
let inner = self.inner.lock().expect("circuit breaker lock poisoned");
inner.state == CircuitState::Closed
&& inner.consecutive_failures == 0
&& inner.last_activity.elapsed() >= idle_threshold
}
#[cfg(test)]
#[expect(clippy::expect_used, reason = "poisoned mutex is unrecoverable")]
pub fn state(&self) -> CircuitState {
self.inner.lock().expect("circuit breaker lock poisoned").state
}
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct PeerKey {
addr: SocketAddr,
sni: String,
}
impl PeerKey {
pub fn new(addr: SocketAddr, sni: impl Into<String>) -> Self {
Self { addr, sni: sni.into() }
}
}
impl std::fmt::Display for PeerKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
if self.sni.is_empty() {
write!(f, "{}", self.addr)
} else {
write!(f, "{} ({})", self.addr, self.sni)
}
}
}
#[derive(Debug)]
pub struct CircuitBreakerRegistry {
breakers: DashMap<PeerKey, CircuitBreaker>,
config: CircuitBreakerConfig,
}
impl CircuitBreakerRegistry {
pub fn new(config: CircuitBreakerConfig) -> Self {
Self {
breakers: DashMap::new(),
config,
}
}
pub fn precheck(&self, peer: &PeerKey) -> bool {
self.breakers.get(peer).is_none_or(|cb| cb.precheck())
}
pub fn try_acquire(&self, peer: PeerKey) -> CircuitCheck {
self.breakers
.entry(peer)
.or_insert_with(|| CircuitBreaker::new(self.config.clone()))
.try_acquire()
}
pub fn record_success(&self, peer: &PeerKey, token: CircuitToken) {
if let Some(cb) = self.breakers.get(peer) {
cb.record_success(token);
}
}
pub fn record_failure(&self, peer: &PeerKey, token: CircuitToken) {
if let Some(cb) = self.breakers.get(peer) {
cb.record_failure(token);
}
}
pub fn evict_idle(&self, idle_threshold: Duration) -> usize {
let stale: Vec<PeerKey> = self
.breakers
.iter()
.filter(|entry| entry.value().is_idle(idle_threshold))
.map(|entry| entry.key().clone())
.collect();
let count = stale.len();
for key in &stale {
self.breakers.remove(key);
}
count
}
#[cfg(test)]
pub fn len(&self) -> usize {
self.breakers.len()
}
#[cfg(test)]
pub fn is_empty(&self) -> bool {
self.breakers.is_empty()
}
}
#[cfg(test)]
#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
#[allow(clippy::unwrap_used, clippy::expect_used, reason = "tests")]
mod tests {
use super::*;
fn config(threshold: u32, recovery_ms: u64, half_open_ms: u64) -> CircuitBreakerConfig {
CircuitBreakerConfig {
threshold,
recovery_window: Duration::from_millis(recovery_ms),
half_open_timeout: Duration::from_millis(half_open_ms),
}
}
#[test]
fn starts_in_closed_state() {
let cb = CircuitBreaker::new(config(3, 30_000, 9_999_000));
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn stays_closed_below_threshold() {
let cb = CircuitBreaker::new(config(3, 30_000, 9_999_000));
let t1 = cb.try_acquire();
record_failure_from_check(&cb, t1);
let t2 = cb.try_acquire();
record_failure_from_check(&cb, t2);
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn trips_to_open_at_threshold() {
let cb = CircuitBreaker::new(config(3, 30_000, 9_999_000));
for _ in 0..3 {
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
}
assert_eq!(cb.state(), CircuitState::Open);
}
#[test]
fn success_resets_failure_count() {
let cb = CircuitBreaker::new(config(3, 30_000, 9_999_000));
let t1 = cb.try_acquire();
record_failure_from_check(&cb, t1);
let t2 = cb.try_acquire();
record_failure_from_check(&cb, t2);
let t3 = cb.try_acquire();
record_success_from_check(&cb, t3);
let t4 = cb.try_acquire();
record_failure_from_check(&cb, t4);
let t5 = cb.try_acquire();
record_failure_from_check(&cb, t5);
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn open_rejects_via_try_acquire() {
let cb = CircuitBreaker::new(config(1, 9_999_000, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
assert!(matches!(cb.try_acquire(), CircuitCheck::Rejected));
}
#[test]
fn half_open_after_recovery_window() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
assert_eq!(cb.state(), CircuitState::Open);
assert!(matches!(cb.try_acquire(), CircuitCheck::Allowed(_)));
assert_eq!(cb.state(), CircuitState::HalfOpen);
}
#[test]
fn half_open_success_transitions_to_closed() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let probe = cb.try_acquire();
record_success_from_check(&cb, probe);
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn half_open_failure_transitions_to_open() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let probe = cb.try_acquire();
record_failure_from_check(&cb, probe);
assert_eq!(cb.state(), CircuitState::Open);
}
#[test]
fn half_open_allows_only_one_probe() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let first = cb.try_acquire();
assert!(matches!(first, CircuitCheck::Allowed(_)));
assert!(matches!(cb.try_acquire(), CircuitCheck::Rejected));
assert!(matches!(cb.try_acquire(), CircuitCheck::Rejected));
}
#[test]
fn open_record_failure_is_noop() {
let cb = CircuitBreaker::new(config(1, 9_999_000, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
assert_eq!(cb.state(), CircuitState::Open);
}
#[test]
fn multiple_successes_keep_closed() {
let cb = CircuitBreaker::new(config(3, 30_000, 9_999_000));
for _ in 0..10 {
let t = cb.try_acquire();
record_success_from_check(&cb, t);
}
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn half_open_timeout_resets_to_open() {
let cb = CircuitBreaker::new(config(1, 0, 0));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let _probe = cb.try_acquire();
let next = cb.try_acquire();
assert!(matches!(next, CircuitCheck::Allowed(_)));
}
#[test]
fn half_open_timeout_does_not_fire_before_expiry() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let _probe = cb.try_acquire();
assert!(matches!(cb.try_acquire(), CircuitCheck::Rejected));
assert_eq!(cb.state(), CircuitState::HalfOpen);
}
#[test]
fn stale_probe_success_ignored() {
let cb = CircuitBreaker::new(config(1, 0, 0));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let stale_probe = cb.try_acquire();
let fresh_probe = cb.try_acquire();
record_success_from_check(&cb, stale_probe);
assert_eq!(
cb.state(),
CircuitState::HalfOpen,
"stale probe success must not close the circuit"
);
record_success_from_check(&cb, fresh_probe);
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn stale_probe_failure_ignored() {
let cb = CircuitBreaker::new(config(1, 0, 0));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let stale_probe = cb.try_acquire();
let fresh_probe = cb.try_acquire();
record_failure_from_check(&cb, stale_probe);
record_success_from_check(&cb, fresh_probe);
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn closed_tokens_always_match_generation() {
let cb = CircuitBreaker::new(config(3, 30_000, 9_999_000));
let t1 = cb.try_acquire();
record_failure_from_check(&cb, t1);
assert_eq!(cb.state(), CircuitState::Closed);
let t2 = cb.try_acquire();
record_failure_from_check(&cb, t2);
assert_eq!(cb.state(), CircuitState::Closed);
let t3 = cb.try_acquire();
record_failure_from_check(&cb, t3);
assert_eq!(cb.state(), CircuitState::Open);
}
#[test]
fn precheck_returns_true_when_closed() {
let cb = CircuitBreaker::new(config(3, 30_000, 9_999_000));
assert!(cb.precheck());
}
#[test]
fn precheck_returns_false_when_open_not_recovered() {
let cb = CircuitBreaker::new(config(1, 9_999_000, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
assert!(!cb.precheck());
}
#[test]
fn precheck_returns_true_when_open_recovered() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
assert!(cb.precheck(), "recovery_window=0 means immediately recoverable");
}
#[test]
fn precheck_returns_false_when_half_open_probe_in_flight() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let _probe = cb.try_acquire();
assert!(!cb.precheck());
}
#[test]
fn precheck_returns_true_when_half_open_probe_timed_out() {
let cb = CircuitBreaker::new(config(1, 0, 0));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
let _probe = cb.try_acquire();
assert!(cb.precheck(), "half_open_timeout=0 means timed out");
}
#[test]
fn precheck_does_not_mutate_state() {
let cb = CircuitBreaker::new(config(1, 0, 9_999_000));
let t = cb.try_acquire();
record_failure_from_check(&cb, t);
assert_eq!(cb.state(), CircuitState::Open);
let _ = cb.precheck();
assert_eq!(cb.state(), CircuitState::Open);
}
fn peer(addr: &str) -> PeerKey {
PeerKey::new(addr.parse().unwrap(), "")
}
fn peer_with_sni(addr: &str, sni: &str) -> PeerKey {
PeerKey::new(addr.parse().unwrap(), sni)
}
#[test]
fn registry_creates_breaker_on_first_access() {
let registry = CircuitBreakerRegistry::new(config(3, 30_000, 9_999_000));
let key = peer("127.0.0.1:8080");
assert!(registry.precheck(&key));
}
#[test]
fn registry_isolates_peers() {
let registry = CircuitBreakerRegistry::new(config(1, 9_999_000, 9_999_000));
let a = peer("127.0.0.1:8080");
let b = peer("127.0.0.1:9090");
let t = registry.try_acquire(a.clone());
record_registry_failure(®istry, &a, t);
assert!(!registry.precheck(&a), "peer a should be open");
assert!(registry.precheck(&b), "peer b should be unaffected");
}
#[test]
fn registry_isolates_peers_by_sni() {
let registry = CircuitBreakerRegistry::new(config(1, 9_999_000, 9_999_000));
let a = peer_with_sni("127.0.0.1:443", "api.example.com");
let b = peer_with_sni("127.0.0.1:443", "web.example.com");
let t = registry.try_acquire(a.clone());
record_registry_failure(®istry, &a, t);
assert!(!registry.precheck(&a), "api peer should be open");
assert!(registry.precheck(&b), "web peer should be unaffected");
}
#[test]
fn registry_propagates_config() {
let registry = CircuitBreakerRegistry::new(config(2, 9_999_000, 9_999_000));
let key = peer("127.0.0.1:8080");
let t1 = registry.try_acquire(key.clone());
record_registry_failure(®istry, &key, t1);
assert!(registry.precheck(&key), "one failure should not trip threshold=2");
let t2 = registry.try_acquire(key.clone());
record_registry_failure(®istry, &key, t2);
assert!(!registry.precheck(&key), "two failures should trip threshold=2");
}
#[test]
fn evict_idle_removes_healthy_idle_entries() {
let registry = CircuitBreakerRegistry::new(config(3, 30_000, 9_999_000));
let a = peer("127.0.0.1:8080");
let b = peer("127.0.0.1:9090");
let ta = registry.try_acquire(a.clone());
record_registry_success(®istry, &a, ta);
let tb = registry.try_acquire(b.clone());
record_registry_success(®istry, &b, tb);
assert_eq!(registry.len(), 2);
let evicted = registry.evict_idle(Duration::ZERO);
assert_eq!(evicted, 2, "both idle entries should be evicted");
assert_eq!(registry.len(), 0);
}
#[test]
fn evict_idle_preserves_active_entries() {
let registry = CircuitBreakerRegistry::new(config(1, 9_999_000, 9_999_000));
let a = peer("127.0.0.1:8080");
let b = peer("127.0.0.1:9090");
let ta = registry.try_acquire(a.clone());
record_registry_failure(®istry, &a, ta);
let tb = registry.try_acquire(b.clone());
record_registry_success(®istry, &b, tb);
let evicted = registry.evict_idle(Duration::ZERO);
assert_eq!(evicted, 1, "only the healthy idle peer should be evicted");
assert_eq!(registry.len(), 1);
assert!(!registry.precheck(&a), "open circuit should survive eviction");
}
#[test]
fn peer_key_display_without_sni() {
let key = peer("127.0.0.1:8080");
assert_eq!(key.to_string(), "127.0.0.1:8080");
}
#[test]
fn peer_key_display_with_sni() {
let key = peer_with_sni("127.0.0.1:443", "api.example.com");
assert_eq!(key.to_string(), "127.0.0.1:443 (api.example.com)");
}
fn record_success_from_check(cb: &CircuitBreaker, check: CircuitCheck) {
if let CircuitCheck::Allowed(token) = check {
cb.record_success(token);
}
}
fn record_failure_from_check(cb: &CircuitBreaker, check: CircuitCheck) {
if let CircuitCheck::Allowed(token) = check {
cb.record_failure(token);
}
}
fn record_registry_success(registry: &CircuitBreakerRegistry, key: &PeerKey, check: CircuitCheck) {
if let CircuitCheck::Allowed(token) = check {
registry.record_success(key, token);
}
}
fn record_registry_failure(registry: &CircuitBreakerRegistry, key: &PeerKey, check: CircuitCheck) {
if let CircuitCheck::Allowed(token) = check {
registry.record_failure(key, token);
}
}
}