boxmaker/crates/loopd/tests/retry.rs

//! Tests for the slot gate and for retrying. Do not edit.

mod support;

use loopd::llama::gate::{GateFull, SlotGate};
use loopd::llama::retry::{backoff_ms, is_retryable};
use loopd::llama::{ChatEvent, ChatMessage, ChatRequest, Client, InferError};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};
use support::{FakeServer, Reply};

const CHAT: &str = "/v1/chat/completions";

fn request(slot: u32) -> ChatRequest {
    ChatRequest {
        slot,
        messages: vec![ChatMessage::User {
            content: "hi".to_string(),
        }],
        tools: vec![],
        thinking: false,
    }
}

fn io(kind: std::io::ErrorKind) -> std::io::Error {
    std::io::Error::from(kind)
}

// ---- the gate on its own ----

#[test]
fn a_free_slot_is_taken_at_once_and_slots_are_independent() {
    let gate = SlotGate::new();
    let mut queued = Vec::new();
    let a = gate.acquire(0, 8, &mut |n| queued.push(n)).unwrap();
    let b = gate.acquire(1, 8, &mut |n| queued.push(n)).unwrap();
    assert!(queued.is_empty(), "nobody had to wait");
    drop(a);
    drop(b);
    let _again = gate.acquire(0, 8, &mut |n| queued.push(n)).unwrap();
    assert!(queued.is_empty());
}

#[test]
fn waiters_are_served_in_order_one_at_a_time() {
    let gate = Arc::new(SlotGate::new());
    let order = Arc::new(Mutex::new(Vec::new()));
    let inside = Arc::new(AtomicUsize::new(0));
    let first = gate.acquire(0, 8, &mut |_| {}).unwrap();
    let mut threads = Vec::new();
    for i in 0..4usize {
        let (gate, order, inside) = (Arc::clone(&gate), Arc::clone(&order), Arc::clone(&inside));
        threads.push(thread::spawn(move || {
            let mut ahead = None;
            let permit = gate.acquire(0, 8, &mut |n| ahead = Some(n)).unwrap();
            assert_eq!(
                inside.fetch_add(1, Ordering::SeqCst),
                0,
                "two permits for one slot at once"
            );
            order.lock().unwrap().push((i, ahead));
            thread::sleep(Duration::from_millis(20));
            inside.fetch_sub(1, Ordering::SeqCst);
            drop(permit);
        }));
        thread::sleep(Duration::from_millis(30)); // so that the arrival order is known
    }
    drop(first);
    for t in threads {
        t.join().unwrap();
    }
    // Each was told how many were ahead of it: the holder plus the earlier waiters.
    let want: Vec<(usize, Option<usize>)> = (0..4).map(|i| (i, Some(i + 1))).collect();
    assert_eq!(*order.lock().unwrap(), want);
}

#[test]
fn a_full_queue_refuses_at_once() {
    let gate = Arc::new(SlotGate::new());
    let _holder = gate.acquire(0, 1, &mut |_| {}).unwrap();
    let waiter = {
        let gate = Arc::clone(&gate);
        thread::spawn(move || drop(gate.acquire(0, 1, &mut |_| {}).unwrap()))
    };
    thread::sleep(Duration::from_millis(50));
    let started = Instant::now();
    assert!(
        matches!(gate.acquire(0, 1, &mut |_| {}), Err(GateFull)),
        "one holder and one waiter is the limit"
    );
    assert!(
        started.elapsed() < Duration::from_millis(50),
        "refusal must not wait"
    );
    assert!(
        gate.acquire(1, 1, &mut |_| {}).is_ok(),
        "another slot is unaffected"
    );
    drop(_holder);
    waiter.join().unwrap();
}

#[test]
fn a_queue_of_zero_means_no_waiting_at_all() {
    let gate = SlotGate::new();
    let _holder = gate.acquire(0, 0, &mut |_| {}).unwrap();
    assert!(matches!(gate.acquire(0, 0, &mut |_| {}), Err(GateFull)));
}

// ---- the gate inside the client ----

#[test]
fn two_requests_for_one_slot_never_overlap_at_the_server() {
    let server = FakeServer::start();
    let size = support::fixture_bytes("http", "turn1.http").len();
    // Each reply takes about 200 ms to arrive.
    server.route(CHAT, vec![Reply::fixture("turn1").trickle(size / 4, 50)]);
    let client = Arc::new(Client::new(support::test_config(&server.socket)));
    let started = Instant::now();
    let second = {
        let client = Arc::clone(&client);
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(30));
            let mut events = Vec::new();
            let done = client.chat(&request(0), &mut |e| events.push(e.clone()));
            (done.map(|d| d.content), events, Instant::now())
        })
    };
    let first = client.chat(&request(0), &mut |_| {}).unwrap();
    let first_done = Instant::now();
    let (second_result, second_events, second_done) = second.join().unwrap();
    assert_eq!(first.content.as_deref(), Some("Blue"));
    assert_eq!(second_result.unwrap().as_deref(), Some("Blue"));
    assert_eq!(second_events.first(), Some(&ChatEvent::Queued { ahead: 1 }));
    assert!(second_done > first_done);
    assert!(
        started.elapsed() >= Duration::from_millis(280),
        "the two ran one after the other"
    );
}

#[test]
fn a_full_queue_is_busy() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("turn1").head_delay(300)]);
    // Both slots report busy, so that the 300 ms of silence reads as a queue and not a stall.
    let busy = r#"[{"id":0,"is_processing":true},{"id":1,"is_processing":true}]"#;
    server.route("/slots", vec![Reply::json(200, busy)]);
    let mut cfg = support::test_config(&server.socket);
    cfg.limits.queue_len = 0;
    let client = Arc::new(Client::new(cfg));
    let holder = {
        let client = Arc::clone(&client);
        thread::spawn(move || client.chat(&request(0), &mut |_| {}).map(|_| ()))
    };
    thread::sleep(Duration::from_millis(80));
    assert!(matches!(
        client.chat(&request(0), &mut |_| {}),
        Err(InferError::Busy)
    ));
    assert!(
        client.chat(&request(1), &mut |_| {}).is_ok(),
        "another slot is free"
    );
    holder.join().unwrap().unwrap();
    assert_eq!(
        server.requests_to(CHAT).len(),
        2,
        "the refused request was never sent"
    );
}

// ---- retry ----

#[test]
fn which_errors_are_retried() {
    use std::io::ErrorKind::ConnectionRefused;
    let http = |status| InferError::Http {
        status,
        body: String::new(),
    };
    let retried = [
        InferError::Connect(io(ConnectionRefused)),
        InferError::StreamClosedEarly,
        InferError::Stalled,
        InferError::LoadTimeout,
        http(503),
    ];
    for e in &retried {
        assert!(
            is_retryable(e),
            "{e:?} means the server went away and may be back"
        );
    }
    let not_retried = [
        InferError::Busy,
        InferError::WaitTimeout,
        InferError::ThinkingOverrun,
        InferError::Protocol("x".to_string()),
        http(400),
        http(404),
        http(500),
    ];
    for e in &not_retried {
        assert!(
            !is_retryable(e),
            "{e:?} would fail again, or has already waited"
        );
    }
}

#[test]
fn backoff_follows_the_schedule_with_a_quarter_of_jitter() {
    let schedule = [2_000, 8_000, 30_000];
    assert_eq!(backoff_ms(&schedule, 1, 0), 2_000);
    assert_eq!(backoff_ms(&schedule, 2, 0), 8_000);
    assert_eq!(backoff_ms(&schedule, 3, 0), 30_000);
    assert_eq!(
        backoff_ms(&schedule, 4, 0),
        30_000,
        "past the end, the last entry"
    );
    assert_eq!(
        backoff_ms(&schedule, 0, 0),
        2_000,
        "retry numbers start at 1; 0 is treated as 1"
    );
    assert_eq!(backoff_ms(&schedule, 1, 1000), 2_500);
    assert_eq!(backoff_ms(&schedule, 1, -1000), 1_500);
    assert_eq!(backoff_ms(&schedule, 2, 500), 9_000);
    assert_eq!(
        backoff_ms(&schedule, 1, 99_999),
        2_500,
        "jitter is clamped to -1000..=1000"
    );
    assert_eq!(backoff_ms(&schedule, 1, i32::MIN), 1_500);
    assert_eq!(backoff_ms(&[], 1, 1000), 0, "no schedule, no wait");
    // Huge values must not overflow or panic.
    let _ = backoff_ms(&[u64::MAX], 1, 1000);
    let _ = backoff_ms(&[u64::MAX], u32::MAX, -1000);
}

fn run_retry(
    server: &FakeServer,
    attempts: u32,
    window_ms: u64,
) -> (Result<Option<String>, InferError>, Vec<ChatEvent>) {
    let mut cfg = support::test_config(&server.socket);
    cfg.limits.retry_attempts = attempts;
    cfg.limits.retry_window_ms = window_ms;
    let mut events = Vec::new();
    let result = Client::new(cfg).chat_with_retry(&request(0), &mut |e| events.push(e.clone()));
    (result.map(|d| d.content), events)
}

fn retrying(events: &[ChatEvent]) -> Vec<(u32, u64)> {
    events
        .iter()
        .filter_map(|e| match e {
            ChatEvent::Retrying {
                attempt, after_ms, ..
            } => Some((*attempt, *after_ms)),
            _ => None,
        })
        .collect()
}

#[test]
fn a_server_that_comes_back_is_survived() {
    let server = FakeServer::start();
    let dead = Reply::fixture("turn1").cut_after(400);
    server.route(CHAT, vec![dead.clone(), dead, Reply::fixture("turn1")]);
    let (result, events) = run_retry(&server, 4, 5_000);
    assert_eq!(result.unwrap().as_deref(), Some("Blue"));
    let retries = retrying(&events);
    assert_eq!(
        retries
            .iter()
            .map(|(attempt, _)| *attempt)
            .collect::<Vec<_>>(),
        vec![2, 3]
    );
    // The test schedule is 10, 20, 30 ms, each moved by at most a quarter.
    assert!((8..=12).contains(&retries[0].1), "{retries:?}");
    assert!((15..=25).contains(&retries[1].1), "{retries:?}");
    let sent = server.requests_to(CHAT);
    assert_eq!(sent.len(), 3);
    assert!(
        sent.iter().all(|r| r.body == sent[0].body),
        "a retry sends the same bytes again"
    );
}

#[test]
fn retries_stop_at_the_attempt_limit() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("turn1").cut_after(400)]);
    let (result, events) = run_retry(&server, 3, 5_000);
    assert!(
        matches!(result, Err(InferError::StreamClosedEarly)),
        "{result:?}"
    );
    assert_eq!(server.requests_to(CHAT).len(), 3, "three attempts in all");
    assert_eq!(retrying(&events).len(), 2);
}

#[test]
fn retries_stop_at_the_time_window() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("turn1").cut_after(400)]);
    // The first retry would wait about 10 ms, which does not fit in a 5 ms window.
    let (result, events) = run_retry(&server, 10, 5);
    assert!(
        matches!(result, Err(InferError::StreamClosedEarly)),
        "{result:?}"
    );
    assert_eq!(server.requests_to(CHAT).len(), 1);
    assert!(retrying(&events).is_empty());
}

#[test]
fn errors_that_would_fail_again_are_not_retried() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("bad_request")]);
    let (result, events) = run_retry(&server, 4, 5_000);
    assert!(
        matches!(result, Err(InferError::Http { status: 400, .. })),
        "{result:?}"
    );
    assert_eq!(server.requests_to(CHAT).len(), 1);
    assert!(retrying(&events).is_empty());
}

#[test]
fn a_loading_server_is_retried() {
    let server = FakeServer::start();
    server.route(
        CHAT,
        vec![
            Reply::json(503, r#"{"error":{"message":"Loading model"}}"#),
            Reply::fixture("turn1"),
        ],
    );
    let (result, events) = run_retry(&server, 4, 5_000);
    assert_eq!(result.unwrap().as_deref(), Some("Blue"));
    let ChatEvent::Retrying { error, .. } = events
        .iter()
        .find(|e| matches!(e, ChatEvent::Retrying { .. }))
        .unwrap()
    else {
        unreachable!()
    };
    assert!(error.contains("503"), "the event says why: {error}");
}