1use std::future::Future;
7use std::time::Duration;
8
9use rand::RngExt;
10use tracing::debug;
11
12use crate::error::{Result, SeerError};
13
14#[derive(Debug, Clone)]
16pub struct RetryPolicy {
17 pub max_attempts: usize,
19 pub initial_delay: Duration,
21 pub max_delay: Duration,
23 pub multiplier: f64,
25 pub jitter: bool,
27}
28
29impl Default for RetryPolicy {
30 fn default() -> Self {
31 Self {
32 max_attempts: 3,
33 initial_delay: Duration::from_millis(100),
34 max_delay: Duration::from_secs(5),
35 multiplier: 2.0,
36 jitter: true,
37 }
38 }
39}
40
41impl RetryPolicy {
42 pub fn new() -> Self {
44 Self::default()
45 }
46
47 pub fn with_max_attempts(mut self, attempts: usize) -> Self {
49 self.max_attempts = attempts.max(1);
50 self
51 }
52
53 pub fn with_initial_delay(mut self, delay: Duration) -> Self {
55 self.initial_delay = delay;
56 self
57 }
58
59 pub fn with_max_delay(mut self, delay: Duration) -> Self {
61 self.max_delay = delay;
62 self
63 }
64
65 pub fn with_multiplier(mut self, multiplier: f64) -> Self {
67 self.multiplier = multiplier.max(1.0);
68 self
69 }
70
71 pub fn with_jitter(mut self, jitter: bool) -> Self {
73 self.jitter = jitter;
74 self
75 }
76
77 pub fn no_retry() -> Self {
79 Self {
80 max_attempts: 1,
81 ..Self::default()
82 }
83 }
84
85 pub fn delay_for_attempt(&self, attempt: usize) -> Duration {
90 let safe_attempt = attempt.min(20) as i32;
93
94 let base_delay = self.initial_delay.as_millis() as f64 * self.multiplier.powi(safe_attempt);
98 let capped_delay = base_delay.min(self.max_delay.as_millis() as f64);
99
100 let final_delay = if self.jitter {
101 let mut rng = rand::rng();
107 let jitter_factor = rng.random_range(0.0..1.0);
108 capped_delay * jitter_factor
109 } else {
110 capped_delay
111 };
112
113 Duration::from_millis(final_delay as u64)
114 }
115}
116
117pub trait RetryClassifier: Send + Sync {
119 fn is_retryable(&self, error: &SeerError) -> bool;
121}
122
123#[derive(Debug, Clone, Default)]
136pub struct NetworkRetryClassifier;
137
138impl NetworkRetryClassifier {
139 pub fn new() -> Self {
140 Self
141 }
142}
143
144impl RetryClassifier for NetworkRetryClassifier {
145 fn is_retryable(&self, error: &SeerError) -> bool {
146 match error {
147 SeerError::Timeout(_) => true,
149 SeerError::WhoisConnectionFailed(_) => true,
150 SeerError::RateLimited(_) => true,
151
152 SeerError::ReqwestError { transient, .. } => *transient,
155
156 SeerError::WhoisError(msg) => {
158 let lower = msg.to_lowercase();
159 lower.contains("connection")
160 || lower.contains("timeout")
161 || lower.contains("refused")
162 || lower.contains("reset")
163 }
164
165 SeerError::RdapError(msg) => {
167 let lower = msg.to_lowercase();
168 lower.contains("status 5")
169 || lower.contains("status 429")
170 || lower.contains("timeout")
171 }
172
173 SeerError::RdapBootstrapError(msg) => {
175 let lower = msg.to_lowercase();
176 lower.contains("timeout") || lower.contains("connection")
177 }
178
179 SeerError::DnsError(msg) => {
181 let lower = msg.to_lowercase();
182 lower.contains("timeout") || lower.contains("temporary")
183 }
184
185 SeerError::HttpError(msg) => {
187 let lower = msg.to_lowercase();
188 lower.contains("timeout")
189 || lower.contains("connection")
190 || lower.contains("status 5")
191 || lower.contains("status 429")
192 }
193
194 SeerError::InvalidDomain(_) => false,
196 SeerError::DomainNotAllowed { .. } => false,
197 SeerError::InvalidIpAddress(_) => false,
198 SeerError::InvalidRecordType(_) => false,
199 SeerError::WhoisServerNotFound(_) => false,
200 SeerError::JsonError(_) => false,
201 SeerError::CertificateError(_) => false,
202 SeerError::SslError(_) => false,
203 SeerError::DnsResolverError(_) => false,
204 SeerError::BulkOperationError { .. } => false,
205 SeerError::LookupFailed { .. } => false,
206 SeerError::ConfigError(_) => false,
207 SeerError::InvalidInput(_) => false,
208 SeerError::RetryExhausted { last_error, .. } => self.is_retryable(last_error),
213 SeerError::Other(_) => false,
214 }
215 }
216}
217
218#[derive(Debug, Clone)]
220pub struct RetryExecutor<C: RetryClassifier> {
221 policy: RetryPolicy,
222 classifier: C,
223}
224
225impl RetryExecutor<NetworkRetryClassifier> {
226 pub fn new(policy: RetryPolicy) -> Self {
228 Self {
229 policy,
230 classifier: NetworkRetryClassifier::new(),
231 }
232 }
233}
234
235impl<C: RetryClassifier> RetryExecutor<C> {
236 pub fn with_classifier(policy: RetryPolicy, classifier: C) -> Self {
238 Self { policy, classifier }
239 }
240
241 pub async fn execute<F, Fut, T>(&self, mut operation: F) -> Result<T>
247 where
248 F: FnMut() -> Fut,
249 Fut: Future<Output = Result<T>>,
250 {
251 let mut last_error: Option<SeerError> = None;
252 let mut attempt = 0;
253
254 while attempt < self.policy.max_attempts {
255 match operation().await {
256 Ok(result) => return Ok(result),
257 Err(e) => {
258 let is_retryable = self.classifier.is_retryable(&e);
259 let attempts_remaining = self.policy.max_attempts - attempt - 1;
260
261 if !is_retryable || attempts_remaining == 0 {
262 if attempt > 0 {
263 debug!(
264 attempt = attempt + 1,
265 max_attempts = self.policy.max_attempts,
266 error = %e,
267 "Operation failed after retries"
268 );
269 }
270 return Err(if attempt > 0 {
271 SeerError::RetryExhausted {
272 attempts: attempt + 1,
273 last_error: Box::new(e),
274 }
275 } else {
276 e
277 });
278 }
279
280 let delay = self.policy.delay_for_attempt(attempt);
281 debug!(
282 attempt = attempt + 1,
283 max_attempts = self.policy.max_attempts,
284 delay_ms = delay.as_millis(),
285 error = %e,
286 "Retrying after transient error"
287 );
288
289 last_error = Some(e);
290 tokio::time::sleep(delay).await;
291 attempt += 1;
292 }
293 }
294 }
295
296 Err(last_error.unwrap_or_else(|| SeerError::Other("retry loop exited unexpectedly".into())))
298 }
299
300 pub async fn execute_once<F, Fut, T>(&self, operation: F) -> Result<T>
303 where
304 F: FnOnce() -> Fut,
305 Fut: Future<Output = Result<T>>,
306 {
307 operation().await
308 }
309}
310
311#[cfg(test)]
312mod tests {
313 use super::*;
314 use std::sync::atomic::{AtomicUsize, Ordering};
315 use std::sync::Arc;
316
317 #[test]
318 fn test_retry_policy_defaults() {
319 let policy = RetryPolicy::default();
320 assert_eq!(policy.max_attempts, 3);
321 assert_eq!(policy.initial_delay, Duration::from_millis(100));
322 assert_eq!(policy.max_delay, Duration::from_secs(5));
323 assert_eq!(policy.multiplier, 2.0);
324 assert!(policy.jitter);
325 }
326
327 #[test]
328 fn test_retry_policy_builder() {
329 let policy = RetryPolicy::new()
330 .with_max_attempts(5)
331 .with_initial_delay(Duration::from_millis(200))
332 .with_max_delay(Duration::from_secs(10))
333 .with_multiplier(3.0)
334 .with_jitter(false);
335
336 assert_eq!(policy.max_attempts, 5);
337 assert_eq!(policy.initial_delay, Duration::from_millis(200));
338 assert_eq!(policy.max_delay, Duration::from_secs(10));
339 assert_eq!(policy.multiplier, 3.0);
340 assert!(!policy.jitter);
341 }
342
343 #[test]
344 fn test_delay_calculation_no_jitter() {
345 let policy = RetryPolicy::new()
346 .with_initial_delay(Duration::from_millis(100))
347 .with_multiplier(2.0)
348 .with_max_delay(Duration::from_secs(10))
349 .with_jitter(false);
350
351 assert_eq!(policy.delay_for_attempt(0), Duration::from_millis(100));
352 assert_eq!(policy.delay_for_attempt(1), Duration::from_millis(200));
353 assert_eq!(policy.delay_for_attempt(2), Duration::from_millis(400));
354 assert_eq!(policy.delay_for_attempt(3), Duration::from_millis(800));
355 }
356
357 #[test]
358 fn test_first_retry_is_jittered() {
359 let policy = RetryPolicy::new()
364 .with_initial_delay(Duration::from_millis(100))
365 .with_jitter(true);
366
367 let delays: Vec<Duration> = (0..64).map(|_| policy.delay_for_attempt(0)).collect();
368 assert!(delays.iter().all(|d| *d <= Duration::from_millis(100)));
370 assert!(
372 delays.iter().any(|d| *d != delays[0]),
373 "attempt-0 delay is constant — jitter not applied to the first retry"
374 );
375 }
376
377 #[test]
378 fn test_delay_capped_at_max() {
379 let policy = RetryPolicy::new()
380 .with_initial_delay(Duration::from_secs(1))
381 .with_multiplier(10.0)
382 .with_max_delay(Duration::from_secs(5))
383 .with_jitter(false);
384
385 assert_eq!(policy.delay_for_attempt(2), Duration::from_secs(5));
387 }
388
389 #[test]
390 fn test_classifier_timeout_is_retryable() {
391 let classifier = NetworkRetryClassifier::new();
392 assert!(classifier.is_retryable(&SeerError::Timeout("test".to_string())));
393 }
394
395 #[test]
396 fn test_classifier_invalid_domain_not_retryable() {
397 let classifier = NetworkRetryClassifier::new();
398 assert!(!classifier.is_retryable(&SeerError::InvalidDomain("test".to_string())));
399 }
400
401 #[test]
402 fn test_classifier_server_not_found_not_retryable() {
403 let classifier = NetworkRetryClassifier::new();
404 assert!(!classifier.is_retryable(&SeerError::WhoisServerNotFound("test".to_string())));
405 }
406
407 #[test]
408 fn test_classifier_rate_limited_is_retryable() {
409 let classifier = NetworkRetryClassifier::new();
410 assert!(classifier.is_retryable(&SeerError::RateLimited("test".to_string())));
411 }
412
413 #[test]
414 fn test_classifier_reads_reqwest_transient_flag() {
415 let classifier = NetworkRetryClassifier::new();
418 assert!(classifier.is_retryable(&SeerError::ReqwestError {
419 message: "HTTP status server error (503 Service Unavailable)".to_string(),
420 transient: true,
421 }));
422 assert!(!classifier.is_retryable(&SeerError::ReqwestError {
423 message: "HTTP status client error (404 Not Found)".to_string(),
424 transient: false,
425 }));
426 }
427
428 #[tokio::test]
429 async fn test_executor_success_on_first_try() {
430 let policy = RetryPolicy::new().with_max_attempts(3);
431 let executor = RetryExecutor::new(policy);
432 let attempts = Arc::new(AtomicUsize::new(0));
433
434 let attempts_clone = attempts.clone();
435 let result: Result<&str> = executor
436 .execute(|| {
437 let a = attempts_clone.clone();
438 async move {
439 a.fetch_add(1, Ordering::SeqCst);
440 Ok("success")
441 }
442 })
443 .await;
444
445 assert!(result.is_ok());
446 assert_eq!(result.unwrap(), "success");
447 assert_eq!(attempts.load(Ordering::SeqCst), 1);
448 }
449
450 #[tokio::test]
451 async fn test_executor_retries_on_transient_error() {
452 let policy = RetryPolicy::new()
453 .with_max_attempts(3)
454 .with_initial_delay(Duration::from_millis(1))
455 .with_jitter(false);
456 let executor = RetryExecutor::new(policy);
457 let attempts = Arc::new(AtomicUsize::new(0));
458
459 let attempts_clone = attempts.clone();
460 let result: Result<&str> = executor
461 .execute(|| {
462 let a = attempts_clone.clone();
463 async move {
464 let count = a.fetch_add(1, Ordering::SeqCst);
465 if count < 2 {
466 Err(SeerError::Timeout("test timeout".to_string()))
467 } else {
468 Ok("success after retries")
469 }
470 }
471 })
472 .await;
473
474 assert!(result.is_ok());
475 assert_eq!(result.unwrap(), "success after retries");
476 assert_eq!(attempts.load(Ordering::SeqCst), 3);
477 }
478
479 #[tokio::test]
480 async fn test_executor_no_retry_on_non_retryable_error() {
481 let policy = RetryPolicy::new()
482 .with_max_attempts(3)
483 .with_initial_delay(Duration::from_millis(1));
484 let executor = RetryExecutor::new(policy);
485 let attempts = Arc::new(AtomicUsize::new(0));
486
487 let attempts_clone = attempts.clone();
488 let result: Result<&str> = executor
489 .execute(|| {
490 let a = attempts_clone.clone();
491 async move {
492 a.fetch_add(1, Ordering::SeqCst);
493 Err(SeerError::InvalidDomain("bad.".to_string()))
494 }
495 })
496 .await;
497
498 assert!(result.is_err());
499 assert_eq!(attempts.load(Ordering::SeqCst), 1);
501 }
502
503 #[tokio::test]
504 async fn test_executor_exhausts_retries() {
505 let policy = RetryPolicy::new()
506 .with_max_attempts(3)
507 .with_initial_delay(Duration::from_millis(1))
508 .with_jitter(false);
509 let executor = RetryExecutor::new(policy);
510 let attempts = Arc::new(AtomicUsize::new(0));
511
512 let attempts_clone = attempts.clone();
513 let result: Result<&str> = executor
514 .execute(|| {
515 let a = attempts_clone.clone();
516 async move {
517 a.fetch_add(1, Ordering::SeqCst);
518 Err(SeerError::Timeout("always fails".to_string()))
519 }
520 })
521 .await;
522
523 assert!(result.is_err());
524 assert_eq!(attempts.load(Ordering::SeqCst), 3);
525
526 match result.unwrap_err() {
528 SeerError::RetryExhausted { attempts, .. } => {
529 assert_eq!(attempts, 3);
530 }
531 other => panic!("Expected RetryExhausted, got {:?}", other),
532 }
533 }
534
535 #[test]
536 fn test_no_retry_policy() {
537 let policy = RetryPolicy::no_retry();
538 assert_eq!(policy.max_attempts, 1);
539 }
540
541 #[test]
542 fn test_delay_overflow_protection() {
543 let policy = RetryPolicy::new()
544 .with_initial_delay(Duration::from_millis(100))
545 .with_multiplier(2.0)
546 .with_max_delay(Duration::from_secs(5))
547 .with_jitter(false);
548
549 let delay_50 = policy.delay_for_attempt(50);
551 let delay_100 = policy.delay_for_attempt(100);
552 let delay_1000 = policy.delay_for_attempt(1000);
553
554 assert!(delay_50 <= Duration::from_secs(5));
556 assert!(delay_100 <= Duration::from_secs(5));
557 assert!(delay_1000 <= Duration::from_secs(5));
558 }
559
560 #[test]
561 fn retry_exhausted_is_retryable_if_inner_is() {
562 let classifier = NetworkRetryClassifier::new();
567
568 let retryable_inner = SeerError::Timeout("inner timed out".to_string());
569 let wrapped_retryable = SeerError::RetryExhausted {
570 attempts: 3,
571 last_error: Box::new(retryable_inner),
572 };
573 assert!(
574 classifier.is_retryable(&wrapped_retryable),
575 "RetryExhausted wrapping a retryable Timeout should be retryable",
576 );
577
578 let non_retryable_inner = SeerError::InvalidDomain("bad.".to_string());
579 let wrapped_non_retryable = SeerError::RetryExhausted {
580 attempts: 3,
581 last_error: Box::new(non_retryable_inner),
582 };
583 assert!(
584 !classifier.is_retryable(&wrapped_non_retryable),
585 "RetryExhausted wrapping a non-retryable InvalidDomain must not be retryable",
586 );
587 }
588}