Files
rustytorch/crates/production/rtx-serving-api/src/resilience.rs
T
2026-03-04 00:08:42 +00:00

769 lines
24 KiB
Rust

//! Resilience patterns for production API
//!
//! This module provides fault tolerance mechanisms including:
//! - Circuit breaker pattern for failing fast when services are unhealthy
//! - Retry logic with exponential backoff and jitter
//! - Timeout handling
//! - Bulkhead pattern for resource isolation
use parking_lot::RwLock;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::time::sleep;
use tracing::{debug, error, info, warn};
/// Circuit breaker states
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
/// Circuit is closed, requests flow through normally
Closed,
/// Circuit is open, requests fail immediately
Open,
/// Circuit is testing if the service has recovered
HalfOpen,
}
impl std::fmt::Display for CircuitState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Closed => write!(f, "closed"),
Self::Open => write!(f, "open"),
Self::HalfOpen => write!(f, "half-open"),
}
}
}
/// Circuit breaker configuration
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
/// Number of failures before opening the circuit
pub failure_threshold: u32,
/// Number of successes in half-open state before closing
pub success_threshold: u32,
/// Duration to wait before transitioning from open to half-open
pub reset_timeout: Duration,
/// Time window for counting failures
pub failure_window: Duration,
}
impl Default for CircuitBreakerConfig {
fn default() -> Self {
Self {
failure_threshold: 5,
success_threshold: 3,
reset_timeout: Duration::from_secs(30),
failure_window: Duration::from_secs(60),
}
}
}
/// Circuit breaker for protecting downstream services
#[derive(Debug)]
pub struct CircuitBreaker {
config: CircuitBreakerConfig,
state: RwLock<CircuitState>,
failure_count: AtomicU32,
success_count: AtomicU32,
last_failure_time: RwLock<Option<Instant>>,
last_state_change: RwLock<Instant>,
total_requests: AtomicU64,
total_failures: AtomicU64,
name: String,
}
impl CircuitBreaker {
/// Create a new circuit breaker with the given configuration
pub fn new(name: impl Into<String>, config: CircuitBreakerConfig) -> Self {
Self {
config,
state: RwLock::new(CircuitState::Closed),
failure_count: AtomicU32::new(0),
success_count: AtomicU32::new(0),
last_failure_time: RwLock::new(None),
last_state_change: RwLock::new(Instant::now()),
total_requests: AtomicU64::new(0),
total_failures: AtomicU64::new(0),
name: name.into(),
}
}
/// Create a circuit breaker with default configuration
pub fn with_defaults(name: impl Into<String>) -> Self {
Self::new(name, CircuitBreakerConfig::default())
}
/// Get current circuit state
pub fn state(&self) -> CircuitState {
*self.state.read()
}
/// Check if a request is allowed to proceed
pub fn allow_request(&self) -> bool {
self.total_requests.fetch_add(1, Ordering::Relaxed);
let mut state = self.state.write();
match *state {
CircuitState::Closed => true,
CircuitState::Open => {
// Check if reset timeout has passed
let last_change = *self.last_state_change.read();
if last_change.elapsed() >= self.config.reset_timeout {
info!(
circuit = %self.name,
"Circuit transitioning from open to half-open"
);
*state = CircuitState::HalfOpen;
*self.last_state_change.write() = Instant::now();
self.success_count.store(0, Ordering::Relaxed);
true
} else {
let remaining_ms = self
.config
.reset_timeout
.checked_sub(last_change.elapsed())
.map_or(0, |d| d.as_millis());
debug!(
circuit = %self.name,
remaining_ms = %remaining_ms,
"Circuit is open, rejecting request"
);
false
}
}
CircuitState::HalfOpen => {
// Allow limited requests in half-open state
true
}
}
}
/// Record a successful request
pub fn record_success(&self) {
let mut state = self.state.write();
match *state {
CircuitState::HalfOpen => {
let successes = self.success_count.fetch_add(1, Ordering::Relaxed) + 1;
if successes >= self.config.success_threshold {
info!(
circuit = %self.name,
successes = successes,
"Circuit transitioning from half-open to closed"
);
*state = CircuitState::Closed;
*self.last_state_change.write() = Instant::now();
self.failure_count.store(0, Ordering::Relaxed);
self.success_count.store(0, Ordering::Relaxed);
}
}
CircuitState::Closed => {
// Reset failure count on success in closed state
self.failure_count.store(0, Ordering::Relaxed);
}
_ => {}
}
}
/// Record a failed request
pub fn record_failure(&self) {
self.total_failures.fetch_add(1, Ordering::Relaxed);
*self.last_failure_time.write() = Some(Instant::now());
let mut state = self.state.write();
match *state {
CircuitState::Closed => {
let failures = self.failure_count.fetch_add(1, Ordering::Relaxed) + 1;
if failures >= self.config.failure_threshold {
warn!(
circuit = %self.name,
failures = failures,
threshold = self.config.failure_threshold,
"Circuit transitioning from closed to open"
);
*state = CircuitState::Open;
*self.last_state_change.write() = Instant::now();
}
}
CircuitState::HalfOpen => {
warn!(
circuit = %self.name,
"Circuit transitioning from half-open to open after failure"
);
*state = CircuitState::Open;
*self.last_state_change.write() = Instant::now();
self.success_count.store(0, Ordering::Relaxed);
}
CircuitState::Open => {
// Already open, nothing to do
}
}
}
/// Get circuit breaker statistics
pub fn stats(&self) -> CircuitBreakerStats {
CircuitBreakerStats {
name: self.name.clone(),
state: self.state(),
total_requests: self.total_requests.load(Ordering::Relaxed),
total_failures: self.total_failures.load(Ordering::Relaxed),
current_failures: self.failure_count.load(Ordering::Relaxed),
current_successes: self.success_count.load(Ordering::Relaxed),
time_in_state: self.last_state_change.read().elapsed(),
}
}
/// Execute a function with circuit breaker protection
pub async fn call<F, T, E>(&self, f: F) -> Result<T, CircuitBreakerError<E>>
where
F: std::future::Future<Output = Result<T, E>>,
{
if !self.allow_request() {
return Err(CircuitBreakerError::CircuitOpen);
}
match f.await {
Ok(result) => {
self.record_success();
Ok(result)
}
Err(e) => {
self.record_failure();
Err(CircuitBreakerError::ServiceError(e))
}
}
}
}
/// Circuit breaker statistics
#[derive(Debug, Clone)]
pub struct CircuitBreakerStats {
/// Circuit breaker name
pub name: String,
/// Current state
pub state: CircuitState,
/// Total requests processed
pub total_requests: u64,
/// Total failures recorded
pub total_failures: u64,
/// Current failure count
pub current_failures: u32,
/// Current success count (in half-open state)
pub current_successes: u32,
/// Time spent in current state
pub time_in_state: Duration,
}
/// Circuit breaker error types
#[derive(Debug)]
pub enum CircuitBreakerError<E> {
/// Circuit is open, request rejected
CircuitOpen,
/// Underlying service error
ServiceError(E),
}
impl<E: std::fmt::Display> std::fmt::Display for CircuitBreakerError<E> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::CircuitOpen => write!(f, "circuit breaker is open"),
Self::ServiceError(e) => write!(f, "service error: {e}"),
}
}
}
impl<E: std::error::Error> std::error::Error for CircuitBreakerError<E> {}
// ============================================================================
// Retry Logic with Exponential Backoff
// ============================================================================
/// Retry configuration
#[derive(Debug, Clone)]
pub struct RetryConfig {
/// Maximum number of retry attempts
pub max_retries: u32,
/// Initial delay between retries
pub initial_delay: Duration,
/// Maximum delay between retries
pub max_delay: Duration,
/// Multiplier for exponential backoff
pub multiplier: f64,
/// Whether to add random jitter
pub add_jitter: bool,
/// Maximum jitter as a fraction of delay (0.0 to 1.0)
pub jitter_factor: f64,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_retries: 3,
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(10),
multiplier: 2.0,
add_jitter: true,
jitter_factor: 0.1,
}
}
}
impl RetryConfig {
/// Create configuration for aggressive retries (more attempts, shorter delays)
#[must_use]
pub fn aggressive() -> Self {
Self {
max_retries: 5,
initial_delay: Duration::from_millis(50),
max_delay: Duration::from_secs(5),
multiplier: 1.5,
add_jitter: true,
jitter_factor: 0.2,
}
}
/// Create configuration for conservative retries (fewer attempts, longer delays)
#[must_use]
pub fn conservative() -> Self {
Self {
max_retries: 2,
initial_delay: Duration::from_millis(500),
max_delay: Duration::from_secs(30),
multiplier: 3.0,
add_jitter: true,
jitter_factor: 0.1,
}
}
/// Calculate delay for a given attempt number
#[must_use]
pub fn delay_for_attempt(&self, attempt: u32) -> Duration {
let base_delay =
self.initial_delay.as_millis() as f64 * self.multiplier.powi(attempt as i32);
let capped_delay = base_delay.min(self.max_delay.as_millis() as f64);
let final_delay = if self.add_jitter {
let jitter_range = capped_delay * self.jitter_factor;
let jitter = (rand::random::<f64>() - 0.5) * 2.0 * jitter_range;
(capped_delay + jitter).max(0.0)
} else {
capped_delay
};
Duration::from_millis(final_delay as u64)
}
}
/// Retry with exponential backoff
pub struct RetryWithBackoff {
config: RetryConfig,
}
impl RetryWithBackoff {
/// Create a new retry handler with the given configuration
#[must_use]
pub fn new(config: RetryConfig) -> Self {
Self { config }
}
/// Create with default configuration
#[must_use]
pub fn with_defaults() -> Self {
Self::new(RetryConfig::default())
}
/// Execute a function with retries
pub async fn retry<F, Fut, T, E>(&self, mut f: F) -> Result<T, RetryError<E>>
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = Result<T, E>>,
E: std::fmt::Display,
{
let mut last_error = None;
for attempt in 0..=self.config.max_retries {
match f().await {
Ok(result) => {
if attempt > 0 {
debug!(
attempt = attempt,
"Retry succeeded after {} attempts", attempt
);
}
return Ok(result);
}
Err(e) => {
if attempt < self.config.max_retries {
let delay = self.config.delay_for_attempt(attempt);
warn!(
attempt = attempt,
max_retries = self.config.max_retries,
delay_ms = delay.as_millis() as u64,
error = %e,
"Retry attempt failed, waiting before next attempt"
);
sleep(delay).await;
}
last_error = Some(e);
}
}
}
error!(
attempts = self.config.max_retries + 1,
"All retry attempts exhausted"
);
Err(RetryError::MaxRetriesExceeded {
attempts: self.config.max_retries + 1,
last_error,
})
}
/// Execute with retries and a predicate to determine if error is retryable
pub async fn retry_if<F, Fut, T, E, P>(
&self,
mut f: F,
should_retry: P,
) -> Result<T, RetryError<E>>
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = Result<T, E>>,
E: std::fmt::Display,
P: Fn(&E) -> bool,
{
let mut last_error = None;
for attempt in 0..=self.config.max_retries {
match f().await {
Ok(result) => return Ok(result),
Err(e) => {
if !should_retry(&e) {
return Err(RetryError::NonRetryableError(e));
}
if attempt < self.config.max_retries {
let delay = self.config.delay_for_attempt(attempt);
warn!(
attempt = attempt,
delay_ms = delay.as_millis() as u64,
error = %e,
"Retryable error, waiting before next attempt"
);
sleep(delay).await;
}
last_error = Some(e);
}
}
}
Err(RetryError::MaxRetriesExceeded {
attempts: self.config.max_retries + 1,
last_error,
})
}
}
/// Retry error types
#[derive(Debug)]
pub enum RetryError<E> {
/// Maximum retries exceeded
MaxRetriesExceeded {
/// Number of attempts made
attempts: u32,
/// Last error encountered
last_error: Option<E>,
},
/// Error is not retryable
NonRetryableError(E),
}
impl<E: std::fmt::Display> std::fmt::Display for RetryError<E> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MaxRetriesExceeded {
attempts,
last_error,
} => {
write!(f, "max retries ({attempts}) exceeded")?;
if let Some(e) = last_error {
write!(f, ", last error: {e}")?;
}
Ok(())
}
Self::NonRetryableError(e) => write!(f, "non-retryable error: {e}"),
}
}
}
impl<E: std::error::Error> std::error::Error for RetryError<E> {}
// ============================================================================
// Timeout Handler
// ============================================================================
/// Execute a future with a timeout
pub async fn with_timeout<F, T>(timeout: Duration, f: F) -> Result<T, TimeoutError>
where
F: std::future::Future<Output = T>,
{
tokio::time::timeout(timeout, f)
.await
.map_err(|_| TimeoutError { timeout })
}
/// Timeout error
#[derive(Debug, Clone)]
pub struct TimeoutError {
/// The timeout duration that was exceeded
pub timeout: Duration,
}
impl std::fmt::Display for TimeoutError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "operation timed out after {:?}", self.timeout)
}
}
impl std::error::Error for TimeoutError {}
// ============================================================================
// Combined Resilience Handler
// ============================================================================
/// Combined resilience handler with circuit breaker, retry, and timeout
pub struct ResilienceHandler {
circuit_breaker: Arc<CircuitBreaker>,
retry: RetryWithBackoff,
timeout: Duration,
}
impl ResilienceHandler {
/// Create a new resilience handler
pub fn new(
name: impl Into<String>,
circuit_config: CircuitBreakerConfig,
retry_config: RetryConfig,
timeout: Duration,
) -> Self {
Self {
circuit_breaker: Arc::new(CircuitBreaker::new(name, circuit_config)),
retry: RetryWithBackoff::new(retry_config),
timeout,
}
}
/// Create with default configurations
pub fn with_defaults(name: impl Into<String>) -> Self {
Self::new(
name,
CircuitBreakerConfig::default(),
RetryConfig::default(),
Duration::from_secs(30),
)
}
/// Get circuit breaker reference
#[must_use]
pub fn circuit_breaker(&self) -> &CircuitBreaker {
&self.circuit_breaker
}
/// Execute a function with full resilience protection
///
/// Returns `ResilienceError<String>` because errors are stringified for logging
/// during retry attempts.
pub async fn call<F, Fut, T, E>(&self, f: F) -> Result<T, ResilienceError<String>>
where
F: Fn() -> Fut + Clone,
Fut: std::future::Future<Output = Result<T, E>>,
E: std::fmt::Display + std::fmt::Debug,
{
// First check circuit breaker
if !self.circuit_breaker.allow_request() {
return Err(ResilienceError::CircuitOpen);
}
// Then try with retries and timeout
let cb = self.circuit_breaker.clone();
let timeout = self.timeout;
let result = self
.retry
.retry(|| {
let f = f.clone();
async move {
with_timeout(timeout, f())
.await
.map_err(|_| "timeout".to_string())?
.map_err(|e| format!("{e}"))
}
})
.await;
match result {
Ok(value) => {
cb.record_success();
Ok(value)
}
Err(RetryError::MaxRetriesExceeded { .. }) => {
cb.record_failure();
Err(ResilienceError::MaxRetriesExceeded)
}
Err(RetryError::NonRetryableError(e)) => {
cb.record_failure();
Err(ResilienceError::ServiceError(e))
}
}
}
}
/// Resilience error types
#[derive(Debug)]
pub enum ResilienceError<E> {
/// Circuit breaker is open
CircuitOpen,
/// Operation timed out
Timeout,
/// Maximum retries exceeded
MaxRetriesExceeded,
/// Underlying service error
ServiceError(E),
}
impl<E: std::fmt::Display> std::fmt::Display for ResilienceError<E> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::CircuitOpen => write!(f, "circuit breaker is open"),
Self::Timeout => write!(f, "operation timed out"),
Self::MaxRetriesExceeded => write!(f, "max retries exceeded"),
Self::ServiceError(e) => write!(f, "service error: {e}"),
}
}
}
impl<E: std::error::Error> std::error::Error for ResilienceError<E> {}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::AtomicU32;
#[test]
fn test_circuit_breaker_closed_state() {
let cb = CircuitBreaker::with_defaults("test");
assert_eq!(cb.state(), CircuitState::Closed);
assert!(cb.allow_request());
}
#[test]
fn test_circuit_breaker_opens_on_failures() {
let config = CircuitBreakerConfig {
failure_threshold: 3,
..Default::default()
};
let cb = CircuitBreaker::new("test", config);
// Record failures up to threshold
cb.record_failure();
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Closed);
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
}
#[test]
fn test_circuit_breaker_rejects_when_open() {
let config = CircuitBreakerConfig {
failure_threshold: 1,
reset_timeout: Duration::from_secs(60),
..Default::default()
};
let cb = CircuitBreaker::new("test", config);
cb.record_failure();
assert_eq!(cb.state(), CircuitState::Open);
assert!(!cb.allow_request());
}
#[test]
fn test_retry_config_delay_calculation() {
let config = RetryConfig {
initial_delay: Duration::from_millis(100),
multiplier: 2.0,
max_delay: Duration::from_secs(10),
add_jitter: false,
..Default::default()
};
assert_eq!(config.delay_for_attempt(0), Duration::from_millis(100));
assert_eq!(config.delay_for_attempt(1), Duration::from_millis(200));
assert_eq!(config.delay_for_attempt(2), Duration::from_millis(400));
}
#[test]
fn test_retry_config_max_delay() {
let config = RetryConfig {
initial_delay: Duration::from_secs(1),
multiplier: 10.0,
max_delay: Duration::from_secs(5),
add_jitter: false,
..Default::default()
};
// 1 * 10^3 = 1000s, but should be capped at 5s
assert_eq!(config.delay_for_attempt(3), Duration::from_secs(5));
}
#[tokio::test]
async fn test_retry_success() {
let retry = RetryWithBackoff::new(RetryConfig {
max_retries: 3,
initial_delay: Duration::from_millis(1),
..Default::default()
});
let counter = AtomicU32::new(0);
let result = retry
.retry(|| {
let count = counter.fetch_add(1, Ordering::Relaxed);
async move {
if count < 2 {
Err::<(), _>("fail")
} else {
Ok(())
}
}
})
.await;
assert!(result.is_ok());
assert_eq!(counter.load(Ordering::Relaxed), 3);
}
#[tokio::test]
async fn test_retry_exhausted() {
let retry = RetryWithBackoff::new(RetryConfig {
max_retries: 2,
initial_delay: Duration::from_millis(1),
..Default::default()
});
let result: Result<(), RetryError<&str>> =
retry.retry(|| async { Err::<(), _>("always fail") }).await;
assert!(matches!(result, Err(RetryError::MaxRetriesExceeded { .. })));
}
#[tokio::test]
async fn test_timeout() {
let result = with_timeout(Duration::from_millis(10), async {
sleep(Duration::from_secs(1)).await;
42
})
.await;
assert!(result.is_err());
}
}