oxigdal_cloud/

retry.rs

//! Retry logic with exponential backoff and circuit breaker pattern
//!
//! This module provides comprehensive retry mechanisms for cloud storage operations,
//! including exponential backoff, jitter, circuit breaker, retry budgets, and idempotency checks.

use std::time::Duration;

use crate::error::{CloudError, Result, RetryError};

/// Default maximum number of retry attempts
pub const DEFAULT_MAX_RETRIES: usize = 3;

/// Default initial backoff duration (100ms)
pub const DEFAULT_INITIAL_BACKOFF: Duration = Duration::from_millis(100);

/// Default maximum backoff duration (30 seconds)
pub const DEFAULT_MAX_BACKOFF: Duration = Duration::from_secs(30);

/// Default backoff multiplier
pub const DEFAULT_BACKOFF_MULTIPLIER: f64 = 2.0;

/// Retry strategy configuration
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of retry attempts
    pub max_retries: usize,
    /// Initial backoff duration
    pub initial_backoff: Duration,
    /// Maximum backoff duration
    pub max_backoff: Duration,
    /// Backoff multiplier for exponential backoff
    pub backoff_multiplier: f64,
    /// Whether to add jitter to backoff
    pub jitter: bool,
    /// Whether to enable circuit breaker
    pub circuit_breaker: bool,
    /// Circuit breaker failure threshold
    pub circuit_breaker_threshold: usize,
    /// Circuit breaker timeout duration
    pub circuit_breaker_timeout: Duration,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: DEFAULT_MAX_RETRIES,
            initial_backoff: DEFAULT_INITIAL_BACKOFF,
            max_backoff: DEFAULT_MAX_BACKOFF,
            backoff_multiplier: DEFAULT_BACKOFF_MULTIPLIER,
            jitter: true,
            circuit_breaker: true,
            circuit_breaker_threshold: 5,
            circuit_breaker_timeout: Duration::from_secs(60),
        }
    }
}

impl RetryConfig {
    /// Creates a new retry configuration
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the maximum number of retries
    #[must_use]
    pub fn with_max_retries(mut self, max_retries: usize) -> Self {
        self.max_retries = max_retries;
        self
    }

    /// Sets the initial backoff duration
    #[must_use]
    pub fn with_initial_backoff(mut self, duration: Duration) -> Self {
        self.initial_backoff = duration;
        self
    }

    /// Sets the maximum backoff duration
    #[must_use]
    pub fn with_max_backoff(mut self, duration: Duration) -> Self {
        self.max_backoff = duration;
        self
    }

    /// Sets the backoff multiplier
    #[must_use]
    pub fn with_backoff_multiplier(mut self, multiplier: f64) -> Self {
        self.backoff_multiplier = multiplier;
        self
    }

    /// Enables or disables jitter
    #[must_use]
    pub fn with_jitter(mut self, jitter: bool) -> Self {
        self.jitter = jitter;
        self
    }

    /// Enables or disables circuit breaker
    #[must_use]
    pub fn with_circuit_breaker(mut self, enabled: bool) -> Self {
        self.circuit_breaker = enabled;
        self
    }
}

/// Circuit breaker state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
    /// Circuit is closed (normal operation)
    Closed,
    /// Circuit is open (failures exceeded threshold)
    Open,
    /// Circuit is half-open (testing if service recovered)
    HalfOpen,
}

/// Circuit breaker for preventing cascading failures
#[derive(Debug)]
pub struct CircuitBreaker {
    /// Current state
    state: CircuitState,
    /// Consecutive failure count
    failure_count: usize,
    /// Failure threshold
    threshold: usize,
    /// Timeout duration
    timeout: Duration,
    /// Last failure time
    last_failure: Option<std::time::Instant>,
}

impl CircuitBreaker {
    /// Creates a new circuit breaker
    #[must_use]
    pub fn new(threshold: usize, timeout: Duration) -> Self {
        Self {
            state: CircuitState::Closed,
            failure_count: 0,
            threshold,
            timeout,
            last_failure: None,
        }
    }

    /// Checks if the circuit breaker allows the operation
    pub fn allow_request(&mut self) -> Result<()> {
        match self.state {
            CircuitState::Closed => Ok(()),
            CircuitState::Open => {
                // Check if timeout has elapsed
                if let Some(last_failure) = self.last_failure {
                    if last_failure.elapsed() >= self.timeout {
                        tracing::info!("Circuit breaker transitioning to half-open state");
                        self.state = CircuitState::HalfOpen;
                        Ok(())
                    } else {
                        Err(CloudError::Retry(RetryError::CircuitBreakerOpen {
                            message: "Circuit breaker is open".to_string(),
                        }))
                    }
                } else {
                    Ok(())
                }
            }
            CircuitState::HalfOpen => Ok(()),
        }
    }

    /// Records a successful operation
    pub fn record_success(&mut self) {
        match self.state {
            CircuitState::Closed => {
                self.failure_count = 0;
            }
            CircuitState::HalfOpen => {
                tracing::info!("Circuit breaker transitioning to closed state");
                self.state = CircuitState::Closed;
                self.failure_count = 0;
            }
            CircuitState::Open => {}
        }
    }

    /// Records a failed operation
    pub fn record_failure(&mut self) {
        self.failure_count += 1;
        self.last_failure = Some(std::time::Instant::now());

        if self.failure_count >= self.threshold && self.state != CircuitState::Open {
            tracing::warn!(
                "Circuit breaker opening after {} failures",
                self.failure_count
            );
            self.state = CircuitState::Open;
        }
    }

    /// Returns the current state of the circuit breaker
    #[must_use]
    pub fn state(&self) -> CircuitState {
        self.state
    }
}

/// Retry budget for limiting retry overhead
#[derive(Debug)]
pub struct RetryBudget {
    /// Available retry tokens
    tokens: usize,
    /// Maximum tokens
    max_tokens: usize,
    /// Token refill rate (tokens per second)
    refill_rate: f64,
    /// Last refill time
    last_refill: std::time::Instant,
}

impl RetryBudget {
    /// Creates a new retry budget
    #[must_use]
    pub fn new(max_tokens: usize, refill_rate: f64) -> Self {
        Self {
            tokens: max_tokens,
            max_tokens,
            refill_rate,
            last_refill: std::time::Instant::now(),
        }
    }

    /// Tries to consume a retry token
    pub fn try_consume(&mut self) -> Result<()> {
        self.refill();

        if self.tokens > 0 {
            self.tokens -= 1;
            Ok(())
        } else {
            Err(CloudError::Retry(RetryError::BudgetExhausted {
                message: "Retry budget exhausted".to_string(),
            }))
        }
    }

    /// Refills tokens based on elapsed time
    fn refill(&mut self) {
        let elapsed = self.last_refill.elapsed();
        let tokens_to_add = (elapsed.as_secs_f64() * self.refill_rate) as usize;

        if tokens_to_add > 0 {
            self.tokens = (self.tokens + tokens_to_add).min(self.max_tokens);
            self.last_refill = std::time::Instant::now();
        }
    }
}

/// Backoff calculator for exponential backoff with jitter
#[derive(Debug)]
pub struct Backoff {
    /// Configuration
    config: RetryConfig,
    /// Current attempt number
    attempt: usize,
    /// Per-instance jitter RNG state (LCG), seeded from process/time/stack
    /// entropy at construction so independent `Backoff` instances -- and
    /// independent process restarts -- do not replay identical jitter
    /// sequences (see `entropy_seed`).
    rng_state: u64,
}

impl Backoff {
    /// Creates a new backoff calculator
    #[must_use]
    pub fn new(config: RetryConfig) -> Self {
        Self {
            config,
            attempt: 0,
            rng_state: Self::entropy_seed(),
        }
    }

    /// Derives a per-instance, non-zero LCG seed from process/time/stack
    /// entropy.
    ///
    /// This is deliberately `std`-only (no `rand`/`getrandom` dependency) per
    /// the Pure Rust / minimal-dependency policy; it does not need to be
    /// cryptographically secure, only to avoid the two failure modes of the
    /// previous fixed-zero static seed: (1) the first jitter draw of every
    /// fresh process being exactly zero, and (2) identical jitter sequences
    /// being replayed across correlated process restarts (e.g. a fleet-wide
    /// redeploy), which defeats jitter's anti-thundering-herd purpose.
    fn entropy_seed() -> u64 {
        use std::sync::atomic::{AtomicU64, Ordering};
        use std::time::{SystemTime, UNIX_EPOCH};

        // A per-process monotonic disambiguator. This is *not* the seed by
        // itself (a predictable counter would be a poor RNG seed on its
        // own) -- it is mixed with time/pid/stack-address entropy below so
        // that two `Backoff`s constructed back-to-back in the same process
        // (faster than the clock's resolution can distinguish) still get
        // different seeds, while cross-process correlation is still broken
        // by the time/pid components varying across restarts.
        static CALL_COUNTER: AtomicU64 = AtomicU64::new(0);
        let counter = CALL_COUNTER.fetch_add(1, Ordering::Relaxed);

        let nanos = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0);
        let pid = u64::from(std::process::id());
        // Mix in a stack address for additional entropy across processes
        // that happen to share a PID namespace view or clock.
        let stack_marker: u8 = 0;
        let addr = std::ptr::addr_of!(stack_marker) as u64;

        // splitmix64-style mixing so the combined bits are well distributed
        // rather than just linearly XORed.
        let mixed = nanos
            ^ pid.wrapping_mul(0x9E37_79B9_7F4A_7C15)
            ^ addr.rotate_left(17)
            ^ counter.wrapping_mul(0xD6E8_FEB8_6659_FD93);
        let mut z = mixed.wrapping_add(0x9E37_79B9_7F4A_7C15);
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^= z >> 31;

        // Avoid a zero seed: it would reproduce the original bug's
        // degenerate first LCG output (`0 * 1664525 + 1013904223`, whose
        // upper 32 bits happen to be zero).
        if z == 0 { 1 } else { z }
    }

    /// Advances the per-instance LCG and returns a value in `[0.0, 1.0)`.
    fn next_random(&mut self) -> f64 {
        let next = self
            .rng_state
            .wrapping_mul(1664525)
            .wrapping_add(1013904223);
        self.rng_state = next;
        (next >> 32) as f64 / u32::MAX as f64
    }

    /// Calculates the next backoff duration
    #[must_use]
    pub fn next(&mut self) -> Duration {
        let base = self.config.initial_backoff.as_secs_f64().mul_add(
            self.config.backoff_multiplier.powi(self.attempt as i32),
            0.0,
        );

        let backoff = if self.config.jitter {
            // Add jitter (0% to 50% of base)
            let jitter_factor = 1.0 + (self.next_random() * 0.5);
            base * jitter_factor
        } else {
            base
        };

        self.attempt += 1;

        Duration::from_secs_f64(backoff.min(self.config.max_backoff.as_secs_f64()))
    }

    /// Resets the backoff state
    ///
    /// Note: this resets the retry attempt counter but deliberately leaves
    /// the RNG state untouched, so a reused `Backoff` keeps drawing fresh
    /// jitter values rather than replaying its sequence from the start.
    pub fn reset(&mut self) {
        self.attempt = 0;
    }
}

/// Determines if an error is retryable
#[must_use]
pub fn is_retryable(error: &CloudError) -> bool {
    match error {
        CloudError::Timeout { .. } => true,
        CloudError::RateLimitExceeded { .. } => true,
        CloudError::Http(http_error) => match http_error {
            crate::error::HttpError::Network { .. } => true,
            crate::error::HttpError::Status { status, .. } => {
                // Retry on 5xx errors and some 4xx errors
                matches!(
                    *status,
                    500 | 502 | 503 | 504 | 408 | 429 // Server errors and rate limiting
                )
            }
            _ => false,
        },
        CloudError::S3(s3_error) => match s3_error {
            crate::error::S3Error::Sdk { .. } => true,
            _ => false,
        },
        CloudError::Azure(azure_error) => match azure_error {
            crate::error::AzureError::Sdk { .. } => true,
            _ => false,
        },
        CloudError::Gcs(gcs_error) => match gcs_error {
            crate::error::GcsError::Sdk { .. } => true,
            _ => false,
        },
        CloudError::Io(_) => true,
        _ => false,
    }
}

/// Retry executor with exponential backoff
#[cfg(feature = "async")]
pub struct RetryExecutor {
    /// Configuration
    config: RetryConfig,
    /// Circuit breaker
    circuit_breaker: Option<CircuitBreaker>,
    /// Retry budget
    retry_budget: Option<RetryBudget>,
}

#[cfg(feature = "async")]
impl RetryExecutor {
    /// Creates a new retry executor
    #[must_use]
    pub fn new(config: RetryConfig) -> Self {
        let circuit_breaker = if config.circuit_breaker {
            Some(CircuitBreaker::new(
                config.circuit_breaker_threshold,
                config.circuit_breaker_timeout,
            ))
        } else {
            None
        };

        let retry_budget = Some(RetryBudget::new(100, 10.0)); // 100 tokens, refill at 10/sec

        Self {
            config,
            circuit_breaker,
            retry_budget,
        }
    }

    /// Executes an async operation with retry logic
    pub async fn execute<F, Fut, T>(&mut self, mut operation: F) -> Result<T>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<T>>,
    {
        // Check circuit breaker
        if let Some(ref mut cb) = self.circuit_breaker {
            cb.allow_request()?;
        }

        let mut backoff = Backoff::new(self.config.clone());
        let mut attempts = 0;

        loop {
            match operation().await {
                Ok(result) => {
                    // Success - record and return
                    if let Some(ref mut cb) = self.circuit_breaker {
                        cb.record_success();
                    }
                    return Ok(result);
                }
                Err(error) => {
                    attempts += 1;

                    // Check if error is retryable
                    if !is_retryable(&error) {
                        tracing::warn!("Non-retryable error: {}", error);
                        return Err(error);
                    }

                    // Check if we've exceeded max retries
                    if attempts > self.config.max_retries {
                        tracing::error!("Max retries ({}) exceeded", self.config.max_retries);
                        if let Some(ref mut cb) = self.circuit_breaker {
                            cb.record_failure();
                        }
                        return Err(CloudError::Retry(RetryError::MaxRetriesExceeded {
                            attempts,
                        }));
                    }

                    // Check retry budget
                    if let Some(ref mut budget) = self.retry_budget {
                        budget.try_consume()?;
                    }

                    // Calculate backoff and wait
                    let delay = backoff.next();
                    tracing::warn!(
                        "Retry attempt {}/{} after {:?}: {}",
                        attempts,
                        self.config.max_retries,
                        delay,
                        error
                    );

                    tokio::time::sleep(delay).await;
                }
            }
        }
    }
}

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

    #[test]
    fn test_retry_config_builder() {
        let config = RetryConfig::new()
            .with_max_retries(5)
            .with_initial_backoff(Duration::from_millis(50))
            .with_backoff_multiplier(3.0)
            .with_jitter(false);

        assert_eq!(config.max_retries, 5);
        assert_eq!(config.initial_backoff, Duration::from_millis(50));
        assert_eq!(config.backoff_multiplier, 3.0);
        assert!(!config.jitter);
    }

    #[test]
    fn test_circuit_breaker_closed() {
        let mut cb = CircuitBreaker::new(3, Duration::from_secs(60));
        assert_eq!(cb.state, CircuitState::Closed);
        assert!(cb.allow_request().is_ok());
    }

    #[test]
    fn test_circuit_breaker_opens() {
        let mut cb = CircuitBreaker::new(3, Duration::from_secs(60));

        // Record 3 failures
        cb.record_failure();
        cb.record_failure();
        cb.record_failure();

        assert_eq!(cb.state, CircuitState::Open);
        assert!(cb.allow_request().is_err());
    }

    #[test]
    fn test_circuit_breaker_half_open() {
        let mut cb = CircuitBreaker::new(3, Duration::from_millis(10));

        // Open the circuit
        cb.record_failure();
        cb.record_failure();
        cb.record_failure();
        assert_eq!(cb.state, CircuitState::Open);

        // Wait for timeout
        std::thread::sleep(Duration::from_millis(20));

        // Should transition to half-open
        assert!(cb.allow_request().is_ok());
        assert_eq!(cb.state, CircuitState::HalfOpen);

        // Success should close the circuit
        cb.record_success();
        assert_eq!(cb.state, CircuitState::Closed);
    }

    #[test]
    fn test_retry_budget() {
        // Use a high refill rate (100 tokens/sec) so we can sleep only 50ms
        // instead of seconds. At 100 tokens/sec, 50ms yields ~5 tokens.
        let mut budget = RetryBudget::new(10, 100.0);

        // Consume 10 tokens
        for _ in 0..10 {
            assert!(budget.try_consume().is_ok());
        }

        // 11th should fail
        assert!(budget.try_consume().is_err());

        // Wait for refill (50ms at 100 tokens/sec = ~5 tokens)
        std::thread::sleep(Duration::from_millis(50));

        // Should have refilled
        assert!(budget.try_consume().is_ok());
    }

    #[test]
    fn test_backoff_exponential() {
        let config = RetryConfig::new()
            .with_initial_backoff(Duration::from_millis(100))
            .with_backoff_multiplier(2.0)
            .with_jitter(false);

        let mut backoff = Backoff::new(config);

        let d1 = backoff.next();
        let d2 = backoff.next();
        let d3 = backoff.next();

        assert!(d1 < d2);
        assert!(d2 < d3);
    }

    #[test]
    fn test_backoff_first_jitter_is_not_deterministic_zero() {
        // Regression test: the old fixed-zero static seed made the very
        // first jitter draw of every fresh process exactly 0.0 (i.e. no
        // jitter at all on the first backoff). With a per-instance,
        // entropy-seeded LCG the first jitter factor should differ from the
        // no-jitter base virtually always.
        let config = RetryConfig::new()
            .with_initial_backoff(Duration::from_millis(1000))
            .with_backoff_multiplier(1.0)
            .with_jitter(true);
        let mut jittered = Backoff::new(config.clone());
        let jittered_first = jittered.next();

        let mut unjittered = Backoff::new(config.with_jitter(false));
        let base_first = unjittered.next();

        assert_ne!(
            jittered_first, base_first,
            "first jittered backoff should not exactly equal the unjittered base"
        );
    }

    #[test]
    fn test_backoff_independent_instances_diverge() {
        // Regression test: two independently constructed `Backoff`s
        // (simulating two process starts) must not replay an identical
        // jitter sequence -- the old implementation used one process-wide
        // static seed shared by every instance.
        let config = RetryConfig::new()
            .with_initial_backoff(Duration::from_millis(1000))
            .with_backoff_multiplier(1.0)
            .with_jitter(true);

        let mut backoff_a = Backoff::new(config.clone());
        let mut backoff_b = Backoff::new(config);

        let sequence_a: Vec<Duration> = (0..5).map(|_| backoff_a.next()).collect();
        let sequence_b: Vec<Duration> = (0..5).map(|_| backoff_b.next()).collect();

        assert_ne!(
            sequence_a, sequence_b,
            "independently constructed Backoff instances must not produce identical jitter sequences"
        );
    }

    #[test]
    fn test_is_retryable() {
        let timeout_error = CloudError::Timeout {
            message: "timeout".to_string(),
        };
        assert!(is_retryable(&timeout_error));

        let rate_limit_error = CloudError::RateLimitExceeded {
            message: "rate limit".to_string(),
        };
        assert!(is_retryable(&rate_limit_error));

        let not_found_error = CloudError::NotFound {
            key: "test".to_string(),
        };
        assert!(!is_retryable(&not_found_error));
    }

    #[cfg(feature = "async")]
    #[tokio::test]
    async fn test_retry_executor_success() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let config = RetryConfig::new().with_max_retries(3);
        let mut executor = RetryExecutor::new(config);

        let attempt = std::sync::Arc::new(AtomicUsize::new(0));
        let attempt_clone = attempt.clone();
        let result = executor
            .execute(|| {
                let attempt = attempt_clone.clone();
                async move {
                    let current = attempt.fetch_add(1, Ordering::SeqCst) + 1;
                    if current < 2 {
                        Err(CloudError::Timeout {
                            message: "timeout".to_string(),
                        })
                    } else {
                        Ok(42)
                    }
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(result.ok(), Some(42));
        assert_eq!(attempt.load(Ordering::SeqCst), 2);
    }

    #[cfg(feature = "async")]
    #[tokio::test]
    async fn test_retry_executor_max_retries() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        let config = RetryConfig::new().with_max_retries(2);
        let mut executor = RetryExecutor::new(config);

        let attempt = std::sync::Arc::new(AtomicUsize::new(0));
        let attempt_clone = attempt.clone();
        let result: Result<i32> = executor
            .execute(|| {
                let attempt = attempt_clone.clone();
                async move {
                    attempt.fetch_add(1, Ordering::SeqCst);
                    Err(CloudError::Timeout {
                        message: "timeout".to_string(),
                    })
                }
            })
            .await;

        assert!(result.is_err());
        assert_eq!(attempt.load(Ordering::SeqCst), 3); // Initial + 2 retries
    }
}