memlink_runtime/

circuit.rs

//! Circuit breaker pattern for fault tolerance.
//!
//! Prevents cascading failures by stopping requests to failing modules
//! and allowing them time to recover.

use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Circuit breaker state.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
    /// Circuit is closed, requests flow normally.
    Closed,
    /// Circuit is open, requests are rejected immediately.
    Open,
    /// Circuit is half-open, testing if service recovered.
    HalfOpen,
}

/// Circuit breaker configuration.
#[derive(Debug, Clone)]
pub struct CircuitConfig {
    /// Number of consecutive failures before opening.
    pub failure_threshold: u32,
    /// Number of successes in half-open state to close.
    pub success_threshold: u32,
    /// Time to wait before transitioning from open to half-open.
    pub open_timeout: Duration,
    /// Time window for counting failures.
    pub failure_window: Duration,
    /// Minimum calls before circuit can open.
    pub min_calls: u32,
}

impl Default for CircuitConfig {
    fn default() -> Self {
        Self {
            failure_threshold: 5,
            success_threshold: 3,
            open_timeout: Duration::from_secs(30),
            failure_window: Duration::from_secs(60),
            min_calls: 10,
        }
    }
}

/// Statistics about circuit breaker state.
#[derive(Debug, Clone)]
pub struct CircuitStats {
    pub state: CircuitState,
    pub consecutive_failures: u32,
    pub consecutive_successes: u32,
    pub total_calls: u64,
    pub total_failures: u64,
    pub total_successes: u64,
    pub total_rejected: u64,
    pub last_failure: Option<Instant>,
    pub last_state_change: Instant,
}

/// Circuit breaker for a module.
#[derive(Debug)]
pub struct CircuitBreaker {
    /// Module name.
    module_name: String,
    /// Current state.
    state: std::sync::Mutex<CircuitState>,
    /// Consecutive failure count.
    consecutive_failures: AtomicU32,
    /// Consecutive success count (in half-open).
    consecutive_successes: AtomicU32,
    /// Total calls.
    total_calls: AtomicU64,
    /// Total failures.
    total_failures: AtomicU64,
    /// Total successes.
    total_successes: AtomicU64,
    /// Total rejected (circuit open).
    total_rejected: AtomicU64,
    /// Timestamp of last failure.
    last_failure: std::sync::Mutex<Option<Instant>>,
    /// Timestamp of last state change.
    last_state_change: std::sync::Mutex<Instant>,
    /// Configuration.
    config: CircuitConfig,
}

impl CircuitBreaker {
    /// Creates a new circuit breaker.
    pub fn new(module_name: String, config: CircuitConfig) -> Self {
        Self {
            module_name,
            state: std::sync::Mutex::new(CircuitState::Closed),
            consecutive_failures: AtomicU32::new(0),
            consecutive_successes: AtomicU32::new(0),
            total_calls: AtomicU64::new(0),
            total_failures: AtomicU64::new(0),
            total_successes: AtomicU64::new(0),
            total_rejected: AtomicU64::new(0),
            last_failure: std::sync::Mutex::new(None),
            last_state_change: std::sync::Mutex::new(Instant::now()),
            config,
        }
    }

    /// Creates a circuit breaker with default configuration.
    pub fn with_defaults(module_name: String) -> Self {
        Self::new(module_name, CircuitConfig::default())
    }

    /// Checks if a request can proceed.
    pub fn can_execute(&self) -> bool {
        let mut state = self.state.lock().unwrap();

        match *state {
            CircuitState::Closed => true,
            CircuitState::Open => {
                // Check if timeout has elapsed
                let elapsed = self.last_state_change.lock().unwrap().elapsed();
                if elapsed >= self.config.open_timeout {
                    // Transition to half-open
                    *state = CircuitState::HalfOpen;
                    *self.last_state_change.lock().unwrap() = Instant::now();
                    self.consecutive_successes.store(0, Ordering::Relaxed);
                    true
                } else {
                    self.total_rejected.fetch_add(1, Ordering::Relaxed);
                    false
                }
            }
            CircuitState::HalfOpen => true,
        }
    }

    /// Records a successful call.
    pub fn record_success(&self) {
        self.total_calls.fetch_add(1, Ordering::Relaxed);
        self.total_successes.fetch_add(1, Ordering::Relaxed);
        self.consecutive_failures.store(0, Ordering::Relaxed);

        let mut state = self.state.lock().unwrap();

        if *state == CircuitState::HalfOpen {
            let successes = self.consecutive_successes.fetch_add(1, Ordering::Relaxed) + 1;
            if successes >= self.config.success_threshold {
                // Transition to closed
                *state = CircuitState::Closed;
                *self.last_state_change.lock().unwrap() = Instant::now();
                self.consecutive_failures.store(0, Ordering::Relaxed);
            }
        }
    }

    /// Records a failed call.
    pub fn record_failure(&self) {
        self.total_calls.fetch_add(1, Ordering::Relaxed);
        self.total_failures.fetch_add(1, Ordering::Relaxed);

        let failures = self.consecutive_failures.fetch_add(1, Ordering::Relaxed) + 1;
        *self.last_failure.lock().unwrap() = Some(Instant::now());

        let mut state = self.state.lock().unwrap();

        if *state == CircuitState::HalfOpen {
            // Any failure in half-open goes back to open
            *state = CircuitState::Open;
            *self.last_state_change.lock().unwrap() = Instant::now();
        } else if *state == CircuitState::Closed {
            // Check if we should open
            if failures >= self.config.failure_threshold {
                let total = self.total_calls.load(Ordering::Relaxed);
                if total >= self.config.min_calls as u64 {
                    *state = CircuitState::Open;
                    *self.last_state_change.lock().unwrap() = Instant::now();
                }
            }
        }
    }

    /// Returns the current state.
    pub fn state(&self) -> CircuitState {
        *self.state.lock().unwrap()
    }

    /// Returns statistics.
    pub fn stats(&self) -> CircuitStats {
        CircuitStats {
            state: self.state(),
            consecutive_failures: self.consecutive_failures.load(Ordering::Relaxed),
            consecutive_successes: self.consecutive_successes.load(Ordering::Relaxed),
            total_calls: self.total_calls.load(Ordering::Relaxed),
            total_failures: self.total_failures.load(Ordering::Relaxed),
            total_successes: self.total_successes.load(Ordering::Relaxed),
            total_rejected: self.total_rejected.load(Ordering::Relaxed),
            last_failure: *self.last_failure.lock().unwrap(),
            last_state_change: *self.last_state_change.lock().unwrap(),
        }
    }

    /// Returns the module name.
    pub fn module_name(&self) -> &str {
        &self.module_name
    }

    /// Forces the circuit to open (for testing/manual intervention).
    pub fn force_open(&self) {
        *self.state.lock().unwrap() = CircuitState::Open;
        *self.last_state_change.lock().unwrap() = Instant::now();
    }

    /// Forces the circuit to close (for testing/manual intervention).
    pub fn force_close(&self) {
        *self.state.lock().unwrap() = CircuitState::Closed;
        *self.last_state_change.lock().unwrap() = Instant::now();
        self.consecutive_failures.store(0, Ordering::Relaxed);
    }

    /// Resets all counters.
    pub fn reset(&self) {
        self.consecutive_failures.store(0, Ordering::Relaxed);
        self.consecutive_successes.store(0, Ordering::Relaxed);
        self.total_calls.store(0, Ordering::Relaxed);
        self.total_failures.store(0, Ordering::Relaxed);
        self.total_successes.store(0, Ordering::Relaxed);
        self.total_rejected.store(0, Ordering::Relaxed);
        *self.last_failure.lock().unwrap() = None;
        self.force_close();
    }
}

/// Registry of circuit breakers for all modules.
#[derive(Debug)]
pub struct CircuitRegistry {
    /// Circuit breakers by module name.
    circuits: DashMap<String, Arc<CircuitBreaker>>,
    /// Default configuration.
    default_config: CircuitConfig,
}

impl CircuitRegistry {
    /// Creates a new circuit registry.
    pub fn new() -> Self {
        Self {
            circuits: DashMap::new(),
            default_config: CircuitConfig::default(),
        }
    }

    /// Creates a registry with custom default configuration.
    pub fn with_config(config: CircuitConfig) -> Self {
        Self {
            circuits: DashMap::new(),
            default_config: config,
        }
    }

    /// Gets or creates a circuit breaker for a module.
    pub fn get_or_create(&self, module_name: &str) -> Arc<CircuitBreaker> {
        self.circuits
            .entry(module_name.to_string())
            .or_insert_with(|| {
                Arc::new(CircuitBreaker::new(
                    module_name.to_string(),
                    self.default_config.clone(),
                ))
            })
            .clone()
    }

    /// Registers a module with custom configuration.
    pub fn register(&self, module_name: &str, config: CircuitConfig) {
        self.circuits.insert(
            module_name.to_string(),
            Arc::new(CircuitBreaker::new(module_name.to_string(), config)),
        );
    }

    /// Checks if a request can proceed.
    pub fn can_execute(&self, module_name: &str) -> bool {
        self.get_or_create(module_name).can_execute()
    }

    /// Records a success.
    pub fn record_success(&self, module_name: &str) {
        self.get_or_create(module_name).record_success();
    }

    /// Records a failure.
    pub fn record_failure(&self, module_name: &str) {
        self.get_or_create(module_name).record_failure();
    }

    /// Returns statistics for a module.
    pub fn stats(&self, module_name: &str) -> Option<CircuitStats> {
        self.circuits.get(module_name).map(|c| c.stats())
    }

    /// Returns all circuit breaker stats.
    pub fn all_stats(&self) -> Vec<(String, CircuitStats)> {
        self.circuits
            .iter()
            .map(|e| (e.key().clone(), e.value().stats()))
            .collect()
    }

    /// Returns the number of registered circuits.
    pub fn circuit_count(&self) -> usize {
        self.circuits.len()
    }

    /// Returns the number of open circuits.
    pub fn open_circuit_count(&self) -> usize {
        self.circuits
            .iter()
            .filter(|e| *e.value().state.lock().unwrap() == CircuitState::Open)
            .count()
    }
}

impl Default for CircuitRegistry {
    fn default() -> Self {
        Self::new()
    }
}

use dashmap::DashMap;

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

    #[test]
    fn test_circuit_closed_initially() {
        let cb = CircuitBreaker::with_defaults("test".to_string());
        assert_eq!(cb.state(), CircuitState::Closed);
        assert!(cb.can_execute());
    }

    #[test]
    fn test_circuit_opens_after_failures() {
        let config = CircuitConfig {
            failure_threshold: 3,
            min_calls: 0,
            ..CircuitConfig::default()
        };
        let cb = CircuitBreaker::new("test".to_string(), config);

        // Record failures
        for _ in 0..3 {
            cb.record_failure();
        }

        assert_eq!(cb.state(), CircuitState::Open);
        assert!(!cb.can_execute());
    }

    #[test]
    fn test_circuit_half_open_after_timeout() {
        let config = CircuitConfig {
            failure_threshold: 1,
            min_calls: 0,
            open_timeout: Duration::from_millis(100),
            ..CircuitConfig::default()
        };
        let cb = CircuitBreaker::new("test".to_string(), config);

        // Open the circuit
        cb.record_failure();
        assert_eq!(cb.state(), CircuitState::Open);

        // Wait for timeout
        std::thread::sleep(Duration::from_millis(150));

        // Should transition to half-open
        assert!(cb.can_execute());
        assert_eq!(cb.state(), CircuitState::HalfOpen);
    }

    #[test]
    fn test_circuit_closes_after_successes() {
        let config = CircuitConfig {
            failure_threshold: 1,
            success_threshold: 2,
            min_calls: 0,
            open_timeout: Duration::from_millis(10),
            ..CircuitConfig::default()
        };
        let cb = CircuitBreaker::new("test".to_string(), config);

        // Open the circuit
        cb.record_failure();

        // Wait for half-open
        std::thread::sleep(Duration::from_millis(50));
        cb.can_execute(); // Triggers transition to half-open

        // Record successes
        cb.record_success();
        cb.record_success();

        assert_eq!(cb.state(), CircuitState::Closed);
    }

    #[test]
    fn test_circuit_reopens_on_failure_in_half_open() {
        let config = CircuitConfig {
            failure_threshold: 1,
            success_threshold: 2,
            min_calls: 0,
            open_timeout: Duration::from_millis(10),
            ..CircuitConfig::default()
        };
        let cb = CircuitBreaker::new("test".to_string(), config);

        // Open the circuit
        cb.record_failure();

        // Wait for half-open
        std::thread::sleep(Duration::from_millis(50));
        cb.can_execute();

        // Record one success
        cb.record_success();

        // Record failure - should go back to open
        cb.record_failure();

        assert_eq!(cb.state(), CircuitState::Open);
    }

    #[test]
    fn test_circuit_registry() {
        let registry = CircuitRegistry::new();

        let cb1 = registry.get_or_create("module1");
        let cb2 = registry.get_or_create("module2");

        assert_eq!(registry.circuit_count(), 2);

        cb1.record_failure();
        cb2.record_success();

        let stats1 = registry.stats("module1").unwrap();
        let stats2 = registry.stats("module2").unwrap();

        assert_eq!(stats1.total_failures, 1);
        assert_eq!(stats2.total_successes, 1);
    }
}