Add the per-slot gate and chat_with_retry
Implemented SlotGate (one request in flight per server slot, FIFO queue, bounded) and chat_with_retry (retry only server-gone errors, jittered backoff within a window and attempt budget). chat acquires the gate and maps GateFull to Busy; Client gained a gate field. Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
This commit is contained in:
@@ -0,0 +1,356 @@
|
||||
//! 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 ¬_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}");
|
||||
}
|
||||
Reference in New Issue
Block a user