Skip to main content

risingwave_common/util/
retry.rs

1// Copyright 2026 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::time::Duration;
16
17use tokio_retry::strategy::ExponentialBackoff;
18
19/// Creates a backoff starting at `initial_delay`, multiplying subsequent delays by
20/// `multiplier`, and capping every delay at `max_delay`.
21///
22/// This wraps [`ExponentialBackoff`], whose `from_millis` argument is the mathematical
23/// base rather than the initial delay.
24pub fn exponential_backoff(
25    initial_delay: Duration,
26    multiplier: u64,
27    max_delay: Duration,
28) -> impl Iterator<Item = Duration> + Clone {
29    let initial_delay = initial_delay.min(max_delay);
30    let initial_delay_ms = initial_delay.as_millis().try_into().unwrap_or(u64::MAX);
31
32    // Keep the confusing low-level API contained in this wrapper.
33    #[expect(
34        clippy::disallowed_methods,
35        reason = "this wrapper is the only permitted caller of the low-level API"
36    )]
37    let remaining_delays = ExponentialBackoff::from_millis(multiplier)
38        .factor(initial_delay_ms)
39        .max_delay(max_delay);
40
41    std::iter::once(initial_delay).chain(remaining_delays)
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    #[test]
49    fn starts_at_initial_delay_and_uses_multiplier() {
50        let delays = exponential_backoff(Duration::from_millis(101), 3, Duration::from_secs(10))
51            .take(4)
52            .collect::<Vec<_>>();
53
54        assert_eq!(
55            delays,
56            [
57                Duration::from_millis(101),
58                Duration::from_millis(303),
59                Duration::from_millis(909),
60                Duration::from_millis(2727),
61            ]
62        );
63    }
64
65    #[test]
66    fn caps_all_delays_at_max_delay() {
67        let delays = exponential_backoff(Duration::from_secs(1), 2, Duration::from_secs(10))
68            .take(6)
69            .collect::<Vec<_>>();
70
71        assert_eq!(
72            delays,
73            [
74                Duration::from_secs(1),
75                Duration::from_secs(2),
76                Duration::from_secs(4),
77                Duration::from_secs(8),
78                Duration::from_secs(10),
79                Duration::from_secs(10),
80            ]
81        );
82    }
83}