1use std::time::Duration;
15
16use crate::error_category::{ErrorCategory, classify_anyhow_error};
17
18#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)]
20pub struct RetryPolicy {
21 pub max_attempts: u32,
23 pub initial_delay: Duration,
24 pub max_delay: Duration,
25 pub multiplier: f64,
26 pub jitter: f64,
27}
28
29impl RetryPolicy {
30 pub fn new(max_attempts: u32, initial_delay: Duration, max_delay: Duration, multiplier: f64) -> Self {
31 Self {
32 max_attempts: max_attempts.max(1),
33 initial_delay,
34 max_delay,
35 multiplier: multiplier.max(1.0),
36 jitter: 0.0,
37 }
38 }
39
40 pub fn from_retries(max_retries: u32, initial_delay: Duration, max_delay: Duration, multiplier: f64) -> Self {
41 Self::new(max_retries.saturating_add(1), initial_delay, max_delay, multiplier)
42 }
43
44 pub fn simple(max_retries: u32, base_delay_ms: u64, max_delay_ms: u64) -> Self {
50 Self::from_retries(max_retries, Duration::from_millis(base_delay_ms), Duration::from_millis(max_delay_ms), 2.0)
51 }
52
53 pub fn delay_for_attempt(&self, attempt_index: u32) -> Duration {
54 let multiplier = self.multiplier.powi(attempt_index as i32);
55 let base_delay = Duration::try_from_secs_f64(self.initial_delay.as_secs_f64() * multiplier)
56 .unwrap_or(self.max_delay)
57 .min(self.max_delay);
58
59 if !self.jitter.is_finite() || self.jitter <= 0.0 {
60 return base_delay;
61 }
62
63 #[allow(clippy::cast_sign_loss)]
64 let max_jitter_ms = (base_delay.as_millis() as f64 * self.jitter)
65 .round()
66 .clamp(0.0, u64::MAX as f64) as u64;
67 if max_jitter_ms == 0 {
68 return base_delay;
69 }
70
71 let offset = (u64::from(attempt_index) * 31) % max_jitter_ms.saturating_add(1);
72 base_delay.saturating_add(Duration::from_millis(offset))
73 }
74
75 pub fn decision_for_category(
76 &self,
77 category: ErrorCategory,
78 attempt_index: u32,
79 retry_after: Option<Duration>,
80 ) -> RetryDecision {
81 let has_remaining_attempts = attempt_index.saturating_add(1) < self.max_attempts;
82 if !category.is_retryable() || !has_remaining_attempts {
83 return RetryDecision {
84 category,
85 retryable: false,
86 delay: None,
87 retry_after,
88 };
89 }
90
91 let delay = retry_after.unwrap_or_else(|| self.delay_for_attempt(attempt_index));
92 RetryDecision {
93 category,
94 retryable: true,
95 delay: Some(delay),
96 retry_after,
97 }
98 }
99
100 pub fn classify_anyhow(&self, error: &anyhow::Error) -> RetryDecision {
107 let category = classify_anyhow_error(error);
108 RetryDecision {
109 category,
110 retryable: category.is_retryable(),
111 delay: None,
112 retry_after: None,
113 }
114 }
115
116 pub fn classify_status(&self, status: u16) -> RetryDecision {
120 let category = match status {
121 429 => ErrorCategory::RateLimit,
122 500 | 502 | 504 => ErrorCategory::Network,
123 503 => ErrorCategory::ServiceUnavailable,
124 401 | 403 => ErrorCategory::Authentication,
125 _ => ErrorCategory::ExecutionError,
126 };
127 RetryDecision {
128 category,
129 retryable: category.is_retryable(),
130 delay: None,
131 retry_after: None,
132 }
133 }
134}
135
136impl Default for RetryPolicy {
137 fn default() -> Self {
138 Self::from_retries(2, Duration::from_secs(1), Duration::from_secs(60), 2.0)
139 }
140}
141
142#[derive(Debug, Clone, PartialEq, Eq)]
144pub struct RetryDecision {
145 pub category: ErrorCategory,
146 pub retryable: bool,
147 pub delay: Option<Duration>,
148 pub retry_after: Option<Duration>,
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154
155 #[test]
156 fn default_policy_allows_two_retries() {
157 let policy = RetryPolicy::default();
158 assert_eq!(policy.max_attempts, 3);
159 assert_eq!(policy.initial_delay, Duration::from_secs(1));
160 assert_eq!(policy.max_delay, Duration::from_secs(60));
161 }
162
163 #[test]
164 fn classify_status_rate_limit() {
165 let policy = RetryPolicy::default();
166 let decision = policy.classify_status(429);
167 assert!(decision.retryable);
168 assert_eq!(decision.category, ErrorCategory::RateLimit);
169 }
170
171 #[test]
172 fn classify_status_server_error() {
173 let policy = RetryPolicy::default();
174 let decision = policy.classify_status(503);
175 assert!(decision.retryable);
176 assert_eq!(decision.category, ErrorCategory::ServiceUnavailable);
177 }
178
179 #[test]
180 fn classify_status_auth_not_retryable() {
181 let policy = RetryPolicy::default();
182 let decision = policy.classify_status(401);
183 assert!(!decision.retryable);
184 assert_eq!(decision.category, ErrorCategory::Authentication);
185 }
186
187 #[test]
188 fn classify_anyhow_network_error() {
189 let policy = RetryPolicy::default();
190 let err = anyhow::anyhow!("connection refused");
191 let decision = policy.classify_anyhow(&err);
192 assert!(decision.retryable);
193 }
194
195 #[test]
196 fn simple_policy_matches_bit_shift_doubling() {
197 let policy = RetryPolicy::simple(10, 1000, 5000);
200 let legacy = |attempt: u32| -> u64 { 1000u64.saturating_mul(1u64 << attempt.min(16)).min(5000) };
201 for attempt in 0..6 {
202 assert_eq!(
203 policy.delay_for_attempt(attempt),
204 Duration::from_millis(legacy(attempt)),
205 "delay mismatch at attempt {attempt}"
206 );
207 }
208 }
209
210 #[test]
211 fn delay_for_attempt_clamps_overflowing_backoff_to_max_delay() {
212 let policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), f64::MAX);
213
214 assert_eq!(policy.delay_for_attempt(2), Duration::from_secs(8));
215 }
216
217 #[test]
218 fn delay_for_attempt_ignores_non_finite_jitter() {
219 let mut policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), 2.0);
220 policy.jitter = f64::INFINITY;
221
222 assert_eq!(policy.delay_for_attempt(1), Duration::from_secs(2));
223 }
224
225 #[test]
226 fn delay_for_attempt_handles_huge_finite_jitter() {
227 let mut policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), 2.0);
228 policy.jitter = f64::MAX;
229
230 assert!(policy.delay_for_attempt(1) >= Duration::from_secs(2));
231 }
232
233 #[test]
234 fn decision_for_category_respects_attempt_budget() {
235 let policy = RetryPolicy::from_retries(1, Duration::from_secs(1), Duration::from_secs(8), 2.0);
236
237 let first = policy.decision_for_category(ErrorCategory::Network, 0, None);
238 assert!(first.retryable);
239 assert_eq!(first.delay, Some(Duration::from_secs(1)));
240
241 let exhausted = policy.decision_for_category(ErrorCategory::Network, 1, None);
242 assert!(!exhausted.retryable);
243 assert!(exhausted.delay.is_none());
244 }
245
246 #[test]
247 fn decision_for_category_prefers_retry_after() {
248 let policy = RetryPolicy::from_retries(3, Duration::from_secs(1), Duration::from_secs(8), 2.0);
249
250 let decision = policy.decision_for_category(ErrorCategory::RateLimit, 0, Some(Duration::from_secs(7)));
251 assert!(decision.retryable);
252 assert_eq!(decision.delay, Some(Duration::from_secs(7)));
253 assert_eq!(decision.retry_after, Some(Duration::from_secs(7)));
254 }
255}