Skip to main content

agent_effects/
retry.rs

1//! Retry policy: how many attempts, and how long to wait between them.
2//!
3//! Whether a failure may be retried at all is decided elsewhere, by its
4//! [`FailureClass`] and [`crate::policy`]. This module only does the counting
5//! and the arithmetic.
6
7use std::time::Duration;
8
9use serde::{Deserialize, Serialize};
10
11use crate::failure::FailureClass;
12
13/// Exponential backoff with optional jitter.
14#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
15pub struct RetryPolicy {
16    /// Total attempts allowed, including the first. `1` disables retries.
17    pub max_attempts: u32,
18    /// Delay before the first retry.
19    pub initial_delay: Duration,
20    /// Upper bound on the computed backoff. A rate limit's `retry_after`
21    /// may exceed it.
22    pub max_delay: Duration,
23    /// Growth factor per retry. Values below `1.0` are treated as `1.0`.
24    pub multiplier: f64,
25    /// Whether to randomize each delay into `[delay / 2, delay]`.
26    pub jitter: bool,
27}
28
29impl RetryPolicy {
30    /// A policy that never retries.
31    pub const NONE: Self = Self {
32        max_attempts: 1,
33        initial_delay: Duration::ZERO,
34        max_delay: Duration::ZERO,
35        multiplier: 1.0,
36        jitter: false,
37    };
38
39    /// Whether another attempt is allowed after `attempts_made` attempts.
40    pub const fn allows_another(&self, attempts_made: u32) -> bool {
41        attempts_made < self.max_attempts
42    }
43
44    /// The backoff before retry number `retry` (0 for the first retry),
45    /// before jitter.
46    pub fn backoff(&self, retry: u32) -> Duration {
47        // `max` maps NaN to 1.0 as well.
48        let multiplier = self.multiplier.max(1.0);
49        let exponent = i32::try_from(retry).unwrap_or(i32::MAX);
50        let secs = self.initial_delay.as_secs_f64() * multiplier.powi(exponent);
51        if secs.is_nan() {
52            // Only `0 * inf`: a zero initial delay stays zero.
53            return Duration::ZERO;
54        }
55        if secs >= self.max_delay.as_secs_f64() {
56            // Also catches the infinity a large exponent produces.
57            return self.max_delay;
58        }
59        Duration::from_secs_f64(secs).min(self.max_delay)
60    }
61
62    /// The delay before retry number `retry`, given the failure that caused
63    /// it and a uniform random `sample` in `[0, 1)` used for jitter.
64    ///
65    /// Taking the sample as an argument keeps this function deterministic.
66    pub fn delay(&self, retry: u32, failure: FailureClass, sample: f64) -> Duration {
67        let backoff = self.backoff(retry);
68        let jittered = if self.jitter {
69            let sample = if sample.is_finite() {
70                sample.clamp(0.0, 1.0)
71            } else {
72                0.0
73            };
74            backoff.mul_f64(0.5 + 0.5 * sample).min(backoff)
75        } else {
76            backoff
77        };
78        match failure.retry_after() {
79            Some(asked) => jittered.max(asked),
80            None => jittered,
81        }
82    }
83}
84
85impl Default for RetryPolicy {
86    fn default() -> Self {
87        Self {
88            max_attempts: 5,
89            initial_delay: Duration::from_secs(1),
90            max_delay: Duration::from_secs(30),
91            multiplier: 2.0,
92            jitter: true,
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use proptest::prelude::*;
100
101    use super::*;
102
103    #[test]
104    fn default_backoff_doubles_up_to_the_cap() {
105        let policy = RetryPolicy::default();
106        let secs: Vec<u64> = (0..7).map(|r| policy.backoff(r).as_secs()).collect();
107        assert_eq!(secs, [1, 2, 4, 8, 16, 30, 30]);
108    }
109
110    #[test]
111    fn rate_limits_may_exceed_the_cap() {
112        let policy = RetryPolicy::default();
113        let failure = FailureClass::RateLimited {
114            retry_after: Some(Duration::from_secs(120)),
115        };
116        assert_eq!(policy.delay(0, failure, 0.0), Duration::from_secs(120));
117    }
118
119    #[test]
120    fn attempt_counting_includes_the_first_attempt() {
121        assert!(!RetryPolicy::NONE.allows_another(1));
122        let policy = RetryPolicy::default();
123        assert!(policy.allows_another(4));
124        assert!(!policy.allows_another(5));
125    }
126
127    proptest! {
128        #[test]
129        fn delays_are_bounded_and_monotone(
130            initial_ms in 0u64..10_000,
131            max_ms in 0u64..600_000,
132            multiplier in prop_oneof![Just(f64::NAN), Just(f64::INFINITY), -10.0f64..10.0],
133            retry in 0u32..u32::MAX,
134            sample in prop_oneof![Just(f64::NAN), -1.0f64..2.0],
135        ) {
136            let policy = RetryPolicy {
137                max_attempts: 5,
138                initial_delay: Duration::from_millis(initial_ms),
139                max_delay: Duration::from_millis(max_ms),
140                multiplier,
141                jitter: true,
142            };
143            let backoff = policy.backoff(retry);
144            prop_assert!(backoff <= policy.max_delay);
145            prop_assert!(policy.backoff(retry.saturating_add(1)) >= backoff);
146            let delay = policy.delay(retry, FailureClass::Transient, sample);
147            prop_assert!(delay <= backoff);
148            // Float rounding may land a nanosecond under the exact half.
149            prop_assert!(delay + Duration::from_nanos(1) >= backoff / 2);
150        }
151    }
152}