use std::time::{Duration, Instant};
const MAX_FAILURES: u32 = 5;
const COOLDOWN_INITIAL: Duration = Duration::from_secs(60);
const COOLDOWN_MAX: Duration = Duration::from_secs(1_800);
const COOLDOWN_MULTIPLIER: u32 = 5;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
Closed,
Open,
HalfOpen,
}
pub struct CircuitBreaker {
state: State,
failure_count: u32,
cooldown_until: Option<Instant>,
backoff: Duration,
}
impl Default for CircuitBreaker {
fn default() -> Self {
Self {
state: State::Closed,
failure_count: 0,
cooldown_until: None,
backoff: COOLDOWN_INITIAL,
}
}
}
impl CircuitBreaker {
pub fn new() -> Self {
Self::default()
}
pub fn can_attempt(&mut self) -> bool {
match self.state {
State::Closed => true,
State::Open => {
if self
.cooldown_until
.is_none_or(|until| Instant::now() >= until)
{
self.state = State::HalfOpen;
true
} else {
false
}
}
State::HalfOpen => true,
}
}
pub fn cooldown_remaining_secs(&self) -> Option<u64> {
match (self.state, self.cooldown_until) {
(State::Open, Some(until)) => {
let now = Instant::now();
if until > now {
Some((until - now).as_secs())
} else {
None
}
}
_ => None,
}
}
pub fn record_success(&mut self) {
self.state = State::Closed;
self.failure_count = 0;
self.cooldown_until = None;
self.backoff = COOLDOWN_INITIAL;
}
pub fn record_failure(&mut self) {
self.failure_count += 1;
match self.state {
State::HalfOpen => {
self.backoff = (self.backoff * COOLDOWN_MULTIPLIER).min(COOLDOWN_MAX);
self.open();
}
State::Closed if self.failure_count >= MAX_FAILURES => {
self.backoff = COOLDOWN_INITIAL;
self.open();
}
_ => {}
}
}
fn open(&mut self) {
self.state = State::Open;
self.cooldown_until = Some(Instant::now() + self.backoff);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn closed_allows_attempts_until_max_failures() {
let mut cb = CircuitBreaker::new();
for _ in 0..MAX_FAILURES - 1 {
assert!(cb.can_attempt());
cb.record_failure();
}
assert!(cb.can_attempt());
cb.record_failure();
assert!(!cb.can_attempt());
}
#[test]
fn success_resets_to_closed() {
let mut cb = CircuitBreaker::new();
for _ in 0..MAX_FAILURES {
cb.record_failure();
}
assert!(!cb.can_attempt());
cb.state = State::HalfOpen;
cb.record_success();
assert_eq!(cb.state, State::Closed);
assert!(cb.can_attempt());
}
#[test]
fn half_open_failure_escalates_backoff() {
let mut cb = CircuitBreaker::new();
for _ in 0..MAX_FAILURES {
cb.record_failure();
}
assert_eq!(cb.backoff, COOLDOWN_INITIAL);
cb.state = State::HalfOpen;
cb.record_failure();
assert_eq!(cb.backoff, COOLDOWN_INITIAL * COOLDOWN_MULTIPLIER);
}
#[test]
fn backoff_caps_at_max() {
let mut cb = CircuitBreaker::new();
for _ in 0..MAX_FAILURES {
cb.record_failure();
}
cb.backoff = COOLDOWN_MAX;
cb.state = State::HalfOpen;
cb.record_failure();
assert_eq!(cb.backoff, COOLDOWN_MAX);
}
}