1use std::future::Future;
9use std::time::Duration;
10
11use crate::Result;
12
13#[derive(Clone, Debug)]
18#[non_exhaustive]
19pub struct RetryConfig {
20 pub max_attempts: u32,
22 pub initial_backoff: Duration,
24 pub max_backoff: Duration,
26 pub attempt_timeout: Option<Duration>,
30}
31
32impl Default for RetryConfig {
33 fn default() -> Self {
34 Self {
35 max_attempts: 3,
36 initial_backoff: Duration::from_millis(100),
37 max_backoff: Duration::from_secs(5),
38 attempt_timeout: None,
39 }
40 }
41}
42
43impl RetryConfig {
44 #[must_use]
46 pub fn with_max_attempts(mut self, max_attempts: u32) -> Self {
47 self.max_attempts = max_attempts;
48 self
49 }
50
51 #[must_use]
53 pub fn with_initial_backoff(mut self, initial_backoff: Duration) -> Self {
54 self.initial_backoff = initial_backoff;
55 self
56 }
57
58 #[must_use]
60 pub fn with_max_backoff(mut self, max_backoff: Duration) -> Self {
61 self.max_backoff = max_backoff;
62 self
63 }
64
65 #[must_use]
68 pub fn with_attempt_timeout(mut self, attempt_timeout: Duration) -> Self {
69 self.attempt_timeout = Some(attempt_timeout);
70 self
71 }
72}
73
74pub async fn run_with_retry<T, F, Fut>(config: Option<&RetryConfig>, mut op: F) -> Result<T>
81where
82 F: FnMut() -> Fut,
83 Fut: Future<Output = Result<T>>,
84{
85 let Some(config) = config else {
86 return op().await;
87 };
88
89 let mut attempt = 1u32;
90 let mut backoff = config.initial_backoff;
91 loop {
92 let outcome = match config.attempt_timeout {
95 Some(timeout) => match tokio::time::timeout(timeout, op()).await {
96 Ok(result) => result,
97 Err(_) => Err(crate::Error::Timeout),
98 },
99 None => op().await,
100 };
101 match outcome {
102 Ok(value) => return Ok(value),
103 Err(err) if err.is_retriable() && attempt < config.max_attempts => {
104 let recommended = err.retry_delay().map(|server| {
119 if server > MAX_SERVER_DELAY {
120 tracing::warn!(
121 recommended = ?server,
122 cap = ?MAX_SERVER_DELAY,
123 "the participant recommended a retry delay above the cap; capping it"
124 );
125 }
126 server.min(MAX_SERVER_DELAY)
127 });
128 let delay = recommended.map_or(backoff, |server| server.max(backoff));
129 let delay = match recommended {
130 Some(minimum) => with_jitter(delay).max(minimum),
131 None => with_jitter(delay),
132 };
133 tokio::time::sleep(delay).await;
134 backoff = (backoff * 2).min(config.max_backoff);
135 attempt += 1;
136 }
137 Err(err) => return Err(err),
138 }
139 }
140}
141
142fn with_jitter(backoff: Duration) -> Duration {
146 let nanos = std::time::SystemTime::now()
147 .duration_since(std::time::UNIX_EPOCH)
148 .map_or(0, |d| d.subsec_nanos());
149 let factor = 0.5 + f64::from(nanos % 1024) / 1024.0;
150 Duration::try_from_secs_f64(backoff.as_secs_f64() * factor).unwrap_or(backoff)
154}
155
156pub const MAX_SERVER_DELAY: Duration = Duration::from_secs(60);
160
161#[cfg(test)]
162#[allow(clippy::unwrap_used)]
163mod tests {
164 use super::*;
165 use crate::Error;
166 use std::cell::Cell;
167
168 fn fast() -> RetryConfig {
169 RetryConfig::default()
170 .with_initial_backoff(Duration::from_millis(1))
171 .with_max_backoff(Duration::from_millis(1))
172 }
173
174 #[tokio::test]
175 async fn retries_retriable_errors_then_succeeds() {
176 let calls = Cell::new(0);
177 let result: Result<u32> = run_with_retry(Some(&fast()), || {
178 calls.set(calls.get() + 1);
179 let n = calls.get();
180 async move { if n < 3 { Err(Error::Timeout) } else { Ok(n) } }
181 })
182 .await;
183
184 assert_eq!(result.unwrap(), 3);
185 assert_eq!(calls.get(), 3);
186 }
187
188 #[tokio::test(start_paused = true)]
189 async fn the_backoff_doubles_each_attempt_up_to_the_cap() {
190 let cfg = RetryConfig::default()
197 .with_max_attempts(6)
198 .with_initial_backoff(Duration::from_millis(10))
199 .with_max_backoff(Duration::from_millis(80));
200
201 let starts = std::cell::RefCell::new(Vec::new());
202 let calls = Cell::new(0u32);
203 let _: Result<u32> = run_with_retry(Some(&cfg), || {
204 starts.borrow_mut().push(tokio::time::Instant::now());
205 calls.set(calls.get() + 1);
206 async move { Err::<u32, _>(Error::Timeout) }
207 })
208 .await;
209
210 let t = starts.borrow();
211 let gaps: Vec<u64> = t
212 .windows(2)
213 .map(|w| u64::try_from((w[1] - w[0]).as_millis()).unwrap_or(u64::MAX))
214 .collect();
215 let expected = [10u64, 20, 40, 80, 80];
216 assert_eq!(gaps.len(), expected.len(), "one gap per retry: {gaps:?}");
217 for (g, e) in gaps.iter().zip(expected) {
218 assert!(
221 *g * 2 >= e && *g * 2 <= e * 3,
222 "gap {g}ms outside jitter band of expected {e}ms; gaps={gaps:?}"
223 );
224 }
225 }
226
227 #[tokio::test(start_paused = true)]
228 async fn it_stops_after_max_attempts_rather_than_looping() {
229 let cfg = RetryConfig::default()
234 .with_max_attempts(4)
235 .with_initial_backoff(Duration::from_millis(1))
236 .with_max_backoff(Duration::from_millis(1));
237 let calls = Cell::new(0u32);
238 let result: Result<u32> = run_with_retry(Some(&cfg), || {
239 calls.set(calls.get() + 1);
240 async move { Err::<u32, _>(Error::Timeout) }
241 })
242 .await;
243 assert!(result.is_err());
244 assert_eq!(
245 calls.get(),
246 4,
247 "exactly max_attempts calls, no more, no fewer"
248 );
249 }
250
251 #[tokio::test(start_paused = true)]
252 async fn a_server_recommended_delay_stretches_the_backoff() {
253 let delay = Duration::from_secs(3);
257 let status = {
258 use tonic_types::{ErrorDetails, StatusExt as _};
259 let mut details = ErrorDetails::new();
260 details.set_retry_info(Some(delay));
261 tonic::Status::with_error_details(tonic::Code::Unavailable, "wait", details)
262 };
263
264 let started = tokio::time::Instant::now();
265 let calls = Cell::new(0);
266 let result: Result<u32> = run_with_retry(Some(&fast()), || {
267 calls.set(calls.get() + 1);
268 let n = calls.get();
269 let err = Error::from(status.clone());
270 async move { if n == 1 { Err(err) } else { Ok(n) } }
271 })
272 .await;
273
274 assert_eq!(result.unwrap(), 2);
275 assert!(
279 started.elapsed() >= delay,
280 "the server's delay is a minimum, slept only {:?}",
281 started.elapsed()
282 );
283 }
284
285 #[tokio::test(start_paused = true)]
286 async fn a_server_recommended_delay_is_capped() {
287 let status = {
290 use tonic_types::{ErrorDetails, StatusExt as _};
291 let mut details = ErrorDetails::new();
292 details.set_retry_info(Some(Duration::from_secs(365 * 86_400)));
293 tonic::Status::with_error_details(tonic::Code::Unavailable, "wait", details)
294 };
295 let started = tokio::time::Instant::now();
296 let calls = Cell::new(0);
297 let result: Result<u32> = run_with_retry(Some(&fast()), || {
298 calls.set(calls.get() + 1);
299 let n = calls.get();
300 let err = Error::from(status.clone());
301 async move { if n == 1 { Err(err) } else { Ok(n) } }
302 })
303 .await;
304 assert_eq!(result.unwrap(), 2);
305 assert!(
306 started.elapsed() >= MAX_SERVER_DELAY,
307 "{:?}",
308 started.elapsed()
309 );
310 assert!(
311 started.elapsed() <= MAX_SERVER_DELAY.mul_f64(1.5) + Duration::from_secs(1),
312 "slept {:?}",
313 started.elapsed()
314 );
315 }
316
317 #[tokio::test(start_paused = true)]
318 async fn an_absurd_json_retry_info_neither_panics_nor_hangs() {
319 for body in [
322 r#"{"errorCategory":1,"retryInfo":"1.8e19 seconds"}"#,
323 r#"{"errorCategory":1,"retryInfo":"1.3e19 seconds"}"#,
324 r#"{"cause":"busy","retryInfo":"18000000000000000000 seconds"}"#,
325 ] {
326 let started = tokio::time::Instant::now();
327 let calls = Cell::new(0);
328 let result: Result<u32> = run_with_retry(Some(&fast()), || {
329 calls.set(calls.get() + 1);
330 let n = calls.get();
331 let err = Error::Http {
332 status: 503,
333 body: body.to_string(),
334 url: None,
335 };
336 async move { if n == 1 { Err(err) } else { Ok(n) } }
337 })
338 .await;
339 assert_eq!(result.unwrap(), 2, "{body}");
340 assert!(
341 started.elapsed() <= MAX_SERVER_DELAY.mul_f64(1.5) + Duration::from_secs(1),
342 "{body}: slept {:?}",
343 started.elapsed()
344 );
345 }
346 }
347
348 #[test]
349 fn jitter_never_panics() {
350 for backoff in [Duration::ZERO, Duration::from_millis(1), Duration::MAX] {
351 let jittered = with_jitter(backoff);
352 assert!(jittered <= Duration::MAX);
353 }
354 }
355
356 #[tokio::test]
357 async fn does_not_retry_non_retriable_errors() {
358 let calls = Cell::new(0);
359 let result: Result<u32> = run_with_retry(Some(&fast()), || {
360 calls.set(calls.get() + 1);
361 async move { Err(Error::InvalidRequest("nope".to_string())) }
362 })
363 .await;
364
365 assert!(result.is_err());
366 assert_eq!(calls.get(), 1, "non-retriable errors are not retried");
367 }
368
369 #[tokio::test]
370 async fn gives_up_after_max_attempts() {
371 let calls = Cell::new(0);
372 let result: Result<u32> = run_with_retry(Some(&fast()), || {
373 calls.set(calls.get() + 1);
374 async move { Err(Error::Timeout) }
375 })
376 .await;
377
378 assert!(result.is_err());
379 assert_eq!(calls.get(), 3, "stops at max_attempts");
380 }
381
382 #[tokio::test]
383 async fn attempt_timeout_bounds_a_hung_attempt() {
384 let calls = Cell::new(0);
385 let config = fast()
386 .with_max_attempts(2)
387 .with_attempt_timeout(Duration::from_millis(5));
388 let result: Result<u32> = run_with_retry(Some(&config), || {
389 calls.set(calls.get() + 1);
390 async move {
391 tokio::time::sleep(Duration::from_secs(30)).await;
394 Ok(1)
395 }
396 })
397 .await;
398
399 assert!(result.is_err(), "every attempt times out");
400 assert_eq!(calls.get(), 2, "the hung attempt is bounded and retried");
401 }
402
403 #[tokio::test]
404 async fn without_config_runs_exactly_once() {
405 let calls = Cell::new(0);
406 let result: Result<u32> = run_with_retry(None, || {
407 calls.set(calls.get() + 1);
408 async move { Err(Error::Timeout) }
409 })
410 .await;
411
412 assert!(result.is_err());
413 assert_eq!(calls.get(), 1);
414 }
415}