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 if attempt == 0 {
91 return self.initial_delay;
92 }
93
94 let safe_attempt = attempt.min(20) as i32;
97
98 let base_delay = self.initial_delay.as_millis() as f64 * self.multiplier.powi(safe_attempt);
99 let capped_delay = base_delay.min(self.max_delay.as_millis() as f64);
100
101 let final_delay = if self.jitter {
102 let mut rng = rand::rng();
108 let jitter_factor = rng.random_range(0.0..1.0);
109 capped_delay * jitter_factor
110 } else {
111 capped_delay
112 };
113
114 Duration::from_millis(final_delay as u64)
115 }
116}
117
118pub trait RetryClassifier: Send + Sync {
120 fn is_retryable(&self, error: &SeerError) -> bool;
122}
123
124#[derive(Debug, Clone, Default)]
137pub struct NetworkRetryClassifier;
138
139impl NetworkRetryClassifier {
140 pub fn new() -> Self {
141 Self
142 }
143}
144
145impl RetryClassifier for NetworkRetryClassifier {
146 fn is_retryable(&self, error: &SeerError) -> bool {
147 match error {
148 SeerError::Timeout(_) => true,
150 SeerError::WhoisConnectionFailed(_) => true,
151 SeerError::RateLimited(_) => true,
152
153 SeerError::ReqwestError { transient, .. } => *transient,
156
157 SeerError::WhoisError(msg) => {
159 let lower = msg.to_lowercase();
160 lower.contains("connection")
161 || lower.contains("timeout")
162 || lower.contains("refused")
163 || lower.contains("reset")
164 }
165
166 SeerError::RdapError(msg) => {
168 let lower = msg.to_lowercase();
169 lower.contains("status 5")
170 || lower.contains("status 429")
171 || lower.contains("timeout")
172 }
173
174 SeerError::RdapBootstrapError(msg) => {
176 let lower = msg.to_lowercase();
177 lower.contains("timeout") || lower.contains("connection")
178 }
179
180 SeerError::DnsError(msg) => {
182 let lower = msg.to_lowercase();
183 lower.contains("timeout") || lower.contains("temporary")
184 }
185
186 SeerError::HttpError(msg) => {
188 let lower = msg.to_lowercase();
189 lower.contains("timeout")
190 || lower.contains("connection")
191 || lower.contains("status 5")
192 || lower.contains("status 429")
193 }
194
195 SeerError::InvalidDomain(_) => false,
197 SeerError::DomainNotAllowed { .. } => false,
198 SeerError::InvalidIpAddress(_) => false,
199 SeerError::InvalidRecordType(_) => false,
200 SeerError::WhoisServerNotFound(_) => false,
201 SeerError::JsonError(_) => false,
202 SeerError::CertificateError(_) => false,
203 SeerError::SslError(_) => false,
204 SeerError::DnsResolverError(_) => false,
205 SeerError::BulkOperationError { .. } => false,
206 SeerError::LookupFailed { .. } => false,
207 SeerError::ConfigError(_) => false,
208 SeerError::InvalidInput(_) => false,
209 SeerError::RetryExhausted { last_error, .. } => self.is_retryable(last_error),
214 SeerError::Other(_) => false,
215 }
216 }
217}
218
219#[derive(Debug, Clone)]
221pub struct RetryExecutor<C: RetryClassifier> {
222 policy: RetryPolicy,
223 classifier: C,
224}
225
226impl RetryExecutor<NetworkRetryClassifier> {
227 pub fn new(policy: RetryPolicy) -> Self {
229 Self {
230 policy,
231 classifier: NetworkRetryClassifier::new(),
232 }
233 }
234}
235
236impl<C: RetryClassifier> RetryExecutor<C> {
237 pub fn with_classifier(policy: RetryPolicy, classifier: C) -> Self {
239 Self { policy, classifier }
240 }
241
242 pub async fn execute<F, Fut, T>(&self, mut operation: F) -> Result<T>
248 where
249 F: FnMut() -> Fut,
250 Fut: Future<Output = Result<T>>,
251 {
252 let mut last_error: Option<SeerError> = None;
253 let mut attempt = 0;
254
255 while attempt < self.policy.max_attempts {
256 match operation().await {
257 Ok(result) => return Ok(result),
258 Err(e) => {
259 let is_retryable = self.classifier.is_retryable(&e);
260 let attempts_remaining = self.policy.max_attempts - attempt - 1;
261
262 if !is_retryable || attempts_remaining == 0 {
263 if attempt > 0 {
264 debug!(
265 attempt = attempt + 1,
266 max_attempts = self.policy.max_attempts,
267 error = %e,
268 "Operation failed after retries"
269 );
270 }
271 return Err(if attempt > 0 {
272 SeerError::RetryExhausted {
273 attempts: attempt + 1,
274 last_error: Box::new(e),
275 }
276 } else {
277 e
278 });
279 }
280
281 let delay = self.policy.delay_for_attempt(attempt);
282 debug!(
283 attempt = attempt + 1,
284 max_attempts = self.policy.max_attempts,
285 delay_ms = delay.as_millis(),
286 error = %e,
287 "Retrying after transient error"
288 );
289
290 last_error = Some(e);
291 tokio::time::sleep(delay).await;
292 attempt += 1;
293 }
294 }
295 }
296
297 Err(last_error.unwrap_or_else(|| SeerError::Other("retry loop exited unexpectedly".into())))
299 }
300
301 pub async fn execute_once<F, Fut, T>(&self, operation: F) -> Result<T>
304 where
305 F: FnOnce() -> Fut,
306 Fut: Future<Output = Result<T>>,
307 {
308 operation().await
309 }
310}
311
312#[cfg(test)]
313mod tests {
314 use super::*;
315 use std::sync::atomic::{AtomicUsize, Ordering};
316 use std::sync::Arc;
317
318 #[test]
319 fn test_retry_policy_defaults() {
320 let policy = RetryPolicy::default();
321 assert_eq!(policy.max_attempts, 3);
322 assert_eq!(policy.initial_delay, Duration::from_millis(100));
323 assert_eq!(policy.max_delay, Duration::from_secs(5));
324 assert_eq!(policy.multiplier, 2.0);
325 assert!(policy.jitter);
326 }
327
328 #[test]
329 fn test_retry_policy_builder() {
330 let policy = RetryPolicy::new()
331 .with_max_attempts(5)
332 .with_initial_delay(Duration::from_millis(200))
333 .with_max_delay(Duration::from_secs(10))
334 .with_multiplier(3.0)
335 .with_jitter(false);
336
337 assert_eq!(policy.max_attempts, 5);
338 assert_eq!(policy.initial_delay, Duration::from_millis(200));
339 assert_eq!(policy.max_delay, Duration::from_secs(10));
340 assert_eq!(policy.multiplier, 3.0);
341 assert!(!policy.jitter);
342 }
343
344 #[test]
345 fn test_delay_calculation_no_jitter() {
346 let policy = RetryPolicy::new()
347 .with_initial_delay(Duration::from_millis(100))
348 .with_multiplier(2.0)
349 .with_max_delay(Duration::from_secs(10))
350 .with_jitter(false);
351
352 assert_eq!(policy.delay_for_attempt(0), Duration::from_millis(100));
353 assert_eq!(policy.delay_for_attempt(1), Duration::from_millis(200));
354 assert_eq!(policy.delay_for_attempt(2), Duration::from_millis(400));
355 assert_eq!(policy.delay_for_attempt(3), Duration::from_millis(800));
356 }
357
358 #[test]
359 fn test_delay_capped_at_max() {
360 let policy = RetryPolicy::new()
361 .with_initial_delay(Duration::from_secs(1))
362 .with_multiplier(10.0)
363 .with_max_delay(Duration::from_secs(5))
364 .with_jitter(false);
365
366 assert_eq!(policy.delay_for_attempt(2), Duration::from_secs(5));
368 }
369
370 #[test]
371 fn test_classifier_timeout_is_retryable() {
372 let classifier = NetworkRetryClassifier::new();
373 assert!(classifier.is_retryable(&SeerError::Timeout("test".to_string())));
374 }
375
376 #[test]
377 fn test_classifier_invalid_domain_not_retryable() {
378 let classifier = NetworkRetryClassifier::new();
379 assert!(!classifier.is_retryable(&SeerError::InvalidDomain("test".to_string())));
380 }
381
382 #[test]
383 fn test_classifier_server_not_found_not_retryable() {
384 let classifier = NetworkRetryClassifier::new();
385 assert!(!classifier.is_retryable(&SeerError::WhoisServerNotFound("test".to_string())));
386 }
387
388 #[test]
389 fn test_classifier_rate_limited_is_retryable() {
390 let classifier = NetworkRetryClassifier::new();
391 assert!(classifier.is_retryable(&SeerError::RateLimited("test".to_string())));
392 }
393
394 #[test]
395 fn test_classifier_reads_reqwest_transient_flag() {
396 let classifier = NetworkRetryClassifier::new();
399 assert!(classifier.is_retryable(&SeerError::ReqwestError {
400 message: "HTTP status server error (503 Service Unavailable)".to_string(),
401 transient: true,
402 }));
403 assert!(!classifier.is_retryable(&SeerError::ReqwestError {
404 message: "HTTP status client error (404 Not Found)".to_string(),
405 transient: false,
406 }));
407 }
408
409 #[tokio::test]
410 async fn test_executor_success_on_first_try() {
411 let policy = RetryPolicy::new().with_max_attempts(3);
412 let executor = RetryExecutor::new(policy);
413 let attempts = Arc::new(AtomicUsize::new(0));
414
415 let attempts_clone = attempts.clone();
416 let result: Result<&str> = executor
417 .execute(|| {
418 let a = attempts_clone.clone();
419 async move {
420 a.fetch_add(1, Ordering::SeqCst);
421 Ok("success")
422 }
423 })
424 .await;
425
426 assert!(result.is_ok());
427 assert_eq!(result.unwrap(), "success");
428 assert_eq!(attempts.load(Ordering::SeqCst), 1);
429 }
430
431 #[tokio::test]
432 async fn test_executor_retries_on_transient_error() {
433 let policy = RetryPolicy::new()
434 .with_max_attempts(3)
435 .with_initial_delay(Duration::from_millis(1))
436 .with_jitter(false);
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 let count = a.fetch_add(1, Ordering::SeqCst);
446 if count < 2 {
447 Err(SeerError::Timeout("test timeout".to_string()))
448 } else {
449 Ok("success after retries")
450 }
451 }
452 })
453 .await;
454
455 assert!(result.is_ok());
456 assert_eq!(result.unwrap(), "success after retries");
457 assert_eq!(attempts.load(Ordering::SeqCst), 3);
458 }
459
460 #[tokio::test]
461 async fn test_executor_no_retry_on_non_retryable_error() {
462 let policy = RetryPolicy::new()
463 .with_max_attempts(3)
464 .with_initial_delay(Duration::from_millis(1));
465 let executor = RetryExecutor::new(policy);
466 let attempts = Arc::new(AtomicUsize::new(0));
467
468 let attempts_clone = attempts.clone();
469 let result: Result<&str> = executor
470 .execute(|| {
471 let a = attempts_clone.clone();
472 async move {
473 a.fetch_add(1, Ordering::SeqCst);
474 Err(SeerError::InvalidDomain("bad.".to_string()))
475 }
476 })
477 .await;
478
479 assert!(result.is_err());
480 assert_eq!(attempts.load(Ordering::SeqCst), 1);
482 }
483
484 #[tokio::test]
485 async fn test_executor_exhausts_retries() {
486 let policy = RetryPolicy::new()
487 .with_max_attempts(3)
488 .with_initial_delay(Duration::from_millis(1))
489 .with_jitter(false);
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::Timeout("always fails".to_string()))
500 }
501 })
502 .await;
503
504 assert!(result.is_err());
505 assert_eq!(attempts.load(Ordering::SeqCst), 3);
506
507 match result.unwrap_err() {
509 SeerError::RetryExhausted { attempts, .. } => {
510 assert_eq!(attempts, 3);
511 }
512 other => panic!("Expected RetryExhausted, got {:?}", other),
513 }
514 }
515
516 #[test]
517 fn test_no_retry_policy() {
518 let policy = RetryPolicy::no_retry();
519 assert_eq!(policy.max_attempts, 1);
520 }
521
522 #[test]
523 fn test_delay_overflow_protection() {
524 let policy = RetryPolicy::new()
525 .with_initial_delay(Duration::from_millis(100))
526 .with_multiplier(2.0)
527 .with_max_delay(Duration::from_secs(5))
528 .with_jitter(false);
529
530 let delay_50 = policy.delay_for_attempt(50);
532 let delay_100 = policy.delay_for_attempt(100);
533 let delay_1000 = policy.delay_for_attempt(1000);
534
535 assert!(delay_50 <= Duration::from_secs(5));
537 assert!(delay_100 <= Duration::from_secs(5));
538 assert!(delay_1000 <= Duration::from_secs(5));
539 }
540
541 #[test]
542 fn retry_exhausted_is_retryable_if_inner_is() {
543 let classifier = NetworkRetryClassifier::new();
548
549 let retryable_inner = SeerError::Timeout("inner timed out".to_string());
550 let wrapped_retryable = SeerError::RetryExhausted {
551 attempts: 3,
552 last_error: Box::new(retryable_inner),
553 };
554 assert!(
555 classifier.is_retryable(&wrapped_retryable),
556 "RetryExhausted wrapping a retryable Timeout should be retryable",
557 );
558
559 let non_retryable_inner = SeerError::InvalidDomain("bad.".to_string());
560 let wrapped_non_retryable = SeerError::RetryExhausted {
561 attempts: 3,
562 last_error: Box::new(non_retryable_inner),
563 };
564 assert!(
565 !classifier.is_retryable(&wrapped_non_retryable),
566 "RetryExhausted wrapping a non-retryable InvalidDomain must not be retryable",
567 );
568 }
569}