fizzy_sdk/resilience/
bulkhead.rs1use std::sync::Arc;
2use std::time::Duration;
3
4use tokio::sync::{OwnedSemaphorePermit, Semaphore};
5
6use crate::error::Error;
7
8#[derive(Debug, Clone)]
13pub struct BulkheadConfig {
14 pub max_concurrent: usize,
16 pub max_wait: Duration,
19}
20
21impl Default for BulkheadConfig {
22 fn default() -> BulkheadConfig {
23 BulkheadConfig {
24 max_concurrent: 10,
25 max_wait: Duration::from_secs(5),
26 }
27 }
28}
29
30pub struct Bulkhead {
34 max_concurrent: usize,
35 max_wait: Duration,
36 permits: Arc<Semaphore>,
37}
38
39pub type BulkheadPermit = OwnedSemaphorePermit;
41
42impl Bulkhead {
43 #[allow(clippy::needless_pass_by_value)] pub fn new(config: BulkheadConfig) -> Bulkhead {
46 let max_concurrent = match config.max_concurrent {
47 0 => BulkheadConfig::default().max_concurrent,
48 max_concurrent => max_concurrent,
49 };
50 Bulkhead {
51 max_concurrent,
52 max_wait: config.max_wait,
53 permits: Arc::new(Semaphore::new(max_concurrent)),
54 }
55 }
56
57 pub async fn acquire(&self) -> Result<BulkheadPermit, Error> {
61 match self.try_acquire() {
62 Some(permit) => Ok(permit),
63 None if self.max_wait.is_zero() => Err(Error::bulkhead_full()),
64 None => tokio::time::timeout(self.max_wait, self.permits.clone().acquire_owned())
65 .await
66 .map_err(|_| Error::bulkhead_full())?
67 .map_err(|_| Error::bulkhead_full()),
68 }
69 }
70
71 pub fn try_acquire(&self) -> Option<BulkheadPermit> {
73 self.permits.clone().try_acquire_owned().ok()
74 }
75
76 pub fn available(&self) -> usize {
78 self.permits.available_permits()
79 }
80
81 pub fn in_use(&self) -> usize {
83 self.max_concurrent - self.available()
84 }
85}
86
87#[cfg(test)]
88#[allow(clippy::unwrap_used)]
89mod tests {
90 use super::*;
91
92 #[tokio::test]
93 async fn a_full_bulkhead_refuses_until_a_call_finishes() {
94 let bulkhead = Bulkhead::new(BulkheadConfig {
95 max_concurrent: 2,
96 max_wait: Duration::ZERO,
97 });
98
99 let first = bulkhead.acquire().await.unwrap();
100 let second = bulkhead.acquire().await.unwrap();
101 assert_eq!(2, bulkhead.in_use());
102 assert_eq!(0, bulkhead.available());
103
104 let refused = bulkhead.acquire().await.unwrap_err();
105 assert_eq!(Some(crate::error::Refusal::BulkheadFull), refused.refusal());
106 assert_eq!("bulkhead is full", refused.message());
107
108 drop(first);
109 assert_eq!(1, bulkhead.in_use());
110
111 let third = bulkhead.acquire().await.unwrap();
112 drop((second, third));
113 assert_eq!(0, bulkhead.in_use());
114 }
115
116 #[test]
117 fn try_acquire_answers_without_waiting() {
118 let bulkhead = Bulkhead::new(BulkheadConfig {
119 max_concurrent: 1,
120 max_wait: Duration::from_secs(5),
121 });
122
123 let permit = bulkhead.try_acquire();
124 assert!(permit.is_some());
125 assert!(bulkhead.try_acquire().is_none());
126
127 drop(permit);
128 assert!(bulkhead.try_acquire().is_some());
129 }
130
131 #[tokio::test]
132 async fn a_caller_that_may_wait_gets_the_permit_the_call_before_it_gives_back() {
133 let bulkhead = Arc::new(Bulkhead::new(BulkheadConfig {
134 max_concurrent: 1,
135 max_wait: Duration::from_millis(500),
136 }));
137
138 let held = bulkhead.acquire().await.unwrap();
139 tokio::spawn(async move {
140 tokio::time::sleep(Duration::from_millis(10)).await;
141 drop(held);
142 });
143
144 assert!(bulkhead.acquire().await.is_ok());
145 }
146
147 #[tokio::test]
148 async fn a_caller_that_waits_too_long_is_refused() {
149 let bulkhead = Bulkhead::new(BulkheadConfig {
150 max_concurrent: 1,
151 max_wait: Duration::from_millis(10),
152 });
153 let _held = bulkhead.acquire().await.unwrap();
154
155 let refused = bulkhead.acquire().await.unwrap_err();
156
157 assert_eq!(Some(crate::error::Refusal::BulkheadFull), refused.refusal());
158 }
159
160 #[test]
161 fn a_config_of_zero_falls_back_to_the_default_width() {
162 let bulkhead = Bulkhead::new(BulkheadConfig {
163 max_concurrent: 0,
164 max_wait: Duration::ZERO,
165 });
166
167 assert_eq!(10, bulkhead.available());
168 }
169}