eventsource_client/
retry.rs
1use std::time::{Duration, Instant};
2
3use rand::{thread_rng, Rng};
4
5pub(crate) trait RetryStrategy {
6 fn next_delay(&mut self, current_time: Instant) -> Duration;
8
9 fn change_base_delay(&mut self, base_delay: Duration);
11
12 fn reset(&mut self, current_time: Instant);
14}
15
16const DEFAULT_RESET_RETRY_INTERVAL: Duration = Duration::from_secs(60);
17
18pub(crate) struct BackoffRetry {
19 base_delay: Duration,
20 max_delay: Duration,
21 backoff_factor: u32,
22 include_jitter: bool,
23
24 reset_interval: Duration,
25 next_delay: Duration,
26 good_since: Option<Instant>,
27}
28
29impl BackoffRetry {
30 pub fn new(
31 base_delay: Duration,
32 max_delay: Duration,
33 backoff_factor: u32,
34 include_jitter: bool,
35 ) -> Self {
36 Self {
37 base_delay,
38 max_delay,
39 backoff_factor,
40 include_jitter,
41 reset_interval: DEFAULT_RESET_RETRY_INTERVAL,
42 next_delay: base_delay,
43 good_since: None,
44 }
45 }
46}
47
48impl RetryStrategy for BackoffRetry {
49 fn next_delay(&mut self, current_time: Instant) -> Duration {
50 let mut current_delay = self.next_delay;
51
52 if let Some(good_since) = self.good_since {
53 if current_time - good_since >= self.reset_interval {
54 current_delay = self.base_delay;
55 }
56 }
57
58 self.good_since = None;
59 self.next_delay = std::cmp::min(self.max_delay, current_delay * self.backoff_factor);
60
61 if self.include_jitter {
62 thread_rng().gen_range(current_delay / 2..=current_delay)
63 } else {
64 current_delay
65 }
66 }
67
68 fn change_base_delay(&mut self, base_delay: Duration) {
69 self.base_delay = base_delay;
70 self.next_delay = self.base_delay;
71 }
72
73 fn reset(&mut self, current_time: Instant) {
74 self.good_since = Some(current_time);
79 }
80}
81
82#[cfg(test)]
83mod tests {
84 use std::ops::Add;
85 use std::time::{Duration, Instant};
86
87 use crate::retry::{BackoffRetry, RetryStrategy};
88
89 #[test]
90 fn test_fixed_retry() {
91 let base = Duration::from_secs(10);
92 let mut retry = BackoffRetry::new(base, Duration::from_secs(30), 1, false);
93 let start = Instant::now() - Duration::from_secs(60);
94
95 assert_eq!(retry.next_delay(start), base);
96 assert_eq!(retry.next_delay(start.add(Duration::from_secs(1))), base);
97 assert_eq!(retry.next_delay(start.add(Duration::from_secs(2))), base);
98 }
99
100 #[test]
101 fn test_able_to_reset_base_delay() {
102 let base = Duration::from_secs(10);
103 let mut retry = BackoffRetry::new(base, Duration::from_secs(30), 1, false);
104 let start = Instant::now();
105
106 assert_eq!(retry.next_delay(start), base);
107 assert_eq!(retry.next_delay(start.add(Duration::from_secs(1))), base);
108
109 let base = Duration::from_secs(3);
110 retry.change_base_delay(base);
111 assert_eq!(retry.next_delay(start.add(Duration::from_secs(2))), base);
112 }
113
114 #[test]
115 fn test_with_backoff() {
116 let base = Duration::from_secs(10);
117 let max = Duration::from_secs(60);
118 let mut retry = BackoffRetry::new(base, max, 2, false);
119 let start = Instant::now() - Duration::from_secs(60);
120
121 assert_eq!(retry.next_delay(start), base);
122 assert_eq!(
123 retry.next_delay(start.add(Duration::from_secs(1))),
124 base * 2
125 );
126 assert_eq!(
127 retry.next_delay(start.add(Duration::from_secs(2))),
128 base * 4
129 );
130 assert_eq!(retry.next_delay(start.add(Duration::from_secs(3))), max);
131 }
132
133 #[test]
134 fn test_with_jitter() {
135 let base = Duration::from_secs(10);
136 let max = Duration::from_secs(60);
137 let mut retry = BackoffRetry::new(base, max, 1, true);
138 let start = Instant::now() - Duration::from_secs(60);
139
140 let delay = retry.next_delay(start);
141 assert!(base / 2 <= delay && delay <= base);
142 }
143
144 #[test]
145 fn test_retry_holds_at_max() {
146 let base = Duration::from_secs(20);
147 let max = Duration::from_secs(30);
148
149 let mut retry = BackoffRetry::new(base, max, 2, false);
150 let start = Instant::now();
151 retry.reset(start);
152
153 let time = start.add(Duration::from_secs(20));
154 let delay = retry.next_delay(time);
155 assert_eq!(delay, base);
156
157 let time = time.add(Duration::from_secs(20));
158 let delay = retry.next_delay(time);
159 assert_eq!(delay, max);
160
161 let time = time.add(Duration::from_secs(20));
162 let delay = retry.next_delay(time);
163 assert_eq!(delay, max);
164 }
165
166 #[test]
167 fn test_reset_interval() {
168 let base = Duration::from_secs(10);
169 let max = Duration::from_secs(60);
170 let reset_interval = Duration::from_secs(45);
171
172 let mut retry = BackoffRetry::new(base, max, 2, false);
174 retry.reset_interval = reset_interval;
175 let start = Instant::now() - Duration::from_secs(60);
176 retry.reset(start);
177
178 let time = start.add(Duration::from_secs(1));
180 let delay = retry.next_delay(time);
181 assert_eq!(delay, base);
182
183 let time = time.add(delay);
186 retry.reset(time);
187
188 let time = time.add(Duration::from_secs(10));
189 let delay = retry.next_delay(time);
190 assert_eq!(delay, base * 2);
191
192 let time = time.add(delay);
195 retry.reset(time);
196
197 let time = time.add(reset_interval);
198 let delay = retry.next_delay(time);
199 assert_eq!(delay, base);
200 }
201}