1use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};
10use std::sync::{Arc, Mutex};
11use std::time::{Duration, Instant};
12
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub enum CircuitState {
16 Closed = 0,
18 Open = 1,
20 HalfOpen = 2,
22}
23
24impl CircuitState {
25 fn from_u8(val: u8) -> Self {
26 match val {
27 0 => CircuitState::Closed,
28 1 => CircuitState::Open,
29 2 => CircuitState::HalfOpen,
30 _ => CircuitState::Closed,
31 }
32 }
33}
34
35#[derive(Debug)]
37pub struct CircuitBreaker {
38 state: Arc<AtomicU8>,
40 failures: Arc<AtomicUsize>,
42 last_failure: Arc<Mutex<Instant>>,
44 half_open_probe_in_flight: Arc<AtomicBool>,
48 config: Arc<CircuitBreakerConfig>,
50}
51
52#[derive(Clone, Debug)]
54pub struct CircuitBreakerConfig {
55 pub failure_threshold: usize,
57 pub open_duration_secs: u64,
59 pub reset_timeout_secs: u64,
61}
62
63impl Default for CircuitBreakerConfig {
64 fn default() -> Self {
65 Self {
66 failure_threshold: 5,
67 open_duration_secs: 30,
68 reset_timeout_secs: 60,
69 }
70 }
71}
72
73impl CircuitBreaker {
74 pub fn new() -> Self {
76 Self::with_config(CircuitBreakerConfig::default())
77 }
78
79 pub fn with_config(config: CircuitBreakerConfig) -> Self {
81 Self {
82 state: Arc::new(AtomicU8::new(0)), failures: Arc::new(AtomicUsize::new(0)),
84 last_failure: Arc::new(Mutex::new(Instant::now())),
85 half_open_probe_in_flight: Arc::new(AtomicBool::new(false)),
86 config: Arc::new(config),
87 }
88 }
89
90 pub fn current_state(&self) -> CircuitState {
92 let state_byte = self.state.load(Ordering::Relaxed);
93 CircuitState::from_u8(state_byte)
94 }
95
96 pub fn should_attempt(&self) -> bool {
104 match self.current_state() {
105 CircuitState::Closed => true,
106 CircuitState::Open => {
107 let last_failure = *self.last_failure.lock().unwrap();
109 let elapsed = last_failure.elapsed();
110
111 if elapsed <= Duration::from_secs(self.config.open_duration_secs) {
112 return false;
113 }
114
115 if self
120 .half_open_probe_in_flight
121 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
122 .is_err()
123 {
124 return false;
125 }
126
127 if self
128 .state
129 .compare_exchange(1, 2, Ordering::AcqRel, Ordering::Relaxed)
130 .is_ok()
131 {
132 true
133 } else {
134 self.half_open_probe_in_flight
138 .store(false, Ordering::Release);
139 false
140 }
141 }
142 CircuitState::HalfOpen => self
143 .half_open_probe_in_flight
144 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
145 .is_ok(),
146 }
147 }
148
149 pub fn record_success(&self) {
151 let last_failure = *self.last_failure.lock().unwrap();
153 if last_failure.elapsed() > Duration::from_secs(self.config.reset_timeout_secs) {
154 self.failures.store(0, Ordering::Relaxed);
155 self.state.store(0, Ordering::Relaxed); self.half_open_probe_in_flight
157 .store(false, Ordering::Release);
158 } else if self.current_state() == CircuitState::HalfOpen {
159 self.failures.store(0, Ordering::Relaxed);
161 self.state.store(0, Ordering::Relaxed);
162 self.half_open_probe_in_flight
163 .store(false, Ordering::Release);
164 }
165 }
166
167 pub fn record_failure(&self) {
169 let failure_count = self.failures.fetch_add(1, Ordering::Relaxed) + 1;
170 *self.last_failure.lock().unwrap() = Instant::now();
171
172 if failure_count >= self.config.failure_threshold
175 || self.current_state() == CircuitState::HalfOpen
176 {
177 self.state.store(1, Ordering::Relaxed); self.half_open_probe_in_flight
179 .store(false, Ordering::Release);
180 }
181 }
182
183 pub fn failure_count(&self) -> usize {
185 self.failures.load(Ordering::Relaxed)
186 }
187
188 pub fn time_since_last_failure(&self) -> Duration {
190 self.last_failure.lock().unwrap().elapsed()
191 }
192
193 pub fn reset(&self) {
195 self.state.store(0, Ordering::Relaxed);
196 self.failures.store(0, Ordering::Relaxed);
197 self.half_open_probe_in_flight
198 .store(false, Ordering::Release);
199 *self.last_failure.lock().unwrap() = Instant::now();
200 }
201}
202
203impl Default for CircuitBreaker {
204 fn default() -> Self {
205 Self::new()
206 }
207}
208
209impl Clone for CircuitBreaker {
210 fn clone(&self) -> Self {
211 Self {
212 state: Arc::clone(&self.state),
213 failures: Arc::clone(&self.failures),
214 last_failure: Arc::clone(&self.last_failure),
215 half_open_probe_in_flight: Arc::clone(&self.half_open_probe_in_flight),
216 config: Arc::clone(&self.config),
217 }
218 }
219}
220
221pub struct JitteredBackoff {
223 attempt: u32,
224 max_backoff_ms: u64,
225}
226
227impl JitteredBackoff {
228 pub fn new() -> Self {
230 Self {
231 attempt: 0,
232 max_backoff_ms: 30_000, }
234 }
235
236 pub fn with_max_backoff(max_backoff_ms: u64) -> Self {
238 Self {
239 attempt: 0,
240 max_backoff_ms,
241 }
242 }
243
244 pub fn backoff_duration(&self) -> Duration {
246 let base_ms = (100u64 * 2_u64.pow(self.attempt.min(10))).min(self.max_backoff_ms);
248
249 use rand::Rng;
251 let mut rng = rand::thread_rng();
252 let jitter_ms = rng.gen_range(0..=(base_ms / 10));
253
254 Duration::from_millis(base_ms + jitter_ms)
255 }
256
257 pub fn next_backoff(&mut self) -> Duration {
259 let duration = self.backoff_duration();
260 self.attempt += 1;
261 duration
262 }
263
264 pub fn reset(&mut self) {
266 self.attempt = 0;
267 }
268}
269
270impl Default for JitteredBackoff {
271 fn default() -> Self {
272 Self::new()
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 use super::*;
279
280 #[test]
281 fn circuit_starts_closed() {
282 let cb = CircuitBreaker::new();
283 assert_eq!(cb.current_state(), CircuitState::Closed);
284 assert!(cb.should_attempt());
285 }
286
287 #[test]
288 fn circuit_opens_after_threshold_failures() {
289 let cb = CircuitBreaker::new();
290 for _ in 0..5 {
291 cb.record_failure();
292 }
293 assert_eq!(cb.current_state(), CircuitState::Open);
294 assert!(!cb.should_attempt());
295 }
296
297 #[test]
298 fn circuit_half_opens_after_timeout() {
299 let config = CircuitBreakerConfig {
300 failure_threshold: 1,
301 open_duration_secs: 0,
302 reset_timeout_secs: 60,
303 };
304 let cb = CircuitBreaker::with_config(config);
305 cb.record_failure();
306 assert_eq!(cb.current_state(), CircuitState::Open);
307
308 std::thread::sleep(Duration::from_millis(100));
309 assert!(cb.should_attempt());
310 assert_eq!(cb.current_state(), CircuitState::HalfOpen);
311 }
312
313 #[test]
314 fn circuit_closes_on_successful_probe() {
315 let config = CircuitBreakerConfig {
316 failure_threshold: 1,
317 open_duration_secs: 0,
318 reset_timeout_secs: 0,
319 };
320 let cb = CircuitBreaker::with_config(config);
321 cb.record_failure();
322 assert_eq!(cb.current_state(), CircuitState::Open);
323
324 std::thread::sleep(Duration::from_millis(100));
326 let can_attempt = cb.should_attempt();
327 assert!(can_attempt);
328 assert_eq!(cb.current_state(), CircuitState::HalfOpen);
329
330 cb.record_success();
332 assert_eq!(cb.current_state(), CircuitState::Closed);
333 assert_eq!(cb.failure_count(), 0);
334 }
335
336 #[test]
337 fn jittered_backoff_increases() {
338 let mut backoff = JitteredBackoff::new();
339 let backoff1 = backoff.next_backoff();
340 let backoff2 = backoff.next_backoff();
341 let backoff3 = backoff.next_backoff();
342
343 assert!(backoff1.as_millis() >= 100 && backoff1.as_millis() <= 220);
344 assert!(backoff2.as_millis() >= 200 && backoff2.as_millis() <= 440);
345 assert!(backoff3.as_millis() >= 400 && backoff3.as_millis() <= 880);
346 }
347
348 #[test]
349 fn jittered_backoff_has_jitter() {
350 let mut backoffs = Vec::new();
351 for _ in 0..10 {
352 let backoff = JitteredBackoff::new().next_backoff();
353 backoffs.push(backoff.as_millis());
354 }
355
356 let min = *backoffs.iter().min().unwrap();
357 let max = *backoffs.iter().max().unwrap();
358 assert!(max > min, "Jitter should produce different values");
359 }
360
361 #[test]
362 fn half_open_grants_probe_to_exactly_one_concurrent_caller() {
363 let config = CircuitBreakerConfig {
364 failure_threshold: 1,
365 open_duration_secs: 0,
366 reset_timeout_secs: 60,
367 };
368 let cb = CircuitBreaker::with_config(config);
369 cb.record_failure();
370 assert_eq!(cb.current_state(), CircuitState::Open);
371 std::thread::sleep(Duration::from_millis(50));
372
373 let granted = Arc::new(AtomicUsize::new(0));
374 let handles: Vec<_> = (0..16)
375 .map(|_| {
376 let cb = cb.clone();
377 let granted = Arc::clone(&granted);
378 std::thread::spawn(move || {
379 if cb.should_attempt() {
380 granted.fetch_add(1, Ordering::Relaxed);
381 }
382 })
383 })
384 .collect();
385 for h in handles {
386 h.join().unwrap();
387 }
388
389 assert_eq!(
390 granted.load(Ordering::Relaxed),
391 1,
392 "exactly one concurrent caller should win the half-open probe slot"
393 );
394 assert_eq!(cb.current_state(), CircuitState::HalfOpen);
395 }
396
397 #[test]
398 fn circuit_breaker_is_cloneable() {
399 let cb1 = CircuitBreaker::new();
400 cb1.record_failure();
401 let cb2 = cb1.clone();
402
403 assert_eq!(cb2.failure_count(), 1);
404 cb2.record_failure();
405 assert_eq!(cb1.failure_count(), 2);
406 }
407}