chio_guards/external/
retry.rs

1//! Retry with deterministic jitter for transient external failures.
2//!
3//! [`retry_with_jitter`] runs an async operation up to `max_retries + 1`
4//! times, sleeping between attempts with a backoff controlled by
5//! [`BackoffStrategy`] and a bounded multiplicative jitter. The sleep uses
6//! [`tokio::time::sleep`], which honors [`tokio::time::pause`] + `advance`
7//! so tests don't depend on wall-clock time.
8//!
9//! Jitter is seeded deterministically from the attempt number by default,
10//! which keeps tests reproducible; callers can override the RNG via
11//! [`retry_with_jitter_rng`].
12
13use std::future::Future;
14use std::time::Duration;
15
16use rand::rngs::StdRng;
17use rand::Rng;
18use rand::SeedableRng;
19
20/// Backoff strategy between retry attempts.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum BackoffStrategy {
23    /// Each attempt sleeps `base_delay * 2^(attempt - 1)` before jitter.
24    Exponential,
25    /// Each attempt sleeps `base_delay` before jitter.
26    Constant,
27    /// Each attempt sleeps `base_delay * attempt` before jitter.
28    Linear,
29}
30
31/// Retry configuration.
32#[derive(Debug, Clone)]
33pub struct RetryConfig {
34    /// Maximum number of retries after the initial attempt. A value of `0`
35    /// means the operation is attempted exactly once.
36    pub max_retries: u32,
37    /// Base delay for the first retry.
38    pub base_delay: Duration,
39    /// Upper bound on the sleep between attempts (before jitter is added).
40    pub max_delay: Duration,
41    /// Fraction of the computed delay to use as bounded multiplicative
42    /// jitter. Must be in `[0.0, 1.0]`; values outside that range are
43    /// clamped.
44    pub jitter_fraction: f64,
45    /// Backoff curve.
46    pub strategy: BackoffStrategy,
47}
48
49impl Default for RetryConfig {
50    fn default() -> Self {
51        Self {
52            max_retries: 3,
53            base_delay: Duration::from_millis(100),
54            max_delay: Duration::from_secs(5),
55            jitter_fraction: 0.25,
56            strategy: BackoffStrategy::Exponential,
57        }
58    }
59}
60
61/// Outcome reported by the caller's operation.
62pub type AttemptResult<T, E> = Result<T, E>;
63
64/// Run `op` with retry + jitter using a deterministic RNG seeded from
65/// `config.max_retries`. For customizable randomness see
66/// [`retry_with_jitter_rng`].
67pub async fn retry_with_jitter<F, Fut, T, E>(config: &RetryConfig, op: F) -> Result<T, E>
68where
69    F: FnMut(u32) -> Fut,
70    Fut: Future<Output = AttemptResult<T, E>>,
71{
72    let seed = u64::from(config.max_retries).wrapping_add(0x9E37_79B9_7F4A_7C15);
73    let rng = StdRng::seed_from_u64(seed);
74    retry_with_jitter_rng(config, rng, op).await
75}
76
77/// Run `op` with retry + jitter using a caller-supplied RNG.
78///
79/// `op` receives the current attempt number (1-indexed).
80pub async fn retry_with_jitter_rng<F, Fut, T, E, R>(
81    config: &RetryConfig,
82    mut rng: R,
83    mut op: F,
84) -> Result<T, E>
85where
86    F: FnMut(u32) -> Fut,
87    Fut: Future<Output = AttemptResult<T, E>>,
88    R: Rng,
89{
90    let total_attempts = config.max_retries.saturating_add(1);
91    let mut last_err: Option<E> = None;
92    for attempt in 1..=total_attempts {
93        match op(attempt).await {
94            Ok(value) => return Ok(value),
95            Err(err) => {
96                last_err = Some(err);
97                if attempt >= total_attempts {
98                    break;
99                }
100                let delay = compute_delay(config, attempt, &mut rng);
101                if !delay.is_zero() {
102                    tokio::time::sleep(delay).await;
103                }
104            }
105        }
106    }
107    match last_err {
108        Some(err) => Err(err),
109        // Unreachable in practice: total_attempts >= 1 so the loop body runs
110        // at least once and either returns Ok or records an error.
111        None => unreachable!("retry loop must have produced at least one result"),
112    }
113}
114
115fn compute_delay<R: Rng>(config: &RetryConfig, attempt: u32, rng: &mut R) -> Duration {
116    let base = config.base_delay.as_secs_f64().max(0.0);
117    let raw = match config.strategy {
118        BackoffStrategy::Constant => base,
119        BackoffStrategy::Linear => base * f64::from(attempt.max(1)),
120        BackoffStrategy::Exponential => {
121            // 2^(attempt - 1). Clamp the exponent to avoid overflow.
122            let exp = attempt.saturating_sub(1).min(30);
123            base * (1u64 << exp) as f64
124        }
125    };
126    let max_secs = config.max_delay.as_secs_f64().max(0.0);
127    let capped = raw.min(max_secs);
128    let jitter = config.jitter_fraction.clamp(0.0, 1.0);
129    let factor = if jitter == 0.0 {
130        1.0
131    } else {
132        1.0 + rng.gen_range(-jitter..=jitter)
133    };
134    let jittered = (capped * factor).max(0.0);
135    Duration::from_secs_f64(jittered)
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use std::sync::atomic::{AtomicU32, Ordering};
142    use std::sync::Arc;
143
144    #[tokio::test(flavor = "current_thread", start_paused = true)]
145    async fn succeeds_on_first_attempt() {
146        let counter = Arc::new(AtomicU32::new(0));
147        let counter_clone = Arc::clone(&counter);
148        let config = RetryConfig::default();
149        let result: Result<u32, &'static str> = retry_with_jitter(&config, |_| {
150            let counter = Arc::clone(&counter_clone);
151            async move {
152                counter.fetch_add(1, Ordering::SeqCst);
153                Ok(42)
154            }
155        })
156        .await;
157        assert_eq!(result, Ok(42));
158        assert_eq!(counter.load(Ordering::SeqCst), 1);
159    }
160
161    #[tokio::test(flavor = "current_thread", start_paused = true)]
162    async fn succeeds_after_retries() {
163        let counter = Arc::new(AtomicU32::new(0));
164        let counter_clone = Arc::clone(&counter);
165        let config = RetryConfig {
166            max_retries: 4,
167            base_delay: Duration::from_millis(10),
168            max_delay: Duration::from_millis(40),
169            jitter_fraction: 0.0,
170            strategy: BackoffStrategy::Exponential,
171        };
172        let result: Result<u32, &'static str> = retry_with_jitter(&config, move |_| {
173            let counter = Arc::clone(&counter_clone);
174            async move {
175                let n = counter.fetch_add(1, Ordering::SeqCst) + 1;
176                if n < 3 {
177                    Err("transient")
178                } else {
179                    Ok(n)
180                }
181            }
182        })
183        .await;
184        assert_eq!(result, Ok(3));
185        assert_eq!(counter.load(Ordering::SeqCst), 3);
186    }
187
188    #[tokio::test(flavor = "current_thread", start_paused = true)]
189    async fn returns_last_error_after_exhausting_retries() {
190        let counter = Arc::new(AtomicU32::new(0));
191        let counter_clone = Arc::clone(&counter);
192        let config = RetryConfig {
193            max_retries: 2,
194            base_delay: Duration::from_millis(1),
195            max_delay: Duration::from_millis(4),
196            jitter_fraction: 0.0,
197            strategy: BackoffStrategy::Constant,
198        };
199        let result: Result<u32, &'static str> = retry_with_jitter(&config, move |_| {
200            let counter = Arc::clone(&counter_clone);
201            async move {
202                counter.fetch_add(1, Ordering::SeqCst);
203                Err("always fails")
204            }
205        })
206        .await;
207        assert_eq!(result, Err("always fails"));
208        assert_eq!(counter.load(Ordering::SeqCst), 3);
209    }
210
211    #[tokio::test(flavor = "current_thread", start_paused = true)]
212    async fn zero_max_retries_runs_once() {
213        let counter = Arc::new(AtomicU32::new(0));
214        let counter_clone = Arc::clone(&counter);
215        let config = RetryConfig {
216            max_retries: 0,
217            base_delay: Duration::from_millis(1),
218            max_delay: Duration::from_millis(1),
219            jitter_fraction: 0.0,
220            strategy: BackoffStrategy::Exponential,
221        };
222        let result: Result<u32, &'static str> = retry_with_jitter(&config, move |_| {
223            let counter = Arc::clone(&counter_clone);
224            async move {
225                counter.fetch_add(1, Ordering::SeqCst);
226                Err("boom")
227            }
228        })
229        .await;
230        assert_eq!(result, Err("boom"));
231        assert_eq!(counter.load(Ordering::SeqCst), 1);
232    }
233
234    #[test]
235    fn compute_delay_caps_at_max_delay() {
236        let config = RetryConfig {
237            max_retries: 10,
238            base_delay: Duration::from_millis(100),
239            max_delay: Duration::from_millis(500),
240            jitter_fraction: 0.0,
241            strategy: BackoffStrategy::Exponential,
242        };
243        let mut rng = StdRng::seed_from_u64(1);
244        // 2^9 * 100ms = 51.2s, should be capped to 500ms.
245        let d = compute_delay(&config, 10, &mut rng);
246        assert_eq!(d, Duration::from_millis(500));
247    }
248}
chio_guards/external/retry.rs

chio_guards/external/
retry.rs
