canton_core/

retry.rs

//! Opt-in retry with exponential backoff.
//!
//! Retries fire only on [`crate::Error::is_retriable`] errors, up to a bounded
//! number of attempts. Callers build the request once (so a retried command
//! keeps the same `command_id` and stays de-duplication-safe) and pass an
//! operation that re-runs the RPC.

use std::future::Future;
use std::time::Duration;

use crate::Result;

/// Retry policy for unary calls.
///
/// Start from [`RetryConfig::default`] and adjust with the fluent setters;
/// `#[non_exhaustive]` so fields can be added without a breaking change.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct RetryConfig {
    /// Maximum number of attempts (including the first). `1` disables retrying.
    pub max_attempts: u32,
    /// Backoff before the first retry.
    pub initial_backoff: Duration,
    /// Upper bound the backoff doubles towards.
    pub max_backoff: Duration,
    /// Optional per-attempt timeout. When set, an attempt that exceeds it is
    /// cancelled and treated as a retriable timeout (bounding a hung call
    /// independently of the channel-level timeout).
    pub attempt_timeout: Option<Duration>,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            initial_backoff: Duration::from_millis(100),
            max_backoff: Duration::from_secs(5),
            attempt_timeout: None,
        }
    }
}

impl RetryConfig {
    /// Set the maximum number of attempts (including the first).
    #[must_use]
    pub fn with_max_attempts(mut self, max_attempts: u32) -> Self {
        self.max_attempts = max_attempts;
        self
    }

    /// Set the initial backoff before the first retry.
    #[must_use]
    pub fn with_initial_backoff(mut self, initial_backoff: Duration) -> Self {
        self.initial_backoff = initial_backoff;
        self
    }

    /// Set the ceiling the backoff doubles towards.
    #[must_use]
    pub fn with_max_backoff(mut self, max_backoff: Duration) -> Self {
        self.max_backoff = max_backoff;
        self
    }

    /// Set a per-attempt timeout: an attempt exceeding it is cancelled and
    /// retried (as a timeout) rather than blocking on the channel-level timeout.
    #[must_use]
    pub fn with_attempt_timeout(mut self, attempt_timeout: Duration) -> Self {
        self.attempt_timeout = Some(attempt_timeout);
        self
    }
}

/// Run `op`, retrying on retriable errors per `config`. With no `config`, runs
/// `op` exactly once.
///
/// # Errors
/// Returns the last error from `op` once attempts are exhausted or a
/// non-retriable error is hit.
pub async fn run_with_retry<T, F, Fut>(config: Option<&RetryConfig>, mut op: F) -> Result<T>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = Result<T>>,
{
    let Some(config) = config else {
        return op().await;
    };

    let mut attempt = 1u32;
    let mut backoff = config.initial_backoff;
    loop {
        // Bound the attempt if a per-attempt timeout is configured; a timeout is
        // a retriable outcome (the op is retried like any transient failure).
        let outcome = match config.attempt_timeout {
            Some(timeout) => match tokio::time::timeout(timeout, op()).await {
                Ok(result) => result,
                Err(_) => Err(crate::Error::Timeout),
            },
            None => op().await,
        };
        match outcome {
            Ok(value) => return Ok(value),
            Err(err) if err.is_retriable() && attempt < config.max_attempts => {
                // Canton attaches a `RetryInfo` recommendation to retryable
                // errors; when the server asks for a longer pause than the
                // local schedule, honour it (it knows why it rejected us).
                //
                // Jitter then spreads the herd — but only upwards. A server
                // recommendation is a *minimum*: `RESOURCE_EXHAUSTED` with
                // "retry in 3s" means the participant will refuse again before
                // then, so coming back at 1.5s spends an attempt on a rejection
                // that was guaranteed. Without a recommendation the local
                // backoff is ours to scatter in either direction.
                let recommended = err.retry_delay();
                let delay = recommended.map_or(backoff, |server| server.max(backoff));
                let delay = match recommended {
                    Some(minimum) => with_jitter(delay).max(minimum),
                    None => with_jitter(delay),
                };
                tokio::time::sleep(delay).await;
                backoff = (backoff * 2).min(config.max_backoff);
                attempt += 1;
            }
            Err(err) => return Err(err),
        }
    }
}

/// Equal jitter (×0.5–1.5) on a backoff delay, so retries from many clients
/// failing at once do not re-arrive in lockstep (thundering herd). Uses the
/// clock's sub-second nanos as a cheap entropy source — no RNG dependency.
fn with_jitter(backoff: Duration) -> Duration {
    let nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_or(0, |d| d.subsec_nanos());
    let factor = 0.5 + f64::from(nanos % 1024) / 1024.0;
    backoff.mul_f64(factor)
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::Error;
    use std::cell::Cell;

    fn fast() -> RetryConfig {
        RetryConfig::default()
            .with_initial_backoff(Duration::from_millis(1))
            .with_max_backoff(Duration::from_millis(1))
    }

    #[tokio::test]
    async fn retries_retriable_errors_then_succeeds() {
        let calls = Cell::new(0);
        let result: Result<u32> = run_with_retry(Some(&fast()), || {
            calls.set(calls.get() + 1);
            let n = calls.get();
            async move { if n < 3 { Err(Error::Timeout) } else { Ok(n) } }
        })
        .await;

        assert_eq!(result.unwrap(), 3);
        assert_eq!(calls.get(), 3);
    }

    #[tokio::test(start_paused = true)]
    async fn the_backoff_doubles_each_attempt_up_to_the_cap() {
        // A schedule with room to grow: 10ms initial, 80ms cap, 5 attempts.
        // Expected sleeps between attempts, before jitter: 10, 20, 40, 80, 80
        // (doubling, then clamped). Jitter scales 0.5-1.5, so assert each
        // observed gap lands in that band around the expected pre-jitter value
        // - which pins the doubling and the cap, the two things a `* -> /` or a
        // missing `.min()` would break.
        let cfg = RetryConfig::default()
            .with_max_attempts(6)
            .with_initial_backoff(Duration::from_millis(10))
            .with_max_backoff(Duration::from_millis(80));

        let starts = std::cell::RefCell::new(Vec::new());
        let calls = Cell::new(0u32);
        let _: Result<u32> = run_with_retry(Some(&cfg), || {
            starts.borrow_mut().push(tokio::time::Instant::now());
            calls.set(calls.get() + 1);
            async move { Err::<u32, _>(Error::Timeout) }
        })
        .await;

        let t = starts.borrow();
        let gaps: Vec<u64> = t
            .windows(2)
            .map(|w| u64::try_from((w[1] - w[0]).as_millis()).unwrap_or(u64::MAX))
            .collect();
        let expected = [10u64, 20, 40, 80, 80];
        assert_eq!(gaps.len(), expected.len(), "one gap per retry: {gaps:?}");
        for (g, e) in gaps.iter().zip(expected) {
            // Jitter scales 0.5-1.5; keep the check in integer arithmetic to
            // pin the doubling and the cap without a float cast.
            assert!(
                *g * 2 >= e && *g * 2 <= e * 3,
                "gap {g}ms outside jitter band of expected {e}ms; gaps={gaps:?}"
            );
        }
    }

    #[tokio::test(start_paused = true)]
    async fn it_stops_after_max_attempts_rather_than_looping() {
        // The mutation pass flagged loop-termination mutants as *hangs*, not
        // clean failures - nothing bounded the attempt count in time. This
        // asserts the exact number of calls, so a broken stop fails loudly
        // instead of spinning until the harness kills it.
        let cfg = RetryConfig::default()
            .with_max_attempts(4)
            .with_initial_backoff(Duration::from_millis(1))
            .with_max_backoff(Duration::from_millis(1));
        let calls = Cell::new(0u32);
        let result: Result<u32> = run_with_retry(Some(&cfg), || {
            calls.set(calls.get() + 1);
            async move { Err::<u32, _>(Error::Timeout) }
        })
        .await;
        assert!(result.is_err());
        assert_eq!(
            calls.get(),
            4,
            "exactly max_attempts calls, no more, no fewer"
        );
    }

    #[tokio::test(start_paused = true)]
    async fn a_server_recommended_delay_stretches_the_backoff() {
        // The op fails once with a Canton-style status carrying
        // `RetryInfo { retry_delay: 3s }`, far above the local 1ms schedule.
        // With the paused clock, the elapsed time measures the actual sleep.
        let delay = Duration::from_secs(3);
        let status = {
            use tonic_types::{ErrorDetails, StatusExt as _};
            let mut details = ErrorDetails::new();
            details.set_retry_info(Some(delay));
            tonic::Status::with_error_details(tonic::Code::Unavailable, "wait", details)
        };

        let started = tokio::time::Instant::now();
        let calls = Cell::new(0);
        let result: Result<u32> = run_with_retry(Some(&fast()), || {
            calls.set(calls.get() + 1);
            let n = calls.get();
            let err = Error::from(status.clone());
            async move { if n == 1 { Err(err) } else { Ok(n) } }
        })
        .await;

        assert_eq!(result.unwrap(), 2);
        // The recommendation is a floor, not a target: jitter may stretch the
        // wait but must never bring the retry back before the participant said
        // it would answer.
        assert!(
            started.elapsed() >= delay,
            "the server's delay is a minimum, slept only {:?}",
            started.elapsed()
        );
    }

    #[tokio::test]
    async fn does_not_retry_non_retriable_errors() {
        let calls = Cell::new(0);
        let result: Result<u32> = run_with_retry(Some(&fast()), || {
            calls.set(calls.get() + 1);
            async move { Err(Error::InvalidRequest("nope".to_string())) }
        })
        .await;

        assert!(result.is_err());
        assert_eq!(calls.get(), 1, "non-retriable errors are not retried");
    }

    #[tokio::test]
    async fn gives_up_after_max_attempts() {
        let calls = Cell::new(0);
        let result: Result<u32> = run_with_retry(Some(&fast()), || {
            calls.set(calls.get() + 1);
            async move { Err(Error::Timeout) }
        })
        .await;

        assert!(result.is_err());
        assert_eq!(calls.get(), 3, "stops at max_attempts");
    }

    #[tokio::test]
    async fn attempt_timeout_bounds_a_hung_attempt() {
        let calls = Cell::new(0);
        let config = fast()
            .with_max_attempts(2)
            .with_attempt_timeout(Duration::from_millis(5));
        let result: Result<u32> = run_with_retry(Some(&config), || {
            calls.set(calls.get() + 1);
            async move {
                // Hangs well past the per-attempt timeout, so each attempt is
                // cancelled and retried as a timeout.
                tokio::time::sleep(Duration::from_secs(30)).await;
                Ok(1)
            }
        })
        .await;

        assert!(result.is_err(), "every attempt times out");
        assert_eq!(calls.get(), 2, "the hung attempt is bounded and retried");
    }

    #[tokio::test]
    async fn without_config_runs_exactly_once() {
        let calls = Cell::new(0);
        let result: Result<u32> = run_with_retry(None, || {
            calls.set(calls.get() + 1);
            async move { Err(Error::Timeout) }
        })
        .await;

        assert!(result.is_err());
        assert_eq!(calls.get(), 1);
    }
}