Rust Tokio Timeout and Retry Backoff

use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use tokio::time::{sleep, timeout};

#[derive(Debug)]
enum FetchError {
    Exhausted { attempts: u32, last: String },
}

async fn flaky_fetch(calls: &AtomicUsize) -> Result<String, String> {
    let call = calls.fetch_add(1, Ordering::SeqCst) + 1;
    match call {
        1 | 2 => {
            sleep(Duration::from_millis(80)).await;
            Err(format!("upstream returned 503 on call {call}"))
        }
        3 => Ok("payload".to_string()),
        other => Err(format!("unexpected call {other}")),
    }
}

async fn retry_with_deadline(
    calls: &AtomicUsize,
    max_attempts: u32,
) -> Result<String, FetchError> {
    let mut last = String::from("no attempt executed");

    for attempt in 1..=max_attempts {
        let budget = Duration::from_millis(120 + u64::from(attempt) * 40);
        match timeout(budget, flaky_fetch(calls)).await {
            Ok(Ok(payload)) => return Ok(payload),
            Ok(Err(message)) => last = message,
            Err(_) => last = format!("attempt {attempt} exceeded {budget:?}"),
        }

        if attempt < max_attempts {
            let backoff = Duration::from_millis(10 * u64::from(attempt));
            sleep(backoff).await;
        }
    }

    Err(FetchError::Exhausted {
        attempts: max_attempts,
        last,
    })
}

#[tokio::main]
async fn main() {
    let calls = AtomicUsize::new(0);
    match retry_with_deadline(&calls, 3).await {
        Ok(payload) => println!(
            "success after {} calls: {payload}",
            calls.load(Ordering::SeqCst)
        ),
        Err(FetchError::Exhausted { attempts, last }) => {
            println!("gave up after {attempts} attempts: {last}");
        }
    }
}