plane_common/
exponential_backoff.rs

1use std::time::{Duration, SystemTime};
2
3pub struct ExponentialBackoff {
4    initial_duration_millis: u128,
5    max_duration: Duration,
6    defer_duration: Duration,
7    multiplier: f64,
8    step: i32,
9    deferred_reset: Option<SystemTime>,
10}
11
12impl ExponentialBackoff {
13    pub fn new(
14        initial_duration: Duration,
15        max_duration: Duration,
16        multiplier: f64,
17        defer_duration: Duration,
18    ) -> Self {
19        let initial_duration_millis = initial_duration.as_millis();
20
21        Self {
22            initial_duration_millis,
23            max_duration,
24            multiplier,
25            step: 0,
26            defer_duration,
27            deferred_reset: None,
28        }
29    }
30
31    /// Reset the backoff, but only if `wait` is not called again for at least `defer_duration`.
32    pub fn defer_reset(&mut self) {
33        self.deferred_reset = Some(SystemTime::now() + self.defer_duration);
34    }
35
36    pub async fn wait(&mut self) {
37        if let Some(deferred_reset) = self.deferred_reset {
38            self.deferred_reset = None;
39            if SystemTime::now() > deferred_reset {
40                self.reset();
41                return;
42            }
43        }
44
45        let duration = self.initial_duration_millis as f64 * self.multiplier.powi(self.step);
46        let duration = Duration::from_millis(duration as u64);
47        let duration = duration.min(self.max_duration);
48        tokio::time::sleep(duration).await;
49
50        self.step += 1;
51    }
52
53    pub fn reset(&mut self) {
54        self.deferred_reset = None;
55        self.step = 0;
56    }
57}
58
59impl Default for ExponentialBackoff {
60    fn default() -> Self {
61        Self::new(
62            Duration::from_secs(1),
63            Duration::from_secs(60),
64            1.1,
65            Duration::from_secs(60),
66        )
67    }
68}