use std::sync::Mutex;
use std::time::{Duration, Instant};
use tracing::{debug, info, warn};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FallbackState {
Primary,
Fallback,
Recovering,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FailureKind {
Transient,
RateLimit,
}
#[derive(Debug, Clone)]
pub(crate) struct FallbackEntry {
name: String,
available: bool,
attempt_count: usize,
max_fail_count: usize,
}
impl FallbackEntry {
pub(crate) fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
available: true,
attempt_count: 0,
max_fail_count: 2,
}
}
#[must_use]
pub(crate) fn with_max_fail_count(mut self, max_fail_count: usize) -> Self {
let new_max = max_fail_count.max(1);
if self.failed() && self.attempt_count < new_max {
self.attempt_count = new_max;
}
self.max_fail_count = new_max;
self
}
pub(crate) fn failed(&self) -> bool {
!self.available || self.attempt_count >= self.max_fail_count
}
pub(crate) fn set_available(&mut self, available: bool) {
self.available = available;
}
pub(crate) fn record_attempt(&mut self) {
self.attempt_count = self.attempt_count.saturating_add(1);
}
pub(crate) fn clear_attempts(&mut self) {
self.attempt_count = 0;
}
}
#[derive(Debug, Clone)]
pub struct FallbackConfig {
pub trip_threshold: usize,
pub recovery_timeout: Duration,
pub recovery_successes_needed: usize,
pub max_fail_count: usize,
}
impl Default for FallbackConfig {
fn default() -> Self {
Self {
trip_threshold: 3,
recovery_timeout: Duration::from_mins(1),
recovery_successes_needed: 2,
max_fail_count: 2,
}
}
}
struct BreakerState {
state: FallbackState,
consecutive_failures: usize,
primary_success_count: usize,
fallback_activated: bool,
original_model: Option<String>,
fallback_models: Vec<FallbackEntry>,
active_fallback: Option<String>,
fallback_switched_at: Option<Instant>,
}
impl Default for BreakerState {
fn default() -> Self {
Self {
state: FallbackState::Primary,
consecutive_failures: 0,
primary_success_count: 0,
fallback_activated: false,
original_model: None,
fallback_models: Vec::new(),
active_fallback: None,
fallback_switched_at: None,
}
}
}
pub struct FallbackManager {
config: FallbackConfig,
state: Mutex<BreakerState>,
}
impl FallbackManager {
#[must_use]
pub fn new(fallback_threshold: usize, primary_resume_threshold: usize) -> Self {
Self {
config: FallbackConfig {
trip_threshold: fallback_threshold,
recovery_successes_needed: primary_resume_threshold,
..FallbackConfig::default()
},
state: Mutex::new(BreakerState::default()),
}
}
#[must_use]
pub fn with_config(mut self, config: FallbackConfig) -> Self {
self.config = config;
self
}
#[must_use]
pub fn config(&self) -> &FallbackConfig {
&self.config
}
#[must_use]
pub fn new_with_config(config: FallbackConfig) -> Self {
Self {
config,
state: Mutex::new(BreakerState::default()),
}
}
#[must_use]
pub fn new_with_fallback(original_model: String, fallback_threshold: usize) -> Self {
let mgr = Self::new(fallback_threshold, 2);
if let Ok(mut state) = mgr.state.lock() {
state.original_model = Some(original_model);
Self::transition_to_fallback_impl(&mut state);
}
mgr
}
pub fn for_model(primary_model: impl Into<String>) -> Self {
let mgr = Self::new(3, 2);
if let Ok(mut state) = mgr.state.lock() {
state.original_model = Some(primary_model.into());
}
mgr
}
pub fn state(&self) -> FallbackState {
self.state
.lock()
.map_or(FallbackState::Primary, |s| s.state)
}
pub fn is_using_fallback(&self) -> bool {
matches!(self.state(), FallbackState::Fallback)
}
pub fn is_fallback_active(&self) -> bool {
self.state.lock().is_ok_and(|s| s.fallback_activated)
}
pub fn consecutive_failures(&self) -> usize {
self.state.lock().map_or(0, |s| s.consecutive_failures)
}
pub fn original_model(&self) -> Option<String> {
self.state
.lock()
.ok()
.and_then(|s| s.original_model.clone())
}
pub fn set_original_model(&self, model: String) {
if let Ok(mut state) = self.state.lock() {
state.original_model = Some(model);
}
}
pub fn fallback_switched_at(&self) -> Option<Instant> {
self.state.lock().ok().and_then(|s| s.fallback_switched_at)
}
pub fn active_model(&self) -> Option<String> {
let Ok(state) = self.state.lock() else {
return None;
};
match state.state {
FallbackState::Primary | FallbackState::Recovering => state.original_model.clone(),
FallbackState::Fallback => state
.active_fallback
.clone()
.or_else(|| state.original_model.clone()),
}
}
pub fn set_fallback_model(&self, model: impl Into<String>) {
if let Ok(mut state) = self.state.lock() {
state.fallback_models.clear();
state
.fallback_models
.push(FallbackEntry::new(model).with_max_fail_count(self.config.max_fail_count));
Self::recompute_active_fallback_impl(&mut state);
}
}
pub fn fallback_model(&self) -> Option<String> {
self.state
.lock()
.ok()
.and_then(|s| s.active_fallback.clone())
}
pub fn fallback_models(&self) -> Vec<String> {
self.state
.lock()
.ok()
.map(|s| s.fallback_models.iter().map(|e| e.name.clone()).collect())
.unwrap_or_default()
}
pub fn add_fallback_model(&self, model: impl Into<String>) {
if let Ok(mut state) = self.state.lock() {
state
.fallback_models
.push(FallbackEntry::new(model).with_max_fail_count(self.config.max_fail_count));
Self::recompute_active_fallback_impl(&mut state);
}
}
pub fn insert_fallback_model(&self, index: usize, model: impl Into<String>) {
if let Ok(mut state) = self.state.lock() {
let entry = FallbackEntry::new(model).with_max_fail_count(self.config.max_fail_count);
if index >= state.fallback_models.len() {
state.fallback_models.push(entry);
} else {
state.fallback_models.insert(index, entry);
}
Self::recompute_active_fallback_impl(&mut state);
}
}
pub fn remove_fallback_model(&self, model: &str) -> bool {
let mut removed = false;
if let Ok(mut state) = self.state.lock() {
if let Some(pos) = state.fallback_models.iter().position(|x| x.name == model) {
state.fallback_models.remove(pos);
removed = true;
}
if removed {
Self::recompute_active_fallback_impl(&mut state);
}
}
removed
}
pub fn set_fallback_models(&self, models: Vec<String>) {
let max_fc = self.config.max_fail_count;
if let Ok(mut state) = self.state.lock() {
state.fallback_models = models
.into_iter()
.map(|name| FallbackEntry::new(name).with_max_fail_count(max_fc))
.collect();
Self::recompute_active_fallback_impl(&mut state);
}
}
pub fn mark_fallback_failed(&self, model: &str) -> bool {
let mut found = false;
if let Ok(mut state) = self.state.lock() {
if let Some(entry) = state.fallback_models.iter_mut().find(|e| e.name == model) {
entry.record_attempt();
found = true;
}
if found {
Self::recompute_active_fallback_impl(&mut state);
}
}
found
}
pub fn clear_fallback_failed(&self, model: &str) -> bool {
let mut found = false;
if let Ok(mut state) = self.state.lock() {
if let Some(entry) = state.fallback_models.iter_mut().find(|e| e.name == model) {
entry.clear_attempts();
found = true;
}
if found {
Self::recompute_active_fallback_impl(&mut state);
}
}
found
}
pub fn clear_all_fallback_failed(&self) {
if let Ok(mut state) = self.state.lock() {
for entry in &mut state.fallback_models {
entry.clear_attempts();
}
Self::recompute_active_fallback_impl(&mut state);
}
}
pub fn failed_fallbacks(&self) -> Vec<String> {
self.state
.lock()
.ok()
.map(|s| {
s.fallback_models
.iter()
.filter(|e| e.failed())
.map(|e| e.name.clone())
.collect()
})
.unwrap_or_default()
}
pub fn available_fallbacks(&self) -> Vec<String> {
self.state
.lock()
.ok()
.map(|s| {
s.fallback_models
.iter()
.filter(|e| !e.failed())
.map(|e| e.name.clone())
.collect()
})
.unwrap_or_default()
}
pub fn set_fallback_available(&self, model: &str, available: bool) -> bool {
let mut found = false;
if let Ok(mut state) = self.state.lock() {
if let Some(entry) = state.fallback_models.iter_mut().find(|e| e.name == model) {
entry.set_available(available);
found = true;
}
if found {
Self::recompute_active_fallback_impl(&mut state);
}
}
found
}
pub fn record_failure(&self, kind: FailureKind) -> bool {
let Ok(mut state) = self.state.lock() else {
return false;
};
match state.state {
FallbackState::Primary => {
state.consecutive_failures = state.consecutive_failures.saturating_add(1);
if state.consecutive_failures >= self.config.trip_threshold {
warn!(
consecutive_failures = state.consecutive_failures,
threshold = self.config.trip_threshold,
"Fallback threshold reached"
);
Self::transition_to_fallback_impl(&mut state);
true
} else {
false
}
}
FallbackState::Fallback => {
let fb_name = state
.active_fallback
.clone()
.unwrap_or_else(|| "unknown".to_string());
warn!(
"Fallback model \"{fb_name}\" also experiencing failures; consider calling mark_fallback_failed(\"{fb_name}\") to skip it"
);
false
}
FallbackState::Recovering => match kind {
FailureKind::RateLimit => {
warn!(
"Primary model rate-limited during recovery test, re-tripping to fallback"
);
Self::transition_to_fallback_impl(&mut state);
false
}
FailureKind::Transient => {
warn!(
consecutive_failures = state.consecutive_failures,
"Transient failure recorded during recovery; staying half-open"
);
false
}
},
}
}
pub fn reset_failure_counter(&self) {
if let Ok(mut state) = self.state.lock() {
state.consecutive_failures = 0;
}
}
pub fn record_success(&self) {
let Ok(mut state) = self.state.lock() else {
return;
};
match state.state {
FallbackState::Primary => {
state.consecutive_failures = 0;
}
FallbackState::Fallback => {}
FallbackState::Recovering => {
state.primary_success_count = state.primary_success_count.saturating_add(1);
debug!(
successes = state.primary_success_count,
threshold = self.config.recovery_successes_needed,
"Primary model success during recovery test"
);
if state.primary_success_count >= self.config.recovery_successes_needed {
Self::transition_to_primary_impl(&mut state);
}
}
}
}
pub fn should_try_resume_primary(&self, min_fallback_duration: Duration) -> bool {
if self.state() != FallbackState::Fallback {
return false;
}
if let Some(switched_at) = self.fallback_switched_at() {
switched_at.elapsed() >= min_fallback_duration
} else {
false
}
}
pub fn transition_to_fallback(&self) {
if let Ok(mut state) = self.state.lock() {
Self::transition_to_fallback_impl(&mut state);
}
}
fn transition_to_fallback_impl(state: &mut BreakerState) {
state.state = FallbackState::Fallback;
state.fallback_activated = true;
state.fallback_switched_at = Some(Instant::now());
state.primary_success_count = 0;
info!("Circuit breaker: transitioned to Fallback state");
}
pub fn transition_to_recovering(&self) {
if let Ok(mut state) = self.state.lock() {
Self::transition_to_recovering_impl(&mut state);
}
}
fn transition_to_recovering_impl(state: &mut BreakerState) {
state.state = FallbackState::Recovering;
state.primary_success_count = 0;
info!("Circuit breaker: transitioned to Recovering state (testing primary)");
}
pub fn transition_to_primary(&self) {
if let Ok(mut state) = self.state.lock() {
Self::transition_to_primary_impl(&mut state);
}
}
fn transition_to_primary_impl(state: &mut BreakerState) {
state.state = FallbackState::Primary;
state.fallback_switched_at = None;
state.primary_success_count = 0;
state.consecutive_failures = 0;
state.fallback_activated = false;
for entry in &mut state.fallback_models {
entry.clear_attempts();
}
Self::recompute_active_fallback_impl(state);
info!("Circuit breaker: transitioned to Primary state (primary model recovered)");
}
pub fn reset(&self) {
if let Ok(mut state) = self.state.lock() {
Self::transition_to_primary_impl(&mut state);
}
}
fn recompute_active_fallback_impl(state: &mut BreakerState) {
state.active_fallback = state
.fallback_models
.iter()
.find(|e| !e.failed())
.map(|e| e.name.clone());
}
}
impl Default for FallbackManager {
fn default() -> Self {
Self::new(3, 2)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_initial_state() {
let mgr = FallbackManager::new(3, 2);
assert_eq!(mgr.state(), FallbackState::Primary);
assert!(!mgr.is_using_fallback());
assert!(!mgr.is_fallback_active());
assert_eq!(mgr.consecutive_failures(), 0);
}
#[test]
fn test_failure_threshold() {
let mgr = FallbackManager::new(3, 2);
assert!(!mgr.record_failure(FailureKind::Transient)); assert!(!mgr.record_failure(FailureKind::Transient)); assert!(mgr.record_failure(FailureKind::Transient)); }
#[test]
fn test_model_failure_triggers_fallback() {
let mgr = FallbackManager::new(3, 2);
assert!(!mgr.record_failure(FailureKind::Transient)); assert!(!mgr.record_failure(FailureKind::Transient)); assert!(mgr.record_failure(FailureKind::Transient)); assert_eq!(mgr.state(), FallbackState::Fallback);
}
#[test]
fn test_recovery() {
let mgr = FallbackManager::new(3, 2);
for _ in 0..3 {
mgr.record_failure(FailureKind::Transient);
}
assert_eq!(mgr.state(), FallbackState::Fallback);
mgr.transition_to_recovering();
assert_eq!(mgr.state(), FallbackState::Recovering);
mgr.record_success(); mgr.record_success(); assert_eq!(mgr.state(), FallbackState::Primary);
}
#[test]
fn test_recovery_failure_goes_back_to_fallback() {
let mgr = FallbackManager::new(3, 2);
for _ in 0..3 {
mgr.record_failure(FailureKind::Transient);
}
mgr.transition_to_recovering();
mgr.record_failure(FailureKind::RateLimit); assert_eq!(mgr.state(), FallbackState::Fallback);
}
#[test]
fn test_should_try_resume_primary() {
let mgr = FallbackManager::new(3, 2);
assert!(!mgr.should_try_resume_primary(Duration::from_secs(10)));
for _ in 0..3 {
mgr.record_failure(FailureKind::Transient);
}
assert!(!mgr.should_try_resume_primary(Duration::from_hours(1)));
assert!(mgr.should_try_resume_primary(Duration::from_secs(0)));
}
#[test]
fn test_new_with_fallback() {
let mgr = FallbackManager::new_with_fallback("llm-70b".into(), 3);
assert!(mgr.is_fallback_active());
assert!(mgr.is_using_fallback());
assert_eq!(mgr.original_model(), Some("llm-70b".to_string()));
}
#[test]
fn test_reset() {
let mgr = FallbackManager::new(3, 2);
for _ in 0..3 {
mgr.record_failure(FailureKind::Transient);
}
assert_eq!(mgr.state(), FallbackState::Fallback);
mgr.reset();
assert_eq!(mgr.state(), FallbackState::Primary);
assert!(!mgr.is_fallback_active());
assert_eq!(mgr.consecutive_failures(), 0);
}
#[test]
fn test_api_failure_does_not_retrip() {
let mgr = FallbackManager::new(3, 2);
for _ in 0..3 {
mgr.record_failure(FailureKind::Transient);
}
mgr.transition_to_fallback();
assert!(!mgr.record_failure(FailureKind::Transient));
}
#[test]
fn test_record_success_resets_on_primary() {
let mgr = FallbackManager::new(3, 2);
mgr.record_failure(FailureKind::Transient);
mgr.record_failure(FailureKind::Transient);
assert_eq!(mgr.consecutive_failures(), 2);
mgr.record_success();
assert_eq!(mgr.consecutive_failures(), 0);
}
#[test]
fn test_for_model() {
let mgr = FallbackManager::for_model("llm-70b");
assert_eq!(mgr.original_model(), Some("llm-70b".to_string()));
assert_eq!(mgr.state(), FallbackState::Primary);
}
#[test]
fn test_concurrent_access() {
use std::sync::Arc;
use std::thread;
let mgr = Arc::new(FallbackManager::new(3, 2));
let mut handles = Vec::new();
for _ in 0..10 {
let mgr = Arc::clone(&mgr);
handles.push(thread::spawn(move || {
mgr.record_failure(FailureKind::Transient);
mgr.record_success();
mgr.state();
mgr.consecutive_failures();
}));
}
for h in handles {
h.join().unwrap();
}
}
#[test]
fn test_consolidated_mutex_fields_are_consistent() {
let mgr = FallbackManager::for_model("primary-model");
mgr.add_fallback_model("fallback-model");
assert_eq!(mgr.active_model(), Some("primary-model".to_string()));
assert!(mgr.fallback_switched_at().is_none());
mgr.transition_to_fallback();
assert_eq!(mgr.active_model(), Some("fallback-model".to_string()));
assert!(
mgr.fallback_switched_at().is_some(),
"switch time should be set after transition"
);
}
#[test]
fn test_consolidated_mutex_clears_fields_together() {
let mgr = FallbackManager::for_model("primary-model");
mgr.add_fallback_model("fallback-model");
mgr.transition_to_fallback();
assert_eq!(mgr.active_model(), Some("fallback-model".to_string()));
assert!(mgr.fallback_switched_at().is_some());
mgr.transition_to_primary();
assert!(
mgr.fallback_switched_at().is_none(),
"switch time should be cleared after transition to primary"
);
}
#[test]
fn test_consolidated_mutex_reset_clears_all() {
let mgr = FallbackManager::for_model("primary-model");
mgr.add_fallback_model("fallback-model");
mgr.record_failure(FailureKind::Transient);
mgr.transition_to_fallback();
assert!(mgr.consecutive_failures() > 0);
assert!(mgr.fallback_switched_at().is_some());
mgr.reset();
assert_eq!(mgr.consecutive_failures(), 0);
assert!(mgr.fallback_switched_at().is_none());
}
#[test]
fn with_max_fail_count_no_padding_when_not_failed() {
let entry = FallbackEntry::new("model-a").with_max_fail_count(5);
assert!(!entry.failed());
assert_eq!(entry.attempt_count, 0);
assert_eq!(entry.max_fail_count, 5);
}
#[test]
fn with_max_fail_count_pads_already_failed_entry() {
let mut entry = FallbackEntry::new("model-b");
entry.record_attempt();
entry.record_attempt();
assert!(entry.failed());
assert_eq!(entry.attempt_count, 2);
let entry = entry.with_max_fail_count(5);
assert_eq!(entry.max_fail_count, 5);
assert_eq!(entry.attempt_count, 5);
assert!(entry.failed());
}
#[test]
fn with_max_fail_count_pads_exactly_to_new_threshold() {
let mut entry = FallbackEntry::new("model-c");
entry.record_attempt();
entry.record_attempt();
assert!(entry.failed());
let entry = entry.with_max_fail_count(3);
assert_eq!(entry.attempt_count, 3);
assert!(entry.failed());
}
#[test]
fn with_max_fail_count_no_padding_when_lowering() {
let mut entry = FallbackEntry::new("model-d");
entry.record_attempt();
entry.record_attempt();
assert!(entry.failed());
let entry = entry.with_max_fail_count(1);
assert_eq!(entry.max_fail_count, 1);
assert_eq!(entry.attempt_count, 2);
assert!(entry.failed());
}
#[test]
fn with_max_fail_count_clamps_to_minimum_one() {
let entry = FallbackEntry::new("model-e").with_max_fail_count(0);
assert_eq!(entry.max_fail_count, 1);
}
}