libdd_trace_utils/send_with_retry/
retry_strategy.rs1use std::time::Duration;
7use tokio::time::sleep;
8
9#[derive(Debug, Clone)]
12#[cfg_attr(test, derive(PartialEq))]
13pub enum RetryBackoffType {
14 Linear,
16 Constant,
18 Exponential,
20}
21
22#[derive(Debug, Clone)]
28#[cfg_attr(test, derive(PartialEq))]
29pub struct RetryStrategy {
30 max_retries: u32,
32 delay_ms: Duration,
34 backoff_type: RetryBackoffType,
36 jitter: Option<Duration>,
38}
39
40impl Default for RetryStrategy {
41 fn default() -> Self {
42 RetryStrategy {
43 max_retries: 5,
44 delay_ms: Duration::from_millis(100),
45 backoff_type: RetryBackoffType::Exponential,
46 jitter: None,
47 }
48 }
49}
50
51impl RetryStrategy {
52 pub fn new(
74 max_retries: u32,
75 delay_ms: u64,
76 backoff_type: RetryBackoffType,
77 jitter: Option<u64>,
78 ) -> RetryStrategy {
79 RetryStrategy {
80 max_retries,
81 delay_ms: Duration::from_millis(delay_ms),
82 backoff_type,
83 jitter: jitter.map(Duration::from_millis),
84 }
85 }
86 pub(crate) async fn delay(&self, attempt: u32) {
95 let delay = match self.backoff_type {
96 RetryBackoffType::Exponential => self.delay_ms * 2u32.pow(attempt - 1),
97 RetryBackoffType::Constant => self.delay_ms,
98 RetryBackoffType::Linear => self.delay_ms + (self.delay_ms * (attempt - 1)),
99 };
100
101 if let Some(jitter) = self.jitter {
102 let jitter = rand::random::<u64>() % jitter.as_millis() as u64;
103 sleep(delay + Duration::from_millis(jitter)).await;
104 } else {
105 sleep(delay).await;
106 }
107 }
108
109 pub(crate) fn max_retries(&self) -> u32 {
111 self.max_retries
112 }
113}
114
115#[cfg(test)]
116mod tests {
118 use super::*;
119 use tokio::time::Instant;
120
121 const RETRY_STRATEGY_TIME_TOLERANCE_MS: u64 = 100;
125
126 #[cfg_attr(miri, ignore)]
127 #[tokio::test]
128 async fn test_retry_strategy_constant() {
129 let retry_strategy = RetryStrategy {
130 max_retries: 5,
131 delay_ms: Duration::from_millis(100),
132 backoff_type: RetryBackoffType::Constant,
133 jitter: None,
134 };
135
136 let start = Instant::now();
137 retry_strategy.delay(1).await;
138 let elapsed = start.elapsed();
139
140 assert!(
141 elapsed >= retry_strategy.delay_ms
142 && elapsed
143 <= retry_strategy.delay_ms
144 + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
145 "Elapsed time of {} ms was not within expected range",
146 elapsed.as_millis()
147 );
148
149 let start = Instant::now();
150 retry_strategy.delay(2).await;
151 let elapsed = start.elapsed();
152
153 assert!(
154 elapsed >= retry_strategy.delay_ms
155 && elapsed
156 <= retry_strategy.delay_ms
157 + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
158 "Elapsed time of {} ms was not within expected range",
159 elapsed.as_millis()
160 );
161 }
162
163 #[cfg_attr(miri, ignore)]
164 #[tokio::test]
165 async fn test_retry_strategy_linear() {
166 let retry_strategy = RetryStrategy {
167 max_retries: 5,
168 delay_ms: Duration::from_millis(100),
169 backoff_type: RetryBackoffType::Linear,
170 jitter: None,
171 };
172
173 let start = Instant::now();
174 retry_strategy.delay(1).await;
175 let elapsed = start.elapsed();
176
177 assert!(
178 elapsed >= retry_strategy.delay_ms
179 && elapsed
180 <= retry_strategy.delay_ms
181 + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
182 "Elapsed time of {} ms was not within expected range",
183 elapsed.as_millis()
184 );
185
186 let start = Instant::now();
187 retry_strategy.delay(3).await;
188 let elapsed = start.elapsed();
189
190 assert!(
193 elapsed >= retry_strategy.delay_ms + (retry_strategy.delay_ms * 2)
194 && elapsed
195 <= retry_strategy.delay_ms
196 + (retry_strategy.delay_ms * 2)
197 + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
198 "Elapsed time of {} ms was not within expected range",
199 elapsed.as_millis()
200 );
201 }
202
203 #[cfg_attr(miri, ignore)]
204 #[tokio::test]
205 async fn test_retry_strategy_exponential() {
206 let retry_strategy = RetryStrategy {
207 max_retries: 5,
208 delay_ms: Duration::from_millis(100),
209 backoff_type: RetryBackoffType::Exponential,
210 jitter: None,
211 };
212
213 let start = Instant::now();
214 retry_strategy.delay(1).await;
215 let elapsed = start.elapsed();
216
217 assert!(
218 elapsed >= retry_strategy.delay_ms
219 && elapsed
220 <= retry_strategy.delay_ms
221 + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
222 "Elapsed time of {} ms was not within expected range",
223 elapsed.as_millis()
224 );
225
226 let start = Instant::now();
227 retry_strategy.delay(3).await;
228 let elapsed = start.elapsed();
229 assert!(
232 elapsed >= retry_strategy.delay_ms * 4
233 && elapsed
234 <= retry_strategy.delay_ms * 4
235 + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
236 "Elapsed time of {} ms was not within expected range",
237 elapsed.as_millis()
238 );
239 }
240
241 #[cfg_attr(miri, ignore)]
242 #[tokio::test]
243 async fn test_retry_strategy_jitter() {
244 let retry_strategy = RetryStrategy {
245 max_retries: 5,
246 delay_ms: Duration::from_millis(100),
247 backoff_type: RetryBackoffType::Constant,
248 jitter: Some(Duration::from_millis(50)),
249 };
250
251 let start = Instant::now();
252 retry_strategy.delay(1).await;
253 let elapsed = start.elapsed();
254
255 assert!(
257 elapsed >= retry_strategy.delay_ms
258 && elapsed
259 <= retry_strategy.delay_ms
260 + retry_strategy.jitter.unwrap()
261 + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
262 "Elapsed time of {} ms was not within expected range",
263 elapsed.as_millis()
264 );
265 }
266
267 #[cfg_attr(miri, ignore)]
268 #[tokio::test]
269 async fn test_retry_strategy_max_retries() {
270 let retry_strategy = RetryStrategy {
271 max_retries: 17,
272 delay_ms: Duration::from_millis(100),
273 backoff_type: RetryBackoffType::Constant,
274 jitter: Some(Duration::from_millis(50)),
275 };
276
277 assert_eq!(
278 retry_strategy.max_retries(),
279 17,
280 "Max retries did not match expected value"
281 );
282 }
283}