lance_core/utils/
backoff.rs

1use rand::Rng;
2use std::time::Duration;
3
4// SPDX-License-Identifier: Apache-2.0
5// SPDX-FileCopyrightText: Copyright The Lance Authors
6
7/// Computes backoff as
8///
9/// ```text
10/// backoff = base^attempt * unit + jitter
11/// ```
12///
13/// The defaults are base=2, unit=50ms, jitter=50ms, min=0ms, max=5s. This gives
14/// a backoff of 50ms, 100ms, 200ms, 400ms, 800ms, 1.6s, 3.2s, 5s, (not including jitter).
15///
16/// You can have non-exponential backoff by setting base=1.
17pub struct Backoff {
18    base: u32,
19    unit: u32,
20    jitter: i32,
21    min: u32,
22    max: u32,
23    attempt: u32,
24}
25
26impl Default for Backoff {
27    fn default() -> Self {
28        Self {
29            base: 2,
30            unit: 50,
31            jitter: 50,
32            min: 0,
33            max: 5000,
34            attempt: 0,
35        }
36    }
37}
38
39impl Backoff {
40    pub fn with_base(self, base: u32) -> Self {
41        Self { base, ..self }
42    }
43
44    pub fn with_jitter(self, jitter: i32) -> Self {
45        Self { jitter, ..self }
46    }
47
48    pub fn with_min(self, min: u32) -> Self {
49        Self { min, ..self }
50    }
51
52    pub fn with_max(self, max: u32) -> Self {
53        Self { max, ..self }
54    }
55
56    pub fn next_backoff(&mut self) -> Duration {
57        let backoff = self
58            .base
59            .saturating_pow(self.attempt)
60            .saturating_mul(self.unit);
61        let jitter = rand::thread_rng().gen_range(-self.jitter..=self.jitter);
62        let backoff = (backoff.saturating_add_signed(jitter)).clamp(self.min, self.max);
63        self.attempt += 1;
64        Duration::from_millis(backoff as u64)
65    }
66
67    pub fn attempt(&self) -> u32 {
68        self.attempt
69    }
70
71    pub fn reset(&mut self) {
72        self.attempt = 0;
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn test_backoff() {
82        let mut backoff = Backoff::default().with_jitter(0);
83        assert_eq!(backoff.next_backoff().as_millis(), 50);
84        assert_eq!(backoff.attempt(), 1);
85        assert_eq!(backoff.next_backoff().as_millis(), 100);
86        assert_eq!(backoff.attempt(), 2);
87        assert_eq!(backoff.next_backoff().as_millis(), 200);
88        assert_eq!(backoff.attempt(), 3);
89        assert_eq!(backoff.next_backoff().as_millis(), 400);
90        assert_eq!(backoff.attempt(), 4);
91    }
92}
lance_core/utils/backoff.rs

lance_core/utils/
backoff.rs
