1use rustlavel_core::{Error, Result};
49use rustlavel_http::Status;
50use std::collections::HashMap;
51use std::collections::VecDeque;
52use std::sync::atomic::{AtomicU32, Ordering};
53use std::sync::{Arc, Mutex};
54use std::time::{Duration, Instant};
55
56const BUCKETS: u64 = 10;
62
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65pub enum State {
66 Closed,
68 Open,
70 HalfOpen,
72}
73
74type FailureRule = Arc<dyn Fn(Status) -> bool + Send + Sync>;
76
77#[derive(Clone)]
78pub struct CircuitBreaker {
79 settings: Settings,
80 hosts: Arc<Mutex<HashMap<String, Circuit>>>,
81}
82
83#[derive(Clone)]
84struct Settings {
85 failure_rate: f64,
86 minimum_calls: u32,
87 window: Duration,
88 reset_after: Duration,
89 probes: u32,
90 is_failure: FailureRule,
91}
92
93impl Default for CircuitBreaker {
94 fn default() -> Self {
95 Self::new()
96 }
97}
98
99impl CircuitBreaker {
100 pub fn new() -> Self {
108 CircuitBreaker {
109 settings: Settings {
110 failure_rate: 0.5,
111 minimum_calls: 20,
112 window: Duration::from_secs(60),
113 reset_after: Duration::from_secs(30),
114 probes: 3,
115 is_failure: Arc::new(|status: Status| status.code() >= 500),
116 },
117 hosts: Arc::new(Mutex::new(HashMap::new())),
118 }
119 }
120
121 pub fn failure_rate(mut self, rate: f64) -> Self {
123 self.settings.failure_rate = rate.clamp(0.0, 1.0);
124 self
125 }
126
127 pub fn minimum_calls(mut self, calls: u32) -> Self {
129 self.settings.minimum_calls = calls.max(1);
130 self
131 }
132
133 pub fn window(mut self, window: Duration) -> Self {
135 self.settings.window = window.max(Duration::from_millis(BUCKETS));
136 self
137 }
138
139 pub fn reset_after(mut self, pause: Duration) -> Self {
141 self.settings.reset_after = pause;
142 self
143 }
144
145 pub fn probes(mut self, probes: u32) -> Self {
147 self.settings.probes = probes.max(1);
148 self
149 }
150
151 pub fn count_failure_when(mut self, rule: impl Fn(Status) -> bool + Send + Sync + 'static) -> Self {
157 self.settings.is_failure = Arc::new(rule);
158 self
159 }
160
161 pub fn state(&self, host: &str) -> State {
163 let mut hosts = self.hosts.lock().unwrap_or_else(|e| e.into_inner());
164 match hosts.get_mut(host) {
165 Some(circuit) => circuit.state(&self.settings, Instant::now()),
166 None => State::Closed,
167 }
168 }
169
170 pub fn acquire(&self, host: &str) -> Result<Permit> {
176 let now = Instant::now();
177 let mut hosts = self.hosts.lock().unwrap_or_else(|e| e.into_inner());
178 let circuit = hosts.entry(host.to_string()).or_default();
179
180 match circuit.state(&self.settings, now) {
181 State::Closed => Ok(Permit::new(self.clone(), host.to_string(), false)),
182 State::Open => {
183 let for_another = self
184 .settings
185 .reset_after
186 .saturating_sub(now.saturating_duration_since(circuit.opened_at.unwrap_or(now)));
187 Err(Error::Unavailable(format!(
188 "{host} is not being called: too many of the last requests to it failed, \
189 so the circuit is open for another {} second(s). Nothing was sent.",
190 for_another.as_secs().max(1)
191 )))
192 }
193 State::HalfOpen => {
194 let took_one = circuit
200 .probes_left
201 .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| {
202 (left > 0).then(|| left - 1)
203 })
204 .is_ok();
205 if took_one {
206 Ok(Permit::new(self.clone(), host.to_string(), true))
207 } else {
208 Err(Error::Unavailable(format!(
209 "{host} is being probed after a failure and is not taking other calls \
210 yet. Nothing was sent."
211 )))
212 }
213 }
214 }
215 }
216
217 pub fn counts_as_failure(&self, status: Status) -> bool {
219 (self.settings.is_failure)(status)
220 }
221
222 fn record(&self, host: &str, was_probe: bool, failed: bool) {
223 let now = Instant::now();
224 let mut hosts = self.hosts.lock().unwrap_or_else(|e| e.into_inner());
225 let Some(circuit) = hosts.get_mut(host) else { return };
226
227 match circuit.state(&self.settings, now) {
230 State::HalfOpen if was_probe => {
231 if failed {
232 circuit.open(now, &self.settings);
233 rustlavel_core::debug!("circuit for {host} opened again: a probe failed");
234 } else {
235 circuit.probe_successes += 1;
236 if circuit.probe_successes >= self.settings.probes {
237 circuit.close(now);
238 rustlavel_core::info!("circuit for {host} closed: the probes succeeded");
239 }
240 }
241 }
242 _ if was_probe => {}
246 _ => {
247 circuit.count(now, &self.settings, failed);
248 if circuit.should_open(&self.settings) {
249 circuit.open(now, &self.settings);
250 rustlavel_core::warn!(
251 "circuit for {host} opened: {:.0}% of the last {} calls failed",
252 circuit.failure_rate() * 100.0,
253 circuit.total()
254 );
255 }
256 }
257 }
258 }
259
260 pub fn reset(&self) {
262 self.hosts.lock().unwrap_or_else(|e| e.into_inner()).clear();
263 }
264}
265
266impl std::fmt::Debug for CircuitBreaker {
267 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
268 f.debug_struct("CircuitBreaker")
269 .field("failure_rate", &self.settings.failure_rate)
270 .field("minimum_calls", &self.settings.minimum_calls)
271 .field("window", &self.settings.window)
272 .field("reset_after", &self.settings.reset_after)
273 .field("probes", &self.settings.probes)
274 .finish()
275 }
276}
277
278pub struct Permit {
286 breaker: CircuitBreaker,
287 host: String,
288 probe: bool,
289 reported: bool,
290}
291
292impl std::fmt::Debug for Permit {
293 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
294 f.debug_struct("Permit").field("host", &self.host).field("probe", &self.probe).finish()
295 }
296}
297
298impl Permit {
299 fn new(breaker: CircuitBreaker, host: String, probe: bool) -> Self {
300 Permit { breaker, host, probe, reported: false }
301 }
302
303 pub fn success(mut self) {
304 self.reported = true;
305 self.breaker.record(&self.host, self.probe, false);
306 }
307
308 pub fn failure(mut self) {
309 self.reported = true;
310 self.breaker.record(&self.host, self.probe, true);
311 }
312
313 pub fn record_status(self, status: Status) {
315 if self.breaker.counts_as_failure(status) {
316 self.failure()
317 } else {
318 self.success()
319 }
320 }
321}
322
323impl Drop for Permit {
324 fn drop(&mut self) {
325 if self.reported || !self.probe {
326 return;
327 }
328 let mut hosts = self.breaker.hosts.lock().unwrap_or_else(|e| e.into_inner());
331 if let Some(circuit) = hosts.get_mut(&self.host) {
332 circuit.probes_left.fetch_add(1, Ordering::SeqCst);
333 }
334 }
335}
336
337#[derive(Debug, Default)]
339struct Circuit {
340 buckets: VecDeque<(u64, u32, u32)>,
342 origin: Option<Instant>,
343 opened_at: Option<Instant>,
344 half_open: bool,
345 probes_left: AtomicU32,
346 probe_successes: u32,
347}
348
349impl Circuit {
350 fn state(&mut self, settings: &Settings, now: Instant) -> State {
352 let Some(opened_at) = self.opened_at else { return State::Closed };
353
354 if self.half_open {
355 return State::HalfOpen;
356 }
357 if now.saturating_duration_since(opened_at) >= settings.reset_after {
358 self.half_open = true;
359 self.probes_left = AtomicU32::new(settings.probes);
360 self.probe_successes = 0;
361 return State::HalfOpen;
362 }
363 State::Open
364 }
365
366 fn open(&mut self, now: Instant, settings: &Settings) {
367 self.opened_at = Some(now);
368 self.half_open = false;
369 self.probe_successes = 0;
370 self.probes_left = AtomicU32::new(settings.probes);
371 self.buckets.clear();
374 }
375
376 fn close(&mut self, _now: Instant) {
377 self.opened_at = None;
378 self.half_open = false;
379 self.probe_successes = 0;
380 self.buckets.clear();
381 }
382
383 fn bucket_of(&mut self, now: Instant, settings: &Settings) -> u64 {
385 let origin = *self.origin.get_or_insert(now);
386 let width = settings.window / BUCKETS as u32;
387 (now.saturating_duration_since(origin).as_nanos() / width.as_nanos().max(1)) as u64
388 }
389
390 fn count(&mut self, now: Instant, settings: &Settings, failed: bool) {
391 let bucket = self.bucket_of(now, settings);
392
393 while let Some(&(number, _, _)) = self.buckets.front() {
395 if number + BUCKETS <= bucket {
396 self.buckets.pop_front();
397 } else {
398 break;
399 }
400 }
401
402 match self.buckets.back_mut() {
403 Some((number, successes, failures)) if *number == bucket => {
404 if failed {
405 *failures += 1
406 } else {
407 *successes += 1
408 }
409 }
410 _ => self.buckets.push_back((bucket, u32::from(!failed), u32::from(failed))),
411 }
412 }
413
414 fn total(&self) -> u32 {
415 self.buckets.iter().map(|(_, s, f)| s + f).sum()
416 }
417
418 fn failures(&self) -> u32 {
419 self.buckets.iter().map(|(_, _, f)| f).sum()
420 }
421
422 fn failure_rate(&self) -> f64 {
423 match self.total() {
424 0 => 0.0,
425 total => f64::from(self.failures()) / f64::from(total),
426 }
427 }
428
429 fn should_open(&self, settings: &Settings) -> bool {
430 self.opened_at.is_none()
431 && self.total() >= settings.minimum_calls
432 && self.failure_rate() >= settings.failure_rate
433 }
434}
435
436#[cfg(test)]
437mod tests {
438 use super::*;
439
440 fn breaker() -> CircuitBreaker {
441 CircuitBreaker::new()
442 .minimum_calls(4)
443 .failure_rate(0.5)
444 .reset_after(Duration::from_millis(60))
445 .probes(2)
446 }
447
448 fn fail(breaker: &CircuitBreaker, host: &str, times: usize) {
449 for _ in 0..times {
450 breaker.acquire(host).expect("closed").failure();
451 }
452 }
453
454 fn succeed(breaker: &CircuitBreaker, host: &str, times: usize) {
455 for _ in 0..times {
456 breaker.acquire(host).expect("closed").success();
457 }
458 }
459
460 #[test]
461 fn a_new_breaker_is_closed_and_lets_everything_through() {
462 let breaker = breaker();
463 assert_eq!(breaker.state("api.example"), State::Closed);
464 succeed(&breaker, "api.example", 50);
465 assert_eq!(breaker.state("api.example"), State::Closed);
466 }
467
468 #[test]
469 fn it_does_not_trip_below_the_minimum_however_bad_the_rate() {
470 let breaker = breaker();
473 fail(&breaker, "api.example", 3);
474 assert_eq!(breaker.state("api.example"), State::Closed);
475 }
476
477 #[test]
478 fn it_trips_once_the_rate_and_the_volume_are_both_reached() {
479 let breaker = breaker();
480 succeed(&breaker, "api.example", 2);
481 fail(&breaker, "api.example", 2);
482 assert_eq!(breaker.state("api.example"), State::Open, "4 calls, half of them failed");
483 }
484
485 #[test]
486 fn a_low_failure_rate_over_many_calls_does_not_trip_it() {
487 let breaker = CircuitBreaker::new().minimum_calls(10).failure_rate(0.5);
490 succeed(&breaker, "api.example", 95);
491 fail(&breaker, "api.example", 5);
492 assert_eq!(breaker.state("api.example"), State::Closed);
493 }
494
495 #[test]
496 fn an_open_breaker_refuses_without_sending_anything() {
497 let breaker = breaker();
498 fail(&breaker, "api.example", 4);
499
500 let error = breaker.acquire("api.example").expect_err("refused");
501 assert!(matches!(error, Error::Unavailable(_)), "{error:?}");
502 let message = error.to_string();
503 assert!(message.contains("api.example"), "{message}");
504 assert!(message.contains("Nothing was sent"), "{message}");
505 }
506
507 #[test]
508 fn breakers_are_kept_per_host() {
509 let breaker = breaker();
510 fail(&breaker, "payments.example", 4);
511
512 assert_eq!(breaker.state("payments.example"), State::Open);
513 assert_eq!(breaker.state("search.example"), State::Closed, "an unrelated host is unaffected");
514 breaker.acquire("search.example").expect("still closed").success();
515 }
516
517 #[tokio::test]
518 async fn after_the_pause_it_probes_and_closes_on_success() {
519 let breaker = breaker();
520 fail(&breaker, "api.example", 4);
521 assert_eq!(breaker.state("api.example"), State::Open);
522
523 tokio::time::sleep(Duration::from_millis(80)).await;
524 assert_eq!(breaker.state("api.example"), State::HalfOpen);
525
526 breaker.acquire("api.example").expect("a probe").success();
527 assert_eq!(breaker.state("api.example"), State::HalfOpen, "one probe of two");
528 breaker.acquire("api.example").expect("a probe").success();
529 assert_eq!(breaker.state("api.example"), State::Closed, "both probes succeeded");
530 }
531
532 #[tokio::test]
533 async fn one_failing_probe_opens_it_again_for_another_pause() {
534 let breaker = breaker();
535 fail(&breaker, "api.example", 4);
536 tokio::time::sleep(Duration::from_millis(80)).await;
537
538 breaker.acquire("api.example").expect("a probe").failure();
539 assert_eq!(breaker.state("api.example"), State::Open, "still not healthy");
540 breaker.acquire("api.example").expect_err("refused again");
541
542 tokio::time::sleep(Duration::from_millis(80)).await;
543 assert_eq!(breaker.state("api.example"), State::HalfOpen, "and it probes again after");
544 }
545
546 #[tokio::test]
547 async fn half_open_lets_through_only_as_many_probes_as_configured() {
548 let breaker = breaker();
549 fail(&breaker, "api.example", 4);
550 tokio::time::sleep(Duration::from_millis(80)).await;
551
552 let first = breaker.acquire("api.example").expect("probe one");
554 let second = breaker.acquire("api.example").expect("probe two");
555 breaker.acquire("api.example").expect_err("the third is refused, not queued");
556
557 first.success();
558 second.success();
559 assert_eq!(breaker.state("api.example"), State::Closed);
560 }
561
562 #[tokio::test]
563 async fn a_probe_that_is_dropped_gives_its_permit_back() {
564 let breaker = breaker();
568 fail(&breaker, "api.example", 4);
569 tokio::time::sleep(Duration::from_millis(80)).await;
570
571 drop(breaker.acquire("api.example").expect("probe one"));
572 drop(breaker.acquire("api.example").expect("probe two"));
573 drop(breaker.acquire("api.example").expect("permits came back"));
574
575 assert_eq!(breaker.state("api.example"), State::HalfOpen, "no outcome was recorded");
576 breaker.acquire("api.example").expect("a probe").success();
577 breaker.acquire("api.example").expect("a probe").success();
578 assert_eq!(breaker.state("api.example"), State::Closed);
579 }
580
581 #[tokio::test]
582 async fn closing_forgets_the_failures_that_opened_it() {
583 let breaker = breaker();
586 fail(&breaker, "api.example", 4);
587 tokio::time::sleep(Duration::from_millis(80)).await;
588 succeed(&breaker, "api.example", 2);
589 assert_eq!(breaker.state("api.example"), State::Closed);
590
591 fail(&breaker, "api.example", 1);
592 assert_eq!(breaker.state("api.example"), State::Closed, "one failure is not four");
593 }
594
595 #[tokio::test]
596 async fn failures_age_out_of_the_window() {
597 let breaker = CircuitBreaker::new()
598 .minimum_calls(4)
599 .failure_rate(0.5)
600 .window(Duration::from_millis(100));
601
602 fail(&breaker, "api.example", 3);
603 assert_eq!(breaker.state("api.example"), State::Closed, "not yet at the minimum");
604
605 tokio::time::sleep(Duration::from_millis(160)).await;
607 fail(&breaker, "api.example", 3);
608 assert_eq!(breaker.state("api.example"), State::Closed, "the old failures aged out");
609 }
610
611 #[test]
612 fn a_4xx_is_not_the_upstreams_fault_and_a_5xx_is() {
613 let breaker = breaker();
614 assert!(!breaker.counts_as_failure(Status::NOT_FOUND));
615 assert!(!breaker.counts_as_failure(Status::UNPROCESSABLE));
616 assert!(!breaker.counts_as_failure(Status::TOO_MANY_REQUESTS));
617 assert!(breaker.counts_as_failure(Status::INTERNAL_ERROR));
618 assert!(breaker.counts_as_failure(Status::SERVICE_UNAVAILABLE));
619
620 for _ in 0..400 {
622 breaker.acquire("api.example").expect("closed").record_status(Status::NOT_FOUND);
623 }
624 assert_eq!(breaker.state("api.example"), State::Closed);
625 }
626
627 #[test]
628 fn the_failure_rule_can_be_replaced() {
629 let breaker = breaker().count_failure_when(|status| status.code() == 429);
630 assert!(breaker.counts_as_failure(Status::TOO_MANY_REQUESTS));
631 assert!(!breaker.counts_as_failure(Status::INTERNAL_ERROR));
632
633 for _ in 0..4 {
634 breaker.acquire("api.example").expect("closed").record_status(Status::TOO_MANY_REQUESTS);
635 }
636 assert_eq!(breaker.state("api.example"), State::Open);
637 }
638
639 #[test]
640 fn reset_forgets_everything() {
641 let breaker = breaker();
642 fail(&breaker, "api.example", 4);
643 assert_eq!(breaker.state("api.example"), State::Open);
644 breaker.reset();
645 assert_eq!(breaker.state("api.example"), State::Closed);
646 }
647
648 #[test]
649 fn an_unavailable_error_is_a_503_and_says_which_dependency() {
650 let breaker = breaker();
651 fail(&breaker, "payments.example", 4);
652 let error = breaker.acquire("payments.example").expect_err("open");
653 assert_eq!(error.status(), 503);
654 assert_eq!(error.title(), "Dependency Unavailable");
655 }
656
657 #[tokio::test]
658 async fn many_tasks_racing_on_one_host_agree_on_the_outcome() {
659 let breaker = CircuitBreaker::new().minimum_calls(100).failure_rate(0.5);
660 let mut tasks = Vec::new();
661 for i in 0..200 {
662 let breaker = breaker.clone();
663 tasks.push(tokio::spawn(async move {
664 if let Ok(permit) = breaker.acquire("api.example") {
665 if i % 2 == 0 { permit.failure() } else { permit.success() }
666 }
667 }));
668 }
669 for task in tasks {
670 task.await.expect("no task panicked");
671 }
672 assert_eq!(breaker.state("api.example"), State::Open);
674 }
675}