Skip to main content

libdd_trace_utils/send_with_retry/
retry_strategy.rs

1// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4//! Types used when calling [`super::send_with_retry`] to configure the retry logic.
5
6use std::time::Duration;
7
8use libdd_capabilities::sleep::SleepCapability;
9
10/// Enum representing the type of backoff to use for the delay between retries.
11#[derive(Debug, Clone)]
12#[cfg_attr(test, derive(PartialEq))]
13pub enum RetryBackoffType {
14    /// Increases the delay by a fixed increment each attempt.
15    Linear,
16    /// The delay is constant for each attempt.
17    Constant,
18    /// The delay is doubled for each attempt.
19    Exponential,
20}
21
22/// Struct representing the retry strategy for sending data.
23///
24/// This struct contains the parameters that define how retries should be handled when sending data.
25/// It includes the maximum number of retries, the delay between retries, the type of backoff to
26/// use, and an optional jitter to add randomness to the delay.
27#[derive(Debug, Clone)]
28#[cfg_attr(test, derive(PartialEq))]
29pub struct RetryStrategy {
30    /// The maximum number of retries to attempt.
31    max_retries: u32,
32    // The minimum delay between retries.
33    delay_ms: Duration,
34    /// The type of backoff to use for the delay between retries.
35    backoff_type: RetryBackoffType,
36    /// An optional jitter to add randomness to the delay.
37    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    /// Creates a new `RetryStrategy` with the specified parameters.
53    ///
54    /// # Arguments
55    ///
56    /// * `max_retries`: The maximum number of retries to attempt.
57    /// * `delay_ms`: The minimum delay between retries, in milliseconds.
58    /// * `backoff_type`: The type of backoff to use for the delay between retries.
59    /// * `jitter`: An optional jitter to add randomness to the delay, in milliseconds.
60    ///
61    /// # Returns
62    ///
63    /// A `RetryStrategy` instance with the specified parameters.
64    ///
65    /// # Examples
66    ///
67    /// ```rust
68    /// use libdd_trace_utils::send_with_retry::{RetryBackoffType, RetryStrategy};
69    /// use std::time::Duration;
70    ///
71    /// let retry_strategy = RetryStrategy::new(5, 100, RetryBackoffType::Exponential, Some(50));
72    /// ```
73    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
84                .filter(|jitter_ms| *jitter_ms != 0)
85                .map(Duration::from_millis),
86        }
87    }
88    /// Delays the next request attempt based on the retry strategy.
89    ///
90    /// If a jitter duration is specified in the retry strategy, a random duration up to the jitter
91    /// value is added to the delay.
92    ///
93    /// # Arguments
94    ///
95    /// * `attempt`: The number of the current attempt (1-indexed).
96    /// * `capabilities`: Provides the sleep capability for the delay.
97    pub(crate) async fn delay<C: SleepCapability>(&self, attempt: u32, capabilities: &C) {
98        capabilities.sleep(self.delay_for_attempt(attempt)).await;
99    }
100
101    fn delay_for_attempt(&self, attempt: u32) -> Duration {
102        let retry_index = attempt.saturating_sub(1);
103        let multiplier = match self.backoff_type {
104            RetryBackoffType::Exponential => 2u32.checked_pow(retry_index).unwrap_or(u32::MAX),
105            RetryBackoffType::Constant => 1,
106            RetryBackoffType::Linear => retry_index.saturating_add(1),
107        };
108        let delay = self
109            .delay_ms
110            .checked_mul(multiplier)
111            .unwrap_or(Duration::MAX);
112
113        if let Some(jitter) = self.jitter {
114            let jitter_ms = u64::try_from(jitter.as_millis()).unwrap_or(u64::MAX);
115            let randomized_ms = rand::random::<u64>() % jitter_ms;
116            delay.saturating_add(Duration::from_millis(randomized_ms))
117        } else {
118            delay
119        }
120    }
121
122    /// Returns the maximum number of retries.
123    pub(crate) fn max_retries(&self) -> u32 {
124        self.max_retries
125    }
126}
127
128#[cfg(test)]
129// For tests RetryStrategy tests the observed delay should be approximate.
130mod tests {
131    use super::*;
132    use libdd_capabilities_impl::NativeSleepCapability;
133    use tokio::time::Instant;
134
135    // This tolerance is on the higher side to account for github's runners not having consistent
136    // performance. It shouldn't impact the quality of the tests since the most important aspect
137    // of the retry logic is we wait a minimum amount of time.
138    const RETRY_STRATEGY_TIME_TOLERANCE_MS: u64 = 100;
139
140    #[cfg_attr(miri, ignore)]
141    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
142    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
143    // with "multi-thread".
144    #[tokio::test(start_paused = true)]
145    async fn test_retry_strategy_constant() {
146        let retry_strategy = RetryStrategy {
147            max_retries: 5,
148            delay_ms: Duration::from_millis(100),
149            backoff_type: RetryBackoffType::Constant,
150            jitter: None,
151        };
152        let capabilities = NativeSleepCapability;
153
154        let start = Instant::now();
155        retry_strategy.delay(1, &capabilities).await;
156        let elapsed = start.elapsed();
157
158        assert!(
159            elapsed >= retry_strategy.delay_ms
160                && elapsed
161                    <= retry_strategy.delay_ms
162                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
163            "Elapsed time of {} ms was not within expected range",
164            elapsed.as_millis()
165        );
166
167        let start = Instant::now();
168        retry_strategy.delay(2, &capabilities).await;
169        let elapsed = start.elapsed();
170
171        assert!(
172            elapsed >= retry_strategy.delay_ms
173                && elapsed
174                    <= retry_strategy.delay_ms
175                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
176            "Elapsed time of {} ms was not within expected range",
177            elapsed.as_millis()
178        );
179    }
180
181    #[cfg_attr(miri, ignore)]
182    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
183    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
184    // with "multi-thread".
185    #[tokio::test(start_paused = true)]
186    async fn test_retry_strategy_linear() {
187        let retry_strategy = RetryStrategy {
188            max_retries: 5,
189            delay_ms: Duration::from_millis(100),
190            backoff_type: RetryBackoffType::Linear,
191            jitter: None,
192        };
193        let capabilities = NativeSleepCapability;
194
195        let start = Instant::now();
196        retry_strategy.delay(1, &capabilities).await;
197        let elapsed = start.elapsed();
198
199        assert!(
200            elapsed >= retry_strategy.delay_ms
201                && elapsed
202                    <= retry_strategy.delay_ms
203                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
204            "Elapsed time of {} ms was not within expected range",
205            elapsed.as_millis()
206        );
207
208        let start = Instant::now();
209        retry_strategy.delay(3, &capabilities).await;
210        let elapsed = start.elapsed();
211
212        // For the Linear strategy, the delay for the 3rd attempt should be delay_ms + (delay_ms *
213        // 2).
214        assert!(
215            elapsed >= retry_strategy.delay_ms + (retry_strategy.delay_ms * 2)
216                && elapsed
217                    <= retry_strategy.delay_ms
218                        + (retry_strategy.delay_ms * 2)
219                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
220            "Elapsed time of {} ms was not within expected range",
221            elapsed.as_millis()
222        );
223    }
224
225    #[cfg_attr(miri, ignore)]
226    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
227    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
228    // with "multi-thread".
229    #[tokio::test(start_paused = true)]
230    async fn test_retry_strategy_exponential() {
231        let retry_strategy = RetryStrategy {
232            max_retries: 5,
233            delay_ms: Duration::from_millis(100),
234            backoff_type: RetryBackoffType::Exponential,
235            jitter: None,
236        };
237        let capabilities = NativeSleepCapability;
238
239        let start = Instant::now();
240        retry_strategy.delay(1, &capabilities).await;
241        let elapsed = start.elapsed();
242
243        assert!(
244            elapsed >= retry_strategy.delay_ms
245                && elapsed
246                    <= retry_strategy.delay_ms
247                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
248            "Elapsed time of {} ms was not within expected range",
249            elapsed.as_millis()
250        );
251
252        let start = Instant::now();
253        retry_strategy.delay(3, &capabilities).await;
254        let elapsed = start.elapsed();
255        // For the Exponential strategy, the delay for the 3rd attempt should be delay_ms * 2^(3-1)
256        // = delay_ms * 4.
257        assert!(
258            elapsed >= retry_strategy.delay_ms * 4
259                && elapsed
260                    <= retry_strategy.delay_ms * 4
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    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
269    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
270    // with "multi-thread".
271    #[tokio::test(start_paused = true)]
272    async fn test_retry_strategy_jitter() {
273        let retry_strategy = RetryStrategy {
274            max_retries: 5,
275            delay_ms: Duration::from_millis(100),
276            backoff_type: RetryBackoffType::Constant,
277            jitter: Some(Duration::from_millis(50)),
278        };
279        let capabilities = NativeSleepCapability;
280
281        let start = Instant::now();
282        retry_strategy.delay(1, &capabilities).await;
283        let elapsed = start.elapsed();
284
285        // The delay should be between delay_ms and delay_ms + jitter
286        assert!(
287            elapsed >= retry_strategy.delay_ms
288                && elapsed
289                    <= retry_strategy.delay_ms
290                        + retry_strategy.jitter.unwrap()
291                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
292            "Elapsed time of {} ms was not within expected range",
293            elapsed.as_millis()
294        );
295    }
296
297    #[cfg_attr(miri, ignore)]
298    #[tokio::test]
299    async fn test_retry_strategy_max_retries() {
300        let retry_strategy = RetryStrategy {
301            max_retries: 17,
302            delay_ms: Duration::from_millis(100),
303            backoff_type: RetryBackoffType::Constant,
304            jitter: Some(Duration::from_millis(50)),
305        };
306
307        assert_eq!(
308            retry_strategy.max_retries(),
309            17,
310            "Max retries did not match expected value"
311        );
312    }
313
314    #[test]
315    fn test_retry_delays() {
316        let retry_strategy = RetryStrategy::new(3, 100, RetryBackoffType::Exponential, None);
317
318        assert_eq!(
319            retry_strategy.delay_for_attempt(1),
320            Duration::from_millis(100)
321        );
322        assert_eq!(
323            retry_strategy.delay_for_attempt(2),
324            Duration::from_millis(200)
325        );
326        assert_eq!(
327            retry_strategy.delay_for_attempt(3),
328            Duration::from_millis(400)
329        );
330    }
331
332    #[test]
333    fn test_zero_jitter_is_supported() {
334        let retry_strategy = RetryStrategy::new(1, 100, RetryBackoffType::Constant, Some(0));
335
336        assert_eq!(retry_strategy.jitter, None);
337        assert_eq!(
338            retry_strategy.delay_for_attempt(1),
339            Duration::from_millis(100)
340        );
341    }
342}