specado_core/
circuit_breaker.rs

1use crate::error::{Error, Result};
2use std::future::Future;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5use tokio::sync::Mutex;
6
7#[derive(Debug)]
8pub struct CircuitBreaker {
9    state: Arc<Mutex<CircuitState>>,
10    failure_threshold: usize,
11    timeout: Duration,
12    half_open_max_requests: usize,
13}
14
15#[derive(Debug)]
16enum CircuitState {
17    Closed {
18        failure_count: usize,
19    },
20    Open {
21        opened_at: Instant,
22    },
23    HalfOpen {
24        trial_count: usize,
25        success_count: usize,
26    },
27}
28
29impl CircuitBreaker {
30    pub fn new(failure_threshold: usize, timeout: Duration, half_open_max_requests: usize) -> Self {
31        Self {
32            state: Arc::new(Mutex::new(CircuitState::Closed { failure_count: 0 })),
33            failure_threshold: failure_threshold.max(1),
34            timeout,
35            half_open_max_requests: half_open_max_requests.max(1),
36        }
37    }
38
39    pub async fn call<F, Fut, T>(&self, operation: F) -> Result<T>
40    where
41        F: FnOnce() -> Fut,
42        Fut: Future<Output = Result<T>> + Send,
43        T: Send,
44    {
45        self.before_call().await?;
46        let outcome = operation().await;
47        match outcome {
48            Ok(value) => {
49                self.on_success().await;
50                Ok(value)
51            }
52            Err(err) => {
53                self.on_failure().await;
54                Err(err)
55            }
56        }
57    }
58
59    async fn before_call(&self) -> Result<()> {
60        let mut state = self.state.lock().await;
61        loop {
62            match &mut *state {
63                CircuitState::Closed { .. } => {
64                    break;
65                }
66                CircuitState::Open { opened_at } => {
67                    if opened_at.elapsed() >= self.timeout {
68                        *state = CircuitState::HalfOpen {
69                            trial_count: 0,
70                            success_count: 0,
71                        };
72                        continue;
73                    } else {
74                        return Err(Error::CircuitBreakerOpen);
75                    }
76                }
77                CircuitState::HalfOpen { trial_count, .. } => {
78                    if *trial_count >= self.half_open_max_requests {
79                        return Err(Error::CircuitBreakerHalfOpen);
80                    }
81                    *trial_count += 1;
82                    break;
83                }
84            }
85        }
86        Ok(())
87    }
88
89    async fn on_success(&self) {
90        let mut state = self.state.lock().await;
91        match &mut *state {
92            CircuitState::Closed { failure_count } => {
93                *failure_count = 0;
94            }
95            CircuitState::HalfOpen { success_count, .. } => {
96                *success_count += 1;
97                if *success_count >= self.half_open_max_requests {
98                    *state = CircuitState::Closed { failure_count: 0 };
99                }
100            }
101            CircuitState::Open { .. } => {
102                *state = CircuitState::Closed { failure_count: 0 };
103            }
104        }
105    }
106
107    async fn on_failure(&self) {
108        let mut state = self.state.lock().await;
109        match &mut *state {
110            CircuitState::Closed { failure_count } => {
111                *failure_count += 1;
112                if *failure_count >= self.failure_threshold {
113                    *state = CircuitState::Open {
114                        opened_at: Instant::now(),
115                    };
116                }
117            }
118            CircuitState::HalfOpen { .. } => {
119                *state = CircuitState::Open {
120                    opened_at: Instant::now(),
121                };
122            }
123            CircuitState::Open { .. } => {}
124        }
125    }
126}
127
128impl Default for CircuitBreaker {
129    fn default() -> Self {
130        Self::new(5, Duration::from_secs(30), 3)
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::error::Error;
138
139    #[tokio::test]
140    async fn opens_after_threshold_failures() {
141        let breaker = CircuitBreaker::new(2, Duration::from_secs(30), 1);
142
143        assert!(breaker
144            .call(|| async { Err::<(), _>(Error::Transform("fail".into())) })
145            .await
146            .is_err());
147        assert!(breaker
148            .call(|| async { Err::<(), _>(Error::Transform("fail".into())) })
149            .await
150            .is_err());
151
152        let err = breaker
153            .call(|| async { Ok::<(), Error>(()) })
154            .await
155            .unwrap_err();
156        assert!(matches!(err, Error::CircuitBreakerOpen));
157    }
158
159    #[tokio::test]
160    async fn half_open_allows_recovery() {
161        let breaker = CircuitBreaker::new(1, Duration::from_millis(0), 1);
162
163        let _ = breaker
164            .call(|| async { Err::<(), _>(Error::Transform("fail".into())) })
165            .await;
166
167        // timeout is zero so breaker should move to half-open and allow a probe request
168        let result = breaker.call(|| async { Ok::<_, Error>("success") }).await;
169        assert!(result.is_ok());
170
171        // subsequent requests should flow normally
172        let result = breaker.call(|| async { Ok::<_, Error>("again") }).await;
173        assert!(result.is_ok());
174    }
175}
specado_core/circuit_breaker.rs

specado_core/
circuit_breaker.rs
