warpdrive_proxy/middleware/

circuit_breaker.rs

//! Circuit breaker middleware
//!
//! Implements the circuit breaker pattern to prevent cascading failures.
//! States: Closed → Open → Half-Open → Closed

use async_trait::async_trait;
use parking_lot::Mutex;
use pingora::http::ResponseHeader;
use pingora::prelude::*;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::{debug, warn};

use super::{Middleware, MiddlewareContext};
use crate::metrics::{CIRCUIT_BREAKER_STATE, UPSTREAM_FAILURES};

/// Circuit breaker state
#[derive(Debug, Clone, Copy, PartialEq)]
enum CircuitState {
    Closed,
    Open,
    HalfOpen,
}

/// Internal circuit breaker state
#[derive(Debug)]
struct CircuitBreakerState {
    state: CircuitState,
    failure_count: u32,
    last_failure_time: Option<Instant>,
    failure_threshold: u32,
    timeout: Duration,
}

impl CircuitBreakerState {
    fn new(failure_threshold: u32, timeout_secs: u64) -> Self {
        Self {
            state: CircuitState::Closed,
            failure_count: 0,
            last_failure_time: None,
            failure_threshold,
            timeout: Duration::from_secs(timeout_secs),
        }
    }

    fn record_success(&mut self) {
        if self.state == CircuitState::HalfOpen {
            // Success in half-open state -> close circuit
            debug!("Circuit breaker: Half-Open → Closed (success)");
            self.state = CircuitState::Closed;
            CIRCUIT_BREAKER_STATE.set(0); // 0 = Closed
        }
        self.failure_count = 0;
        self.last_failure_time = None;
    }

    fn record_failure(&mut self) {
        self.failure_count += 1;
        self.last_failure_time = Some(Instant::now());

        if self.failure_count >= self.failure_threshold && self.state == CircuitState::Closed {
            warn!(
                "Circuit breaker: Closed → Open (failures: {})",
                self.failure_count
            );
            self.state = CircuitState::Open;
            CIRCUIT_BREAKER_STATE.set(1); // 1 = Open
        }
    }

    fn is_request_allowed(&mut self) -> bool {
        match self.state {
            CircuitState::Closed => true,
            CircuitState::Open => {
                // Check if timeout has elapsed
                if let Some(last_failure) = self.last_failure_time {
                    if last_failure.elapsed() >= self.timeout {
                        debug!("Circuit breaker: Open → Half-Open (timeout elapsed)");
                        self.state = CircuitState::HalfOpen;
                        self.failure_count = 0;
                        CIRCUIT_BREAKER_STATE.set(2); // 2 = Half-Open
                        return true;
                    }
                }
                false
            }
            CircuitState::HalfOpen => true,
        }
    }
}

/// Circuit breaker middleware
///
/// Tracks upstream failures and opens the circuit when threshold is exceeded.
/// While open, requests are rejected with 503 to prevent overwhelming the failing upstream.
pub struct CircuitBreakerMiddleware {
    /// Internal circuit breaker state
    state: Arc<Mutex<CircuitBreakerState>>,
    /// Whether circuit breaker is enabled
    enabled: bool,
}

impl CircuitBreakerMiddleware {
    /// Create a new circuit breaker middleware
    ///
    /// # Arguments
    ///
    /// * `enabled` - Whether to enable circuit breaking
    /// * `failure_threshold` - Number of consecutive failures before opening
    /// * `timeout_secs` - Seconds to wait before trying half-open state
    pub fn new(enabled: bool, failure_threshold: u32, timeout_secs: u64) -> Self {
        debug!(
            "Circuit breaker initialized: enabled={}, threshold={}, timeout={}s",
            enabled, failure_threshold, timeout_secs
        );

        // Initialize circuit breaker state metric to Closed (0)
        if enabled {
            CIRCUIT_BREAKER_STATE.set(0);
        }

        Self {
            state: Arc::new(Mutex::new(CircuitBreakerState::new(
                failure_threshold,
                timeout_secs,
            ))),
            enabled,
        }
    }
}

#[async_trait]
impl Middleware for CircuitBreakerMiddleware {
    /// Check circuit state before forwarding request
    async fn request_filter(
        &self,
        session: &mut Session,
        _ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        if !self.enabled {
            return Ok(());
        }

        // Check if request is allowed (lock scope)
        let allowed = {
            let mut state = self.state.lock();
            state.is_request_allowed()
        }; // Lock released here

        if !allowed {
            warn!("Circuit breaker: request rejected (circuit OPEN)");

            // Return 503 Service Unavailable
            session.respond_error(503).await?;

            return Err(Error::explain(
                ErrorType::HTTPStatus(503),
                "Circuit breaker open - upstream unavailable",
            ));
        }

        debug!("Circuit breaker: request allowed");
        Ok(())
    }

    /// Record success/failure based on upstream response
    async fn response_filter(
        &self,
        _session: &mut Session,
        upstream_response: &mut ResponseHeader,
        _ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        if !self.enabled {
            return Ok(());
        }

        let status = upstream_response.status.as_u16();
        let mut state = self.state.lock();

        // Consider 5xx responses as failures
        if status >= 500 {
            warn!("Circuit breaker: upstream failure (status {})", status);

            // Record upstream failure metric
            UPSTREAM_FAILURES
                .with_label_values(&[&status.to_string()])
                .inc();

            state.record_failure();
        } else {
            debug!("Circuit breaker: upstream success (status {})", status);
            state.record_success();
        }

        Ok(())
    }
}

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

    #[test]
    fn test_circuit_breaker_creation() {
        let cb = CircuitBreakerMiddleware::new(true, 5, 60);
        assert!(cb.enabled);
    }

    #[test]
    fn test_circuit_breaker_disabled() {
        let cb = CircuitBreakerMiddleware::new(false, 5, 60);
        assert!(!cb.enabled);
    }

    #[test]
    fn test_circuit_breaker_state_transitions() {
        let cb = CircuitBreakerMiddleware::new(true, 3, 1);
        let mut state = cb.state.lock();

        // Should start closed
        assert!(state.is_request_allowed());

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

        // Circuit should be open after threshold
        assert_eq!(state.state, CircuitState::Open);
        assert!(!state.is_request_allowed());
    }
}