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    #[tokio::test]
132    async fn test_retry_strategy_constant() {
133        let retry_strategy = RetryStrategy {
134            max_retries: 5,
135            delay_ms: Duration::from_millis(100),
136            backoff_type: RetryBackoffType::Constant,
137            jitter: None,
138        };
139        let capabilities = NativeSleepCapability;
140
141        let start = Instant::now();
142        retry_strategy.delay(1, &capabilities).await;
143        let elapsed = start.elapsed();
144
145        assert!(
146            elapsed >= retry_strategy.delay_ms
147                && elapsed
148                    <= retry_strategy.delay_ms
149                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
150            "Elapsed time of {} ms was not within expected range",
151            elapsed.as_millis()
152        );
153
154        let start = Instant::now();
155        retry_strategy.delay(2, &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
168    #[cfg_attr(miri, ignore)]
169    #[tokio::test]
170    async fn test_retry_strategy_linear() {
171        let retry_strategy = RetryStrategy {
172            max_retries: 5,
173            delay_ms: Duration::from_millis(100),
174            backoff_type: RetryBackoffType::Linear,
175            jitter: None,
176        };
177        let capabilities = NativeSleepCapability;
178
179        let start = Instant::now();
180        retry_strategy.delay(1, &capabilities).await;
181        let elapsed = start.elapsed();
182
183        assert!(
184            elapsed >= retry_strategy.delay_ms
185                && elapsed
186                    <= retry_strategy.delay_ms
187                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
188            "Elapsed time of {} ms was not within expected range",
189            elapsed.as_millis()
190        );
191
192        let start = Instant::now();
193        retry_strategy.delay(3, &capabilities).await;
194        let elapsed = start.elapsed();
195
196        // For the Linear strategy, the delay for the 3rd attempt should be delay_ms + (delay_ms *
197        // 2).
198        assert!(
199            elapsed >= retry_strategy.delay_ms + (retry_strategy.delay_ms * 2)
200                && elapsed
201                    <= retry_strategy.delay_ms
202                        + (retry_strategy.delay_ms * 2)
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
209    #[cfg_attr(miri, ignore)]
210    #[tokio::test]
211    async fn test_retry_strategy_exponential() {
212        let retry_strategy = RetryStrategy {
213            max_retries: 5,
214            delay_ms: Duration::from_millis(100),
215            backoff_type: RetryBackoffType::Exponential,
216            jitter: None,
217        };
218        let capabilities = NativeSleepCapability;
219
220        let start = Instant::now();
221        retry_strategy.delay(1, &capabilities).await;
222        let elapsed = start.elapsed();
223
224        assert!(
225            elapsed >= retry_strategy.delay_ms
226                && elapsed
227                    <= retry_strategy.delay_ms
228                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
229            "Elapsed time of {} ms was not within expected range",
230            elapsed.as_millis()
231        );
232
233        let start = Instant::now();
234        retry_strategy.delay(3, &capabilities).await;
235        let elapsed = start.elapsed();
236        // For the Exponential strategy, the delay for the 3rd attempt should be delay_ms * 2^(3-1)
237        // = delay_ms * 4.
238        assert!(
239            elapsed >= retry_strategy.delay_ms * 4
240                && elapsed
241                    <= retry_strategy.delay_ms * 4
242                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
243            "Elapsed time of {} ms was not within expected range",
244            elapsed.as_millis()
245        );
246    }
247
248    #[cfg_attr(miri, ignore)]
249    #[tokio::test]
250    async fn test_retry_strategy_jitter() {
251        let retry_strategy = RetryStrategy {
252            max_retries: 5,
253            delay_ms: Duration::from_millis(100),
254            backoff_type: RetryBackoffType::Constant,
255            jitter: Some(Duration::from_millis(50)),
256        };
257        let capabilities = NativeSleepCapability;
258
259        let start = Instant::now();
260        retry_strategy.delay(1, &capabilities).await;
261        let elapsed = start.elapsed();
262
263        // The delay should be between delay_ms and delay_ms + jitter
264        assert!(
265            elapsed >= retry_strategy.delay_ms
266                && elapsed
267                    <= retry_strategy.delay_ms
268                        + retry_strategy.jitter.unwrap()
269                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
270            "Elapsed time of {} ms was not within expected range",
271            elapsed.as_millis()
272        );
273    }
274
275    #[cfg_attr(miri, ignore)]
276    #[tokio::test]
277    async fn test_retry_strategy_max_retries() {
278        let retry_strategy = RetryStrategy {
279            max_retries: 17,
280            delay_ms: Duration::from_millis(100),
281            backoff_type: RetryBackoffType::Constant,
282            jitter: Some(Duration::from_millis(50)),
283        };
284
285        assert_eq!(
286            retry_strategy.max_retries(),
287            17,
288            "Max retries did not match expected value"
289        );
290    }
291}