1use crate::error::Result;
2use std::future::Future;
3use std::time::Duration;
4
5#[derive(Debug, Clone)]
6pub struct RetryPolicy {
7 max_attempts: usize,
8 base_delay: Duration,
9 max_delay: Duration,
10}
11
12impl RetryPolicy {
13 pub fn new(max_attempts: usize, base_delay: Duration, max_delay: Duration) -> Self {
14 Self {
15 max_attempts: max_attempts.max(1),
16 base_delay,
17 max_delay,
18 }
19 }
20
21 pub fn max_attempts(&self) -> usize {
22 self.max_attempts
23 }
24
25 fn backoff_delay(&self, attempt: usize) -> Duration {
26 if attempt <= 1 {
27 return self.base_delay.min(self.max_delay);
28 }
29
30 let mut delay = self.base_delay;
31 for _ in 1..attempt {
32 delay = delay.checked_mul(2).unwrap_or(self.max_delay);
33 if delay >= self.max_delay {
34 return self.max_delay;
35 }
36 }
37
38 delay.min(self.max_delay)
39 }
40
41 pub async fn execute<F, Fut, T>(&self, mut operation: F) -> Result<T>
42 where
43 F: FnMut() -> Fut,
44 Fut: Future<Output = Result<T>> + Send,
45 T: Send,
46 {
47 let mut attempt = 0;
48 loop {
49 attempt += 1;
50 match operation().await {
51 Ok(value) => return Ok(value),
52 Err(err) => {
53 if attempt >= self.max_attempts {
54 return Err(err);
55 }
56 let delay = self.backoff_delay(attempt);
57 if !delay.is_zero() {
58 tokio::time::sleep(delay).await;
59 } else {
60 tokio::task::yield_now().await;
61 }
62 }
63 }
64 }
65 }
66}
67
68impl Default for RetryPolicy {
69 fn default() -> Self {
70 Self {
71 max_attempts: 3,
72 base_delay: Duration::from_millis(100),
73 max_delay: Duration::from_secs(10),
74 }
75 }
76}
77
78#[cfg(test)]
79mod tests {
80 use super::*;
81 use crate::error::Error;
82 use std::sync::{Arc, Mutex};
83
84 #[tokio::test]
85 async fn retries_until_success() {
86 let policy = RetryPolicy::new(3, Duration::from_millis(0), Duration::from_millis(0));
87 let attempts = Arc::new(Mutex::new(0));
88 let result = policy
89 .execute({
90 let attempts = attempts.clone();
91 move || {
92 let attempts = attempts.clone();
93 Box::pin(async move {
94 let mut guard = attempts.lock().unwrap();
95 *guard += 1;
96 if *guard < 2 {
97 Err(Error::Transform("fail".into()))
98 } else {
99 Ok("ok")
100 }
101 })
102 }
103 })
104 .await;
105
106 assert_eq!(*attempts.lock().unwrap(), 2);
107 assert_eq!(result.unwrap(), "ok");
108 }
109
110 #[tokio::test]
111 async fn stops_after_max_attempts() {
112 let policy = RetryPolicy::new(2, Duration::from_millis(0), Duration::from_millis(0));
113 let attempts = Arc::new(Mutex::new(0));
114 let result: Result<()> = policy
115 .execute({
116 let attempts = attempts.clone();
117 move || {
118 let attempts = attempts.clone();
119 Box::pin(async move {
120 let mut guard = attempts.lock().unwrap();
121 *guard += 1;
122 Err(Error::Transform("still failing".into()))
123 })
124 }
125 })
126 .await;
127
128 assert!(result.is_err());
129 assert_eq!(*attempts.lock().unwrap(), 2);
130 }
131
132 #[test]
133 fn backoff_caps_at_max_delay() {
134 let policy = RetryPolicy::new(5, Duration::from_millis(100), Duration::from_millis(350));
135
136 assert_eq!(policy.backoff_delay(1), Duration::from_millis(100));
137 assert_eq!(policy.backoff_delay(2), Duration::from_millis(200));
138 assert_eq!(policy.backoff_delay(3), Duration::from_millis(350));
139 assert_eq!(policy.backoff_delay(4), Duration::from_millis(350));
140 }
141}