use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy)]
pub struct BreakerConfig {
pub window: usize,
pub failure_threshold: usize,
pub cooldown: Duration,
}
impl Default for BreakerConfig {
fn default() -> Self {
Self {
window: 10,
failure_threshold: 5,
cooldown: Duration::from_secs(30),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
Closed,
Open,
HalfOpen,
}
#[derive(Debug)]
struct Entry {
state: State,
window: VecDeque<bool>,
opened_at: Option<Instant>,
}
impl Entry {
fn new() -> Self {
Self {
state: State::Closed,
window: VecDeque::new(),
opened_at: None,
}
}
fn failures(&self) -> usize {
self.window.iter().filter(|&&f| f).count()
}
}
#[derive(Clone)]
pub struct CircuitBreaker {
config: BreakerConfig,
entries: Arc<Mutex<HashMap<String, Entry>>>,
}
impl CircuitBreaker {
pub fn new(config: BreakerConfig) -> Self {
Self {
config,
entries: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn check_at(&self, package: &str, now: Instant) -> Result<(), String> {
let mut entries = self.entries.lock().unwrap();
let entry = entries.entry(package.to_string()).or_insert_with(Entry::new);
if entry.state == State::Open {
let cooled = entry
.opened_at
.map(|t| now.duration_since(t) >= self.config.cooldown)
.unwrap_or(true);
if cooled {
entry.state = State::HalfOpen;
return Ok(());
}
return Err(format!(
"package '{package}' circuit is open (too many recent failures); \
rejecting call until cooldown elapses"
));
}
Ok(())
}
pub fn record_at(&self, package: &str, success: bool, now: Instant) {
let mut entries = self.entries.lock().unwrap();
let entry = entries.entry(package.to_string()).or_insert_with(Entry::new);
match entry.state {
State::HalfOpen => {
if success {
entry.state = State::Closed;
entry.window.clear();
entry.opened_at = None;
} else {
entry.state = State::Open;
entry.opened_at = Some(now);
}
}
State::Closed => {
entry.window.push_back(!success);
while entry.window.len() > self.config.window {
entry.window.pop_front();
}
if entry.failures() >= self.config.failure_threshold {
entry.state = State::Open;
entry.opened_at = Some(now);
}
}
State::Open => {
}
}
}
pub fn is_open(&self, package: &str) -> bool {
self.entries
.lock()
.unwrap()
.get(package)
.map(|e| e.state == State::Open)
.unwrap_or(false)
}
pub fn reset(&self, package: &str) {
if let Some(entry) = self.entries.lock().unwrap().get_mut(package) {
entry.state = State::Closed;
entry.window.clear();
entry.opened_at = None;
}
}
pub fn check(&self, package: &str) -> Result<(), String> {
self.check_at(package, Instant::now())
}
pub fn record(&self, package: &str, success: bool) {
self.record_at(package, success, Instant::now())
}
}
impl Default for CircuitBreaker {
fn default() -> Self {
Self::new(BreakerConfig::default())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn cfg() -> BreakerConfig {
BreakerConfig {
window: 5,
failure_threshold: 3,
cooldown: Duration::from_secs(10),
}
}
#[test]
fn healthy_package_stays_closed() {
let cb = CircuitBreaker::new(cfg());
let t = Instant::now();
for _ in 0..20 {
assert!(cb.check_at("p", t).is_ok());
cb.record_at("p", true, t);
}
assert!(!cb.is_open("p"));
}
#[test]
fn trips_open_after_threshold_failures() {
let cb = CircuitBreaker::new(cfg());
let t = Instant::now();
for _ in 0..3 {
assert!(cb.check_at("bad", t).is_ok());
cb.record_at("bad", false, t);
}
assert!(cb.is_open("bad"));
assert!(cb.check_at("bad", t).is_err());
}
#[test]
fn occasional_failures_below_threshold_stay_closed() {
let cb = CircuitBreaker::new(cfg());
let t = Instant::now();
for i in 0..15 {
cb.check_at("p", t).ok();
let success = i % 4 != 0; cb.record_at("p", success, t);
}
assert!(!cb.is_open("p"));
}
#[test]
fn halfopen_probe_success_closes() {
let cb = CircuitBreaker::new(cfg());
let t0 = Instant::now();
for _ in 0..3 {
cb.check_at("p", t0).ok();
cb.record_at("p", false, t0);
}
assert!(cb.check_at("p", t0).is_err()); let t1 = t0 + Duration::from_secs(11); assert!(cb.check_at("p", t1).is_ok()); cb.record_at("p", true, t1); assert!(!cb.is_open("p"));
assert!(cb.check_at("p", t1).is_ok());
}
#[test]
fn halfopen_probe_failure_reopens() {
let cb = CircuitBreaker::new(cfg());
let t0 = Instant::now();
for _ in 0..3 {
cb.check_at("p", t0).ok();
cb.record_at("p", false, t0);
}
let t1 = t0 + Duration::from_secs(11);
assert!(cb.check_at("p", t1).is_ok()); cb.record_at("p", false, t1); assert!(cb.is_open("p"));
assert!(cb.check_at("p", t1).is_err()); }
#[test]
fn reset_clears_open_state() {
let cb = CircuitBreaker::new(cfg());
let t = Instant::now();
for _ in 0..3 {
cb.check_at("p", t).ok();
cb.record_at("p", false, t);
}
assert!(cb.is_open("p"));
cb.reset("p");
assert!(!cb.is_open("p"));
assert!(cb.check_at("p", t).is_ok());
}
}