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 pub async fn execute_once<F, Fut, T>(&self, operation: F) -> Result<T>
309 where
310 F: FnOnce() -> Fut,
311 Fut: Future<Output = Result<T>>,
312 {
313 operation().await
314 }
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320 use std::sync::atomic::{AtomicUsize, Ordering};
321 use std::sync::Arc;
322
323 #[test]
324 fn test_retry_policy_defaults() {
325 let policy = RetryPolicy::default();
326 assert_eq!(policy.max_attempts, 3);
327 assert_eq!(policy.initial_delay, Duration::from_millis(100));
328 assert_eq!(policy.max_delay, Duration::from_secs(5));
329 assert_eq!(policy.multiplier, 2.0);
330 assert!(policy.jitter);
331 }
332
333 #[test]
334 fn test_retry_policy_builder() {
335 let policy = RetryPolicy::new()
336 .with_max_attempts(5)
337 .with_initial_delay(Duration::from_millis(200))
338 .with_max_delay(Duration::from_secs(10))
339 .with_multiplier(3.0)
340 .with_jitter(false);
341
342 assert_eq!(policy.max_attempts, 5);
343 assert_eq!(policy.initial_delay, Duration::from_millis(200));
344 assert_eq!(policy.max_delay, Duration::from_secs(10));
345 assert_eq!(policy.multiplier, 3.0);
346 assert!(!policy.jitter);
347 }
348
349 #[test]
350 fn test_delay_calculation_no_jitter() {
351 let policy = RetryPolicy::new()
352 .with_initial_delay(Duration::from_millis(100))
353 .with_multiplier(2.0)
354 .with_max_delay(Duration::from_secs(10))
355 .with_jitter(false);
356
357 assert_eq!(policy.delay_for_attempt(0), Duration::from_millis(100));
358 assert_eq!(policy.delay_for_attempt(1), Duration::from_millis(200));
359 assert_eq!(policy.delay_for_attempt(2), Duration::from_millis(400));
360 assert_eq!(policy.delay_for_attempt(3), Duration::from_millis(800));
361 }
362
363 #[test]
364 fn test_first_retry_is_jittered() {
365 let policy = RetryPolicy::new()
370 .with_initial_delay(Duration::from_millis(100))
371 .with_jitter(true);
372
373 let delays: Vec<Duration> = (0..64).map(|_| policy.delay_for_attempt(0)).collect();
374 assert!(delays.iter().all(|d| *d <= Duration::from_millis(100)));
376 assert!(
378 delays.iter().any(|d| *d != delays[0]),
379 "attempt-0 delay is constant — jitter not applied to the first retry"
380 );
381 }
382
383 #[test]
384 fn test_delay_capped_at_max() {
385 let policy = RetryPolicy::new()
386 .with_initial_delay(Duration::from_secs(1))
387 .with_multiplier(10.0)
388 .with_max_delay(Duration::from_secs(5))
389 .with_jitter(false);
390
391 assert_eq!(policy.delay_for_attempt(2), Duration::from_secs(5));
393 }
394
395 #[test]
396 fn test_classifier_timeout_is_retryable() {
397 let classifier = NetworkRetryClassifier::new();
398 assert!(classifier.is_retryable(&SeerError::Timeout("test".to_string())));
399 }
400
401 #[test]
402 fn test_classifier_invalid_domain_not_retryable() {
403 let classifier = NetworkRetryClassifier::new();
404 assert!(!classifier.is_retryable(&SeerError::InvalidDomain("test".to_string())));
405 }
406
407 #[test]
408 fn test_classifier_server_not_found_not_retryable() {
409 let classifier = NetworkRetryClassifier::new();
410 assert!(!classifier.is_retryable(&SeerError::WhoisServerNotFound("test".to_string())));
411 }
412
413 #[test]
414 fn test_classifier_rate_limited_is_retryable() {
415 let classifier = NetworkRetryClassifier::new();
416 assert!(classifier.is_retryable(&SeerError::RateLimited("test".to_string())));
417 }
418
419 #[test]
420 fn test_classifier_reads_reqwest_transient_flag() {
421 let classifier = NetworkRetryClassifier::new();
424 assert!(classifier.is_retryable(&SeerError::ReqwestError {
425 message: "HTTP status server error (503 Service Unavailable)".to_string(),
426 transient: true,
427 }));
428 assert!(!classifier.is_retryable(&SeerError::ReqwestError {
429 message: "HTTP status client error (404 Not Found)".to_string(),
430 transient: false,
431 }));
432 }
433
434 #[tokio::test]
435 async fn test_executor_success_on_first_try() {
436 let policy = RetryPolicy::new().with_max_attempts(3);
437 let executor = RetryExecutor::new(policy);
438 let attempts = Arc::new(AtomicUsize::new(0));
439
440 let attempts_clone = attempts.clone();
441 let result: Result<&str> = executor
442 .execute(|| {
443 let a = attempts_clone.clone();
444 async move {
445 a.fetch_add(1, Ordering::SeqCst);
446 Ok("success")
447 }
448 })
449 .await;
450
451 assert!(result.is_ok());
452 assert_eq!(result.unwrap(), "success");
453 assert_eq!(attempts.load(Ordering::SeqCst), 1);
454 }
455
456 #[tokio::test]
457 async fn test_executor_retries_on_transient_error() {
458 let policy = RetryPolicy::new()
459 .with_max_attempts(3)
460 .with_initial_delay(Duration::from_millis(1))
461 .with_jitter(false);
462 let executor = RetryExecutor::new(policy);
463 let attempts = Arc::new(AtomicUsize::new(0));
464
465 let attempts_clone = attempts.clone();
466 let result: Result<&str> = executor
467 .execute(|| {
468 let a = attempts_clone.clone();
469 async move {
470 let count = a.fetch_add(1, Ordering::SeqCst);
471 if count < 2 {
472 Err(SeerError::Timeout("test timeout".to_string()))
473 } else {
474 Ok("success after retries")
475 }
476 }
477 })
478 .await;
479
480 assert!(result.is_ok());
481 assert_eq!(result.unwrap(), "success after retries");
482 assert_eq!(attempts.load(Ordering::SeqCst), 3);
483 }
484
485 #[tokio::test]
486 async fn test_executor_no_retry_on_non_retryable_error() {
487 let policy = RetryPolicy::new()
488 .with_max_attempts(3)
489 .with_initial_delay(Duration::from_millis(1));
490 let executor = RetryExecutor::new(policy);
491 let attempts = Arc::new(AtomicUsize::new(0));
492
493 let attempts_clone = attempts.clone();
494 let result: Result<&str> = executor
495 .execute(|| {
496 let a = attempts_clone.clone();
497 async move {
498 a.fetch_add(1, Ordering::SeqCst);
499 Err(SeerError::InvalidDomain("bad.".to_string()))
500 }
501 })
502 .await;
503
504 assert!(result.is_err());
505 assert_eq!(attempts.load(Ordering::SeqCst), 1);
507 }
508
509 #[tokio::test]
510 async fn test_executor_exhausts_retries() {
511 let policy = RetryPolicy::new()
512 .with_max_attempts(3)
513 .with_initial_delay(Duration::from_millis(1))
514 .with_jitter(false);
515 let executor = RetryExecutor::new(policy);
516 let attempts = Arc::new(AtomicUsize::new(0));
517
518 let attempts_clone = attempts.clone();
519 let result: Result<&str> = executor
520 .execute(|| {
521 let a = attempts_clone.clone();
522 async move {
523 a.fetch_add(1, Ordering::SeqCst);
524 Err(SeerError::Timeout("always fails".to_string()))
525 }
526 })
527 .await;
528
529 assert!(result.is_err());
530 assert_eq!(attempts.load(Ordering::SeqCst), 3);
531
532 match result.unwrap_err() {
534 SeerError::RetryExhausted { attempts, .. } => {
535 assert_eq!(attempts, 3);
536 }
537 other => panic!("Expected RetryExhausted, got {:?}", other),
538 }
539 }
540
541 #[test]
542 fn test_no_retry_policy() {
543 let policy = RetryPolicy::no_retry();
544 assert_eq!(policy.max_attempts, 1);
545 }
546
547 #[test]
548 fn test_delay_overflow_protection() {
549 let policy = RetryPolicy::new()
550 .with_initial_delay(Duration::from_millis(100))
551 .with_multiplier(2.0)
552 .with_max_delay(Duration::from_secs(5))
553 .with_jitter(false);
554
555 let delay_50 = policy.delay_for_attempt(50);
557 let delay_100 = policy.delay_for_attempt(100);
558 let delay_1000 = policy.delay_for_attempt(1000);
559
560 assert!(delay_50 <= Duration::from_secs(5));
562 assert!(delay_100 <= Duration::from_secs(5));
563 assert!(delay_1000 <= Duration::from_secs(5));
564 }
565
566 #[test]
567 fn retry_exhausted_is_retryable_if_inner_is() {
568 let classifier = NetworkRetryClassifier::new();
573
574 let retryable_inner = SeerError::Timeout("inner timed out".to_string());
575 let wrapped_retryable = SeerError::RetryExhausted {
576 attempts: 3,
577 last_error: Box::new(retryable_inner),
578 };
579 assert!(
580 classifier.is_retryable(&wrapped_retryable),
581 "RetryExhausted wrapping a retryable Timeout should be retryable",
582 );
583
584 let non_retryable_inner = SeerError::InvalidDomain("bad.".to_string());
585 let wrapped_non_retryable = SeerError::RetryExhausted {
586 attempts: 3,
587 last_error: Box::new(non_retryable_inner),
588 };
589 assert!(
590 !classifier.is_retryable(&wrapped_non_retryable),
591 "RetryExhausted wrapping a non-retryable InvalidDomain must not be retryable",
592 );
593 }
594
595 #[test]
596 fn bootstrap_all_registries_failed_is_retryable() {
597 let classifier = NetworkRetryClassifier::new();
601 assert!(classifier.is_retryable(&SeerError::RdapBootstrapError(
602 "all IANA bootstrap registries failed".to_string()
603 )));
604 assert!(!classifier.is_retryable(&SeerError::RdapBootstrapError(
606 "no RDAP server for example.ru".to_string()
607 )));
608 assert!(!classifier.is_retryable(&SeerError::RdapBootstrapError(
609 "bootstrap refresh throttled and no cache available".to_string()
610 )));
611 }
612}