feat: 引入 aether-runtime/cache/data/http/testkit 基础 crate,完善并发门控与审计系统

新增 crate:
- aether-runtime: 服务运行时基础设施(并发门控、分布式并发、指标、队列、优雅关闭、tracing)
- aether-cache: 通用 TTL 缓存与命名空间抽象
- aether-data: 数据访问层(PostgreSQL/Redis 后端、repository 模式)
- aether-http: HTTP 客户端封装(重试、配置)
- aether-testkit: 集成测试工具集(gateway/executor/hub/proxy fixture、等待、负载测试)

gateway 扩展:
- 引入 audit 模块(shadow 执行审计、决策链路追踪、请求审计 bundle)
- 引入 cache 模块(AuthContext 缓存、direct-plan bypass 缓存)
- 引入 data 模块(auth/candidates/config/usage/video_tasks 数据访问)
- 集成 ConcurrencyGate/DistributedConcurrencyGate 请求门控
- 新增本地 auth 拒绝、过载响应构建器
- 补充 control/auth_cache/video/concurrency 集成测试

aether-proxy 扩展:
- AppState 集成 stream_gate / distributed_stream_gate 并发门控
- 新增 ProxyAdmissionError 及准入拒绝流程
- stream_handler 补充门控饱和/不可用场景测试
- 配置与注册客户端逻辑完善

aether-hub 扩展:
- main.rs 引入运行时初始化、指标端点、健康检查
- local_relay 重构为 lib.rs 暴露公共接口
This commit is contained in:
fawney19
2026-03-24 15:12:56 +08:00
parent eaf8475f9e
commit b5a0070023
157 changed files with 22097 additions and 448 deletions

View File

@@ -0,0 +1,27 @@
[package]
name = "aether-testkit"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Shared integration test helpers for Aether Rust services"
[dependencies]
async-stream.workspace = true
aether-data.workspace = true
aether-contracts.workspace = true
aether-executor.workspace = true
aether-gateway.workspace = true
aether-hub.workspace = true
aether-http.workspace = true
aether-runtime.workspace = true
axum.workspace = true
bytes.workspace = true
http.workspace = true
futures-util.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
tokio.workspace = true
tokio-tungstenite = { version = "0.28", features = ["rustls-tls-webpki-roots"] }

View File

@@ -0,0 +1,809 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::path::PathBuf;
use std::time::{Duration, Instant};
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_hub::protocol;
use aether_testkit::{
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for, run_http_load_probe,
ExecutorHarness, ExecutorHarnessConfig, GatewayHarness, GatewayHarnessConfig,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, HubHarness,
HubHarnessConfig, SpawnedServer,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::any;
use axum::{extract::Request, Json, Router};
use futures_util::{SinkExt, StreamExt};
use reqwest::Method;
use serde::Serialize;
use serde_json::json;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message;
#[derive(Debug, Clone)]
struct CapacityCurveBaselineConfig {
points: Vec<usize>,
requests_per_point_multiplier: usize,
sync_delay: Duration,
stream_chunk_delay: Duration,
hub_hold: Duration,
timeout: Duration,
saturation_latency_multiplier: u64,
output_path: Option<PathBuf>,
}
impl Default for CapacityCurveBaselineConfig {
fn default() -> Self {
Self {
points: vec![8, 16, 32, 64, 128, 256],
requests_per_point_multiplier: 8,
sync_delay: Duration::from_millis(75),
stream_chunk_delay: Duration::from_millis(25),
hub_hold: Duration::from_millis(75),
timeout: Duration::from_secs(10),
saturation_latency_multiplier: 4,
output_path: None,
}
}
}
#[derive(Debug, Serialize)]
struct CapacityCurveBaselineReport {
suite: &'static str,
gateway_sync: CapacityCurveScenarioReport,
gateway_stream: CapacityCurveScenarioReport,
executor_sync: CapacityCurveScenarioReport,
executor_stream: CapacityCurveScenarioReport,
hub_tunnel_stream: CapacityCurveScenarioReport,
}
#[derive(Debug, Serialize)]
struct CapacityCurveScenarioReport {
name: String,
gate: String,
latency_budget_ms: u64,
points: Vec<CapacityCurvePointResult>,
saturation_point: Option<CapacityCurveSaturationPoint>,
}
#[derive(Debug, Serialize)]
struct CapacityCurvePointResult {
limit: usize,
concurrency: usize,
total_requests: usize,
duration_ms: u64,
successful_requests: usize,
rejected_requests: usize,
failed_requests: usize,
throughput_rps: u64,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
metrics: GateMetricSnapshot,
}
#[derive(Debug, Serialize)]
struct CapacityCurveSaturationPoint {
limit: usize,
concurrency: usize,
reason: String,
p95_ms: u64,
rejected_requests: usize,
failed_requests: usize,
high_watermark: u64,
}
#[derive(Debug, Serialize)]
struct GateMetricSnapshot {
in_flight: u64,
available_permits: u64,
high_watermark: u64,
rejected_total: u64,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("capacity-curve-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &CapacityCurveBaselineConfig,
) -> Result<CapacityCurveBaselineReport, Box<dyn std::error::Error>> {
let upstream = SpawnedServer::start(build_delayed_upstream(
config.sync_delay,
config.stream_chunk_delay,
))
.await?;
Ok(CapacityCurveBaselineReport {
suite: "capacity_curve_baseline",
gateway_sync: run_gateway_curve(
"gateway_proxy_sync",
"gateway_requests",
false,
upstream.base_url(),
config,
)
.await?,
gateway_stream: run_gateway_curve(
"gateway_proxy_stream",
"gateway_requests",
true,
upstream.base_url(),
config,
)
.await?,
executor_sync: run_executor_curve(
"executor_sync",
"executor_requests",
false,
upstream.base_url(),
config,
)
.await?,
executor_stream: run_executor_curve(
"executor_stream",
"executor_requests",
true,
upstream.base_url(),
config,
)
.await?,
hub_tunnel_stream: run_hub_curve("hub_tunnel_stream", "hub_requests", config).await?,
})
}
async fn run_gateway_curve(
scenario_name: &str,
gate_name: &str,
stream: bool,
upstream_base_url: &str,
config: &CapacityCurveBaselineConfig,
) -> Result<CapacityCurveScenarioReport, Box<dyn std::error::Error>> {
let latency_budget_ms = scenario_latency_budget_ms(
if stream {
config.stream_chunk_delay.saturating_mul(3u32)
} else {
config.sync_delay
},
config.saturation_latency_multiplier,
);
let mut points = Vec::new();
for limit in &config.points {
let gateway = GatewayHarness::start(GatewayHarnessConfig {
upstream_base_url: upstream_base_url.to_string(),
control_base_url: None,
executor_base_url: None,
max_in_flight_requests: Some(*limit),
distributed_request_gate: None,
})
.await?;
let total_requests = total_requests_for_limit(*limit, config.requests_per_point_multiplier);
let probe = chat_probe_config(
format!("{}/v1/chat/completions", gateway.base_url()),
stream,
total_requests,
*limit,
config.timeout,
);
let started_at = Instant::now();
let result = run_http_load_probe(&probe)
.await
.map_err(std::io::Error::other)?;
let duration_ms = started_at.elapsed().as_millis() as u64;
let metrics = capture_gate_metrics(
&format!("{}/_gateway/metrics", gateway.base_url()),
gate_name,
)
.await?;
points.push(capacity_point(
*limit,
total_requests,
duration_ms,
result,
metrics,
));
}
Ok(CapacityCurveScenarioReport {
name: scenario_name.to_string(),
gate: gate_name.to_string(),
latency_budget_ms,
saturation_point: detect_saturation_point(&points, latency_budget_ms),
points,
})
}
async fn run_executor_curve(
scenario_name: &str,
gate_name: &str,
stream: bool,
upstream_base_url: &str,
config: &CapacityCurveBaselineConfig,
) -> Result<CapacityCurveScenarioReport, Box<dyn std::error::Error>> {
let latency_budget_ms = scenario_latency_budget_ms(
if stream {
config.stream_chunk_delay.saturating_mul(3u32)
} else {
config.sync_delay
},
config.saturation_latency_multiplier,
);
let mut points = Vec::new();
for limit in &config.points {
let executor = ExecutorHarness::start(ExecutorHarnessConfig {
max_in_flight_requests: Some(*limit),
distributed_request_gate: None,
})
.await?;
let total_requests = total_requests_for_limit(*limit, config.requests_per_point_multiplier);
let probe = execution_probe_config(
format!(
"{}/v1/execute/{}",
executor.base_url(),
if stream { "stream" } else { "sync" }
),
execution_plan(format!("{upstream_base_url}/v1/chat/completions"), stream),
total_requests,
*limit,
config.timeout,
);
let started_at = Instant::now();
let result = run_http_load_probe(&probe)
.await
.map_err(std::io::Error::other)?;
let duration_ms = started_at.elapsed().as_millis() as u64;
let metrics =
capture_gate_metrics(&format!("{}/metrics", executor.base_url()), gate_name).await?;
points.push(capacity_point(
*limit,
total_requests,
duration_ms,
result,
metrics,
));
}
Ok(CapacityCurveScenarioReport {
name: scenario_name.to_string(),
gate: gate_name.to_string(),
latency_budget_ms,
saturation_point: detect_saturation_point(&points, latency_budget_ms),
points,
})
}
async fn run_hub_curve(
scenario_name: &str,
gate_name: &str,
config: &CapacityCurveBaselineConfig,
) -> Result<CapacityCurveScenarioReport, Box<dyn std::error::Error>> {
let latency_budget_ms =
scenario_latency_budget_ms(config.hub_hold, config.saturation_latency_multiplier);
let mut points = Vec::new();
for limit in &config.points {
let relay_concurrency = (*limit).saturating_sub(1).max(1);
let hub = HubHarness::start(HubHarnessConfig {
max_streams: (*limit).max(128),
ping_interval: Duration::from_secs(15),
idle_timeout: Duration::ZERO,
outbound_queue_capacity: 128,
max_in_flight_requests: Some(*limit),
distributed_request_gate: None,
})
.await?;
let peer = connect_protocol_peer(hub.base_url(), config.hub_hold).await?;
let total_requests =
total_requests_for_limit(relay_concurrency, config.requests_per_point_multiplier);
let probe = HttpLoadProbeConfig {
url: format!("{}/local/relay/node-baseline", hub.base_url()),
method: Method::POST,
headers: BTreeMap::from([(
"content-type".to_string(),
"application/octet-stream".to_string(),
)]),
body: Some(relay_envelope()),
total_requests,
concurrency: relay_concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
};
let started_at = Instant::now();
let result = run_http_load_probe(&probe)
.await
.map_err(std::io::Error::other)?;
let duration_ms = started_at.elapsed().as_millis() as u64;
let metrics =
capture_gate_metrics(&format!("{}/metrics", hub.base_url()), gate_name).await?;
points.push(capacity_point(
*limit,
total_requests,
duration_ms,
result,
metrics,
));
drop(peer);
}
Ok(CapacityCurveScenarioReport {
name: scenario_name.to_string(),
gate: gate_name.to_string(),
latency_budget_ms,
saturation_point: detect_saturation_point(&points, latency_budget_ms),
points,
})
}
fn capacity_point(
limit: usize,
total_requests: usize,
duration_ms: u64,
result: HttpLoadProbeResult,
metrics: GateMetricSnapshot,
) -> CapacityCurvePointResult {
let rejected_requests = result.status_counts.get(&503).copied().unwrap_or_default();
let successful_requests = result
.status_counts
.iter()
.filter(|(status, _)| **status >= 200 && **status < 300)
.map(|(_, count)| *count)
.sum::<usize>();
let throughput_rps = if duration_ms == 0 {
successful_requests as u64
} else {
((successful_requests as u64) * 1_000) / duration_ms.max(1)
};
CapacityCurvePointResult {
limit,
concurrency: result.concurrency,
total_requests,
duration_ms,
successful_requests,
rejected_requests,
failed_requests: result.failed_requests,
throughput_rps,
p50_ms: result.p50_ms,
p95_ms: result.p95_ms,
max_ms: result.max_ms,
mean_ms: result.mean_ms,
metrics,
}
}
fn detect_saturation_point(
points: &[CapacityCurvePointResult],
latency_budget_ms: u64,
) -> Option<CapacityCurveSaturationPoint> {
points.iter().find_map(|point| {
let reason = if point.failed_requests > 0 {
Some("failures_observed")
} else if point.rejected_requests > 0 {
Some("admission_rejections_observed")
} else if point.p95_ms > latency_budget_ms {
Some("latency_budget_exceeded")
} else {
None
}?;
Some(CapacityCurveSaturationPoint {
limit: point.limit,
concurrency: point.concurrency,
reason: reason.to_string(),
p95_ms: point.p95_ms,
rejected_requests: point.rejected_requests,
failed_requests: point.failed_requests,
high_watermark: point.metrics.high_watermark,
})
})
}
async fn capture_gate_metrics(
metrics_url: &str,
gate_name: &str,
) -> Result<GateMetricSnapshot, Box<dyn std::error::Error>> {
let samples = fetch_prometheus_samples(metrics_url)
.await
.map_err(std::io::Error::other)?;
Ok(GateMetricSnapshot {
in_flight: find_metric_value_u64(&samples, "concurrency_in_flight", &[("gate", gate_name)])
.unwrap_or_default(),
available_permits: find_metric_value_u64(
&samples,
"concurrency_available_permits",
&[("gate", gate_name)],
)
.unwrap_or_default(),
high_watermark: find_metric_value_u64(
&samples,
"concurrency_high_watermark",
&[("gate", gate_name)],
)
.unwrap_or_default(),
rejected_total: find_metric_value_u64(
&samples,
"concurrency_rejected_total",
&[("gate", gate_name)],
)
.unwrap_or_default(),
})
}
fn scenario_latency_budget_ms(base: Duration, multiplier: u64) -> u64 {
(base.as_millis() as u64).saturating_mul(multiplier.max(1))
}
fn total_requests_for_limit(limit: usize, multiplier: usize) -> usize {
limit.saturating_mul(multiplier.max(1))
}
fn execution_probe_config(
url: String,
plan: ExecutionPlan,
total_requests: usize,
concurrency: usize,
timeout: Duration,
) -> HttpLoadProbeConfig {
HttpLoadProbeConfig {
url,
method: Method::POST,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(
serde_json::to_vec(&plan).expect("execution plan should serialize for capacity curve"),
),
total_requests,
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
}
}
fn chat_probe_config(
url: String,
stream: bool,
total_requests: usize,
concurrency: usize,
timeout: Duration,
) -> HttpLoadProbeConfig {
HttpLoadProbeConfig {
url,
method: Method::POST,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(
serde_json::to_vec(&json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": stream,
}))
.expect("chat body should serialize"),
),
total_requests,
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
}
}
fn execution_plan(url: String, stream: bool) -> ExecutionPlan {
ExecutionPlan {
request_id: if stream {
"capacity-curve-stream-request".to_string()
} else {
"capacity-curve-sync-request".to_string()
},
candidate_id: Some(if stream {
"capacity-curve-stream-candidate".to_string()
} else {
"capacity-curve-sync-candidate".to_string()
}),
provider_name: Some("openai".to_string()),
provider_id: "provider-capacity".to_string(),
endpoint_id: "endpoint-capacity".to_string(),
key_id: "key-capacity".to_string(),
method: "POST".to_string(),
url,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": stream,
})),
stream,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(2_000),
read_ms: Some(10_000),
first_byte_ms: Some(5_000),
total_ms: Some(10_000),
..ExecutionTimeouts::default()
}),
}
}
fn build_delayed_upstream(sync_delay: Duration, stream_chunk_delay: Duration) -> Router {
Router::new().route(
"/v1/chat/completions",
any(move |request: Request| {
let sync_delay = sync_delay;
let stream_chunk_delay = stream_chunk_delay;
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX)
.await
.expect("capacity upstream body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).unwrap_or_else(|_| json!({}));
let stream = payload
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false);
if stream {
let body = async_stream::stream! {
tokio::time::sleep(stream_chunk_delay).await;
yield Ok::<_, Infallible>(Bytes::from_static(
b"data: {\"id\":\"chunk-1\",\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n",
));
tokio::time::sleep(stream_chunk_delay).await;
yield Ok::<_, Infallible>(Bytes::from_static(
b"data: {\"id\":\"chunk-2\",\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n",
));
tokio::time::sleep(stream_chunk_delay).await;
yield Ok::<_, Infallible>(Bytes::from_static(b"data: [DONE]\n\n"));
};
Response::builder()
.status(StatusCode::OK)
.header(http::header::CONTENT_TYPE, "text/event-stream")
.body(Body::from_stream(body))
.expect("capacity upstream stream response should build")
} else {
tokio::time::sleep(sync_delay).await;
Json(json!({
"id": "chatcmpl-capacity",
"object": "chat.completion",
"model": payload.get("model").and_then(|value| value.as_str()).unwrap_or("gpt-5"),
"choices": [{"message": {"role": "assistant", "content": "hello"}}]
}))
.into_response()
}
}
}),
)
}
fn relay_envelope() -> Vec<u8> {
let meta = protocol::RequestMeta {
method: "POST".to_string(),
url: "https://capacity.example/v1/chat/completions".to_string(),
headers: std::collections::HashMap::from([(
"content-type".to_string(),
"application/json".to_string(),
)]),
timeout: 30,
};
let meta_json = serde_json::to_vec(&meta).expect("hub relay metadata should serialize");
let body = br#"{"model":"gpt-5","messages":[{"role":"user","content":"hello"}]}"#;
let mut envelope = Vec::with_capacity(4 + meta_json.len() + body.len());
envelope.extend_from_slice(&(meta_json.len() as u32).to_be_bytes());
envelope.extend_from_slice(&meta_json);
envelope.extend_from_slice(body);
envelope
}
async fn connect_protocol_peer(
hub_base_url: &str,
hold: Duration,
) -> Result<tokio::task::JoinHandle<()>, Box<dyn std::error::Error>> {
let ws_url = format!("{}/proxy", hub_base_url.replace("http://", "ws://"));
let request = ws_url.into_client_request()?;
let mut request = request;
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static("node-baseline"));
request.headers_mut().insert(
"x-node-name",
http::HeaderValue::from_static("proxy-baseline"),
);
request.headers_mut().insert(
"x-tunnel-max-streams",
http::HeaderValue::from_static("512"),
);
let (socket, _response) = tokio_tungstenite::connect_async(request).await?;
let (mut sink, mut stream) = socket.split();
Ok(tokio::spawn(async move {
while let Some(message) = stream.next().await {
let Ok(message) = message else {
break;
};
match message {
Message::Binary(data) => {
if handle_binary_frame(&mut sink, data.to_vec(), hold)
.await
.is_err()
{
break;
}
}
Message::Ping(payload) => {
if sink.send(Message::Pong(payload)).await.is_err() {
break;
}
}
Message::Close(_) => break,
_ => {}
}
}
let _ = sink.close().await;
}))
}
async fn handle_binary_frame<S>(
sink: &mut S,
data: Vec<u8>,
hold: Duration,
) -> Result<(), tokio_tungstenite::tungstenite::Error>
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin,
{
let Some(header) = protocol::FrameHeader::parse(&data) else {
return Ok(());
};
match header.msg_type {
protocol::PING => {
let payload = protocol::frame_payload_by_header(&data, &header).unwrap_or(&[]);
sink.send(Message::Binary(protocol::encode_pong(payload).into()))
.await?;
}
protocol::REQUEST_HEADERS => {
let payload = protocol::decode_payload(&data, &header).unwrap_or_default();
let _ = serde_json::from_slice::<protocol::RequestMeta>(&payload);
}
protocol::REQUEST_BODY => {
if header.flags & protocol::FLAG_END_STREAM != 0 {
tokio::time::sleep(hold).await;
let response_meta = protocol::ResponseMeta {
status: 200,
headers: vec![(
"content-type".to_string(),
"text/plain; charset=utf-8".to_string(),
)],
};
let response_meta_json =
serde_json::to_vec(&response_meta).expect("response metadata should serialize");
sink.send(Message::Binary(
protocol::encode_frame(
header.stream_id,
protocol::RESPONSE_HEADERS,
0,
&response_meta_json,
)
.into(),
))
.await?;
for chunk in [
b"capacity-".as_slice(),
b"tunnel-".as_slice(),
b"stream".as_slice(),
] {
sink.send(Message::Binary(
protocol::encode_frame(header.stream_id, protocol::RESPONSE_BODY, 0, chunk)
.into(),
))
.await?;
}
sink.send(Message::Binary(
protocol::encode_frame(header.stream_id, protocol::STREAM_END, 0, &[]).into(),
))
.await?;
}
}
_ => {}
}
Ok(())
}
fn parse_args(
args: Vec<String>,
) -> Result<CapacityCurveBaselineConfig, Box<dyn std::error::Error>> {
let mut config = CapacityCurveBaselineConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--points" => {
config.points = next_value(&mut iter, "--points")?
.split(',')
.filter(|value| !value.trim().is_empty())
.map(|value| value.trim().parse::<usize>())
.collect::<Result<Vec<_>, _>>()?;
}
"--requests-per-point-multiplier" => {
config.requests_per_point_multiplier =
next_value(&mut iter, "--requests-per-point-multiplier")?.parse()?
}
"--sync-delay-ms" => {
config.sync_delay =
Duration::from_millis(next_value(&mut iter, "--sync-delay-ms")?.parse()?)
}
"--stream-chunk-delay-ms" => {
config.stream_chunk_delay = Duration::from_millis(
next_value(&mut iter, "--stream-chunk-delay-ms")?.parse()?,
)
}
"--hub-hold-ms" => {
config.hub_hold =
Duration::from_millis(next_value(&mut iter, "--hub-hold-ms")?.parse()?)
}
"--timeout-ms" => {
config.timeout =
Duration::from_millis(next_value(&mut iter, "--timeout-ms")?.parse()?)
}
"--saturation-latency-multiplier" => {
config.saturation_latency_multiplier =
next_value(&mut iter, "--saturation-latency-multiplier")?.parse()?
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
if config.points.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"capacity curve requires at least one point",
)
.into());
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin capacity_curve_baseline -- [--points 8,16,32,64,128,256] [--requests-per-point-multiplier 8] [--sync-delay-ms 75] [--stream-chunk-delay-ms 25] [--hub-hold-ms 75] [--timeout-ms 10000] [--saturation-latency-multiplier 4] [--output /tmp/capacity_curve_baseline.json]"
);
}

View File

@@ -0,0 +1,785 @@
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_data::postgres::{
DatabaseRecordId, PostgresLeaseClaimOptions, PostgresLeaseClaimSpec, PostgresLeaseRunnerConfig,
PostgresPoolConfig,
};
use aether_data::redis::{
RedisClientConfig, RedisConsumerGroup, RedisConsumerName, RedisLockLease, RedisLockRunner,
RedisLockRunnerConfig, RedisStreamName, RedisStreamReclaimConfig, RedisStreamRunner,
RedisStreamRunnerConfig,
};
use aether_data::{PostgresBackend, RedisBackend};
use aether_testkit::{init_test_runtime_for, ManagedPostgresServer, ManagedRedisServer};
use futures_util::stream::{self, StreamExt};
use serde::Serialize;
#[derive(Debug, Clone)]
struct DependencyPressureBaselineConfig {
redis_lock_total: usize,
redis_lock_concurrency: usize,
redis_stream_total: usize,
redis_stream_concurrency: usize,
redis_reclaim_total: usize,
redis_reclaim_min_idle: Duration,
postgres_rows: usize,
postgres_lease_cycles: usize,
postgres_lease_concurrency: usize,
postgres_lease_batch_size: usize,
postgres_lease_ms: u64,
timeout: Duration,
output_path: Option<PathBuf>,
redis_url: Option<String>,
postgres_url: Option<String>,
}
impl Default for DependencyPressureBaselineConfig {
fn default() -> Self {
Self {
redis_lock_total: 1_000,
redis_lock_concurrency: 20,
redis_stream_total: 2_000,
redis_stream_concurrency: 20,
redis_reclaim_total: 256,
redis_reclaim_min_idle: Duration::from_millis(100),
postgres_rows: 512,
postgres_lease_cycles: 128,
postgres_lease_concurrency: 16,
postgres_lease_batch_size: 16,
postgres_lease_ms: 250,
timeout: Duration::from_secs(10),
output_path: None,
redis_url: None,
postgres_url: None,
}
}
}
#[derive(Debug, Clone, Serialize)]
struct OperationSummary {
total_calls: usize,
total_items: usize,
failed_calls: usize,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
}
#[derive(Debug, Clone, Serialize)]
struct RedisLockPressureReport {
acquire: OperationSummary,
renew: OperationSummary,
release: OperationSummary,
}
#[derive(Debug, Clone, Serialize)]
struct RedisStreamPressureReport {
append: OperationSummary,
read_group: OperationSummary,
reclaim: OperationSummary,
ack: OperationSummary,
}
#[derive(Debug, Clone, Serialize)]
struct PostgresLeasePressureReport {
claim: OperationSummary,
renew: OperationSummary,
release: OperationSummary,
}
#[derive(Debug, Clone, Serialize)]
struct DependencyPressureBaselineReport {
suite: &'static str,
redis_url: String,
postgres_url: String,
redis_lock: RedisLockPressureReport,
redis_stream: RedisStreamPressureReport,
postgres_lease: PostgresLeasePressureReport,
}
#[derive(Default)]
struct SummaryCollector {
latencies_ms: tokio::sync::Mutex<Vec<u64>>,
total_items: std::sync::atomic::AtomicUsize,
failed_calls: std::sync::atomic::AtomicUsize,
total_calls: std::sync::atomic::AtomicUsize,
}
impl SummaryCollector {
async fn record(&self, elapsed: Duration, items: usize, failed: bool) {
self.latencies_ms
.lock()
.await
.push(elapsed.as_millis() as u64);
self.total_calls
.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
self.total_items
.fetch_add(items, std::sync::atomic::Ordering::AcqRel);
if failed {
self.failed_calls
.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
}
}
async fn summarize(&self) -> OperationSummary {
let mut latencies = self.latencies_ms.lock().await.clone();
latencies.sort_unstable();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
OperationSummary {
total_calls: self.total_calls.load(std::sync::atomic::Ordering::Acquire),
total_items: self.total_items.load(std::sync::atomic::Ordering::Acquire),
failed_calls: self.failed_calls.load(std::sync::atomic::Ordering::Acquire),
p50_ms,
p95_ms,
max_ms,
mean_ms,
}
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("dependency-pressure-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &DependencyPressureBaselineConfig,
) -> Result<DependencyPressureBaselineReport, Box<dyn std::error::Error>> {
let managed_redis = if config.redis_url.is_none() {
Some(ManagedRedisServer::start().await?)
} else {
None
};
let managed_postgres = if config.postgres_url.is_none() {
Some(ManagedPostgresServer::start().await?)
} else {
None
};
let redis_url = config
.redis_url
.clone()
.or_else(|| {
managed_redis
.as_ref()
.map(|server| server.redis_url().to_string())
})
.expect("redis url should resolve");
let postgres_url = config
.postgres_url
.clone()
.or_else(|| {
managed_postgres
.as_ref()
.map(|server| server.database_url().to_string())
})
.expect("postgres url should resolve");
let redis_backend = RedisBackend::from_config(RedisClientConfig {
url: redis_url.clone(),
key_prefix: Some(format!("aether-dependency-pressure-{}", std::process::id())),
})?;
let postgres_backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: postgres_url.clone(),
min_connections: 1,
max_connections: (config.postgres_lease_concurrency as u32).saturating_add(8),
acquire_timeout_ms: config.timeout.as_millis() as u64,
idle_timeout_ms: 60_000,
max_lifetime_ms: 10 * 60_000,
statement_cache_capacity: 128,
require_ssl: false,
})?;
bootstrap_postgres_lease_table(postgres_backend.pool_clone(), config).await?;
let lock_runner = redis_backend.lock_runner(RedisLockRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
default_ttl_ms: 5_000,
})?;
let stream_runner = redis_backend.stream_runner(RedisStreamRunnerConfig {
command_timeout_ms: Some(config.timeout.as_millis() as u64),
read_block_ms: Some(10),
read_count: 64,
})?;
let lease_runner = postgres_backend.lease_runner(PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(config.timeout.as_millis() as u64),
lock_timeout_ms: Some(1_000),
})?;
let redis_lock = benchmark_redis_lock(&redis_backend, &lock_runner, config).await?;
let redis_stream = benchmark_redis_stream(&redis_backend, &stream_runner, config).await?;
let postgres_lease = benchmark_postgres_lease(&lease_runner, config).await?;
Ok(DependencyPressureBaselineReport {
suite: "dependency_pressure_baseline",
redis_url,
postgres_url,
redis_lock,
redis_stream,
postgres_lease,
})
}
async fn bootstrap_postgres_lease_table(
pool: sqlx::PgPool,
config: &DependencyPressureBaselineConfig,
) -> Result<(), Box<dyn std::error::Error>> {
sqlx::query("DROP TABLE IF EXISTS baseline_lease_jobs")
.execute(&pool)
.await?;
sqlx::query(
"CREATE TABLE baseline_lease_jobs (
id TEXT PRIMARY KEY,
status TEXT NOT NULL,
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
lease_owner TEXT,
lease_expires_at TIMESTAMPTZ
)",
)
.execute(&pool)
.await?;
let mut builder = sqlx::QueryBuilder::new("INSERT INTO baseline_lease_jobs (id, status) ");
builder.push_values(0..config.postgres_rows, |mut row, index| {
row.push_bind(format!("job-{index:05}")).push_bind("ready");
});
builder.build().execute(&pool).await?;
Ok(())
}
async fn benchmark_redis_lock(
backend: &RedisBackend,
runner: &RedisLockRunner,
config: &DependencyPressureBaselineConfig,
) -> Result<RedisLockPressureReport, Box<dyn std::error::Error>> {
let acquire = Arc::new(SummaryCollector::default());
let renew = Arc::new(SummaryCollector::default());
let release = Arc::new(SummaryCollector::default());
let next = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let keyspace = backend.keyspace();
stream::iter(0..config.redis_lock_concurrency)
.for_each_concurrent(config.redis_lock_concurrency, |_| {
let runner = runner.clone();
let acquire = acquire.clone();
let renew = renew.clone();
let release = release.clone();
let next = next.clone();
let keyspace = keyspace.clone();
async move {
loop {
let index = next.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
if index >= config.redis_lock_total {
break;
}
let owner = format!("lock-owner-{index}");
let key = keyspace.lock_key(&format!("dependency-pressure-{index}"));
let acquire_started = Instant::now();
match runner.try_acquire(&key, &owner, None).await {
Ok(Some(lease)) => {
acquire.record(acquire_started.elapsed(), 1, false).await;
record_redis_lock_follow_up(&runner, &lease, &renew, &release).await;
}
Ok(None) => {
acquire.record(acquire_started.elapsed(), 0, true).await;
}
Err(_) => {
acquire.record(acquire_started.elapsed(), 0, true).await;
}
}
}
}
})
.await;
Ok(RedisLockPressureReport {
acquire: acquire.summarize().await,
renew: renew.summarize().await,
release: release.summarize().await,
})
}
async fn record_redis_lock_follow_up(
runner: &RedisLockRunner,
lease: &RedisLockLease,
renew: &SummaryCollector,
release: &SummaryCollector,
) {
let renew_started = Instant::now();
let renew_ok = runner.renew(lease, None).await.unwrap_or(false);
renew
.record(renew_started.elapsed(), usize::from(renew_ok), !renew_ok)
.await;
let release_started = Instant::now();
let release_ok = runner.release(lease).await.unwrap_or(false);
release
.record(
release_started.elapsed(),
usize::from(release_ok),
!release_ok,
)
.await;
}
async fn benchmark_redis_stream(
backend: &RedisBackend,
runner: &RedisStreamRunner,
config: &DependencyPressureBaselineConfig,
) -> Result<RedisStreamPressureReport, Box<dyn std::error::Error>> {
let stream = backend.keyspace().stream_name("dependency-pressure");
let group = RedisConsumerGroup("dependency-group".to_string());
let consumer_a = RedisConsumerName("consumer-a".to_string());
let consumer_b = RedisConsumerName("consumer-b".to_string());
runner
.ensure_consumer_group(&stream, &group, "0-0")
.await
.map_err(std::io::Error::other)?;
let append = benchmark_redis_stream_append(runner, &stream, config).await?;
let (read_group, drained_ids) =
benchmark_redis_stream_read_group(runner, &stream, &group, &consumer_a, config).await?;
let ack = SummaryCollector::default();
benchmark_redis_stream_ack_into(&ack, runner, &stream, &group, &drained_ids).await?;
let (reclaim, reclaimed_ids) =
benchmark_redis_stream_reclaim(runner, &stream, &group, &consumer_a, &consumer_b, config)
.await?;
benchmark_redis_stream_ack_into(&ack, runner, &stream, &group, &reclaimed_ids).await?;
Ok(RedisStreamPressureReport {
append,
read_group,
reclaim,
ack: ack.summarize().await,
})
}
async fn benchmark_redis_stream_append(
runner: &RedisStreamRunner,
stream: &RedisStreamName,
config: &DependencyPressureBaselineConfig,
) -> Result<OperationSummary, Box<dyn std::error::Error>> {
let collector = Arc::new(SummaryCollector::default());
let next = Arc::new(std::sync::atomic::AtomicUsize::new(0));
stream::iter(0..config.redis_stream_concurrency)
.for_each_concurrent(config.redis_stream_concurrency, |_| {
let runner = runner.clone();
let stream = stream.clone();
let collector = collector.clone();
let next = next.clone();
async move {
loop {
let index = next.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
if index >= config.redis_stream_total {
break;
}
let started = Instant::now();
let result = runner
.append_json(
&stream,
"payload",
&serde_json::json!({
"job_id": index,
"kind": "dependency-pressure",
}),
)
.await;
collector
.record(
started.elapsed(),
usize::from(result.is_ok()),
result.is_err(),
)
.await;
}
}
})
.await;
Ok(collector.summarize().await)
}
async fn benchmark_redis_stream_read_group(
runner: &RedisStreamRunner,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
config: &DependencyPressureBaselineConfig,
) -> Result<(OperationSummary, Vec<String>), Box<dyn std::error::Error>> {
let collector = SummaryCollector::default();
let mut ids = Vec::with_capacity(config.redis_stream_total);
while ids.len() < config.redis_stream_total {
let started = Instant::now();
match runner.read_group(stream, group, consumer).await {
Ok(entries) => {
let item_count = entries.len();
ids.extend(entries.into_iter().map(|entry| entry.id));
collector.record(started.elapsed(), item_count, false).await;
}
Err(_) => {
collector.record(started.elapsed(), 0, true).await;
}
}
}
Ok((collector.summarize().await, ids))
}
async fn benchmark_redis_stream_reclaim(
runner: &RedisStreamRunner,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
consumer_a: &RedisConsumerName,
consumer_b: &RedisConsumerName,
config: &DependencyPressureBaselineConfig,
) -> Result<(OperationSummary, Vec<String>), Box<dyn std::error::Error>> {
for index in 0..config.redis_reclaim_total {
runner
.append_json(
stream,
"payload",
&serde_json::json!({
"job_id": format!("reclaim-{index}"),
"kind": "dependency-pressure",
}),
)
.await
.map_err(std::io::Error::other)?;
}
let mut pending_ids = Vec::with_capacity(config.redis_reclaim_total);
while pending_ids.len() < config.redis_reclaim_total {
let entries = runner
.read_group(stream, group, consumer_a)
.await
.map_err(std::io::Error::other)?;
pending_ids.extend(entries.into_iter().map(|entry| entry.id));
}
tokio::time::sleep(config.redis_reclaim_min_idle + Duration::from_millis(20)).await;
let collector = SummaryCollector::default();
let mut reclaimed_ids = Vec::with_capacity(config.redis_reclaim_total);
let mut next_start_id = "0-0".to_string();
while reclaimed_ids.len() < config.redis_reclaim_total {
let started = Instant::now();
match runner
.claim_stale(
stream,
group,
consumer_b,
&next_start_id,
RedisStreamReclaimConfig {
min_idle_ms: config.redis_reclaim_min_idle.as_millis() as u64,
count: 64,
},
)
.await
{
Ok(result) => {
next_start_id = result.next_start_id.clone();
let item_count = result.entries.len();
reclaimed_ids.extend(result.entries.into_iter().map(|entry| entry.id));
collector.record(started.elapsed(), item_count, false).await;
}
Err(_) => {
collector.record(started.elapsed(), 0, true).await;
}
}
}
Ok((collector.summarize().await, reclaimed_ids))
}
async fn benchmark_redis_stream_ack_into(
collector: &SummaryCollector,
runner: &RedisStreamRunner,
stream: &RedisStreamName,
group: &RedisConsumerGroup,
ids: &[String],
) -> Result<(), Box<dyn std::error::Error>> {
for chunk in ids.chunks(64) {
let started = Instant::now();
match runner.ack(stream, group, chunk).await {
Ok(acked) => {
collector.record(started.elapsed(), acked, false).await;
}
Err(_) => {
collector.record(started.elapsed(), 0, true).await;
}
}
}
Ok(())
}
async fn benchmark_postgres_lease(
runner: &aether_data::postgres::PostgresLeaseRunner,
config: &DependencyPressureBaselineConfig,
) -> Result<PostgresLeasePressureReport, Box<dyn std::error::Error>> {
let spec = PostgresLeaseClaimSpec {
table: "baseline_lease_jobs",
id_column: "id",
lease_owner_column: "lease_owner",
lease_expires_at_column: "lease_expires_at",
eligibility_predicate_sql: "status = 'ready'",
order_by_sql: "id ASC",
};
let claim_options = PostgresLeaseClaimOptions {
batch_size: config.postgres_lease_batch_size,
lease_ms: config.postgres_lease_ms,
};
let claim = Arc::new(SummaryCollector::default());
let renew = Arc::new(SummaryCollector::default());
let release = Arc::new(SummaryCollector::default());
let next_cycle = Arc::new(std::sync::atomic::AtomicUsize::new(0));
stream::iter(0..config.postgres_lease_concurrency)
.for_each_concurrent(config.postgres_lease_concurrency, |worker_index| {
let runner = runner.clone();
let spec = spec.clone();
let claim = claim.clone();
let renew = renew.clone();
let release = release.clone();
let next_cycle = next_cycle.clone();
async move {
let owner = format!("lease-owner-{worker_index}");
loop {
let cycle = next_cycle.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
if cycle >= config.postgres_lease_cycles {
break;
}
let claim_started = Instant::now();
match runner.claim_ids(&spec, claim_options, &owner).await {
Ok(ids) => {
let item_count = ids.len();
claim
.record(claim_started.elapsed(), item_count, false)
.await;
if ids.is_empty() {
tokio::time::sleep(Duration::from_millis(5)).await;
continue;
}
record_postgres_lease_follow_up(
&runner,
&spec,
&owner,
&ids,
config.postgres_lease_ms,
&renew,
&release,
)
.await;
}
Err(_) => {
claim.record(claim_started.elapsed(), 0, true).await;
}
}
}
}
})
.await;
Ok(PostgresLeasePressureReport {
claim: claim.summarize().await,
renew: renew.summarize().await,
release: release.summarize().await,
})
}
async fn record_postgres_lease_follow_up(
runner: &aether_data::postgres::PostgresLeaseRunner,
spec: &PostgresLeaseClaimSpec,
owner: &str,
ids: &[DatabaseRecordId],
lease_ms: u64,
renew: &SummaryCollector,
release: &SummaryCollector,
) {
let renew_started = Instant::now();
match runner.renew_ids(spec, ids, owner, lease_ms).await {
Ok(renewed) => {
let item_count = renewed.len();
renew
.record(renew_started.elapsed(), item_count, false)
.await;
}
Err(_) => {
renew.record(renew_started.elapsed(), 0, true).await;
}
}
let release_started = Instant::now();
match runner.release_ids(spec, ids, owner).await {
Ok(released) => {
let item_count = released.len();
release
.record(release_started.elapsed(), item_count, false)
.await;
}
Err(_) => {
release.record(release_started.elapsed(), 0, true).await;
}
}
}
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
}
let max_ms = *latencies.last().unwrap_or(&0);
let mean_ms = latencies.iter().sum::<u64>() / latencies.len() as u64;
let p50_ms = percentile(latencies, 50);
let p95_ms = percentile(latencies, 95);
(p50_ms, p95_ms, max_ms, mean_ms)
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {
if latencies.is_empty() {
return 0;
}
let last_index = latencies.len() - 1;
let rank = ((last_index as f64) * (percentile as f64 / 100.0)).round() as usize;
latencies[rank.min(last_index)]
}
fn parse_args(
args: Vec<String>,
) -> Result<DependencyPressureBaselineConfig, Box<dyn std::error::Error>> {
let mut config = DependencyPressureBaselineConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--redis-lock-total" => {
config.redis_lock_total = iter
.next()
.ok_or("missing value for --redis-lock-total")?
.parse()?;
}
"--redis-lock-concurrency" => {
config.redis_lock_concurrency = iter
.next()
.ok_or("missing value for --redis-lock-concurrency")?
.parse()?;
}
"--redis-stream-total" => {
config.redis_stream_total = iter
.next()
.ok_or("missing value for --redis-stream-total")?
.parse()?;
}
"--redis-stream-concurrency" => {
config.redis_stream_concurrency = iter
.next()
.ok_or("missing value for --redis-stream-concurrency")?
.parse()?;
}
"--redis-reclaim-total" => {
config.redis_reclaim_total = iter
.next()
.ok_or("missing value for --redis-reclaim-total")?
.parse()?;
}
"--redis-reclaim-min-idle-ms" => {
config.redis_reclaim_min_idle = Duration::from_millis(
iter.next()
.ok_or("missing value for --redis-reclaim-min-idle-ms")?
.parse()?,
);
}
"--postgres-rows" => {
config.postgres_rows = iter
.next()
.ok_or("missing value for --postgres-rows")?
.parse()?;
}
"--postgres-lease-cycles" => {
config.postgres_lease_cycles = iter
.next()
.ok_or("missing value for --postgres-lease-cycles")?
.parse()?;
}
"--postgres-lease-concurrency" => {
config.postgres_lease_concurrency = iter
.next()
.ok_or("missing value for --postgres-lease-concurrency")?
.parse()?;
}
"--postgres-lease-batch-size" => {
config.postgres_lease_batch_size = iter
.next()
.ok_or("missing value for --postgres-lease-batch-size")?
.parse()?;
}
"--postgres-lease-ms" => {
config.postgres_lease_ms = iter
.next()
.ok_or("missing value for --postgres-lease-ms")?
.parse()?;
}
"--timeout-ms" => {
config.timeout = Duration::from_millis(
iter.next()
.ok_or("missing value for --timeout-ms")?
.parse()?,
);
}
"--output" => {
config.output_path = Some(PathBuf::from(
iter.next().ok_or("missing value for --output")?,
));
}
"--redis-url" => {
config.redis_url = Some(iter.next().ok_or("missing value for --redis-url")?);
}
"--postgres-url" => {
config.postgres_url = Some(iter.next().ok_or("missing value for --postgres-url")?);
}
other => {
return Err(format!("unknown argument: {other}").into());
}
}
}
validate_config(&config)?;
Ok(config)
}
fn validate_config(
config: &DependencyPressureBaselineConfig,
) -> Result<(), Box<dyn std::error::Error>> {
if config.redis_lock_total == 0
|| config.redis_lock_concurrency == 0
|| config.redis_stream_total == 0
|| config.redis_stream_concurrency == 0
|| config.redis_reclaim_total == 0
|| config.postgres_rows == 0
|| config.postgres_lease_cycles == 0
|| config.postgres_lease_concurrency == 0
|| config.postgres_lease_batch_size == 0
|| config.postgres_lease_ms == 0
|| config.timeout.is_zero()
{
return Err("all dependency pressure baseline numeric settings must be positive".into());
}
Ok(())
}

View File

@@ -0,0 +1,696 @@
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_data::postgres::{
PostgresLeaseClaimOptions, PostgresLeaseClaimSpec, PostgresLeaseRunnerConfig,
PostgresPoolConfig, PostgresTransactionOptions,
};
use aether_data::redis::{RedisClientConfig, RedisLockRunnerConfig};
use aether_data::{PostgresBackend, RedisBackend};
use aether_testkit::{
init_test_runtime_for, reserve_local_port, HubHarness, HubHarnessConfig, ManagedPostgresServer,
ManagedRedisServer,
};
use futures_util::{FutureExt, StreamExt};
use serde::Serialize;
use tokio::sync::{oneshot, Mutex};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message;
#[derive(Debug, Clone)]
struct FailureRecoveryBaselineConfig {
redis_attempts: usize,
redis_concurrency: usize,
redis_restart_delay: Duration,
redis_downtime: Duration,
postgres_statement_timeout: Duration,
postgres_sleep: Duration,
hub_attempts: usize,
hub_concurrency: usize,
hub_hold: Duration,
hub_restart_delay: Duration,
hub_downtime: Duration,
timeout: Duration,
output_path: Option<PathBuf>,
redis_url: Option<String>,
postgres_url: Option<String>,
}
impl Default for FailureRecoveryBaselineConfig {
fn default() -> Self {
Self {
redis_attempts: 400,
redis_concurrency: 10,
redis_restart_delay: Duration::from_millis(200),
redis_downtime: Duration::from_millis(150),
postgres_statement_timeout: Duration::from_millis(50),
postgres_sleep: Duration::from_millis(200),
hub_attempts: 60,
hub_concurrency: 4,
hub_hold: Duration::from_millis(50),
hub_restart_delay: Duration::from_millis(200),
hub_downtime: Duration::from_millis(150),
timeout: Duration::from_secs(10),
output_path: None,
redis_url: None,
postgres_url: None,
}
}
}
#[derive(Debug, Clone, Copy)]
enum FaultPhase {
Pre,
During,
Post,
}
#[derive(Debug, Default, Clone, Serialize)]
struct PhaseCounts {
pre_successes: usize,
pre_failures: usize,
during_successes: usize,
during_failures: usize,
post_successes: usize,
post_failures: usize,
}
#[derive(Debug, Clone, Serialize)]
struct RecoverySummary {
total_attempts: usize,
successful_attempts: usize,
failed_attempts: usize,
recovered_after_restart_ms: Option<u64>,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
phase_counts: PhaseCounts,
}
#[derive(Debug, Clone, Serialize)]
struct PostgresSlowQueryRecoveryReport {
slow_query_timed_out: bool,
slow_query_latency_ms: u64,
recovery_claim_succeeded: bool,
recovery_claim_latency_ms: u64,
recovery_claimed_items: usize,
}
#[derive(Debug, Clone, Serialize)]
struct FailureRecoveryBaselineReport {
suite: &'static str,
redis_url: String,
postgres_url: String,
redis_restart: RecoverySummary,
postgres_slow_query: PostgresSlowQueryRecoveryReport,
hub_restart: RecoverySummary,
}
#[derive(Default)]
struct RecoveryCollector {
latencies_ms: Mutex<Vec<u64>>,
phase_counts: Mutex<PhaseCounts>,
successful_attempts: AtomicUsize,
failed_attempts: AtomicUsize,
}
impl RecoveryCollector {
async fn record(&self, phase: FaultPhase, success: bool, latency: Duration) {
self.latencies_ms
.lock()
.await
.push(latency.as_millis() as u64);
let mut counts = self.phase_counts.lock().await;
match (phase, success) {
(FaultPhase::Pre, true) => counts.pre_successes += 1,
(FaultPhase::Pre, false) => counts.pre_failures += 1,
(FaultPhase::During, true) => counts.during_successes += 1,
(FaultPhase::During, false) => counts.during_failures += 1,
(FaultPhase::Post, true) => counts.post_successes += 1,
(FaultPhase::Post, false) => counts.post_failures += 1,
}
if success {
self.successful_attempts.fetch_add(1, Ordering::AcqRel);
} else {
self.failed_attempts.fetch_add(1, Ordering::AcqRel);
}
}
async fn summarize(&self, recovered_after_restart_ms: Option<u64>) -> RecoverySummary {
let mut latencies = self.latencies_ms.lock().await.clone();
latencies.sort_unstable();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let phase_counts = self.phase_counts.lock().await.clone();
RecoverySummary {
total_attempts: self.successful_attempts.load(Ordering::Acquire)
+ self.failed_attempts.load(Ordering::Acquire),
successful_attempts: self.successful_attempts.load(Ordering::Acquire),
failed_attempts: self.failed_attempts.load(Ordering::Acquire),
recovered_after_restart_ms,
p50_ms,
p95_ms,
max_ms,
mean_ms,
phase_counts,
}
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("failure-recovery-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &FailureRecoveryBaselineConfig,
) -> Result<FailureRecoveryBaselineReport, Box<dyn std::error::Error>> {
let managed_redis = if config.redis_url.is_none() {
Some(ManagedRedisServer::start().await?)
} else {
None
};
let managed_postgres = if config.postgres_url.is_none() {
Some(ManagedPostgresServer::start().await?)
} else {
None
};
let redis_url = config
.redis_url
.clone()
.or_else(|| {
managed_redis
.as_ref()
.map(|server| server.redis_url().to_string())
})
.expect("redis url should resolve");
let postgres_url = config
.postgres_url
.clone()
.or_else(|| {
managed_postgres
.as_ref()
.map(|server| server.database_url().to_string())
})
.expect("postgres url should resolve");
let redis_server = Arc::new(Mutex::new(
managed_redis.ok_or("failure recovery baseline requires managed redis")?,
));
let postgres_server =
managed_postgres.ok_or("failure recovery baseline requires managed postgres")?;
let redis_restart = benchmark_redis_restart_recovery(redis_server.clone(), config).await?;
let postgres_slow_query =
benchmark_postgres_slow_query_recovery(postgres_server.database_url(), config).await?;
let hub_restart = benchmark_hub_restart_recovery(config).await?;
Ok(FailureRecoveryBaselineReport {
suite: "failure_recovery_baseline",
redis_url,
postgres_url,
redis_restart,
postgres_slow_query,
hub_restart,
})
}
async fn benchmark_redis_restart_recovery(
redis_server: Arc<Mutex<ManagedRedisServer>>,
config: &FailureRecoveryBaselineConfig,
) -> Result<RecoverySummary, Box<dyn std::error::Error>> {
let redis_url = redis_server.lock().await.redis_url().to_string();
let backend = RedisBackend::from_config(RedisClientConfig {
url: redis_url,
key_prefix: Some(format!("aether-failure-recovery-{}", std::process::id())),
})?;
let runner = backend.lock_runner(RedisLockRunnerConfig {
command_timeout_ms: Some(250),
default_ttl_ms: 1_000,
})?;
let keyspace = backend.keyspace();
let collector = Arc::new(RecoveryCollector::default());
let next_attempt = Arc::new(AtomicUsize::new(0));
let phase = Arc::new(AtomicUsize::new(0));
let recovered_after_restart_ms = Arc::new(AtomicU64::new(0));
let absolute_restart_started = Arc::new(Mutex::new(None::<Instant>));
let restart_phase = phase.clone();
let restart_started = absolute_restart_started.clone();
let server_for_restart = redis_server.clone();
let redis_restart_delay = config.redis_restart_delay;
let redis_downtime = config.redis_downtime;
let restart_task = tokio::spawn(async move {
tokio::time::sleep(redis_restart_delay).await;
restart_phase.store(1, Ordering::Release);
*restart_started.lock().await = Some(Instant::now());
{
let mut server = server_for_restart.lock().await;
server.stop().map_err(std::io::Error::other)?;
}
tokio::time::sleep(redis_downtime).await;
{
let mut server = server_for_restart.lock().await;
server
.restart()
.await
.map_err(|err| std::io::Error::other(err.to_string()))?;
}
restart_phase.store(2, Ordering::Release);
Ok::<(), std::io::Error>(())
});
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..config.redis_concurrency {
let runner = runner.clone();
let keyspace = keyspace.clone();
let collector = collector.clone();
let next_attempt = next_attempt.clone();
let phase = phase.clone();
let restart_started = absolute_restart_started.clone();
let recovered_after_restart_ms = recovered_after_restart_ms.clone();
let total_attempts = config.redis_attempts;
tasks.spawn(async move {
loop {
let current = next_attempt.fetch_add(1, Ordering::AcqRel);
if current >= total_attempts {
break;
}
let current_phase = classify_phase(phase.load(Ordering::Acquire));
let key = keyspace.lock_key(&format!("recovery-lock-{current}"));
let owner = format!("redis-owner-{current}");
let started = Instant::now();
let success = match runner.try_acquire(&key, &owner, Some(1_000)).await {
Ok(Some(lease)) => runner.release(&lease).await.unwrap_or(false),
Ok(None) => false,
Err(_) => false,
};
if success && matches!(current_phase, FaultPhase::Post) {
if let Some(restart_started_at) = *restart_started.lock().await {
let elapsed = restart_started_at.elapsed().as_millis() as u64;
let _ = recovered_after_restart_ms.compare_exchange(
0,
elapsed,
Ordering::AcqRel,
Ordering::Acquire,
);
}
}
collector
.record(current_phase, success, started.elapsed())
.await;
tokio::time::sleep(Duration::from_millis(10)).await;
}
});
}
while let Some(result) = tasks.join_next().await {
result.map_err(std::io::Error::other)?;
}
restart_task
.await
.map_err(|err| format!("redis restart task failed: {err}"))?
.map_err(std::io::Error::other)?;
Ok(collector
.summarize(load_optional_atomic_u64(&recovered_after_restart_ms))
.await)
}
async fn benchmark_postgres_slow_query_recovery(
postgres_url: &str,
config: &FailureRecoveryBaselineConfig,
) -> Result<PostgresSlowQueryRecoveryReport, Box<dyn std::error::Error>> {
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: postgres_url.to_string(),
min_connections: 1,
max_connections: 8,
acquire_timeout_ms: config.timeout.as_millis() as u64,
idle_timeout_ms: 60_000,
max_lifetime_ms: 10 * 60_000,
statement_cache_capacity: 64,
require_ssl: false,
})?;
bootstrap_failure_recovery_lease_table(backend.pool_clone()).await?;
let transaction_runner = backend.transaction_runner();
let lease_runner = backend.lease_runner(PostgresLeaseRunnerConfig {
statement_timeout_ms: Some(config.timeout.as_millis() as u64),
lock_timeout_ms: Some(1_000),
})?;
let postgres_sleep_secs = config.postgres_sleep.as_secs_f64();
let slow_query_started = Instant::now();
let slow_query_timed_out = transaction_runner
.run(
PostgresTransactionOptions {
statement_timeout_ms: Some(config.postgres_statement_timeout.as_millis() as u64),
..PostgresTransactionOptions::read_write()
},
|tx| {
async move {
sqlx::query("SELECT pg_sleep($1::double precision)")
.bind(postgres_sleep_secs)
.execute(&mut **tx)
.await?;
Ok(())
}
.boxed()
},
)
.await
.is_err();
let slow_query_latency_ms = slow_query_started.elapsed().as_millis() as u64;
let recovery_claim_started = Instant::now();
let claimed_ids = lease_runner
.claim_ids(
&PostgresLeaseClaimSpec {
table: "baseline_failure_lease_jobs",
id_column: "id",
lease_owner_column: "lease_owner",
lease_expires_at_column: "lease_expires_at",
eligibility_predicate_sql: "status = 'ready'",
order_by_sql: "id ASC",
},
PostgresLeaseClaimOptions {
batch_size: 8,
lease_ms: 250,
},
"recovery-owner",
)
.await
.unwrap_or_default();
let recovery_claim_latency_ms = recovery_claim_started.elapsed().as_millis() as u64;
Ok(PostgresSlowQueryRecoveryReport {
slow_query_timed_out,
slow_query_latency_ms,
recovery_claim_succeeded: !claimed_ids.is_empty(),
recovery_claim_latency_ms,
recovery_claimed_items: claimed_ids.len(),
})
}
async fn bootstrap_failure_recovery_lease_table(
pool: sqlx::PgPool,
) -> Result<(), Box<dyn std::error::Error>> {
sqlx::query("DROP TABLE IF EXISTS baseline_failure_lease_jobs")
.execute(&pool)
.await?;
sqlx::query(
"CREATE TABLE baseline_failure_lease_jobs (
id TEXT PRIMARY KEY,
status TEXT NOT NULL,
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
lease_owner TEXT,
lease_expires_at TIMESTAMPTZ
)",
)
.execute(&pool)
.await?;
let mut builder =
sqlx::QueryBuilder::new("INSERT INTO baseline_failure_lease_jobs (id, status) ");
builder.push_values(0..32, |mut row, index| {
row.push_bind(format!("recovery-job-{index:03}"))
.push_bind("ready");
});
builder.build().execute(&pool).await?;
Ok(())
}
async fn benchmark_hub_restart_recovery(
config: &FailureRecoveryBaselineConfig,
) -> Result<RecoverySummary, Box<dyn std::error::Error>> {
let port = reserve_local_port()?;
let hub_config = HubHarnessConfig::default();
let initial_hub = HubHarness::start_on_port(hub_config.clone(), port).await?;
let ws_url = format!("ws://127.0.0.1:{port}/proxy");
let collector = Arc::new(RecoveryCollector::default());
let next_attempt = Arc::new(AtomicUsize::new(0));
let phase = Arc::new(AtomicUsize::new(0));
let recovered_after_restart_ms = Arc::new(AtomicU64::new(0));
let restart_started = Arc::new(Mutex::new(None::<Instant>));
let (done_tx, done_rx) = oneshot::channel::<()>();
let hub_restart_delay = config.hub_restart_delay;
let hub_downtime = config.hub_downtime;
let phase_for_restart = phase.clone();
let restart_started_for_task = restart_started.clone();
let restart_task = tokio::spawn(async move {
tokio::time::sleep(hub_restart_delay).await;
phase_for_restart.store(1, Ordering::Release);
*restart_started_for_task.lock().await = Some(Instant::now());
drop(initial_hub);
tokio::time::sleep(hub_downtime).await;
let restarted_hub = start_hub_on_port_retry(hub_config, port).await?;
phase_for_restart.store(2, Ordering::Release);
let _ = done_rx.await;
drop(restarted_hub);
Ok::<(), String>(())
});
let mut workers = tokio::task::JoinSet::new();
for worker_index in 0..config.hub_concurrency {
let ws_url = ws_url.clone();
let next_attempt = next_attempt.clone();
let collector = collector.clone();
let phase = phase.clone();
let recovered_after_restart_ms = recovered_after_restart_ms.clone();
let restart_started = restart_started.clone();
let timeout = config.timeout;
let hold = config.hub_hold;
let total_attempts = config.hub_attempts;
workers.spawn(async move {
loop {
let current = next_attempt.fetch_add(1, Ordering::AcqRel);
if current >= total_attempts {
break;
}
let current_phase = classify_phase(phase.load(Ordering::Acquire));
let mut request = ws_url
.clone()
.into_client_request()
.map_err(|err| format!("failed to build websocket request: {err}"))?;
request.headers_mut().insert(
"x-node-id",
format!("recovery-node-{worker_index}-{current}")
.parse()
.map_err(|err| format!("failed to build x-node-id header: {err}"))?,
);
request.headers_mut().insert(
"x-node-name",
format!("recovery-node-{worker_index}-{current}")
.parse()
.map_err(|err| format!("failed to build x-node-name header: {err}"))?,
);
let started = Instant::now();
let success =
match tokio::time::timeout(timeout, tokio_tungstenite::connect_async(request))
.await
{
Ok(Ok((mut ws, _))) => {
tokio::time::sleep(hold).await;
let _ = ws.close(None).await;
while let Some(message) = ws.next().await {
if matches!(message, Ok(Message::Close(_))) || message.is_err() {
break;
}
}
true
}
_ => false,
};
if success && matches!(current_phase, FaultPhase::Post) {
if let Some(restart_started_at) = *restart_started.lock().await {
let elapsed = restart_started_at.elapsed().as_millis() as u64;
let _ = recovered_after_restart_ms.compare_exchange(
0,
elapsed,
Ordering::AcqRel,
Ordering::Acquire,
);
}
}
collector
.record(current_phase, success, started.elapsed())
.await;
tokio::time::sleep(Duration::from_millis(10)).await;
}
Ok::<(), String>(())
});
}
while let Some(result) = workers.join_next().await {
result
.map_err(|err| format!("hub recovery worker task failed: {err}"))?
.map_err(|err| format!("hub recovery worker failed: {err}"))?;
}
let _ = done_tx.send(());
restart_task
.await
.map_err(|err| format!("hub restart task failed: {err}"))?
.map_err(std::io::Error::other)?;
Ok(collector
.summarize(load_optional_atomic_u64(&recovered_after_restart_ms))
.await)
}
async fn start_hub_on_port_retry(
config: HubHarnessConfig,
port: u16,
) -> Result<HubHarness, String> {
let mut attempts = 0usize;
loop {
match HubHarness::start_on_port(config.clone(), port).await {
Ok(hub) => return Ok(hub),
Err(err) => {
attempts += 1;
if attempts >= 20 {
return Err(format!("failed to restart hub on fixed port {port}: {err}"));
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
}
}
fn classify_phase(value: usize) -> FaultPhase {
match value {
0 => FaultPhase::Pre,
1 => FaultPhase::During,
_ => FaultPhase::Post,
}
}
fn load_optional_atomic_u64(value: &AtomicU64) -> Option<u64> {
match value.load(Ordering::Acquire) {
0 => None,
millis => Some(millis),
}
}
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
}
let max_ms = *latencies.last().unwrap_or(&0);
let mean_ms = latencies.iter().sum::<u64>() / latencies.len() as u64;
let p50_ms = percentile(latencies, 50);
let p95_ms = percentile(latencies, 95);
(p50_ms, p95_ms, max_ms, mean_ms)
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {
if latencies.is_empty() {
return 0;
}
let last_index = latencies.len() - 1;
let rank = ((last_index as f64) * (percentile as f64 / 100.0)).round() as usize;
latencies[rank.min(last_index)]
}
fn parse_args(
args: Vec<String>,
) -> Result<FailureRecoveryBaselineConfig, Box<dyn std::error::Error>> {
let mut config = FailureRecoveryBaselineConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--redis-attempts" => {
config.redis_attempts = next_value(&mut iter, "--redis-attempts")?.parse()?
}
"--redis-concurrency" => {
config.redis_concurrency = next_value(&mut iter, "--redis-concurrency")?.parse()?
}
"--redis-restart-delay-ms" => {
config.redis_restart_delay = Duration::from_millis(
next_value(&mut iter, "--redis-restart-delay-ms")?.parse()?,
)
}
"--redis-downtime-ms" => {
config.redis_downtime =
Duration::from_millis(next_value(&mut iter, "--redis-downtime-ms")?.parse()?)
}
"--postgres-statement-timeout-ms" => {
config.postgres_statement_timeout = Duration::from_millis(
next_value(&mut iter, "--postgres-statement-timeout-ms")?.parse()?,
)
}
"--postgres-sleep-ms" => {
config.postgres_sleep =
Duration::from_millis(next_value(&mut iter, "--postgres-sleep-ms")?.parse()?)
}
"--hub-attempts" => {
config.hub_attempts = next_value(&mut iter, "--hub-attempts")?.parse()?
}
"--hub-concurrency" => {
config.hub_concurrency = next_value(&mut iter, "--hub-concurrency")?.parse()?
}
"--hub-hold-ms" => {
config.hub_hold =
Duration::from_millis(next_value(&mut iter, "--hub-hold-ms")?.parse()?)
}
"--hub-restart-delay-ms" => {
config.hub_restart_delay =
Duration::from_millis(next_value(&mut iter, "--hub-restart-delay-ms")?.parse()?)
}
"--hub-downtime-ms" => {
config.hub_downtime =
Duration::from_millis(next_value(&mut iter, "--hub-downtime-ms")?.parse()?)
}
"--timeout-ms" => {
config.timeout =
Duration::from_millis(next_value(&mut iter, "--timeout-ms")?.parse()?)
}
"--redis-url" => config.redis_url = Some(next_value(&mut iter, "--redis-url")?),
"--postgres-url" => {
config.postgres_url = Some(next_value(&mut iter, "--postgres-url")?)
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
other => return Err(format!("unknown argument: {other}").into()),
}
}
validate_config(&config)?;
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next()
.ok_or_else(|| format!("missing value for {flag}").into())
}
fn validate_config(
config: &FailureRecoveryBaselineConfig,
) -> Result<(), Box<dyn std::error::Error>> {
if config.redis_attempts == 0
|| config.redis_concurrency == 0
|| config.postgres_statement_timeout.is_zero()
|| config.postgres_sleep.is_zero()
|| config.hub_attempts == 0
|| config.hub_concurrency == 0
|| config.timeout.is_zero()
{
return Err("all failure recovery baseline numeric settings must be positive".into());
}
Ok(())
}

View File

@@ -0,0 +1,90 @@
use std::time::Duration;
use aether_testkit::{run_http_load_probe, HttpLoadProbeConfig};
use reqwest::Method;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = parse_args(std::env::args().skip(1).collect())?;
let result = run_http_load_probe(&config)
.await
.map_err(std::io::Error::other)?;
println!("{}", serde_json::to_string_pretty(&result)?);
Ok(())
}
fn parse_args(args: Vec<String>) -> Result<HttpLoadProbeConfig, Box<dyn std::error::Error>> {
let mut url: Option<String> = None;
let mut total_requests: Option<usize> = None;
let mut concurrency: Option<usize> = None;
let mut timeout_ms: Option<u64> = None;
let mut method = Method::GET;
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--url" => url = Some(next_value(&mut iter, "--url")?),
"--requests" => total_requests = Some(next_value(&mut iter, "--requests")?.parse()?),
"--concurrency" => concurrency = Some(next_value(&mut iter, "--concurrency")?.parse()?),
"--timeout-ms" => timeout_ms = Some(next_value(&mut iter, "--timeout-ms")?.parse()?),
"--method" => {
method = Method::from_bytes(next_value(&mut iter, "--method")?.as_bytes())?
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
let mut config = HttpLoadProbeConfig::default();
config.url = url.ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::InvalidInput, "missing required --url")
})?;
config.total_requests = total_requests.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"missing required --requests",
)
})?;
config.concurrency = concurrency.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"missing required --concurrency",
)
})?;
config.method = method;
if let Some(timeout_ms) = timeout_ms {
config.timeout = Duration::from_millis(timeout_ms);
}
config
.validate()
.map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin http_load_probe -- --url <URL> --requests <N> --concurrency <N> [--method GET] [--timeout-ms 30000]"
);
}

View File

@@ -0,0 +1,263 @@
use std::path::PathBuf;
use std::time::Duration;
use aether_hub::protocol;
use aether_testkit::{
init_test_runtime_for, run_http_load_probe, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
HttpLoadProbeResult, HubHarness, HubHarnessConfig,
};
use futures_util::{SinkExt, StreamExt};
use reqwest::Method;
use serde::Serialize;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message;
#[derive(Debug, Clone)]
struct HubTunnelBaselineConfig {
total_requests: usize,
concurrency: usize,
timeout: Duration,
output_path: Option<PathBuf>,
}
impl Default for HubTunnelBaselineConfig {
fn default() -> Self {
Self {
total_requests: 200,
concurrency: 20,
timeout: Duration::from_secs(10),
output_path: None,
}
}
}
#[derive(Debug, Serialize)]
struct HubTunnelBaselineReport {
suite: &'static str,
scenario: HttpLoadProbeResult,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("hub-tunnel-stream-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &HubTunnelBaselineConfig,
) -> Result<HubTunnelBaselineReport, Box<dyn std::error::Error>> {
let hub = HubHarness::start(HubHarnessConfig::default()).await?;
let peer = connect_protocol_peer(hub.base_url()).await?;
let result = run_http_load_probe(&HttpLoadProbeConfig {
url: format!("{}/local/relay/node-baseline", hub.base_url()),
method: Method::POST,
headers: std::collections::BTreeMap::from([(
"content-type".to_string(),
"application/octet-stream".to_string(),
)]),
body: Some(relay_envelope()),
total_requests: config.total_requests,
concurrency: config.concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
})
.await
.map_err(std::io::Error::other)?;
drop(peer);
Ok(HubTunnelBaselineReport {
suite: "hub_tunnel_stream_baseline",
scenario: result,
})
}
fn relay_envelope() -> Vec<u8> {
let meta = protocol::RequestMeta {
method: "POST".to_string(),
url: "https://baseline.example/v1/chat/completions".to_string(),
headers: std::collections::HashMap::from([(
"content-type".to_string(),
"application/json".to_string(),
)]),
timeout: 30,
};
let meta_json = serde_json::to_vec(&meta).expect("hub relay metadata should serialize");
let body = br#"{"model":"gpt-5","messages":[{"role":"user","content":"hello"}]}"#;
let mut envelope = Vec::with_capacity(4 + meta_json.len() + body.len());
envelope.extend_from_slice(&(meta_json.len() as u32).to_be_bytes());
envelope.extend_from_slice(&meta_json);
envelope.extend_from_slice(body);
envelope
}
async fn connect_protocol_peer(
hub_base_url: &str,
) -> Result<tokio::task::JoinHandle<()>, Box<dyn std::error::Error>> {
let ws_url = format!("{}/proxy", hub_base_url.replace("http://", "ws://"));
let request = ws_url.into_client_request()?;
let mut request = request;
request
.headers_mut()
.insert("x-node-id", http::HeaderValue::from_static("node-baseline"));
request.headers_mut().insert(
"x-node-name",
http::HeaderValue::from_static("proxy-baseline"),
);
request.headers_mut().insert(
"x-tunnel-max-streams",
http::HeaderValue::from_static("128"),
);
let (socket, _response) = tokio_tungstenite::connect_async(request).await?;
let (mut sink, mut stream) = socket.split();
Ok(tokio::spawn(async move {
while let Some(message) = stream.next().await {
let Ok(message) = message else {
break;
};
match message {
Message::Binary(data) => {
if handle_binary_frame(&mut sink, data.to_vec()).await.is_err() {
break;
}
}
Message::Ping(payload) => {
if sink.send(Message::Pong(payload)).await.is_err() {
break;
}
}
Message::Close(_) => break,
_ => {}
}
}
let _ = sink.close().await;
}))
}
async fn handle_binary_frame<S>(
sink: &mut S,
data: Vec<u8>,
) -> Result<(), tokio_tungstenite::tungstenite::Error>
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin,
{
let Some(header) = protocol::FrameHeader::parse(&data) else {
return Ok(());
};
match header.msg_type {
protocol::PING => {
let payload = protocol::frame_payload_by_header(&data, &header).unwrap_or(&[]);
sink.send(Message::Binary(protocol::encode_pong(payload).into()))
.await?;
}
protocol::REQUEST_HEADERS => {
let payload = protocol::decode_payload(&data, &header).unwrap_or_default();
let _ = serde_json::from_slice::<protocol::RequestMeta>(&payload);
}
protocol::REQUEST_BODY => {
if header.flags & protocol::FLAG_END_STREAM != 0 {
let response_meta = protocol::ResponseMeta {
status: 200,
headers: vec![(
"content-type".to_string(),
"text/plain; charset=utf-8".to_string(),
)],
};
let response_meta_json =
serde_json::to_vec(&response_meta).expect("response metadata should serialize");
sink.send(Message::Binary(
protocol::encode_frame(
header.stream_id,
protocol::RESPONSE_HEADERS,
0,
&response_meta_json,
)
.into(),
))
.await?;
for chunk in [
b"baseline-".as_slice(),
b"tunnel-".as_slice(),
b"stream".as_slice(),
] {
sink.send(Message::Binary(
protocol::encode_frame(header.stream_id, protocol::RESPONSE_BODY, 0, chunk)
.into(),
))
.await?;
}
sink.send(Message::Binary(
protocol::encode_frame(header.stream_id, protocol::STREAM_END, 0, &[]).into(),
))
.await?;
}
}
_ => {}
}
Ok(())
}
fn parse_args(args: Vec<String>) -> Result<HubTunnelBaselineConfig, Box<dyn std::error::Error>> {
let mut config = HubTunnelBaselineConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--requests" => config.total_requests = next_value(&mut iter, "--requests")?.parse()?,
"--concurrency" => {
config.concurrency = next_value(&mut iter, "--concurrency")?.parse()?
}
"--timeout-ms" => {
config.timeout =
Duration::from_millis(next_value(&mut iter, "--timeout-ms")?.parse()?)
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin hub_tunnel_stream_baseline -- [--requests 200] [--concurrency 20] [--timeout-ms 10000] [--output /tmp/hub_tunnel_baseline.json]"
);
}

View File

@@ -0,0 +1,656 @@
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_runtime::{DistributedConcurrencyGate, RedisDistributedConcurrencyConfig};
use aether_testkit::{
init_test_runtime_for, run_multi_url_http_load_probe, ExecutorHarness, ExecutorHarnessConfig,
GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
HubHarness, HubHarnessConfig, ManagedRedisServer, MultiUrlHttpLoadProbeResult, SpawnedServer,
};
use axum::body::to_bytes;
use axum::extract::Request;
use axum::response::IntoResponse;
use axum::routing::any;
use axum::{Json, Router};
use futures_util::StreamExt;
use reqwest::Method;
use serde::Serialize;
use serde_json::json;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::Message;
#[derive(Debug, Clone)]
struct MultiInstanceAdmissionBaselineConfig {
gateway_requests: usize,
gateway_concurrency: usize,
executor_requests: usize,
executor_concurrency: usize,
hub_attempts: usize,
hub_concurrency: usize,
hub_hold: Duration,
upstream_delay: Duration,
request_limit: usize,
hub_request_limit: usize,
timeout: Duration,
output_path: Option<PathBuf>,
redis_url: Option<String>,
}
impl Default for MultiInstanceAdmissionBaselineConfig {
fn default() -> Self {
Self {
gateway_requests: 200,
gateway_concurrency: 20,
executor_requests: 200,
executor_concurrency: 20,
hub_attempts: 40,
hub_concurrency: 10,
hub_hold: Duration::from_millis(100),
upstream_delay: Duration::from_millis(100),
request_limit: 8,
hub_request_limit: 4,
timeout: Duration::from_secs(10),
output_path: None,
redis_url: None,
}
}
}
#[derive(Debug, Serialize)]
struct MultiInstanceAdmissionBaselineReport {
suite: &'static str,
redis_url: String,
gateway_sync: MultiUrlHttpLoadProbeResult,
executor_sync: MultiUrlHttpLoadProbeResult,
hub_proxy: WebSocketAdmissionProbeResult,
}
#[derive(Debug, Clone, Serialize)]
struct WebSocketAdmissionProbeResult {
target_urls: Vec<String>,
target_attempt_counts: BTreeMap<String, usize>,
total_attempts: usize,
concurrency: usize,
completed_attempts: usize,
failed_attempts: usize,
rejected_attempts: usize,
successful_attempts: usize,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
status_counts: BTreeMap<u16, usize>,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("multi-instance-admission-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &MultiInstanceAdmissionBaselineConfig,
) -> Result<MultiInstanceAdmissionBaselineReport, Box<dyn std::error::Error>> {
let managed_redis = if config.redis_url.is_none() {
Some(ManagedRedisServer::start().await?)
} else {
None
};
let redis_url = config
.redis_url
.clone()
.or_else(|| {
managed_redis
.as_ref()
.map(|server| server.redis_url().to_string())
})
.expect("redis url should be resolved");
let upstream = SpawnedServer::start(build_delayed_upstream(config.upstream_delay)).await?;
let (gateway_urls, _gateways) =
start_gateway_pair(&redis_url, upstream.base_url(), config).await?;
let (executor_urls, _executors) =
start_executor_pair(&redis_url, upstream.base_url(), config).await?;
let (hub_urls, _hubs) = start_hub_pair(&redis_url, config).await?;
let gateway_sync = run_multi_url_http_load_probe(
&gateway_sync_probe_config(&gateway_urls, config),
&gateway_urls,
)
.await
.map_err(std::io::Error::other)?;
let executor_sync = run_multi_url_http_load_probe(
&executor_sync_probe_config(&executor_urls, upstream.base_url(), config),
&executor_urls,
)
.await
.map_err(std::io::Error::other)?;
let hub_proxy = run_hub_proxy_connection_probe(&hub_urls, config)
.await
.map_err(std::io::Error::other)?;
Ok(MultiInstanceAdmissionBaselineReport {
suite: "multi_instance_admission_baseline",
redis_url,
gateway_sync,
executor_sync,
hub_proxy,
})
}
async fn start_gateway_pair(
redis_url: &str,
upstream_base_url: &str,
config: &MultiInstanceAdmissionBaselineConfig,
) -> Result<(Vec<String>, Vec<GatewayHarness>), Box<dyn std::error::Error>> {
let gate_a = distributed_request_gate(
"gateway_requests_distributed",
config.request_limit,
redis_url,
"gateway-a",
)?;
let gate_b = distributed_request_gate(
"gateway_requests_distributed",
config.request_limit,
redis_url,
"gateway-a",
)?;
let gateway_a = GatewayHarness::start(GatewayHarnessConfig {
upstream_base_url: upstream_base_url.to_string(),
control_base_url: None,
executor_base_url: None,
max_in_flight_requests: None,
distributed_request_gate: Some(gate_a),
})
.await?;
let gateway_b = GatewayHarness::start(GatewayHarnessConfig {
upstream_base_url: upstream_base_url.to_string(),
control_base_url: None,
executor_base_url: None,
max_in_flight_requests: None,
distributed_request_gate: Some(gate_b),
})
.await?;
Ok((
vec![
format!("{}/v1/chat/completions", gateway_a.base_url()),
format!("{}/v1/chat/completions", gateway_b.base_url()),
],
vec![gateway_a, gateway_b],
))
}
async fn start_executor_pair(
redis_url: &str,
upstream_base_url: &str,
config: &MultiInstanceAdmissionBaselineConfig,
) -> Result<(Vec<String>, Vec<ExecutorHarness>), Box<dyn std::error::Error>> {
let gate_a = distributed_request_gate(
"executor_requests_distributed",
config.request_limit,
redis_url,
"executor-a",
)?;
let gate_b = distributed_request_gate(
"executor_requests_distributed",
config.request_limit,
redis_url,
"executor-a",
)?;
let executor_a = ExecutorHarness::start(ExecutorHarnessConfig {
max_in_flight_requests: None,
distributed_request_gate: Some(gate_a),
})
.await?;
let executor_b = ExecutorHarness::start(ExecutorHarnessConfig {
max_in_flight_requests: None,
distributed_request_gate: Some(gate_b),
})
.await?;
let _ = upstream_base_url;
Ok((
vec![
format!("{}/v1/execute/sync", executor_a.base_url()),
format!("{}/v1/execute/sync", executor_b.base_url()),
],
vec![executor_a, executor_b],
))
}
async fn start_hub_pair(
redis_url: &str,
config: &MultiInstanceAdmissionBaselineConfig,
) -> Result<(Vec<String>, Vec<HubHarness>), Box<dyn std::error::Error>> {
let gate_a = distributed_request_gate(
"hub_requests_distributed",
config.hub_request_limit,
redis_url,
"hub-a",
)?;
let gate_b = distributed_request_gate(
"hub_requests_distributed",
config.hub_request_limit,
redis_url,
"hub-a",
)?;
let hub_a = HubHarness::start(HubHarnessConfig {
distributed_request_gate: Some(gate_a),
..HubHarnessConfig::default()
})
.await?;
let hub_b = HubHarness::start(HubHarnessConfig {
distributed_request_gate: Some(gate_b),
..HubHarnessConfig::default()
})
.await?;
Ok((
vec![
format!("{}/proxy", hub_a.base_url().replace("http://", "ws://")),
format!("{}/proxy", hub_b.base_url().replace("http://", "ws://")),
],
vec![hub_a, hub_b],
))
}
fn distributed_request_gate(
name: &'static str,
limit: usize,
redis_url: &str,
key_scope: &str,
) -> Result<DistributedConcurrencyGate, Box<dyn std::error::Error>> {
Ok(DistributedConcurrencyGate::new_redis(
name,
limit,
RedisDistributedConcurrencyConfig {
url: redis_url.to_string(),
key_prefix: Some(format!(
"aether-baseline-{}-{name}-{key_scope}",
std::process::id()
)),
lease_ttl_ms: 30_000,
renew_interval_ms: 10_000,
command_timeout_ms: Some(1_000),
},
)?)
}
fn gateway_sync_probe_config(
urls: &[String],
config: &MultiInstanceAdmissionBaselineConfig,
) -> HttpLoadProbeConfig {
let mut probe = chat_probe_config(
urls[0].clone(),
config.gateway_requests,
config.gateway_concurrency,
config.timeout,
);
probe.response_mode = HttpLoadProbeResponseMode::FullBody;
probe
}
fn executor_sync_probe_config(
urls: &[String],
upstream_base_url: &str,
config: &MultiInstanceAdmissionBaselineConfig,
) -> HttpLoadProbeConfig {
let _ = urls;
HttpLoadProbeConfig {
url: urls[0].clone(),
method: Method::POST,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(
serde_json::to_vec(&execution_plan(format!(
"{upstream_base_url}/v1/chat/completions"
)))
.expect("execution plan should serialize"),
),
total_requests: config.executor_requests,
concurrency: config.executor_concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
}
}
fn chat_probe_config(
url: String,
total_requests: usize,
concurrency: usize,
timeout: Duration,
) -> HttpLoadProbeConfig {
HttpLoadProbeConfig {
url,
method: Method::POST,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(
serde_json::to_vec(&json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": false,
}))
.expect("chat body should serialize"),
),
total_requests,
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
}
}
fn execution_plan(url: String) -> ExecutionPlan {
ExecutionPlan {
request_id: "multi-instance-sync-request".to_string(),
candidate_id: Some("multi-instance-sync-candidate".to_string()),
provider_name: Some("openai".to_string()),
provider_id: "provider-baseline".to_string(),
endpoint_id: "endpoint-baseline".to_string(),
key_id: "key-baseline".to_string(),
method: "POST".to_string(),
url,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": false,
})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(2_000),
read_ms: Some(10_000),
first_byte_ms: Some(5_000),
total_ms: Some(10_000),
..ExecutionTimeouts::default()
}),
}
}
fn build_delayed_upstream(delay: Duration) -> Router {
Router::new().route(
"/v1/chat/completions",
any(move |request: Request| {
let delay = delay;
async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX)
.await
.expect("fake upstream body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).unwrap_or_else(|_| json!({}));
tokio::time::sleep(delay).await;
Json(json!({
"id": "chatcmpl-distributed",
"object": "chat.completion",
"model": payload.get("model").and_then(|value| value.as_str()).unwrap_or("gpt-5"),
"choices": [{"message": {"role": "assistant", "content": "hello"}}]
}))
.into_response()
}
}),
)
}
async fn run_hub_proxy_connection_probe(
urls: &[String],
config: &MultiInstanceAdmissionBaselineConfig,
) -> Result<WebSocketAdmissionProbeResult, String> {
if urls.is_empty() {
return Err("hub proxy connection probe requires at least one target url".to_string());
}
let next_attempt = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let latencies_ms = Arc::new(tokio::sync::Mutex::new(Vec::with_capacity(
config.hub_attempts,
)));
let target_attempt_counts = Arc::new(tokio::sync::Mutex::new(BTreeMap::<String, usize>::new()));
let status_counts = Arc::new(tokio::sync::Mutex::new(BTreeMap::<u16, usize>::new()));
let failed_attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let rejected_attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let successful_attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let completed_attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let mut workers = tokio::task::JoinSet::new();
for worker_index in 0..config.hub_concurrency {
let urls = urls.to_vec();
let next_attempt = Arc::clone(&next_attempt);
let latencies_ms = Arc::clone(&latencies_ms);
let target_attempt_counts = Arc::clone(&target_attempt_counts);
let status_counts = Arc::clone(&status_counts);
let failed_attempts = Arc::clone(&failed_attempts);
let rejected_attempts = Arc::clone(&rejected_attempts);
let successful_attempts = Arc::clone(&successful_attempts);
let completed_attempts = Arc::clone(&completed_attempts);
let timeout = config.timeout;
let hold = config.hub_hold;
let total_attempts = config.hub_attempts;
workers.spawn(async move {
loop {
let current = next_attempt.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
if current >= total_attempts {
break;
}
let url = urls[current % urls.len()].clone();
{
let mut counts = target_attempt_counts.lock().await;
*counts.entry(url.clone()).or_insert(0) += 1;
}
let mut request = url
.into_client_request()
.map_err(|err| format!("failed to build websocket request: {err}"))?;
request.headers_mut().insert(
"x-node-id",
format!("baseline-node-{worker_index}-{current}")
.parse()
.map_err(|err| format!("failed to build x-node-id header: {err}"))?,
);
request.headers_mut().insert(
"x-node-name",
format!("baseline-node-{worker_index}-{current}")
.parse()
.map_err(|err| format!("failed to build x-node-name header: {err}"))?,
);
let started_at = Instant::now();
match tokio::time::timeout(timeout, tokio_tungstenite::connect_async(request)).await
{
Ok(Ok((mut ws, _response))) => {
{
let mut counts = status_counts.lock().await;
*counts.entry(101).or_insert(0) += 1;
}
latencies_ms
.lock()
.await
.push(started_at.elapsed().as_millis() as u64);
successful_attempts.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
completed_attempts.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
tokio::time::sleep(hold).await;
let _ = ws.close(None).await;
while let Some(message) = ws.next().await {
if matches!(message, Ok(Message::Close(_))) || message.is_err() {
break;
}
}
}
Ok(Err(tokio_tungstenite::tungstenite::Error::Http(response))) => {
let status = response.status().as_u16();
{
let mut counts = status_counts.lock().await;
*counts.entry(status).or_insert(0) += 1;
}
latencies_ms
.lock()
.await
.push(started_at.elapsed().as_millis() as u64);
if status == 503 {
rejected_attempts.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
} else {
failed_attempts.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
}
completed_attempts.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
}
Ok(Err(_)) | Err(_) => {
latencies_ms
.lock()
.await
.push(started_at.elapsed().as_millis() as u64);
failed_attempts.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
completed_attempts.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
}
}
}
Ok::<(), String>(())
});
}
while let Some(result) = workers.join_next().await {
result
.map_err(|err| format!("hub admission worker task failed: {err}"))?
.map_err(|err| format!("hub admission worker failed: {err}"))?;
}
let mut latencies = latencies_ms.lock().await.clone();
latencies.sort_unstable();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
let target_attempt_counts = target_attempt_counts.lock().await.clone();
let status_counts = status_counts.lock().await.clone();
Ok(WebSocketAdmissionProbeResult {
target_urls: urls.to_vec(),
target_attempt_counts,
total_attempts: config.hub_attempts,
concurrency: config.hub_concurrency,
completed_attempts: completed_attempts.load(std::sync::atomic::Ordering::Acquire),
failed_attempts: failed_attempts.load(std::sync::atomic::Ordering::Acquire),
rejected_attempts: rejected_attempts.load(std::sync::atomic::Ordering::Acquire),
successful_attempts: successful_attempts.load(std::sync::atomic::Ordering::Acquire),
p50_ms,
p95_ms,
max_ms,
mean_ms,
status_counts,
})
}
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
}
let max_ms = *latencies.last().unwrap_or(&0);
let mean_ms = latencies.iter().sum::<u64>() / latencies.len() as u64;
let p50_ms = percentile(latencies, 50);
let p95_ms = percentile(latencies, 95);
(p50_ms, p95_ms, max_ms, mean_ms)
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {
if latencies.is_empty() {
return 0;
}
let last_index = latencies.len() - 1;
let rank = ((last_index as f64) * (percentile as f64 / 100.0)).round() as usize;
latencies[rank.min(last_index)]
}
fn parse_args(
args: Vec<String>,
) -> Result<MultiInstanceAdmissionBaselineConfig, Box<dyn std::error::Error>> {
let mut config = MultiInstanceAdmissionBaselineConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--gateway-requests" => {
config.gateway_requests = next_value(&mut iter, "--gateway-requests")?.parse()?
}
"--gateway-concurrency" => {
config.gateway_concurrency =
next_value(&mut iter, "--gateway-concurrency")?.parse()?
}
"--executor-requests" => {
config.executor_requests = next_value(&mut iter, "--executor-requests")?.parse()?
}
"--executor-concurrency" => {
config.executor_concurrency =
next_value(&mut iter, "--executor-concurrency")?.parse()?
}
"--hub-attempts" => {
config.hub_attempts = next_value(&mut iter, "--hub-attempts")?.parse()?
}
"--hub-concurrency" => {
config.hub_concurrency = next_value(&mut iter, "--hub-concurrency")?.parse()?
}
"--hub-hold-ms" => {
config.hub_hold =
Duration::from_millis(next_value(&mut iter, "--hub-hold-ms")?.parse()?)
}
"--upstream-delay-ms" => {
config.upstream_delay =
Duration::from_millis(next_value(&mut iter, "--upstream-delay-ms")?.parse()?)
}
"--request-limit" => {
config.request_limit = next_value(&mut iter, "--request-limit")?.parse()?
}
"--hub-request-limit" => {
config.hub_request_limit = next_value(&mut iter, "--hub-request-limit")?.parse()?
}
"--timeout-ms" => {
config.timeout =
Duration::from_millis(next_value(&mut iter, "--timeout-ms")?.parse()?)
}
"--redis-url" => config.redis_url = Some(next_value(&mut iter, "--redis-url")?),
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin multi_instance_admission_baseline -- [--gateway-requests 200] [--gateway-concurrency 20] [--executor-requests 200] [--executor-concurrency 20] [--hub-attempts 40] [--hub-concurrency 10] [--hub-hold-ms 100] [--upstream-delay-ms 100] [--request-limit 8] [--hub-request-limit 4] [--redis-url redis://127.0.0.1:6379/0] [--output /tmp/multi_instance_admission_baseline.json]"
);
}

View File

@@ -0,0 +1,431 @@
use std::path::PathBuf;
use std::time::{Duration, Instant};
use aether_data::redis::{
RedisClientConfig, RedisConsumerGroup, RedisConsumerName, RedisStreamReclaimConfig,
RedisStreamRunnerConfig,
};
use aether_data::RedisBackend;
use aether_testkit::{init_test_runtime_for, ManagedRedisServer};
use serde::Serialize;
#[derive(Debug, Clone)]
struct RedisWorkerBaselineConfig {
append_total: usize,
append_concurrency: usize,
reclaim_total: usize,
reclaim_min_idle: Duration,
output_path: Option<PathBuf>,
redis_url: Option<String>,
}
impl Default for RedisWorkerBaselineConfig {
fn default() -> Self {
Self {
append_total: 1_000,
append_concurrency: 20,
reclaim_total: 128,
reclaim_min_idle: Duration::from_millis(100),
output_path: None,
redis_url: None,
}
}
}
#[derive(Debug, Serialize)]
struct OperationSummary {
total_calls: usize,
total_items: usize,
failed_calls: usize,
p50_ms: u64,
p95_ms: u64,
max_ms: u64,
mean_ms: u64,
}
#[derive(Debug, Serialize)]
struct RedisWorkerBaselineReport {
suite: &'static str,
redis_url: String,
append: OperationSummary,
read_group: OperationSummary,
reclaim: OperationSummary,
ack: OperationSummary,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("redis-worker-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &RedisWorkerBaselineConfig,
) -> Result<RedisWorkerBaselineReport, Box<dyn std::error::Error>> {
let managed_redis = if config.redis_url.is_none() {
Some(ManagedRedisServer::start().await?)
} else {
None
};
let redis_url = config
.redis_url
.clone()
.or_else(|| {
managed_redis
.as_ref()
.map(|server| server.redis_url().to_string())
})
.expect("redis url should be resolved");
let backend = RedisBackend::from_config(RedisClientConfig {
url: redis_url.clone(),
key_prefix: Some(format!("aether-baseline-{}", std::process::id())),
})?;
let stream = backend.keyspace().stream_name("worker-baseline");
let group = RedisConsumerGroup("worker-group".to_string());
let consumer_a = RedisConsumerName("consumer-a".to_string());
let consumer_b = RedisConsumerName("consumer-b".to_string());
let runner = backend.stream_runner(RedisStreamRunnerConfig {
command_timeout_ms: Some(2_000),
read_block_ms: Some(10),
read_count: 64,
})?;
runner
.ensure_consumer_group(&stream, &group, "0-0")
.await
.map_err(std::io::Error::other)?;
let append = benchmark_append(&runner, &stream, config).await?;
let (read_group, drained_ids) =
benchmark_read_group(&runner, &stream, &group, &consumer_a, config).await?;
let ack_read = benchmark_ack(&runner, &stream, &group, &drained_ids).await?;
let (reclaim, reclaimed_ids) =
benchmark_reclaim(&runner, &stream, &group, &consumer_a, &consumer_b, config).await?;
let ack_reclaim = benchmark_ack(&runner, &stream, &group, &reclaimed_ids).await?;
Ok(RedisWorkerBaselineReport {
suite: "redis_worker_baseline",
redis_url,
append,
read_group,
reclaim,
ack: combine_summaries(&ack_read, &ack_reclaim),
})
}
async fn benchmark_append(
runner: &aether_data::redis::RedisStreamRunner,
stream: &aether_data::redis::RedisStreamName,
config: &RedisWorkerBaselineConfig,
) -> Result<OperationSummary, Box<dyn std::error::Error>> {
let next = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let latencies = std::sync::Arc::new(tokio::sync::Mutex::new(Vec::with_capacity(
config.append_total,
)));
let failed = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let mut tasks = tokio::task::JoinSet::new();
for _ in 0..config.append_concurrency {
let runner = runner.clone();
let stream = stream.clone();
let next = next.clone();
let latencies = latencies.clone();
let failed = failed.clone();
let total = config.append_total;
tasks.spawn(async move {
loop {
let current = next.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
if current >= total {
break;
}
let started = Instant::now();
let result = runner
.append_json(
&stream,
"payload",
&serde_json::json!({
"job_id": current,
"kind": "baseline",
}),
)
.await;
latencies
.lock()
.await
.push(started.elapsed().as_millis() as u64);
if result.is_err() {
failed.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
}
}
});
}
while let Some(result) = tasks.join_next().await {
result.map_err(std::io::Error::other)?;
}
let samples = latencies.lock().await.clone();
Ok(summarize_operation(
samples,
config.append_total,
failed.load(std::sync::atomic::Ordering::Acquire),
))
}
async fn benchmark_read_group(
runner: &aether_data::redis::RedisStreamRunner,
stream: &aether_data::redis::RedisStreamName,
group: &RedisConsumerGroup,
consumer: &RedisConsumerName,
config: &RedisWorkerBaselineConfig,
) -> Result<(OperationSummary, Vec<String>), Box<dyn std::error::Error>> {
let mut latencies = Vec::new();
let mut failed = 0usize;
let mut ids = Vec::with_capacity(config.append_total);
while ids.len() < config.append_total {
let started = Instant::now();
match runner.read_group(stream, group, consumer).await {
Ok(entries) => {
latencies.push(started.elapsed().as_millis() as u64);
ids.extend(entries.into_iter().map(|entry| entry.id));
}
Err(_) => {
latencies.push(started.elapsed().as_millis() as u64);
failed += 1;
}
}
}
let summary = summarize_operation(latencies, ids.len(), failed);
Ok((summary, ids))
}
async fn benchmark_reclaim(
runner: &aether_data::redis::RedisStreamRunner,
stream: &aether_data::redis::RedisStreamName,
group: &RedisConsumerGroup,
consumer_a: &RedisConsumerName,
consumer_b: &RedisConsumerName,
config: &RedisWorkerBaselineConfig,
) -> Result<(OperationSummary, Vec<String>), Box<dyn std::error::Error>> {
for index in 0..config.reclaim_total {
runner
.append_json(
stream,
"payload",
&serde_json::json!({
"job_id": format!("reclaim-{index}"),
"kind": "baseline",
}),
)
.await
.map_err(std::io::Error::other)?;
}
let mut pending_ids = Vec::with_capacity(config.reclaim_total);
while pending_ids.len() < config.reclaim_total {
let entries = runner
.read_group(stream, group, consumer_a)
.await
.map_err(std::io::Error::other)?;
pending_ids.extend(entries.into_iter().map(|entry| entry.id));
}
tokio::time::sleep(config.reclaim_min_idle + Duration::from_millis(20)).await;
let mut latencies = Vec::new();
let mut failed = 0usize;
let mut reclaimed_ids = Vec::with_capacity(config.reclaim_total);
let mut next_start_id = "0-0".to_string();
while reclaimed_ids.len() < config.reclaim_total {
let started = Instant::now();
match runner
.claim_stale(
stream,
group,
consumer_b,
&next_start_id,
RedisStreamReclaimConfig {
min_idle_ms: config.reclaim_min_idle.as_millis() as u64,
count: 64,
},
)
.await
{
Ok(result) => {
latencies.push(started.elapsed().as_millis() as u64);
next_start_id = result.next_start_id;
reclaimed_ids.extend(result.entries.into_iter().map(|entry| entry.id));
if next_start_id == "0-0" && reclaimed_ids.len() < config.reclaim_total {
failed += 1;
break;
}
}
Err(_) => {
latencies.push(started.elapsed().as_millis() as u64);
failed += 1;
}
}
}
let summary = summarize_operation(latencies, reclaimed_ids.len(), failed);
Ok((summary, reclaimed_ids))
}
async fn benchmark_ack(
runner: &aether_data::redis::RedisStreamRunner,
stream: &aether_data::redis::RedisStreamName,
group: &RedisConsumerGroup,
ids: &[String],
) -> Result<OperationSummary, Box<dyn std::error::Error>> {
let mut latencies = Vec::new();
let mut failed = 0usize;
let mut acked = 0usize;
for chunk in ids.chunks(64) {
let started = Instant::now();
match runner.ack(stream, group, chunk).await {
Ok(count) => {
latencies.push(started.elapsed().as_millis() as u64);
acked += count;
}
Err(_) => {
latencies.push(started.elapsed().as_millis() as u64);
failed += 1;
}
}
}
Ok(summarize_operation(latencies, acked, failed))
}
fn summarize_operation(
latencies: Vec<u64>,
total_items: usize,
failed_calls: usize,
) -> OperationSummary {
if latencies.is_empty() {
return OperationSummary {
total_calls: 0,
total_items,
failed_calls,
p50_ms: 0,
p95_ms: 0,
max_ms: 0,
mean_ms: 0,
};
}
let mut sorted = latencies;
sorted.sort_unstable();
let total_calls = sorted.len();
let max_ms = *sorted.last().unwrap_or(&0);
let mean_ms = sorted.iter().sum::<u64>() / total_calls as u64;
let p50_ms = percentile(&sorted, 50);
let p95_ms = percentile(&sorted, 95);
OperationSummary {
total_calls,
total_items,
failed_calls,
p50_ms,
p95_ms,
max_ms,
mean_ms,
}
}
fn combine_summaries(first: &OperationSummary, second: &OperationSummary) -> OperationSummary {
let total_calls = first.total_calls + second.total_calls;
let total_items = first.total_items + second.total_items;
let failed_calls = first.failed_calls + second.failed_calls;
let max_ms = first.max_ms.max(second.max_ms);
let mean_ms = if total_calls == 0 {
0
} else {
((first.mean_ms * first.total_calls as u64) + (second.mean_ms * second.total_calls as u64))
/ total_calls as u64
};
OperationSummary {
total_calls,
total_items,
failed_calls,
p50_ms: first.p50_ms.min(second.p50_ms),
p95_ms: first.p95_ms.max(second.p95_ms),
max_ms,
mean_ms,
}
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {
if latencies.is_empty() {
return 0;
}
let last_index = latencies.len() - 1;
let rank = ((last_index as f64) * (percentile as f64 / 100.0)).round() as usize;
latencies[rank.min(last_index)]
}
fn parse_args(args: Vec<String>) -> Result<RedisWorkerBaselineConfig, Box<dyn std::error::Error>> {
let mut config = RedisWorkerBaselineConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--append-total" => {
config.append_total = next_value(&mut iter, "--append-total")?.parse()?
}
"--append-concurrency" => {
config.append_concurrency =
next_value(&mut iter, "--append-concurrency")?.parse()?
}
"--reclaim-total" => {
config.reclaim_total = next_value(&mut iter, "--reclaim-total")?.parse()?
}
"--reclaim-min-idle-ms" => {
config.reclaim_min_idle =
Duration::from_millis(next_value(&mut iter, "--reclaim-min-idle-ms")?.parse()?)
}
"--redis-url" => config.redis_url = Some(next_value(&mut iter, "--redis-url")?),
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin redis_worker_baseline -- [--append-total 1000] [--append-concurrency 20] [--reclaim-total 128] [--reclaim-min-idle-ms 100] [--redis-url redis://127.0.0.1:6379/0] [--output /tmp/redis_worker_baseline.json]"
);
}

View File

@@ -0,0 +1,373 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::path::PathBuf;
use std::time::Duration;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_testkit::{
init_test_runtime_for, run_http_load_probe, ExecutorHarness, ExecutorHarnessConfig,
GatewayHarness, GatewayHarnessConfig, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
HttpLoadProbeResult, SpawnedServer,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::any;
use axum::{extract::Request, Json, Router};
use reqwest::Method;
use serde::Serialize;
use serde_json::json;
#[derive(Debug, Clone)]
struct SingleInstanceBaselineConfig {
sync_requests: usize,
sync_concurrency: usize,
stream_requests: usize,
stream_concurrency: usize,
timeout: Duration,
output_path: Option<PathBuf>,
}
impl Default for SingleInstanceBaselineConfig {
fn default() -> Self {
Self {
sync_requests: 200,
sync_concurrency: 20,
stream_requests: 100,
stream_concurrency: 10,
timeout: Duration::from_secs(10),
output_path: None,
}
}
}
#[derive(Debug, Serialize)]
struct NamedBaselineResult {
name: String,
result: HttpLoadProbeResult,
}
#[derive(Debug, Serialize)]
struct SingleInstanceBaselineReport {
suite: &'static str,
scenarios: Vec<NamedBaselineResult>,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
init_test_runtime_for("single-instance-baseline");
let config = parse_args(std::env::args().skip(1).collect())?;
let report = run_suite(&config).await?;
let raw = serde_json::to_string_pretty(&report)?;
println!("{raw}");
if let Some(path) = config.output_path.as_ref() {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, format!("{raw}\n"))?;
}
Ok(())
}
async fn run_suite(
config: &SingleInstanceBaselineConfig,
) -> Result<SingleInstanceBaselineReport, Box<dyn std::error::Error>> {
let upstream = SpawnedServer::start(build_fake_upstream()).await?;
let gateway = GatewayHarness::start(GatewayHarnessConfig::new(upstream.base_url())).await?;
let executor = ExecutorHarness::start(ExecutorHarnessConfig::default()).await?;
let gateway_sync = run_http_load_probe(&gateway_sync_probe_config(gateway.base_url(), config))
.await
.map_err(std::io::Error::other)?;
let gateway_stream =
run_http_load_probe(&gateway_stream_probe_config(gateway.base_url(), config))
.await
.map_err(std::io::Error::other)?;
let executor_sync = run_http_load_probe(&executor_sync_probe_config(
executor.base_url(),
upstream.base_url(),
config,
))
.await
.map_err(std::io::Error::other)?;
let executor_stream = run_http_load_probe(&executor_stream_probe_config(
executor.base_url(),
upstream.base_url(),
config,
))
.await
.map_err(std::io::Error::other)?;
Ok(SingleInstanceBaselineReport {
suite: "single_instance_baseline",
scenarios: vec![
NamedBaselineResult {
name: "gateway_proxy_sync".to_string(),
result: gateway_sync,
},
NamedBaselineResult {
name: "gateway_proxy_stream".to_string(),
result: gateway_stream,
},
NamedBaselineResult {
name: "executor_sync".to_string(),
result: executor_sync,
},
NamedBaselineResult {
name: "executor_stream".to_string(),
result: executor_stream,
},
],
})
}
fn gateway_sync_probe_config(
gateway_base_url: &str,
config: &SingleInstanceBaselineConfig,
) -> HttpLoadProbeConfig {
let mut probe = chat_probe_config(
format!("{gateway_base_url}/v1/chat/completions"),
false,
config.sync_requests,
config.sync_concurrency,
config.timeout,
);
probe.response_mode = HttpLoadProbeResponseMode::FullBody;
probe
}
fn gateway_stream_probe_config(
gateway_base_url: &str,
config: &SingleInstanceBaselineConfig,
) -> HttpLoadProbeConfig {
let mut probe = chat_probe_config(
format!("{gateway_base_url}/v1/chat/completions"),
true,
config.stream_requests,
config.stream_concurrency,
config.timeout,
);
probe.response_mode = HttpLoadProbeResponseMode::FullBody;
probe
}
fn executor_sync_probe_config(
executor_base_url: &str,
upstream_base_url: &str,
config: &SingleInstanceBaselineConfig,
) -> HttpLoadProbeConfig {
execution_probe_config(
format!("{executor_base_url}/v1/execute/sync"),
execution_plan(format!("{upstream_base_url}/v1/chat/completions"), false),
config.sync_requests,
config.sync_concurrency,
config.timeout,
)
}
fn executor_stream_probe_config(
executor_base_url: &str,
upstream_base_url: &str,
config: &SingleInstanceBaselineConfig,
) -> HttpLoadProbeConfig {
execution_probe_config(
format!("{executor_base_url}/v1/execute/stream"),
execution_plan(format!("{upstream_base_url}/v1/chat/completions"), true),
config.stream_requests,
config.stream_concurrency,
config.timeout,
)
}
fn execution_probe_config(
url: String,
plan: ExecutionPlan,
total_requests: usize,
concurrency: usize,
timeout: Duration,
) -> HttpLoadProbeConfig {
HttpLoadProbeConfig {
url,
method: Method::POST,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(
serde_json::to_vec(&plan).expect("execution plan should serialize for load probe"),
),
total_requests,
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
}
}
fn chat_probe_config(
url: String,
stream: bool,
total_requests: usize,
concurrency: usize,
timeout: Duration,
) -> HttpLoadProbeConfig {
HttpLoadProbeConfig {
url,
method: Method::POST,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
body: Some(
serde_json::to_vec(&json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": stream,
}))
.expect("chat body should serialize"),
),
total_requests,
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
}
}
fn execution_plan(url: String, stream: bool) -> ExecutionPlan {
ExecutionPlan {
request_id: if stream {
"baseline-stream-request".to_string()
} else {
"baseline-sync-request".to_string()
},
candidate_id: Some(if stream {
"baseline-stream-candidate".to_string()
} else {
"baseline-sync-candidate".to_string()
}),
provider_name: Some("openai".to_string()),
provider_id: "provider-baseline".to_string(),
endpoint_id: "endpoint-baseline".to_string(),
key_id: "key-baseline".to_string(),
method: "POST".to_string(),
url,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": stream,
})),
stream,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(2_000),
read_ms: Some(10_000),
first_byte_ms: Some(5_000),
total_ms: Some(10_000),
..ExecutionTimeouts::default()
}),
}
}
fn build_fake_upstream() -> Router {
Router::new().route(
"/v1/chat/completions",
any(|request: Request| async move {
let (_parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX)
.await
.expect("fake upstream body should read");
let payload: serde_json::Value =
serde_json::from_slice(&raw_body).unwrap_or_else(|_| json!({}));
let stream = payload
.get("stream")
.and_then(|value| value.as_bool())
.unwrap_or(false);
if stream {
let body = futures_util::stream::iter([
Ok::<_, Infallible>(Bytes::from_static(
b"data: {\"id\":\"chunk-1\",\"choices\":[{\"delta\":{\"content\":\"hel\"}}]}\n\n",
)),
Ok::<_, Infallible>(Bytes::from_static(
b"data: {\"id\":\"chunk-2\",\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n",
)),
Ok::<_, Infallible>(Bytes::from_static(b"data: [DONE]\n\n")),
]);
Response::builder()
.status(StatusCode::OK)
.header(http::header::CONTENT_TYPE, "text/event-stream")
.body(Body::from_stream(body))
.expect("fake upstream stream response should build")
} else {
Json(json!({
"id": "chatcmpl-baseline",
"object": "chat.completion",
"model": payload.get("model").and_then(|value| value.as_str()).unwrap_or("gpt-5"),
"choices": [{"message": {"role": "assistant", "content": "hello"}}]
}))
.into_response()
}
}),
)
}
fn parse_args(
args: Vec<String>,
) -> Result<SingleInstanceBaselineConfig, Box<dyn std::error::Error>> {
let mut config = SingleInstanceBaselineConfig::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--sync-requests" => {
config.sync_requests = next_value(&mut iter, "--sync-requests")?.parse()?
}
"--sync-concurrency" => {
config.sync_concurrency = next_value(&mut iter, "--sync-concurrency")?.parse()?
}
"--stream-requests" => {
config.stream_requests = next_value(&mut iter, "--stream-requests")?.parse()?
}
"--stream-concurrency" => {
config.stream_concurrency =
next_value(&mut iter, "--stream-concurrency")?.parse()?
}
"--timeout-ms" => {
config.timeout =
Duration::from_millis(next_value(&mut iter, "--timeout-ms")?.parse()?)
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
}
other => {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("unknown argument: {other}"),
)
.into());
}
}
}
Ok(config)
}
fn next_value(
iter: &mut impl Iterator<Item = String>,
flag: &str,
) -> Result<String, Box<dyn std::error::Error>> {
iter.next().ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!("missing value for {flag}"),
)
.into()
})
}
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin single_instance_baseline -- [--sync-requests 200] [--sync-concurrency 20] [--stream-requests 100] [--stream-concurrency 10] [--timeout-ms 10000] [--output /tmp/baseline.json]"
);
}

View File

@@ -0,0 +1,52 @@
use aether_executor::server::build_router_with_request_gates;
use aether_runtime::DistributedConcurrencyGate;
use crate::server::SpawnedServer;
#[derive(Debug, Clone, Default)]
pub struct ExecutorHarnessConfig {
pub max_in_flight_requests: Option<usize>,
pub distributed_request_gate: Option<DistributedConcurrencyGate>,
}
#[derive(Debug)]
pub struct ExecutorHarness {
server: SpawnedServer,
}
impl ExecutorHarness {
pub async fn start(config: ExecutorHarnessConfig) -> Result<Self, String> {
Self::start_with_server(config, None).await
}
pub async fn start_on_port(config: ExecutorHarnessConfig, port: u16) -> Result<Self, String> {
Self::start_with_server(config, Some(port)).await
}
async fn start_with_server(
config: ExecutorHarnessConfig,
port: Option<u16>,
) -> Result<Self, String> {
let router = build_router_with_request_gates(
config.max_in_flight_requests,
config.distributed_request_gate,
);
let server = match port {
Some(port) => SpawnedServer::start_on_port(port, router)
.await
.map_err(|err| format!("failed to start executor harness: {err}"))?,
None => SpawnedServer::start(router)
.await
.map_err(|err| format!("failed to start executor harness: {err}"))?,
};
Ok(Self { server })
}
pub fn base_url(&self) -> &str {
self.server.base_url()
}
pub fn port(&self) -> u16 {
self.server.port()
}
}

View File

@@ -0,0 +1,3 @@
pub fn test_trace_id(prefix: &str) -> String {
format!("{prefix}-test-trace")
}

View File

@@ -0,0 +1,76 @@
use aether_gateway::{build_router_with_state, AppState};
use aether_runtime::DistributedConcurrencyGate;
use crate::server::SpawnedServer;
#[derive(Debug, Clone)]
pub struct GatewayHarnessConfig {
pub upstream_base_url: String,
pub control_base_url: Option<String>,
pub executor_base_url: Option<String>,
pub max_in_flight_requests: Option<usize>,
pub distributed_request_gate: Option<DistributedConcurrencyGate>,
}
impl GatewayHarnessConfig {
pub fn new(upstream_base_url: impl Into<String>) -> Self {
Self {
upstream_base_url: upstream_base_url.into(),
control_base_url: None,
executor_base_url: None,
max_in_flight_requests: None,
distributed_request_gate: None,
}
}
}
#[derive(Debug)]
pub struct GatewayHarness {
server: SpawnedServer,
}
impl GatewayHarness {
pub async fn start(config: GatewayHarnessConfig) -> Result<Self, String> {
Self::start_with_server(config, None).await
}
pub async fn start_on_port(config: GatewayHarnessConfig, port: u16) -> Result<Self, String> {
Self::start_with_server(config, Some(port)).await
}
async fn start_with_server(
config: GatewayHarnessConfig,
port: Option<u16>,
) -> Result<Self, String> {
let mut state = AppState::new_with_executor(
config.upstream_base_url,
config.control_base_url,
config.executor_base_url,
)
.map_err(|err| format!("failed to build gateway harness state: {err}"))?;
if let Some(limit) = config.max_in_flight_requests {
state = state.with_request_concurrency_limit(limit);
}
if let Some(gate) = config.distributed_request_gate {
state = state.with_distributed_request_concurrency_gate(gate);
}
let router = build_router_with_state(state);
let server = match port {
Some(port) => SpawnedServer::start_on_port(port, router)
.await
.map_err(|err| format!("failed to start gateway harness: {err}"))?,
None => SpawnedServer::start(router)
.await
.map_err(|err| format!("failed to start gateway harness: {err}"))?,
};
Ok(Self { server })
}
pub fn base_url(&self) -> &str {
self.server.base_url()
}
pub fn port(&self) -> u16 {
self.server.port()
}
}

View File

@@ -0,0 +1,18 @@
use aether_http::{build_http_client, HttpClientConfig};
pub fn json_body(value: serde_json::Value) -> serde_json::Value {
value
}
pub fn test_http_client_config() -> HttpClientConfig {
HttpClientConfig {
connect_timeout_ms: Some(1_000),
request_timeout_ms: Some(5_000),
user_agent: Some("aether-testkit".to_string()),
..HttpClientConfig::default()
}
}
pub fn test_http_client() -> reqwest::Client {
build_http_client(&test_http_client_config()).expect("failed to build test HTTP client")
}

View File

@@ -0,0 +1,83 @@
use std::time::Duration;
use aether_hub::{build_router_with_state, AppState, ConnConfig, ControlPlaneClient};
use aether_runtime::DistributedConcurrencyGate;
use crate::server::SpawnedServer;
#[derive(Debug, Clone)]
pub struct HubHarnessConfig {
pub max_streams: usize,
pub ping_interval: Duration,
pub idle_timeout: Duration,
pub outbound_queue_capacity: usize,
pub max_in_flight_requests: Option<usize>,
pub distributed_request_gate: Option<DistributedConcurrencyGate>,
}
impl Default for HubHarnessConfig {
fn default() -> Self {
Self {
max_streams: 128,
ping_interval: Duration::from_secs(15),
idle_timeout: Duration::ZERO,
outbound_queue_capacity: 128,
max_in_flight_requests: None,
distributed_request_gate: None,
}
}
}
#[derive(Debug)]
pub struct HubHarness {
server: SpawnedServer,
}
impl HubHarness {
pub async fn start(config: HubHarnessConfig) -> Result<Self, String> {
Self::start_with_server(config, None).await
}
pub async fn start_on_port(config: HubHarnessConfig, port: u16) -> Result<Self, String> {
Self::start_with_server(config, Some(port)).await
}
async fn start_with_server(
config: HubHarnessConfig,
port: Option<u16>,
) -> Result<Self, String> {
let state = AppState::new(
ControlPlaneClient::disabled(),
ConnConfig {
ping_interval: config.ping_interval,
idle_timeout: config.idle_timeout,
outbound_queue_capacity: config.outbound_queue_capacity,
},
config.max_streams,
)
.with_request_concurrency_limit(config.max_in_flight_requests);
let state = if let Some(gate) = config.distributed_request_gate {
state.with_distributed_request_gate(gate)
} else {
state
};
let router = build_router_with_state(state);
let server = match port {
Some(port) => SpawnedServer::start_on_port(port, router)
.await
.map_err(|err| format!("failed to start hub harness: {err}"))?,
None => SpawnedServer::start(router)
.await
.map_err(|err| format!("failed to start hub harness: {err}"))?,
};
Ok(Self { server })
}
pub fn base_url(&self) -> &str {
self.server.base_url()
}
pub fn port(&self) -> u16 {
self.server.port()
}
}

View File

@@ -0,0 +1,30 @@
mod executor;
mod fixtures;
mod gateway;
mod http;
mod hub;
mod load;
mod metrics;
mod postgres;
mod redis;
mod server;
mod tracing;
mod wait;
pub use executor::{ExecutorHarness, ExecutorHarnessConfig};
pub use fixtures::test_trace_id;
pub use gateway::{GatewayHarness, GatewayHarnessConfig};
pub use http::{json_body, test_http_client, test_http_client_config};
pub use hub::{HubHarness, HubHarnessConfig};
pub use load::{
run_http_load_probe, run_multi_url_http_load_probe, HttpLoadProbeConfig,
HttpLoadProbeResponseMode, HttpLoadProbeResult, MultiUrlHttpLoadProbeResult,
};
pub use metrics::{
fetch_prometheus_samples, find_metric_value_u64, parse_prometheus_samples, PrometheusSample,
};
pub use postgres::ManagedPostgresServer;
pub use redis::ManagedRedisServer;
pub use server::{reserve_local_port, SpawnedServer};
pub use tracing::{init_test_runtime, init_test_runtime_for, test_runtime_config};
pub use wait::wait_until;

View File

@@ -0,0 +1,353 @@
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use http::{HeaderMap, HeaderName, HeaderValue};
use reqwest::{Client, Method};
use tokio::sync::Mutex;
#[derive(Debug, Clone, Copy, serde::Serialize, PartialEq, Eq)]
pub enum HttpLoadProbeResponseMode {
HeadersOnly,
FullBody,
}
impl Default for HttpLoadProbeResponseMode {
fn default() -> Self {
Self::HeadersOnly
}
}
#[derive(Debug, Clone)]
pub struct HttpLoadProbeConfig {
pub url: String,
pub method: Method,
pub headers: BTreeMap<String, String>,
pub body: Option<Vec<u8>>,
pub total_requests: usize,
pub concurrency: usize,
pub timeout: Duration,
pub response_mode: HttpLoadProbeResponseMode,
}
impl Default for HttpLoadProbeConfig {
fn default() -> Self {
Self {
url: String::new(),
method: Method::GET,
headers: BTreeMap::new(),
body: None,
total_requests: 100,
concurrency: 10,
timeout: Duration::from_secs(30),
response_mode: HttpLoadProbeResponseMode::HeadersOnly,
}
}
}
impl HttpLoadProbeConfig {
pub fn validate(&self) -> Result<(), String> {
if self.url.trim().is_empty() {
return Err("load probe url cannot be empty".to_string());
}
if self.total_requests == 0 {
return Err("load probe total_requests must be positive".to_string());
}
if self.concurrency == 0 {
return Err("load probe concurrency must be positive".to_string());
}
if self.timeout.is_zero() {
return Err("load probe timeout must be positive".to_string());
}
Ok(())
}
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
pub struct HttpLoadProbeResult {
pub url: String,
pub method: String,
pub response_mode: HttpLoadProbeResponseMode,
pub total_requests: usize,
pub concurrency: usize,
pub completed_requests: usize,
pub failed_requests: usize,
pub p50_ms: u64,
pub p95_ms: u64,
pub max_ms: u64,
pub mean_ms: u64,
pub status_counts: BTreeMap<u16, usize>,
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
pub struct MultiUrlHttpLoadProbeResult {
pub target_urls: Vec<String>,
pub target_request_counts: BTreeMap<String, usize>,
pub method: String,
pub response_mode: HttpLoadProbeResponseMode,
pub total_requests: usize,
pub concurrency: usize,
pub completed_requests: usize,
pub failed_requests: usize,
pub p50_ms: u64,
pub p95_ms: u64,
pub max_ms: u64,
pub mean_ms: u64,
pub status_counts: BTreeMap<u16, usize>,
}
pub async fn run_http_load_probe(
config: &HttpLoadProbeConfig,
) -> Result<HttpLoadProbeResult, String> {
config.validate()?;
run_http_load_probe_against_urls(config, std::slice::from_ref(&config.url))
.await
.map(|result| HttpLoadProbeResult {
url: result
.target_urls
.into_iter()
.next()
.unwrap_or_else(|| config.url.clone()),
method: result.method,
response_mode: result.response_mode,
total_requests: result.total_requests,
concurrency: result.concurrency,
completed_requests: result.completed_requests,
failed_requests: result.failed_requests,
p50_ms: result.p50_ms,
p95_ms: result.p95_ms,
max_ms: result.max_ms,
mean_ms: result.mean_ms,
status_counts: result.status_counts,
})
}
pub async fn run_multi_url_http_load_probe(
config: &HttpLoadProbeConfig,
urls: &[String],
) -> Result<MultiUrlHttpLoadProbeResult, String> {
config.validate()?;
if urls.is_empty() {
return Err("multi-url load probe requires at least one target url".to_string());
}
run_http_load_probe_against_urls(config, urls).await
}
async fn run_http_load_probe_against_urls(
config: &HttpLoadProbeConfig,
urls: &[String],
) -> Result<MultiUrlHttpLoadProbeResult, String> {
let client = Client::builder()
.timeout(config.timeout)
.build()
.map_err(|err| format!("failed to build load probe http client: {err}"))?;
let total_requests = config.total_requests;
let request_headers = build_headers(&config.headers)?;
let request_body = config.body.clone().map(Arc::new);
let response_mode = config.response_mode;
let next_request = Arc::new(AtomicUsize::new(0));
let latencies_ms = Arc::new(Mutex::new(Vec::with_capacity(config.total_requests)));
let status_counts = Arc::new(Mutex::new(BTreeMap::<u16, usize>::new()));
let target_request_counts = Arc::new(Mutex::new(BTreeMap::<String, usize>::new()));
let failed_requests = Arc::new(AtomicUsize::new(0));
let completed_requests = Arc::new(AtomicUsize::new(0));
let mut workers = tokio::task::JoinSet::new();
for _ in 0..config.concurrency {
let client = client.clone();
let next_request = Arc::clone(&next_request);
let latencies_ms = Arc::clone(&latencies_ms);
let status_counts = Arc::clone(&status_counts);
let target_request_counts = Arc::clone(&target_request_counts);
let failed_requests = Arc::clone(&failed_requests);
let completed_requests = Arc::clone(&completed_requests);
let method = config.method.clone();
let urls = urls.to_vec();
let request_headers = request_headers.clone();
let request_body = request_body.clone();
workers.spawn(async move {
loop {
let current = next_request.fetch_add(1, Ordering::AcqRel);
if current >= total_requests {
break;
}
let started_at = Instant::now();
let url = urls[current % urls.len()].clone();
let mut request = client.request(method.clone(), &url);
for (name, value) in request_headers.iter() {
request = request.header(name, value);
}
if let Some(body) = request_body.as_ref() {
request = request.body(body.as_ref().clone());
}
match request.send().await {
Ok(response) => {
let status = response.status().as_u16();
let body_result = match response_mode {
HttpLoadProbeResponseMode::HeadersOnly => Ok(()),
HttpLoadProbeResponseMode::FullBody => {
response.bytes().await.map(|_| ()).map_err(|_| ())
}
};
if body_result.is_ok() {
let mut counts = status_counts.lock().await;
*counts.entry(status).or_insert(0) += 1;
drop(counts);
let mut target_counts = target_request_counts.lock().await;
*target_counts.entry(url).or_insert(0) += 1;
} else {
failed_requests.fetch_add(1, Ordering::AcqRel);
}
let latency_ms = started_at.elapsed().as_millis() as u64;
latencies_ms.lock().await.push(latency_ms);
completed_requests.fetch_add(1, Ordering::AcqRel);
}
Err(_) => {
let latency_ms = started_at.elapsed().as_millis() as u64;
latencies_ms.lock().await.push(latency_ms);
failed_requests.fetch_add(1, Ordering::AcqRel);
completed_requests.fetch_add(1, Ordering::AcqRel);
}
}
}
});
}
while let Some(result) = workers.join_next().await {
result.map_err(|err| format!("load probe worker task failed: {err}"))?;
}
let status_counts = status_counts.lock().await.clone();
let target_request_counts = target_request_counts.lock().await.clone();
let mut latencies = latencies_ms.lock().await.clone();
latencies.sort_unstable();
let (p50_ms, p95_ms, max_ms, mean_ms) = summarize_latencies(&latencies);
Ok(MultiUrlHttpLoadProbeResult {
target_urls: urls.to_vec(),
target_request_counts,
method: config.method.as_str().to_string(),
response_mode: config.response_mode,
total_requests: config.total_requests,
concurrency: config.concurrency,
completed_requests: completed_requests.load(Ordering::Acquire),
failed_requests: failed_requests.load(Ordering::Acquire),
p50_ms,
p95_ms,
max_ms,
mean_ms,
status_counts,
})
}
fn build_headers(headers: &BTreeMap<String, String>) -> Result<HeaderMap, String> {
let mut result = HeaderMap::new();
for (name, value) in headers {
let name = HeaderName::try_from(name.as_str())
.map_err(|err| format!("invalid load probe header name `{name}`: {err}"))?;
let value = HeaderValue::from_str(value)
.map_err(|err| format!("invalid load probe header value for `{name}`: {err}"))?;
result.insert(name, value);
}
Ok(result)
}
fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64) {
if latencies.is_empty() {
return (0, 0, 0, 0);
}
let max_ms = *latencies.last().unwrap_or(&0);
let mean_ms = latencies.iter().sum::<u64>() / latencies.len() as u64;
let p50_ms = percentile(latencies, 50);
let p95_ms = percentile(latencies, 95);
(p50_ms, p95_ms, max_ms, mean_ms)
}
fn percentile(latencies: &[u64], percentile: u8) -> u64 {
if latencies.is_empty() {
return 0;
}
let last_index = latencies.len() - 1;
let rank = ((last_index as f64) * (percentile as f64 / 100.0)).round() as usize;
latencies[rank.min(last_index)]
}
#[cfg(test)]
mod tests {
use super::{
build_headers, summarize_latencies, HttpLoadProbeConfig, HttpLoadProbeResponseMode,
};
use reqwest::Method;
use std::collections::BTreeMap;
use std::time::Duration;
#[test]
fn validates_probe_config() {
assert!(HttpLoadProbeConfig {
url: String::new(),
..HttpLoadProbeConfig::default()
}
.validate()
.is_err());
assert!(HttpLoadProbeConfig {
total_requests: 0,
..HttpLoadProbeConfig::default()
}
.validate()
.is_err());
assert!(HttpLoadProbeConfig {
concurrency: 0,
..HttpLoadProbeConfig::default()
}
.validate()
.is_err());
assert!(HttpLoadProbeConfig {
timeout: Duration::ZERO,
..HttpLoadProbeConfig::default()
}
.validate()
.is_err());
}
#[test]
fn summarizes_latency_distribution() {
let (p50_ms, p95_ms, max_ms, mean_ms) =
summarize_latencies(&[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]);
assert_eq!(p50_ms, 60);
assert_eq!(p95_ms, 100);
assert_eq!(max_ms, 100);
assert_eq!(mean_ms, 55);
}
#[test]
fn default_probe_config_is_reasonable() {
let config = HttpLoadProbeConfig::default();
assert_eq!(config.method, Method::GET);
assert!(config.headers.is_empty());
assert!(config.body.is_none());
assert_eq!(config.total_requests, 100);
assert_eq!(config.concurrency, 10);
assert_eq!(config.timeout, Duration::from_secs(30));
assert_eq!(config.response_mode, HttpLoadProbeResponseMode::HeadersOnly);
}
#[test]
fn validates_probe_headers() {
let mut headers = BTreeMap::new();
headers.insert("x-aether-test".to_string(), "ok".to_string());
let built = build_headers(&headers).expect("headers should build");
assert_eq!(
built
.get("x-aether-test")
.and_then(|value| value.to_str().ok()),
Some("ok")
);
let invalid = BTreeMap::from([("bad header".to_string(), "ok".to_string())]);
assert!(build_headers(&invalid).is_err());
}
}

View File

@@ -0,0 +1,136 @@
use std::collections::BTreeMap;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PrometheusSample {
pub name: String,
pub labels: BTreeMap<String, String>,
pub value: String,
}
pub async fn fetch_prometheus_samples(url: &str) -> Result<Vec<PrometheusSample>, String> {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(5))
.build()
.map_err(|err| format!("failed to build metrics http client: {err}"))?;
let response = client
.get(url)
.send()
.await
.map_err(|err| format!("failed to fetch metrics from {url}: {err}"))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|err| format!("failed to read metrics body from {url}: {err}"))?;
if !status.is_success() {
return Err(format!("metrics endpoint {url} returned {status}: {body}"));
}
Ok(parse_prometheus_samples(&body))
}
pub fn parse_prometheus_samples(text: &str) -> Vec<PrometheusSample> {
text.lines()
.filter_map(parse_prometheus_line)
.collect::<Vec<_>>()
}
pub fn find_metric_value_u64(
samples: &[PrometheusSample],
metric_name: &str,
labels: &[(&str, &str)],
) -> Option<u64> {
samples
.iter()
.find(|sample| {
metric_name_matches(&sample.name, metric_name) && labels_match(sample, labels)
})
.and_then(|sample| sample.value.parse::<u64>().ok())
}
fn metric_name_matches(actual: &str, expected: &str) -> bool {
actual == expected
|| actual
.rsplit_once('_')
.map(|(_, suffix)| suffix == expected)
.unwrap_or(false)
|| actual.ends_with(&format!("_{expected}"))
}
fn labels_match(sample: &PrometheusSample, labels: &[(&str, &str)]) -> bool {
labels
.iter()
.all(|(key, value)| sample.labels.get(*key).map(|current| current.as_str()) == Some(*value))
}
fn parse_prometheus_line(line: &str) -> Option<PrometheusSample> {
let trimmed = line.trim();
if trimmed.is_empty() || trimmed.starts_with('#') {
return None;
}
let (metric, value) = trimmed.rsplit_once(' ')?;
let (name, labels) = if let Some((name, raw_labels)) = metric
.split_once('{')
.and_then(|(name, rest)| rest.strip_suffix('}').map(|labels| (name, labels)))
{
(name.to_string(), parse_labels(raw_labels))
} else {
(metric.to_string(), BTreeMap::new())
};
Some(PrometheusSample {
name,
labels,
value: value.to_string(),
})
}
fn parse_labels(raw: &str) -> BTreeMap<String, String> {
let mut labels = BTreeMap::new();
for pair in raw.split(',').filter(|pair| !pair.is_empty()) {
if let Some((key, value)) = pair.split_once('=') {
labels.insert(
key.trim().to_string(),
value
.trim()
.trim_matches('"')
.replace("\\\"", "\"")
.replace("\\n", "\n")
.replace("\\\\", "\\"),
);
}
}
labels
}
#[cfg(test)]
mod tests {
use super::{find_metric_value_u64, parse_prometheus_samples};
#[test]
fn parses_prometheus_samples_with_labels() {
let samples = parse_prometheus_samples(
r#"
# HELP aether_gateway_concurrency_in_flight Current number of in-flight operations.
# TYPE aether_gateway_concurrency_in_flight gauge
aether_gateway_concurrency_in_flight{gate="gateway_requests"} 7
aether_gateway_concurrency_rejected_total{gate="gateway_requests"} 12
"#,
);
assert_eq!(
find_metric_value_u64(
&samples,
"concurrency_in_flight",
&[("gate", "gateway_requests")]
),
Some(7)
);
assert_eq!(
find_metric_value_u64(
&samples,
"concurrency_rejected_total",
&[("gate", "gateway_requests")]
),
Some(12)
);
}
}

View File

@@ -0,0 +1,148 @@
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use sqlx::{Connection, PgConnection};
use crate::wait_until;
#[derive(Debug)]
pub struct ManagedPostgresServer {
child: Option<Child>,
postgres_bin: String,
port: u16,
workdir: PathBuf,
data_dir: PathBuf,
database_url: String,
}
impl ManagedPostgresServer {
pub async fn start() -> Result<Self, Box<dyn std::error::Error>> {
let port = reserve_local_port()?;
let workdir = std::env::temp_dir().join(format!(
"aether-postgres-baseline-{}-{}",
std::process::id(),
port
));
let data_dir = workdir.join("data");
std::fs::create_dir_all(&workdir)?;
let initdb_bin = std::env::var("AETHER_INITDB_BIN")
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "initdb".to_string());
let postgres_bin = std::env::var("AETHER_POSTGRES_BIN")
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "postgres".to_string());
let init_output = Command::new(&initdb_bin)
.arg("-D")
.arg(&data_dir)
.arg("-U")
.arg("aether")
.arg("--auth=trust")
.arg("--encoding=UTF8")
.arg("--no-instructions")
.output()?;
if !init_output.status.success() {
return Err(std::io::Error::other(format!(
"initdb failed: {}",
String::from_utf8_lossy(&init_output.stderr)
))
.into());
}
let database_url = format!("postgres://aether@127.0.0.1:{port}/postgres");
let mut server = Self {
child: None,
postgres_bin,
port,
workdir,
data_dir,
database_url,
};
server.restart().await?;
Ok(server)
}
pub fn database_url(&self) -> &str {
&self.database_url
}
pub fn port(&self) -> u16 {
self.port
}
pub fn stop(&mut self) -> Result<(), std::io::Error> {
if let Some(mut child) = self.child.take() {
let _ = child.kill();
let _ = child.wait();
}
Ok(())
}
pub async fn restart(&mut self) -> Result<(), Box<dyn std::error::Error>> {
self.stop()?;
let log_path = self.workdir.join("postgres.log");
let stdout = std::fs::File::create(&log_path)?;
let stderr = stdout.try_clone()?;
let child = Command::new(&self.postgres_bin)
.arg("-D")
.arg(&self.data_dir)
.arg("-h")
.arg("127.0.0.1")
.arg("-p")
.arg(self.port.to_string())
.arg("-F")
.arg("-c")
.arg("fsync=off")
.arg("-c")
.arg("synchronous_commit=off")
.arg("-c")
.arg("full_page_writes=off")
.stdout(Stdio::from(stdout))
.stderr(Stdio::from(stderr))
.spawn()?;
self.child = Some(child);
let database_url = self.database_url.clone();
let ready = wait_until(
std::time::Duration::from_secs(10),
std::time::Duration::from_millis(50),
|| {
let database_url = database_url.clone();
async move {
match PgConnection::connect(&database_url).await {
Ok(connection) => connection.close().await.is_ok(),
Err(_) => false,
}
}
},
)
.await;
if !ready {
self.stop()?;
let logs = std::fs::read_to_string(&log_path).unwrap_or_default();
return Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
format!("timed out waiting for local postgres; logs:\n{logs}"),
)
.into());
}
Ok(())
}
}
impl Drop for ManagedPostgresServer {
fn drop(&mut self) {
let _ = self.stop();
let _ = std::fs::remove_dir_all(&self.workdir);
}
}
fn reserve_local_port() -> Result<u16, std::io::Error> {
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let port = listener.local_addr()?.port();
drop(listener);
Ok(port)
}

View File

@@ -0,0 +1,110 @@
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use crate::wait_until;
#[derive(Debug)]
pub struct ManagedRedisServer {
child: Option<Child>,
binary: String,
port: u16,
workdir: PathBuf,
redis_url: String,
}
impl ManagedRedisServer {
pub async fn start() -> Result<Self, Box<dyn std::error::Error>> {
let port = reserve_local_port()?;
let workdir = std::env::temp_dir().join(format!(
"aether-redis-baseline-{}-{}",
std::process::id(),
port
));
std::fs::create_dir_all(&workdir)?;
let binary = std::env::var("AETHER_REDIS_SERVER_BIN")
.ok()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "redis-server".to_string());
let redis_url = format!("redis://127.0.0.1:{port}/0");
let mut server = Self {
child: None,
binary,
port,
workdir,
redis_url,
};
server.restart().await?;
Ok(server)
}
pub fn redis_url(&self) -> &str {
&self.redis_url
}
pub fn port(&self) -> u16 {
self.port
}
pub fn stop(&mut self) -> Result<(), std::io::Error> {
if let Some(mut child) = self.child.take() {
let _ = child.kill();
let _ = child.wait();
}
Ok(())
}
pub async fn restart(&mut self) -> Result<(), Box<dyn std::error::Error>> {
self.stop()?;
let child = Command::new(&self.binary)
.arg("--save")
.arg("")
.arg("--appendonly")
.arg("no")
.arg("--port")
.arg(self.port.to_string())
.arg("--dir")
.arg(&self.workdir)
.arg("--bind")
.arg("127.0.0.1")
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()?;
self.child = Some(child);
let port = self.port;
let ready = wait_until(
std::time::Duration::from_secs(5),
std::time::Duration::from_millis(50),
|| async move {
tokio::net::TcpStream::connect(("127.0.0.1", port))
.await
.is_ok()
},
)
.await;
if !ready {
self.stop()?;
return Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"timed out waiting for local redis-server",
)
.into());
}
Ok(())
}
}
impl Drop for ManagedRedisServer {
fn drop(&mut self) {
let _ = self.stop();
let _ = std::fs::remove_dir_all(&self.workdir);
}
}
fn reserve_local_port() -> Result<u16, std::io::Error> {
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let port = listener.local_addr()?.port();
drop(listener);
Ok(port)
}

View File

@@ -0,0 +1,64 @@
use std::fmt;
use std::net::SocketAddr;
use axum::Router;
pub struct SpawnedServer {
base_url: String,
port: u16,
handle: tokio::task::JoinHandle<()>,
}
impl SpawnedServer {
pub async fn start(app: Router) -> Result<Self, std::io::Error> {
let port = reserve_local_port()?;
Self::start_on_port(port, app).await
}
pub async fn start_on_port(port: u16, app: Router) -> Result<Self, std::io::Error> {
let listener = tokio::net::TcpListener::bind(("127.0.0.1", port)).await?;
let addr = listener.local_addr()?;
let handle = tokio::spawn(async move {
axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.await
.expect("spawned server should run");
});
Ok(Self {
base_url: format!("http://{addr}"),
port: addr.port(),
handle,
})
}
pub fn base_url(&self) -> &str {
&self.base_url
}
pub fn port(&self) -> u16 {
self.port
}
}
impl fmt::Debug for SpawnedServer {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SpawnedServer")
.field("base_url", &self.base_url)
.finish_non_exhaustive()
}
}
impl Drop for SpawnedServer {
fn drop(&mut self) {
self.handle.abort();
}
}
pub fn reserve_local_port() -> Result<u16, std::io::Error> {
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let port = listener.local_addr()?.port();
drop(listener);
Ok(port)
}

View File

@@ -0,0 +1,14 @@
use aether_runtime::{init_service_runtime, ServiceRuntimeConfig};
pub fn init_test_runtime() {
init_test_runtime_for("aether-testkit");
}
pub fn test_runtime_config(service_name: &'static str) -> ServiceRuntimeConfig {
ServiceRuntimeConfig::new(service_name, "aether_testkit=debug")
.with_metrics_namespace("aether_testkit")
}
pub fn init_test_runtime_for(service_name: &'static str) {
let _ = init_service_runtime(test_runtime_config(service_name));
}

View File

@@ -0,0 +1,50 @@
use std::future::Future;
use std::time::Duration;
pub async fn wait_until<F, Fut>(
timeout: Duration,
poll_interval: Duration,
mut predicate: F,
) -> bool
where
F: FnMut() -> Fut,
Fut: Future<Output = bool>,
{
let deadline = tokio::time::Instant::now() + timeout;
loop {
if predicate().await {
return true;
}
if tokio::time::Instant::now() >= deadline {
return false;
}
tokio::time::sleep(poll_interval).await;
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use super::wait_until;
#[tokio::test]
async fn returns_true_when_predicate_eventually_passes() {
let flag = Arc::new(AtomicBool::new(false));
let background = flag.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
background.store(true, Ordering::Release);
});
let ready = wait_until(Duration::from_millis(100), Duration::from_millis(5), || {
let flag = flag.clone();
async move { flag.load(Ordering::Acquire) }
})
.await;
assert!(ready);
}
}