seer_core/

retry.rs

//! Retry logic with exponential backoff for transient failures.
//!
//! This module provides configurable retry policies and executors for handling
//! transient network failures in WHOIS and RDAP lookups.

use std::future::Future;
use std::time::Duration;

use rand::RngExt;
use tracing::debug;

use crate::error::{Result, SeerError};

/// Configuration for retry behavior with exponential backoff.
#[derive(Debug, Clone)]
pub struct RetryPolicy {
    /// Maximum number of attempts (including the initial attempt).
    pub max_attempts: usize,
    /// Initial delay before the first retry.
    pub initial_delay: Duration,
    /// Maximum delay between retries (caps exponential growth).
    pub max_delay: Duration,
    /// Multiplier for exponential backoff (delay *= multiplier after each retry).
    pub multiplier: f64,
    /// Whether to add random jitter to delays to avoid thundering herd.
    pub jitter: bool,
}

impl Default for RetryPolicy {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            initial_delay: Duration::from_millis(100),
            max_delay: Duration::from_secs(5),
            multiplier: 2.0,
            jitter: true,
        }
    }
}

impl RetryPolicy {
    /// Creates a new retry policy with default settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the maximum number of attempts.
    pub fn with_max_attempts(mut self, attempts: usize) -> Self {
        self.max_attempts = attempts.max(1);
        self
    }

    /// Sets the initial delay before the first retry.
    pub fn with_initial_delay(mut self, delay: Duration) -> Self {
        self.initial_delay = delay;
        self
    }

    /// Sets the maximum delay between retries.
    pub fn with_max_delay(mut self, delay: Duration) -> Self {
        self.max_delay = delay;
        self
    }

    /// Sets the multiplier for exponential backoff.
    pub fn with_multiplier(mut self, multiplier: f64) -> Self {
        self.multiplier = multiplier.max(1.0);
        self
    }

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

    /// Creates a policy that disables retries (single attempt only).
    pub fn no_retry() -> Self {
        Self {
            max_attempts: 1,
            ..Self::default()
        }
    }

    /// Calculates the delay for a given attempt number (0-indexed).
    ///
    /// The attempt number is internally capped to prevent integer overflow
    /// in the exponential calculation.
    pub fn delay_for_attempt(&self, attempt: usize) -> Duration {
        // Cap attempt to prevent overflow in powi() - 20 attempts with multiplier 2.0
        // gives 2^20 = ~1 million, which is safe for f64 and reasonable for delays
        let safe_attempt = attempt.min(20) as i32;

        // powi(0) == 1.0, so attempt 0 naturally yields initial_delay — no
        // short-circuit, which would skip the jitter branch below on the very
        // attempt where concurrent callers are most likely in lockstep.
        let base_delay = self.initial_delay.as_millis() as f64 * self.multiplier.powi(safe_attempt);
        let capped_delay = base_delay.min(self.max_delay.as_millis() as f64);

        let final_delay = if self.jitter {
            // Full jitter (AWS "Exponential Backoff and Jitter"):
            // sleep ∈ [0, capped_delay]. Half-jitter ([0.5..1.0]) clusters
            // retries in the upper half of the window, which leaves
            // measurable thundering-herd behaviour after a brief outage.
            // Full jitter spreads retries uniformly across the window.
            let mut rng = rand::rng();
            let jitter_factor = rng.random_range(0.0..1.0);
            capped_delay * jitter_factor
        } else {
            capped_delay
        };

        Duration::from_millis(final_delay as u64)
    }
}

/// Whether `error` is transient, so a WHOIS/RDAP operation is worth retrying.
///
/// Retryable:
/// - Timeouts
/// - Connection failures (IO errors)
/// - Rate limiting (429)
/// - Server errors (5xx)
///
/// Not retryable:
/// - Invalid input (domain, IP, record type)
/// - Server not found
/// - Parse errors (JSON, WHOIS format)
pub fn is_retryable(error: &SeerError) -> bool {
    match error {
        // Transient errors - worth retrying
        SeerError::Timeout(_) => true,
        SeerError::WhoisConnectionFailed(_) => true,
        SeerError::RateLimited(_) => true,

        // Transiency was already classified at the From<reqwest::Error>
        // boundary (see error.rs); just read the stored flag.
        SeerError::ReqwestError { transient, .. } => *transient,

        // WHOIS errors might be transient if they're connection-related
        SeerError::WhoisError(msg) => {
            let lower = msg.to_lowercase();
            lower.contains("connection")
                || lower.contains("timeout")
                || lower.contains("refused")
                || lower.contains("reset")
        }

        // RDAP errors might be transient server errors
        SeerError::RdapError(msg) => {
            let lower = msg.to_lowercase();
            lower.contains("status 5") || lower.contains("status 429") || lower.contains("timeout")
        }

        // Bootstrap errors could be transient if IANA is temporarily
        // unavailable. "all IANA bootstrap registries failed" is exactly
        // that case (every fetch errored), so it must be retryable —
        // otherwise the bootstrap loader's retry policy never retries.
        // "no RDAP server for ..." and "throttled and no cache" are not.
        SeerError::RdapBootstrapError(msg) => {
            let lower = msg.to_lowercase();
            lower.contains("timeout")
                || lower.contains("connection")
                || lower.contains("registries failed")
        }

        // DNS errors can be transient
        SeerError::DnsError(msg) => {
            let lower = msg.to_lowercase();
            lower.contains("timeout") || lower.contains("temporary")
        }

        // HTTP errors might be transient (server errors 5xx or 429 Too Many Requests)
        SeerError::HttpError(msg) => {
            let lower = msg.to_lowercase();
            lower.contains("timeout")
                || lower.contains("connection")
                || lower.contains("status 5")
                || lower.contains("status 429")
        }

        // Not retryable - permanent failures
        SeerError::InvalidDomain(_) => false,
        SeerError::DomainNotAllowed { .. } => false,
        SeerError::InvalidIpAddress(_) => false,
        SeerError::InvalidRecordType(_) => false,
        SeerError::WhoisServerNotFound(_) => false,
        SeerError::JsonError(_) => false,
        SeerError::CertificateError(_) => false,
        SeerError::SslError(_) => false,
        SeerError::DnsResolverError(_) => false,
        SeerError::LookupFailed { .. } => false,
        SeerError::ConfigError(_) => false,
        SeerError::InvalidInput(_) => false,
        // Transparent pass-through: if a prior attempt was wrapped into
        // RetryExhausted upstream, defer the retryable decision to the
        // underlying cause so a caller layering retries still sees the
        // true fault classification instead of a non-retryable wrapper.
        SeerError::RetryExhausted { last_error, .. } => is_retryable(last_error),
        SeerError::Other(_) => false,
    }
}

/// Executes operations with retry logic using exponential backoff.
#[derive(Debug, Clone)]
pub struct RetryExecutor {
    policy: RetryPolicy,
}

impl RetryExecutor {
    /// Creates a new executor for `policy`.
    pub fn new(policy: RetryPolicy) -> Self {
        Self { policy }
    }

    /// Executes an async operation with retry logic.
    ///
    /// The operation will be retried up to `max_attempts` times if it fails
    /// with a retryable error (see [`is_retryable`]). Delays between retries
    /// follow exponential backoff with optional jitter.
    pub async fn execute<F, Fut, T>(&self, mut operation: F) -> Result<T>
    where
        F: FnMut() -> Fut,
        Fut: Future<Output = Result<T>>,
    {
        self.execute_with_delay_hint(|| {
            let attempt = operation();
            async move { attempt.await.map_err(|e| (e, None)) }
        })
        .await
    }

    /// Like [`execute`](Self::execute), but each failure may carry a
    /// server-suggested delay (an HTTP 429 `Retry-After`). A hint replaces the
    /// backoff for that retry, capped at the policy's
    /// [`max_delay`](RetryPolicy::max_delay) so a hostile or misconfigured
    /// header cannot pin the caller.
    pub async fn execute_with_delay_hint<F, Fut, T>(&self, mut operation: F) -> Result<T>
    where
        F: FnMut() -> Fut,
        Fut: Future<Output = std::result::Result<T, (SeerError, Option<Duration>)>>,
    {
        let mut attempt = 0;
        loop {
            let (e, hint) = match operation().await {
                Ok(result) => return Ok(result),
                Err(failure) => failure,
            };
            let attempts_remaining = self.policy.max_attempts.saturating_sub(attempt + 1);
            if !is_retryable(&e) || attempts_remaining == 0 {
                if attempt == 0 {
                    return Err(e);
                }
                debug!(
                    attempt = attempt + 1,
                    max_attempts = self.policy.max_attempts,
                    error = %e,
                    "Operation failed after retries"
                );
                return Err(SeerError::RetryExhausted {
                    attempts: attempt + 1,
                    last_error: Box::new(e),
                });
            }

            let delay = match hint {
                Some(hint) => hint.min(self.policy.max_delay),
                None => self.policy.delay_for_attempt(attempt),
            };
            debug!(
                attempt = attempt + 1,
                max_attempts = self.policy.max_attempts,
                delay_ms = delay.as_millis(),
                error = %e,
                "Retrying after transient error"
            );
            tokio::time::sleep(delay).await;
            attempt += 1;
        }
    }
}

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

    #[test]
    fn test_retry_policy_defaults() {
        let policy = RetryPolicy::default();
        assert_eq!(policy.max_attempts, 3);
        assert_eq!(policy.initial_delay, Duration::from_millis(100));
        assert_eq!(policy.max_delay, Duration::from_secs(5));
        assert_eq!(policy.multiplier, 2.0);
        assert!(policy.jitter);
    }

    #[test]
    fn test_retry_policy_builder() {
        let policy = RetryPolicy::new()
            .with_max_attempts(5)
            .with_initial_delay(Duration::from_millis(200))
            .with_max_delay(Duration::from_secs(10))
            .with_multiplier(3.0)
            .with_jitter(false);

        assert_eq!(policy.max_attempts, 5);
        assert_eq!(policy.initial_delay, Duration::from_millis(200));
        assert_eq!(policy.max_delay, Duration::from_secs(10));
        assert_eq!(policy.multiplier, 3.0);
        assert!(!policy.jitter);
    }

    #[test]
    fn test_delay_calculation_no_jitter() {
        let policy = RetryPolicy::new()
            .with_initial_delay(Duration::from_millis(100))
            .with_multiplier(2.0)
            .with_max_delay(Duration::from_secs(10))
            .with_jitter(false);

        assert_eq!(policy.delay_for_attempt(0), Duration::from_millis(100));
        assert_eq!(policy.delay_for_attempt(1), Duration::from_millis(200));
        assert_eq!(policy.delay_for_attempt(2), Duration::from_millis(400));
        assert_eq!(policy.delay_for_attempt(3), Duration::from_millis(800));
    }

    #[test]
    fn test_first_retry_is_jittered() {
        // The first retry (attempt 0) is exactly where a shared outage or
        // 429 hits every concurrent caller at once, so full jitter matters
        // most there. A fixed `initial_delay` short-circuit produced a
        // thundering herd of lockstep retries (2026-07-11 review).
        let policy = RetryPolicy::new()
            .with_initial_delay(Duration::from_millis(100))
            .with_jitter(true);

        let delays: Vec<Duration> = (0..64).map(|_| policy.delay_for_attempt(0)).collect();
        // Full jitter: every sample within [0, initial_delay] ...
        assert!(delays.iter().all(|d| *d <= Duration::from_millis(100)));
        // ... and actually randomized (64 identical samples ≈ impossible).
        assert!(
            delays.iter().any(|d| *d != delays[0]),
            "attempt-0 delay is constant — jitter not applied to the first retry"
        );
    }

    #[test]
    fn test_delay_capped_at_max() {
        let policy = RetryPolicy::new()
            .with_initial_delay(Duration::from_secs(1))
            .with_multiplier(10.0)
            .with_max_delay(Duration::from_secs(5))
            .with_jitter(false);

        // 1s * 10^2 = 100s, but capped at 5s
        assert_eq!(policy.delay_for_attempt(2), Duration::from_secs(5));
    }

    #[test]
    fn test_classifier_timeout_is_retryable() {
        assert!(is_retryable(&SeerError::Timeout("test".to_string())));
    }

    #[test]
    fn test_classifier_invalid_domain_not_retryable() {
        assert!(!is_retryable(&SeerError::InvalidDomain("test".to_string())));
    }

    #[test]
    fn test_classifier_server_not_found_not_retryable() {
        assert!(!is_retryable(&SeerError::WhoisServerNotFound(
            "test".to_string()
        )));
    }

    #[test]
    fn test_classifier_rate_limited_is_retryable() {
        assert!(is_retryable(&SeerError::RateLimited("test".to_string())));
    }

    #[test]
    fn test_classifier_reads_reqwest_transient_flag() {
        // Transiency is classified once at the From<reqwest::Error> boundary
        // (tested in error.rs); the classifier only reads the stored flag.
        assert!(is_retryable(&SeerError::ReqwestError {
            message: "HTTP status server error (503 Service Unavailable)".to_string(),
            transient: true,
        }));
        assert!(!is_retryable(&SeerError::ReqwestError {
            message: "HTTP status client error (404 Not Found)".to_string(),
            transient: false,
        }));
    }

    #[tokio::test]
    async fn test_executor_success_on_first_try() {
        let policy = RetryPolicy::new().with_max_attempts(3);
        let executor = RetryExecutor::new(policy);
        let attempts = Arc::new(AtomicUsize::new(0));

        let attempts_clone = attempts.clone();
        let result: Result<&str> = executor
            .execute(|| {
                let a = attempts_clone.clone();
                async move {
                    a.fetch_add(1, Ordering::SeqCst);
                    Ok("success")
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "success");
        assert_eq!(attempts.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_executor_retries_on_transient_error() {
        let policy = RetryPolicy::new()
            .with_max_attempts(3)
            .with_initial_delay(Duration::from_millis(1))
            .with_jitter(false);
        let executor = RetryExecutor::new(policy);
        let attempts = Arc::new(AtomicUsize::new(0));

        let attempts_clone = attempts.clone();
        let result: Result<&str> = executor
            .execute(|| {
                let a = attempts_clone.clone();
                async move {
                    let count = a.fetch_add(1, Ordering::SeqCst);
                    if count < 2 {
                        Err(SeerError::Timeout("test timeout".to_string()))
                    } else {
                        Ok("success after retries")
                    }
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "success after retries");
        assert_eq!(attempts.load(Ordering::SeqCst), 3);
    }

    #[tokio::test]
    async fn test_executor_no_retry_on_non_retryable_error() {
        let policy = RetryPolicy::new()
            .with_max_attempts(3)
            .with_initial_delay(Duration::from_millis(1));
        let executor = RetryExecutor::new(policy);
        let attempts = Arc::new(AtomicUsize::new(0));

        let attempts_clone = attempts.clone();
        let result: Result<&str> = executor
            .execute(|| {
                let a = attempts_clone.clone();
                async move {
                    a.fetch_add(1, Ordering::SeqCst);
                    Err(SeerError::InvalidDomain("bad.".to_string()))
                }
            })
            .await;

        assert!(result.is_err());
        // Should only attempt once since InvalidDomain is not retryable
        assert_eq!(attempts.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_executor_exhausts_retries() {
        let policy = RetryPolicy::new()
            .with_max_attempts(3)
            .with_initial_delay(Duration::from_millis(1))
            .with_jitter(false);
        let executor = RetryExecutor::new(policy);
        let attempts = Arc::new(AtomicUsize::new(0));

        let attempts_clone = attempts.clone();
        let result: Result<&str> = executor
            .execute(|| {
                let a = attempts_clone.clone();
                async move {
                    a.fetch_add(1, Ordering::SeqCst);
                    Err(SeerError::Timeout("always fails".to_string()))
                }
            })
            .await;

        assert!(result.is_err());
        assert_eq!(attempts.load(Ordering::SeqCst), 3);

        // Check that we get RetryExhausted error
        match result.unwrap_err() {
            SeerError::RetryExhausted { attempts, .. } => {
                assert_eq!(attempts, 3);
            }
            other => panic!("Expected RetryExhausted, got {:?}", other),
        }
    }

    #[tokio::test(start_paused = true)]
    async fn delay_hint_replaces_backoff_and_is_capped_at_max_delay() {
        // A server hint (Retry-After) wins over the backoff, but never waits
        // past `max_delay`: the 60s hint here is honored as 2s.
        let policy = RetryPolicy::new()
            .with_max_attempts(2)
            .with_initial_delay(Duration::from_millis(1))
            .with_max_delay(Duration::from_secs(2))
            .with_jitter(false);
        let attempts = Arc::new(AtomicUsize::new(0));
        let start = tokio::time::Instant::now();
        let a = attempts.clone();
        let result: Result<&str> = RetryExecutor::new(policy)
            .execute_with_delay_hint(|| {
                let a = a.clone();
                async move {
                    if a.fetch_add(1, Ordering::SeqCst) == 0 {
                        Err((
                            SeerError::RateLimited("429".to_string()),
                            Some(Duration::from_secs(60)),
                        ))
                    } else {
                        Ok("ok")
                    }
                }
            })
            .await;
        assert_eq!(result.unwrap(), "ok");
        assert_eq!(attempts.load(Ordering::SeqCst), 2);
        assert_eq!(start.elapsed(), Duration::from_secs(2));
    }

    #[tokio::test(start_paused = true)]
    async fn short_delay_hint_is_honored_over_backoff() {
        let policy = RetryPolicy::new()
            .with_max_attempts(2)
            .with_initial_delay(Duration::from_secs(4))
            .with_jitter(false);
        let start = tokio::time::Instant::now();
        let first = Arc::new(AtomicUsize::new(0));
        let result: Result<()> = RetryExecutor::new(policy)
            .execute_with_delay_hint(|| {
                let first = first.clone();
                async move {
                    if first.fetch_add(1, Ordering::SeqCst) == 0 {
                        Err((
                            SeerError::RateLimited("429".to_string()),
                            Some(Duration::from_secs(1)),
                        ))
                    } else {
                        Ok(())
                    }
                }
            })
            .await;
        assert!(result.is_ok());
        assert_eq!(start.elapsed(), Duration::from_secs(1));
    }

    #[test]
    fn test_no_retry_policy() {
        let policy = RetryPolicy::no_retry();
        assert_eq!(policy.max_attempts, 1);
    }

    #[test]
    fn test_delay_overflow_protection() {
        let policy = RetryPolicy::new()
            .with_initial_delay(Duration::from_millis(100))
            .with_multiplier(2.0)
            .with_max_delay(Duration::from_secs(5))
            .with_jitter(false);

        // Test with very large attempt numbers - should not panic or produce invalid durations
        let delay_50 = policy.delay_for_attempt(50);
        let delay_100 = policy.delay_for_attempt(100);
        let delay_1000 = policy.delay_for_attempt(1000);

        // All should be capped at max_delay due to our overflow protection
        assert!(delay_50 <= Duration::from_secs(5));
        assert!(delay_100 <= Duration::from_secs(5));
        assert!(delay_1000 <= Duration::from_secs(5));
    }

    #[test]
    fn retry_exhausted_is_retryable_if_inner_is() {
        // Regression for H5: the retry classifier must look through a
        // RetryExhausted wrapper so a caller layering retries can still see
        // the true underlying fault classification instead of collapsing to
        // a non-retryable wrapper.

        let retryable_inner = SeerError::Timeout("inner timed out".to_string());
        let wrapped_retryable = SeerError::RetryExhausted {
            attempts: 3,
            last_error: Box::new(retryable_inner),
        };
        assert!(
            is_retryable(&wrapped_retryable),
            "RetryExhausted wrapping a retryable Timeout should be retryable",
        );

        let non_retryable_inner = SeerError::InvalidDomain("bad.".to_string());
        let wrapped_non_retryable = SeerError::RetryExhausted {
            attempts: 3,
            last_error: Box::new(non_retryable_inner),
        };
        assert!(
            !is_retryable(&wrapped_non_retryable),
            "RetryExhausted wrapping a non-retryable InvalidDomain must not be retryable",
        );
    }

    #[test]
    fn bootstrap_all_registries_failed_is_retryable() {
        // The RDAP bootstrap loader retries through RetryExecutor; if its
        // "every IANA registry fetch failed" error isn't classified as
        // transient, the retry policy never gets a second attempt.
        assert!(is_retryable(&SeerError::RdapBootstrapError(
            "all IANA bootstrap registries failed".to_string()
        )));
        // Permanent bootstrap outcomes stay non-retryable.
        assert!(!is_retryable(&SeerError::RdapBootstrapError(
            "no RDAP server for example.ru".to_string()
        )));
        assert!(!is_retryable(&SeerError::RdapBootstrapError(
            "bootstrap refresh throttled and no cache available".to_string()
        )));
    }
}