#![allow(non_snake_case)]
use crate::core::RiResult;
use std::collections::HashMap as FxHashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub enum RiCircuitBreakerState {
Closed,
Open,
HalfOpen,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiCircuitBreakerConfig {
pub failure_threshold: u32,
pub success_threshold: u32,
pub timeout_seconds: u64,
pub monitoring_period_seconds: u64,
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiCircuitBreakerConfig {
#[new]
fn py_new() -> Self {
Self::default()
}
#[staticmethod]
fn py_new_with_values(failure_threshold: u32, success_threshold: u32, timeout_seconds: u64, monitoring_period_seconds: u64) -> Self {
Self {
failure_threshold,
success_threshold,
timeout_seconds,
monitoring_period_seconds,
}
}
fn get_failure_threshold(&self) -> u32 {
self.failure_threshold
}
fn set_failure_threshold(&mut self, value: u32) {
self.failure_threshold = value;
}
fn get_success_threshold(&self) -> u32 {
self.success_threshold
}
fn set_success_threshold(&mut self, value: u32) {
self.success_threshold = value;
}
fn get_timeout_seconds(&self) -> u64 {
self.timeout_seconds
}
fn set_timeout_seconds(&mut self, value: u64) {
self.timeout_seconds = value;
}
fn get_monitoring_period_seconds(&self) -> u64 {
self.monitoring_period_seconds
}
fn set_monitoring_period_seconds(&mut self, value: u64) {
self.monitoring_period_seconds = value;
}
}
impl Default for RiCircuitBreakerConfig {
fn default() -> Self {
Self {
failure_threshold: 5,
success_threshold: 3,
timeout_seconds: 60,
monitoring_period_seconds: 30,
}
}
}
#[derive(Debug)]
struct CircuitBreakerStats {
state: RwLock<RiCircuitBreakerState>,
failure_count: AtomicUsize,
success_count: AtomicUsize,
last_failure_time: RwLock<Option<Instant>>,
last_state_change: RwLock<Instant>,
consecutive_failures: AtomicUsize,
consecutive_successes: AtomicUsize,
}
impl CircuitBreakerStats {
#[allow(dead_code)]
fn new() -> Self {
Self {
state: RwLock::new(RiCircuitBreakerState::Closed),
failure_count: AtomicUsize::new(0),
success_count: AtomicUsize::new(0),
last_failure_time: RwLock::new(None),
last_state_change: RwLock::new(Instant::now()),
consecutive_failures: AtomicUsize::new(0),
consecutive_successes: AtomicUsize::new(0),
}
}
async fn record_success(&self, config: &RiCircuitBreakerConfig) {
self.success_count.fetch_add(1, Ordering::Relaxed);
self.consecutive_failures.store(0, Ordering::Relaxed);
self.consecutive_successes.fetch_add(1, Ordering::Relaxed);
let state = self.state.read().await;
match *state {
RiCircuitBreakerState::HalfOpen => {
let successes = self.consecutive_successes.load(Ordering::Relaxed);
if successes >= config.success_threshold as usize {
drop(state);
let mut state_write = self.state.write().await;
*state_write = RiCircuitBreakerState::Closed;
self.consecutive_successes.store(0, Ordering::Relaxed);
self.failure_count.store(0, Ordering::Relaxed);
*self.last_state_change.write().await = Instant::now();
}
}
RiCircuitBreakerState::Open => {
}
RiCircuitBreakerState::Closed => {
self.consecutive_failures.store(0, Ordering::Relaxed);
}
}
}
async fn record_failure(&self, config: &RiCircuitBreakerConfig) {
self.failure_count.fetch_add(1, Ordering::Relaxed);
self.consecutive_successes.store(0, Ordering::Relaxed);
self.consecutive_failures.fetch_add(1, Ordering::Relaxed);
*self.last_failure_time.write().await = Some(Instant::now());
let state = self.state.read().await;
match *state {
RiCircuitBreakerState::Closed => {
let failures = self.consecutive_failures.load(Ordering::Relaxed);
if failures >= config.failure_threshold as usize {
drop(state);
let mut state_write = self.state.write().await;
*state_write = RiCircuitBreakerState::Open;
self.consecutive_failures.store(0, Ordering::Relaxed);
*self.last_state_change.write().await = Instant::now();
}
}
RiCircuitBreakerState::HalfOpen => {
drop(state);
let mut state_write = self.state.write().await;
*state_write = RiCircuitBreakerState::Open;
*self.last_state_change.write().await = Instant::now();
}
RiCircuitBreakerState::Open => {
}
}
}
async fn should_attempt_reset(&self, config: &RiCircuitBreakerConfig) -> bool {
let state = self.state.read().await;
if let RiCircuitBreakerState::Open = *state {
let last_change = *self.last_state_change.read().await;
Instant::now().duration_since(last_change) >= Duration::from_secs(config.timeout_seconds)
} else {
false
}
}
async fn transition_to_half_open(&self) {
let mut state = self.state.write().await;
*state = RiCircuitBreakerState::HalfOpen;
*self.last_state_change.write().await = Instant::now();
}
async fn get_state(&self) -> RiCircuitBreakerState {
self.state.read().await.clone()
}
#[allow(dead_code)]
fn get_stats(&self) -> RiCircuitBreakerMetrics {
let state_str = match self.state.blocking_read().clone() {
RiCircuitBreakerState::Closed => "Closed",
RiCircuitBreakerState::Open => "Open",
RiCircuitBreakerState::HalfOpen => "HalfOpen",
};
RiCircuitBreakerMetrics {
state: state_str.to_string(),
failure_count: self.failure_count.load(Ordering::Relaxed),
success_count: self.success_count.load(Ordering::Relaxed),
consecutive_failures: self.consecutive_failures.load(Ordering::Relaxed),
consecutive_successes: self.consecutive_successes.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiCircuitBreakerMetrics {
pub state: String,
pub failure_count: usize,
pub success_count: usize,
pub consecutive_failures: usize,
pub consecutive_successes: usize,
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiCircuitBreakerMetrics {
#[new]
fn py_new(state: String, failure_count: usize, success_count: usize, consecutive_failures: usize, consecutive_successes: usize) -> Self {
Self {
state,
failure_count,
success_count,
consecutive_failures,
consecutive_successes,
}
}
fn get_state(&self) -> &str {
&self.state
}
fn get_failure_count(&self) -> usize {
self.failure_count
}
fn get_success_count(&self) -> usize {
self.success_count
}
fn get_consecutive_failures(&self) -> usize {
self.consecutive_failures
}
fn get_consecutive_successes(&self) -> usize {
self.consecutive_successes
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiCircuitBreaker {
config: RiCircuitBreakerConfig,
stats: Arc<CircuitBreakerStats>,
}
impl RiCircuitBreaker {
pub fn new(config: RiCircuitBreakerConfig) -> Self {
Self {
config,
stats: Arc::new(CircuitBreakerStats::new()),
}
}
pub fn allow_request(&self) -> bool {
let state = futures::executor::block_on(async {
let state = self.stats.state.read().await;
state.clone()
});
match state {
RiCircuitBreakerState::Closed => true,
RiCircuitBreakerState::Open => {
let last_change = futures::executor::block_on(async {
let guard = self.stats.last_state_change.read().await;
*guard
});
if last_change.elapsed() >= Duration::from_secs(self.config.timeout_seconds) {
futures::executor::block_on(async {
let mut state = self.stats.state.write().await;
*state = RiCircuitBreakerState::HalfOpen;
*self.stats.last_state_change.write().await = Instant::now();
});
true
} else {
false
}
}
RiCircuitBreakerState::HalfOpen => true,
}
}
pub fn record_success(&self) {
futures::executor::block_on(async {
self.stats.record_success(&self.config).await;
});
}
pub fn record_failure(&self) {
futures::executor::block_on(async {
self.stats.record_failure(&self.config).await;
});
}
pub async fn execute<F, R>(&self, operation: F) -> RiResult<R>
where
F: std::future::Future<Output = RiResult<R>>,
{
if !self.allow_request() {
return Err(crate::core::RiError::ServiceMesh("Circuit breaker is open".to_string()));
}
match operation.await {
Ok(result) => {
self.record_success();
Ok(result)
}
Err(error) => {
self.record_failure();
Err(error)
}
}
}
pub fn get_state(&self) -> RiCircuitBreakerState {
futures::executor::block_on(async {
self.stats.get_state().await
})
}
pub fn get_stats(&self) -> RiCircuitBreakerMetrics {
let state_str = match futures::executor::block_on(async {
let state = self.stats.state.read().await;
state.clone()
}) {
RiCircuitBreakerState::Closed => "Closed",
RiCircuitBreakerState::Open => "Open",
RiCircuitBreakerState::HalfOpen => "HalfOpen",
};
RiCircuitBreakerMetrics {
state: state_str.to_string(),
failure_count: self.stats.failure_count.load(Ordering::Relaxed),
success_count: self.stats.success_count.load(Ordering::Relaxed),
consecutive_failures: self.stats.consecutive_failures.load(Ordering::Relaxed),
consecutive_successes: self.stats.consecutive_successes.load(Ordering::Relaxed),
}
}
pub fn get_config(&self) -> RiCircuitBreakerConfig {
self.config.clone()
}
pub fn reset(&self) {
futures::executor::block_on(async move {
let mut state = self.stats.state.write().await;
*state = RiCircuitBreakerState::Closed;
self.stats.failure_count.store(0, Ordering::Relaxed);
self.stats.success_count.store(0, Ordering::Relaxed);
self.stats.consecutive_failures.store(0, Ordering::Relaxed);
self.stats.consecutive_successes.store(0, Ordering::Relaxed);
*self.stats.last_state_change.write().await = Instant::now();
});
}
pub fn force_open(&self) {
futures::executor::block_on(async move {
let mut state = self.stats.state.write().await;
*state = RiCircuitBreakerState::Open;
*self.stats.last_state_change.write().await = Instant::now();
});
}
pub fn force_close(&self) {
futures::executor::block_on(async move {
let mut state = self.stats.state.write().await;
*state = RiCircuitBreakerState::Closed;
*self.stats.last_state_change.write().await = Instant::now();
});
}
pub fn failure_rate(&self) -> f64 {
let failures = self.stats.failure_count.load(Ordering::Relaxed);
let successes = self.stats.success_count.load(Ordering::Relaxed);
let total = failures + successes;
if total == 0 {
0.0
} else {
failures as f64 / total as f64
}
}
pub fn success_rate(&self) -> f64 {
let failures = self.stats.failure_count.load(Ordering::Relaxed);
let successes = self.stats.success_count.load(Ordering::Relaxed);
let total = failures + successes;
if total == 0 {
1.0
} else {
successes as f64 / total as f64
}
}
pub fn total_requests(&self) -> usize {
self.stats.failure_count.load(Ordering::Relaxed) + self.stats.success_count.load(Ordering::Relaxed)
}
pub fn is_open(&self) -> bool {
let state = futures::executor::block_on(self.stats.state.read());
matches!(*state, RiCircuitBreakerState::Open)
}
pub fn is_closed(&self) -> bool {
let state = futures::executor::block_on(self.stats.state.read());
matches!(*state, RiCircuitBreakerState::Closed)
}
pub fn is_half_open(&self) -> bool {
let state = futures::executor::block_on(self.stats.state.read());
matches!(*state, RiCircuitBreakerState::HalfOpen)
}
}
pub struct RiAdvancedCircuitBreaker {
config: RiCircuitBreakerConfig,
stats_by_error: RwLock<FxHashMap<String, Arc<CircuitBreakerStats>>>,
default_stats: Arc<CircuitBreakerStats>,
}
impl RiAdvancedCircuitBreaker {
pub fn new(config: RiCircuitBreakerConfig) -> Self {
Self {
config,
stats_by_error: RwLock::new(FxHashMap::default()),
default_stats: Arc::new(CircuitBreakerStats::new()),
}
}
async fn get_stats_for_error(&self, error_type: Option<&str>) -> Arc<CircuitBreakerStats> {
match error_type {
Some(error_type) => {
let mut stats_map = self.stats_by_error.write().await;
stats_map.entry(error_type.to_string())
.or_insert_with(|| Arc::new(CircuitBreakerStats::new()))
.clone()
}
None => self.default_stats.clone(),
}
}
pub async fn record_success_with_type(&self, error_type: Option<&str>) {
let stats = self.get_stats_for_error(error_type).await;
stats.record_success(&self.config).await;
}
pub async fn record_failure_with_type(&self, error_type: Option<&str>) {
let stats = self.get_stats_for_error(error_type).await;
stats.record_failure(&self.config).await;
}
pub async fn allow_request_for_type(&self, error_type: Option<&str>) -> bool {
let stats = self.get_stats_for_error(error_type).await;
let state = stats.get_state().await;
match state {
RiCircuitBreakerState::Closed => true,
RiCircuitBreakerState::Open => {
if stats.should_attempt_reset(&self.config).await {
stats.transition_to_half_open().await;
true
} else {
false
}
}
RiCircuitBreakerState::HalfOpen => true,
}
}
}