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(),
};
}
}
}
}
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,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
Closed,
Open,
HalfOpen,
}
#[derive(Debug)]
pub enum CircuitError<E> {
Failed { error: E },
Open,
}
impl<E> CircuitError<E> {
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,
}
}
}