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;
7use tokio::time::sleep;
8
9/// Enum representing the type of backoff to use for the delay between retries.
10/// ```
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.map(Duration::from_millis),
84        }
85    }
86    /// Delays the next request attempt based on the retry strategy.
87    ///
88    /// If a jitter duration is specified in the retry strategy, a random duration up to the jitter
89    /// value is added to the delay.
90    ///
91    /// # Arguments
92    ///
93    /// * `attempt`: The number of the current attempt (1-indexed).
94    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    /// Returns the maximum number of retries.
110    pub(crate) fn max_retries(&self) -> u32 {
111        self.max_retries
112    }
113}
114
115#[cfg(test)]
116// For tests RetryStrategy tests the observed delay should be approximate.
117mod tests {
118    use super::*;
119    use tokio::time::Instant;
120
121    // This tolerance is on the higher side to account for github's runners not having consistent
122    // performance. It shouldn't impact the quality of the tests since the most important aspect
123    // of the retry logic is we wait a minimum amount of time.
124    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        // For the Linear strategy, the delay for the 3rd attempt should be delay_ms + (delay_ms *
191        // 2).
192        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        // For the Exponential strategy, the delay for the 3rd attempt should be delay_ms * 2^(3-1)
230        // = delay_ms * 4.
231        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        // The delay should be between delay_ms and delay_ms + jitter
256        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}