use rustlavel_core::{Error, Result};
use rustlavel_http::Status;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
const BUCKETS: u64 = 10;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum State {
Closed,
Open,
HalfOpen,
}
type FailureRule = Arc<dyn Fn(Status) -> bool + Send + Sync>;
#[derive(Clone)]
pub struct CircuitBreaker {
settings: Settings,
hosts: Arc<Mutex<HashMap<String, Circuit>>>,
}
#[derive(Clone)]
struct Settings {
failure_rate: f64,
minimum_calls: u32,
window: Duration,
reset_after: Duration,
probes: u32,
is_failure: FailureRule,
}
impl Default for CircuitBreaker {
fn default() -> Self {
Self::new()
}
}
impl CircuitBreaker {
pub fn new() -> Self {
CircuitBreaker {
settings: Settings {
failure_rate: 0.5,
minimum_calls: 20,
window: Duration::from_secs(60),
reset_after: Duration::from_secs(30),
probes: 3,
is_failure: Arc::new(|status: Status| status.code() >= 500),
},
hosts: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn failure_rate(mut self, rate: f64) -> Self {
self.settings.failure_rate = rate.clamp(0.0, 1.0);
self
}
pub fn minimum_calls(mut self, calls: u32) -> Self {
self.settings.minimum_calls = calls.max(1);
self
}
pub fn window(mut self, window: Duration) -> Self {
self.settings.window = window.max(Duration::from_millis(BUCKETS));
self
}
pub fn reset_after(mut self, pause: Duration) -> Self {
self.settings.reset_after = pause;
self
}
pub fn probes(mut self, probes: u32) -> Self {
self.settings.probes = probes.max(1);
self
}
pub fn count_failure_when(mut self, rule: impl Fn(Status) -> bool + Send + Sync + 'static) -> Self {
self.settings.is_failure = Arc::new(rule);
self
}
pub fn state(&self, host: &str) -> State {
let mut hosts = self.hosts.lock().unwrap_or_else(|e| e.into_inner());
match hosts.get_mut(host) {
Some(circuit) => circuit.state(&self.settings, Instant::now()),
None => State::Closed,
}
}
pub fn acquire(&self, host: &str) -> Result<Permit> {
let now = Instant::now();
let mut hosts = self.hosts.lock().unwrap_or_else(|e| e.into_inner());
let circuit = hosts.entry(host.to_string()).or_default();
match circuit.state(&self.settings, now) {
State::Closed => Ok(Permit::new(self.clone(), host.to_string(), false)),
State::Open => {
let for_another = self
.settings
.reset_after
.saturating_sub(now.saturating_duration_since(circuit.opened_at.unwrap_or(now)));
Err(Error::Unavailable(format!(
"{host} is not being called: too many of the last requests to it failed, \
so the circuit is open for another {} second(s). Nothing was sent.",
for_another.as_secs().max(1)
)))
}
State::HalfOpen => {
let took_one = circuit
.probes_left
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| {
(left > 0).then(|| left - 1)
})
.is_ok();
if took_one {
Ok(Permit::new(self.clone(), host.to_string(), true))
} else {
Err(Error::Unavailable(format!(
"{host} is being probed after a failure and is not taking other calls \
yet. Nothing was sent."
)))
}
}
}
}
pub fn counts_as_failure(&self, status: Status) -> bool {
(self.settings.is_failure)(status)
}
fn record(&self, host: &str, was_probe: bool, failed: bool) {
let now = Instant::now();
let mut hosts = self.hosts.lock().unwrap_or_else(|e| e.into_inner());
let Some(circuit) = hosts.get_mut(host) else { return };
match circuit.state(&self.settings, now) {
State::HalfOpen if was_probe => {
if failed {
circuit.open(now, &self.settings);
rustlavel_core::debug!("circuit for {host} opened again: a probe failed");
} else {
circuit.probe_successes += 1;
if circuit.probe_successes >= self.settings.probes {
circuit.close(now);
rustlavel_core::info!("circuit for {host} closed: the probes succeeded");
}
}
}
_ if was_probe => {}
_ => {
circuit.count(now, &self.settings, failed);
if circuit.should_open(&self.settings) {
circuit.open(now, &self.settings);
rustlavel_core::warn!(
"circuit for {host} opened: {:.0}% of the last {} calls failed",
circuit.failure_rate() * 100.0,
circuit.total()
);
}
}
}
}
pub fn reset(&self) {
self.hosts.lock().unwrap_or_else(|e| e.into_inner()).clear();
}
}
impl std::fmt::Debug for CircuitBreaker {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CircuitBreaker")
.field("failure_rate", &self.settings.failure_rate)
.field("minimum_calls", &self.settings.minimum_calls)
.field("window", &self.settings.window)
.field("reset_after", &self.settings.reset_after)
.field("probes", &self.settings.probes)
.finish()
}
}
pub struct Permit {
breaker: CircuitBreaker,
host: String,
probe: bool,
reported: bool,
}
impl std::fmt::Debug for Permit {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Permit").field("host", &self.host).field("probe", &self.probe).finish()
}
}
impl Permit {
fn new(breaker: CircuitBreaker, host: String, probe: bool) -> Self {
Permit { breaker, host, probe, reported: false }
}
pub fn success(mut self) {
self.reported = true;
self.breaker.record(&self.host, self.probe, false);
}
pub fn failure(mut self) {
self.reported = true;
self.breaker.record(&self.host, self.probe, true);
}
pub fn record_status(self, status: Status) {
if self.breaker.counts_as_failure(status) {
self.failure()
} else {
self.success()
}
}
}
impl Drop for Permit {
fn drop(&mut self) {
if self.reported || !self.probe {
return;
}
let mut hosts = self.breaker.hosts.lock().unwrap_or_else(|e| e.into_inner());
if let Some(circuit) = hosts.get_mut(&self.host) {
circuit.probes_left.fetch_add(1, Ordering::SeqCst);
}
}
}
#[derive(Debug, Default)]
struct Circuit {
buckets: VecDeque<(u64, u32, u32)>,
origin: Option<Instant>,
opened_at: Option<Instant>,
half_open: bool,
probes_left: AtomicU32,
probe_successes: u32,
}
impl Circuit {
fn state(&mut self, settings: &Settings, now: Instant) -> State {
let Some(opened_at) = self.opened_at else { return State::Closed };
if self.half_open {
return State::HalfOpen;
}
if now.saturating_duration_since(opened_at) >= settings.reset_after {
self.half_open = true;
self.probes_left = AtomicU32::new(settings.probes);
self.probe_successes = 0;
return State::HalfOpen;
}
State::Open
}
fn open(&mut self, now: Instant, settings: &Settings) {
self.opened_at = Some(now);
self.half_open = false;
self.probe_successes = 0;
self.probes_left = AtomicU32::new(settings.probes);
self.buckets.clear();
}
fn close(&mut self, _now: Instant) {
self.opened_at = None;
self.half_open = false;
self.probe_successes = 0;
self.buckets.clear();
}
fn bucket_of(&mut self, now: Instant, settings: &Settings) -> u64 {
let origin = *self.origin.get_or_insert(now);
let width = settings.window / BUCKETS as u32;
(now.saturating_duration_since(origin).as_nanos() / width.as_nanos().max(1)) as u64
}
fn count(&mut self, now: Instant, settings: &Settings, failed: bool) {
let bucket = self.bucket_of(now, settings);
while let Some(&(number, _, _)) = self.buckets.front() {
if number + BUCKETS <= bucket {
self.buckets.pop_front();
} else {
break;
}
}
match self.buckets.back_mut() {
Some((number, successes, failures)) if *number == bucket => {
if failed {
*failures += 1
} else {
*successes += 1
}
}
_ => self.buckets.push_back((bucket, u32::from(!failed), u32::from(failed))),
}
}
fn total(&self) -> u32 {
self.buckets.iter().map(|(_, s, f)| s + f).sum()
}
fn failures(&self) -> u32 {
self.buckets.iter().map(|(_, _, f)| f).sum()
}
fn failure_rate(&self) -> f64 {
match self.total() {
0 => 0.0,
total => f64::from(self.failures()) / f64::from(total),
}
}
fn should_open(&self, settings: &Settings) -> bool {
self.opened_at.is_none()
&& self.total() >= settings.minimum_calls
&& self.failure_rate() >= settings.failure_rate
}
}
#[cfg(test)]
mod tests {
use super::*;
fn breaker() -> CircuitBreaker {
CircuitBreaker::new()
.minimum_calls(4)
.failure_rate(0.5)
.reset_after(Duration::from_millis(60))
.probes(2)
}
fn fail(breaker: &CircuitBreaker, host: &str, times: usize) {
for _ in 0..times {
breaker.acquire(host).expect("closed").failure();
}
}
fn succeed(breaker: &CircuitBreaker, host: &str, times: usize) {
for _ in 0..times {
breaker.acquire(host).expect("closed").success();
}
}
#[test]
fn a_new_breaker_is_closed_and_lets_everything_through() {
let breaker = breaker();
assert_eq!(breaker.state("api.example"), State::Closed);
succeed(&breaker, "api.example", 50);
assert_eq!(breaker.state("api.example"), State::Closed);
}
#[test]
fn it_does_not_trip_below_the_minimum_however_bad_the_rate() {
let breaker = breaker();
fail(&breaker, "api.example", 3);
assert_eq!(breaker.state("api.example"), State::Closed);
}
#[test]
fn it_trips_once_the_rate_and_the_volume_are_both_reached() {
let breaker = breaker();
succeed(&breaker, "api.example", 2);
fail(&breaker, "api.example", 2);
assert_eq!(breaker.state("api.example"), State::Open, "4 calls, half of them failed");
}
#[test]
fn a_low_failure_rate_over_many_calls_does_not_trip_it() {
let breaker = CircuitBreaker::new().minimum_calls(10).failure_rate(0.5);
succeed(&breaker, "api.example", 95);
fail(&breaker, "api.example", 5);
assert_eq!(breaker.state("api.example"), State::Closed);
}
#[test]
fn an_open_breaker_refuses_without_sending_anything() {
let breaker = breaker();
fail(&breaker, "api.example", 4);
let error = breaker.acquire("api.example").expect_err("refused");
assert!(matches!(error, Error::Unavailable(_)), "{error:?}");
let message = error.to_string();
assert!(message.contains("api.example"), "{message}");
assert!(message.contains("Nothing was sent"), "{message}");
}
#[test]
fn breakers_are_kept_per_host() {
let breaker = breaker();
fail(&breaker, "payments.example", 4);
assert_eq!(breaker.state("payments.example"), State::Open);
assert_eq!(breaker.state("search.example"), State::Closed, "an unrelated host is unaffected");
breaker.acquire("search.example").expect("still closed").success();
}
#[tokio::test]
async fn after_the_pause_it_probes_and_closes_on_success() {
let breaker = breaker();
fail(&breaker, "api.example", 4);
assert_eq!(breaker.state("api.example"), State::Open);
tokio::time::sleep(Duration::from_millis(80)).await;
assert_eq!(breaker.state("api.example"), State::HalfOpen);
breaker.acquire("api.example").expect("a probe").success();
assert_eq!(breaker.state("api.example"), State::HalfOpen, "one probe of two");
breaker.acquire("api.example").expect("a probe").success();
assert_eq!(breaker.state("api.example"), State::Closed, "both probes succeeded");
}
#[tokio::test]
async fn one_failing_probe_opens_it_again_for_another_pause() {
let breaker = breaker();
fail(&breaker, "api.example", 4);
tokio::time::sleep(Duration::from_millis(80)).await;
breaker.acquire("api.example").expect("a probe").failure();
assert_eq!(breaker.state("api.example"), State::Open, "still not healthy");
breaker.acquire("api.example").expect_err("refused again");
tokio::time::sleep(Duration::from_millis(80)).await;
assert_eq!(breaker.state("api.example"), State::HalfOpen, "and it probes again after");
}
#[tokio::test]
async fn half_open_lets_through_only_as_many_probes_as_configured() {
let breaker = breaker();
fail(&breaker, "api.example", 4);
tokio::time::sleep(Duration::from_millis(80)).await;
let first = breaker.acquire("api.example").expect("probe one");
let second = breaker.acquire("api.example").expect("probe two");
breaker.acquire("api.example").expect_err("the third is refused, not queued");
first.success();
second.success();
assert_eq!(breaker.state("api.example"), State::Closed);
}
#[tokio::test]
async fn a_probe_that_is_dropped_gives_its_permit_back() {
let breaker = breaker();
fail(&breaker, "api.example", 4);
tokio::time::sleep(Duration::from_millis(80)).await;
drop(breaker.acquire("api.example").expect("probe one"));
drop(breaker.acquire("api.example").expect("probe two"));
drop(breaker.acquire("api.example").expect("permits came back"));
assert_eq!(breaker.state("api.example"), State::HalfOpen, "no outcome was recorded");
breaker.acquire("api.example").expect("a probe").success();
breaker.acquire("api.example").expect("a probe").success();
assert_eq!(breaker.state("api.example"), State::Closed);
}
#[tokio::test]
async fn closing_forgets_the_failures_that_opened_it() {
let breaker = breaker();
fail(&breaker, "api.example", 4);
tokio::time::sleep(Duration::from_millis(80)).await;
succeed(&breaker, "api.example", 2);
assert_eq!(breaker.state("api.example"), State::Closed);
fail(&breaker, "api.example", 1);
assert_eq!(breaker.state("api.example"), State::Closed, "one failure is not four");
}
#[tokio::test]
async fn failures_age_out_of_the_window() {
let breaker = CircuitBreaker::new()
.minimum_calls(4)
.failure_rate(0.5)
.window(Duration::from_millis(100));
fail(&breaker, "api.example", 3);
assert_eq!(breaker.state("api.example"), State::Closed, "not yet at the minimum");
tokio::time::sleep(Duration::from_millis(160)).await;
fail(&breaker, "api.example", 3);
assert_eq!(breaker.state("api.example"), State::Closed, "the old failures aged out");
}
#[test]
fn a_4xx_is_not_the_upstreams_fault_and_a_5xx_is() {
let breaker = breaker();
assert!(!breaker.counts_as_failure(Status::NOT_FOUND));
assert!(!breaker.counts_as_failure(Status::UNPROCESSABLE));
assert!(!breaker.counts_as_failure(Status::TOO_MANY_REQUESTS));
assert!(breaker.counts_as_failure(Status::INTERNAL_ERROR));
assert!(breaker.counts_as_failure(Status::SERVICE_UNAVAILABLE));
for _ in 0..400 {
breaker.acquire("api.example").expect("closed").record_status(Status::NOT_FOUND);
}
assert_eq!(breaker.state("api.example"), State::Closed);
}
#[test]
fn the_failure_rule_can_be_replaced() {
let breaker = breaker().count_failure_when(|status| status.code() == 429);
assert!(breaker.counts_as_failure(Status::TOO_MANY_REQUESTS));
assert!(!breaker.counts_as_failure(Status::INTERNAL_ERROR));
for _ in 0..4 {
breaker.acquire("api.example").expect("closed").record_status(Status::TOO_MANY_REQUESTS);
}
assert_eq!(breaker.state("api.example"), State::Open);
}
#[test]
fn reset_forgets_everything() {
let breaker = breaker();
fail(&breaker, "api.example", 4);
assert_eq!(breaker.state("api.example"), State::Open);
breaker.reset();
assert_eq!(breaker.state("api.example"), State::Closed);
}
#[test]
fn an_unavailable_error_is_a_503_and_says_which_dependency() {
let breaker = breaker();
fail(&breaker, "payments.example", 4);
let error = breaker.acquire("payments.example").expect_err("open");
assert_eq!(error.status(), 503);
assert_eq!(error.title(), "Dependency Unavailable");
}
#[tokio::test]
async fn many_tasks_racing_on_one_host_agree_on_the_outcome() {
let breaker = CircuitBreaker::new().minimum_calls(100).failure_rate(0.5);
let mut tasks = Vec::new();
for i in 0..200 {
let breaker = breaker.clone();
tasks.push(tokio::spawn(async move {
if let Ok(permit) = breaker.acquire("api.example") {
if i % 2 == 0 { permit.failure() } else { permit.success() }
}
}));
}
for task in tasks {
task.await.expect("no task panicked");
}
assert_eq!(breaker.state("api.example"), State::Open);
}
}