tokio_retry2/strategy/
exponential_backoff.rs

1use std::iter::Iterator;
2use tokio::time::Duration;
3
4/// A retry strategy driven by exponential back-off.
5///
6/// The power corresponds to the number of past attempts.
7#[derive(Debug, Clone)]
8pub struct ExponentialBackoff {
9    current: u64,
10    base: u64,
11    factor: u64,
12    max_delay: Option<Duration>,
13}
14
15impl ExponentialBackoff {
16    /// Constructs a new exponential back-off strategy,
17    /// given a base duration in milliseconds.
18    ///
19    /// The resulting duration is calculated by taking the base to the `n`-th power,
20    /// where `n` denotes the number of past attempts.
21    pub const fn from_millis(base: u64) -> Self {
22        ExponentialBackoff {
23            current: base,
24            base,
25            factor: 1u64,
26            max_delay: None,
27        }
28    }
29
30    /// A multiplicative factor that will be applied to the retry delay.
31    ///
32    /// For example, using a factor of `1000` will make each delay in units of seconds.
33    ///
34    /// Default factor is `1`.
35    pub const fn factor(mut self, factor: u64) -> ExponentialBackoff {
36        self.factor = factor;
37        self
38    }
39
40    /// Apply a maximum delay. No single retry delay will be longer than this `Duration`.
41    pub const fn max_delay(mut self, duration: Duration) -> ExponentialBackoff {
42        self.max_delay = Some(duration);
43        self
44    }
45
46    /// Apply a maximum delay. No single retry delay will be longer than this `Duration::from_millis`.
47    pub const fn max_delay_millis(mut self, duration: u64) -> ExponentialBackoff {
48        self.max_delay = Some(Duration::from_millis(duration));
49        self
50    }
51}
52
53impl Iterator for ExponentialBackoff {
54    type Item = Duration;
55
56    fn next(&mut self) -> Option<Duration> {
57        // set delay duration by applying factor
58        let duration = if let Some(duration) = self.current.checked_mul(self.factor) {
59            Duration::from_millis(duration)
60        } else {
61            Duration::from_millis(u64::MAX)
62        };
63
64        // check if we reached max delay
65        if let Some(ref max_delay) = self.max_delay {
66            if duration > *max_delay {
67                #[cfg(feature = "tracing")]
68                tracing::warn!("`max_delay` for strategy reached");
69                return Some(*max_delay);
70            }
71        }
72
73        if let Some(next) = self.current.checked_mul(self.base) {
74            self.current = next;
75        } else {
76            self.current = u64::MAX;
77        }
78
79        Some(duration)
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn returns_some_exponential_base_10() {
89        let mut s = ExponentialBackoff::from_millis(10);
90
91        assert_eq!(s.next(), Some(Duration::from_millis(10)));
92        assert_eq!(s.next(), Some(Duration::from_millis(100)));
93        assert_eq!(s.next(), Some(Duration::from_millis(1000)));
94    }
95
96    #[test]
97    fn returns_some_exponential_base_2() {
98        let mut s = ExponentialBackoff::from_millis(2);
99
100        assert_eq!(s.next(), Some(Duration::from_millis(2)));
101        assert_eq!(s.next(), Some(Duration::from_millis(4)));
102        assert_eq!(s.next(), Some(Duration::from_millis(8)));
103    }
104
105    #[test]
106    fn saturates_at_maximum_value() {
107        let mut s = ExponentialBackoff::from_millis(u64::MAX - 1);
108
109        assert_eq!(s.next(), Some(Duration::from_millis(u64::MAX - 1)));
110        assert_eq!(s.next(), Some(Duration::from_millis(u64::MAX)));
111        assert_eq!(s.next(), Some(Duration::from_millis(u64::MAX)));
112    }
113
114    #[test]
115    fn can_use_factor_to_get_seconds() {
116        let factor = 1000;
117        let mut s = ExponentialBackoff::from_millis(2).factor(factor);
118
119        assert_eq!(s.next(), Some(Duration::from_secs(2)));
120        assert_eq!(s.next(), Some(Duration::from_secs(4)));
121        assert_eq!(s.next(), Some(Duration::from_secs(8)));
122    }
123
124    #[test]
125    fn stops_increasing_at_max_delay() {
126        let mut s = ExponentialBackoff::from_millis(2).max_delay(Duration::from_millis(4));
127
128        assert_eq!(s.next(), Some(Duration::from_millis(2)));
129        assert_eq!(s.next(), Some(Duration::from_millis(4)));
130        assert_eq!(s.next(), Some(Duration::from_millis(4)));
131    }
132
133    #[test]
134    fn returns_max_when_max_less_than_base() {
135        let mut s = ExponentialBackoff::from_millis(20).max_delay(Duration::from_millis(10));
136
137        assert_eq!(s.next(), Some(Duration::from_millis(10)));
138        assert_eq!(s.next(), Some(Duration::from_millis(10)));
139    }
140}