i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 10:57:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improve gateway transport and usage runtime

Browse Source
This commit is contained in:
elky
2026-06-25 22:36:27 +08:00
parent d336d1a7fa
commit 6f00e9fc67
112 changed files with 12456 additions and 1387 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-testkit/src/bin/capacity_curve_baseline.rs
+3
View File
@@ -331,6 +331,7 @@ async fn run_tunnel_curve(
concurrency: relay_concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
};
let started_at = Instant::now();
let result = run_http_load_probe(&probe)
@@ -480,6 +481,7 @@ fn execution_probe_config(
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
}
}
@@ -506,6 +508,7 @@ fn chat_probe_config(
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
}
}
crates/aether-testkit/src/bin/gateway_openai_chat_c_compare.rs
+919
View File
@@ -0,0 +1,919 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_gateway::build_router_with_state;
use aether_gateway::testkit::{build_openai_chat_pressure_state, OpenAiChatPressureStateConfig};
use aether_testkit::{
fetch_prometheus_samples, run_http_load_probe, run_multi_url_http_load_probe,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult,
MultiUrlHttpLoadProbeResult, PrometheusSample, SpawnedServer,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::extract::State;
use axum::http::header::{AUTHORIZATION, CACHE_CONTROL, CONTENT_TYPE};
use axum::http::{Request, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::Router;
use reqwest::Method;
use serde::Serialize;
const DEFAULT_REQUESTS: usize = 10_000;
const DEFAULT_C1K: usize = 1_000;
const DEFAULT_C1W: usize = 10_000;
const DEFAULT_TARGETS: usize = 4;
const DEFAULT_CLIENT_SHARDS: usize = 512;
const DEFAULT_TIMEOUT_MS: u64 = 60_000;
const DEFAULT_WARMUP_REQUESTS: usize = 1_000;
const DEFAULT_WARMUP_CONCURRENCY: usize = 50;
const BODY_LIMIT: usize = 1024 * 1024;
#[derive(Debug, Clone)]
struct Config {
requests: usize,
c1k: usize,
c1w: usize,
targets: usize,
client_shards: usize,
timeout_ms: u64,
warmup_requests: usize,
warmup_concurrency: usize,
first_body_hold_ms: u64,
output_path: Option<PathBuf>,
}
impl Default for Config {
fn default() -> Self {
Self {
requests: DEFAULT_REQUESTS,
c1k: DEFAULT_C1K,
c1w: DEFAULT_C1W,
targets: DEFAULT_TARGETS,
client_shards: DEFAULT_CLIENT_SHARDS,
timeout_ms: DEFAULT_TIMEOUT_MS,
warmup_requests: DEFAULT_WARMUP_REQUESTS,
warmup_concurrency: DEFAULT_WARMUP_CONCURRENCY,
first_body_hold_ms: 0,
output_path: None,
}
}
}
#[derive(Debug, Serialize)]
struct CompareReport {
suite: &'static str,
gateway_base_url: String,
gateway_metrics_url: String,
upstream_targets: Vec<String>,
config: ReportConfig,
warmup: Option<HttpLoadProbeResult>,
c1k: RunReport,
c1w: RunReport,
mock_only: MockOnlyCompareReport,
}
#[derive(Debug, Serialize)]
struct ReportConfig {
requests: usize,
c1k: usize,
c1w: usize,
targets: usize,
client_shards: usize,
timeout_ms: u64,
warmup_requests: usize,
warmup_concurrency: usize,
first_body_hold_ms: u64,
}
#[derive(Debug, Serialize)]
struct RunReport {
label: String,
load: HttpLoadProbeResult,
gateway_metrics: GatewayMetricDelta,
mock_metrics: Vec<MockMetricDelta>,
}
#[derive(Debug, Serialize)]
struct MockOnlyCompareReport {
c1k: MultiUrlHttpLoadProbeResult,
c1w: MultiUrlHttpLoadProbeResult,
}
#[derive(Debug, Default, Serialize)]
struct GatewayMetricDelta {
stream_pre_first_byte_spawn_total: u64,
request_candidate_queue_dropped_total: u64,
raw_candidates_scanned_total: u64,
payload_build_selected_total: u64,
payload_build_prefetch_avoided_total: u64,
selected_rank_sum: u64,
model_directive_cache_hit_total: u64,
model_directive_cache_miss_total: u64,
redaction_request_cache_hit_total: u64,
redaction_request_cache_miss_total: u64,
target_raw_seen_total: BTreeMap<String, u64>,
target_preselect_total: BTreeMap<String, u64>,
target_selected_total: BTreeMap<String, u64>,
target_max_in_flight: BTreeMap<String, u64>,
target_saturated_total: BTreeMap<String, u64>,
stage_deltas: BTreeMap<String, StageMetricDelta>,
}
#[derive(Debug, Default, Serialize)]
struct StageMetricDelta {
count: u64,
sum_ms: u64,
max_ms_after: u64,
p95_bucket_ms: Option<u64>,
p99_bucket_ms: Option<u64>,
}
#[derive(Debug, Clone, Default, Serialize)]
struct MockMetricSnapshot {
requests_total: u64,
completed_total: u64,
in_flight: u64,
max_in_flight: u64,
first_chunk_yield_total: u64,
request_body_read_sum_ms: u64,
request_body_read_max_ms: u64,
response_header_to_first_chunk_sum_ms: u64,
response_header_to_first_chunk_max_ms: u64,
}
#[derive(Debug, Serialize)]
struct MockMetricDelta {
target: String,
requests_total: u64,
completed_total: u64,
in_flight_after: u64,
max_in_flight_after: u64,
first_chunk_yield_total: u64,
request_body_read_sum_ms: u64,
request_body_read_max_ms_after: u64,
response_header_to_first_chunk_sum_ms: u64,
response_header_to_first_chunk_max_ms_after: u64,
}
#[derive(Debug, Default)]
struct MockTargetMetrics {
requests_total: AtomicU64,
completed_total: AtomicU64,
in_flight: AtomicUsize,
max_in_flight: AtomicUsize,
first_chunk_yield_total: AtomicU64,
request_body_read_sum_ms: AtomicU64,
request_body_read_max_ms: AtomicU64,
response_header_to_first_chunk_sum_ms: AtomicU64,
response_header_to_first_chunk_max_ms: AtomicU64,
}
#[derive(Clone)]
struct MockTargetState {
label: String,
metrics: Arc<MockTargetMetrics>,
first_chunk_delay: Duration,
}
struct CompletionGuard {
metrics: Option<Arc<MockTargetMetrics>>,
}
impl CompletionGuard {
fn new(metrics: Arc<MockTargetMetrics>) -> Self {
Self {
metrics: Some(metrics),
}
}
}
impl Drop for CompletionGuard {
fn drop(&mut self) {
if let Some(metrics) = self.metrics.take() {
record_mock_completed(&metrics);
}
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = parse_args(std::env::args().skip(1).collect())?;
if config.targets == 0 {
return Err("targets must be positive".into());
}
let mut mock_servers = Vec::with_capacity(config.targets);
let mut mock_metrics = Vec::with_capacity(config.targets);
for index in 0..config.targets {
let metrics = Arc::new(MockTargetMetrics::default());
let state = MockTargetState {
label: format!("target-{index}"),
metrics: Arc::clone(&metrics),
first_chunk_delay: Duration::ZERO,
};
let server = SpawnedServer::start(mock_router(state)).await?;
mock_servers.push(server);
mock_metrics.push(metrics);
}
let upstream_targets = mock_servers
.iter()
.map(|server| format!("{}/v1", server.base_url()))
.collect::<Vec<_>>();
let mock_chat_urls = mock_servers
.iter()
.map(|server| format!("{}/v1/chat/completions", server.base_url()))
.collect::<Vec<_>>();
let mock_health_url = format!("{}/health", mock_servers[0].base_url());
let gateway_state = build_openai_chat_pressure_state(OpenAiChatPressureStateConfig::new(
upstream_targets.clone(),
))
.map_err(std::io::Error::other)?;
let gateway = SpawnedServer::start(build_router_with_state(gateway_state)).await?;
let gateway_url = format!("{}/v1/chat/completions", gateway.base_url());
let gateway_health_url = format!("{}/_gateway/health", gateway.base_url());
let gateway_metrics_url = format!("{}/_gateway/metrics", gateway.base_url());
let warmup = if config.warmup_requests > 0 {
Some(
run_http_load_probe(&load_config(
&gateway_url,
&gateway_health_url,
&config,
config.warmup_requests,
config.warmup_concurrency,
0,
))
.await
.map_err(std::io::Error::other)?,
)
} else {
None
};
let before_c1k = fetch_prometheus_samples(&gateway_metrics_url)
.await
.map_err(std::io::Error::other)?;
let mock_before_c1k = snapshot_mock_metrics(&mock_metrics);
let c1k_load = run_http_load_probe(&load_config(
&gateway_url,
&gateway_health_url,
&config,
config.requests,
config.c1k,
config.first_body_hold_ms,
))
.await
.map_err(std::io::Error::other)?;
let after_c1k = fetch_prometheus_samples(&gateway_metrics_url)
.await
.map_err(std::io::Error::other)?;
let mock_after_c1k = snapshot_mock_metrics(&mock_metrics);
let before_c1w = after_c1k.clone();
let mock_before_c1w = mock_after_c1k.clone();
let c1w_load = run_http_load_probe(&load_config(
&gateway_url,
&gateway_health_url,
&config,
config.requests,
config.c1w,
config.first_body_hold_ms,
))
.await
.map_err(std::io::Error::other)?;
let after_c1w = fetch_prometheus_samples(&gateway_metrics_url)
.await
.map_err(std::io::Error::other)?;
let mock_after_c1w = snapshot_mock_metrics(&mock_metrics);
let mock_only_c1k = run_multi_url_http_load_probe(
&load_config(
&mock_chat_urls[0],
&mock_health_url,
&config,
config.requests,
config.c1k,
config.first_body_hold_ms,
),
&mock_chat_urls,
)
.await
.map_err(std::io::Error::other)?;
let mock_only_c1w = run_multi_url_http_load_probe(
&load_config(
&mock_chat_urls[0],
&mock_health_url,
&config,
config.requests,
config.c1w,
config.first_body_hold_ms,
),
&mock_chat_urls,
)
.await
.map_err(std::io::Error::other)?;
let report = CompareReport {
suite: "gateway_openai_chat_c_compare",
gateway_base_url: gateway.base_url().to_string(),
gateway_metrics_url,
upstream_targets,
config: ReportConfig {
requests: config.requests,
c1k: config.c1k,
c1w: config.c1w,
targets: config.targets,
client_shards: config.client_shards,
timeout_ms: config.timeout_ms,
warmup_requests: config.warmup_requests,
warmup_concurrency: config.warmup_concurrency,
first_body_hold_ms: config.first_body_hold_ms,
},
warmup,
c1k: RunReport {
label: format!("c{}", config.c1k),
load: c1k_load,
gateway_metrics: gateway_metric_delta(&before_c1k, &after_c1k),
mock_metrics: mock_metric_delta(&mock_before_c1k, &mock_after_c1k),
},
c1w: RunReport {
label: format!("c{}", config.c1w),
load: c1w_load,
gateway_metrics: gateway_metric_delta(&before_c1w, &after_c1w),
mock_metrics: mock_metric_delta(&mock_before_c1w, &mock_after_c1w),
},
mock_only: MockOnlyCompareReport {
c1k: mock_only_c1k,
c1w: mock_only_c1w,
},
};
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(())
}
fn load_config(
gateway_url: &str,
gateway_health_url: &str,
config: &Config,
requests: usize,
concurrency: usize,
first_body_hold_ms: u64,
) -> HttpLoadProbeConfig {
let mut headers = BTreeMap::new();
headers.insert(
CONTENT_TYPE.as_str().to_string(),
"application/json".to_string(),
);
headers.insert(
AUTHORIZATION.as_str().to_string(),
"Bearer sk-aether-openai-chat-pressure".to_string(),
);
HttpLoadProbeConfig {
url: gateway_url.to_string(),
warmup_url: Some(gateway_health_url.to_string()),
method: Method::POST,
headers,
body: Some(
serde_json::json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
})
.to_string()
.into_bytes(),
),
total_requests: requests,
concurrency,
warmup_connections: 0,
timeout: Duration::from_millis(config.timeout_ms),
connect_timeout: Some(Duration::from_millis(10_000)),
response_mode: HttpLoadProbeResponseMode::FirstBodyByte,
client_shards: config.client_shards,
pool_max_idle_per_host: Some(config.requests.max(concurrency).max(1024)),
start_ramp: Duration::ZERO,
http1_only: true,
http2_prior_knowledge: false,
first_body_hold: Duration::from_millis(first_body_hold_ms),
}
}
fn mock_router(state: MockTargetState) -> Router {
Router::new()
.route("/health", get(|| async { "ok" }))
.route("/v1/chat/completions", post(mock_chat_completions))
.with_state(state)
}
async fn mock_chat_completions(
State(state): State<MockTargetState>,
request: Request<Body>,
) -> Response {
record_mock_started(&state.metrics);
let body_started_at = Instant::now();
let body = match to_bytes(request.into_body(), BODY_LIMIT).await {
Ok(body) => body,
Err(err) => {
record_mock_completed(&state.metrics);
return (
StatusCode::BAD_REQUEST,
format!("failed to read mock request body: {err}"),
)
.into_response();
}
};
record_mock_request_body_read(&state.metrics, body_started_at.elapsed());
let wants_stream = request_wants_stream(&body);
let model = request_model(&body).unwrap_or_else(|| "gpt-5-upstream".to_string());
if !wants_stream {
record_mock_completed(&state.metrics);
return axum::Json(serde_json::json!({
"id": "chatcmpl-pressure-mock",
"object": "chat.completion",
"model": model,
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop"
}]
}))
.into_response();
}
let metrics = Arc::clone(&state.metrics);
let first_chunk_delay = state.first_chunk_delay;
let label = state.label.clone();
let stream = async_stream::stream! {
let _guard = CompletionGuard::new(Arc::clone(&metrics));
let header_created_at = Instant::now();
if !first_chunk_delay.is_zero() {
tokio::time::sleep(first_chunk_delay).await;
}
record_mock_first_chunk_yield(&metrics, header_created_at.elapsed());
let chunk = serde_json::json!({
"id": format!("chatcmpl-pressure-mock-{label}"),
"object": "chat.completion.chunk",
"model": model,
"choices": [{
"index": 0,
"delta": {"role": "assistant", "content": "x"},
"finish_reason": null
}]
});
yield Ok::<Bytes, Infallible>(Bytes::from(format!("data: {chunk}\n\n")));
yield Ok::<Bytes, Infallible>(Bytes::from_static(
b"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
));
yield Ok::<Bytes, Infallible>(Bytes::from_static(b"data: [DONE]\n\n"));
};
let mut response = Response::new(Body::from_stream(stream));
response.headers_mut().insert(
CONTENT_TYPE,
"text/event-stream; charset=utf-8"
.parse()
.expect("valid content-type"),
);
response.headers_mut().insert(
CACHE_CONTROL,
"no-cache".parse().expect("valid cache-control"),
);
response
}
fn request_wants_stream(body: &[u8]) -> bool {
serde_json::from_slice::<serde_json::Value>(body)
.ok()
.and_then(|value| value.get("stream").and_then(serde_json::Value::as_bool))
.unwrap_or(false)
}
fn request_model(body: &[u8]) -> Option<String> {
serde_json::from_slice::<serde_json::Value>(body)
.ok()
.and_then(|value| {
value
.get("model")
.and_then(serde_json::Value::as_str)
.map(str::to_string)
})
}
fn record_mock_started(metrics: &MockTargetMetrics) {
metrics.requests_total.fetch_add(1, Ordering::AcqRel);
let in_flight = metrics.in_flight.fetch_add(1, Ordering::AcqRel) + 1;
metrics.max_in_flight.fetch_max(in_flight, Ordering::AcqRel);
}
fn record_mock_completed(metrics: &MockTargetMetrics) {
metrics.completed_total.fetch_add(1, Ordering::AcqRel);
metrics.in_flight.fetch_sub(1, Ordering::AcqRel);
}
fn record_mock_request_body_read(metrics: &MockTargetMetrics, elapsed: Duration) {
let elapsed_ms = elapsed.as_millis() as u64;
metrics
.request_body_read_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
metrics
.request_body_read_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
}
fn record_mock_first_chunk_yield(metrics: &MockTargetMetrics, elapsed: Duration) {
let elapsed_ms = elapsed.as_millis() as u64;
metrics
.first_chunk_yield_total
.fetch_add(1, Ordering::AcqRel);
metrics
.response_header_to_first_chunk_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
metrics
.response_header_to_first_chunk_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
}
fn snapshot_mock_metrics(metrics: &[Arc<MockTargetMetrics>]) -> Vec<MockMetricSnapshot> {
metrics
.iter()
.map(|metrics| MockMetricSnapshot {
requests_total: metrics.requests_total.load(Ordering::Acquire),
completed_total: metrics.completed_total.load(Ordering::Acquire),
in_flight: metrics.in_flight.load(Ordering::Acquire) as u64,
max_in_flight: metrics.max_in_flight.load(Ordering::Acquire) as u64,
first_chunk_yield_total: metrics.first_chunk_yield_total.load(Ordering::Acquire),
request_body_read_sum_ms: metrics.request_body_read_sum_ms.load(Ordering::Acquire),
request_body_read_max_ms: metrics.request_body_read_max_ms.load(Ordering::Acquire),
response_header_to_first_chunk_sum_ms: metrics
.response_header_to_first_chunk_sum_ms
.load(Ordering::Acquire),
response_header_to_first_chunk_max_ms: metrics
.response_header_to_first_chunk_max_ms
.load(Ordering::Acquire),
})
.collect()
}
fn mock_metric_delta(
before: &[MockMetricSnapshot],
after: &[MockMetricSnapshot],
) -> Vec<MockMetricDelta> {
after
.iter()
.enumerate()
.map(|(index, after)| {
let before = before.get(index).cloned().unwrap_or_default();
MockMetricDelta {
target: format!("target-{index}"),
requests_total: after.requests_total.saturating_sub(before.requests_total),
completed_total: after.completed_total.saturating_sub(before.completed_total),
in_flight_after: after.in_flight,
max_in_flight_after: after.max_in_flight,
first_chunk_yield_total: after
.first_chunk_yield_total
.saturating_sub(before.first_chunk_yield_total),
request_body_read_sum_ms: after
.request_body_read_sum_ms
.saturating_sub(before.request_body_read_sum_ms),
request_body_read_max_ms_after: after.request_body_read_max_ms,
response_header_to_first_chunk_sum_ms: after
.response_header_to_first_chunk_sum_ms
.saturating_sub(before.response_header_to_first_chunk_sum_ms),
response_header_to_first_chunk_max_ms_after: after
.response_header_to_first_chunk_max_ms,
}
})
.collect()
}
fn gateway_metric_delta(
before: &[PrometheusSample],
after: &[PrometheusSample],
) -> GatewayMetricDelta {
let stage_names = [
"frontdoor_handler_queue",
"frontdoor_admission",
"frontdoor_context",
"frontdoor_body_buffer",
"frontdoor_owner_forward",
"frontdoor_auth_model",
"frontdoor_rpm",
"frontdoor_rpm_system_default",
"frontdoor_rpm_runtime_check",
"frontdoor_rpm_memory_fallback",
"frontdoor_local_ai_public",
"frontdoor_execute_stream",
"frontdoor_to_stream_response_ready",
"frontdoor_to_stream_body_first_poll",
"frontdoor_to_stream_first_client_yield",
"openai_chat_stream_target_select",
"openai_chat_payload_parts_prepare",
"openai_chat_payload_model_directives",
"openai_chat_payload_redaction",
"openai_chat_payload_auth_prepare",
"openai_chat_payload_body_build",
"stream_candidate_execute",
"upstream_execution_gate_wait",
"stream_upstream_target_admission",
"stream_upstream_headers",
"direct_reqwest_request_send",
"direct_send_headers",
"direct_h2c_response_headers_wait",
"stream_first_data",
"stream_first_client_yield",
"stream_total",
];
GatewayMetricDelta {
stream_pre_first_byte_spawn_total: counter_delta(
before,
after,
"stream_pre_first_byte_spawn_total",
&[],
),
request_candidate_queue_dropped_total: counter_delta(
before,
after,
"request_candidate_queue_dropped_total",
&[],
),
raw_candidates_scanned_total: counter_delta(
before,
after,
"openai_chat_stream_target_select_raw_candidates_scanned_total",
&[],
),
payload_build_selected_total: counter_delta(
before,
after,
"openai_chat_stream_payload_build_selected_total",
&[],
),
payload_build_prefetch_avoided_total: counter_delta(
before,
after,
"openai_chat_stream_payload_build_prefetch_avoided_total",
&[],
),
selected_rank_sum: counter_delta(
before,
after,
"openai_chat_stream_target_select_selected_rank_sum",
&[],
),
model_directive_cache_hit_total: counter_delta(
before,
after,
"openai_chat_model_directive_cache_hit_total",
&[],
),
model_directive_cache_miss_total: counter_delta(
before,
after,
"openai_chat_model_directive_cache_miss_total",
&[],
),
redaction_request_cache_hit_total: counter_delta(
before,
after,
"chat_pii_redaction_request_cache_hit_total",
&[],
),
redaction_request_cache_miss_total: counter_delta(
before,
after,
"chat_pii_redaction_request_cache_miss_total",
&[],
),
target_raw_seen_total: metric_delta_by_target(
before,
after,
"upstream_target_raw_seen_total",
),
target_preselect_total: metric_delta_by_target(
before,
after,
"upstream_target_preselect_total",
),
target_selected_total: metric_delta_by_target(
before,
after,
"upstream_target_selected_total",
),
target_max_in_flight: metric_after_by_target(after, "upstream_target_max_in_flight"),
target_saturated_total: metric_delta_by_target(
before,
after,
"upstream_target_saturated_total",
),
stage_deltas: stage_names
.iter()
.map(|stage| {
(
(*stage).to_string(),
StageMetricDelta {
count: counter_delta(
before,
after,
"gateway_stage_latency_count",
&[("stage", stage)],
),
sum_ms: counter_delta(
before,
after,
"gateway_stage_latency_sum_ms",
&[("stage", stage)],
),
max_ms_after: metric_value(
after,
"gateway_stage_latency_max_ms",
&[("stage", stage)],
)
.unwrap_or_default(),
p95_bucket_ms: stage_latency_bucket_quantile(before, after, stage, 0.95),
p99_bucket_ms: stage_latency_bucket_quantile(before, after, stage, 0.99),
},
)
})
.collect(),
}
}
fn stage_latency_bucket_quantile(
before: &[PrometheusSample],
after: &[PrometheusSample],
stage: &str,
quantile: f64,
) -> Option<u64> {
let count = counter_delta(
before,
after,
"gateway_stage_latency_count",
&[("stage", stage)],
);
if count == 0 {
return None;
}
let threshold = ((count as f64) * quantile).ceil().max(1.0) as u64;
let mut buckets = after
.iter()
.filter(|sample| {
metric_name_matches(&sample.name, "gateway_stage_latency_bucket")
&& labels_match(sample, &[("stage", stage)])
})
.filter_map(|sample| {
let le_ms = sample.labels.get("le_ms")?.parse::<u64>().ok()?;
let after_value = sample.value.parse::<u64>().ok()?;
let before_value = metric_value(
before,
"gateway_stage_latency_bucket",
&[("stage", stage), ("le_ms", sample.labels.get("le_ms")?)],
)
.unwrap_or_default();
Some((le_ms, after_value.saturating_sub(before_value)))
})
.collect::<Vec<_>>();
buckets.sort_unstable_by_key(|(le_ms, _)| *le_ms);
buckets
.into_iter()
.find_map(|(le_ms, value)| (value >= threshold).then_some(le_ms))
}
fn counter_delta(
before: &[PrometheusSample],
after: &[PrometheusSample],
metric_name: &str,
labels: &[(&str, &str)],
) -> u64 {
let before = metric_value(before, metric_name, labels).unwrap_or_default();
let after = metric_value(after, metric_name, labels).unwrap_or_default();
after.saturating_sub(before)
}
fn metric_value(
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_delta_by_target(
before: &[PrometheusSample],
after: &[PrometheusSample],
metric_name: &str,
) -> BTreeMap<String, u64> {
let mut values = BTreeMap::new();
for sample in after
.iter()
.filter(|sample| metric_name_matches(&sample.name, metric_name))
{
let Some(target) = sample.labels.get("target") else {
continue;
};
let before_value =
metric_value(before, metric_name, &[("target", target)]).unwrap_or_default();
let after_value = sample.value.parse::<u64>().unwrap_or_default();
values.insert(target.clone(), after_value.saturating_sub(before_value));
}
values
}
fn metric_after_by_target(
samples: &[PrometheusSample],
metric_name: &str,
) -> BTreeMap<String, u64> {
samples
.iter()
.filter(|sample| metric_name_matches(&sample.name, metric_name))
.filter_map(|sample| {
Some((
sample.labels.get("target")?.clone(),
sample.value.parse::<u64>().ok()?,
))
})
.collect()
}
fn metric_name_matches(actual: &str, expected: &str) -> bool {
actual == expected || actual.ends_with(&format!("_{expected}"))
}
fn labels_match(sample: &PrometheusSample, labels: &[(&str, &str)]) -> bool {
labels
.iter()
.all(|(key, value)| sample.labels.get(*key).map(String::as_str) == Some(*value))
}
fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
let mut config = Config::default();
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
match arg.as_str() {
"--requests" => config.requests = next_value(&mut iter, "--requests")?.parse()?,
"--c1k" => config.c1k = next_value(&mut iter, "--c1k")?.parse()?,
"--c1w" => config.c1w = next_value(&mut iter, "--c1w")?.parse()?,
"--targets" => config.targets = next_value(&mut iter, "--targets")?.parse()?,
"--client-shards" => {
config.client_shards = next_value(&mut iter, "--client-shards")?.parse()?
}
"--timeout-ms" => config.timeout_ms = next_value(&mut iter, "--timeout-ms")?.parse()?,
"--warmup-requests" => {
config.warmup_requests = next_value(&mut iter, "--warmup-requests")?.parse()?
}
"--warmup-concurrency" => {
config.warmup_concurrency =
next_value(&mut iter, "--warmup-concurrency")?.parse()?
}
"--first-body-hold-ms" => {
config.first_body_hold_ms =
next_value(&mut iter, "--first-body-hold-ms")?.parse()?
}
"--output" => {
config.output_path = Some(PathBuf::from(next_value(&mut iter, "--output")?))
}
"--help" | "-h" => {
print_help();
std::process::exit(0);
}
other => return Err(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(|| format!("missing value for {flag}").into())
}
fn print_help() {
println!(
"gateway_openai_chat_c_compare \
[--requests N] [--c1k N] [--c1w N] [--targets N] \
[--client-shards N] [--timeout-ms N] \
[--warmup-requests N] [--warmup-concurrency N] \
[--first-body-hold-ms N] [--output PATH]"
);
}
crates/aether-testkit/src/bin/gateway_tunnel_stream_baseline.rs
+1
View File
@@ -152,6 +152,7 @@ async fn run_suite(
concurrency: config.concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
})
.await
.map_err(std::io::Error::other)?;
crates/aether-testkit/src/bin/mock_openai_upstream.rs
+345 -45
View File
@@ -1,10 +1,11 @@
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::time::{Duration, Instant};
use axum::body::{Body, Bytes};
use axum::body::{to_bytes, Body, Bytes};
use axum::extract::State;
use axum::http::{header, HeaderValue, StatusCode};
use axum::response::{IntoResponse, Response};
@@ -20,8 +21,11 @@ struct Config {
chunk_delay: Duration,
payload_bytes: usize,
status: StatusCode,
assume_stream: bool,
}
const MAX_MOCK_REQUEST_BODY_BYTES: usize = 1024 * 1024;
impl Default for Config {
fn default() -> Self {
Self {
@@ -33,6 +37,7 @@ impl Default for Config {
chunk_delay: Duration::from_millis(20),
payload_bytes: 32,
status: StatusCode::OK,
assume_stream: false,
}
}
}
@@ -43,41 +48,85 @@ struct Metrics {
completed_total: AtomicU64,
in_flight: AtomicU64,
max_in_flight: AtomicU64,
request_body_read_sum_ms: AtomicU64,
request_body_read_max_ms: AtomicU64,
accepted_to_response_header_sum_ms: AtomicU64,
accepted_to_response_header_max_ms: AtomicU64,
first_chunk_yield_total: AtomicU64,
response_header_to_first_chunk_sum_ms: AtomicU64,
response_header_to_first_chunk_max_ms: AtomicU64,
binds: BTreeMap<String, Arc<BindMetrics>>,
}
#[derive(Debug, Default)]
struct BindMetrics {
requests_total: AtomicU64,
completed_total: AtomicU64,
in_flight: AtomicU64,
max_in_flight: AtomicU64,
request_body_read_sum_ms: AtomicU64,
request_body_read_max_ms: AtomicU64,
accepted_to_response_header_sum_ms: AtomicU64,
accepted_to_response_header_max_ms: AtomicU64,
first_chunk_yield_total: AtomicU64,
response_header_to_first_chunk_sum_ms: AtomicU64,
response_header_to_first_chunk_max_ms: AtomicU64,
}
#[derive(Debug, Clone)]
struct App {
config: Config,
metrics: Arc<Metrics>,
bind_label: Arc<str>,
}
struct RequestCompletionGuard {
app: App,
completed: bool,
}
impl RequestCompletionGuard {
fn new(app: App) -> Self {
Self {
app,
completed: false,
}
}
fn complete(mut self) {
self.complete_once();
}
fn complete_once(&mut self) {
if self.completed {
return;
}
record_request_completed(&self.app);
self.completed = true;
}
}
impl Drop for RequestCompletionGuard {
fn drop(&mut self) {
self.complete_once();
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = parse_args(std::env::args().skip(1).collect())?;
let app_state = App {
config: config.clone(),
metrics: Arc::new(Metrics::default()),
};
let metrics = Arc::new(Metrics::for_binds(&config.binds));
let app = Router::new()
.route("/health", get(health))
.route("/metrics", get(metrics))
.route("/v1/chat/completions", post(chat_completions))
.route("/chat/completions", post(chat_completions))
.route("/v1/responses", post(responses))
.route("/responses", post(responses))
.with_state(app_state);
serve_listeners(&config.binds, app).await?;
serve_listeners(&config, metrics).await?;
Ok(())
}
async fn serve_listeners(
binds: &[SocketAddr],
app: Router,
config: &Config,
metrics: Arc<Metrics>,
) -> Result<(), Box<dyn std::error::Error>> {
let mut listeners = Vec::with_capacity(binds.len());
for bind in binds {
let mut listeners = Vec::with_capacity(config.binds.len());
for bind in &config.binds {
listeners.push((*bind, tokio::net::TcpListener::bind(bind).await?));
}
if listeners.len() == 1 {
@@ -85,6 +134,7 @@ async fn serve_listeners(
.into_iter()
.next()
.ok_or_else(|| std::io::Error::other("mock upstream listener set is empty"))?;
let app = build_router(config.clone(), Arc::clone(&metrics), bind);
eprintln!("mock OpenAI upstream listening on http://{bind}");
axum::serve(listener, app).await?;
return Ok(());
@@ -92,7 +142,7 @@ async fn serve_listeners(
let mut servers = tokio::task::JoinSet::new();
for (bind, listener) in listeners {
let app = app.clone();
let app = build_router(config.clone(), Arc::clone(&metrics), bind);
eprintln!("mock OpenAI upstream listening on http://{bind}");
servers.spawn(async move { axum::serve(listener, app).await });
}
@@ -105,12 +155,43 @@ async fn serve_listeners(
Ok(())
}
fn build_router(config: Config, shared_metrics: Arc<Metrics>, bind: SocketAddr) -> Router {
Router::new()
.route("/health", get(health))
.route("/metrics", get(metrics))
.route("/v1/chat/completions", post(chat_completions))
.route("/chat/completions", post(chat_completions))
.route("/v1/responses", post(responses))
.route("/responses", post(responses))
.with_state(App {
config,
metrics: shared_metrics,
bind_label: Arc::from(bind.to_string()),
})
}
impl Metrics {
fn for_binds(binds: &[SocketAddr]) -> Self {
let mut metrics = Self::default();
for bind in binds {
metrics
.binds
.insert(bind.to_string(), Arc::new(BindMetrics::default()));
}
metrics
}
fn bind_metrics(&self, bind: &str) -> Option<Arc<BindMetrics>> {
self.binds.get(bind).cloned()
}
}
async fn health() -> impl IntoResponse {
(StatusCode::OK, "ok\n")
}
async fn metrics(State(app): State<App>) -> Response {
let body = format!(
let mut body = format!(
concat!(
"# HELP mock_upstream_requests_total Total requests accepted by the mock upstream.\n",
"# TYPE mock_upstream_requests_total counter\n",
@@ -123,24 +204,127 @@ async fn metrics(State(app): State<App>) -> Response {
"mock_upstream_in_flight {}\n",
"# HELP mock_upstream_max_in_flight Maximum in-flight requests/streams observed.\n",
"# TYPE mock_upstream_max_in_flight gauge\n",
"mock_upstream_max_in_flight {}\n"
"mock_upstream_max_in_flight {}\n",
"# HELP mock_upstream_first_chunk_yield_total Total first stream chunks yielded by the mock upstream.\n",
"# TYPE mock_upstream_first_chunk_yield_total counter\n",
"mock_upstream_first_chunk_yield_total {}\n",
"# HELP mock_upstream_request_body_read_sum_ms Total milliseconds spent reading request bodies after handler entry.\n",
"# TYPE mock_upstream_request_body_read_sum_ms counter\n",
"mock_upstream_request_body_read_sum_ms {}\n",
"# HELP mock_upstream_request_body_read_max_ms Maximum milliseconds spent reading a request body after handler entry.\n",
"# TYPE mock_upstream_request_body_read_max_ms gauge\n",
"mock_upstream_request_body_read_max_ms {}\n",
"# HELP mock_upstream_accepted_to_response_header_sum_ms Total milliseconds from request acceptance to response construction.\n",
"# TYPE mock_upstream_accepted_to_response_header_sum_ms counter\n",
"mock_upstream_accepted_to_response_header_sum_ms {}\n",
"# HELP mock_upstream_accepted_to_response_header_max_ms Maximum milliseconds from request acceptance to response construction.\n",
"# TYPE mock_upstream_accepted_to_response_header_max_ms gauge\n",
"mock_upstream_accepted_to_response_header_max_ms {}\n",
"# HELP mock_upstream_response_header_to_first_chunk_sum_ms Total milliseconds from response construction to first stream chunk yield.\n",
"# TYPE mock_upstream_response_header_to_first_chunk_sum_ms counter\n",
"mock_upstream_response_header_to_first_chunk_sum_ms {}\n",
"# HELP mock_upstream_response_header_to_first_chunk_max_ms Maximum milliseconds from response construction to first stream chunk yield.\n",
"# TYPE mock_upstream_response_header_to_first_chunk_max_ms gauge\n",
"mock_upstream_response_header_to_first_chunk_max_ms {}\n"
),
app.metrics.requests_total.load(Ordering::Acquire),
app.metrics.completed_total.load(Ordering::Acquire),
app.metrics.in_flight.load(Ordering::Acquire),
app.metrics.max_in_flight.load(Ordering::Acquire),
app.metrics.first_chunk_yield_total.load(Ordering::Acquire),
app.metrics
.request_body_read_sum_ms
.load(Ordering::Acquire),
app.metrics
.request_body_read_max_ms
.load(Ordering::Acquire),
app.metrics
.accepted_to_response_header_sum_ms
.load(Ordering::Acquire),
app.metrics
.accepted_to_response_header_max_ms
.load(Ordering::Acquire),
app.metrics
.response_header_to_first_chunk_sum_ms
.load(Ordering::Acquire),
app.metrics
.response_header_to_first_chunk_max_ms
.load(Ordering::Acquire),
);
for (bind, metrics) in &app.metrics.binds {
let bind = prometheus_label_value(bind);
body.push_str(&format!(
concat!(
"mock_upstream_requests_total{{bind=\"{}\"}} {}\n",
"mock_upstream_completed_total{{bind=\"{}\"}} {}\n",
"mock_upstream_in_flight{{bind=\"{}\"}} {}\n",
"mock_upstream_max_in_flight{{bind=\"{}\"}} {}\n",
"mock_upstream_first_chunk_yield_total{{bind=\"{}\"}} {}\n",
"mock_upstream_request_body_read_sum_ms{{bind=\"{}\"}} {}\n",
"mock_upstream_request_body_read_max_ms{{bind=\"{}\"}} {}\n",
"mock_upstream_accepted_to_response_header_sum_ms{{bind=\"{}\"}} {}\n",
"mock_upstream_accepted_to_response_header_max_ms{{bind=\"{}\"}} {}\n",
"mock_upstream_response_header_to_first_chunk_sum_ms{{bind=\"{}\"}} {}\n",
"mock_upstream_response_header_to_first_chunk_max_ms{{bind=\"{}\"}} {}\n"
),
bind,
metrics.requests_total.load(Ordering::Acquire),
bind,
metrics.completed_total.load(Ordering::Acquire),
bind,
metrics.in_flight.load(Ordering::Acquire),
bind,
metrics.max_in_flight.load(Ordering::Acquire),
bind,
metrics.first_chunk_yield_total.load(Ordering::Acquire),
bind,
metrics.request_body_read_sum_ms.load(Ordering::Acquire),
bind,
metrics.request_body_read_max_ms.load(Ordering::Acquire),
bind,
metrics
.accepted_to_response_header_sum_ms
.load(Ordering::Acquire),
bind,
metrics
.accepted_to_response_header_max_ms
.load(Ordering::Acquire),
bind,
metrics
.response_header_to_first_chunk_sum_ms
.load(Ordering::Acquire),
bind,
metrics
.response_header_to_first_chunk_max_ms
.load(Ordering::Acquire),
));
}
(StatusCode::OK, body).into_response()
}
async fn chat_completions(State(app): State<App>, body: Bytes) -> Response {
let stream = request_wants_stream(&body);
record_request_started(&app.metrics);
async fn chat_completions(State(app): State<App>, request: axum::extract::Request) -> Response {
let request_started_at = record_request_started(&app);
if app.config.status != StatusCode::OK {
record_request_completed(&app.metrics);
record_response_header_created(&app, request_started_at.elapsed());
record_request_completed(&app);
return (app.config.status, "mock upstream error\n").into_response();
}
if app.config.assume_stream {
record_response_header_created(&app, request_started_at.elapsed());
return build_chat_sse_response(app);
}
let body = match read_request_body(&app, request).await {
Ok(body) => body,
Err(response) => {
record_response_header_created(&app, request_started_at.elapsed());
record_request_completed(&app);
return response;
}
};
let stream = request_wants_stream(&body);
if stream {
record_response_header_created(&app, request_started_at.elapsed());
return build_chat_sse_response(app);
}
let payload = json!({
@@ -162,18 +346,34 @@ async fn chat_completions(State(app): State<App>, body: Bytes) -> Response {
"total_tokens": app.config.chunks.max(1) + 1
}
});
record_request_completed(&app.metrics);
record_response_header_created(&app, request_started_at.elapsed());
record_request_completed(&app);
axum::Json(payload).into_response()
}
async fn responses(State(app): State<App>, body: Bytes) -> Response {
let stream = request_wants_stream(&body);
record_request_started(&app.metrics);
async fn responses(State(app): State<App>, request: axum::extract::Request) -> Response {
let request_started_at = record_request_started(&app);
if app.config.status != StatusCode::OK {
record_request_completed(&app.metrics);
record_response_header_created(&app, request_started_at.elapsed());
record_request_completed(&app);
return (app.config.status, "mock upstream error\n").into_response();
}
if app.config.assume_stream {
record_response_header_created(&app, request_started_at.elapsed());
return build_responses_sse_response(app);
}
let body = match read_request_body(&app, request).await {
Ok(body) => body,
Err(response) => {
record_response_header_created(&app, request_started_at.elapsed());
record_request_completed(&app);
return response;
}
};
let stream = request_wants_stream(&body);
if stream {
record_response_header_created(&app, request_started_at.elapsed());
return build_responses_sse_response(app);
}
let payload = json!({
@@ -196,13 +396,15 @@ async fn responses(State(app): State<App>, body: Bytes) -> Response {
"total_tokens": app.config.chunks.max(1) + 1
}
});
record_request_completed(&app.metrics);
record_response_header_created(&app, request_started_at.elapsed());
record_request_completed(&app);
axum::Json(payload).into_response()
}
fn build_chat_sse_response(app: App) -> Response {
let metrics = Arc::clone(&app.metrics);
let response_created_at = Instant::now();
let config = app.config.clone();
let mut completion = Some(RequestCompletionGuard::new(app.clone()));
let stream = async_stream::stream! {
if !config.first_byte_delay.is_zero() {
tokio::time::sleep(config.first_byte_delay).await;
@@ -224,24 +426,31 @@ fn build_chat_sse_response(app: App) -> Response {
"finish_reason": serde_json::Value::Null
}]
});
if index == 0 {
record_first_chunk_yield(&app, response_created_at.elapsed());
}
yield Ok::<Bytes, Infallible>(Bytes::from(format!("data: {payload}\n\n")));
}
yield Ok::<Bytes, Infallible>(Bytes::from(
"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
));
yield Ok::<Bytes, Infallible>(Bytes::from("data: [DONE]\n\n"));
record_request_completed(&metrics);
if let Some(completion) = completion.take() {
completion.complete();
}
};
sse_response(Body::from_stream(stream))
}
fn build_responses_sse_response(app: App) -> Response {
let metrics = Arc::clone(&app.metrics);
let response_created_at = Instant::now();
let config = app.config.clone();
let mut completion = Some(RequestCompletionGuard::new(app.clone()));
let stream = async_stream::stream! {
if !config.first_byte_delay.is_zero() {
tokio::time::sleep(config.first_byte_delay).await;
}
record_first_chunk_yield(&app, response_created_at.elapsed());
yield Ok::<Bytes, Infallible>(Bytes::from(
"event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_mock\",\"status\":\"in_progress\"}}\n\n",
));
@@ -262,7 +471,9 @@ fn build_responses_sse_response(app: App) -> Response {
"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_mock\",\"status\":\"completed\"}}\n\n",
));
yield Ok::<Bytes, Infallible>(Bytes::from("data: [DONE]\n\n"));
record_request_completed(&metrics);
if let Some(completion) = completion.take() {
completion.complete();
}
};
sse_response(Body::from_stream(stream))
}
@@ -278,6 +489,21 @@ fn sse_response(body: Body) -> Response {
response
}
async fn read_request_body(app: &App, request: axum::extract::Request) -> Result<Bytes, Response> {
let body_started_at = Instant::now();
let body = to_bytes(request.into_body(), MAX_MOCK_REQUEST_BODY_BYTES)
.await
.map_err(|err| {
(
StatusCode::BAD_REQUEST,
format!("mock upstream failed to read request body: {err}\n"),
)
.into_response()
});
record_request_body_read(app, body_started_at.elapsed());
body
}
fn request_wants_stream(body: &[u8]) -> bool {
serde_json::from_slice::<serde_json::Value>(body)
.ok()
@@ -303,15 +529,86 @@ fn mock_payload(bytes: usize) -> String {
"x".repeat(bytes)
}
fn record_request_started(metrics: &Metrics) {
metrics.requests_total.fetch_add(1, Ordering::AcqRel);
let in_flight = metrics.in_flight.fetch_add(1, Ordering::AcqRel) + 1;
metrics.max_in_flight.fetch_max(in_flight, Ordering::AcqRel);
fn record_request_started(app: &App) -> Instant {
app.metrics.requests_total.fetch_add(1, Ordering::AcqRel);
let in_flight = app.metrics.in_flight.fetch_add(1, Ordering::AcqRel) + 1;
app.metrics
.max_in_flight
.fetch_max(in_flight, Ordering::AcqRel);
if let Some(bind) = app.metrics.bind_metrics(app.bind_label.as_ref()) {
bind.requests_total.fetch_add(1, Ordering::AcqRel);
let in_flight = bind.in_flight.fetch_add(1, Ordering::AcqRel) + 1;
bind.max_in_flight.fetch_max(in_flight, Ordering::AcqRel);
}
Instant::now()
}
fn record_request_completed(metrics: &Metrics) {
metrics.completed_total.fetch_add(1, Ordering::AcqRel);
metrics.in_flight.fetch_sub(1, Ordering::AcqRel);
fn record_request_completed(app: &App) {
app.metrics.completed_total.fetch_add(1, Ordering::AcqRel);
app.metrics.in_flight.fetch_sub(1, Ordering::AcqRel);
if let Some(bind) = app.metrics.bind_metrics(app.bind_label.as_ref()) {
bind.completed_total.fetch_add(1, Ordering::AcqRel);
bind.in_flight.fetch_sub(1, Ordering::AcqRel);
}
}
fn record_first_chunk_yield(app: &App, elapsed: Duration) {
let elapsed_ms = elapsed.as_millis() as u64;
app.metrics
.first_chunk_yield_total
.fetch_add(1, Ordering::AcqRel);
app.metrics
.response_header_to_first_chunk_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
app.metrics
.response_header_to_first_chunk_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
if let Some(bind) = app.metrics.bind_metrics(app.bind_label.as_ref()) {
bind.first_chunk_yield_total.fetch_add(1, Ordering::AcqRel);
bind.response_header_to_first_chunk_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
bind.response_header_to_first_chunk_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
}
}
fn record_request_body_read(app: &App, elapsed: Duration) {
let elapsed_ms = elapsed.as_millis() as u64;
app.metrics
.request_body_read_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
app.metrics
.request_body_read_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
if let Some(bind) = app.metrics.bind_metrics(app.bind_label.as_ref()) {
bind.request_body_read_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
bind.request_body_read_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
}
}
fn record_response_header_created(app: &App, elapsed: Duration) {
let elapsed_ms = elapsed.as_millis() as u64;
app.metrics
.accepted_to_response_header_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
app.metrics
.accepted_to_response_header_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
if let Some(bind) = app.metrics.bind_metrics(app.bind_label.as_ref()) {
bind.accepted_to_response_header_sum_ms
.fetch_add(elapsed_ms, Ordering::AcqRel);
bind.accepted_to_response_header_max_ms
.fetch_max(elapsed_ms, Ordering::AcqRel);
}
}
fn prometheus_label_value(value: &str) -> String {
value
.replace('\\', r"\\")
.replace('\n', r"\n")
.replace('"', r#"\""#)
}
fn current_unix_secs() -> u64 {
@@ -350,6 +647,9 @@ fn parse_args(args: Vec<String>) -> Result<Config, Box<dyn std::error::Error>> {
let status = next_value(&mut iter, "--status")?.parse::<u16>()?;
config.status = StatusCode::from_u16(status)?;
}
"--assume-stream" => {
config.assume_stream = true;
}
"--help" | "-h" => {
print_usage();
std::process::exit(0);
@@ -388,6 +688,6 @@ fn next_value(
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin mock_openai_upstream -- [--bind 127.0.0.1:18181]... [--chunks 8] [--first-byte-delay-ms 0] [--chunk-delay-ms 20] [--payload-bytes 32] [--status 200]"
"usage: cargo run -p aether-testkit --bin mock_openai_upstream -- [--bind 127.0.0.1:18181]... [--chunks 8] [--first-byte-delay-ms 0] [--chunk-delay-ms 20] [--payload-bytes 32] [--status 200] [--assume-stream]"
);
}
crates/aether-testkit/src/bin/multi_instance_admission_baseline.rs
+2
View File
@@ -352,6 +352,7 @@ fn execution_runtime_sync_probe_config(
concurrency: config.execution_runtime_concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
}
}
@@ -377,6 +378,7 @@ fn chat_probe_config(
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
}
}
crates/aether-testkit/src/bin/multi_instance_owner_relay_baseline.rs
+2
View File
@@ -164,6 +164,7 @@ async fn run_suite(
concurrency: config.concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
})
.await
.map_err(std::io::Error::other)?;
@@ -180,6 +181,7 @@ async fn run_suite(
concurrency: config.concurrency,
timeout: config.timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
})
.await
.map_err(std::io::Error::other)?;
crates/aether-testkit/src/bin/single_instance_baseline.rs
+2
View File
@@ -197,6 +197,7 @@ fn execution_probe_config(
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
}
}
@@ -223,6 +224,7 @@ fn chat_probe_config(
concurrency,
timeout,
response_mode: HttpLoadProbeResponseMode::FullBody,
..HttpLoadProbeConfig::default()
}
}
crates/aether-testkit/src/load.rs
+85 -5
View File
@@ -11,6 +11,14 @@ use tokio::sync::Mutex;
use crate::runtime::{BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot};
const MAX_ERROR_SAMPLES: usize = 32;
const FIRST_BODY_BACKGROUND_DRAIN_CHUNKS_ENV: &str =
"AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS";
const FIRST_BODY_BACKGROUND_DRAIN_MS_ENV: &str = "AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_MS";
const FIRST_BODY_AUTO_CLIENT_SHARDS_MAX_ENV: &str =
"AETHER_TESTKIT_FIRST_BODY_AUTO_CLIENT_SHARDS_MAX";
const DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS: usize = 64;
const DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_MS: u64 = 100;
const DEFAULT_FIRST_BODY_AUTO_CLIENT_SHARDS_MAX: usize = 512;
#[derive(Debug, Clone, Copy, Default, serde::Serialize, PartialEq, Eq)]
pub enum HttpLoadProbeResponseMode {
@@ -263,7 +271,8 @@ async fn run_http_load_probe_against_urls(
config: &HttpLoadProbeConfig,
urls: &[String],
) -> Result<MultiUrlHttpLoadProbeResult, String> {
let clients = Arc::new(build_probe_clients(config)?);
let effective_client_shards = effective_probe_client_shards(config);
let clients = Arc::new(build_probe_clients(config, effective_client_shards)?);
let total_requests = config.total_requests;
let request_headers = build_headers(&config.headers)?;
let request_body = config.body.clone().map(Arc::new);
@@ -423,7 +432,7 @@ async fn run_http_load_probe_against_urls(
concurrency: config.concurrency,
warmup_connections: config.warmup_connections,
warmup_url: config.warmup_url.clone(),
client_shards: config.client_shards,
client_shards: effective_client_shards,
start_ramp_ms: config.start_ramp.as_millis() as u64,
pool_max_idle_per_host: config.pool_max_idle_per_host,
http1_only: config.http1_only,
@@ -454,9 +463,29 @@ async fn run_http_load_probe_against_urls(
})
}
fn build_probe_clients(config: &HttpLoadProbeConfig) -> Result<Vec<Client>, String> {
let mut clients = Vec::with_capacity(config.client_shards);
for _ in 0..config.client_shards {
fn effective_probe_client_shards(config: &HttpLoadProbeConfig) -> usize {
let configured = config.client_shards.max(1);
if configured != 1
|| config.http1_only
|| !config.http2_prior_knowledge
|| config.response_mode != HttpLoadProbeResponseMode::FirstBodyByte
{
return configured;
}
let max_auto = env_usize(
FIRST_BODY_AUTO_CLIENT_SHARDS_MAX_ENV,
DEFAULT_FIRST_BODY_AUTO_CLIENT_SHARDS_MAX,
)
.max(1);
config.concurrency.max(1).min(max_auto)
}
fn build_probe_clients(
config: &HttpLoadProbeConfig,
client_shards: usize,
) -> Result<Vec<Client>, String> {
let mut clients = Vec::with_capacity(client_shards);
for _ in 0..client_shards {
let mut builder = Client::builder().timeout(config.timeout);
if let Some(connect_timeout) = config.connect_timeout {
builder = builder.connect_timeout(connect_timeout);
@@ -664,6 +693,7 @@ async fn observe_response_body(
if !first_body_hold.is_zero() {
tokio::time::sleep(first_body_hold).await;
}
drain_first_body_response_tail(response).await?;
Ok(BodyObservation {
first_body_latency_ms: Some(first_body_latency_ms),
})
@@ -693,6 +723,56 @@ async fn observe_response_body(
}
}
async fn drain_first_body_response_tail(
mut response: reqwest::Response,
) -> Result<(), ClassifiedLoadError> {
let max_chunks = env_usize(
FIRST_BODY_BACKGROUND_DRAIN_CHUNKS_ENV,
DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS,
);
if max_chunks == 0 {
return Ok(());
}
let timeout = Duration::from_millis(env_u64(
FIRST_BODY_BACKGROUND_DRAIN_MS_ENV,
DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_MS,
));
let drain = async {
for _ in 0..max_chunks {
let Some(chunk) = response
.chunk()
.await
.map_err(|err| classify_reqwest_error("body", &err))?
else {
break;
};
drop(chunk);
}
Ok(())
};
if timeout.is_zero() {
return drain.await;
}
match tokio::time::timeout(timeout, drain).await {
Ok(result) => result,
Err(_) => Ok(()),
}
}
fn env_usize(key: &str, default_value: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(default_value)
}
fn env_u64(key: &str, default_value: u64) -> u64 {
std::env::var(key)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(default_value)
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
struct BodyObservation {
first_body_latency_ms: Option<u64>,