eventsource_client/
retry.rs

1use std::time::{Duration, Instant};
2
3use rand::{thread_rng, Rng};
4
5pub(crate) trait RetryStrategy {
6    /// Return the next amount of time a failed request should delay before re-attempting.
7    fn next_delay(&mut self, current_time: Instant) -> Duration;
8
9    /// Modify the strategy's default base delay.
10    fn change_base_delay(&mut self, base_delay: Duration);
11
12    /// Used to indicate to the strategy that it can reset as a successful connection has been made.
13    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        // While the external application has indicated success, we don't actually want to reset the
75        // retry policy just yet. Instead, we want to record the time it was successful. Then when
76        // we calculate the next delay, we can reset the strategy ONLY when it has been at least
77        // DEFAULT_RESET_RETRY_INTERVAL seconds.
78        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        // Prepare a retry strategy that has succeeded at a specific point.
173        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        // Verify that calculating the next delay returns as expected
179        let time = start.add(Duration::from_secs(1));
180        let delay = retry.next_delay(time);
181        assert_eq!(delay, base);
182
183        // Verify resetting the last known good time doesn't change the retry policy since it hasn't
184        // exceeded the retry interval.
185        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        // And finally check that if we exceed the reset interval, the retry strategy will default
193        // back to base.
194        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}