pub const DEFAULT_FAILURE_THRESHOLD: u64 = 5;
pub const DEFAULT_SUCCESS_THRESHOLD: u64 = 2;
pub const DEFAULT_TIMEOUT_SECS: u64 = 60;
pub const DEFAULT_HALF_OPEN_MAX_CALLS: u64 = 3;
use crate::error::{CircuitBreakerStats, CircuitState, FlowGuardError};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{info, trace, warn};
#[cfg(feature = "circuit-breaker")]
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
pub failure_threshold: u64,
pub success_threshold: u64,
pub timeout: Duration,
pub half_open_max_calls: u64,
}
impl Default for CircuitBreakerConfig {
fn default() -> Self {
Self {
failure_threshold: DEFAULT_FAILURE_THRESHOLD,
success_threshold: DEFAULT_SUCCESS_THRESHOLD,
timeout: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
half_open_max_calls: DEFAULT_HALF_OPEN_MAX_CALLS,
}
}
}
impl CircuitBreakerConfig {
pub fn new(failure_threshold: u64, success_threshold: u64, timeout: Duration) -> Self {
Self {
failure_threshold,
success_threshold,
timeout,
half_open_max_calls: 3,
}
}
pub fn half_open_max_calls(mut self, max_calls: u64) -> Self {
self.half_open_max_calls = max_calls;
self
}
}
#[cfg(feature = "circuit-breaker")]
pub struct CircuitBreaker {
state: Arc<RwLock<CircuitState>>,
failure_count: Arc<AtomicU64>,
success_count: Arc<AtomicU64>,
total_calls: Arc<AtomicU64>,
last_failure_time: Arc<RwLock<Option<Instant>>>,
last_state_change: Arc<RwLock<Option<Instant>>>,
half_open_calls: Arc<AtomicU64>,
config: CircuitBreakerConfig,
}
impl CircuitBreaker {
pub fn new(config: CircuitBreakerConfig) -> Self {
info!(
"创建熔断器: failure_threshold={}, success_threshold={}, timeout={:?}",
config.failure_threshold, config.success_threshold, config.timeout
);
Self {
state: Arc::new(RwLock::new(CircuitState::Closed)),
failure_count: Arc::new(AtomicU64::new(0)),
success_count: Arc::new(AtomicU64::new(0)),
total_calls: Arc::new(AtomicU64::new(0)),
last_failure_time: Arc::new(RwLock::new(None)),
last_state_change: Arc::new(RwLock::new(Some(Instant::now()))),
half_open_calls: Arc::new(AtomicU64::new(0)),
config,
}
}
}
impl Default for CircuitBreaker {
fn default() -> Self {
Self::new(CircuitBreakerConfig::default())
}
}
impl CircuitBreaker {
pub async fn execute<F, Fut, T>(&self, operation: F) -> Result<T, FlowGuardError>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<T, FlowGuardError>>,
{
self.total_calls.fetch_add(1, Ordering::Relaxed);
let state = self.state.read().await;
match *state {
CircuitState::Open => {
let last_failure = self.last_failure_time.read().await;
if let Some(last_failure) = *last_failure {
if last_failure.elapsed() >= self.config.timeout {
drop(state);
self.transition_to_half_open().await;
} else {
drop(state);
warn!("熔断器打开,拒绝请求");
return Err(FlowGuardError::LimitError(
"熔断器打开,请求被拒绝".to_string(),
));
}
}
}
CircuitState::HalfOpen => {
let calls = self.half_open_calls.load(Ordering::Relaxed);
if calls >= self.config.half_open_max_calls {
drop(state);
warn!("半开状态调用次数已达上限,拒绝请求");
return Err(FlowGuardError::LimitError(
"半开状态调用次数已达上限".to_string(),
));
}
self.half_open_calls.fetch_add(1, Ordering::Relaxed);
drop(state);
}
CircuitState::Closed => {
drop(state);
}
}
let result = operation().await;
match result {
Ok(value) => {
self.on_success().await;
Ok(value)
}
Err(e) => {
self.on_failure().await;
Err(e)
}
}
}
async fn on_success(&self) {
let state = self.state.read().await;
match *state {
CircuitState::Closed => {
self.failure_count.store(0, Ordering::Relaxed);
self.success_count.fetch_add(1, Ordering::Relaxed);
trace!("操作成功(关闭状态)");
}
CircuitState::HalfOpen => {
let success_count = self.success_count.fetch_add(1, Ordering::Relaxed) + 1;
if success_count >= self.config.success_threshold {
drop(state);
self.transition_to_closed().await;
} else {
trace!(
"操作成功(半开状态): {}/{}",
success_count,
self.config.success_threshold
);
}
}
CircuitState::Open => {
warn!("熔断器打开状态下收到成功响应");
}
}
}
async fn on_failure(&self) {
let state = self.state.read().await;
match *state {
CircuitState::Closed => {
let failure_count = self.failure_count.fetch_add(1, Ordering::Relaxed) + 1;
*self.last_failure_time.write().await = Some(Instant::now());
if failure_count >= self.config.failure_threshold {
drop(state);
self.transition_to_open().await;
} else {
trace!(
"操作失败(关闭状态): {}/{}",
failure_count,
self.config.failure_threshold
);
}
}
CircuitState::HalfOpen => {
drop(state);
self.transition_to_open().await;
}
CircuitState::Open => {
warn!("熔断器打开状态下收到失败响应");
}
}
}
async fn transition_to_open(&self) {
let old_state = *self.state.read().await;
if old_state != CircuitState::Open {
*self.state.write().await = CircuitState::Open;
*self.last_state_change.write().await = Some(Instant::now());
self.success_count.store(0, Ordering::Relaxed);
self.half_open_calls.store(0, Ordering::Relaxed);
warn!(
"熔断器状态变更: {:?} -> Open (failure_count={})",
old_state,
self.failure_count.load(Ordering::Relaxed)
);
}
}
async fn transition_to_half_open(&self) {
let old_state = *self.state.read().await;
if old_state != CircuitState::HalfOpen {
*self.state.write().await = CircuitState::HalfOpen;
*self.last_state_change.write().await = Some(Instant::now());
self.success_count.store(0, Ordering::Relaxed);
self.half_open_calls.store(1, Ordering::Relaxed);
info!("熔断器状态变更: {:?} -> HalfOpen", old_state);
}
}
async fn transition_to_closed(&self) {
let old_state = *self.state.read().await;
if old_state != CircuitState::Closed {
*self.state.write().await = CircuitState::Closed;
*self.last_state_change.write().await = Some(Instant::now());
self.failure_count.store(0, Ordering::Relaxed);
self.success_count.store(0, Ordering::Relaxed);
self.half_open_calls.store(0, Ordering::Relaxed);
info!("熔断器状态变更: {:?} -> Closed", old_state);
}
}
pub async fn is_open(&self) -> bool {
let state = self.state.read().await;
*state == CircuitState::Open
}
pub async fn is_half_open(&self) -> bool {
let state = self.state.read().await;
*state == CircuitState::HalfOpen
}
pub async fn is_closed(&self) -> bool {
let state = self.state.read().await;
*state == CircuitState::Closed
}
pub async fn get_state(&self) -> CircuitState {
*self.state.read().await
}
pub async fn reset(&self) {
info!("重置熔断器");
*self.state.write().await = CircuitState::Closed;
self.failure_count.store(0, Ordering::Relaxed);
self.success_count.store(0, Ordering::Relaxed);
self.total_calls.store(0, Ordering::Relaxed);
*self.last_failure_time.write().await = None;
*self.last_state_change.write().await = Some(Instant::now());
self.half_open_calls.store(0, Ordering::Relaxed);
}
pub async fn get_stats(&self) -> CircuitBreakerStats {
let state = *self.state.read().await;
let last_failure = self.last_failure_time.read().await;
let last_state_change = self.last_state_change.read().await;
CircuitBreakerStats {
state,
failure_count: self.failure_count.load(Ordering::Relaxed),
success_count: self.success_count.load(Ordering::Relaxed),
total_calls: self.total_calls.load(Ordering::Relaxed),
last_failure_time: last_failure.and_then(|t| {
let elapsed = t.elapsed();
let duration = chrono::Duration::from_std(elapsed).ok()?;
Some(chrono::Utc::now() - duration)
}),
last_state_change: last_state_change.and_then(|t| {
let elapsed = t.elapsed();
let duration = chrono::Duration::from_std(elapsed).ok()?;
Some(chrono::Utc::now() - duration)
}),
}
}
pub fn config(&self) -> &CircuitBreakerConfig {
&self.config
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_circuit_breaker_config_default() {
let config = CircuitBreakerConfig::default();
assert_eq!(config.failure_threshold, 5);
assert_eq!(config.success_threshold, 2);
assert_eq!(config.timeout, Duration::from_secs(60));
assert_eq!(config.half_open_max_calls, 3);
}
#[test]
fn test_circuit_breaker_config_new() {
let config = CircuitBreakerConfig::new(10, 3, Duration::from_secs(120));
assert_eq!(config.failure_threshold, 10);
assert_eq!(config.success_threshold, 3);
assert_eq!(config.timeout, Duration::from_secs(120));
}
#[test]
fn test_circuit_breaker_config_builder() {
let config =
CircuitBreakerConfig::new(5, 2, Duration::from_secs(60)).half_open_max_calls(5);
assert_eq!(config.half_open_max_calls, 5);
}
#[tokio::test]
async fn test_circuit_breaker_initial_state() {
let breaker = CircuitBreaker::default();
assert!(breaker.is_closed().await);
assert!(!breaker.is_open().await);
assert!(!breaker.is_half_open().await);
let stats = breaker.get_stats().await;
assert_eq!(stats.state, CircuitState::Closed);
assert_eq!(stats.failure_count, 0);
assert_eq!(stats.success_count, 0);
assert_eq!(stats.total_calls, 0);
}
#[tokio::test]
async fn test_circuit_breaker_success() {
let breaker = CircuitBreaker::default();
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_ok());
let stats = breaker.get_stats().await;
assert_eq!(stats.success_count, 1);
assert_eq!(stats.failure_count, 0);
assert_eq!(stats.total_calls, 1);
assert!(breaker.is_closed().await);
}
#[tokio::test]
async fn test_circuit_breaker_failure() {
let config = CircuitBreakerConfig::new(3, 2, Duration::from_secs(60));
let breaker = CircuitBreaker::new(config);
let result = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
assert!(result.is_err());
let stats = breaker.get_stats().await;
assert_eq!(stats.failure_count, 1);
assert!(breaker.is_closed().await);
let result = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
assert!(result.is_err());
let stats = breaker.get_stats().await;
assert_eq!(stats.failure_count, 2);
assert!(breaker.is_closed().await);
let result = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
assert!(result.is_err());
let stats = breaker.get_stats().await;
assert_eq!(stats.failure_count, 3);
assert!(breaker.is_open().await);
}
#[tokio::test]
async fn test_circuit_breaker_open_rejects_requests() {
let config = CircuitBreakerConfig::new(2, 2, Duration::from_secs(60));
let breaker = CircuitBreaker::new(config);
for _ in 0..2 {
let _ = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
}
assert!(breaker.is_open().await);
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("熔断器打开"));
}
#[tokio::test]
async fn test_circuit_breaker_half_open_recovery() {
let config = CircuitBreakerConfig::new(2, 2, Duration::from_millis(100));
let breaker = CircuitBreaker::new(config);
for _ in 0..2 {
let _ = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
}
assert!(breaker.is_open().await);
tokio::time::sleep(Duration::from_millis(150)).await;
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_ok());
assert!(breaker.is_half_open().await);
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_ok());
assert!(breaker.is_closed().await);
}
#[tokio::test]
async fn test_circuit_breaker_half_open_failure() {
let config = CircuitBreakerConfig::new(2, 2, Duration::from_millis(100));
let breaker = CircuitBreaker::new(config);
for _ in 0..2 {
let _ = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
}
assert!(breaker.is_open().await);
tokio::time::sleep(Duration::from_millis(150)).await;
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_ok());
assert!(breaker.is_half_open().await);
let result = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
assert!(result.is_err());
assert!(breaker.is_open().await);
}
#[tokio::test]
async fn test_circuit_breaker_reset() {
let config = CircuitBreakerConfig::new(2, 2, Duration::from_secs(60));
let breaker = CircuitBreaker::new(config);
for _ in 0..2 {
let _ = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
}
assert!(breaker.is_open().await);
breaker.reset().await;
assert!(breaker.is_closed().await);
let stats = breaker.get_stats().await;
assert_eq!(stats.failure_count, 0);
assert_eq!(stats.success_count, 0);
assert_eq!(stats.total_calls, 0);
}
#[tokio::test]
async fn test_circuit_breaker_get_state() {
let breaker = CircuitBreaker::default();
assert_eq!(breaker.get_state().await, CircuitState::Closed);
}
#[tokio::test]
async fn test_circuit_breaker_get_stats() {
let breaker = CircuitBreaker::default();
let _ = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
let stats = breaker.get_stats().await;
assert_eq!(stats.state, CircuitState::Closed);
assert_eq!(stats.success_count, 1);
assert_eq!(stats.total_calls, 1);
assert!(stats.last_state_change.is_some());
}
#[tokio::test]
async fn test_circuit_breaker_half_open_max_calls() {
let config =
CircuitBreakerConfig::new(2, 3, Duration::from_millis(100)).half_open_max_calls(2);
let breaker = CircuitBreaker::new(config);
for _ in 0..2 {
let _ = breaker
.execute(|| async {
Err::<(), FlowGuardError>(FlowGuardError::LimitError("test error".to_string()))
})
.await;
}
assert!(breaker.is_open().await);
tokio::time::sleep(Duration::from_millis(150)).await;
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_ok());
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_ok());
let result = breaker
.execute(|| async { Ok::<(), FlowGuardError>(()) })
.await;
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("半开状态调用次数已达上限"));
}
#[tokio::test]
async fn test_circuit_breaker_config() {
let config = CircuitBreakerConfig::new(10, 5, Duration::from_secs(30));
let breaker = CircuitBreaker::new(config);
let breaker_config = breaker.config();
assert_eq!(breaker_config.failure_threshold, 10);
assert_eq!(breaker_config.success_threshold, 5);
assert_eq!(breaker_config.timeout, Duration::from_secs(30));
}
}