foxhunt/common/src/error_recovery.rs

//! Consolidated error recovery and retry strategies for HFT systems
//!
//! This module provides unified retry strategies, circuit breakers, and error recovery
//! patterns used across all Foxhunt services for consistent error handling.

use crate::error::{CommonError, ErrorCategory, RetryStrategy, ErrorSeverity};
use serde::{Deserialize, Serialize};
use std::time::{Duration, Instant};
use tokio::time::sleep;
use tracing::{error, warn, info, debug};

/// HFT-optimized retry configuration with circuit breaker support
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetryConfig {
    /// Maximum number of retry attempts
    pub max_attempts: u32,
    /// Base delay between retries
    pub base_delay: Duration,
    /// Maximum delay cap for exponential backoff
    pub max_delay: Duration,
    /// Circuit breaker failure threshold
    pub circuit_breaker_threshold: u32,
    /// Circuit breaker timeout before reset attempt
    pub circuit_breaker_timeout: Duration,
    /// Enable jitter for exponential backoff
    pub enable_jitter: bool,
    /// HFT-specific nanosecond precision delays
    pub hft_precision_mode: bool,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            base_delay: Duration::from_millis(100),
            max_delay: Duration::from_secs(30),
            circuit_breaker_threshold: 5,
            circuit_breaker_timeout: Duration::from_secs(60),
            enable_jitter: true,
            hft_precision_mode: false,
        }
    }
}

impl RetryConfig {
    /// Create HFT-optimized configuration for low-latency operations
    pub fn hft_optimized() -> Self {
        Self {
            max_attempts: 3,
            base_delay: Duration::from_micros(50), // 50µs base delay
            max_delay: Duration::from_millis(5),   // 5ms max delay
            circuit_breaker_threshold: 10,
            circuit_breaker_timeout: Duration::from_secs(30),
            enable_jitter: false, // No jitter for HFT - predictable timing
            hft_precision_mode: true,
        }
    }

    /// Create configuration for network operations
    pub fn network_optimized() -> Self {
        Self {
            max_attempts: 5,
            base_delay: Duration::from_millis(500),
            max_delay: Duration::from_secs(10),
            circuit_breaker_threshold: 3,
            circuit_breaker_timeout: Duration::from_secs(30),
            enable_jitter: true,
            hft_precision_mode: false,
        }
    }

    /// Create configuration for database operations
    pub fn database_optimized() -> Self {
        Self {
            max_attempts: 7,
            base_delay: Duration::from_millis(250),
            max_delay: Duration::from_secs(5),
            circuit_breaker_threshold: 5,
            circuit_breaker_timeout: Duration::from_secs(45),
            enable_jitter: true,
            hft_precision_mode: false,
        }
    }
}

/// Circuit breaker states for error recovery
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum CircuitBreakerState {
    /// Circuit is closed - operations proceed normally
    Closed,
    /// Circuit is open - operations fail immediately
    Open,
    /// Circuit is half-open - testing if service has recovered
    HalfOpen,
}

/// Circuit breaker for managing service failures
#[derive(Debug)]
pub struct CircuitBreaker {
    state: CircuitBreakerState,
    failure_count: u32,
    last_failure_time: Option<Instant>,
    config: RetryConfig,
}

impl CircuitBreaker {
    /// Create new circuit breaker with configuration
    pub fn new(config: RetryConfig) -> Self {
        Self {
            state: CircuitBreakerState::Closed,
            failure_count: 0,
            last_failure_time: None,
            config,
        }
    }

    /// Check if operation should be allowed
    pub fn can_proceed(&mut self) -> bool {
        match self.state {
            CircuitBreakerState::Closed => true,
            CircuitBreakerState::Open => {
                if let Some(last_failure) = self.last_failure_time {
                    if last_failure.elapsed() >= self.config.circuit_breaker_timeout {
                        info!("Circuit breaker transitioning to half-open state");
                        self.state = CircuitBreakerState::HalfOpen;
                        return true;
                    }
                }
                false
            }
            CircuitBreakerState::HalfOpen => true,
        }
    }

    /// Record successful operation
    pub fn record_success(&mut self) {
        if self.state == CircuitBreakerState::HalfOpen {
            info!("Circuit breaker closing after successful operation");
            self.state = CircuitBreakerState::Closed;
            self.failure_count = 0;
            self.last_failure_time = None;
        }
    }

    /// Record failed operation
    pub fn record_failure(&mut self) {
        self.failure_count += 1;
        self.last_failure_time = Some(Instant::now());

        match self.state {
            CircuitBreakerState::Closed => {
                if self.failure_count >= self.config.circuit_breaker_threshold {
                    warn!("Circuit breaker opening after {} failures", self.failure_count);
                    self.state = CircuitBreakerState::Open;
                }
            }
            CircuitBreakerState::HalfOpen => {
                warn!("Circuit breaker reopening after failure in half-open state");
                self.state = CircuitBreakerState::Open;
            }
            CircuitBreakerState::Open => {
                // Already open, just update timestamp
            }
        }
    }

    /// Get current state
    pub fn state(&self) -> CircuitBreakerState {
        self.state.clone()
    }

    /// Get current failure count
    pub fn failure_count(&self) -> u32 {
        self.failure_count
    }
}

/// Retry executor with circuit breaker and telemetry
pub struct RetryExecutor {
    circuit_breaker: CircuitBreaker,
    config: RetryConfig,
    operation_name: String,
}

impl RetryExecutor {
    /// Create new retry executor
    pub fn new<S: Into<String>>(operation_name: S, config: RetryConfig) -> Self {
        Self {
            circuit_breaker: CircuitBreaker::new(config.clone()),
            config,
            operation_name: operation_name.into(),
        }
    }

    /// Execute operation with retry logic and circuit breaker
    pub async fn execute<F, Fut, T>(&mut self, mut operation: F) -> Result<T, CommonError>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<T, CommonError>>,
    {
        let mut attempt = 0;
        let mut last_error = None;

        while attempt < self.config.max_attempts {
            // Check circuit breaker
            if !self.circuit_breaker.can_proceed() {
                return Err(CommonError::service_unavailable(
                    &self.operation_name,
                    "Circuit breaker is open"
                ));
            }

            attempt += 1;
            debug!("Executing {} attempt {}/{}", self.operation_name, attempt, self.config.max_attempts);

            match operation().await {
                Ok(result) => {
                    if attempt > 1 {
                        info!("Operation {} succeeded on attempt {}", self.operation_name, attempt);
                    }
                    self.circuit_breaker.record_success();
                    return Ok(result);
                }
                Err(error) => {
                    last_error = Some(error.clone());
                    self.circuit_breaker.record_failure();

                    // Check if error is retryable
                    if !error.is_retryable() {
                        warn!("Operation {} failed with non-retryable error: {}", self.operation_name, error);
                        return Err(error);
                    }

                    if attempt < self.config.max_attempts {
                        let delay = self.calculate_delay(attempt, &error);
                        warn!("Operation {} failed on attempt {}, retrying in {:?}: {}", 
                              self.operation_name, attempt, delay, error);
                        sleep(delay).await;
                    }
                }
            }
        }

        // All retries exhausted
        let final_error = last_error.unwrap_or_else(|| {
            CommonError::internal(format!("Operation {} failed after {} attempts", self.operation_name, self.config.max_attempts))
        });

        error!("Operation {} failed after {} attempts: {}", self.operation_name, self.config.max_attempts, final_error);
        Err(final_error)
    }

    /// Calculate delay for retry attempt based on error type and configuration
    fn calculate_delay(&self, attempt: u32, error: &CommonError) -> Duration {
        let base_delay = match error.retry_strategy() {
            RetryStrategy::Immediate => Duration::from_millis(0),
            RetryStrategy::Linear { base_delay_ms } => Duration::from_millis(base_delay_ms * u64::from(attempt)),
            RetryStrategy::Exponential { base_delay_ms, max_delay_ms } => {
                let delay_ms = base_delay_ms * 2_u64.pow(attempt.saturating_sub(1).min(10));
                let capped_delay = delay_ms.min(max_delay_ms);
                Duration::from_millis(capped_delay)
            }
            RetryStrategy::HftCustom { initial_delay_ns, .. } => {
                if self.config.hft_precision_mode {
                    Duration::from_nanos(initial_delay_ns * u64::from(attempt))
                } else {
                    Duration::from_micros((initial_delay_ns / 1000) * u64::from(attempt))
                }
            }
            RetryStrategy::CircuitBreaker => self.config.circuit_breaker_timeout,
            RetryStrategy::NoRetry => Duration::from_millis(0), // Should not reach here
        };

        let final_delay = base_delay.min(self.config.max_delay);

        // Add jitter for non-HFT operations to prevent thundering herd
        if self.config.enable_jitter && !self.config.hft_precision_mode {
            self.add_jitter(final_delay)
        } else {
            final_delay
        }
    }

    /// Add jitter to delay to prevent thundering herd problem
    fn add_jitter(&self, delay: Duration) -> Duration {
        use rand::Rng;
        let jitter_range = delay.as_millis() / 10; // ±10% jitter
        let jitter_offset = rand::thread_rng().gen_range(0..=jitter_range * 2);
        let jitter_delay = jitter_offset.saturating_sub(jitter_range);
        
        delay.saturating_add(Duration::from_millis(jitter_delay as u64))
    }

    /// Get circuit breaker state
    pub fn circuit_breaker_state(&self) -> CircuitBreakerState {
        self.circuit_breaker.state()
    }

    /// Get circuit breaker failure count
    pub fn circuit_breaker_failure_count(&self) -> u32 {
        self.circuit_breaker.failure_count()
    }
}

/// Error recovery policy based on error characteristics
#[derive(Debug, Clone)]
pub struct ErrorRecoveryPolicy {
    /// Default retry configuration
    pub default_config: RetryConfig,
    /// Category-specific configurations
    pub category_configs: std::collections::HashMap<ErrorCategory, RetryConfig>,
    /// Severity-specific overrides
    pub severity_overrides: std::collections::HashMap<ErrorSeverity, RetryConfig>,
}

impl Default for ErrorRecoveryPolicy {
    fn default() -> Self {
        let mut category_configs = std::collections::HashMap::new();
        category_configs.insert(ErrorCategory::Network, RetryConfig::network_optimized());
        category_configs.insert(ErrorCategory::Database, RetryConfig::database_optimized());
        category_configs.insert(ErrorCategory::Trading, RetryConfig::hft_optimized());
        category_configs.insert(ErrorCategory::Risk, RetryConfig {
            max_attempts: 1, // Risk errors should generally not be retried
            ..RetryConfig::default()
        });

        let mut severity_overrides = std::collections::HashMap::new();
        severity_overrides.insert(ErrorSeverity::Critical, RetryConfig {
            max_attempts: 1, // Critical errors should not be retried
            ..RetryConfig::default()
        });

        Self {
            default_config: RetryConfig::default(),
            category_configs,
            severity_overrides,
        }
    }
}

impl ErrorRecoveryPolicy {
    /// Get retry configuration for a specific error
    pub fn get_config_for_error(&self, error: &CommonError) -> RetryConfig {
        // Check severity overrides first
        if let Some(config) = self.severity_overrides.get(&error.severity()) {
            return config.clone();
        }

        // Check category-specific configuration
        if let Some(config) = self.category_configs.get(&error.category()) {
            return config.clone();
        }

        // Fall back to default
        self.default_config.clone()
    }

    /// Create retry executor for a specific error type
    pub fn create_executor<S: Into<String>>(&self, operation_name: S, error: &CommonError) -> RetryExecutor {
        let config = self.get_config_for_error(error);
        RetryExecutor::new(operation_name, config)
    }
}

/// Convenience function to execute operation with automatic retry based on error characteristics
pub async fn retry_with_policy<F, Fut, T>(
    operation_name: &str,
    policy: &ErrorRecoveryPolicy,
    sample_error: &CommonError,
    operation: F,
) -> Result<T, CommonError>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = Result<T, CommonError>>,
{
    let mut executor = policy.create_executor(operation_name, sample_error);
    executor.execute(operation).await
}

/// Macro for easy retry execution with automatic policy detection
#[macro_export]
macro_rules! retry_operation {
    ($name:expr, $operation:expr) => {{
        use $crate::error_recovery::{ErrorRecoveryPolicy, retry_with_policy};
        
        let policy = ErrorRecoveryPolicy::default();
        
        // Execute once to get error type for policy detection
        let sample_result = $operation().await;
        match sample_result {
            Ok(result) => Ok(result),
            Err(sample_error) => {
                retry_with_policy($name, &policy, &sample_error, $operation).await
            }
        }
    }};
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::test;

    #[test]
    async fn test_circuit_breaker_transitions() {
        let config = RetryConfig {
            circuit_breaker_threshold: 2,
            circuit_breaker_timeout: Duration::from_millis(100),
            ..RetryConfig::default()
        };
        
        let mut circuit_breaker = CircuitBreaker::new(config);
        
        // Should start closed
        assert_eq!(circuit_breaker.state(), CircuitBreakerState::Closed);
        assert!(circuit_breaker.can_proceed());
        
        // Record failures to open circuit
        circuit_breaker.record_failure();
        assert_eq!(circuit_breaker.state(), CircuitBreakerState::Closed);
        
        circuit_breaker.record_failure();
        assert_eq!(circuit_breaker.state(), CircuitBreakerState::Open);
        assert!(!circuit_breaker.can_proceed());
        
        // Wait for timeout and transition to half-open
        tokio::time::sleep(Duration::from_millis(150)).await;
        assert!(circuit_breaker.can_proceed());
        assert_eq!(circuit_breaker.state(), CircuitBreakerState::HalfOpen);
        
        // Record success to close circuit
        circuit_breaker.record_success();
        assert_eq!(circuit_breaker.state(), CircuitBreakerState::Closed);
    }

    #[test]
    async fn test_retry_executor_success() {
        let config = RetryConfig::default();
        let mut executor = RetryExecutor::new("test_operation", config);
        
        let mut attempt_count = 0;
        let result = executor.execute(|| async {
            attempt_count += 1;
            if attempt_count == 2 {
                Ok("success")
            } else {
                Err(CommonError::network("Connection failed"))
            }
        }).await;
        
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "success");
        assert_eq!(attempt_count, 2);
    }

    #[test]
    async fn test_retry_executor_non_retryable_error() {
        let config = RetryConfig::default();
        let mut executor = RetryExecutor::new("test_operation", config);
        
        let mut attempt_count = 0;
        let result = executor.execute(|| async {
            attempt_count += 1;
            Err(CommonError::authentication("Invalid token"))
        }).await;
        
        assert!(result.is_err());
        assert_eq!(attempt_count, 1); // Should not retry authentication errors
    }

    #[test]
    async fn test_hft_retry_config() {
        let config = RetryConfig::hft_optimized();
        assert!(config.hft_precision_mode);
        assert_eq!(config.base_delay, Duration::from_micros(50));
        assert_eq!(config.max_delay, Duration::from_millis(5));
        assert!(!config.enable_jitter); // No jitter for HFT
    }

    #[test]
    fn test_error_recovery_policy() {
        let policy = ErrorRecoveryPolicy::default();
        
        let network_error = CommonError::network("Connection failed");
        let config = policy.get_config_for_error(&network_error);
        assert_eq!(config.max_attempts, 5); // Network optimized
        
        let trading_error = CommonError::trading("Order rejected");
        let config = policy.get_config_for_error(&trading_error);
        assert!(config.hft_precision_mode); // HFT optimized
        
        let critical_error = CommonError::authentication("Invalid token");
        let config = policy.get_config_for_error(&critical_error);
        assert_eq!(config.max_attempts, 1); // Critical errors should not retry
    }

    #[test]
    fn test_retry_delay_calculation() {
        let config = RetryConfig::default();
        let mut executor = RetryExecutor::new("test", config);
        
        let network_error = CommonError::network("Connection failed");
        let delay1 = executor.calculate_delay(1, &network_error);
        let delay2 = executor.calculate_delay(2, &network_error);
        
        // Exponential backoff should increase delay
        assert!(delay2 > delay1);
    }
}