use std::{
sync::Mutex,
time::{Duration, Instant},
};
use imagegen_bridge_core::{BridgeError, ErrorCode};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CircuitBreakerConfig {
pub enabled: bool,
pub failure_threshold: u32,
pub open_duration: Duration,
pub half_open_max_calls: u32,
pub success_threshold: u32,
}
impl Default for CircuitBreakerConfig {
fn default() -> Self {
Self {
enabled: true,
failure_threshold: 5,
open_duration: Duration::from_secs(3 * 60),
half_open_max_calls: 1,
success_threshold: 1,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
Closed,
Open,
HalfOpen,
}
impl CircuitState {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Closed => "closed",
Self::Open => "open",
Self::HalfOpen => "half_open",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CircuitBreakerSnapshot {
pub state: CircuitState,
pub consecutive_failures: u32,
pub retry_after_ms: u64,
pub rejected_calls: u64,
pub transitions: u64,
}
#[derive(Debug)]
pub(crate) struct CircuitBreaker {
config: CircuitBreakerConfig,
inner: Mutex<Inner>,
}
#[derive(Debug)]
struct Inner {
phase: Phase,
epoch: u64,
rejected_calls: u64,
transitions: u64,
}
#[derive(Debug)]
enum Phase {
Closed { consecutive_failures: u32 },
Open { until: Instant },
HalfOpen { in_flight: u32, successes: u32 },
}
#[derive(Debug, Clone, Copy)]
enum PermitKind {
Closed(u64),
HalfOpen(u64),
Disabled,
}
#[derive(Debug, Clone, Copy)]
enum Outcome {
Healthy,
Failed,
Neutral,
}
#[derive(Debug)]
pub(crate) struct CircuitPermit<'a> {
breaker: &'a CircuitBreaker,
kind: PermitKind,
finished: bool,
}
impl CircuitBreaker {
pub(crate) fn new(config: CircuitBreakerConfig) -> Result<Self, BridgeError> {
if config.enabled
&& (config.failure_threshold == 0
|| config.open_duration.is_zero()
|| config.half_open_max_calls == 0
|| config.success_threshold == 0)
{
return Err(BridgeError::new(
ErrorCode::Configuration,
"enabled circuit-breaker limits must be greater than zero",
));
}
Ok(Self {
config,
inner: Mutex::new(Inner {
phase: Phase::Closed {
consecutive_failures: 0,
},
epoch: 0,
rejected_calls: 0,
transitions: 0,
}),
})
}
pub(crate) fn acquire(&self, provider: &str) -> Result<CircuitPermit<'_>, BridgeError> {
if !self.config.enabled {
return Ok(CircuitPermit {
breaker: self,
kind: PermitKind::Disabled,
finished: false,
});
}
let now = Instant::now();
let mut inner = self
.inner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Phase::Open { until } = inner.phase
&& now >= until
{
inner.phase = Phase::HalfOpen {
in_flight: 0,
successes: 0,
};
inner.epoch = inner.epoch.saturating_add(1);
inner.transitions = inner.transitions.saturating_add(1);
}
let epoch = inner.epoch;
let kind = match &mut inner.phase {
Phase::Closed { .. } => PermitKind::Closed(epoch),
Phase::HalfOpen { in_flight, .. } if *in_flight < self.config.half_open_max_calls => {
*in_flight = in_flight.saturating_add(1);
PermitKind::HalfOpen(epoch)
}
Phase::Open { until } => {
let retry_after_ms = duration_ms(until.saturating_duration_since(now));
inner.rejected_calls = inner.rejected_calls.saturating_add(1);
return Err(open_error(provider, retry_after_ms));
}
Phase::HalfOpen { .. } => {
inner.rejected_calls = inner.rejected_calls.saturating_add(1);
return Err(open_error(provider, duration_ms(self.config.open_duration))
.with_detail("circuit_state", "half_open"));
}
};
drop(inner);
Ok(CircuitPermit {
breaker: self,
kind,
finished: false,
})
}
pub(crate) fn snapshot(&self) -> CircuitBreakerSnapshot {
let now = Instant::now();
let inner = self
.inner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let (state, consecutive_failures, retry_after_ms) = match inner.phase {
Phase::Closed {
consecutive_failures,
} => (CircuitState::Closed, consecutive_failures, 0),
Phase::Open { until } => (
CircuitState::Open,
self.config.failure_threshold,
duration_ms(until.saturating_duration_since(now)),
),
Phase::HalfOpen { .. } => (CircuitState::HalfOpen, self.config.failure_threshold, 0),
};
CircuitBreakerSnapshot {
state,
consecutive_failures,
retry_after_ms,
rejected_calls: inner.rejected_calls,
transitions: inner.transitions,
}
}
fn complete(&self, kind: PermitKind, result: Result<(), &BridgeError>) {
if matches!(kind, PermitKind::Disabled) {
return;
}
let outcome = match result {
Ok(()) => Outcome::Healthy,
Err(error) if counts_as_failure(error) => Outcome::Failed,
Err(_) => Outcome::Neutral,
};
let mut inner = self
.inner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let permit_epoch = match kind {
PermitKind::Closed(epoch) | PermitKind::HalfOpen(epoch) => epoch,
PermitKind::Disabled => return,
};
if permit_epoch != inner.epoch {
return;
}
match (&mut inner.phase, kind, outcome) {
(
Phase::Closed {
consecutive_failures,
},
PermitKind::Closed(_),
Outcome::Healthy,
) => {
*consecutive_failures = 0;
}
(
Phase::Closed {
consecutive_failures,
},
PermitKind::Closed(_),
Outcome::Failed,
) => {
*consecutive_failures = consecutive_failures.saturating_add(1);
if *consecutive_failures >= self.config.failure_threshold {
inner.phase = Phase::Open {
until: Instant::now() + self.config.open_duration,
};
inner.epoch = inner.epoch.saturating_add(1);
inner.transitions = inner.transitions.saturating_add(1);
}
}
(
Phase::HalfOpen {
in_flight,
successes,
},
PermitKind::HalfOpen(_),
Outcome::Healthy,
) => {
*in_flight = in_flight.saturating_sub(1);
*successes = successes.saturating_add(1);
if *successes >= self.config.success_threshold && *in_flight == 0 {
inner.phase = Phase::Closed {
consecutive_failures: 0,
};
inner.epoch = inner.epoch.saturating_add(1);
inner.transitions = inner.transitions.saturating_add(1);
}
}
(
Phase::HalfOpen { .. },
PermitKind::HalfOpen(_),
Outcome::Failed | Outcome::Neutral,
) => {
inner.phase = Phase::Open {
until: Instant::now() + self.config.open_duration,
};
inner.epoch = inner.epoch.saturating_add(1);
inner.transitions = inner.transitions.saturating_add(1);
}
_ => {}
}
}
fn abandon(&self, kind: PermitKind) {
let PermitKind::HalfOpen(epoch) = kind else {
return;
};
let mut inner = self
.inner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if epoch == inner.epoch
&& let Phase::HalfOpen { in_flight, .. } = &mut inner.phase
{
*in_flight = in_flight.saturating_sub(1);
}
}
}
impl CircuitPermit<'_> {
pub(crate) fn finish<T>(mut self, result: &Result<T, BridgeError>) {
self.breaker
.complete(self.kind, result.as_ref().map(|_| ()));
self.finished = true;
}
}
impl Drop for CircuitPermit<'_> {
fn drop(&mut self) {
if !self.finished {
self.breaker.abandon(self.kind);
}
}
}
fn counts_as_failure(error: &BridgeError) -> bool {
matches!(
error.code,
ErrorCode::RateLimited
| ErrorCode::Overloaded
| ErrorCode::Timeout
| ErrorCode::Upstream
| ErrorCode::Protocol
)
}
fn open_error(provider: &str, retry_after_ms: u64) -> BridgeError {
BridgeError::new(ErrorCode::Overloaded, "provider circuit breaker is open")
.retryable(true)
.with_provider(provider)
.with_detail("circuit_state", "open")
.with_detail("retry_after_ms", retry_after_ms)
}
fn duration_ms(value: Duration) -> u64 {
u64::try_from(value.as_millis()).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
fn config() -> CircuitBreakerConfig {
CircuitBreakerConfig {
failure_threshold: 2,
open_duration: Duration::from_millis(20),
..CircuitBreakerConfig::default()
}
}
#[test]
fn opens_after_threshold_and_recovers_through_half_open() {
let breaker = CircuitBreaker::new(config()).unwrap();
for _ in 0..2 {
let permit = breaker.acquire("fake").unwrap();
let result: Result<(), BridgeError> =
Err(BridgeError::new(ErrorCode::Upstream, "injected"));
permit.finish(&result);
}
assert_eq!(breaker.snapshot().state, CircuitState::Open);
let rejected = breaker.acquire("fake").unwrap_err();
assert_eq!(rejected.details["circuit_state"], "open");
std::thread::sleep(Duration::from_millis(25));
let permit = breaker.acquire("fake").unwrap();
assert_eq!(breaker.snapshot().state, CircuitState::HalfOpen);
permit.finish(&Ok::<_, BridgeError>(()));
assert_eq!(breaker.snapshot().state, CircuitState::Closed);
}
#[test]
fn slow_success_and_request_errors_do_not_open_the_circuit() {
let breaker = CircuitBreaker::new(config()).unwrap();
for code in [
ErrorCode::SafetyRejected,
ErrorCode::InvalidRequest,
ErrorCode::Cancelled,
] {
let permit = breaker.acquire("fake").unwrap();
permit.finish(&Err::<(), _>(BridgeError::new(
code,
"non-provider failure",
)));
}
let permit = breaker.acquire("fake").unwrap();
std::thread::sleep(Duration::from_millis(30));
permit.finish(&Ok::<_, BridgeError>(()));
assert_eq!(breaker.snapshot().state, CircuitState::Closed);
}
#[test]
fn stale_closed_permit_cannot_contaminate_a_recovered_epoch() {
let policy = CircuitBreakerConfig {
failure_threshold: 1,
open_duration: Duration::from_millis(5),
..CircuitBreakerConfig::default()
};
let breaker = CircuitBreaker::new(policy).unwrap();
let stale = breaker.acquire("fake").unwrap();
let opener = breaker.acquire("fake").unwrap();
opener.finish(&Err::<(), _>(BridgeError::new(ErrorCode::Upstream, "open")));
std::thread::sleep(Duration::from_millis(8));
breaker
.acquire("fake")
.unwrap()
.finish(&Ok::<_, BridgeError>(()));
stale.finish(&Err::<(), _>(BridgeError::new(
ErrorCode::Upstream,
"stale",
)));
assert_eq!(breaker.snapshot().state, CircuitState::Closed);
assert_eq!(breaker.snapshot().consecutive_failures, 0);
}
#[test]
fn concurrent_half_open_probes_settle_before_closing() {
let policy = CircuitBreakerConfig {
failure_threshold: 1,
open_duration: Duration::from_millis(5),
half_open_max_calls: 2,
success_threshold: 1,
..CircuitBreakerConfig::default()
};
let breaker = CircuitBreaker::new(policy).unwrap();
breaker
.acquire("fake")
.unwrap()
.finish(&Err::<(), _>(BridgeError::new(ErrorCode::Upstream, "open")));
std::thread::sleep(Duration::from_millis(8));
let healthy = breaker.acquire("fake").unwrap();
let failed = breaker.acquire("fake").unwrap();
healthy.finish(&Ok::<_, BridgeError>(()));
assert_eq!(breaker.snapshot().state, CircuitState::HalfOpen);
failed.finish(&Err::<(), _>(BridgeError::new(
ErrorCode::Upstream,
"probe failed",
)));
assert_eq!(breaker.snapshot().state, CircuitState::Open);
}
}