1mod bulkhead;
41mod circuit_breaker;
42mod rate_limit;
43
44use std::collections::HashMap;
45use std::fmt;
46use std::sync::{Arc, Mutex};
47use std::time::{Duration, Instant};
48
49use async_trait::async_trait;
50
51use crate::client::ClientBuilder;
52use crate::error::{Error, ErrorCode};
53use crate::observability::{Hooks, OperationInfo, OperationState, RequestInfo, RequestResult};
54
55pub use bulkhead::{Bulkhead, BulkheadConfig, BulkheadPermit};
56pub use circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
57pub use rate_limit::{RateLimitConfig, RateLimiter};
58
59#[derive(Debug, Clone)]
62pub struct ResilienceConfig {
63 pub circuit_breaker: Option<CircuitBreakerConfig>,
65 pub bulkhead: Option<BulkheadConfig>,
67 pub rate_limit: Option<RateLimitConfig>,
69}
70
71impl Default for ResilienceConfig {
72 fn default() -> ResilienceConfig {
73 ResilienceConfig {
74 circuit_breaker: Some(CircuitBreakerConfig::default()),
75 bulkhead: Some(BulkheadConfig::default()),
76 rate_limit: Some(RateLimitConfig::default()),
77 }
78 }
79}
80
81impl ResilienceConfig {
82 pub fn none() -> ResilienceConfig {
85 ResilienceConfig {
86 circuit_breaker: None,
87 bulkhead: None,
88 rate_limit: None,
89 }
90 }
91}
92
93impl ClientBuilder {
94 pub fn resilience(mut self, config: ResilienceConfig) -> ClientBuilder {
97 self.hooks = Arc::new(ResilienceHooks::new(self.hooks.clone(), config));
98 self
99 }
100
101 pub fn circuit_breaker(self, config: CircuitBreakerConfig) -> ClientBuilder {
103 self.resilience(ResilienceConfig {
104 circuit_breaker: Some(config),
105 ..ResilienceConfig::none()
106 })
107 }
108
109 pub fn bulkhead(self, config: BulkheadConfig) -> ClientBuilder {
111 self.resilience(ResilienceConfig {
112 bulkhead: Some(config),
113 ..ResilienceConfig::none()
114 })
115 }
116
117 pub fn rate_limit(self, config: RateLimitConfig) -> ClientBuilder {
119 self.resilience(ResilienceConfig {
120 rate_limit: Some(config),
121 ..ResilienceConfig::none()
122 })
123 }
124}
125
126pub fn should_trip_circuit(error: &Error) -> bool {
135 match error.code() {
136 _ if error.refusal().is_some() || error.is_cancelled() => false,
137 ErrorCode::RateLimit => false,
138 ErrorCode::Network => true,
139 _ => error.http_status().is_some_and(|status| status >= 500),
140 }
141}
142
143#[derive(Clone)]
145pub struct Clock(Arc<dyn Fn() -> Instant + Send + Sync>);
146
147impl Clock {
148 pub fn new(now: impl Fn() -> Instant + Send + Sync + 'static) -> Clock {
150 Clock(Arc::new(now))
151 }
152
153 pub fn now(&self) -> Instant {
155 (self.0)()
156 }
157}
158
159impl Default for Clock {
160 fn default() -> Clock {
161 Clock::new(Instant::now)
162 }
163}
164
165impl fmt::Debug for Clock {
166 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
167 f.write_str("Clock")
168 }
169}
170
171pub struct ResilienceHooks {
176 inner: Arc<dyn Hooks>,
177 circuit_breakers: Option<Registry<CircuitBreaker>>,
178 bulkheads: Option<Registry<Bulkhead>>,
179 rate_limiter: Option<RateLimiter>,
180 pending: Mutex<HashMap<String, Vec<BulkheadPermit>>>,
187}
188
189impl ResilienceHooks {
190 pub fn new(inner: Arc<dyn Hooks>, config: ResilienceConfig) -> ResilienceHooks {
192 ResilienceHooks {
193 inner,
194 circuit_breakers: config
195 .circuit_breaker
196 .map(|config| Registry::new(move || CircuitBreaker::new(config.clone()))),
197 bulkheads: config
198 .bulkhead
199 .map(|config| Registry::new(move || Bulkhead::new(config.clone()))),
200 rate_limiter: config.rate_limit.map(RateLimiter::new),
201 pending: Mutex::default(),
202 }
203 }
204
205 async fn admit(&self, scope: &str) -> Result<Option<BulkheadPermit>, Error> {
212 if let Some(breakers) = &self.circuit_breakers
213 && !breakers.get(scope).allow()
214 {
215 return Err(Error::circuit_open());
216 }
217
218 let permit = match &self.bulkheads {
219 Some(bulkheads) => Some(bulkheads.get(scope).acquire().await?),
220 None => None,
221 };
222
223 if let Some(limiter) = &self.rate_limiter
224 && !limiter.allow()
225 {
226 return Err(Error::rate_limited());
227 }
228
229 Ok(permit)
230 }
231
232 fn record(&self, scope: &str, outcome: Result<(), &Error>) {
233 if let Some(breakers) = &self.circuit_breakers {
234 let breaker = breakers.get(scope);
235 match outcome {
236 Ok(()) => breaker.record_success(),
237 Err(error) if should_trip_circuit(error) => breaker.record_failure(),
238 Err(_) => {}
239 }
240 }
241 }
242}
243
244#[async_trait]
245impl Hooks for ResilienceHooks {
246 async fn on_operation_gate(&self, op: &OperationInfo) -> Result<(), Error> {
247 let scope = scope_of(op);
248 let permit = self.admit(&scope).await?;
249 self.inner.on_operation_gate(op).await?;
250 if let Some(permit) = permit {
251 self.pending
252 .lock()
253 .unwrap_or_else(std::sync::PoisonError::into_inner)
254 .entry(scope)
255 .or_default()
256 .push(permit);
257 }
258 Ok(())
259 }
260
261 fn on_operation_abandoned(&self, op: &OperationInfo) {
264 let permit = self
265 .pending
266 .lock()
267 .unwrap_or_else(std::sync::PoisonError::into_inner)
268 .get_mut(&scope_of(op))
269 .and_then(Vec::pop);
270 drop(permit);
271 self.inner.on_operation_abandoned(op);
272 }
273
274 fn on_operation_start(&self, op: &OperationInfo) -> OperationState {
275 let permit = self
276 .pending
277 .lock()
278 .unwrap_or_else(std::sync::PoisonError::into_inner)
279 .get_mut(&scope_of(op))
280 .and_then(Vec::pop);
281 Some(Box::new(Held {
282 permit,
283 inner: self.inner.on_operation_start(op),
284 }))
285 }
286
287 fn on_operation_end(
288 &self,
289 op: &OperationInfo,
290 state: OperationState,
291 outcome: Result<(), &Error>,
292 duration: Duration,
293 ) {
294 let held = match state.and_then(|state| state.downcast::<Held>().ok()) {
295 Some(held) => *held,
296 None => Held::default(),
297 };
298 drop(held.permit);
299 self.record(&scope_of(op), outcome);
300 self.inner
301 .on_operation_end(op, held.inner, outcome, duration);
302 }
303
304 fn on_request_start(&self, info: &RequestInfo) {
305 self.inner.on_request_start(info);
306 }
307
308 fn on_request_end(&self, info: &RequestInfo, result: &RequestResult<'_>) {
312 if let Some(limiter) = &self.rate_limiter {
313 let asked = result.retry_after.filter(|seconds| *seconds > 0);
314 match result.status.map(|status| status.as_u16()) {
315 Some(429) => limiter.set_retry_after_in(Duration::from_secs(asked.unwrap_or(60))),
316 Some(503) => {
317 if let Some(seconds) = asked {
318 limiter.set_retry_after_in(Duration::from_secs(seconds));
319 }
320 }
321 _ => {}
322 }
323 }
324 self.inner.on_request_end(info, result);
325 }
326
327 fn on_retry(&self, info: &RequestInfo, next_attempt: u32, cause: &Error) {
328 self.inner.on_retry(info, next_attempt, cause);
329 }
330}
331
332#[derive(Default)]
335struct Held {
336 permit: Option<BulkheadPermit>,
337 inner: OperationState,
338}
339
340fn scope_of(op: &OperationInfo) -> String {
341 format!("{}.{}", op.service, op.operation)
342}
343
344struct Registry<T> {
347 build: Box<dyn Fn() -> T + Send + Sync>,
348 entries: Mutex<HashMap<String, Arc<T>>>,
349}
350
351impl<T> Registry<T> {
352 fn new(build: impl Fn() -> T + Send + Sync + 'static) -> Registry<T> {
353 Registry {
354 build: Box::new(build),
355 entries: Mutex::default(),
356 }
357 }
358
359 fn get(&self, scope: &str) -> Arc<T> {
360 self.entries
361 .lock()
362 .unwrap_or_else(std::sync::PoisonError::into_inner)
363 .entry(scope.to_string())
364 .or_insert_with(|| Arc::new((self.build)()))
365 .clone()
366 }
367}
368
369#[cfg(test)]
371pub(crate) fn test_clock() -> (Clock, Arc<Mutex<Instant>>) {
372 let now = Arc::new(Mutex::new(Instant::now()));
373 let reading = now.clone();
374 (
375 Clock::new(move || {
376 *reading
377 .lock()
378 .unwrap_or_else(std::sync::PoisonError::into_inner)
379 }),
380 now,
381 )
382}
383
384#[cfg(test)]
385pub(crate) fn advance(clock: &Arc<Mutex<Instant>>, elapsed: Duration) {
386 let mut now = clock
387 .lock()
388 .unwrap_or_else(std::sync::PoisonError::into_inner);
389 *now += elapsed;
390}
391
392#[cfg(test)]
393#[allow(clippy::unwrap_used)]
394mod tests {
395 use std::borrow::Cow;
396
397 use super::*;
398 use crate::error::Refusal;
399
400 #[test]
401 fn the_default_config_installs_every_layer() {
402 let config = ResilienceConfig::default();
403
404 assert!(config.circuit_breaker.is_some());
405 assert!(config.bulkhead.is_some());
406 assert!(config.rate_limit.is_some());
407 }
408
409 #[test]
410 fn only_what_fizzy_failed_to_answer_trips_the_breaker() {
411 let cases = [
412 (Error::circuit_open(), false),
413 (Error::bulkhead_full(), false),
414 (Error::rate_limited(), false),
415 (Error::rate_limit(Some(3)), false),
416 (Error::cancelled(), false),
417 (
418 Error::network(std::io::Error::other("connection refused")),
419 true,
420 ),
421 (Error::api(500, "boom"), true),
422 (Error::api(503, "unavailable"), true),
423 (Error::api(400, "bad request"), false),
424 (Error::auth("authentication required"), false),
425 (Error::usage("bad argument"), false),
426 (Error::not_found("board", "1"), false),
427 ];
428
429 for (error, expected) in cases {
430 assert_eq!(expected, should_trip_circuit(&error), "{error}");
431 }
432 }
433
434 #[test]
435 fn a_registry_keeps_one_breaker_per_scope() {
436 let registry = Registry::new(|| CircuitBreaker::new(CircuitBreakerConfig::default()));
437
438 assert!(Arc::ptr_eq(
439 ®istry.get("scope1"),
440 ®istry.get("scope1")
441 ));
442 assert!(!Arc::ptr_eq(
443 ®istry.get("scope1"),
444 ®istry.get("scope2")
445 ));
446 }
447
448 #[test]
449 fn a_registry_keeps_one_bulkhead_per_scope() {
450 let registry = Registry::new(|| Bulkhead::new(BulkheadConfig::default()));
451
452 assert!(Arc::ptr_eq(
453 ®istry.get("scope1"),
454 ®istry.get("scope1")
455 ));
456 assert!(!Arc::ptr_eq(
457 ®istry.get("scope1"),
458 ®istry.get("scope2")
459 ));
460 }
461
462 #[tokio::test]
463 async fn a_breaker_only_config_refuses_once_the_scope_has_failed_enough() {
464 let hooks = hooks(ResilienceConfig {
465 circuit_breaker: Some(CircuitBreakerConfig {
466 failure_threshold: 2,
467 ..CircuitBreakerConfig::default()
468 }),
469 ..ResilienceConfig::none()
470 });
471
472 fail(&hooks, &operation("Boards", "ListBoards")).await;
473 assert!(
474 hooks
475 .on_operation_gate(&operation("Boards", "ListBoards"))
476 .await
477 .is_ok()
478 );
479 fail(&hooks, &operation("Boards", "ListBoards")).await;
480
481 let refused = hooks
482 .on_operation_gate(&operation("Boards", "ListBoards"))
483 .await
484 .unwrap_err();
485 assert_eq!(Some(Refusal::CircuitOpen), refused.refusal());
486 assert_eq!("circuit breaker is open", refused.message());
487 assert!(
488 hooks
489 .on_operation_gate(&operation("Boards", "GetBoard"))
490 .await
491 .is_ok()
492 );
493 }
494
495 #[tokio::test]
496 async fn a_bulkhead_with_no_wait_refuses_a_scope_that_is_already_busy() {
497 let hooks = hooks(ResilienceConfig {
498 bulkhead: Some(BulkheadConfig {
499 max_concurrent: 1,
500 max_wait: Duration::ZERO,
501 }),
502 ..ResilienceConfig::none()
503 });
504 let op = operation("Boards", "ListBoards");
505
506 hooks.on_operation_gate(&op).await.unwrap();
507 let state = hooks.on_operation_start(&op);
508
509 let refused = hooks.on_operation_gate(&op).await.unwrap_err();
510 assert_eq!(Some(Refusal::BulkheadFull), refused.refusal());
511 assert_eq!("bulkhead is full", refused.message());
512 assert!(
513 hooks
514 .on_operation_gate(&operation("Boards", "GetBoard"))
515 .await
516 .is_ok()
517 );
518
519 hooks.on_operation_end(&op, state, Ok(()), Duration::ZERO);
520 assert!(hooks.on_operation_gate(&op).await.is_ok());
521 }
522
523 #[tokio::test]
524 async fn a_limiter_only_config_refuses_once_the_budget_is_spent() {
525 let hooks = hooks(ResilienceConfig {
526 rate_limit: Some(RateLimitConfig {
527 requests_per_second: 0.0001,
528 burst_size: 1,
529 ..RateLimitConfig::default()
530 }),
531 ..ResilienceConfig::none()
532 });
533 let op = operation("Boards", "ListBoards");
534
535 hooks.on_operation_gate(&op).await.unwrap();
536
537 let refused = hooks.on_operation_gate(&op).await.unwrap_err();
538 assert_eq!(ErrorCode::RateLimit, refused.code());
539 assert_eq!("rate limit exceeded", refused.message());
540 assert_eq!(None, refused.http_status());
541 }
542
543 #[tokio::test]
546 async fn a_refused_call_gives_back_the_permit_it_took() {
547 let hooks = hooks(ResilienceConfig {
548 bulkhead: Some(BulkheadConfig {
549 max_concurrent: 1,
550 max_wait: Duration::ZERO,
551 }),
552 rate_limit: Some(RateLimitConfig {
553 requests_per_second: 0.0001,
554 burst_size: 1,
555 ..RateLimitConfig::default()
556 }),
557 ..ResilienceConfig::none()
558 });
559 let op = operation("Boards", "ListBoards");
560
561 hooks.on_operation_gate(&op).await.unwrap();
562 let state = hooks.on_operation_start(&op);
563 hooks.on_operation_end(&op, state, Ok(()), Duration::ZERO);
564
565 assert_eq!(
566 ErrorCode::RateLimit,
567 hooks.on_operation_gate(&op).await.unwrap_err().code()
568 );
569 assert_eq!(
570 1,
571 hooks
572 .bulkheads
573 .as_ref()
574 .unwrap()
575 .get("Boards.ListBoards")
576 .available()
577 );
578 }
579
580 #[tokio::test]
583 async fn a_call_refused_below_gives_back_the_permit_too() {
584 let hooks = ResilienceHooks::new(
585 Arc::new(Refusing),
586 ResilienceConfig {
587 bulkhead: Some(BulkheadConfig {
588 max_concurrent: 1,
589 max_wait: Duration::ZERO,
590 }),
591 ..ResilienceConfig::none()
592 },
593 );
594 let op = operation("Boards", "ListBoards");
595
596 let refused = hooks.on_operation_gate(&op).await.unwrap_err();
597
598 assert_eq!(ErrorCode::Usage, refused.code());
599 assert_eq!(
600 1,
601 hooks
602 .bulkheads
603 .as_ref()
604 .unwrap()
605 .get("Boards.ListBoards")
606 .available()
607 );
608 }
609
610 #[tokio::test]
611 async fn the_hooks_underneath_still_hear_everything() {
612 let recorder = Arc::new(Recorder::default());
613 let hooks = ResilienceHooks::new(recorder.clone(), ResilienceConfig::default());
614 let op = operation("Boards", "ListBoards");
615
616 hooks.on_operation_gate(&op).await.unwrap();
617 let state = hooks.on_operation_start(&op);
618 hooks.on_operation_end(&op, state, Ok(()), Duration::ZERO);
619
620 assert_eq!(
621 vec!["gate", "start", "end carrying its own"],
622 recorder.entries()
623 );
624 }
625
626 fn hooks(config: ResilienceConfig) -> ResilienceHooks {
627 ResilienceHooks::new(Arc::new(crate::observability::NoopHooks), config)
628 }
629
630 async fn fail(hooks: &ResilienceHooks, op: &OperationInfo) {
631 hooks.on_operation_gate(op).await.unwrap();
632 let state = hooks.on_operation_start(op);
633 hooks.on_operation_end(op, state, Err(&Error::api(500, "boom")), Duration::ZERO);
634 }
635
636 fn operation(service: &'static str, operation: &'static str) -> OperationInfo {
637 OperationInfo {
638 service: Cow::Borrowed(service),
639 operation: Cow::Borrowed(operation),
640 resource_type: Cow::Borrowed("board"),
641 is_mutation: false,
642 resource_id: None,
643 }
644 }
645
646 #[derive(Default)]
647 struct Recorder {
648 entries: Mutex<Vec<String>>,
649 }
650
651 impl Recorder {
652 fn entries(&self) -> Vec<String> {
653 self.entries
654 .lock()
655 .unwrap_or_else(std::sync::PoisonError::into_inner)
656 .clone()
657 }
658
659 fn record(&self, entry: &str) {
660 self.entries
661 .lock()
662 .unwrap_or_else(std::sync::PoisonError::into_inner)
663 .push(entry.to_string());
664 }
665 }
666
667 #[async_trait]
668 impl Hooks for Recorder {
669 async fn on_operation_gate(&self, _op: &OperationInfo) -> Result<(), Error> {
670 self.record("gate");
671 Ok(())
672 }
673
674 fn on_operation_start(&self, _op: &OperationInfo) -> OperationState {
675 self.record("start");
676 Some(Box::new("its own"))
677 }
678
679 fn on_operation_end(
680 &self,
681 _op: &OperationInfo,
682 state: OperationState,
683 _outcome: Result<(), &Error>,
684 _duration: Duration,
685 ) {
686 let carried = match state.and_then(|state| state.downcast::<&str>().ok()) {
687 Some(carried) => *carried,
688 None => "nothing",
689 };
690 self.record(&format!("end carrying {carried}"));
691 }
692 }
693
694 struct Refusing;
696
697 #[async_trait]
698 impl Hooks for Refusing {
699 async fn on_operation_gate(&self, _op: &OperationInfo) -> Result<(), Error> {
700 Err(Error::usage("blocked"))
701 }
702 }
703}