374 lines
11 KiB
Rust
374 lines
11 KiB
Rust
//! Load balancing integration tests
|
|
//!
|
|
//! Tests for various load balancing policies: round_robin, random, cache_aware, etc.
|
|
|
|
use axum::{
|
|
body::Body,
|
|
extract::Request,
|
|
http::{header::CONTENT_TYPE, StatusCode},
|
|
};
|
|
use serde_json::json;
|
|
use smg::config::{CircuitBreakerConfig, RetryConfig};
|
|
use tower::ServiceExt;
|
|
|
|
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
|
|
|
|
#[cfg(test)]
|
|
mod round_robin_tests {
|
|
use super::*;
|
|
|
|
/// Test that round robin distributes requests evenly across workers
|
|
#[tokio::test]
|
|
async fn test_round_robin_distribution() {
|
|
let config = TestRouterConfig::round_robin(3100);
|
|
let ctx =
|
|
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19001, 3))
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
let num_requests = 30;
|
|
let mut success_count = 0;
|
|
|
|
for i in 0..num_requests {
|
|
let payload = json!({
|
|
"text": format!("Request {}", i),
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
if resp.status() == StatusCode::OK {
|
|
success_count += 1;
|
|
}
|
|
}
|
|
|
|
// All requests should succeed with round robin across 3 healthy workers
|
|
assert_eq!(
|
|
success_count, num_requests,
|
|
"All requests should succeed with round robin"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test round robin with one worker failing
|
|
#[tokio::test]
|
|
async fn test_round_robin_with_failing_worker() {
|
|
let config = TestRouterConfig::round_robin_with_retry(
|
|
3101,
|
|
RetryConfig {
|
|
max_retries: 2,
|
|
initial_backoff_ms: 10,
|
|
max_backoff_ms: 50,
|
|
..Default::default()
|
|
},
|
|
);
|
|
let ctx = AppTestContext::new_with_config(
|
|
config,
|
|
vec![
|
|
TestWorkerConfig::flaky(19004, 1.0), // Always fail
|
|
TestWorkerConfig::healthy(19005), // Always succeed
|
|
],
|
|
)
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// With retries enabled, requests should eventually succeed
|
|
// by being retried on the healthy worker
|
|
let payload = json!({
|
|
"text": "Test with retries",
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
assert_eq!(
|
|
resp.status(),
|
|
StatusCode::OK,
|
|
"Request should succeed via retry on healthy worker"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod random_tests {
|
|
use super::*;
|
|
|
|
/// Test that random policy distributes requests across workers
|
|
#[tokio::test]
|
|
async fn test_random_distribution() {
|
|
let config = TestRouterConfig::random(3102);
|
|
let ctx =
|
|
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19010, 2))
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
let num_requests = 20;
|
|
let mut success_count = 0;
|
|
|
|
for i in 0..num_requests {
|
|
let payload = json!({
|
|
"text": format!("Random request {}", i),
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
if resp.status() == StatusCode::OK {
|
|
success_count += 1;
|
|
}
|
|
}
|
|
|
|
// All requests should succeed with random policy
|
|
assert_eq!(
|
|
success_count, num_requests,
|
|
"All requests should succeed with random policy"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod cache_aware_tests {
|
|
use super::*;
|
|
|
|
/// Test cache-aware routing uses consistent hashing
|
|
#[tokio::test]
|
|
async fn test_cache_aware_consistent_routing() {
|
|
let config = TestRouterConfig::cache_aware(3103);
|
|
let ctx =
|
|
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19020, 2))
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Same prompt should route to same worker (consistent hashing)
|
|
let same_prompt = "Hello, cache-aware routing test!";
|
|
let mut worker_ids: Vec<Option<String>> = Vec::new();
|
|
|
|
for _ in 0..5 {
|
|
let payload = json!({
|
|
"text": same_prompt,
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// Extract worker ID from response header if available
|
|
let worker_id = resp
|
|
.headers()
|
|
.get("x-worker-id")
|
|
.map(|v| v.to_str().unwrap().to_string());
|
|
worker_ids.push(worker_id);
|
|
}
|
|
|
|
// All requests should succeed
|
|
assert_eq!(worker_ids.len(), 5);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test cache-aware routing with different prompts
|
|
#[tokio::test]
|
|
async fn test_cache_aware_different_prompts() {
|
|
let config = TestRouterConfig::cache_aware(3104);
|
|
let ctx =
|
|
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19022, 2))
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
let prompts = vec![
|
|
"First unique prompt",
|
|
"Second unique prompt",
|
|
"Third unique prompt",
|
|
"Fourth unique prompt",
|
|
];
|
|
|
|
let mut success_count = 0;
|
|
for prompt in prompts {
|
|
let payload = json!({
|
|
"text": prompt,
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
if resp.status() == StatusCode::OK {
|
|
success_count += 1;
|
|
}
|
|
}
|
|
|
|
// All requests should succeed
|
|
assert_eq!(success_count, 4);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod worker_health_tests {
|
|
use super::*;
|
|
|
|
/// Test that failing workers are avoided via circuit breaker/retry
|
|
#[tokio::test]
|
|
async fn test_skip_failing_workers() {
|
|
let config = TestRouterConfig::round_robin_with_reliability(
|
|
3105,
|
|
RetryConfig {
|
|
max_retries: 3,
|
|
initial_backoff_ms: 10,
|
|
max_backoff_ms: 50,
|
|
..Default::default()
|
|
},
|
|
CircuitBreakerConfig {
|
|
failure_threshold: 2,
|
|
success_threshold: 1,
|
|
timeout_duration_secs: 2,
|
|
window_duration_secs: 10,
|
|
},
|
|
);
|
|
let ctx = AppTestContext::new_with_config(
|
|
config,
|
|
vec![
|
|
TestWorkerConfig::flaky(19030, 1.0), // Always fails
|
|
TestWorkerConfig::healthy(19031), // Always succeeds
|
|
],
|
|
)
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Requests should succeed by retrying on healthy worker
|
|
// or by circuit breaker opening on failing worker
|
|
for i in 0..10 {
|
|
let payload = json!({
|
|
"text": format!("Request {}", i),
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
}
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test behavior when all workers are unhealthy
|
|
#[tokio::test]
|
|
async fn test_all_workers_unhealthy() {
|
|
let ctx = AppTestContext::new(vec![]).await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let payload = json!({
|
|
"text": "Test with no workers",
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
// Should return service unavailable when no workers available
|
|
assert!(
|
|
resp.status() == StatusCode::SERVICE_UNAVAILABLE
|
|
|| resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
|
|
"Expected 503 or 500 when no workers available, got {}",
|
|
resp.status()
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod worker_response_delay_tests {
|
|
use super::*;
|
|
|
|
/// Test that response delays are handled correctly
|
|
#[tokio::test]
|
|
async fn test_worker_with_delay() {
|
|
let config = TestRouterConfig::random(3106);
|
|
let ctx = AppTestContext::new_with_config(
|
|
config,
|
|
vec![TestWorkerConfig::slow(19040, 100)], // 100ms delay
|
|
)
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
let start = std::time::Instant::now();
|
|
let payload = json!({
|
|
"text": "Test with delay",
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.oneshot(req).await.unwrap();
|
|
let elapsed = start.elapsed();
|
|
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
// Response should take at least 100ms due to configured delay
|
|
assert!(
|
|
elapsed.as_millis() >= 100,
|
|
"Response should be delayed by at least 100ms, got {}ms",
|
|
elapsed.as_millis()
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
}
|