kube_runtime/utils/
backoff_reset_timer.rs

1use std::time::{Duration, Instant};
2
3use backoff::{backoff::Backoff, Clock, SystemClock};
4
5/// A [`Backoff`] wrapper that resets after a fixed duration has elapsed.
6pub struct ResetTimerBackoff<B, C = SystemClock> {
7    backoff: B,
8    clock: C,
9    last_backoff: Option<Instant>,
10    reset_duration: Duration,
11}
12
13impl<B: Backoff> ResetTimerBackoff<B> {
14    pub fn new(backoff: B, reset_duration: Duration) -> Self {
15        Self::new_with_custom_clock(backoff, reset_duration, SystemClock {})
16    }
17}
18
19impl<B: Backoff, C: Clock> ResetTimerBackoff<B, C> {
20    fn new_with_custom_clock(backoff: B, reset_duration: Duration, clock: C) -> Self {
21        Self {
22            backoff,
23            clock,
24            last_backoff: None,
25            reset_duration,
26        }
27    }
28}
29
30impl<B: Backoff, C: Clock> Backoff for ResetTimerBackoff<B, C> {
31    fn next_backoff(&mut self) -> Option<Duration> {
32        if let Some(last_backoff) = self.last_backoff {
33            if self.clock.now() > last_backoff + self.reset_duration {
34                tracing::debug!(
35                    ?last_backoff,
36                    reset_duration = ?self.reset_duration,
37                    "Resetting backoff, since reset duration has expired"
38                );
39                self.backoff.reset();
40            }
41        }
42        self.last_backoff = Some(self.clock.now());
43        self.backoff.next_backoff()
44    }
45
46    fn reset(&mut self) {
47        // Do not even bother trying to reset here, since `next_backoff` will take care of this when the timer expires.
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use backoff::{backoff::Backoff, Clock};
54    use tokio::time::advance;
55
56    use super::ResetTimerBackoff;
57    use crate::utils::stream_backoff::tests::LinearBackoff;
58    use std::time::{Duration, Instant};
59
60    #[tokio::test]
61    async fn should_reset_when_timer_expires() {
62        tokio::time::pause();
63        let mut backoff = ResetTimerBackoff::new_with_custom_clock(
64            LinearBackoff::new(Duration::from_secs(2)),
65            Duration::from_secs(60),
66            TokioClock,
67        );
68        assert_eq!(backoff.next_backoff(), Some(Duration::from_secs(2)));
69        advance(Duration::from_secs(40)).await;
70        assert_eq!(backoff.next_backoff(), Some(Duration::from_secs(4)));
71        advance(Duration::from_secs(40)).await;
72        assert_eq!(backoff.next_backoff(), Some(Duration::from_secs(6)));
73        advance(Duration::from_secs(80)).await;
74        assert_eq!(backoff.next_backoff(), Some(Duration::from_secs(2)));
75        advance(Duration::from_secs(80)).await;
76        assert_eq!(backoff.next_backoff(), Some(Duration::from_secs(2)));
77    }
78
79    struct TokioClock;
80
81    impl Clock for TokioClock {
82        fn now(&self) -> Instant {
83            tokio::time::Instant::now().into_std()
84        }
85    }
86}