halldyll_core/fetch/

circuit_breaker.rs

//! Circuit Breaker - Prevents cascade failures per domain
//!
//! Implements the circuit breaker pattern to protect against failing domains:
//! - **Closed**: Normal operation, requests pass through
//! - **Open**: Domain is failing, requests are rejected immediately
//! - **Half-Open**: Testing if domain recovered
//!
//! ## Usage
//!
//! ```rust,ignore
//! let breaker = CircuitBreaker::new(CircuitBreakerConfig::default());
//! 
//! // Before making a request
//! if !breaker.allow_request("example.com") {
//!     return Err(Error::CircuitOpen);
//! }
//!
//! // After request
//! match result {
//!     Ok(_) => breaker.record_success("example.com"),
//!     Err(e) => breaker.record_failure("example.com"),
//! }
//! ```

use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{Duration, Instant};

/// Circuit breaker configuration
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
    /// Failure threshold before opening circuit
    pub failure_threshold: u32,
    /// Success threshold to close circuit from half-open
    pub success_threshold: u32,
    /// Duration the circuit stays open before going to half-open
    pub open_duration: Duration,
    /// Time window to count failures
    pub failure_window: Duration,
    /// Timeout considered as failure
    pub timeout_as_failure: bool,
    /// 5xx errors considered as failure
    pub server_error_as_failure: bool,
    /// 429 rate limit considered as failure
    pub rate_limit_as_failure: bool,
}

impl Default for CircuitBreakerConfig {
    fn default() -> Self {
        Self {
            failure_threshold: 5,
            success_threshold: 2,
            open_duration: Duration::from_secs(30),
            failure_window: Duration::from_secs(60),
            timeout_as_failure: true,
            server_error_as_failure: true,
            rate_limit_as_failure: false, // Rate limits are expected, not failures
        }
    }
}

impl CircuitBreakerConfig {
    /// Production preset - more tolerant, longer recovery
    pub fn production() -> Self {
        Self {
            failure_threshold: 10,
            success_threshold: 3,
            open_duration: Duration::from_secs(60),
            failure_window: Duration::from_secs(120),
            timeout_as_failure: true,
            server_error_as_failure: true,
            rate_limit_as_failure: false,
        }
    }

    /// Aggressive preset - quick to open, quick to recover
    pub fn aggressive() -> Self {
        Self {
            failure_threshold: 3,
            success_threshold: 1,
            open_duration: Duration::from_secs(15),
            failure_window: Duration::from_secs(30),
            timeout_as_failure: true,
            server_error_as_failure: true,
            rate_limit_as_failure: true,
        }
    }
}

/// Circuit state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
    /// Circuit is closed, requests pass through
    Closed,
    /// Circuit is open, requests are rejected
    Open,
    /// Testing if domain has recovered
    HalfOpen,
}

/// Per-domain circuit state
#[derive(Debug)]
struct DomainCircuit {
    state: CircuitState,
    failures: Vec<Instant>,
    successes_in_half_open: u32,
    opened_at: Option<Instant>,
    last_failure: Option<Instant>,
}

impl DomainCircuit {
    fn new() -> Self {
        Self {
            state: CircuitState::Closed,
            failures: Vec::new(),
            successes_in_half_open: 0,
            opened_at: None,
            last_failure: None,
        }
    }

    /// Count recent failures within the window
    fn recent_failures(&self, window: Duration) -> u32 {
        let cutoff = Instant::now() - window;
        self.failures.iter().filter(|&&t| t > cutoff).count() as u32
    }

    /// Clean up old failures
    fn cleanup_old_failures(&mut self, window: Duration) {
        let cutoff = Instant::now() - window;
        self.failures.retain(|&t| t > cutoff);
    }
}

/// Circuit breaker for multiple domains
pub struct CircuitBreaker {
    config: CircuitBreakerConfig,
    circuits: RwLock<HashMap<String, DomainCircuit>>,
}

impl CircuitBreaker {
    /// Create new circuit breaker
    pub fn new(config: CircuitBreakerConfig) -> Self {
        Self {
            config,
            circuits: RwLock::new(HashMap::new()),
        }
    }

    /// Create with default config
    pub fn default_config() -> Self {
        Self::new(CircuitBreakerConfig::default())
    }

    /// Check if a request to this domain is allowed
    pub fn allow_request(&self, domain: &str) -> bool {
        let mut circuits = self.circuits.write().unwrap();
        let circuit = circuits.entry(domain.to_string()).or_insert_with(DomainCircuit::new);

        match circuit.state {
            CircuitState::Closed => true,
            CircuitState::Open => {
                // Check if we should transition to half-open
                if let Some(opened_at) = circuit.opened_at {
                    if opened_at.elapsed() >= self.config.open_duration {
                        circuit.state = CircuitState::HalfOpen;
                        circuit.successes_in_half_open = 0;
                        true
                    } else {
                        false
                    }
                } else {
                    false
                }
            }
            CircuitState::HalfOpen => true,
        }
    }

    /// Record a successful request
    pub fn record_success(&self, domain: &str) {
        let mut circuits = self.circuits.write().unwrap();
        if let Some(circuit) = circuits.get_mut(domain) {
            match circuit.state {
                CircuitState::HalfOpen => {
                    circuit.successes_in_half_open += 1;
                    if circuit.successes_in_half_open >= self.config.success_threshold {
                        // Transition back to closed
                        circuit.state = CircuitState::Closed;
                        circuit.failures.clear();
                        circuit.opened_at = None;
                        circuit.successes_in_half_open = 0;
                    }
                }
                CircuitState::Closed => {
                    // Nothing special, just cleanup old failures
                    circuit.cleanup_old_failures(self.config.failure_window);
                }
                CircuitState::Open => {
                    // Shouldn't happen, but handle gracefully
                }
            }
        }
    }

    /// Record a failed request
    pub fn record_failure(&self, domain: &str) {
        let mut circuits = self.circuits.write().unwrap();
        let circuit = circuits.entry(domain.to_string()).or_insert_with(DomainCircuit::new);

        circuit.failures.push(Instant::now());
        circuit.last_failure = Some(Instant::now());
        circuit.cleanup_old_failures(self.config.failure_window);

        match circuit.state {
            CircuitState::Closed => {
                if circuit.recent_failures(self.config.failure_window) >= self.config.failure_threshold {
                    // Open the circuit
                    circuit.state = CircuitState::Open;
                    circuit.opened_at = Some(Instant::now());
                }
            }
            CircuitState::HalfOpen => {
                // Any failure in half-open reopens the circuit
                circuit.state = CircuitState::Open;
                circuit.opened_at = Some(Instant::now());
                circuit.successes_in_half_open = 0;
            }
            CircuitState::Open => {
                // Already open, refresh the timer
                circuit.opened_at = Some(Instant::now());
            }
        }
    }

    /// Record a timeout (may or may not count as failure based on config)
    pub fn record_timeout(&self, domain: &str) {
        if self.config.timeout_as_failure {
            self.record_failure(domain);
        }
    }

    /// Record a server error (5xx)
    pub fn record_server_error(&self, domain: &str) {
        if self.config.server_error_as_failure {
            self.record_failure(domain);
        }
    }

    /// Record a rate limit (429)
    pub fn record_rate_limit(&self, domain: &str) {
        if self.config.rate_limit_as_failure {
            self.record_failure(domain);
        }
    }

    /// Get circuit state for a domain
    pub fn get_state(&self, domain: &str) -> CircuitState {
        let circuits = self.circuits.read().unwrap();
        circuits.get(domain).map(|c| c.state).unwrap_or(CircuitState::Closed)
    }

    /// Get all open circuits (for monitoring)
    pub fn get_open_circuits(&self) -> Vec<String> {
        let circuits = self.circuits.read().unwrap();
        circuits
            .iter()
            .filter(|(_, c)| c.state == CircuitState::Open)
            .map(|(domain, _)| domain.clone())
            .collect()
    }

    /// Reset circuit for a domain
    pub fn reset(&self, domain: &str) {
        let mut circuits = self.circuits.write().unwrap();
        circuits.remove(domain);
    }

    /// Reset all circuits
    pub fn reset_all(&self) {
        let mut circuits = self.circuits.write().unwrap();
        circuits.clear();
    }

    /// Get circuit statistics
    pub fn stats(&self) -> CircuitBreakerStats {
        let circuits = self.circuits.read().unwrap();
        let total = circuits.len();
        let open = circuits.values().filter(|c| c.state == CircuitState::Open).count();
        let half_open = circuits.values().filter(|c| c.state == CircuitState::HalfOpen).count();
        let closed = circuits.values().filter(|c| c.state == CircuitState::Closed).count();

        CircuitBreakerStats {
            total_domains: total,
            open_circuits: open,
            half_open_circuits: half_open,
            closed_circuits: closed,
        }
    }
}

/// Circuit breaker statistics
#[derive(Debug, Clone)]
pub struct CircuitBreakerStats {
    /// Total tracked domains
    pub total_domains: usize,
    /// Number of open circuits
    pub open_circuits: usize,
    /// Number of half-open circuits
    pub half_open_circuits: usize,
    /// Number of closed circuits
    pub closed_circuits: usize,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_circuit_starts_closed() {
        let breaker = CircuitBreaker::default_config();
        assert!(breaker.allow_request("example.com"));
        assert_eq!(breaker.get_state("example.com"), CircuitState::Closed);
    }

    #[test]
    fn test_circuit_opens_after_failures() {
        let config = CircuitBreakerConfig {
            failure_threshold: 3,
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        // Record failures
        for _ in 0..3 {
            breaker.record_failure("example.com");
        }

        assert_eq!(breaker.get_state("example.com"), CircuitState::Open);
        assert!(!breaker.allow_request("example.com"));
    }

    #[test]
    fn test_circuit_transitions_to_half_open() {
        let config = CircuitBreakerConfig {
            failure_threshold: 2,
            open_duration: Duration::from_millis(10),
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        // Open the circuit
        breaker.record_failure("example.com");
        breaker.record_failure("example.com");
        assert_eq!(breaker.get_state("example.com"), CircuitState::Open);

        // Wait for open duration
        std::thread::sleep(Duration::from_millis(15));

        // Should transition to half-open on next request
        assert!(breaker.allow_request("example.com"));
        assert_eq!(breaker.get_state("example.com"), CircuitState::HalfOpen);
    }

    #[test]
    fn test_circuit_closes_after_successes() {
        let config = CircuitBreakerConfig {
            failure_threshold: 2,
            success_threshold: 2,
            open_duration: Duration::from_millis(10),
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        // Open the circuit
        breaker.record_failure("example.com");
        breaker.record_failure("example.com");

        // Wait and transition to half-open
        std::thread::sleep(Duration::from_millis(15));
        breaker.allow_request("example.com");

        // Record successes
        breaker.record_success("example.com");
        breaker.record_success("example.com");

        assert_eq!(breaker.get_state("example.com"), CircuitState::Closed);
    }

    #[test]
    fn test_stats() {
        let config = CircuitBreakerConfig {
            failure_threshold: 2,
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        // Create some circuits
        breaker.allow_request("good.com");
        breaker.record_failure("bad.com");
        breaker.record_failure("bad.com");

        let stats = breaker.stats();
        assert_eq!(stats.total_domains, 2);
        assert_eq!(stats.open_circuits, 1);
        assert_eq!(stats.closed_circuits, 1);
    }
}