git_lock/
backoff.rs

1use std::time::Duration;
2
3fn randomize(backoff_ms: usize) -> usize {
4    let new_value = (fastrand::usize(750..=1250) * backoff_ms) / 1000;
5    if new_value == 0 {
6        backoff_ms
7    } else {
8        new_value
9    }
10}
11
12/// A utility to calculate steps for exponential backoff similar to how it's done in `git`.
13pub struct Exponential<Fn> {
14    multiplier: usize,
15    max_multiplier: usize,
16    exponent: usize,
17    transform: Fn,
18}
19
20impl Default for Exponential<fn(usize) -> usize> {
21    fn default() -> Self {
22        Exponential {
23            multiplier: 1,
24            max_multiplier: 1000,
25            exponent: 1,
26            transform: std::convert::identity,
27        }
28    }
29}
30
31impl Exponential<fn(usize) -> usize> {
32    /// Create a new exponential backoff iterator that backs off in randomized, ever increasing steps.
33    pub fn default_with_random() -> Self {
34        Exponential {
35            multiplier: 1,
36            max_multiplier: 1000,
37            exponent: 1,
38            transform: randomize,
39        }
40    }
41}
42
43impl<Transform> Exponential<Transform>
44where
45    Transform: Fn(usize) -> usize,
46{
47    /// Return an iterator that yields `Duration` instances to sleep on until `time` is depleted.
48    pub fn until_no_remaining(&mut self, time: Duration) -> impl Iterator<Item = Duration> + '_ {
49        let mut elapsed = Duration::default();
50        let mut stop_next_iteration = false;
51        self.take_while(move |d| {
52            if stop_next_iteration {
53                false
54            } else {
55                elapsed += *d;
56                if elapsed > time {
57                    stop_next_iteration = true;
58                }
59                true
60            }
61        })
62    }
63}
64
65impl<Transform> Iterator for Exponential<Transform>
66where
67    Transform: Fn(usize) -> usize,
68{
69    type Item = Duration;
70
71    fn next(&mut self) -> Option<Self::Item> {
72        let wait = Duration::from_millis((self.transform)(self.multiplier) as u64);
73
74        self.multiplier += 2 * self.exponent + 1;
75        if self.multiplier > self.max_multiplier {
76            self.multiplier = self.max_multiplier;
77        } else {
78            self.exponent += 1;
79        }
80        Some(wait)
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use std::convert::TryInto;
87
88    use super::*;
89
90    const EXPECTED_TILL_SECOND: &[usize] = &[
91        1usize, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484, 529,
92        576, 625, 676, 729, 784, 841, 900, 961, 1000, 1000,
93    ];
94
95    #[test]
96    fn random_exponential_produces_values_in_the_correct_range() {
97        let mut num_identities = 0;
98        for (actual, expected) in Exponential::default_with_random().zip(EXPECTED_TILL_SECOND) {
99            let actual: usize = actual.as_millis().try_into().unwrap();
100            if actual == *expected {
101                num_identities += 1;
102            }
103            assert!(
104                actual * 1000 >= (expected - 1) * 750,
105                "value too small: {actual} < {expected}"
106            );
107            assert!(
108                actual * 1000 <= (expected + 1) * 1250,
109                "value too big: {actual} > {expected}"
110            );
111        }
112        assert!(
113            num_identities < EXPECTED_TILL_SECOND.len(),
114            "too many untransformed values: {num_identities}"
115        );
116    }
117
118    #[test]
119    fn how_many_iterations_for_a_second_of_waittime() {
120        let max = Duration::from_millis(1000);
121        assert_eq!(Exponential::default().until_no_remaining(max).count(), 14);
122        assert_eq!(
123            Exponential::default()
124                .until_no_remaining(max)
125                .reduce(|acc, n| acc + n)
126                .unwrap(),
127            Duration::from_millis(1015),
128            "a little overshoot"
129        );
130    }
131
132    #[test]
133    fn output_with_default_settings() {
134        assert_eq!(
135            Exponential::default().take(33).collect::<Vec<_>>(),
136            EXPECTED_TILL_SECOND
137                .iter()
138                .map(|n| Duration::from_millis(*n as u64))
139                .collect::<Vec<_>>()
140        );
141    }
142}