tokio_retry2/strategy/
exponential_backoff.rs1use std::iter::Iterator;
2use tokio::time::Duration;
3
4#[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 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 pub const fn factor(mut self, factor: u64) -> ExponentialBackoff {
36 self.factor = factor;
37 self
38 }
39
40 pub const fn max_delay(mut self, duration: Duration) -> ExponentialBackoff {
42 self.max_delay = Some(duration);
43 self
44 }
45
46 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 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 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}