klukai_types/
backoff.rs

1use std::{iter, time};
2
3use rand::{Rng, rng};
4
5/// Exponential backoff.
6#[derive(Debug, Clone)]
7pub struct Backoff {
8    retries: u32,
9    min: time::Duration,
10    max: time::Duration,
11    jitter: f32,
12    factor: u32,
13}
14
15impl Backoff {
16    /// Create a new instance.
17    #[inline]
18    pub fn new(retries: u32) -> Self {
19        Self {
20            retries,
21            min: time::Duration::from_millis(100),
22            max: time::Duration::from_secs(10),
23            jitter: 0.3,
24            factor: 2,
25        }
26    }
27
28    /// Set the min and max durations.
29    #[inline]
30    pub fn timeout_range(mut self, min: time::Duration, max: time::Duration) -> Self {
31        self.min = min;
32        self.max = max;
33        self
34    }
35
36    /// Set the amount of jitter per backoff.
37    ///
38    /// ## Panics
39    /// This method panics if jitter <= 0 or jitter >= 1
40    #[inline]
41    pub fn jitter(mut self, jitter: f32) -> Self {
42        assert!(
43            jitter > 0f32 && jitter < 1f32,
44            "<exponential-backoff>: jitter must be greater than 0 and less than 1."
45        );
46        self.jitter = jitter;
47        self
48    }
49
50    /// Set the growth factor for each iteration of the backoff.
51    #[inline]
52    pub fn factor(mut self, factor: u32) -> Self {
53        self.factor = factor;
54        self
55    }
56
57    /// Create an iterator.
58    #[inline]
59    pub fn iter(self) -> Iter {
60        Iter::new(self)
61    }
62}
63
64/// Immutable iterator.
65#[derive(Debug, Clone)]
66pub struct Iter {
67    inner: Backoff,
68    retry_count: u32,
69    last_exponent: u32,
70}
71
72impl Iter {
73    #[inline]
74    pub fn new(inner: Backoff) -> Self {
75        Self {
76            inner,
77            retry_count: 0,
78            last_exponent: 0,
79        }
80    }
81
82    pub fn retry_count(&self) -> u32 {
83        self.retry_count
84    }
85}
86
87impl iter::Iterator for Iter {
88    type Item = time::Duration;
89
90    #[inline]
91    fn next(&mut self) -> Option<Self::Item> {
92        // if  {
93        if self.inner.retries != 0 && self.retry_count >= self.inner.retries {
94            return None;
95        }
96
97        // Create exponential duration.
98        // Prevent multiply overflow by calling pow every time.
99        if self.inner.min * self.last_exponent < self.inner.max {
100            self.last_exponent = self.inner.factor.pow(self.retry_count);
101        }
102
103        let mut duration = self.inner.min * self.last_exponent;
104
105        self.retry_count += 1;
106
107        // Apply jitter. Uses multiples of 100 to prevent relying on floats.
108        let jitter_factor = (self.inner.jitter * 100f32) as u32;
109        let random: u32 = rng().random_range(0..(jitter_factor * 2));
110        duration *= 100;
111        if random < jitter_factor {
112            let jitter = (duration * random) / 100;
113            duration -= jitter;
114        } else {
115            let jitter = (duration * (random / 2)) / 100;
116            duration += jitter;
117        };
118        duration /= 100;
119
120        // Make sure it doesn't exceed upper / lower bounds.
121        duration = duration.min(self.inner.max);
122        duration = duration.max(self.inner.min);
123
124        Some(duration)
125    }
126}
127
128impl iter::FusedIterator for Iter {}
129
130#[cfg(test)]
131mod tests {
132    use std::time::Duration;
133
134    use super::*;
135
136    #[test]
137    fn test_backoff() {
138        let boff = Backoff::new(15)
139            .timeout_range(Duration::from_millis(10), Duration::from_secs(1))
140            .iter();
141
142        let mut total = Duration::default();
143        for dur in boff {
144            println!("{dur:?}");
145            total += dur;
146        }
147        println!("total: {total:?}");
148    }
149}
klukai_types/backoff.rs

klukai_types/
backoff.rs
