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    pub(super) fn failure(&self) {
101        let mut state = self.state.lock();
102        state.failures = state.failures.saturating_add(1);
103        state.trial_in_flight = false;
104        if state.failures >= self.config.threshold {
105            state.opened_at = Some(Instant::now());
106        }
107    }
108}
109
110/// Releases a recovery trial if its local call is cancelled or abandoned.
111pub(super) struct Attempt<'a> {
112    breaker: &'a CircuitBreaker,
113    recovery: bool,
114    completed: bool,
115}
116
117impl Attempt<'_> {
118    pub(super) fn success(mut self) {
119        self.completed = true;
120        self.breaker.success();
121    }
122
123    pub(super) fn failure(mut self) {
124        self.completed = true;
125        self.breaker.failure();
126    }
127}
128
129impl Drop for Attempt<'_> {
130    fn drop(&mut self) {
131        if self.recovery && !self.completed {
132            self.breaker.state.lock().trial_in_flight = false;
133        }
134    }
135}
136
137/// Errors for which retrying on the capable tier is safe.
138pub(super) fn fallback_eligible(error: &LibsyError) -> bool {
139    matches!(
140        error,
141        LibsyError::ClientCall { source, .. } if match source {
142            LlmClientError::ContextWindowExceeded { .. }
143                | LlmClientError::Transport { .. }
144                | LlmClientError::Timeout { .. } => true,
145            LlmClientError::UpstreamHttp { status, .. } =>
146                matches!(
147                    *status,
148                    StatusCode::FORBIDDEN
149                        | StatusCode::REQUEST_TIMEOUT
150                        | StatusCode::TOO_MANY_REQUESTS
151                ) || status.is_server_error(),
152            _ => false,
153        }
154    )
155}
156
157/// Only endpoint unavailability contributes to breaker health.
158pub(super) fn endpoint_failure(error: &LibsyError) -> bool {
159    matches!(
160        error,
161        LibsyError::ClientCall { source, .. } if match source {
162            LlmClientError::Transport { .. } | LlmClientError::Timeout { .. } => true,
163            LlmClientError::UpstreamHttp { status, .. } => matches!(
164                *status,
165                StatusCode::BAD_GATEWAY
166                    | StatusCode::SERVICE_UNAVAILABLE
167                    | StatusCode::GATEWAY_TIMEOUT
168            ),
169            _ => false,
170        }
171    )
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[test]
179    fn recovery_trials_are_exclusive_and_can_be_retried_after_failure_or_drop() {
180        let breaker = CircuitBreaker::new(BreakerConfig {
181            threshold: 1,
182            cooldown: Duration::ZERO,
183        });
184        breaker.failure();
185        let trial = breaker.attempt().expect("recovery trial");
186        assert!(breaker.attempt().is_none());
187        drop(trial);
188        breaker.attempt().expect("released trial").failure();
189        breaker.attempt().expect("retry after failure").success();
190        assert!(breaker.state.lock().opened_at.is_none());
191        assert!(breaker.attempt().is_some());
192    }
193
194    #[test]
195    fn recovery_waits_for_the_cooldown() {
196        let breaker = CircuitBreaker::new(BreakerConfig {
197            threshold: 1,
198            cooldown: Duration::from_secs(30),
199        });
200        breaker.failure();
201        assert!(breaker.attempt().is_none());
202    }
203}
switchyard_libsy/algorithms/vgr/safety.rs

switchyard_libsy/algorithms/vgr/
safety.rs
