phytrace_sdk/reliability/

retry.rs

//! Retry logic with exponential backoff and jitter.

use rand::Rng;
use std::time::Duration;

use crate::core::config::RetryConfig;
use crate::error::{PhyTraceError, PhyTraceResult};
use crate::transport::traits::SendResult;

/// Retry handler with exponential backoff.
#[derive(Debug, Clone)]
pub struct RetryHandler {
    config: RetryConfig,
}

impl RetryHandler {
    /// Create a new retry handler from configuration.
    pub fn new(config: &RetryConfig) -> Self {
        Self {
            config: config.clone(),
        }
    }

    /// Create with default configuration.
    pub fn default_handler() -> Self {
        Self::new(&RetryConfig::default())
    }

    /// Execute an async operation with retries.
    pub async fn execute<F, Fut, T>(&self, mut operation: F) -> PhyTraceResult<T>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = PhyTraceResult<T>>,
    {
        let mut attempt = 0;
        let mut last_error = None;

        while attempt <= self.config.max_retries {
            match operation().await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    if !self.is_retryable(&e) {
                        return Err(e);
                    }

                    last_error = Some(e);
                    attempt += 1;

                    if attempt <= self.config.max_retries {
                        let delay = self.calculate_delay(attempt);
                        tokio::time::sleep(delay).await;
                    }
                }
            }
        }

        Err(last_error
            .unwrap_or_else(|| PhyTraceError::Transport("Max retries exceeded".to_string())))
    }

    /// Execute with retries and return attempt count.
    pub async fn execute_with_stats<F, Fut, T>(
        &self,
        mut operation: F,
    ) -> (PhyTraceResult<T>, RetryStats)
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = PhyTraceResult<T>>,
    {
        let mut stats = RetryStats::default();
        let mut last_error = None;

        while stats.attempts <= self.config.max_retries {
            stats.attempts += 1;

            match operation().await {
                Ok(result) => {
                    stats.succeeded = true;
                    return (Ok(result), stats);
                }
                Err(e) => {
                    if !self.is_retryable(&e) {
                        return (Err(e), stats);
                    }

                    last_error = Some(e);
                    stats.retries += 1;

                    if stats.attempts <= self.config.max_retries {
                        let delay = self.calculate_delay(stats.attempts);
                        stats.total_delay_ms += delay.as_millis() as u64;
                        tokio::time::sleep(delay).await;
                    }
                }
            }
        }

        (
            Err(last_error
                .unwrap_or_else(|| PhyTraceError::Transport("Max retries exceeded".to_string()))),
            stats,
        )
    }

    /// Check if a send result indicates a retryable failure.
    pub fn should_retry(&self, result: &SendResult) -> bool {
        !result.success && result.is_retryable()
    }

    /// Calculate delay for a given attempt number.
    pub fn calculate_delay(&self, attempt: u32) -> Duration {
        let base_delay = self.config.initial_backoff_ms as f64;
        let multiplier = self.config.backoff_multiplier;

        // Exponential backoff
        let delay = base_delay * multiplier.powi(attempt.saturating_sub(1) as i32);

        // Cap at max backoff
        let delay = delay.min(self.config.max_backoff_ms as f64);

        // Add jitter if enabled
        let delay = if self.config.jitter {
            let jitter = rand::thread_rng().gen_range(0.0..=0.3);
            delay * (1.0 + jitter)
        } else {
            delay
        };

        Duration::from_millis(delay as u64)
    }

    /// Check if an error is retryable.
    fn is_retryable(&self, error: &PhyTraceError) -> bool {
        matches!(
            error,
            PhyTraceError::Transport(_) | PhyTraceError::Timeout(_)
        )
    }

    /// Get the maximum number of retries.
    pub fn max_retries(&self) -> u32 {
        self.config.max_retries
    }

    /// Get the initial backoff duration.
    pub fn initial_backoff(&self) -> Duration {
        Duration::from_millis(self.config.initial_backoff_ms)
    }
}

/// Statistics from a retry operation.
#[derive(Debug, Clone, Default)]
pub struct RetryStats {
    /// Number of attempts made.
    pub attempts: u32,
    /// Number of retries (attempts - 1 if successful on first try).
    pub retries: u32,
    /// Total delay time in milliseconds.
    pub total_delay_ms: u64,
    /// Whether the operation eventually succeeded.
    pub succeeded: bool,
}

impl RetryStats {
    /// Check if any retries were needed.
    pub fn had_retries(&self) -> bool {
        self.retries > 0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU32, Ordering};

    #[test]
    fn test_calculate_delay() {
        let config = RetryConfig {
            initial_backoff_ms: 100,
            max_backoff_ms: 10000,
            backoff_multiplier: 2.0,
            jitter: false,
            max_retries: 5,
        };
        let handler = RetryHandler::new(&config);

        // Without jitter, delays should be predictable
        assert_eq!(handler.calculate_delay(1), Duration::from_millis(100));
        assert_eq!(handler.calculate_delay(2), Duration::from_millis(200));
        assert_eq!(handler.calculate_delay(3), Duration::from_millis(400));
        assert_eq!(handler.calculate_delay(4), Duration::from_millis(800));
    }

    #[test]
    fn test_delay_capped() {
        let config = RetryConfig {
            initial_backoff_ms: 1000,
            max_backoff_ms: 5000,
            backoff_multiplier: 2.0,
            jitter: false,
            max_retries: 10,
        };
        let handler = RetryHandler::new(&config);

        // Should be capped at max
        assert_eq!(handler.calculate_delay(5), Duration::from_millis(5000));
        assert_eq!(handler.calculate_delay(6), Duration::from_millis(5000));
    }

    #[tokio::test]
    async fn test_execute_success_first_try() {
        let handler = RetryHandler::default_handler();

        let result = handler
            .execute(|| async { Ok::<_, PhyTraceError>(42) })
            .await;

        assert_eq!(result.unwrap(), 42);
    }

    #[tokio::test]
    async fn test_execute_success_after_retry() {
        let handler = RetryHandler::new(&RetryConfig {
            max_retries: 3,
            initial_backoff_ms: 1, // Minimal delay for testing
            ..Default::default()
        });

        let counter = AtomicU32::new(0);

        let result = handler
            .execute(|| async {
                let count = counter.fetch_add(1, Ordering::Relaxed);
                if count < 2 {
                    Err(PhyTraceError::Transport("temporary failure".to_string()))
                } else {
                    Ok(42)
                }
            })
            .await;

        assert_eq!(result.unwrap(), 42);
        assert_eq!(counter.load(Ordering::Relaxed), 3);
    }

    #[tokio::test]
    async fn test_execute_all_retries_fail() {
        let handler = RetryHandler::new(&RetryConfig {
            max_retries: 2,
            initial_backoff_ms: 1,
            ..Default::default()
        });

        let counter = AtomicU32::new(0);

        let result = handler
            .execute(|| async {
                counter.fetch_add(1, Ordering::Relaxed);
                Err::<i32, _>(PhyTraceError::Transport("persistent failure".to_string()))
            })
            .await;

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

    #[tokio::test]
    async fn test_execute_with_stats() {
        let handler = RetryHandler::new(&RetryConfig {
            max_retries: 3,
            initial_backoff_ms: 1,
            ..Default::default()
        });

        let counter = AtomicU32::new(0);

        let (result, stats) = handler
            .execute_with_stats(|| async {
                let count = counter.fetch_add(1, Ordering::Relaxed);
                if count < 2 {
                    Err(PhyTraceError::Transport("temporary".to_string()))
                } else {
                    Ok(42)
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(stats.attempts, 3);
        assert_eq!(stats.retries, 2);
        assert!(stats.succeeded);
    }

    #[test]
    fn test_should_retry() {
        let handler = RetryHandler::default_handler();

        let retryable = SendResult::failure(Some(500), "Server error");
        assert!(handler.should_retry(&retryable));

        let not_retryable = SendResult::failure(Some(400), "Bad request");
        assert!(!handler.should_retry(&not_retryable));

        let success = SendResult::success(50);
        assert!(!handler.should_retry(&success));
    }
}