use reqwest::{Error, Response, StatusCode}; use std::future::Future; use std::time::Duration; use rand::RngExt; /// Errors that can be retried #[derive(Debug)] pub enum RetryableError { RateLimit(StatusCode, Option), // Status code and optional Retry-After seconds ServerError(StatusCode), NonRetryable(StatusCode), RequestError(Error), } impl From for RetryableError { fn from(err: Error) -> Self { Self::RequestError(err) } } /// Check HTTP response status and determine if/how to retry /// /// Returns `Ok(response)` for successful responses, or `Err` with appropriate retry error pub fn check_response_status(response: Response) -> Result { match response.status() { StatusCode::TOO_MANY_REQUESTS | StatusCode::FORBIDDEN => { // Extract Retry-After header if present let retry_after = response.headers() .get("retry-after") .and_then(|v| v.to_str().ok()) .and_then(|s| s.parse::().ok()); Err(RetryableError::RateLimit(response.status(), retry_after)) } StatusCode::INTERNAL_SERVER_ERROR | StatusCode::BAD_GATEWAY | StatusCode::SERVICE_UNAVAILABLE | StatusCode::GATEWAY_TIMEOUT => { Err(RetryableError::ServerError(response.status())) } _ if response.status().is_success() => { Ok(response) } _ => { // Non-retryable error (4xx client errors like 404, 401, etc.) Err(RetryableError::NonRetryable(response.status())) } } } impl std::fmt::Display for RetryableError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::RateLimit(status, retry_after) => { if let Some(seconds) = retry_after { write!(f, "Rate limit error: {status} (retry after {seconds} seconds)") } else { write!(f, "Rate limit error: {status}") } } Self::ServerError(status) => write!(f, "Server error: {status}"), Self::NonRetryable(status) => write!(f, "HTTP error: {status}"), Self::RequestError(e) => write!(f, "Request error: {e}"), } } } impl std::error::Error for RetryableError {} /// Configuration for retry behavior #[derive(Debug, Clone)] pub struct RetryConfig { pub max_retries: u32, pub base_delay_ms: u64, pub max_delay_ms: u64, pub jitter_factor: f64, } impl Default for RetryConfig { fn default() -> Self { Self { max_retries: 5, base_delay_ms: 1000, // 1 second base max_delay_ms: 32000, // 32 seconds max jitter_factor: 0.3, // ±30% jitter } } } impl RetryConfig { /// Calculate delay with exponential backoff and jitter #[must_use] pub fn calculate_delay(&self, attempt: u32) -> Duration { // Exponential backoff: base * 2^attempt let exponential_delay = self.base_delay_ms * 2u64.pow(attempt); // Cap at max delay let capped_delay = exponential_delay.min(self.max_delay_ms); // Add jitter: random value between (1 - jitter_factor) and (1 + jitter_factor) let mut rng = rand::rng(); let jitter_range = 1.0 - self.jitter_factor..=1.0 + self.jitter_factor; let jitter_multiplier = rng.random_range(jitter_range); let final_delay = (capped_delay as f64 * jitter_multiplier) as u64; Duration::from_millis(final_delay) } } /// Legacy delay calculation function (kept for backward compatibility with tests) #[cfg(test)] #[must_use] pub const fn calculate_delay(n: u32) -> u64 { let base: u64 = 30; base * 2u64.pow(n) } /// Retry a future with exponential backoff and jitter pub async fn retry_with_backoff( config: &RetryConfig, mut operation: F, ) -> Result where F: FnMut() -> Fut, Fut: Future>, { let mut attempt = 0; loop { match operation().await { Ok(result) => return Ok(result), Err(err) => { // Check if error is retryable let should_retry = matches!( err, RetryableError::RateLimit(_, _) | RetryableError::ServerError(_) ); if !should_retry || attempt >= config.max_retries { log_final_failure(&err, attempt); return Err(err); } // Calculate delay - use Retry-After header if available for rate limits let delay = match &err { RetryableError::RateLimit(_, Some(retry_after_secs)) => { // Use the server's requested delay, but apply jitter and cap let retry_after_ms = retry_after_secs * 1000; let capped = retry_after_ms.min(config.max_delay_ms); // Add jitter to avoid thundering herd let mut rng = rand::rng(); let jitter_range = 1.0 - config.jitter_factor..=1.0 + config.jitter_factor; let jitter_multiplier = rng.random_range(jitter_range); let final_delay = (capped as f64 * jitter_multiplier) as u64; Duration::from_millis(final_delay) } _ => config.calculate_delay(attempt) }; log_retry_attempt(&err, attempt, config.max_retries, &delay); tokio::time::sleep(delay).await; attempt += 1; } } } } /// Log details about a final failure after exhausting retries. fn log_final_failure(err: &RetryableError, attempt: u32) { match err { RetryableError::RateLimit(status, retry_after) => { if let Some(seconds) = retry_after { println!("Rate limit exceeded after {attempt} retries: {status} (server requested {seconds} second wait)"); } else { println!("Rate limit exceeded after {attempt} retries: {status}"); } } RetryableError::ServerError(status) => { println!("Server error after {attempt} retries: {status}"); } RetryableError::NonRetryable(status) => { println!("Non-retryable error: {status}"); } RetryableError::RequestError(e) => { println!("Request error after {attempt} retries: {e}"); } } } /// Log details about a retry attempt. fn log_retry_attempt(err: &RetryableError, attempt: u32, max_retries: u32, delay: &Duration) { match err { RetryableError::RateLimit(status, retry_after) => { if let Some(seconds) = retry_after { println!( "Rate limit hit ({status}), server requested {seconds} second wait, retrying in {:.2}s...", delay.as_secs_f64() ); } else { println!( "Rate limit hit ({status}), attempt {}/{max_retries}, retrying in {:.2}s...", attempt + 1, delay.as_secs_f64() ); } } RetryableError::ServerError(status) => { println!( "Server error ({status}), attempt {}/{max_retries}, retrying in {:.2}s...", attempt + 1, delay.as_secs_f64() ); } _ => {} } } #[cfg(test)] #[allow(clippy::unwrap_used, clippy::expect_used)] mod tests { use super::*; #[test] fn test_retry_config_defaults() { let config = RetryConfig::default(); assert_eq!(config.max_retries, 5); assert_eq!(config.base_delay_ms, 1000); assert_eq!(config.max_delay_ms, 32000); assert!((config.jitter_factor - 0.3).abs() < f64::EPSILON); } #[test] fn test_retry_config_calculate_delay_exponential() { let config = RetryConfig::default(); // Test exponential growth (without jitter for base check) let delay0 = config.calculate_delay(0); let delay1 = config.calculate_delay(1); let delay2 = config.calculate_delay(2); // Delays should be in expected range considering jitter (±30%) assert!(delay0.as_millis() >= 700 && delay0.as_millis() <= 1300); // ~1000ms assert!(delay1.as_millis() >= 1400 && delay1.as_millis() <= 2600); // ~2000ms assert!(delay2.as_millis() >= 2800 && delay2.as_millis() <= 5200); // ~4000ms } #[test] fn test_retry_config_max_delay_cap() { let config = RetryConfig::default(); // Very large attempt should be capped at max_delay_ms let delay = config.calculate_delay(10); // Would be 1024s without cap // Should be capped around 32s with jitter (32000ms * 1.3 = 41600ms) assert!(delay.as_millis() <= 41600); } #[test] fn test_retry_config_jitter_variability() { let config = RetryConfig::default(); // Calculate multiple delays for same attempt let delays: Vec = (0..10) .map(|_| config.calculate_delay(1)) .collect(); // Not all delays should be identical (jitter adds randomness) let first = delays[0]; let has_variation = delays.iter().any(|d| *d != first); assert!(has_variation, "Jitter should cause variation in delays"); } #[test] fn test_retry_config_custom() { let config = RetryConfig { max_retries: 3, base_delay_ms: 500, max_delay_ms: 10000, jitter_factor: 0.5, }; assert_eq!(config.max_retries, 3); assert_eq!(config.base_delay_ms, 500); let delay = config.calculate_delay(0); // Should be around 500ms with ±50% jitter (250-750ms) assert!(delay.as_millis() >= 250 && delay.as_millis() <= 750); } #[tokio::test] async fn test_retry_with_backoff_success_on_first_try() { use std::sync::{Arc, atomic::{AtomicUsize, Ordering}}; let config = RetryConfig::default(); let call_count = Arc::new(AtomicUsize::new(0)); let call_count_clone = call_count.clone(); let result = retry_with_backoff(&config, move || { let count = call_count_clone.clone(); async move { count.fetch_add(1, Ordering::SeqCst); Ok::<_, RetryableError>(42) } }).await; assert!(result.is_ok()); assert_eq!(result.unwrap(), 42); assert_eq!(call_count.load(Ordering::SeqCst), 1); } #[tokio::test] async fn test_retry_with_backoff_success_after_retry() { use std::sync::{Arc, atomic::{AtomicUsize, Ordering}}; let config = RetryConfig { max_retries: 3, base_delay_ms: 10, // Short delay for testing max_delay_ms: 100, jitter_factor: 0.1, }; let call_count = Arc::new(AtomicUsize::new(0)); let call_count_clone = call_count.clone(); let result = retry_with_backoff(&config, move || { let count = call_count_clone.clone(); async move { let current = count.fetch_add(1, Ordering::SeqCst) + 1; if current < 3 { Err(RetryableError::ServerError(StatusCode::INTERNAL_SERVER_ERROR)) } else { Ok::<_, RetryableError>(current) } } }).await; assert!(result.is_ok()); assert_eq!(result.unwrap(), 3); assert_eq!(call_count.load(Ordering::SeqCst), 3); } #[tokio::test] async fn test_retry_with_backoff_non_retryable_error() { use std::sync::{Arc, atomic::{AtomicUsize, Ordering}}; let config = RetryConfig::default(); let call_count = Arc::new(AtomicUsize::new(0)); let call_count_clone = call_count.clone(); let result = retry_with_backoff(&config, move || { let count = call_count_clone.clone(); async move { count.fetch_add(1, Ordering::SeqCst); Err::(RetryableError::NonRetryable(StatusCode::NOT_FOUND)) } }).await; assert!(result.is_err()); assert_eq!(call_count.load(Ordering::SeqCst), 1); // Should not retry } #[tokio::test] async fn test_retry_with_backoff_max_retries_exceeded() { use std::sync::{Arc, atomic::{AtomicUsize, Ordering}}; let config = RetryConfig { max_retries: 2, base_delay_ms: 10, max_delay_ms: 100, jitter_factor: 0.1, }; let call_count = Arc::new(AtomicUsize::new(0)); let call_count_clone = call_count.clone(); let result = retry_with_backoff(&config, move || { let count = call_count_clone.clone(); async move { count.fetch_add(1, Ordering::SeqCst); Err::(RetryableError::RateLimit(StatusCode::TOO_MANY_REQUESTS, None)) } }).await; assert!(result.is_err()); assert_eq!(call_count.load(Ordering::SeqCst), 3); // Initial attempt + 2 retries } #[tokio::test] async fn test_retry_with_backoff_respects_retry_after_header() { use std::sync::{Arc, atomic::{AtomicUsize, Ordering}}; use std::time::Instant; let config = RetryConfig { max_retries: 3, base_delay_ms: 10, max_delay_ms: 5000, // Allow up to 5 seconds to not cap our test jitter_factor: 0.0, // No jitter for precise timing test }; let call_count = Arc::new(AtomicUsize::new(0)); let call_count_clone = call_count.clone(); let start = Instant::now(); let result = retry_with_backoff(&config, move || { let count = call_count_clone.clone(); async move { let current = count.fetch_add(1, Ordering::SeqCst) + 1; if current == 1 { // First call fails with Retry-After of 1 second (1000ms) Err(RetryableError::RateLimit(StatusCode::TOO_MANY_REQUESTS, Some(1))) } else { // Second call succeeds Ok::<_, RetryableError>(42) } } }).await; let elapsed = start.elapsed(); assert!(result.is_ok()); assert_eq!(result.unwrap(), 42); assert_eq!(call_count.load(Ordering::SeqCst), 2); // Should wait approximately 1 second (1000ms) as specified by Retry-After // Allow some tolerance for test execution assert!(elapsed.as_millis() >= 900 && elapsed.as_millis() <= 1100, "Expected ~1000ms delay, got {}ms", elapsed.as_millis()); } #[test] fn test_retryable_error_display_with_retry_after() { let error = RetryableError::RateLimit(StatusCode::TOO_MANY_REQUESTS, Some(60)); assert_eq!(format!("{error}"), "Rate limit error: 429 Too Many Requests (retry after 60 seconds)"); } #[test] fn test_retryable_error_display_without_retry_after() { let error = RetryableError::RateLimit(StatusCode::TOO_MANY_REQUESTS, None); assert_eq!(format!("{error}"), "Rate limit error: 429 Too Many Requests"); } #[test] fn test_check_response_status_success() { // Success responses (2xx) should return Ok(response) // This is tested via integration tests } #[test] fn test_status_code_categorization() { // Test that we correctly categorize different status codes // 429 and 403 -> RateLimit assert!(matches!( StatusCode::TOO_MANY_REQUESTS, StatusCode::TOO_MANY_REQUESTS )); assert!(matches!( StatusCode::FORBIDDEN, StatusCode::FORBIDDEN )); // 500-503 -> ServerError assert!(matches!( StatusCode::INTERNAL_SERVER_ERROR, StatusCode::INTERNAL_SERVER_ERROR )); assert!(matches!( StatusCode::BAD_GATEWAY, StatusCode::BAD_GATEWAY )); // 2xx -> Success assert!(StatusCode::OK.is_success()); assert!(StatusCode::CREATED.is_success()); // Other 4xx -> NonRetryable assert!(StatusCode::NOT_FOUND.is_client_error()); assert!(StatusCode::UNAUTHORIZED.is_client_error()); } #[test] fn test_calculate_delay_zero() { // First retry should be 30 seconds (30 * 2^0) assert_eq!(calculate_delay(0), 30); } #[test] fn test_calculate_delay_one() { // Second retry should be 60 seconds (30 * 2^1) assert_eq!(calculate_delay(1), 60); } #[test] fn test_calculate_delay_two() { // Third retry should be 120 seconds (30 * 2^2) assert_eq!(calculate_delay(2), 120); } #[test] fn test_calculate_delay_three() { // Fourth retry should be 240 seconds (30 * 2^3) assert_eq!(calculate_delay(3), 240); } #[test] fn test_calculate_delay_exponential_growth() { // Verify exponential backoff pattern let delays: Vec = (0..5).map(calculate_delay).collect(); assert_eq!(delays, vec![30, 60, 120, 240, 480]); } #[test] fn test_calculate_delay_large_values() { // Test that large values don't overflow let delay = calculate_delay(10); assert_eq!(delay, 30720); // 30 * 2^10 = 30720 } #[test] fn test_calculate_delay_consistency() { // Verify that calling the function multiple times with the same input // produces the same output (no random behavior) let delay1 = calculate_delay(5); let delay2 = calculate_delay(5); assert_eq!(delay1, delay2); } }