use std::sync::atomic::{AtomicU32, Ordering};
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
Closed,
Open,
HalfOpen,
}
pub struct CircuitBreaker {
state: std::sync::Mutex<CircuitState>,
failure_count: AtomicU32,
success_count: AtomicU32,
failure_threshold: u32,
success_threshold: u32,
timeout: Duration,
last_failure: std::sync::Mutex<Option<Instant>>,
}
impl CircuitBreaker {
pub fn new(failure_threshold: u32, success_threshold: u32, timeout: Duration) -> Self {
Self {
state: std::sync::Mutex::new(CircuitState::Closed),
failure_count: AtomicU32::new(0),
success_count: AtomicU32::new(0),
failure_threshold,
success_threshold,
timeout,
last_failure: std::sync::Mutex::new(None),
}
}
pub fn state(&self) -> CircuitState {
let state = *self.state.lock().unwrap();
match state {
CircuitState::Open => {
if let Some(last) = *self.last_failure.lock().unwrap() {
if last.elapsed() >= self.timeout {
return CircuitState::HalfOpen;
}
}
CircuitState::Open
}
_ => state,
}
}
pub fn allow_request(&self) -> bool {
match self.state() {
CircuitState::Closed => true,
CircuitState::HalfOpen => true,
CircuitState::Open => false,
}
}
pub fn record_success(&self) {
let observed = self.state();
let mut state = self.state.lock().unwrap();
match observed {
CircuitState::HalfOpen => {
let count = self.success_count.fetch_add(1, Ordering::SeqCst) + 1;
if count >= self.success_threshold {
*state = CircuitState::Closed;
self.failure_count.store(0, Ordering::SeqCst);
self.success_count.store(0, Ordering::SeqCst);
}
}
CircuitState::Closed => {
self.failure_count.store(0, Ordering::SeqCst);
}
_ => {}
}
}
pub fn record_failure(&self) {
let mut state = self.state.lock().unwrap();
*self.last_failure.lock().unwrap() = Some(Instant::now());
match *state {
CircuitState::Closed => {
let count = self.failure_count.fetch_add(1, Ordering::SeqCst) + 1;
if count >= self.failure_threshold {
*state = CircuitState::Open;
}
}
CircuitState::HalfOpen => {
*state = CircuitState::Open;
self.success_count.store(0, Ordering::SeqCst);
}
_ => {}
}
}
pub fn reset(&self) {
*self.state.lock().unwrap() = CircuitState::Closed;
self.failure_count.store(0, Ordering::SeqCst);
self.success_count.store(0, Ordering::SeqCst);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn starts_closed() {
let cb = CircuitBreaker::new(3, 2, Duration::from_secs(1));
assert_eq!(cb.state(), CircuitState::Closed);
assert!(cb.allow_request());
}
#[test]
fn opens_after_threshold() {
let cb = CircuitBreaker::new(3, 2, Duration::from_secs(1));
cb.record_failure();
cb.record_failure();
assert!(cb.allow_request());
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
assert!(!cb.allow_request());
}
#[test]
fn success_resets_failure_count() {
let cb = CircuitBreaker::new(3, 2, Duration::from_secs(1));
cb.record_failure();
cb.record_failure();
cb.record_success();
cb.record_failure();
cb.record_failure();
assert!(cb.allow_request()); }
#[test]
fn reset_clears_state() {
let cb = CircuitBreaker::new(1, 1, Duration::from_secs(1));
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
cb.reset();
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn half_open_after_timeout() {
let cb = CircuitBreaker::new(1, 1, Duration::from_millis(10));
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
std::thread::sleep(Duration::from_millis(20));
assert_eq!(cb.state(), CircuitState::HalfOpen);
}
#[test]
fn half_open_success_closes() {
let cb = CircuitBreaker::new(1, 1, Duration::from_millis(10));
cb.record_failure();
std::thread::sleep(Duration::from_millis(20));
assert_eq!(cb.state(), CircuitState::HalfOpen);
cb.record_success();
assert_eq!(cb.state(), CircuitState::Closed);
}
#[test]
fn half_open_failure_reopens() {
let cb = CircuitBreaker::new(1, 1, Duration::from_millis(10));
cb.record_failure();
std::thread::sleep(Duration::from_millis(20));
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
}
}