runifold_model/
circuit.rs

1use std::{
2    sync::{Arc, Mutex},
3    time::Duration,
4};
5
6use runifold_core::Instant;
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9
10use crate::{ModelError, ModelErrorKind, ModelRef};
11
12/// Clock used by model-routing resilience policy.
13///
14/// Applications may inject a deterministic implementation for tests.
15pub trait RouterClock: Send + Sync {
16    /// Returns the current monotonic time.
17    fn now(&self) -> Instant;
18}
19
20/// Monotonic system clock.
21#[derive(Clone, Copy, Debug, Default)]
22pub struct SystemRouterClock;
23
24impl RouterClock for SystemRouterClock {
25    fn now(&self) -> Instant {
26        Instant::now()
27    }
28}
29
30/// Invalid circuit-breaker configuration.
31#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
32#[non_exhaustive]
33pub enum CircuitBreakerConfigError {
34    /// A circuit could never open.
35    #[error("circuit-breaker failure threshold must be greater than zero")]
36    ZeroFailureThreshold,
37    /// An open circuit would immediately expire.
38    #[error("circuit-breaker cooldown must be greater than zero")]
39    ZeroCooldown,
40}
41
42/// Per-route circuit-breaker policy.
43#[derive(Clone, Debug, Eq, PartialEq)]
44pub struct CircuitBreakerConfig {
45    failure_threshold: u32,
46    cooldown: Duration,
47    counted_kinds: Vec<ModelErrorKind>,
48}
49
50impl Default for CircuitBreakerConfig {
51    fn default() -> Self {
52        Self {
53            failure_threshold: 5,
54            cooldown: Duration::from_secs(30),
55            counted_kinds: vec![
56                ModelErrorKind::Transport,
57                ModelErrorKind::Provider,
58                ModelErrorKind::Protocol,
59                ModelErrorKind::StreamState,
60            ],
61        }
62    }
63}
64
65impl CircuitBreakerConfig {
66    /// Creates a breaker that counts transport, provider, protocol, and stream
67    /// state failures.
68    ///
69    /// # Errors
70    ///
71    /// Returns [`CircuitBreakerConfigError`] when the threshold or cooldown is
72    /// zero.
73    pub fn new(
74        failure_threshold: u32,
75        cooldown: Duration,
76    ) -> Result<Self, CircuitBreakerConfigError> {
77        if failure_threshold == 0 {
78            return Err(CircuitBreakerConfigError::ZeroFailureThreshold);
79        }
80        if cooldown.is_zero() {
81            return Err(CircuitBreakerConfigError::ZeroCooldown);
82        }
83        Ok(Self {
84            failure_threshold,
85            cooldown,
86            counted_kinds: vec![
87                ModelErrorKind::Transport,
88                ModelErrorKind::Provider,
89                ModelErrorKind::Protocol,
90                ModelErrorKind::StreamState,
91            ],
92        })
93    }
94
95    /// Replaces the failure kinds counted by this breaker.
96    #[must_use]
97    pub fn counted_kinds(mut self, kinds: impl IntoIterator<Item = ModelErrorKind>) -> Self {
98        self.counted_kinds.clear();
99        for kind in kinds {
100            if !self.counted_kinds.contains(&kind) {
101                self.counted_kinds.push(kind);
102            }
103        }
104        self
105    }
106
107    /// Returns the consecutive counted-failure threshold.
108    pub const fn failure_threshold(&self) -> u32 {
109        self.failure_threshold
110    }
111
112    /// Returns how long an opened route remains unavailable before probing.
113    pub const fn cooldown(&self) -> Duration {
114        self.cooldown
115    }
116
117    /// Returns error kinds that contribute to opening the circuit.
118    pub fn failure_kinds(&self) -> &[ModelErrorKind] {
119        &self.counted_kinds
120    }
121
122    pub(crate) fn counts(&self, error: &ModelError) -> bool {
123        error.kind != ModelErrorKind::Cancelled && self.counted_kinds.contains(&error.kind)
124    }
125}
126
127#[cfg(test)]
128mod config_tests {
129    use std::time::Duration;
130
131    use super::CircuitBreakerConfig;
132    use crate::ModelErrorKind;
133
134    #[test]
135    fn default_breaker_is_bounded_and_counts_runtime_failures() {
136        let config = CircuitBreakerConfig::default();
137
138        assert_eq!(config.failure_threshold(), 5);
139        assert_eq!(config.cooldown(), Duration::from_secs(30));
140        assert!(config.failure_kinds().contains(&ModelErrorKind::Transport));
141        assert!(config.failure_kinds().contains(&ModelErrorKind::Protocol));
142    }
143}
144
145/// Public route-health state.
146#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
147#[serde(rename_all = "snake_case")]
148#[non_exhaustive]
149pub enum CircuitState {
150    /// Requests may use the route.
151    Closed,
152    /// Requests skip the route during its cooldown.
153    Open,
154    /// Exactly one recovery probe is currently using the route.
155    HalfOpen,
156}
157
158/// Point-in-time health for one physical model route.
159#[derive(Clone, Debug, Eq, PartialEq)]
160pub struct ModelRouteHealth {
161    /// Stable route name.
162    pub route: String,
163    /// Physical model target.
164    pub target: ModelRef,
165    /// Current circuit state.
166    pub state: CircuitState,
167    /// Consecutive counted failures in the current generation.
168    pub consecutive_failures: u32,
169    /// Remaining cooldown for an open route.
170    pub retry_after: Option<Duration>,
171}
172
173#[derive(Debug)]
174pub(crate) struct BreakerState {
175    generation: u64,
176    phase: BreakerPhase,
177}
178
179#[derive(Debug)]
180enum BreakerPhase {
181    Closed { failures: u32 },
182    Open { until: Option<Instant> },
183    HalfOpen,
184}
185
186impl Default for BreakerState {
187    fn default() -> Self {
188        Self {
189            generation: 0,
190            phase: BreakerPhase::Closed { failures: 0 },
191        }
192    }
193}
194
195pub(crate) type SharedBreakerState = Arc<Mutex<BreakerState>>;
196
197pub(crate) enum RoutePermit {
198    Disabled,
199    Acquired(BreakerPermit),
200    Rejected,
201}
202
203pub(crate) struct BreakerPermit {
204    state: SharedBreakerState,
205    config: CircuitBreakerConfig,
206    clock: Arc<dyn RouterClock>,
207    generation: u64,
208    probe: bool,
209    resolved: bool,
210}
211
212impl BreakerPermit {
213    pub(crate) const fn is_probe(&self) -> bool {
214        self.probe
215    }
216
217    pub(crate) fn success(mut self) {
218        let mut state = lock(&self.state);
219        if state.generation == self.generation {
220            state.generation = state.generation.wrapping_add(1);
221            state.phase = BreakerPhase::Closed { failures: 0 };
222        }
223        self.resolved = true;
224    }
225
226    pub(crate) fn failure(mut self, error: &ModelError) {
227        if self.config.counts(error) {
228            record_counted_failure(
229                &self.state,
230                &self.config,
231                self.clock.now(),
232                self.generation,
233                self.probe,
234            );
235        } else if self.probe {
236            reopen(
237                &self.state,
238                self.clock.now(),
239                self.config.cooldown,
240                self.generation,
241            );
242        }
243        self.resolved = true;
244    }
245}
246
247impl Drop for BreakerPermit {
248    fn drop(&mut self) {
249        if self.probe && !self.resolved {
250            reopen(
251                &self.state,
252                self.clock.now(),
253                self.config.cooldown,
254                self.generation,
255            );
256        }
257    }
258}
259
260pub(crate) fn acquire(
261    state: &SharedBreakerState,
262    config: Option<&CircuitBreakerConfig>,
263    clock: &Arc<dyn RouterClock>,
264) -> RoutePermit {
265    let Some(config) = config else {
266        return RoutePermit::Disabled;
267    };
268    let now = clock.now();
269    let mut state_guard = lock(state);
270    let generation = state_guard.generation;
271    let probe = match state_guard.phase {
272        BreakerPhase::Closed { .. } => false,
273        BreakerPhase::Open { until: Some(until) } if now >= until => {
274            state_guard.phase = BreakerPhase::HalfOpen;
275            true
276        }
277        BreakerPhase::Open { .. } | BreakerPhase::HalfOpen => return RoutePermit::Rejected,
278    };
279    drop(state_guard);
280    RoutePermit::Acquired(BreakerPermit {
281        state: state.clone(),
282        config: config.clone(),
283        clock: clock.clone(),
284        generation,
285        probe,
286        resolved: false,
287    })
288}
289
290pub(crate) fn snapshot(
291    state: &SharedBreakerState,
292    route: String,
293    target: ModelRef,
294    config: Option<&CircuitBreakerConfig>,
295    now: Instant,
296) -> ModelRouteHealth {
297    let state = lock(state);
298    let (health, failures, retry_after) = match state.phase {
299        BreakerPhase::Closed { failures } => (CircuitState::Closed, failures, None),
300        BreakerPhase::Open { until } => (
301            CircuitState::Open,
302            0,
303            until.map(|until| until.saturating_duration_since(now)),
304        ),
305        BreakerPhase::HalfOpen => (CircuitState::HalfOpen, 0, None),
306    };
307    if config.is_none() {
308        return ModelRouteHealth {
309            route,
310            target,
311            state: CircuitState::Closed,
312            consecutive_failures: 0,
313            retry_after: None,
314        };
315    }
316    ModelRouteHealth {
317        route,
318        target,
319        state: health,
320        consecutive_failures: failures,
321        retry_after,
322    }
323}
324
325fn record_counted_failure(
326    state: &SharedBreakerState,
327    config: &CircuitBreakerConfig,
328    now: Instant,
329    generation: u64,
330    probe: bool,
331) {
332    let mut state = lock(state);
333    if state.generation != generation {
334        return;
335    }
336    if probe {
337        state.generation = state.generation.wrapping_add(1);
338        state.phase = BreakerPhase::Open {
339            until: now.checked_add(config.cooldown),
340        };
341        return;
342    }
343    let BreakerPhase::Closed { failures } = &mut state.phase else {
344        return;
345    };
346    *failures = failures.saturating_add(1);
347    if *failures >= config.failure_threshold {
348        state.generation = state.generation.wrapping_add(1);
349        state.phase = BreakerPhase::Open {
350            until: now.checked_add(config.cooldown),
351        };
352    }
353}
354
355fn reopen(state: &SharedBreakerState, now: Instant, cooldown: Duration, generation: u64) {
356    let mut state = lock(state);
357    if state.generation == generation {
358        state.generation = state.generation.wrapping_add(1);
359        state.phase = BreakerPhase::Open {
360            until: now.checked_add(cooldown),
361        };
362    }
363}
364
365fn lock(state: &SharedBreakerState) -> std::sync::MutexGuard<'_, BreakerState> {
366    state
367        .lock()
368        .unwrap_or_else(std::sync::PoisonError::into_inner)
369}
370
371#[cfg(test)]
372mod tests {
373    use std::{
374        sync::{Arc, Mutex},
375        time::{Duration, Instant},
376    };
377
378    use runifold_core::RetrySafety;
379
380    use crate::{ModelError, ModelErrorKind, ModelRef};
381
382    use super::{
383        BreakerState, CircuitBreakerConfig, CircuitState, RoutePermit, RouterClock,
384        SharedBreakerState, acquire, snapshot,
385    };
386
387    struct ManualClock {
388        now: Mutex<Instant>,
389    }
390
391    impl ManualClock {
392        fn new() -> Self {
393            Self {
394                now: Mutex::new(Instant::now()),
395            }
396        }
397
398        fn advance(&self, duration: Duration) {
399            let mut now = self
400                .now
401                .lock()
402                .unwrap_or_else(std::sync::PoisonError::into_inner);
403            *now += duration;
404        }
405    }
406
407    impl RouterClock for ManualClock {
408        fn now(&self) -> Instant {
409            *self
410                .now
411                .lock()
412                .unwrap_or_else(std::sync::PoisonError::into_inner)
413        }
414    }
415
416    fn state() -> SharedBreakerState {
417        Arc::new(Mutex::new(BreakerState::default()))
418    }
419
420    fn config(threshold: u32) -> CircuitBreakerConfig {
421        CircuitBreakerConfig::new(threshold, Duration::from_secs(10)).unwrap()
422    }
423
424    fn failure() -> ModelError {
425        let mut error = ModelError::local(ModelErrorKind::Transport, "failure");
426        error.retry_safety = RetrySafety::Safe;
427        error
428    }
429
430    fn permit(
431        state: &SharedBreakerState,
432        config: &CircuitBreakerConfig,
433        clock: &Arc<dyn RouterClock>,
434    ) -> super::BreakerPermit {
435        match acquire(state, Some(config), clock) {
436            RoutePermit::Acquired(permit) => permit,
437            RoutePermit::Disabled | RoutePermit::Rejected => panic!("expected route permit"),
438        }
439    }
440
441    fn health(
442        state: &SharedBreakerState,
443        config: &CircuitBreakerConfig,
444        clock: &Arc<dyn RouterClock>,
445    ) -> super::ModelRouteHealth {
446        snapshot(
447            state,
448            "route".into(),
449            ModelRef::new("test", "model"),
450            Some(config),
451            clock.now(),
452        )
453    }
454
455    #[test]
456    fn threshold_opens_then_one_successful_probe_closes() {
457        let state = state();
458        let clock_impl = Arc::new(ManualClock::new());
459        let clock: Arc<dyn RouterClock> = clock_impl.clone();
460        let config = config(2);
461
462        permit(&state, &config, &clock).failure(&failure());
463        assert_eq!(health(&state, &config, &clock).consecutive_failures, 1);
464        permit(&state, &config, &clock).failure(&failure());
465        assert_eq!(health(&state, &config, &clock).state, CircuitState::Open);
466        assert!(matches!(
467            acquire(&state, Some(&config), &clock),
468            RoutePermit::Rejected
469        ));
470
471        clock_impl.advance(config.cooldown());
472        let probe = permit(&state, &config, &clock);
473        assert_eq!(
474            health(&state, &config, &clock).state,
475            CircuitState::HalfOpen
476        );
477        assert!(matches!(
478            acquire(&state, Some(&config), &clock),
479            RoutePermit::Rejected
480        ));
481        probe.success();
482
483        let health = health(&state, &config, &clock);
484        assert_eq!(health.state, CircuitState::Closed);
485        assert_eq!(health.consecutive_failures, 0);
486    }
487
488    #[test]
489    fn stale_failure_cannot_overwrite_a_newer_success_generation() {
490        let state = state();
491        let clock_impl = Arc::new(ManualClock::new());
492        let clock: Arc<dyn RouterClock> = clock_impl.clone();
493        let config = config(1);
494        let delayed_permit = permit(&state, &config, &clock);
495        let opener = permit(&state, &config, &clock);
496
497        opener.failure(&failure());
498        clock_impl.advance(config.cooldown());
499        permit(&state, &config, &clock).success();
500        delayed_permit.failure(&failure());
501
502        assert_eq!(health(&state, &config, &clock).state, CircuitState::Closed);
503    }
504
505    #[test]
506    fn abandoned_half_open_probe_reopens_the_route() {
507        let state = state();
508        let clock_impl = Arc::new(ManualClock::new());
509        let clock: Arc<dyn RouterClock> = clock_impl.clone();
510        let config = config(1);
511
512        permit(&state, &config, &clock).failure(&failure());
513        clock_impl.advance(config.cooldown());
514        let probe = permit(&state, &config, &clock);
515        drop(probe);
516
517        let health = health(&state, &config, &clock);
518        assert_eq!(health.state, CircuitState::Open);
519        assert_eq!(health.retry_after, Some(config.cooldown()));
520    }
521
522    #[test]
523    fn non_counted_failure_does_not_damage_a_closed_route() {
524        let state = state();
525        let clock: Arc<dyn RouterClock> = Arc::new(ManualClock::new());
526        let config = config(1);
527        let error = ModelError::local(ModelErrorKind::InvalidRequest, "caller error");
528
529        permit(&state, &config, &clock).failure(&error);
530
531        let health = health(&state, &config, &clock);
532        assert_eq!(health.state, CircuitState::Closed);
533        assert_eq!(health.consecutive_failures, 0);
534    }
535}
runifold_model/circuit.rs

runifold_model/
circuit.rs
