nobreak/

lib.rs

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

pub fn breaker() -> CircuitBreaker {
    builder().build()
}

pub async fn call_async<T, TFut, TOk, TError>(f: T) -> Result<TOk, CircuitError<TError>>
where
    T: FnOnce() -> TFut,
    TFut: Future<Output = Result<TOk, TError>>,
{
    builder().build().call_async(f).await
}

pub fn call<T, TOk, TError>(f: T) -> Result<TOk, CircuitError<TError>>
where
    T: FnOnce() -> Result<TOk, TError>,
{
    builder().build().call(f)
}

pub fn builder() -> CircuitBreakerBuilder {
    CircuitBreakerBuilder {
        failure_threshold: 5,
        success_threshold: 2,
        open_duration: Duration::from_secs(30),
    }
}

pub struct CircuitBreakerBuilder {
    failure_threshold: usize,
    success_threshold: usize,
    open_duration: Duration,
}

impl CircuitBreakerBuilder {
    pub fn with_failure_threshold(mut self, threshold: usize) -> Self {
        self.failure_threshold = threshold;
        self
    }

    pub fn with_success_threshold(mut self, threshold: usize) -> Self {
        self.success_threshold = threshold;
        self
    }

    pub fn with_open_duration(mut self, duration: Duration) -> Self {
        self.open_duration = duration;
        self
    }

    pub fn build(self) -> CircuitBreaker {
        CircuitBreaker {
            failure_threshold: self.failure_threshold,
            success_threshold: self.success_threshold,
            open_duration: self.open_duration,
            inner: Mutex::new(Inner {
                state: State::Closed,
                failure_count: 0,
                success_count: 0,
            }),
        }
    }
}

pub struct CircuitBreaker {
    failure_threshold: usize,
    success_threshold: usize,
    open_duration: Duration,
    inner: Mutex<Inner>,
}

struct Inner {
    state: State,
    failure_count: usize,
    success_count: usize,
}

enum State {
    Closed,
    Open { since: Instant },
    HalfOpen,
}

impl CircuitBreaker {
    pub fn call<T, TOk, TError>(&self, f: T) -> Result<TOk, CircuitError<TError>>
    where
        T: FnOnce() -> Result<TOk, TError>,
    {
        {
            let mut inner = self.inner.lock().unwrap();
            match inner.state {
                State::Open { since } => {
                    if since.elapsed() >= self.open_duration {
                        inner.state = State::HalfOpen;
                        inner.success_count = 0;
                    } else {
                        return Err(CircuitError::Open);
                    }
                }
                State::Closed | State::HalfOpen => {}
            }
        }

        match f() {
            Ok(val) => {
                self.record_success();
                Ok(val)
            }
            Err(err) => {
                self.record_failure();
                Err(CircuitError::Failed { error: err })
            }
        }
    }

    pub async fn call_async<T, TFut, TOk, TError>(&self, f: T) -> Result<TOk, CircuitError<TError>>
    where
        T: FnOnce() -> TFut,
        TFut: Future<Output = Result<TOk, TError>>,
    {
        {
            let mut inner = self.inner.lock().unwrap();
            match inner.state {
                State::Open { since } => {
                    if since.elapsed() >= self.open_duration {
                        inner.state = State::HalfOpen;
                        inner.success_count = 0;
                    } else {
                        return Err(CircuitError::Open);
                    }
                }
                State::Closed | State::HalfOpen => {}
            }
        }

        match f().await {
            Ok(val) => {
                self.record_success();
                Ok(val)
            }
            Err(err) => {
                self.record_failure();
                Err(CircuitError::Failed { error: err })
            }
        }
    }

    fn record_success(&self) {
        let mut inner = self.inner.lock().unwrap();
        match inner.state {
            State::HalfOpen => {
                inner.success_count += 1;
                if inner.success_count >= self.success_threshold {
                    inner.state = State::Closed;
                    inner.failure_count = 0;
                    inner.success_count = 0;
                }
            }
            _ => {
                inner.failure_count = 0;
            }
        }
    }

    fn record_failure(&self) {
        let mut inner = self.inner.lock().unwrap();
        match inner.state {
            State::HalfOpen => {
                inner.state = State::Open {
                    since: Instant::now(),
                };
                inner.failure_count = 0;
                inner.success_count = 0;
            }
            _ => {
                inner.failure_count += 1;
                if inner.failure_count >= self.failure_threshold {
                    inner.state = State::Open {
                        since: Instant::now(),
                    };
                }
            }
        }
    }

    /// Returns the current state of the circuit breaker.
    pub fn state(&self) -> CircuitState {
        let inner = self.inner.lock().unwrap();
        match inner.state {
            State::Closed => CircuitState::Closed,
            State::Open { since } => {
                if since.elapsed() >= self.open_duration {
                    CircuitState::HalfOpen
                } else {
                    CircuitState::Open
                }
            }
            State::HalfOpen => CircuitState::HalfOpen,
        }
    }
}

/// The observable state of the circuit breaker.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
    Closed,
    Open,
    HalfOpen,
}

#[derive(Debug)]
pub enum CircuitError<E> {
    /// The operation was executed but returned an error.
    Failed { error: E },
    /// The circuit is open; the operation was not executed.
    Open,
}

impl<E> CircuitError<E> {
    /// Extract the underlying error, if one exists.
    pub fn into_inner(self) -> Option<E> {
        match self {
            CircuitError::Failed { error } => Some(error),
            CircuitError::Open => None,
        }
    }
}

impl<E: fmt::Display> fmt::Display for CircuitError<E> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CircuitError::Failed { error } => write!(f, "{error}"),
            CircuitError::Open => write!(f, "circuit breaker is open"),
        }
    }
}

impl<E: std::error::Error + 'static> std::error::Error for CircuitError<E> {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            CircuitError::Failed { error } => Some(error),
            CircuitError::Open => None,
        }
    }
}