alkhttp/tests/retry_after_budget.rs

//! Budget-aware `Retry-After` wire coverage (review-002 CLI-01): a
//! counting responder that always answers `429` with a large
//! `Retry-After` must not extend the caller's wall time past
//! `max_total_retry_duration` + one attempt's request time. The
//! middleware-level truncation is pinned here through the real
//! `SharedHttpClient` stack; the per-URL throttle map semantics for
//! separate logical requests stay in `src/client/retry_after.rs`.

use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use alkhttp::client::{HttpClientConfig, SharedHttpClient};

struct CountingResponder {
    addr: std::net::SocketAddr,
    hits: Arc<AtomicU32>,
}

impl CountingResponder {
    async fn spawn(status_line: &'static str, extra_headers: &'static str) -> Self {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .expect("responder binds");
        let addr = listener.local_addr().expect("responder address");
        let hits = Arc::new(AtomicU32::new(0));
        let hits_loop = Arc::clone(&hits);
        tokio::spawn(async move {
            loop {
                let Ok((mut sock, _)) = listener.accept().await else {
                    return;
                };
                hits_loop.fetch_add(1, Ordering::SeqCst);
                let body = "";
                let response = format!(
                    "{status_line}\r\ncontent-length: {}\r\nconnection: close\r\n{extra_headers}\r\n",
                    body.len()
                );
                let _ = tokio::io::AsyncWriteExt::write_all(&mut sock, response.as_bytes()).await;
                tokio::io::AsyncWriteExt::shutdown(&mut sock).await.ok();
            }
        });
        Self { addr, hits }
    }

    fn hits(&self) -> u32 {
        self.hits.load(Ordering::SeqCst)
    }
}

fn throttling_config() -> HttpClientConfig {
    HttpClientConfig {
        request_timeout: Some(Duration::from_secs(5)),
        connect_timeout: Some(Duration::from_secs(2)),
        read_timeout: Some(Duration::from_secs(2)),
        retry: alkhttp::client::RetryConfig {
            max_retries: 50,
            initial_backoff: Duration::from_millis(10),
            max_retry_interval: Duration::from_millis(50),
        },
        max_total_retry_duration: Duration::from_secs(1),
        retry_after_ceiling: Duration::from_secs(300),
        ..HttpClientConfig::default()
    }
}

#[tokio::test]
async fn always_throttled_request_is_bounded_by_the_retry_budget() {
    let responder =
        CountingResponder::spawn("HTTP/1.1 429 Too Many Requests", "retry-after: 300\r\n").await;
    let http = SharedHttpClient::new(throttling_config()).expect("client builds");
    let started = Instant::now();
    let response = http
        .client()
        .get(format!("http://{}/throttled", responder.addr))
        .send()
        .await
        .expect("429 is a delivered response, not a transport error");
    let wall = started.elapsed();
    assert_eq!(response.status(), 429);
    assert!(
        responder.hits() >= 2,
        "the request retries into the throttle (hits: {})",
        responder.hits()
    );
    assert!(
        wall < Duration::from_secs(3),
        "wall time must be bounded by max_total_retry_duration + one attempt, took {wall:?} over {} hits",
        responder.hits()
    );
    assert!(
        responder.hits() <= 60,
        "the attempt cap is generous but not infinite, got {}",
        responder.hits()
    );
}

#[tokio::test]
async fn separate_logical_request_still_honors_the_recorded_window() {
    let responder =
        CountingResponder::spawn("HTTP/1.1 429 Too Many Requests", "retry-after: 1\r\n").await;
    let http = SharedHttpClient::new(throttling_config()).expect("client builds");
    let url = format!("http://{}/window", responder.addr);

    let first_started = Instant::now();
    let first = http.client().get(&url).send().await.expect("first 429");
    let first_wall = first_started.elapsed();
    assert_eq!(first.status(), 429);
    assert!(
        first_wall >= Duration::from_millis(700),
        "the first logical request sleeps within its own budget, took {first_wall:?}"
    );

    let second_started = Instant::now();
    let second = http.client().get(&url).send().await.expect("second 429");
    let second_wall = second_started.elapsed();
    assert_eq!(second.status(), 429);
    assert!(
        second_wall <= Duration::from_millis(1100),
        "a fresh logical request honors the recorded throttle window (shortened by the budget clamp), took {second_wall:?}"
    );
    assert!(
        responder.hits() >= 2,
        "both logical requests reached the upstream"
    );
}