use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CircuitState {
Closed = 0,
Open = 1,
HalfOpen = 2,
}
impl CircuitState {
fn from_u8(val: u8) -> Self {
match val {
0 => CircuitState::Closed,
1 => CircuitState::Open,
2 => CircuitState::HalfOpen,
_ => CircuitState::Closed,
}
}
}
#[derive(Debug)]
pub struct CircuitBreaker {
state: Arc<AtomicU8>,
failures: Arc<AtomicUsize>,
last_failure: Arc<Mutex<Instant>>,
half_open_probe_in_flight: Arc<AtomicBool>,
config: Arc<CircuitBreakerConfig>,
}
#[derive(Clone, Debug)]
pub struct CircuitBreakerConfig {
pub failure_threshold: usize,
pub open_duration_secs: u64,
pub reset_timeout_secs: u64,
}
impl Default for CircuitBreakerConfig {
fn default() -> Self {
Self {
failure_threshold: 5,
open_duration_secs: 30,
reset_timeout_secs: 60,
}
}
}
impl CircuitBreaker {
pub fn new() -> Self {
Self::with_config(CircuitBreakerConfig::default())
}
pub fn with_config(config: CircuitBreakerConfig) -> Self {
Self {
state: Arc::new(AtomicU8::new(0)), failures: Arc::new(AtomicUsize::new(0)),
last_failure: Arc::new(Mutex::new(Instant::now())),
half_open_probe_in_flight: Arc::new(AtomicBool::new(false)),
config: Arc::new(config),
}
}
pub fn current_state(&self) -> CircuitState {
let state_byte = self.state.load(Ordering::Relaxed);
CircuitState::from_u8(state_byte)
}
pub fn should_attempt(&self) -> bool {
match self.current_state() {
CircuitState::Closed => true,
CircuitState::Open => {
let last_failure = *self.last_failure.lock().unwrap();
let elapsed = last_failure.elapsed();
if elapsed <= Duration::from_secs(self.config.open_duration_secs) {
return false;
}
if self
.half_open_probe_in_flight
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
.is_err()
{
return false;
}
if self
.state
.compare_exchange(1, 2, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
true
} else {
self.half_open_probe_in_flight
.store(false, Ordering::Release);
false
}
}
CircuitState::HalfOpen => self
.half_open_probe_in_flight
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
.is_ok(),
}
}
pub fn record_success(&self) {
let last_failure = *self.last_failure.lock().unwrap();
if last_failure.elapsed() > Duration::from_secs(self.config.reset_timeout_secs) {
self.failures.store(0, Ordering::Relaxed);
self.state.store(0, Ordering::Relaxed); self.half_open_probe_in_flight
.store(false, Ordering::Release);
} else if self.current_state() == CircuitState::HalfOpen {
self.failures.store(0, Ordering::Relaxed);
self.state.store(0, Ordering::Relaxed);
self.half_open_probe_in_flight
.store(false, Ordering::Release);
}
}
pub fn record_failure(&self) {
let failure_count = self.failures.fetch_add(1, Ordering::Relaxed) + 1;
*self.last_failure.lock().unwrap() = Instant::now();
if failure_count >= self.config.failure_threshold
|| self.current_state() == CircuitState::HalfOpen
{
self.state.store(1, Ordering::Relaxed); self.half_open_probe_in_flight
.store(false, Ordering::Release);
}
}
pub fn failure_count(&self) -> usize {
self.failures.load(Ordering::Relaxed)
}
pub fn time_since_last_failure(&self) -> Duration {
self.last_failure.lock().unwrap().elapsed()
}
pub fn reset(&self) {
self.state.store(0, Ordering::Relaxed);
self.failures.store(0, Ordering::Relaxed);
self.half_open_probe_in_flight
.store(false, Ordering::Release);
*self.last_failure.lock().unwrap() = Instant::now();
}
}
impl Default for CircuitBreaker {
fn default() -> Self {
Self::new()
}
}
impl Clone for CircuitBreaker {
fn clone(&self) -> Self {
Self {
state: Arc::clone(&self.state),
failures: Arc::clone(&self.failures),
last_failure: Arc::clone(&self.last_failure),
half_open_probe_in_flight: Arc::clone(&self.half_open_probe_in_flight),
config: Arc::clone(&self.config),
}
}
}
pub struct JitteredBackoff {
attempt: u32,
max_backoff_ms: u64,
}
impl JitteredBackoff {
pub fn new() -> Self {
Self {
attempt: 0,
max_backoff_ms: 30_000, }
}
pub fn with_max_backoff(max_backoff_ms: u64) -> Self {
Self {
attempt: 0,
max_backoff_ms,
}
}
pub fn backoff_duration(&self) -> Duration {
let base_ms = (100u64 * 2_u64.pow(self.attempt.min(10))).min(self.max_backoff_ms);
use rand::Rng;
let mut rng = rand::thread_rng();
let jitter_ms = rng.gen_range(0..=(base_ms / 10));
Duration::from_millis(base_ms + jitter_ms)
}
pub fn next_backoff(&mut self) -> Duration {
let duration = self.backoff_duration();
self.attempt += 1;
duration
}
pub fn reset(&mut self) {
self.attempt = 0;
}
}
impl Default for JitteredBackoff {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn circuit_starts_closed() {
let cb = CircuitBreaker::new();
assert_eq!(cb.current_state(), CircuitState::Closed);
assert!(cb.should_attempt());
}
#[test]
fn circuit_opens_after_threshold_failures() {
let cb = CircuitBreaker::new();
for _ in 0..5 {
cb.record_failure();
}
assert_eq!(cb.current_state(), CircuitState::Open);
assert!(!cb.should_attempt());
}
#[test]
fn circuit_half_opens_after_timeout() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
open_duration_secs: 0,
reset_timeout_secs: 60,
};
let cb = CircuitBreaker::with_config(config);
cb.record_failure();
assert_eq!(cb.current_state(), CircuitState::Open);
std::thread::sleep(Duration::from_millis(100));
assert!(cb.should_attempt());
assert_eq!(cb.current_state(), CircuitState::HalfOpen);
}
#[test]
fn circuit_closes_on_successful_probe() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
open_duration_secs: 0,
reset_timeout_secs: 0,
};
let cb = CircuitBreaker::with_config(config);
cb.record_failure();
assert_eq!(cb.current_state(), CircuitState::Open);
std::thread::sleep(Duration::from_millis(100));
let can_attempt = cb.should_attempt();
assert!(can_attempt);
assert_eq!(cb.current_state(), CircuitState::HalfOpen);
cb.record_success();
assert_eq!(cb.current_state(), CircuitState::Closed);
assert_eq!(cb.failure_count(), 0);
}
#[test]
fn jittered_backoff_increases() {
let mut backoff = JitteredBackoff::new();
let backoff1 = backoff.next_backoff();
let backoff2 = backoff.next_backoff();
let backoff3 = backoff.next_backoff();
assert!(backoff1.as_millis() >= 100 && backoff1.as_millis() <= 220);
assert!(backoff2.as_millis() >= 200 && backoff2.as_millis() <= 440);
assert!(backoff3.as_millis() >= 400 && backoff3.as_millis() <= 880);
}
#[test]
fn jittered_backoff_has_jitter() {
let mut backoffs = Vec::new();
for _ in 0..10 {
let backoff = JitteredBackoff::new().next_backoff();
backoffs.push(backoff.as_millis());
}
let min = *backoffs.iter().min().unwrap();
let max = *backoffs.iter().max().unwrap();
assert!(max > min, "Jitter should produce different values");
}
#[test]
fn half_open_grants_probe_to_exactly_one_concurrent_caller() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
open_duration_secs: 0,
reset_timeout_secs: 60,
};
let cb = CircuitBreaker::with_config(config);
cb.record_failure();
assert_eq!(cb.current_state(), CircuitState::Open);
std::thread::sleep(Duration::from_millis(50));
let granted = Arc::new(AtomicUsize::new(0));
let handles: Vec<_> = (0..16)
.map(|_| {
let cb = cb.clone();
let granted = Arc::clone(&granted);
std::thread::spawn(move || {
if cb.should_attempt() {
granted.fetch_add(1, Ordering::Relaxed);
}
})
})
.collect();
for h in handles {
h.join().unwrap();
}
assert_eq!(
granted.load(Ordering::Relaxed),
1,
"exactly one concurrent caller should win the half-open probe slot"
);
assert_eq!(cb.current_state(), CircuitState::HalfOpen);
}
#[test]
fn circuit_breaker_is_cloneable() {
let cb1 = CircuitBreaker::new();
cb1.record_failure();
let cb2 = cb1.clone();
assert_eq!(cb2.failure_count(), 1);
cb2.record_failure();
assert_eq!(cb1.failure_count(), 2);
}
}