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) => {
179 let lower = msg.to_lowercase();
180 lower.contains("timeout")
181 || lower.contains("connection")
182 || lower.contains("registries failed")
183 }
184
185 SeerError::DnsError(msg) => {
187 let lower = msg.to_lowercase();
188 lower.contains("timeout") || lower.contains("temporary")
189 }
190
191 SeerError::HttpError(msg) => {
193 let lower = msg.to_lowercase();
194 lower.contains("timeout")
195 || lower.contains("connection")
196 || lower.contains("status 5")
197 || lower.contains("status 429")
198 }
199
200 SeerError::InvalidDomain(_) => false,
202 SeerError::DomainNotAllowed { .. } => false,
203 SeerError::InvalidIpAddress(_) => false,
204 SeerError::InvalidRecordType(_) => false,
205 SeerError::WhoisServerNotFound(_) => false,
206 SeerError::JsonError(_) => false,
207 SeerError::CertificateError(_) => false,
208 SeerError::SslError(_) => false,
209 SeerError::DnsResolverError(_) => false,
210 SeerError::BulkOperationError { .. } => false,
211 SeerError::LookupFailed { .. } => false,
212 SeerError::ConfigError(_) => false,
213 SeerError::InvalidInput(_) => false,
214 SeerError::RetryExhausted { last_error, .. } => self.is_retryable(last_error),
219 SeerError::Other(_) => false,
220 }
221 }
222}
223
224#[derive(Debug, Clone)]
226pub struct RetryExecutor<C: RetryClassifier> {
227 policy: RetryPolicy,
228 classifier: C,
229}
230
231impl RetryExecutor<NetworkRetryClassifier> {
232 pub fn new(policy: RetryPolicy) -> Self {
234 Self {
235 policy,
236 classifier: NetworkRetryClassifier::new(),
237 }
238 }
239}
240
241impl<C: RetryClassifier> RetryExecutor<C> {
242 pub fn with_classifier(policy: RetryPolicy, classifier: C) -> Self {
244 Self { policy, classifier }
245 }
246
247 pub async fn execute<F, Fut, T>(&self, mut operation: F) -> Result<T>
253 where
254 F: FnMut() -> Fut,
255 Fut: Future<Output = Result<T>>,
256 {
257 let mut last_error: Option<SeerError> = None;
258 let mut attempt = 0;
259
260 while attempt < self.policy.max_attempts {
261 match operation().await {
262 Ok(result) => return Ok(result),
263 Err(e) => {
264 let is_retryable = self.classifier.is_retryable(&e);
265 let attempts_remaining = self.policy.max_attempts - attempt - 1;
266
267 if !is_retryable || attempts_remaining == 0 {
268 if attempt > 0 {
269 debug!(
270 attempt = attempt + 1,
271 max_attempts = self.policy.max_attempts,
272 error = %e,
273 "Operation failed after retries"
274 );
275 }
276 return Err(if attempt > 0 {
277 SeerError::RetryExhausted {
278 attempts: attempt + 1,
279 last_error: Box::new(e),
280 }
281 } else {
282 e
283 });
284 }
285
286 let delay = self.policy.delay_for_attempt(attempt);
287 debug!(
288 attempt = attempt + 1,
289 max_attempts = self.policy.max_attempts,
290 delay_ms = delay.as_millis(),
291 error = %e,
292 "Retrying after transient error"
293 );
294
295 last_error = Some(e);
296 tokio::time::sleep(delay).await;
297 attempt += 1;
298 }
299 }
300 }
301
302 Err(last_error.unwrap_or_else(|| SeerError::Other("retry loop exited unexpectedly".into())))
304 }
305}
306
307#[cfg(test)]
308mod tests {
309 use super::*;
310 use std::sync::atomic::{AtomicUsize, Ordering};
311 use std::sync::Arc;
312
313 #[test]
314 fn test_retry_policy_defaults() {
315 let policy = RetryPolicy::default();
316 assert_eq!(policy.max_attempts, 3);
317 assert_eq!(policy.initial_delay, Duration::from_millis(100));
318 assert_eq!(policy.max_delay, Duration::from_secs(5));
319 assert_eq!(policy.multiplier, 2.0);
320 assert!(policy.jitter);
321 }
322
323 #[test]
324 fn test_retry_policy_builder() {
325 let policy = RetryPolicy::new()
326 .with_max_attempts(5)
327 .with_initial_delay(Duration::from_millis(200))
328 .with_max_delay(Duration::from_secs(10))
329 .with_multiplier(3.0)
330 .with_jitter(false);
331
332 assert_eq!(policy.max_attempts, 5);
333 assert_eq!(policy.initial_delay, Duration::from_millis(200));
334 assert_eq!(policy.max_delay, Duration::from_secs(10));
335 assert_eq!(policy.multiplier, 3.0);
336 assert!(!policy.jitter);
337 }
338
339 #[test]
340 fn test_delay_calculation_no_jitter() {
341 let policy = RetryPolicy::new()
342 .with_initial_delay(Duration::from_millis(100))
343 .with_multiplier(2.0)
344 .with_max_delay(Duration::from_secs(10))
345 .with_jitter(false);
346
347 assert_eq!(policy.delay_for_attempt(0), Duration::from_millis(100));
348 assert_eq!(policy.delay_for_attempt(1), Duration::from_millis(200));
349 assert_eq!(policy.delay_for_attempt(2), Duration::from_millis(400));
350 assert_eq!(policy.delay_for_attempt(3), Duration::from_millis(800));
351 }
352
353 #[test]
354 fn test_first_retry_is_jittered() {
355 let policy = RetryPolicy::new()
360 .with_initial_delay(Duration::from_millis(100))
361 .with_jitter(true);
362
363 let delays: Vec<Duration> = (0..64).map(|_| policy.delay_for_attempt(0)).collect();
364 assert!(delays.iter().all(|d| *d <= Duration::from_millis(100)));
366 assert!(
368 delays.iter().any(|d| *d != delays[0]),
369 "attempt-0 delay is constant — jitter not applied to the first retry"
370 );
371 }
372
373 #[test]
374 fn test_delay_capped_at_max() {
375 let policy = RetryPolicy::new()
376 .with_initial_delay(Duration::from_secs(1))
377 .with_multiplier(10.0)
378 .with_max_delay(Duration::from_secs(5))
379 .with_jitter(false);
380
381 assert_eq!(policy.delay_for_attempt(2), Duration::from_secs(5));
383 }
384
385 #[test]
386 fn test_classifier_timeout_is_retryable() {
387 let classifier = NetworkRetryClassifier::new();
388 assert!(classifier.is_retryable(&SeerError::Timeout("test".to_string())));
389 }
390
391 #[test]
392 fn test_classifier_invalid_domain_not_retryable() {
393 let classifier = NetworkRetryClassifier::new();
394 assert!(!classifier.is_retryable(&SeerError::InvalidDomain("test".to_string())));
395 }
396
397 #[test]
398 fn test_classifier_server_not_found_not_retryable() {
399 let classifier = NetworkRetryClassifier::new();
400 assert!(!classifier.is_retryable(&SeerError::WhoisServerNotFound("test".to_string())));
401 }
402
403 #[test]
404 fn test_classifier_rate_limited_is_retryable() {
405 let classifier = NetworkRetryClassifier::new();
406 assert!(classifier.is_retryable(&SeerError::RateLimited("test".to_string())));
407 }
408
409 #[test]
410 fn test_classifier_reads_reqwest_transient_flag() {
411 let classifier = NetworkRetryClassifier::new();
414 assert!(classifier.is_retryable(&SeerError::ReqwestError {
415 message: "HTTP status server error (503 Service Unavailable)".to_string(),
416 transient: true,
417 }));
418 assert!(!classifier.is_retryable(&SeerError::ReqwestError {
419 message: "HTTP status client error (404 Not Found)".to_string(),
420 transient: false,
421 }));
422 }
423
424 #[tokio::test]
425 async fn test_executor_success_on_first_try() {
426 let policy = RetryPolicy::new().with_max_attempts(3);
427 let executor = RetryExecutor::new(policy);
428 let attempts = Arc::new(AtomicUsize::new(0));
429
430 let attempts_clone = attempts.clone();
431 let result: Result<&str> = executor
432 .execute(|| {
433 let a = attempts_clone.clone();
434 async move {
435 a.fetch_add(1, Ordering::SeqCst);
436 Ok("success")
437 }
438 })
439 .await;
440
441 assert!(result.is_ok());
442 assert_eq!(result.unwrap(), "success");
443 assert_eq!(attempts.load(Ordering::SeqCst), 1);
444 }
445
446 #[tokio::test]
447 async fn test_executor_retries_on_transient_error() {
448 let policy = RetryPolicy::new()
449 .with_max_attempts(3)
450 .with_initial_delay(Duration::from_millis(1))
451 .with_jitter(false);
452 let executor = RetryExecutor::new(policy);
453 let attempts = Arc::new(AtomicUsize::new(0));
454
455 let attempts_clone = attempts.clone();
456 let result: Result<&str> = executor
457 .execute(|| {
458 let a = attempts_clone.clone();
459 async move {
460 let count = a.fetch_add(1, Ordering::SeqCst);
461 if count < 2 {
462 Err(SeerError::Timeout("test timeout".to_string()))
463 } else {
464 Ok("success after retries")
465 }
466 }
467 })
468 .await;
469
470 assert!(result.is_ok());
471 assert_eq!(result.unwrap(), "success after retries");
472 assert_eq!(attempts.load(Ordering::SeqCst), 3);
473 }
474
475 #[tokio::test]
476 async fn test_executor_no_retry_on_non_retryable_error() {
477 let policy = RetryPolicy::new()
478 .with_max_attempts(3)
479 .with_initial_delay(Duration::from_millis(1));
480 let executor = RetryExecutor::new(policy);
481 let attempts = Arc::new(AtomicUsize::new(0));
482
483 let attempts_clone = attempts.clone();
484 let result: Result<&str> = executor
485 .execute(|| {
486 let a = attempts_clone.clone();
487 async move {
488 a.fetch_add(1, Ordering::SeqCst);
489 Err(SeerError::InvalidDomain("bad.".to_string()))
490 }
491 })
492 .await;
493
494 assert!(result.is_err());
495 assert_eq!(attempts.load(Ordering::SeqCst), 1);
497 }
498
499 #[tokio::test]
500 async fn test_executor_exhausts_retries() {
501 let policy = RetryPolicy::new()
502 .with_max_attempts(3)
503 .with_initial_delay(Duration::from_millis(1))
504 .with_jitter(false);
505 let executor = RetryExecutor::new(policy);
506 let attempts = Arc::new(AtomicUsize::new(0));
507
508 let attempts_clone = attempts.clone();
509 let result: Result<&str> = executor
510 .execute(|| {
511 let a = attempts_clone.clone();
512 async move {
513 a.fetch_add(1, Ordering::SeqCst);
514 Err(SeerError::Timeout("always fails".to_string()))
515 }
516 })
517 .await;
518
519 assert!(result.is_err());
520 assert_eq!(attempts.load(Ordering::SeqCst), 3);
521
522 match result.unwrap_err() {
524 SeerError::RetryExhausted { attempts, .. } => {
525 assert_eq!(attempts, 3);
526 }
527 other => panic!("Expected RetryExhausted, got {:?}", other),
528 }
529 }
530
531 #[tokio::test]
532 async fn test_executor_with_custom_classifier() {
533 struct RetryEverything;
536 impl RetryClassifier for RetryEverything {
537 fn is_retryable(&self, _: &SeerError) -> bool {
538 true
539 }
540 }
541
542 let policy = RetryPolicy::new()
543 .with_max_attempts(3)
544 .with_initial_delay(Duration::from_millis(1))
545 .with_jitter(false);
546 let executor = RetryExecutor::with_classifier(policy, RetryEverything);
547 let attempts = Arc::new(AtomicUsize::new(0));
548
549 let attempts_clone = attempts.clone();
550 let result: Result<&str> = executor
551 .execute(|| {
552 let a = attempts_clone.clone();
553 async move {
554 a.fetch_add(1, Ordering::SeqCst);
555 Err(SeerError::InvalidDomain("bad.".to_string()))
556 }
557 })
558 .await;
559
560 assert!(matches!(
561 result,
562 Err(SeerError::RetryExhausted { attempts: 3, .. })
563 ));
564 assert_eq!(attempts.load(Ordering::SeqCst), 3);
565 }
566
567 #[test]
568 fn test_no_retry_policy() {
569 let policy = RetryPolicy::no_retry();
570 assert_eq!(policy.max_attempts, 1);
571 }
572
573 #[test]
574 fn test_delay_overflow_protection() {
575 let policy = RetryPolicy::new()
576 .with_initial_delay(Duration::from_millis(100))
577 .with_multiplier(2.0)
578 .with_max_delay(Duration::from_secs(5))
579 .with_jitter(false);
580
581 let delay_50 = policy.delay_for_attempt(50);
583 let delay_100 = policy.delay_for_attempt(100);
584 let delay_1000 = policy.delay_for_attempt(1000);
585
586 assert!(delay_50 <= Duration::from_secs(5));
588 assert!(delay_100 <= Duration::from_secs(5));
589 assert!(delay_1000 <= Duration::from_secs(5));
590 }
591
592 #[test]
593 fn retry_exhausted_is_retryable_if_inner_is() {
594 let classifier = NetworkRetryClassifier::new();
599
600 let retryable_inner = SeerError::Timeout("inner timed out".to_string());
601 let wrapped_retryable = SeerError::RetryExhausted {
602 attempts: 3,
603 last_error: Box::new(retryable_inner),
604 };
605 assert!(
606 classifier.is_retryable(&wrapped_retryable),
607 "RetryExhausted wrapping a retryable Timeout should be retryable",
608 );
609
610 let non_retryable_inner = SeerError::InvalidDomain("bad.".to_string());
611 let wrapped_non_retryable = SeerError::RetryExhausted {
612 attempts: 3,
613 last_error: Box::new(non_retryable_inner),
614 };
615 assert!(
616 !classifier.is_retryable(&wrapped_non_retryable),
617 "RetryExhausted wrapping a non-retryable InvalidDomain must not be retryable",
618 );
619 }
620
621 #[test]
622 fn bootstrap_all_registries_failed_is_retryable() {
623 let classifier = NetworkRetryClassifier::new();
627 assert!(classifier.is_retryable(&SeerError::RdapBootstrapError(
628 "all IANA bootstrap registries failed".to_string()
629 )));
630 assert!(!classifier.is_retryable(&SeerError::RdapBootstrapError(
632 "no RDAP server for example.ru".to_string()
633 )));
634 assert!(!classifier.is_retryable(&SeerError::RdapBootstrapError(
635 "bootstrap refresh throttled and no cache available".to_string()
636 )));
637 }
638}