1use std::time::Duration;
2
3pub const RETRY_INITIAL_DELAY_MS: u64 = 2000;
4pub const RETRY_MAX_DELAY_MS: u64 = 30_000;
5pub const RETRY_BACKOFF_FACTOR: u64 = 2;
6pub const MAX_RATE_LIMIT_RETRIES: u32 = 3;
7
8#[derive(Debug, Clone, Copy)]
9pub struct BackoffOutcome {
10 pub wait_ms: u64,
11 pub exceeds_budget: bool,
12}
13
14pub fn should_retry_status(status: u16) -> bool {
15 matches!(status, 429 | 500 | 502 | 503 | 504)
16}
17
18pub fn compute_backoff_delay(attempt: u32, retry_after: Option<&str>) -> BackoffOutcome {
19 if let Some(raw) = retry_after
20 && let Ok(raw_secs) = raw.parse::<f64>()
21 {
22 let target_ms = (raw_secs * 1000.0).ceil() as u64;
23 return BackoffOutcome {
24 wait_ms: target_ms.min(RETRY_MAX_DELAY_MS),
25 exceeds_budget: target_ms > RETRY_MAX_DELAY_MS,
26 };
27 }
28
29 let mut exp =
30 RETRY_INITIAL_DELAY_MS.saturating_mul(RETRY_BACKOFF_FACTOR.saturating_pow(attempt));
31 if exp > RETRY_MAX_DELAY_MS {
32 exp = RETRY_MAX_DELAY_MS;
33 }
34 let jitter = exp / 2;
35 let wait_ms = (exp / 2) + (jitter / 2);
36 BackoffOutcome {
37 wait_ms,
38 exceeds_budget: false,
39 }
40}
41
42pub async fn sleep(ms: u64) {
43 tokio::time::sleep(Duration::from_millis(ms)).await;
44}
45
46#[cfg(test)]
47pub async fn retry_on_statuses<T, E, F>(mut next: F) -> Result<T, E>
48where
49 E: std::fmt::Debug,
50 F: FnMut(u32) -> Result<T, E>,
51{
52 let mut attempt = 0;
53 loop {
54 attempt += 1;
55 if attempt > MAX_RATE_LIMIT_RETRIES + 1 {
56 break;
57 }
58 match next(attempt) {
59 Ok(value) => return Ok(value),
60 Err(err) if attempt <= MAX_RATE_LIMIT_RETRIES + 1 => {
61 if attempt > MAX_RATE_LIMIT_RETRIES {
62 return Err(err);
63 }
64 sleep(compute_backoff_delay(attempt, None).wait_ms).await;
65 }
66 Err(err) => return Err(err),
67 }
68 }
69 unreachable!()
70}