gitdataai/libs/transport/circuit_breaker.rs

use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;

#[derive(Clone)]
pub struct CircuitBreaker {
    state: Arc<Mutex<CircuitState>>,
    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<Instant>,
    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<F, T, E>(&self, f: F) -> Result<T, CircuitBreakerError<E>>
    where
        F: std::future::Future<Output = Result<T, E>>,
    {
        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<E> {
    Open,
    Inner(E),
}

impl<E: std::fmt::Display> std::fmt::Display for CircuitBreakerError<E> {
    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<E: std::error::Error> std::error::Error for CircuitBreakerError<E> {}