libdd-trace-utils 12.0.0

// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

//! Types used when calling [`super::send_with_retry`] to configure the retry logic.

use std::time::Duration;

use libdd_capabilities::sleep::SleepCapability;

/// Enum representing the type of backoff to use for the delay between retries.
#[derive(Debug, Clone)]
#[cfg_attr(test, derive(PartialEq))]
pub enum RetryBackoffType {
    /// Increases the delay by a fixed increment each attempt.
    Linear,
    /// The delay is constant for each attempt.
    Constant,
    /// The delay is doubled for each attempt.
    Exponential,
}

/// Struct representing the retry strategy for sending data.
///
/// This struct contains the parameters that define how retries should be handled when sending data.
/// It includes the maximum number of retries, the delay between retries, the type of backoff to
/// use, and an optional jitter to add randomness to the delay.
#[derive(Debug, Clone)]
#[cfg_attr(test, derive(PartialEq))]
pub struct RetryStrategy {
    /// The maximum number of retries to attempt.
    max_retries: u32,
    // The minimum delay between retries.
    delay_ms: Duration,
    /// The type of backoff to use for the delay between retries.
    backoff_type: RetryBackoffType,
    /// An optional jitter to add randomness to the delay.
    jitter: Option<Duration>,
}

impl Default for RetryStrategy {
    fn default() -> Self {
        RetryStrategy {
            max_retries: 5,
            delay_ms: Duration::from_millis(100),
            backoff_type: RetryBackoffType::Exponential,
            jitter: None,
        }
    }
}

impl RetryStrategy {
    /// Creates a new `RetryStrategy` with the specified parameters.
    ///
    /// # Arguments
    ///
    /// * `max_retries`: The maximum number of retries to attempt.
    /// * `delay_ms`: The minimum delay between retries, in milliseconds.
    /// * `backoff_type`: The type of backoff to use for the delay between retries.
    /// * `jitter`: An optional jitter to add randomness to the delay, in milliseconds.
    ///
    /// # Returns
    ///
    /// A `RetryStrategy` instance with the specified parameters.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use libdd_trace_utils::send_with_retry::{RetryBackoffType, RetryStrategy};
    /// use std::time::Duration;
    ///
    /// let retry_strategy = RetryStrategy::new(5, 100, RetryBackoffType::Exponential, Some(50));
    /// ```
    pub fn new(
        max_retries: u32,
        delay_ms: u64,
        backoff_type: RetryBackoffType,
        jitter: Option<u64>,
    ) -> RetryStrategy {
        RetryStrategy {
            max_retries,
            delay_ms: Duration::from_millis(delay_ms),
            backoff_type,
            jitter: jitter
                .filter(|jitter_ms| *jitter_ms != 0)
                .map(Duration::from_millis),
        }
    }
    /// Delays the next request attempt based on the retry strategy.
    ///
    /// If a jitter duration is specified in the retry strategy, a random duration up to the jitter
    /// value is added to the delay.
    ///
    /// # Arguments
    ///
    /// * `attempt`: The number of the current attempt (1-indexed).
    /// * `capabilities`: Provides the sleep capability for the delay.
    pub(crate) async fn delay<C: SleepCapability>(&self, attempt: u32, capabilities: &C) {
        capabilities.sleep(self.delay_for_attempt(attempt)).await;
    }

    fn delay_for_attempt(&self, attempt: u32) -> Duration {
        let retry_index = attempt.saturating_sub(1);
        let multiplier = match self.backoff_type {
            RetryBackoffType::Exponential => 2u32.checked_pow(retry_index).unwrap_or(u32::MAX),
            RetryBackoffType::Constant => 1,
            RetryBackoffType::Linear => retry_index.saturating_add(1),
        };
        let delay = self
            .delay_ms
            .checked_mul(multiplier)
            .unwrap_or(Duration::MAX);

        if let Some(jitter) = self.jitter {
            let jitter_ms = u64::try_from(jitter.as_millis()).unwrap_or(u64::MAX);
            let randomized_ms = rand::random::<u64>() % jitter_ms;
            delay.saturating_add(Duration::from_millis(randomized_ms))
        } else {
            delay
        }
    }

    /// Returns the maximum number of retries.
    pub(crate) fn max_retries(&self) -> u32 {
        self.max_retries
    }
}

#[cfg(test)]
// For tests RetryStrategy tests the observed delay should be approximate.
mod tests {
    use super::*;
    use libdd_capabilities_impl::NativeSleepCapability;
    use tokio::time::Instant;

    // This tolerance is on the higher side to account for github's runners not having consistent
    // performance. It shouldn't impact the quality of the tests since the most important aspect
    // of the retry logic is we wait a minimum amount of time.
    const RETRY_STRATEGY_TIME_TOLERANCE_MS: u64 = 100;

    #[cfg_attr(miri, ignore)]
    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
    // with "multi-thread".
    #[tokio::test(start_paused = true)]
    async fn test_retry_strategy_constant() {
        let retry_strategy = RetryStrategy {
            max_retries: 5,
            delay_ms: Duration::from_millis(100),
            backoff_type: RetryBackoffType::Constant,
            jitter: None,
        };
        let capabilities = NativeSleepCapability;

        let start = Instant::now();
        retry_strategy.delay(1, &capabilities).await;
        let elapsed = start.elapsed();

        assert!(
            elapsed >= retry_strategy.delay_ms
                && elapsed
                    <= retry_strategy.delay_ms
                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
            "Elapsed time of {} ms was not within expected range",
            elapsed.as_millis()
        );

        let start = Instant::now();
        retry_strategy.delay(2, &capabilities).await;
        let elapsed = start.elapsed();

        assert!(
            elapsed >= retry_strategy.delay_ms
                && elapsed
                    <= retry_strategy.delay_ms
                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
            "Elapsed time of {} ms was not within expected range",
            elapsed.as_millis()
        );
    }

    #[cfg_attr(miri, ignore)]
    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
    // with "multi-thread".
    #[tokio::test(start_paused = true)]
    async fn test_retry_strategy_linear() {
        let retry_strategy = RetryStrategy {
            max_retries: 5,
            delay_ms: Duration::from_millis(100),
            backoff_type: RetryBackoffType::Linear,
            jitter: None,
        };
        let capabilities = NativeSleepCapability;

        let start = Instant::now();
        retry_strategy.delay(1, &capabilities).await;
        let elapsed = start.elapsed();

        assert!(
            elapsed >= retry_strategy.delay_ms
                && elapsed
                    <= retry_strategy.delay_ms
                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
            "Elapsed time of {} ms was not within expected range",
            elapsed.as_millis()
        );

        let start = Instant::now();
        retry_strategy.delay(3, &capabilities).await;
        let elapsed = start.elapsed();

        // For the Linear strategy, the delay for the 3rd attempt should be delay_ms + (delay_ms *
        // 2).
        assert!(
            elapsed >= retry_strategy.delay_ms + (retry_strategy.delay_ms * 2)
                && elapsed
                    <= retry_strategy.delay_ms
                        + (retry_strategy.delay_ms * 2)
                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
            "Elapsed time of {} ms was not within expected range",
            elapsed.as_millis()
        );
    }

    #[cfg_attr(miri, ignore)]
    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
    // with "multi-thread".
    #[tokio::test(start_paused = true)]
    async fn test_retry_strategy_exponential() {
        let retry_strategy = RetryStrategy {
            max_retries: 5,
            delay_ms: Duration::from_millis(100),
            backoff_type: RetryBackoffType::Exponential,
            jitter: None,
        };
        let capabilities = NativeSleepCapability;

        let start = Instant::now();
        retry_strategy.delay(1, &capabilities).await;
        let elapsed = start.elapsed();

        assert!(
            elapsed >= retry_strategy.delay_ms
                && elapsed
                    <= retry_strategy.delay_ms
                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
            "Elapsed time of {} ms was not within expected range",
            elapsed.as_millis()
        );

        let start = Instant::now();
        retry_strategy.delay(3, &capabilities).await;
        let elapsed = start.elapsed();
        // For the Exponential strategy, the delay for the 3rd attempt should be delay_ms * 2^(3-1)
        // = delay_ms * 4.
        assert!(
            elapsed >= retry_strategy.delay_ms * 4
                && elapsed
                    <= retry_strategy.delay_ms * 4
                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
            "Elapsed time of {} ms was not within expected range",
            elapsed.as_millis()
        );
    }

    #[cfg_attr(miri, ignore)]
    // Enabling "paused" tokio mode so the tokio's inner clock is auto-advanced to the nearest
    // pending deadline. Warning: this only works in current-thread runtime, it can't be combined
    // with "multi-thread".
    #[tokio::test(start_paused = true)]
    async fn test_retry_strategy_jitter() {
        let retry_strategy = RetryStrategy {
            max_retries: 5,
            delay_ms: Duration::from_millis(100),
            backoff_type: RetryBackoffType::Constant,
            jitter: Some(Duration::from_millis(50)),
        };
        let capabilities = NativeSleepCapability;

        let start = Instant::now();
        retry_strategy.delay(1, &capabilities).await;
        let elapsed = start.elapsed();

        // The delay should be between delay_ms and delay_ms + jitter
        assert!(
            elapsed >= retry_strategy.delay_ms
                && elapsed
                    <= retry_strategy.delay_ms
                        + retry_strategy.jitter.unwrap()
                        + Duration::from_millis(RETRY_STRATEGY_TIME_TOLERANCE_MS),
            "Elapsed time of {} ms was not within expected range",
            elapsed.as_millis()
        );
    }

    #[cfg_attr(miri, ignore)]
    #[tokio::test]
    async fn test_retry_strategy_max_retries() {
        let retry_strategy = RetryStrategy {
            max_retries: 17,
            delay_ms: Duration::from_millis(100),
            backoff_type: RetryBackoffType::Constant,
            jitter: Some(Duration::from_millis(50)),
        };

        assert_eq!(
            retry_strategy.max_retries(),
            17,
            "Max retries did not match expected value"
        );
    }

    #[test]
    fn test_retry_delays() {
        let retry_strategy = RetryStrategy::new(3, 100, RetryBackoffType::Exponential, None);

        assert_eq!(
            retry_strategy.delay_for_attempt(1),
            Duration::from_millis(100)
        );
        assert_eq!(
            retry_strategy.delay_for_attempt(2),
            Duration::from_millis(200)
        );
        assert_eq!(
            retry_strategy.delay_for_attempt(3),
            Duration::from_millis(400)
        );
    }

    #[test]
    fn test_zero_jitter_is_supported() {
        let retry_strategy = RetryStrategy::new(1, 100, RetryBackoffType::Constant, Some(0));

        assert_eq!(retry_strategy.jitter, None);
        assert_eq!(
            retry_strategy.delay_for_attempt(1),
            Duration::from_millis(100)
        );
    }
}