gateway_core/
circuit.rs

1use std::{
2    collections::HashMap,
3    sync::Mutex,
4    time::{Duration, Instant},
5};
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum CircuitState {
9    Closed,
10    Open,
11    HalfOpen,
12}
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum CircuitDecision {
16    Allow { probe: bool },
17    Skip,
18}
19
20#[derive(Debug, Clone, Copy)]
21struct Circuit {
22    state: CircuitState,
23    failures: u32,
24    phase_started: Option<Instant>,
25}
26
27impl Default for Circuit {
28    fn default() -> Self {
29        Self {
30            state: CircuitState::Closed,
31            failures: 0,
32            phase_started: None,
33        }
34    }
35}
36
37pub struct CircuitBreaker {
38    threshold: u32,
39    cooldown: Duration,
40    circuits: Mutex<HashMap<String, Circuit>>,
41}
42
43impl CircuitBreaker {
44    pub fn new(threshold: u32, cooldown: Duration) -> Self {
45        Self {
46            threshold: threshold.max(1),
47            cooldown,
48            circuits: Mutex::new(HashMap::new()),
49        }
50    }
51
52    pub fn allow(&self, provider: &str) -> CircuitDecision {
53        self.allow_at(provider, Instant::now())
54    }
55
56    pub fn record_success(&self, provider: &str) {
57        let mut circuits = self.lock();
58        let circuit = circuits.entry(provider.to_owned()).or_default();
59        circuit.state = CircuitState::Closed;
60        circuit.failures = 0;
61        circuit.phase_started = None;
62    }
63
64    pub fn record_failure(&self, provider: &str) {
65        self.record_failure_at(provider, Instant::now());
66    }
67
68    pub fn state(&self, provider: &str) -> CircuitState {
69        self.lock()
70            .get(provider)
71            .map_or(CircuitState::Closed, |circuit| circuit.state)
72    }
73
74    pub fn snapshot(&self) -> Vec<(String, CircuitState)> {
75        self.lock()
76            .iter()
77            .map(|(provider, circuit)| (provider.clone(), circuit.state))
78            .collect()
79    }
80
81    fn allow_at(&self, provider: &str, now: Instant) -> CircuitDecision {
82        let mut circuits = self.lock();
83        let circuit = circuits.entry(provider.to_owned()).or_default();
84        match circuit.state {
85            CircuitState::Closed => CircuitDecision::Allow { probe: false },
86            CircuitState::Open if elapsed(circuit.phase_started, now) >= self.cooldown => {
87                circuit.state = CircuitState::HalfOpen;
88                circuit.phase_started = Some(now);
89                CircuitDecision::Allow { probe: true }
90            }
91            CircuitState::HalfOpen if elapsed(circuit.phase_started, now) >= self.cooldown => {
92                circuit.phase_started = Some(now);
93                CircuitDecision::Allow { probe: true }
94            }
95            CircuitState::Open | CircuitState::HalfOpen => CircuitDecision::Skip,
96        }
97    }
98
99    fn record_failure_at(&self, provider: &str, now: Instant) {
100        let mut circuits = self.lock();
101        let circuit = circuits.entry(provider.to_owned()).or_default();
102        match circuit.state {
103            CircuitState::Closed => {
104                circuit.failures = circuit.failures.saturating_add(1);
105                if circuit.failures >= self.threshold {
106                    circuit.state = CircuitState::Open;
107                    circuit.phase_started = Some(now);
108                }
109            }
110            CircuitState::HalfOpen => {
111                circuit.state = CircuitState::Open;
112                circuit.failures = self.threshold;
113                circuit.phase_started = Some(now);
114            }
115            CircuitState::Open => {}
116        }
117    }
118
119    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, Circuit>> {
120        self.circuits
121            .lock()
122            .unwrap_or_else(std::sync::PoisonError::into_inner)
123    }
124}
125
126fn elapsed(started: Option<Instant>, now: Instant) -> Duration {
127    started.map_or(Duration::MAX, |started| {
128        now.saturating_duration_since(started)
129    })
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn opens_skips_probes_and_recovers() {
138        let breaker = CircuitBreaker::new(2, Duration::from_secs(10));
139        let now = Instant::now();
140        breaker.record_failure_at("openai", now);
141        breaker.record_failure_at("openai", now);
142        assert_eq!(breaker.allow_at("openai", now), CircuitDecision::Skip);
143        assert_eq!(
144            breaker.allow_at("openai", now + Duration::from_secs(10)),
145            CircuitDecision::Allow { probe: true }
146        );
147        breaker.record_success("openai");
148        assert_eq!(breaker.state("openai"), CircuitState::Closed);
149    }
150}
gateway_core/circuit.rs

gateway_core/
circuit.rs
