mysten_common/
backoff.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{iter::Iterator, time::Duration};
5
6use rand::Rng as _;
7
8/// Creates a generator which yields an approximately exponential series of durations, as back-off delays.
9/// Jitters are added to each delay by default to prevent thundering herd on retries.
10///
11/// The API is inspired by tokio-retry::strategy::ExponentialBackoff for ease of use.
12/// But bugs in the original implementation have been fixed.
13///
14/// ```rust,no_run
15/// use std::time::Duration;
16/// use mysten_common::backoff::ExponentialBackoff;
17///
18/// // Basic example:
19/// let mut backoff = ExponentialBackoff::new(Duration::from_millis(100), Duration::from_secs(10));
20/// for (attempt, delay) in backoff.enumerate() {
21///     println!("Attempt {attempt}: Delay: {:?}", delay);
22/// }
23///
24/// // Specifying backoff factor and maximum jitter:
25/// let mut backoff = ExponentialBackoff::new(Duration::from_secs(5), Duration::from_secs(60))
26///     .factor(2.0)
27///     .max_jitter(Duration::from_secs(1));
28/// loop {
29///     // next() should always return a Some(Duration).
30///     let delay = backoff.next().unwrap();
31///     println!("Delay: {:?}", delay);
32/// }
33/// ```
34#[derive(Debug, Clone)]
35pub struct ExponentialBackoff {
36    initial_delay: Duration,
37    next: Duration,
38    factor: f64,
39    max_delay: Duration,
40    max_jitter: Duration,
41}
42
43impl ExponentialBackoff {
44    /// Constructs a new exponential backoff generator, specifying the maximum delay.
45    pub fn new(initial_delay: Duration, max_delay: Duration) -> ExponentialBackoff {
46        ExponentialBackoff {
47            initial_delay,
48            next: initial_delay,
49            factor: 1.2,
50            max_delay,
51            max_jitter: initial_delay,
52        }
53    }
54
55    /// Resets the backoff to the initial delay.
56    pub fn reset(&mut self) {
57        self.next = self.initial_delay;
58    }
59
60    /// Sets the approximate ratio of consecutive backoff delays, before jitters are applied.
61    /// Setting this to Duration::ZERO disables jittering.
62    ///
63    /// Default is 1.2.
64    pub fn factor(mut self, factor: f64) -> ExponentialBackoff {
65        self.factor = factor;
66        self
67    }
68
69    /// Sets the maximum jitter per delay.
70    ///
71    /// Default is the initial delay.
72    pub fn max_jitter(mut self, max_jitter: Duration) -> ExponentialBackoff {
73        self.max_jitter = max_jitter;
74        self
75    }
76}
77
78impl Iterator for ExponentialBackoff {
79    type Item = Duration;
80
81    /// Yields backoff delays. Never terminates.
82    fn next(&mut self) -> Option<Duration> {
83        let current = self.next;
84
85        let jitter = if self.max_jitter.is_zero() {
86            Duration::ZERO
87        } else {
88            Duration::from_secs_f64(
89                rand::thread_rng().gen_range(0.0..self.max_jitter.as_secs_f64()),
90            )
91        };
92        self.next = current
93            .mul_f64(self.factor)
94            .min(self.max_delay)
95            .saturating_add(jitter);
96
97        Some(current)
98    }
99}
100
101#[test]
102fn test_exponential_backoff_default() {
103    let mut backoff = ExponentialBackoff::new(Duration::from_millis(50), Duration::from_secs(10));
104
105    let bounds = vec![
106        (Duration::from_millis(50), Duration::from_millis(100)),
107        (Duration::from_millis(60), Duration::from_millis(170)),
108    ];
109    for (lower, upper) in bounds {
110        let delay = backoff.next().unwrap();
111        assert!(delay >= lower && delay <= upper);
112    }
113}
114
115#[test]
116fn test_exponential_backoff_base_100_factor_2_no_jitter() {
117    let mut backoff = ExponentialBackoff::new(Duration::from_millis(100), Duration::from_secs(10))
118        .factor(2.0)
119        .max_jitter(Duration::ZERO);
120
121    assert_eq!(backoff.next(), Some(Duration::from_millis(100)));
122    assert_eq!(backoff.next(), Some(Duration::from_millis(200)));
123    assert_eq!(backoff.next(), Some(Duration::from_millis(400)));
124    assert_eq!(backoff.next(), Some(Duration::from_millis(800)));
125}
126
127#[test]
128fn test_exponential_backoff_max_delay() {
129    let mut backoff = ExponentialBackoff::new(Duration::from_millis(200), Duration::from_secs(1))
130        .factor(3.0)
131        .max_jitter(Duration::ZERO);
132
133    assert_eq!(backoff.next(), Some(Duration::from_millis(200)));
134    assert_eq!(backoff.next(), Some(Duration::from_millis(600)));
135
136    for _ in 0..10 {
137        assert_eq!(backoff.next(), Some(Duration::from_secs(1)));
138    }
139}
140
141#[test]
142fn test_exponential_backoff_reset() {
143    let mut backoff = ExponentialBackoff::new(Duration::from_millis(100), Duration::from_secs(10))
144        .factor(2.0)
145        .max_jitter(Duration::ZERO);
146
147    assert_eq!(backoff.next(), Some(Duration::from_millis(100)));
148    assert_eq!(backoff.next(), Some(Duration::from_millis(200)));
149    assert_eq!(backoff.next(), Some(Duration::from_millis(400)));
150
151    backoff.reset();
152
153    assert_eq!(backoff.next(), Some(Duration::from_millis(100)));
154    assert_eq!(backoff.next(), Some(Duration::from_millis(200)));
155}
mysten_common/backoff.rs

mysten_common/
backoff.rs
