graceful_worker/

backoff.rs

//! How long a failing operation waits before trying again — and how a
//! stop cuts the wait short.
//!
//! # What this adds over the other backoff crates
//!
//! `backoff`, `backon` and `exponential-backoff` compute a schedule and
//! sleep it. None of them takes a shutdown signal, so a worker that must
//! stop promptly has to wrap every wait in its own `select!` against
//! whatever it uses for cancellation, and get that right at each call site.
//!
//! [`Backoff::wait`] takes a [`Watcher`] and returns `false` when a stop
//! arrived mid-wait, so the loop reads:
//!
//! ```
//! # #[tokio::main(flavor = "current_thread")]
//! # async fn main() {
//! use graceful_worker::{Backoff, Shutdown};
//!
//! let shutdown = Shutdown::new();
//! let watcher = shutdown.watcher();
//! let mut backoff = Backoff::new();
//!
//! shutdown.stop();
//!
//! // A fifteen-minute delay, abandoned at once rather than in fifteen
//! // minutes.
//! for _ in 0..20 {
//!     backoff.failed();
//! }
//! assert!(!backoff.wait(&watcher).await);
//! # }
//! ```
//!
//! # On slicing, honestly
//!
//! [`Backoff::wait`] sleeps the delay in bounded slices ([`with_slice`]).
//! **That is not what makes the wait interruptible.** Each slice is a
//! [`Watcher::sleep`], which is a `select!` against the cancellation token
//! and so returns the moment a stop arrives — a single unsliced sleep would
//! be exactly as responsive.
//!
//! Slicing therefore buys nothing for shutdown latency here, and an earlier
//! version of this documentation claimed otherwise. It is kept because it
//! bounds the granularity of the wait for anyone who wants that — a
//! progress tick, a slice-boundary check of some other condition — and
//! because [`Duration::ZERO`] disables it. If you have no such need, leave
//! it alone: the default costs one extra timer per minute of waiting.
//!
//! The integration with [`Watcher`] is the reason to use this crate. The
//! slicing is a knob.
//!
//! [`with_slice`]: Backoff::with_slice
//! [`Watcher::sleep`]: crate::Watcher::sleep
//!
//! # Why any of this is configurable
//!
//! The defaults are a working set, not a recommendation: five seconds,
//! doubling to a fifteen-minute ceiling, sliced at a minute. What matters
//! for a given system is the ceiling and the platform's grace period, and
//! only the caller knows either.

use std::time::Duration;

use crate::shutdown::Watcher;

/// The wait before the first retry, by default.
pub const DEFAULT_INITIAL_DELAY: Duration = Duration::from_secs(5);

/// The longest wait between attempts, by default.
pub const DEFAULT_MAX_DELAY: Duration = Duration::from_secs(900);

/// The longest single sleep, by default.
///
/// Also the worst case for how long a stop goes unnoticed during a wait.
pub const DEFAULT_SLICE: Duration = Duration::from_secs(60);

/// What the delay is multiplied by after each failure, by default.
pub const DEFAULT_FACTOR: u32 = 2;

/// The retry state of one loop.
///
/// A loop holds one of these, tells it when an attempt succeeded or failed,
/// and asks it to wait. Nothing else in the loop needs to know about
/// multiplying, ceilings or slicing.
///
/// # Examples
///
/// ```
/// use std::time::Duration;
/// use graceful_worker::Backoff;
///
/// let mut backoff = Backoff::new();
/// assert_eq!(backoff.delay(), Duration::from_secs(5));
///
/// backoff.failed();
/// assert_eq!(backoff.delay(), Duration::from_secs(10));
/// backoff.failed();
/// assert_eq!(backoff.delay(), Duration::from_secs(20));
///
/// // Any success puts it straight back to the start.
/// backoff.succeeded();
/// assert_eq!(backoff.delay(), Duration::from_secs(5));
/// assert_eq!(backoff.attempt(), 0);
/// ```
///
/// A schedule of your own:
///
/// ```
/// use std::time::Duration;
/// use graceful_worker::Backoff;
///
/// let backoff = Backoff::new()
///     .with_initial_delay(Duration::from_millis(100))
///     .with_max_delay(Duration::from_secs(30))
///     .with_slice(Duration::from_secs(5))
///     .with_factor(3);
///
/// assert_eq!(backoff.delay(), Duration::from_millis(100));
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Backoff {
    /// What the next wait will be.
    delay: Duration,
    /// How many consecutive failures there have been.
    attempt: u32,
    /// Where a reset returns to.
    initial: Duration,
    /// The ceiling.
    max: Duration,
    /// The longest single sleep. Zero means do not slice.
    slice: Duration,
    /// The multiplier applied after each failure.
    factor: u32,
}

impl Backoff {
    /// A fresh backoff on the default schedule.
    #[must_use]
    pub const fn new() -> Self {
        Self {
            delay: DEFAULT_INITIAL_DELAY,
            attempt: 0,
            initial: DEFAULT_INITIAL_DELAY,
            max: DEFAULT_MAX_DELAY,
            slice: DEFAULT_SLICE,
            factor: DEFAULT_FACTOR,
        }
    }

    /// Set the first delay, and reset to it.
    ///
    /// Clamped to the ceiling, so an initial delay longer than the maximum
    /// is the maximum rather than a schedule that counts downwards.
    #[must_use]
    pub const fn with_initial_delay(mut self, initial: Duration) -> Self {
        self.initial = initial;
        self.delay = initial;
        self.clamp()
    }

    /// Set the ceiling.
    #[must_use]
    pub const fn with_max_delay(mut self, max: Duration) -> Self {
        self.max = max;
        self.clamp()
    }

    /// Set the longest single sleep [`Backoff::wait`] performs.
    ///
    /// This is the worst case for how long a stop goes unnoticed during a
    /// wait, so it should be comfortably shorter than the platform's grace
    /// period. [`Duration::ZERO`] disables slicing, which restores the
    /// behaviour of every other backoff crate — including its drawback.
    #[must_use]
    pub const fn with_slice(mut self, slice: Duration) -> Self {
        self.slice = slice;
        self
    }

    /// Set the multiplier applied after each failure.
    ///
    /// A factor of 1 is a constant delay. Zero is treated as 1, because a
    /// backoff that multiplies by zero would retry instantly forever.
    #[must_use]
    pub const fn with_factor(mut self, factor: u32) -> Self {
        self.factor = if factor == 0 { 1 } else { factor };
        self
    }

    /// Keep the current and initial delays within the ceiling.
    const fn clamp(mut self) -> Self {
        if self.initial.as_nanos() > self.max.as_nanos() {
            self.initial = self.max;
        }
        if self.delay.as_nanos() > self.max.as_nanos() {
            self.delay = self.max;
        }
        self
    }

    /// How long to wait before the next attempt.
    #[must_use]
    pub const fn delay(&self) -> Duration {
        self.delay
    }

    /// How many consecutive failures there have been.
    ///
    /// Useful as a dimension on a retry metric.
    #[must_use]
    pub const fn attempt(&self) -> u32 {
        self.attempt
    }

    /// Whether anything has failed since the last success.
    ///
    /// A loop uses this to decide whether an empty queue is unremarkable or
    /// worth a log line.
    #[must_use]
    pub const fn is_retrying(&self) -> bool {
        self.attempt > 0
    }

    /// The ceiling this backoff will not exceed.
    #[must_use]
    pub const fn max_delay(&self) -> Duration {
        self.max
    }

    /// Record a failure: multiply the delay, up to the ceiling.
    ///
    /// Returns the delay that was in force *for this failure*, which is
    /// what a metric should record — not the multiplied value the next one
    /// will use.
    pub const fn failed(&mut self) -> Duration {
        let current = self.delay;

        // Saturating rather than wrapping: at the ceiling this is a no-op,
        // and no arithmetic here can overflow into a tiny delay.
        let raised = self.delay.saturating_mul(self.factor);
        self.delay = if raised.as_nanos() > self.max.as_nanos() {
            self.max
        } else {
            raised
        };
        self.attempt = self.attempt.saturating_add(1);

        current
    }

    /// Record a success: back to the initial delay, attempt count zero.
    ///
    /// Any success resets it, not just a run of them.
    pub const fn succeeded(&mut self) {
        self.delay = self.initial;
        self.attempt = 0;
    }

    /// The slice to sleep for, given how much of a delay is left.
    ///
    /// Never zero while `remaining` is non-zero, which is what stops
    /// [`Backoff::wait`] spinning.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::time::Duration;
    /// use graceful_worker::Backoff;
    ///
    /// let backoff = Backoff::new();
    ///
    /// // A short remainder is slept in one go.
    /// assert_eq!(backoff.sleep_slice(Duration::from_secs(3)), Duration::from_secs(3));
    /// // A long one is sliced, so shutdown stays responsive.
    /// assert_eq!(backoff.sleep_slice(Duration::from_secs(900)), Duration::from_secs(60));
    /// ```
    #[must_use]
    pub const fn sleep_slice(&self, remaining: Duration) -> Duration {
        if self.slice.is_zero() || remaining.as_nanos() <= self.slice.as_nanos() {
            remaining
        } else {
            self.slice
        }
    }

    /// Wait out the current delay, in slices, unless a stop arrives first.
    ///
    /// Returns `true` if the whole delay elapsed, `false` if it was cut
    /// short — in which case the loop should stop rather than retry.
    ///
    /// This is the method the crate exists for. See the module
    /// documentation.
    ///
    /// # Examples
    ///
    /// The usual loop shape. Marked `no_run` because running it would do
    /// exactly what it says: wait five real seconds.
    ///
    /// ```no_run
    /// # #[tokio::main(flavor = "current_thread")]
    /// # async fn main() {
    /// use graceful_worker::{Backoff, Shutdown};
    ///
    /// let shutdown = Shutdown::new();
    /// let watcher = shutdown.watcher();
    /// let mut backoff = Backoff::new();
    ///
    /// while watcher.is_running() {
    ///     let worked = false; // ... do a unit of work ...
    ///     if worked {
    ///         backoff.succeeded();
    ///     } else {
    ///         backoff.failed();
    ///         if !backoff.wait(&watcher).await {
    ///             break; // asked to stop mid-wait
    ///         }
    ///     }
    ///     # break;
    /// }
    /// # }
    /// ```
    pub async fn wait(&self, watcher: &Watcher) -> bool {
        let mut remaining = self.delay;

        while !remaining.is_zero() {
            let slice = self.sleep_slice(remaining);
            if !watcher.sleep(slice).await {
                return false;
            }
            remaining = remaining.saturating_sub(slice);
        }

        true
    }
}

impl Default for Backoff {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use crate::shutdown::Shutdown;

    use super::*;

    #[test]
    fn the_default_schedule_starts_at_five_seconds() {
        assert_eq!(Backoff::new().delay(), DEFAULT_INITIAL_DELAY);
        assert_eq!(Backoff::new().max_delay(), DEFAULT_MAX_DELAY);
    }

    #[test]
    fn the_delay_multiplies_on_each_failure() {
        let mut backoff = Backoff::new();
        let seen: Vec<u64> = (0..6).map(|_| backoff.failed().as_secs()).collect();
        assert_eq!(seen, [5, 10, 20, 40, 80, 160]);
    }

    #[test]
    fn the_delay_stops_at_the_ceiling() {
        let mut backoff = Backoff::new();
        for _ in 0..50 {
            backoff.failed();
        }
        assert_eq!(backoff.delay(), DEFAULT_MAX_DELAY);
    }

    #[test]
    fn the_ceiling_is_not_overshot_on_the_way_up() {
        // 10, 20, 40, 80, 160, 320, 640, and then a doubling that would
        // reach 1280.
        let mut backoff = Backoff::new();
        for _ in 0..7 {
            backoff.failed();
        }
        assert_eq!(backoff.delay(), Duration::from_secs(640));
        backoff.failed();
        assert_eq!(backoff.delay(), DEFAULT_MAX_DELAY, "clamped rather than 1280");
    }

    #[test]
    fn any_success_resets_it_completely() {
        let mut backoff = Backoff::new();
        for _ in 0..10 {
            backoff.failed();
        }
        assert!(backoff.is_retrying());

        backoff.succeeded();

        assert_eq!(backoff.delay(), DEFAULT_INITIAL_DELAY);
        assert_eq!(backoff.attempt(), 0);
        assert!(!backoff.is_retrying());
    }

    #[test]
    fn failed_reports_the_delay_that_applied_not_the_next_one() {
        // A metric should chart the actual schedule.
        let mut backoff = Backoff::new();
        assert_eq!(backoff.failed(), Duration::from_secs(5));
        assert_eq!(backoff.delay(), Duration::from_secs(10));
    }

    #[test]
    fn the_schedule_is_configurable() {
        let mut backoff = Backoff::new()
            .with_initial_delay(Duration::from_millis(100))
            .with_max_delay(Duration::from_secs(1))
            .with_factor(3);

        assert_eq!(backoff.delay(), Duration::from_millis(100));
        assert_eq!(backoff.failed(), Duration::from_millis(100));
        assert_eq!(backoff.delay(), Duration::from_millis(300));
        backoff.failed();
        assert_eq!(backoff.delay(), Duration::from_millis(900));
        backoff.failed();
        assert_eq!(backoff.delay(), Duration::from_secs(1), "clamped");

        backoff.succeeded();
        assert_eq!(backoff.delay(), Duration::from_millis(100));
    }

    #[test]
    fn an_initial_longer_than_the_maximum_is_clamped() {
        // Otherwise the schedule would count downwards.
        let backoff = Backoff::new()
            .with_max_delay(Duration::from_secs(2))
            .with_initial_delay(Duration::from_secs(30));

        assert_eq!(backoff.delay(), Duration::from_secs(2));
    }

    #[test]
    fn a_factor_of_one_is_a_constant_delay() {
        let mut backoff = Backoff::new().with_factor(1);
        for _ in 0..5 {
            assert_eq!(backoff.failed(), DEFAULT_INITIAL_DELAY);
        }
    }

    #[test]
    fn a_factor_of_zero_is_treated_as_one() {
        // A backoff that multiplied by zero would retry instantly forever.
        let mut backoff = Backoff::new().with_factor(0);
        backoff.failed();
        assert_eq!(backoff.delay(), DEFAULT_INITIAL_DELAY);
    }

    #[test]
    fn a_long_wait_is_slept_in_slices() {
        let backoff = Backoff::new();
        assert_eq!(backoff.sleep_slice(DEFAULT_MAX_DELAY), DEFAULT_SLICE);
        assert_eq!(backoff.sleep_slice(Duration::from_secs(61)), DEFAULT_SLICE);
        assert_eq!(backoff.sleep_slice(Duration::from_secs(60)), DEFAULT_SLICE);
        assert_eq!(
            backoff.sleep_slice(Duration::from_secs(1)),
            Duration::from_secs(1)
        );
        assert_eq!(backoff.sleep_slice(Duration::ZERO), Duration::ZERO);
    }

    #[test]
    fn a_zero_slice_disables_slicing() {
        let backoff = Backoff::new().with_slice(Duration::ZERO);
        assert_eq!(backoff.sleep_slice(DEFAULT_MAX_DELAY), DEFAULT_MAX_DELAY);
    }

    #[test]
    fn a_slice_is_never_zero_while_time_remains() {
        // What stops `wait` spinning.
        for slice in [Duration::ZERO, Duration::from_secs(1), Duration::from_secs(90)] {
            let backoff = Backoff::new().with_slice(slice);
            assert!(!backoff.sleep_slice(Duration::from_secs(30)).is_zero());
        }
    }

    #[tokio::test(start_paused = true)]
    async fn waiting_serves_the_whole_delay_when_nothing_stops_it() {
        let shutdown = Shutdown::new();
        let watcher = shutdown.watcher();
        let mut backoff = Backoff::new();
        for _ in 0..20 {
            backoff.failed();
        }
        assert_eq!(backoff.delay(), DEFAULT_MAX_DELAY);

        assert!(backoff.wait(&watcher).await, "fifteen minutes, in slices");
    }

    #[tokio::test(start_paused = true)]
    async fn waiting_is_abandoned_when_a_stop_arrives() {
        // The reason this crate exists: a fifteen-minute backoff must not
        // outlast a thirty-second grace period.
        let shutdown = Shutdown::new();
        let watcher = shutdown.watcher();
        let mut backoff = Backoff::new();
        for _ in 0..20 {
            backoff.failed();
        }

        let waiting = tokio::spawn(async move { backoff.wait(&watcher).await });
        tokio::task::yield_now().await;
        shutdown.stop();

        assert!(!waiting.await.expect("the waiting task"));
    }

    #[tokio::test(start_paused = true)]
    async fn a_stop_is_noticed_within_one_slice() {
        // The property, stated as an elapsed-time assertion rather than as
        // a description of the loop.
        let shutdown = Shutdown::new();
        let watcher = shutdown.watcher();
        let backoff = Backoff::new().with_slice(Duration::from_secs(10));
        let mut backoff = backoff;
        for _ in 0..20 {
            backoff.failed();
        }

        let started = tokio::time::Instant::now();
        let waiting = tokio::spawn(async move { backoff.wait(&watcher).await });

        // Let one slice pass, then ask to stop partway through the next.
        tokio::time::sleep(Duration::from_secs(15)).await;
        shutdown.stop();
        assert!(!waiting.await.expect("the waiting task"));

        let elapsed = started.elapsed();
        assert!(
            elapsed < Duration::from_secs(30),
            "a stop must be noticed within a slice, not after the whole \
             delay; took {elapsed:?}"
        );
    }

    #[tokio::test(start_paused = true)]
    async fn slicing_is_not_what_makes_the_wait_interruptible() {
        // The corrected claim, as a test. With slicing switched off
        // entirely the wait is a single 15-minute sleep, and a stop still
        // ends it at once -- because `Watcher::sleep` is a `select!`
        // against the token, not because the sleep was chopped up.
        let shutdown = Shutdown::new();
        let watcher = shutdown.watcher();
        let mut backoff = Backoff::new().with_slice(Duration::ZERO);
        for _ in 0..20 {
            backoff.failed();
        }
        assert_eq!(backoff.delay(), DEFAULT_MAX_DELAY);
        assert_eq!(
            backoff.sleep_slice(DEFAULT_MAX_DELAY),
            DEFAULT_MAX_DELAY,
            "one unsliced sleep"
        );

        let started = tokio::time::Instant::now();
        let waiting = tokio::spawn(async move { backoff.wait(&watcher).await });
        tokio::time::sleep(Duration::from_secs(5)).await;
        shutdown.stop();

        assert!(!waiting.await.expect("the waiting task"));
        assert!(
            started.elapsed() < Duration::from_secs(30),
            "an unsliced wait is abandoned just as promptly; took {:?}",
            started.elapsed()
        );
    }

    #[tokio::test(start_paused = true)]
    async fn a_zero_delay_waits_for_nothing() {
        let shutdown = Shutdown::new();
        let watcher = shutdown.watcher();
        let backoff = Backoff::new().with_initial_delay(Duration::ZERO);
        assert!(backoff.wait(&watcher).await);
    }

    #[test]
    fn a_default_backoff_is_a_new_one() {
        assert_eq!(Backoff::default(), Backoff::new());
    }
}