Skip to main content

pgtask_core/
retry.rs

1use std::time::Duration;
2
3use rand::Rng as _;
4use serde::{Deserialize, Serialize};
5
6#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
7pub enum RetryPolicy {
8    Never,
9    Fixed {
10        delay: Duration,
11    },
12    Exponential {
13        base_delay: Duration,
14        factor: u32,
15        max_delay: Duration,
16    },
17}
18
19impl RetryPolicy {
20    pub fn delay_for(self, attempt: u16) -> Option<Duration> {
21        match self {
22            Self::Never => None,
23            Self::Fixed { delay } => Some(delay),
24            Self::Exponential {
25                base_delay,
26                factor,
27                max_delay,
28            } => {
29                let multiplier = factor.saturating_pow(u32::from(attempt.saturating_sub(1)));
30                let ceiling = base_delay.saturating_mul(multiplier).min(max_delay);
31                let ceiling_nanos = u64::try_from(ceiling.as_nanos()).unwrap_or(u64::MAX);
32                Some(Duration::from_nanos(rand::rng().random_range(0..=ceiling_nanos)))
33            }
34        }
35    }
36}
37
38impl Default for RetryPolicy {
39    fn default() -> Self {
40        Self::Exponential {
41            base_delay: Duration::from_secs(1),
42            factor: 2,
43            max_delay: Duration::from_mins(1),
44        }
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use std::time::Duration;
51
52    use super::RetryPolicy;
53
54    #[test]
55    fn policies_calculate_bounded_delays() {
56        assert_eq!(RetryPolicy::Never.delay_for(1), None);
57        assert_eq!(
58            RetryPolicy::Fixed {
59                delay: Duration::from_secs(3)
60            }
61            .delay_for(10),
62            Some(Duration::from_secs(3))
63        );
64        let exponential = RetryPolicy::Exponential {
65            base_delay: Duration::from_secs(2),
66            factor: 3,
67            max_delay: Duration::from_secs(20),
68        };
69        for _ in 0..100 {
70            assert!(exponential.delay_for(1).unwrap() <= Duration::from_secs(2));
71            assert!(exponential.delay_for(2).unwrap() <= Duration::from_secs(6));
72            assert!(exponential.delay_for(10).unwrap() <= Duration::from_secs(20));
73        }
74    }
75}