rama_utils/backoff/
exponential.rs

1use core::fmt::{self, Display};
2use core::time::Duration;
3
4use super::Backoff;
5use crate::rng::{HasherRng, Rng};
6
7use parking_lot::Mutex;
8
9/// A jittered [exponential backoff] strategy.
10///
11/// The backoff duration will increase exponentially for every subsequent
12/// backoff, up to a maximum duration. A small amount of [random jitter] is
13/// added to each backoff duration, in order to avoid retry spikes.
14///
15/// [exponential backoff]: https://en.wikipedia.org/wiki/Exponential_backoff
16/// [random jitter]: https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
17pub struct ExponentialBackoff<F, R = HasherRng> {
18    min: Duration,
19    max: Duration,
20    jitter: f64,
21    rng_creator: F,
22    state: Mutex<ExponentialBackoffState<R>>,
23}
24
25impl<F: fmt::Debug, R: fmt::Debug> fmt::Debug for ExponentialBackoff<F, R> {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("ExponentialBackoff")
28            .field("min", &self.min)
29            .field("max", &self.max)
30            .field("jitter", &self.jitter)
31            .field("rng_creator", &self.rng_creator)
32            .field("state", &self.state)
33            .finish()
34    }
35}
36
37impl<F, R> Clone for ExponentialBackoff<F, R>
38where
39    R: Rng + Clone,
40    F: Fn() -> R + Clone,
41{
42    fn clone(&self) -> Self {
43        Self {
44            min: self.min,
45            max: self.max,
46            jitter: self.jitter,
47            rng_creator: self.rng_creator.clone(),
48            state: Mutex::new(ExponentialBackoffState {
49                rng: (self.rng_creator)(),
50                iterations: 0,
51            }),
52        }
53    }
54}
55
56impl Clone for ExponentialBackoff<(), HasherRng> {
57    fn clone(&self) -> Self {
58        Self {
59            min: self.min,
60            max: self.max,
61            jitter: self.jitter,
62            rng_creator: (),
63            state: Mutex::new(ExponentialBackoffState {
64                rng: HasherRng::default(),
65                iterations: 0,
66            }),
67        }
68    }
69}
70
71struct ExponentialBackoffState<R = HasherRng> {
72    rng: R,
73    iterations: u32,
74}
75
76impl<R: fmt::Debug> fmt::Debug for ExponentialBackoffState<R> {
77    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78        f.debug_struct("ExponentialBackoffState")
79            .field("rng", &self.rng)
80            .field("iterations", &self.iterations)
81            .finish()
82    }
83}
84
85impl<F, R> ExponentialBackoff<F, R>
86where
87    R: Rng + Clone,
88    F: Fn() -> R + Clone,
89{
90    /// Create a new `ExponentialBackoff`.
91    ///
92    /// # Error
93    ///
94    /// Returns a config validation error if:
95    /// - `min` > `max`
96    /// - `max` == 0
97    /// - `jitter` < `0.0`
98    /// - `jitter` > `100.0`
99    /// - `jitter` is NaN
100    pub fn new(
101        min: Duration,
102        max: Duration,
103        jitter: f64,
104        rng_creator: F,
105    ) -> Result<Self, InvalidBackoff> {
106        let rng = rng_creator();
107        Self::new_inner(min, max, jitter, rng_creator, rng)
108    }
109}
110
111impl<F, R> ExponentialBackoff<F, R> {
112    fn new_inner(
113        min: Duration,
114        max: Duration,
115        jitter: f64,
116        rng_creator: F,
117        rng: R,
118    ) -> Result<Self, InvalidBackoff> {
119        if min > max {
120            return Err(InvalidBackoff("maximum must not be less than minimum"));
121        }
122        if max == Duration::from_millis(0) {
123            return Err(InvalidBackoff("maximum must be non-zero"));
124        }
125        if jitter < 0.0 {
126            return Err(InvalidBackoff("jitter must not be negative"));
127        }
128        if jitter > 100.0 {
129            return Err(InvalidBackoff("jitter must not be greater than 100"));
130        }
131        if jitter.is_nan() {
132            return Err(InvalidBackoff("jitter must not be NaN"));
133        }
134        Ok(Self {
135            min,
136            max,
137            jitter,
138            rng_creator,
139            state: Mutex::new(ExponentialBackoffState { rng, iterations: 0 }),
140        })
141    }
142}
143
144impl<F, R: Rng> ExponentialBackoff<F, R> {
145    fn base(&self) -> Duration {
146        debug_assert!(
147            self.min <= self.max,
148            "maximum backoff must not be less than minimum backoff"
149        );
150        debug_assert!(
151            self.max > Duration::from_millis(0),
152            "Maximum backoff must be non-zero"
153        );
154        self.min
155            .checked_mul(2_u32.saturating_pow(self.state.lock().iterations))
156            .unwrap_or(self.max)
157            .min(self.max)
158    }
159
160    /// Returns a random, uniform duration on `[0, base*self.jitter]` no greater
161    /// than `self.max`.
162    fn jitter(&self, base: Duration) -> Option<Duration> {
163        if self.jitter <= 0.0 {
164            None
165        } else {
166            let jitter_factor = self.state.lock().rng.next_f64();
167            debug_assert!(
168                jitter_factor > 0.0,
169                "rng returns values between 0.0 and 1.0"
170            );
171            let rand_jitter = jitter_factor * self.jitter;
172            let secs = (base.as_secs() as f64) * rand_jitter;
173            let nanos = (base.subsec_nanos() as f64) * rand_jitter;
174            #[expect(
175                clippy::unchecked_time_subtraction,
176                reason = "base is produced by self.base() which caps at self.max via .min(self.max), so this subtraction cannot underflow"
177            )]
178            let remaining = self.max - base;
179            let result = Duration::new(secs as u64, nanos as u32);
180            if remaining.is_zero() || result.is_zero() {
181                None
182            } else {
183                Some(result.min(remaining))
184            }
185        }
186    }
187}
188
189impl<F, R> Backoff for ExponentialBackoff<F, R>
190where
191    R: Rng,
192    F: Send + Sync + 'static,
193{
194    async fn next_backoff(&self) -> bool {
195        let base = self.base();
196        let Some(jitter) = self.jitter(base) else {
197            self.reset().await;
198            return false;
199        };
200
201        let next = base + jitter;
202
203        self.state.lock().iterations += 1;
204
205        tokio::time::sleep(next).await;
206        true
207    }
208
209    async fn reset(&self) {
210        self.state.lock().iterations = 0;
211    }
212}
213
214impl Default for ExponentialBackoff<(), HasherRng> {
215    fn default() -> Self {
216        Self {
217            min: Duration::from_millis(50),
218            max: Duration::from_secs(3),
219            jitter: 0.99,
220            rng_creator: (),
221            state: Mutex::new(ExponentialBackoffState {
222                rng: HasherRng::default(),
223                iterations: 0,
224            }),
225        }
226    }
227}
228
229/// Backoff validation error.
230#[derive(Debug)]
231pub struct InvalidBackoff(&'static str);
232
233impl Display for InvalidBackoff {
234    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
235        write!(f, "invalid backoff: {}", self.0)
236    }
237}
238
239impl core::error::Error for InvalidBackoff {}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use quickcheck::*;
245
246    #[tokio::test]
247    async fn backoff_default() {
248        let backoff = ExponentialBackoff::default();
249        assert!(backoff.next_backoff().await);
250    }
251
252    #[tokio::test]
253    async fn backoff_reset() {
254        let backoff = ExponentialBackoff::default();
255        assert!(backoff.next_backoff().await);
256        assert!(backoff.state.lock().iterations == 1);
257        backoff.reset().await;
258        assert!(backoff.state.lock().iterations == 0);
259    }
260
261    #[tokio::test]
262    async fn backoff_clone() {
263        let backoff = ExponentialBackoff::default();
264
265        assert!(backoff.state.lock().iterations == 0);
266        assert!(backoff.next_backoff().await);
267        assert!(backoff.state.lock().iterations == 1);
268
269        let cloned = backoff.clone();
270        assert!(cloned.state.lock().iterations == 0);
271        assert!(backoff.state.lock().iterations == 1);
272
273        assert!(cloned.next_backoff().await);
274        assert!(cloned.state.lock().iterations == 1);
275        assert!(backoff.state.lock().iterations == 1);
276    }
277
278    #[test]
279    fn jitter_must_be_finite() {
280        let min = Duration::from_millis(0);
281        let max = Duration::from_millis(1);
282        let rng = HasherRng::default();
283
284        for n in [f64::INFINITY, f64::NEG_INFINITY, f64::NAN] {
285            let result = ExponentialBackoff::new(min, max, n, || rng.clone());
286            assert!(
287                matches!(result, Err(InvalidBackoff(_))),
288                "{n} should be an invalid jitter",
289            );
290        }
291    }
292
293    quickcheck! {
294        fn backoff_base_first(min_ms: u64, max_ms: u64) -> TestResult {
295            let min = Duration::from_millis(min_ms);
296            let max = Duration::from_millis(max_ms);
297            let Ok(backoff) = ExponentialBackoff::new(min, max, 0.0, HasherRng::default) else {
298                return TestResult::discard();
299            };
300
301            let delay = backoff.base();
302            TestResult::from_bool(min == delay)
303        }
304
305        fn backoff_base(min_ms: u64, max_ms: u64, iterations: u32) -> TestResult {
306            let min = Duration::from_millis(min_ms);
307            let max = Duration::from_millis(max_ms);
308            let Ok(backoff) = ExponentialBackoff::new(min, max, 0.0, HasherRng::default) else {
309                return TestResult::discard();
310            };
311
312            backoff.state.lock().iterations = iterations;
313            let delay = backoff.base();
314            TestResult::from_bool(min <= delay && delay <= max)
315        }
316
317        fn backoff_jitter(base_ms: u64, max_ms: u64, jitter: f64) -> TestResult {
318            let base = Duration::from_millis(base_ms);
319            let max = Duration::from_millis(max_ms);
320            let Ok(backoff) = ExponentialBackoff::new(base, max, jitter, HasherRng::default) else {
321                return TestResult::discard();
322            };
323
324            let j = backoff.jitter(base);
325            if jitter == 0.0 || base_ms == 0 || max_ms == base_ms {
326                TestResult::from_bool(j.is_none())
327            } else {
328                TestResult::from_bool(j.is_some())
329            }
330        }
331    }
332}
rama_utils/backoff/exponential.rs

rama_utils/backoff/
exponential.rs
