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.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    /// * `capabilities`: Provides the sleep capability for the delay.
95    pub(crate) async fn delay<C: SleepCapability>(&self, attempt: u32, capabilities: &C) {
96        let delay = match self.backoff_type {
97            RetryBackoffType::Exponential => self.delay_ms * 2u32.pow(attempt - 1),
98            RetryBackoffType::Constant => self.delay_ms,
99            RetryBackoffType::Linear => self.delay_ms + (self.delay_ms * (attempt - 1)),
100        };
101
102        if let Some(jitter) = self.jitter {
103            let jitter = rand::random::<u64>() % jitter.as_millis() as u64;
104            capabilities
105                .sleep(delay + Duration::from_millis(jitter))
106                .await;
107        } else {
108            capabilities.sleep(delay).await;
109        }
110    }
111
112    /// Returns the maximum number of retries.
113    pub(crate) fn max_retries(&self) -> u32 {
114        self.max_retries
115    }
116}
117
118#[cfg(test)]
119// For tests RetryStrategy tests the observed delay should be approximate.
120mod tests {
121    use super::*;
122    use libdd_capabilities_impl::NativeSleepCapability;
123    use tokio::time::Instant;
124
125    // This tolerance is on the higher side to account for github's runners not having consistent
126    // performance. It shouldn't impact the quality of the tests since the most important aspect
127    // of the retry logic is we wait a minimum amount of time.
128    const RETRY_STRATEGY_TIME_TOLERANCE_MS: u64 = 100;
129
130    #[cfg_attr(miri, ignore)]
131    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
132    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
133    // with "multi-thread".
134    #[tokio::test(start_paused = true)]
135    async fn test_retry_strategy_constant() {
136        let retry_strategy = RetryStrategy {
137            max_retries: 5,
138            delay_ms: Duration::from_millis(100),
139            backoff_type: RetryBackoffType::Constant,
140            jitter: None,
141        };
142        let capabilities = NativeSleepCapability;
143
144        let start = Instant::now();
145        retry_strategy.delay(1, &capabilities).await;
146        let elapsed = start.elapsed();
147
148        assert!(
149            elapsed >= retry_strategy.delay_ms
150                && elapsed
151                    <= retry_strategy.delay_ms
152                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
153            "Elapsed time of {} ms was not within expected range",
154            elapsed.as_millis()
155        );
156
157        let start = Instant::now();
158        retry_strategy.delay(2, &capabilities).await;
159        let elapsed = start.elapsed();
160
161        assert!(
162            elapsed >= retry_strategy.delay_ms
163                && elapsed
164                    <= retry_strategy.delay_ms
165                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
166            "Elapsed time of {} ms was not within expected range",
167            elapsed.as_millis()
168        );
169    }
170
171    #[cfg_attr(miri, ignore)]
172    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
173    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
174    // with "multi-thread".
175    #[tokio::test(start_paused = true)]
176    async fn test_retry_strategy_linear() {
177        let retry_strategy = RetryStrategy {
178            max_retries: 5,
179            delay_ms: Duration::from_millis(100),
180            backoff_type: RetryBackoffType::Linear,
181            jitter: None,
182        };
183        let capabilities = NativeSleepCapability;
184
185        let start = Instant::now();
186        retry_strategy.delay(1, &capabilities).await;
187        let elapsed = start.elapsed();
188
189        assert!(
190            elapsed >= retry_strategy.delay_ms
191                && elapsed
192                    <= retry_strategy.delay_ms
193                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
194            "Elapsed time of {} ms was not within expected range",
195            elapsed.as_millis()
196        );
197
198        let start = Instant::now();
199        retry_strategy.delay(3, &capabilities).await;
200        let elapsed = start.elapsed();
201
202        // For the Linear strategy, the delay for the 3rd attempt should be delay_ms + (delay_ms *
203        // 2).
204        assert!(
205            elapsed >= retry_strategy.delay_ms + (retry_strategy.delay_ms * 2)
206                && elapsed
207                    <= retry_strategy.delay_ms
208                        + (retry_strategy.delay_ms * 2)
209                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
210            "Elapsed time of {} ms was not within expected range",
211            elapsed.as_millis()
212        );
213    }
214
215    #[cfg_attr(miri, ignore)]
216    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
217    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
218    // with "multi-thread".
219    #[tokio::test(start_paused = true)]
220    async fn test_retry_strategy_exponential() {
221        let retry_strategy = RetryStrategy {
222            max_retries: 5,
223            delay_ms: Duration::from_millis(100),
224            backoff_type: RetryBackoffType::Exponential,
225            jitter: None,
226        };
227        let capabilities = NativeSleepCapability;
228
229        let start = Instant::now();
230        retry_strategy.delay(1, &capabilities).await;
231        let elapsed = start.elapsed();
232
233        assert!(
234            elapsed >= retry_strategy.delay_ms
235                && elapsed
236                    <= retry_strategy.delay_ms
237                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
238            "Elapsed time of {} ms was not within expected range",
239            elapsed.as_millis()
240        );
241
242        let start = Instant::now();
243        retry_strategy.delay(3, &capabilities).await;
244        let elapsed = start.elapsed();
245        // For the Exponential strategy, the delay for the 3rd attempt should be delay_ms * 2^(3-1)
246        // = delay_ms * 4.
247        assert!(
248            elapsed >= retry_strategy.delay_ms * 4
249                && elapsed
250                    <= retry_strategy.delay_ms * 4
251                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
252            "Elapsed time of {} ms was not within expected range",
253            elapsed.as_millis()
254        );
255    }
256
257    #[cfg_attr(miri, ignore)]
258    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
259    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
260    // with "multi-thread".
261    #[tokio::test(start_paused = true)]
262    async fn test_retry_strategy_jitter() {
263        let retry_strategy = RetryStrategy {
264            max_retries: 5,
265            delay_ms: Duration::from_millis(100),
266            backoff_type: RetryBackoffType::Constant,
267            jitter: Some(Duration::from_millis(50)),
268        };
269        let capabilities = NativeSleepCapability;
270
271        let start = Instant::now();
272        retry_strategy.delay(1, &capabilities).await;
273        let elapsed = start.elapsed();
274
275        // The delay should be between delay_ms and delay_ms + jitter
276        assert!(
277            elapsed >= retry_strategy.delay_ms
278                && elapsed
279                    <= retry_strategy.delay_ms
280                        + retry_strategy.jitter.unwrap()
281                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
282            "Elapsed time of {} ms was not within expected range",
283            elapsed.as_millis()
284        );
285    }
286
287    #[cfg_attr(miri, ignore)]
288    #[tokio::test]
289    async fn test_retry_strategy_max_retries() {
290        let retry_strategy = RetryStrategy {
291            max_retries: 17,
292            delay_ms: Duration::from_millis(100),
293            backoff_type: RetryBackoffType::Constant,
294            jitter: Some(Duration::from_millis(50)),
295        };
296
297        assert_eq!(
298            retry_strategy.max_retries(),
299            17,
300            "Max retries did not match expected value"
301        );
302    }
303}