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

runifold_model/
circuit.rs
