switchyard_libsy/algorithms/vgr/
safety.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::time::{Duration, Instant};
7
8use http::StatusCode;
9use parking_lot::Mutex;
10use switchyard_protocol::LlmClientError;
11
12use crate::LibsyError;
13
14/// Operator-controlled stop for local attempts.
15#[derive(Clone, Debug, Default)]
16pub struct KillSwitch(Arc<AtomicBool>);
17
18impl KillSwitch {
19    /// Creates a released switch.
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Stops local attempts for subsequent requests.
25    pub fn engage(&self) {
26        self.0.store(true, Ordering::Relaxed);
27    }
28
29    /// Allows local attempts again.
30    pub fn release(&self) {
31        self.0.store(false, Ordering::Relaxed);
32    }
33
34    /// Reports whether local attempts are stopped.
35    pub fn is_engaged(&self) -> bool {
36        self.0.load(Ordering::Relaxed)
37    }
38}
39
40/// Circuit-breaker configuration.
41#[derive(Clone, Copy, Debug)]
42pub struct BreakerConfig {
43    /// Consecutive endpoint failures required to open.
44    pub threshold: u32,
45    /// Delay before one half-open trial is admitted.
46    pub cooldown: Duration,
47}
48
49impl Default for BreakerConfig {
50    fn default() -> Self {
51        Self {
52            threshold: 5,
53            cooldown: Duration::from_secs(30),
54        }
55    }
56}
57
58#[derive(Default)]
59struct BreakerState {
60    failures: u32,
61    opened_at: Option<Instant>,
62    trial_in_flight: bool,
63}
64
65pub(super) struct CircuitBreaker {
66    config: BreakerConfig,
67    state: Mutex<BreakerState>,
68}
69
70impl CircuitBreaker {
71    pub(super) fn new(config: BreakerConfig) -> Self {
72        Self {
73            config,
74            state: Mutex::new(BreakerState::default()),
75        }
76    }
77
78    pub(super) fn attempt(&self) -> Option<Attempt<'_>> {
79        let mut state = self.state.lock();
80        let recovery = if let Some(opened_at) = state.opened_at {
81            if state.trial_in_flight || opened_at.elapsed() < self.config.cooldown {
82                return None;
83            }
84            state.trial_in_flight = true;
85            true
86        } else {
87            false
88        };
89        Some(Attempt {
90            breaker: self,
91            recovery,
92            completed: false,
93        })
94    }
95
96    pub(super) fn success(&self) {
97        *self.state.lock() = BreakerState::default();
98    }
99
100    /// Whether repeated failures have opened the breaker; it stays open until a trial succeeds.
101    pub(super) fn is_open(&self) -> bool {
102        self.state.lock().opened_at.is_some()
103    }
104
105    pub(super) fn failure(&self) {
106        let mut state = self.state.lock();
107        state.failures = state.failures.saturating_add(1);
108        state.trial_in_flight = false;
109        if state.failures >= self.config.threshold {
110            state.opened_at = Some(Instant::now());
111        }
112    }
113}
114
115/// Releases a recovery trial if its local call is cancelled or abandoned.
116pub(super) struct Attempt<'a> {
117    breaker: &'a CircuitBreaker,
118    recovery: bool,
119    completed: bool,
120}
121
122impl Attempt<'_> {
123    pub(super) fn success(mut self) {
124        self.completed = true;
125        self.breaker.success();
126    }
127
128    pub(super) fn failure(mut self) {
129        self.completed = true;
130        self.breaker.failure();
131    }
132}
133
134impl Drop for Attempt<'_> {
135    fn drop(&mut self) {
136        if self.recovery && !self.completed {
137            self.breaker.state.lock().trial_in_flight = false;
138        }
139    }
140}
141
142/// Errors for which retrying on the capable tier is safe.
143pub(super) fn fallback_eligible(error: &LibsyError) -> bool {
144    matches!(
145        error,
146        LibsyError::ClientCall { source, .. } if match source {
147            LlmClientError::ContextWindowExceeded { .. }
148                | LlmClientError::Transport { .. }
149                | LlmClientError::Timeout { .. }
150                | LlmClientError::TemporarilyUnavailable => true,
151            LlmClientError::UpstreamHttp { status, .. } =>
152                matches!(
153                    *status,
154                    StatusCode::FORBIDDEN
155                        | StatusCode::REQUEST_TIMEOUT
156                        | StatusCode::TOO_MANY_REQUESTS
157                ) || status.is_server_error(),
158            _ => false,
159        }
160    )
161}
162
163/// Only endpoint unavailability contributes to breaker health.
164pub(super) fn endpoint_failure(error: &LibsyError) -> bool {
165    matches!(
166        error,
167        LibsyError::ClientCall { source, .. } if match source {
168            LlmClientError::Transport { .. }
169            | LlmClientError::Timeout { .. }
170            | LlmClientError::TemporarilyUnavailable => true,
171            LlmClientError::UpstreamHttp { status, .. } => matches!(
172                *status,
173                StatusCode::BAD_GATEWAY
174                    | StatusCode::SERVICE_UNAVAILABLE
175                    | StatusCode::GATEWAY_TIMEOUT
176            ),
177            _ => false,
178        }
179    )
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn recovery_trials_are_exclusive_and_can_be_retried_after_failure_or_drop() {
188        let breaker = CircuitBreaker::new(BreakerConfig {
189            threshold: 1,
190            cooldown: Duration::ZERO,
191        });
192        breaker.failure();
193        let trial = breaker.attempt().expect("recovery trial");
194        assert!(breaker.attempt().is_none());
195        drop(trial);
196        breaker.attempt().expect("released trial").failure();
197        breaker.attempt().expect("retry after failure").success();
198        assert!(breaker.state.lock().opened_at.is_none());
199        assert!(breaker.attempt().is_some());
200    }
201
202    #[test]
203    fn recovery_waits_for_the_cooldown() {
204        let breaker = CircuitBreaker::new(BreakerConfig {
205            threshold: 1,
206            cooldown: Duration::from_secs(30),
207        });
208        breaker.failure();
209        assert!(breaker.attempt().is_none());
210    }
211}
switchyard_libsy/algorithms/vgr/safety.rs

switchyard_libsy/algorithms/vgr/
safety.rs
