do_memory_core/retry/
mod.rs

1pub mod budget;
2pub mod limiter;
3
4use std::future::Future;
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::time::Duration;
8use tokio::time::sleep;
9use tracing::warn;
10
11pub use budget::RetryBudget;
12pub use limiter::ConcurrencyLimiter;
13
14pub trait Retryable {
15    fn is_recoverable(&self) -> bool;
16}
17
18impl Retryable for crate::error::Error {
19    fn is_recoverable(&self) -> bool {
20        self.is_recoverable()
21    }
22}
23
24#[derive(Debug, Default)]
25pub struct RetryMetrics {
26    total: AtomicU64,
27    success: AtomicU64,
28    failure: AtomicU64,
29}
30
31impl RetryMetrics {
32    #[must_use]
33    pub fn new() -> Self {
34        Self {
35            total: AtomicU64::new(0),
36            success: AtomicU64::new(0),
37            failure: AtomicU64::new(0),
38        }
39    }
40
41    pub fn record_retry(&self, succeeded: bool) {
42        self.total.fetch_add(1, Ordering::SeqCst);
43        if succeeded {
44            self.success.fetch_add(1, Ordering::SeqCst);
45        } else {
46            self.failure.fetch_add(1, Ordering::SeqCst);
47        }
48    }
49
50    #[must_use]
51    pub fn total(&self) -> u64 {
52        self.total.load(Ordering::SeqCst)
53    }
54
55    #[must_use]
56    pub fn success_count(&self) -> u64 {
57        self.success.load(Ordering::SeqCst)
58    }
59
60    #[must_use]
61    pub fn failure_count(&self) -> u64 {
62        self.failure.load(Ordering::SeqCst)
63    }
64}
65
66#[derive(Debug, Clone)]
67pub struct RetryConfig {
68    pub max_retries: u32,
69    pub base_delay: Duration,
70    pub max_delay: Duration,
71    pub jitter_factor: f64,
72    pub max_retry_budget: Option<u32>,
73    pub max_concurrent_retries: Option<usize>,
74    pub budget_window: Option<Duration>,
75    pub retry_queue_timeout: Option<Duration>,
76}
77
78impl Default for RetryConfig {
79    fn default() -> Self {
80        Self {
81            max_retries: 3,
82            base_delay: Duration::from_millis(100),
83            max_delay: Duration::from_secs(5),
84            jitter_factor: 0.25,
85            max_retry_budget: None,
86            max_concurrent_retries: None,
87            budget_window: None,
88            retry_queue_timeout: None,
89        }
90    }
91}
92
93impl RetryConfig {
94    #[must_use]
95    pub fn new() -> Self {
96        Self::default()
97    }
98
99    #[must_use]
100    pub fn with_max_retries(mut self, max_retries: u32) -> Self {
101        self.max_retries = max_retries;
102        self
103    }
104
105    #[must_use]
106    pub fn with_base_delay(mut self, base_delay: Duration) -> Self {
107        self.base_delay = base_delay;
108        self
109    }
110
111    #[must_use]
112    pub fn with_max_delay(mut self, max_delay: Duration) -> Self {
113        self.max_delay = max_delay;
114        self
115    }
116
117    #[must_use]
118    pub fn with_jitter(mut self, factor: f64) -> Self {
119        self.jitter_factor = factor;
120        self
121    }
122
123    #[must_use]
124    pub fn with_retry_budget(mut self, budget: u32) -> Self {
125        self.max_retry_budget = Some(budget);
126        self
127    }
128
129    /// Set the maximum number of concurrent retries process-wide.
130    #[must_use]
131    pub fn with_max_concurrent_retries(mut self, max: usize) -> Self {
132        self.max_concurrent_retries = Some(max);
133        self
134    }
135
136    /// Set the sliding window duration for the shared retry budget.
137    #[must_use]
138    pub fn with_budget_window(mut self, window: Duration) -> Self {
139        self.budget_window = Some(window);
140        self
141    }
142
143    /// Set a timeout for the retry queue (prevents unbounded queuing).
144    #[must_use]
145    pub fn with_retry_queue_timeout(mut self, timeout: Duration) -> Self {
146        self.retry_queue_timeout = Some(timeout);
147        self
148    }
149}
150
151pub struct RetryPolicy {
152    config: RetryConfig,
153    metrics: Option<RetryMetrics>,
154    retry_budget: Option<NonZeroBudget>,
155    shared_budget: Option<Arc<RetryBudget>>,
156    limiter: Option<Arc<ConcurrencyLimiter>>,
157}
158
159struct NonZeroBudget {
160    remaining: u32,
161}
162
163impl RetryPolicy {
164    #[must_use]
165    pub fn new() -> Self {
166        Self {
167            config: RetryConfig::default(),
168            metrics: None,
169            retry_budget: None,
170            shared_budget: None,
171            limiter: None,
172        }
173    }
174
175    #[must_use]
176    pub fn with_config(config: RetryConfig) -> Self {
177        let retry_budget = config
178            .max_retry_budget
179            .map(|b| NonZeroBudget { remaining: b });
180
181        let shared_budget = config
182            .budget_window
183            .zip(config.max_retry_budget)
184            .map(|(window, tokens)| RetryBudget::new(tokens, window));
185
186        let limiter = config.max_concurrent_retries.map(ConcurrencyLimiter::new);
187
188        Self {
189            config,
190            metrics: None,
191            retry_budget,
192            shared_budget,
193            limiter,
194        }
195    }
196
197    #[must_use]
198    pub fn with_metrics(mut self, metrics: RetryMetrics) -> Self {
199        self.metrics = Some(metrics);
200        self
201    }
202
203    #[must_use]
204    pub fn with_retry_budget(mut self, budget: u32) -> Self {
205        self.retry_budget = Some(NonZeroBudget { remaining: budget });
206        self
207    }
208
209    /// Attach a shared process-wide retry budget.
210    #[must_use]
211    pub fn with_shared_budget(mut self, budget: Arc<RetryBudget>) -> Self {
212        self.shared_budget = Some(budget);
213        self
214    }
215
216    /// Attach a concurrency limiter for retries.
217    #[must_use]
218    pub fn with_limiter(mut self, limiter: Arc<ConcurrencyLimiter>) -> Self {
219        self.limiter = Some(limiter);
220        self
221    }
222
223    fn calculate_delay(&self, attempt: u32) -> Duration {
224        let exp_delay = self.config.base_delay * (2u32.pow(attempt.saturating_sub(1)));
225        let delay = std::cmp::min(exp_delay, self.config.max_delay);
226
227        if self.config.jitter_factor > 0.0 {
228            let jitter_range = delay.as_millis() as f64 * self.config.jitter_factor;
229            let jitter = (rand::random::<f64>() - 0.5) * 2.0 * jitter_range;
230            let adjusted_ms = (delay.as_millis() as f64 + jitter).max(0.0);
231            Duration::from_millis(adjusted_ms as u64)
232        } else {
233            delay
234        }
235    }
236
237    fn can_retry(&mut self) -> bool {
238        if let Some(ref mut budget) = self.retry_budget {
239            if budget.remaining == 0 {
240                return false;
241            }
242            budget.remaining = budget.remaining.saturating_sub(1);
243        }
244        if let Some(ref shared) = self.shared_budget {
245            if !shared.acquire() {
246                return false;
247            }
248        }
249        true
250    }
251
252    fn record_success(&self, attempt: u32) {
253        if attempt > 0 {
254            if let Some(ref metrics) = self.metrics {
255                metrics.record_retry(true);
256            }
257        }
258    }
259
260    fn record_failure(&self, attempt: u32) {
261        if attempt > 0 {
262            if let Some(ref metrics) = self.metrics {
263                metrics.record_retry(false);
264            }
265        }
266    }
267
268    pub async fn execute<F, T, E, Fut>(&mut self, operation: F) -> Result<T, E>
269    where
270        F: Fn() -> Fut,
271        Fut: Future<Output = Result<T, E>>,
272        E: Retryable + std::error::Error + Send + Sync + 'static,
273        E: std::fmt::Debug,
274    {
275        let _permit = if let Some(ref limiter) = self.limiter {
276            Some(limiter.acquire().await)
277        } else {
278            None
279        };
280
281        let mut attempt = 0;
282
283        loop {
284            match operation().await {
285                Ok(result) => {
286                    self.record_success(attempt);
287                    return Ok(result);
288                }
289                Err(e) => {
290                    let error_ref = &e;
291                    let is_recoverable = error_ref.is_recoverable();
292
293                    if !is_recoverable || !self.can_retry() || attempt >= self.config.max_retries {
294                        return Err(e);
295                    }
296
297                    attempt += 1;
298                    let delay = self.calculate_delay(attempt);
299
300                    self.record_failure(attempt);
301
302                    warn!(
303                        "Retry attempt {}/{} failed: {:?}, retrying in {:?}",
304                        attempt, self.config.max_retries, error_ref, delay
305                    );
306
307                    sleep(delay).await;
308                }
309            }
310        }
311    }
312
313    pub fn execute_sync<F, T, E>(mut self, operation: F) -> Result<T, E>
314    where
315        F: Fn() -> Result<T, E>,
316        E: Retryable + std::error::Error + Send + Sync + 'static,
317        E: std::fmt::Debug,
318    {
319        let mut attempt = 0;
320
321        loop {
322            match operation() {
323                Ok(result) => {
324                    self.record_success(attempt);
325                    return Ok(result);
326                }
327                Err(e) => {
328                    let is_recoverable = e.is_recoverable();
329
330                    if !is_recoverable || !self.can_retry() || attempt >= self.config.max_retries {
331                        return Err(e);
332                    }
333
334                    attempt += 1;
335                    let delay = self.calculate_delay(attempt);
336
337                    self.record_failure(attempt);
338
339                    warn!(
340                        "Retry attempt {}/{} failed: {:?}, retrying in {:?}",
341                        attempt, self.config.max_retries, e, delay
342                    );
343
344                    std::thread::sleep(delay);
345                }
346            }
347        }
348    }
349}
350
351impl Default for RetryPolicy {
352    fn default() -> Self {
353        Self::new()
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use std::sync::Arc;
360    use std::sync::atomic::{AtomicUsize, Ordering};
361    use std::time::Duration;
362    use tokio::time::timeout;
363
364    use super::{RetryConfig, RetryPolicy};
365
366    #[derive(Debug)]
367    struct TestError(bool);
368
369    impl std::error::Error for TestError {}
370
371    impl std::fmt::Display for TestError {
372        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
373            write!(f, "TestError({})", self.0)
374        }
375    }
376
377    impl super::Retryable for TestError {
378        fn is_recoverable(&self) -> bool {
379            self.0
380        }
381    }
382
383    #[allow(clippy::excessive_nesting)]
384    #[tokio::test]
385    async fn test_retry_success_first_attempt() {
386        let call_count = AtomicUsize::new(0);
387        let mut policy = RetryPolicy::new();
388
389        let result = policy
390            .execute(|| {
391                let count = call_count.fetch_add(1, Ordering::SeqCst);
392                async move {
393                    if count == 0 {
394                        Ok("success")
395                    } else {
396                        Err(TestError(true))
397                    }
398                }
399            })
400            .await;
401
402        assert_eq!(result.unwrap(), "success");
403        assert_eq!(call_count.load(Ordering::SeqCst), 1);
404    }
405
406    #[allow(clippy::excessive_nesting)]
407    #[tokio::test]
408    async fn test_retry_success_after_failures() {
409        let call_count = AtomicUsize::new(0);
410        let mut policy = RetryPolicy::with_config(
411            RetryConfig::new()
412                .with_max_retries(3)
413                .with_base_delay(Duration::from_millis(10)),
414        );
415
416        let result = policy
417            .execute(|| {
418                let count = call_count.fetch_add(1, Ordering::SeqCst);
419                async move {
420                    if count < 2 {
421                        Err(TestError(true))
422                    } else {
423                        Ok("success")
424                    }
425                }
426            })
427            .await;
428
429        assert_eq!(result.unwrap(), "success");
430        assert_eq!(call_count.load(Ordering::SeqCst), 3);
431    }
432
433    #[allow(clippy::excessive_nesting)]
434    #[tokio::test]
435    async fn test_retry_non_recoverable_error() {
436        let call_count = AtomicUsize::new(0);
437        let mut policy = RetryPolicy::with_config(RetryConfig::new().with_max_retries(3));
438
439        let result = policy
440            .execute(|| {
441                let count = call_count.fetch_add(1, Ordering::SeqCst);
442                async move {
443                    if count < 5 {
444                        Err(TestError(false))
445                    } else {
446                        Ok("success")
447                    }
448                }
449            })
450            .await;
451
452        assert!(result.is_err());
453        assert_eq!(call_count.load(Ordering::SeqCst), 1);
454    }
455
456    #[allow(clippy::excessive_nesting)]
457    #[tokio::test]
458    async fn test_retry_max_retries_exceeded() {
459        let call_count = AtomicUsize::new(0);
460        let mut policy = RetryPolicy::with_config(
461            RetryConfig::new()
462                .with_max_retries(2)
463                .with_base_delay(Duration::from_millis(5)),
464        );
465
466        let result = policy
467            .execute(|| {
468                call_count.fetch_add(1, Ordering::SeqCst);
469                async move { Err::<(), _>(TestError(true)) }
470            })
471            .await;
472
473        assert!(result.is_err());
474        assert_eq!(call_count.load(Ordering::SeqCst), 3);
475    }
476
477    #[allow(clippy::excessive_nesting)]
478    fn make_failing_closure(
479        cc: Arc<AtomicUsize>,
480    ) -> impl Fn() -> std::pin::Pin<
481        Box<dyn std::future::Future<Output = Result<&'static str, TestError>> + Send>,
482    > {
483        move || {
484            cc.fetch_add(1, Ordering::SeqCst);
485            let cc2 = cc.clone();
486            Box::pin(async move {
487                if cc2.load(Ordering::SeqCst) < 3 {
488                    Err(TestError(true))
489                } else {
490                    Ok("success")
491                }
492            })
493        }
494    }
495
496    #[allow(clippy::excessive_nesting)]
497    #[tokio::test]
498    async fn test_retry_with_budget() {
499        let call_count = Arc::new(AtomicUsize::new(0));
500        let mut policy =
501            RetryPolicy::with_config(RetryConfig::new().with_max_retries(10)).with_retry_budget(2);
502
503        let result = policy
504            .execute(make_failing_closure(call_count.clone()))
505            .await;
506
507        // Budget of 2 allows initial + 2 retries = 3 attempts
508        // Closure fails on attempts 1-2, succeeds on attempt 3
509        assert_eq!(result.unwrap(), "success");
510        assert_eq!(call_count.load(Ordering::SeqCst), 3);
511    }
512
513    #[tokio::test]
514    async fn test_retry_sync() {
515        let call_count = AtomicUsize::new(0);
516        let policy = RetryPolicy::with_config(
517            RetryConfig::new()
518                .with_max_retries(3)
519                .with_base_delay(Duration::from_millis(10)),
520        );
521
522        let result = policy.execute_sync(|| {
523            let count = call_count.fetch_add(1, Ordering::SeqCst);
524            if count < 2 {
525                Err(TestError(true))
526            } else {
527                Ok("success")
528            }
529        });
530
531        assert_eq!(result.unwrap(), "success");
532        assert_eq!(call_count.load(Ordering::SeqCst), 3);
533    }
534
535    #[allow(clippy::excessive_nesting)]
536    #[tokio::test]
537    async fn test_retry_with_jitter() {
538        let call_count = Arc::new(AtomicUsize::new(0));
539        let mut policy = RetryPolicy::with_config(
540            RetryConfig::new()
541                .with_max_retries(3)
542                .with_base_delay(Duration::from_millis(100))
543                .with_jitter(0.5),
544        );
545
546        let start = std::time::Instant::now();
547        let result = policy
548            .execute(make_failing_closure(call_count.clone()))
549            .await;
550
551        let elapsed = start.elapsed();
552        assert_eq!(result.unwrap(), "success");
553        assert_eq!(call_count.load(Ordering::SeqCst), 3);
554        assert!(elapsed >= Duration::from_millis(100));
555        assert!(elapsed < Duration::from_millis(500));
556    }
557
558    #[tokio::test]
559    async fn test_retry_timeout() {
560        let mut policy = RetryPolicy::with_config(
561            RetryConfig::new()
562                .with_max_retries(10)
563                .with_base_delay(Duration::from_secs(10)),
564        );
565
566        let result = timeout(
567            Duration::from_millis(100),
568            policy.execute(|| async move {
569                tokio::time::sleep(Duration::from_secs(1)).await;
570                Ok::<&str, TestError>("success")
571            }),
572        )
573        .await;
574
575        assert!(result.is_err());
576    }
577}
do_memory_core/retry/mod.rs

do_memory_core/retry/
mod.rs
