//! 供应商路由器模块 //! //! 负责选择和管理代理目标供应商,实现智能故障转移 use crate::database::Database; use crate::error::AppError; use crate::provider::Provider; use crate::proxy::circuit_breaker::{AllowResult, CircuitBreaker, CircuitBreakerConfig}; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::RwLock; /// 供应商路由器 pub struct ProviderRouter { /// 数据库连接 db: Arc, /// 熔断器管理器 - key 格式: "app_type:provider_id" circuit_breakers: Arc>>>, } impl ProviderRouter { /// 创建新的供应商路由器 pub fn new(db: Arc) -> Self { Self { db, circuit_breakers: Arc::new(RwLock::new(HashMap::new())), } } /// 选择可用的供应商(支持故障转移) /// /// 返回按优先级排序的可用供应商列表: /// - 故障转移关闭时:仅返回当前供应商 /// - 故障转移开启时:完全按照故障转移队列顺序返回,忽略当前供应商设置 pub async fn select_providers(&self, app_type: &str) -> Result, AppError> { let mut result = Vec::new(); let mut total_providers = 0usize; let mut circuit_open_count = 0usize; // 检查该应用的自动故障转移开关是否开启(从 proxy_config 表读取) let auto_failover_enabled = match self.db.get_proxy_config_for_app(app_type).await { Ok(config) => { let enabled = config.auto_failover_enabled; log::info!("[{app_type}] Failover enabled from proxy_config: {enabled}"); enabled } Err(e) => { log::error!( "[{app_type}] Failed to read proxy_config for auto_failover_enabled: {e}, defaulting to disabled" ); false } }; if auto_failover_enabled { // 故障转移开启:使用 in_failover_queue 标记的供应商,按 sort_index 排序 let failover_providers = self.db.get_failover_providers(app_type)?; total_providers = failover_providers.len(); log::debug!("[{app_type}] Found {total_providers} failover queue provider(s)"); log::info!( "[{app_type}] Failover enabled, using queue order ({total_providers} items)" ); for provider in failover_providers { // 检查熔断器状态 let circuit_key = format!("{}:{}", app_type, provider.id); let breaker = self.get_or_create_circuit_breaker(&circuit_key).await; let state = breaker.get_state().await; if breaker.is_available().await { log::debug!( "[{}] Queue provider available: {} ({}) (state: {:?})", app_type, provider.name, provider.id, state ); log::info!( "[{}] Queue provider available: {} ({}) at sort_index {:?}", app_type, provider.name, provider.id, provider.sort_index ); result.push(provider); } else { circuit_open_count += 1; log::debug!( "[{}] Queue provider {} circuit breaker open (state: {:?}), skipping", app_type, provider.name, state ); } } } else { // 故障转移关闭:仅使用当前供应商,跳过熔断器检查 // 原因:单 Provider 场景下,熔断器打开会导致所有请求失败,用户体验差 log::info!("[{app_type}] Failover disabled, using current provider only (circuit breaker bypassed)"); if let Some(current_id) = self.db.get_current_provider(app_type)? { if let Some(current) = self.db.get_provider_by_id(¤t_id, app_type)? { log::info!( "[{}] Current provider: {} ({})", app_type, current.name, current.id ); total_providers = 1; result.push(current); } else { log::debug!( "[{app_type}] Current provider id {current_id} not found in database" ); } } else { log::debug!("[{app_type}] No current provider configured"); } } if result.is_empty() { // 区分两种情况:全部熔断 vs 未配置供应商 if total_providers > 0 && circuit_open_count == total_providers { log::warn!("[{app_type}] 所有 {total_providers} 个供应商均已熔断,无可用渠道"); return Err(AppError::AllProvidersCircuitOpen); } else { log::warn!("[{app_type}] 未配置供应商或故障转移队列为空"); return Err(AppError::NoProvidersConfigured); } } log::info!( "[{}] Provider chain: {} provider(s) available", app_type, result.len() ); Ok(result) } /// 请求执行前获取熔断器“放行许可” /// /// - Closed:直接放行 /// - Open:超时到达后切到 HalfOpen 并放行一次探测 /// - HalfOpen:按限流规则放行探测 /// /// 注意:调用方必须在请求结束后通过 `record_result()` 释放 HalfOpen 名额, /// 否则会导致该 Provider 长时间无法进入探测状态。 pub async fn allow_provider_request(&self, provider_id: &str, app_type: &str) -> AllowResult { let circuit_key = format!("{app_type}:{provider_id}"); let breaker = self.get_or_create_circuit_breaker(&circuit_key).await; breaker.allow_request().await } /// 记录供应商请求结果 pub async fn record_result( &self, provider_id: &str, app_type: &str, used_half_open_permit: bool, success: bool, error_msg: Option, ) -> Result<(), AppError> { // 1. 按应用独立获取熔断器配置(用于更新健康状态和判断是否禁用) let failure_threshold = match self.db.get_proxy_config_for_app(app_type).await { Ok(app_config) => app_config.circuit_failure_threshold, Err(e) => { log::warn!( "Failed to load circuit config for {app_type}, using default threshold: {e}" ); 5 // 默认值 } }; // 2. 更新熔断器状态 let circuit_key = format!("{app_type}:{provider_id}"); let breaker = self.get_or_create_circuit_breaker(&circuit_key).await; if success { breaker.record_success(used_half_open_permit).await; log::debug!("Provider {provider_id} request succeeded"); } else { breaker.record_failure(used_half_open_permit).await; log::warn!( "Provider {} request failed: {}", provider_id, error_msg.as_deref().unwrap_or("Unknown error") ); } // 3. 更新数据库健康状态(使用配置的阈值) self.db .update_provider_health_with_threshold( provider_id, app_type, success, error_msg.clone(), failure_threshold, ) .await?; Ok(()) } /// 重置熔断器(手动恢复) pub async fn reset_circuit_breaker(&self, circuit_key: &str) { let breakers = self.circuit_breakers.read().await; if let Some(breaker) = breakers.get(circuit_key) { log::info!("Manually resetting circuit breaker for {circuit_key}"); breaker.reset().await; } } /// 重置指定供应商的熔断器 pub async fn reset_provider_breaker(&self, provider_id: &str, app_type: &str) { let circuit_key = format!("{app_type}:{provider_id}"); self.reset_circuit_breaker(&circuit_key).await; } /// 更新所有熔断器的配置(热更新) /// /// 当用户在 UI 中修改熔断器配置后调用此方法, /// 所有现有的熔断器会立即使用新配置 pub async fn update_all_configs(&self, config: CircuitBreakerConfig) { let breakers = self.circuit_breakers.read().await; let count = breakers.len(); for breaker in breakers.values() { breaker.update_config(config.clone()).await; } log::info!("已更新 {count} 个熔断器的配置"); } /// 获取熔断器状态 #[allow(dead_code)] pub async fn get_circuit_breaker_stats( &self, provider_id: &str, app_type: &str, ) -> Option { let circuit_key = format!("{app_type}:{provider_id}"); let breakers = self.circuit_breakers.read().await; if let Some(breaker) = breakers.get(&circuit_key) { Some(breaker.get_stats().await) } else { None } } /// 获取或创建熔断器 async fn get_or_create_circuit_breaker(&self, key: &str) -> Arc { // 先尝试读锁获取 { let breakers = self.circuit_breakers.read().await; if let Some(breaker) = breakers.get(key) { return breaker.clone(); } } // 如果不存在,获取写锁创建 let mut breakers = self.circuit_breakers.write().await; // 双重检查,防止竞争条件 if let Some(breaker) = breakers.get(key) { return breaker.clone(); } // 从 key 中提取 app_type (格式: "app_type:provider_id") let app_type = key.split(':').next().unwrap_or("claude"); // 按应用独立读取熔断器配置 let config = match self.db.get_proxy_config_for_app(app_type).await { Ok(app_config) => { log::debug!( "Loading circuit breaker config for {key} (app={app_type}): \ failure_threshold={}, success_threshold={}, timeout={}s", app_config.circuit_failure_threshold, app_config.circuit_success_threshold, app_config.circuit_timeout_seconds ); crate::proxy::circuit_breaker::CircuitBreakerConfig { failure_threshold: app_config.circuit_failure_threshold, success_threshold: app_config.circuit_success_threshold, timeout_seconds: app_config.circuit_timeout_seconds as u64, error_rate_threshold: app_config.circuit_error_rate_threshold, min_requests: app_config.circuit_min_requests, } } Err(e) => { log::warn!( "Failed to load circuit breaker config for {key} (app={app_type}): {e}, using default" ); crate::proxy::circuit_breaker::CircuitBreakerConfig::default() } }; log::debug!("Creating new circuit breaker for {key} with config: {config:?}"); let breaker = Arc::new(CircuitBreaker::new(config)); breakers.insert(key.to_string(), breaker.clone()); breaker } } #[cfg(test)] mod tests { use super::*; use crate::database::Database; use serde_json::json; #[tokio::test] async fn test_provider_router_creation() { let db = Arc::new(Database::memory().unwrap()); let router = ProviderRouter::new(db); let breaker = router.get_or_create_circuit_breaker("claude:test").await; assert!(breaker.allow_request().await.allowed); } #[tokio::test] async fn test_failover_disabled_uses_current_provider() { let db = Arc::new(Database::memory().unwrap()); let provider_a = Provider::with_id("a".to_string(), "Provider A".to_string(), json!({}), None); let provider_b = Provider::with_id("b".to_string(), "Provider B".to_string(), json!({}), None); db.save_provider("claude", &provider_a).unwrap(); db.save_provider("claude", &provider_b).unwrap(); db.set_current_provider("claude", "a").unwrap(); db.add_to_failover_queue("claude", "b").unwrap(); let router = ProviderRouter::new(db.clone()); let providers = router.select_providers("claude").await.unwrap(); assert_eq!(providers.len(), 1); assert_eq!(providers[0].id, "a"); } #[tokio::test] async fn test_failover_enabled_uses_queue_order() { let db = Arc::new(Database::memory().unwrap()); // 设置 sort_index 来控制顺序:b=1, a=2 let mut provider_a = Provider::with_id("a".to_string(), "Provider A".to_string(), json!({}), None); provider_a.sort_index = Some(2); let mut provider_b = Provider::with_id("b".to_string(), "Provider B".to_string(), json!({}), None); provider_b.sort_index = Some(1); db.save_provider("claude", &provider_a).unwrap(); db.save_provider("claude", &provider_b).unwrap(); db.set_current_provider("claude", "a").unwrap(); db.add_to_failover_queue("claude", "b").unwrap(); db.add_to_failover_queue("claude", "a").unwrap(); // 启用自动故障转移(使用新的 proxy_config API) let mut config = db.get_proxy_config_for_app("claude").await.unwrap(); config.auto_failover_enabled = true; db.update_proxy_config_for_app(config).await.unwrap(); let router = ProviderRouter::new(db.clone()); let providers = router.select_providers("claude").await.unwrap(); assert_eq!(providers.len(), 2); // 按 sort_index 排序:b(1) 在前,a(2) 在后 assert_eq!(providers[0].id, "b"); assert_eq!(providers[1].id, "a"); } #[tokio::test] async fn test_select_providers_does_not_consume_half_open_permit() { let db = Arc::new(Database::memory().unwrap()); db.update_circuit_breaker_config(&CircuitBreakerConfig { failure_threshold: 1, timeout_seconds: 0, ..Default::default() }) .await .unwrap(); let provider_a = Provider::with_id("a".to_string(), "Provider A".to_string(), json!({}), None); let provider_b = Provider::with_id("b".to_string(), "Provider B".to_string(), json!({}), None); db.save_provider("claude", &provider_a).unwrap(); db.save_provider("claude", &provider_b).unwrap(); db.add_to_failover_queue("claude", "a").unwrap(); db.add_to_failover_queue("claude", "b").unwrap(); // 启用自动故障转移(使用新的 proxy_config API) let mut config = db.get_proxy_config_for_app("claude").await.unwrap(); config.auto_failover_enabled = true; db.update_proxy_config_for_app(config).await.unwrap(); let router = ProviderRouter::new(db.clone()); router .record_result("b", "claude", false, false, Some("fail".to_string())) .await .unwrap(); let providers = router.select_providers("claude").await.unwrap(); assert_eq!(providers.len(), 2); assert!(router.allow_provider_request("b", "claude").await.allowed); } }