use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::Mutex; #[derive(Clone)] pub struct CircuitBreaker { state: Arc>, config: CircuitConfig, } #[derive(Clone)] struct CircuitConfig { failure_threshold: u32, success_threshold: u32, timeout: Duration, half_open_max_calls: u32, } struct CircuitState { status: CircuitStatus, failure_count: u32, success_count: u32, last_failure_time: Option, half_open_calls: u32, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum CircuitStatus { Closed, Open, HalfOpen, } impl CircuitBreaker { pub fn new() -> Self { Self { state: Arc::new(Mutex::new(CircuitState { status: CircuitStatus::Closed, failure_count: 0, success_count: 0, last_failure_time: None, half_open_calls: 0, })), config: CircuitConfig { failure_threshold: 5, success_threshold: 2, timeout: Duration::from_secs(60), half_open_max_calls: 3, }, } } pub async fn call(&self, f: F) -> Result> where F: std::future::Future>, { let mut state = self.state.lock().await; match state.status { CircuitStatus::Open => { if let Some(last_failure) = state.last_failure_time { if last_failure.elapsed() > self.config.timeout { state.status = CircuitStatus::HalfOpen; state.half_open_calls = 0; } else { return Err(CircuitBreakerError::Open); } } } CircuitStatus::HalfOpen => { if state.half_open_calls >= self.config.half_open_max_calls { return Err(CircuitBreakerError::Open); } state.half_open_calls += 1; } CircuitStatus::Closed => {} } drop(state); match f.await { Ok(result) => { self.on_success().await; Ok(result) } Err(e) => { self.on_failure().await; Err(CircuitBreakerError::Inner(e)) } } } async fn on_success(&self) { let mut state = self.state.lock().await; state.failure_count = 0; if state.status == CircuitStatus::HalfOpen { state.success_count += 1; if state.success_count >= self.config.success_threshold { state.status = CircuitStatus::Closed; state.success_count = 0; } } } async fn on_failure(&self) { let mut state = self.state.lock().await; state.failure_count += 1; state.last_failure_time = Some(Instant::now()); if state.status == CircuitStatus::HalfOpen { state.status = CircuitStatus::Open; state.success_count = 0; } else if state.failure_count >= self.config.failure_threshold { state.status = CircuitStatus::Open; } } pub async fn is_open(&self) -> bool { let state = self.state.lock().await; state.status == CircuitStatus::Open } } #[derive(Debug)] pub enum CircuitBreakerError { Open, Inner(E), } impl std::fmt::Display for CircuitBreakerError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { CircuitBreakerError::Open => write!(f, "Circuit breaker is open"), CircuitBreakerError::Inner(e) => write!(f, "{}", e), } } } impl std::error::Error for CircuitBreakerError {}