Refactor tunnel stability protocol

This commit is contained in:
elky
2026-06-01 01:36:49 +08:00
parent 392353ffff
commit 37413c0211
21 changed files with 1960 additions and 132 deletions
@@ -1,11 +1,12 @@
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::time::Duration;
use aether_gateway::tunnel_protocol as protocol;
use aether_testkit::{
fetch_prometheus_samples, find_metric_value_u64, init_test_runtime_for, run_http_load_probe,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, TunnelHarness,
TunnelHarnessConfig,
HttpLoadProbeConfig, HttpLoadProbeResponseMode, HttpLoadProbeResult, PrometheusSample,
TunnelHarness, TunnelHarnessConfig,
};
use futures_util::{SinkExt, StreamExt};
use reqwest::Method;
@@ -20,6 +21,11 @@ const TUNNEL_RELAY_PATH_PREFIX: &str = "/api/internal/tunnel/relay";
struct GatewayTunnelBaselineConfig {
total_requests: usize,
concurrency: usize,
request_body_bytes: usize,
tunnel_connections: usize,
outbound_queue_capacity: usize,
close_one_connection_after: Option<Duration>,
require_acceptance: bool,
timeout: Duration,
output_path: Option<PathBuf>,
}
@@ -29,6 +35,11 @@ impl Default for GatewayTunnelBaselineConfig {
Self {
total_requests: 200,
concurrency: 20,
request_body_bytes: 6 * 1024 * 1024,
tunnel_connections: 4,
outbound_queue_capacity: 512,
close_one_connection_after: None,
require_acceptance: false,
timeout: Duration::from_secs(10),
output_path: None,
}
@@ -38,8 +49,21 @@ impl Default for GatewayTunnelBaselineConfig {
#[derive(Debug, Serialize)]
struct GatewayTunnelBaselineReport {
suite: &'static str,
config: GatewayTunnelEffectiveConfig,
scenario: HttpLoadProbeResult,
tunnel_metrics: TunnelMetricsSnapshot,
acceptance: AcceptanceReport,
}
#[derive(Debug, Serialize)]
struct GatewayTunnelEffectiveConfig {
total_requests: usize,
concurrency: usize,
request_body_bytes: usize,
tunnel_connections: usize,
outbound_queue_capacity: usize,
close_one_connection_after_ms: Option<u64>,
timeout_ms: u64,
}
#[derive(Debug, Serialize)]
@@ -54,8 +78,28 @@ struct TunnelMetricsSnapshot {
proxy_connection_congested_total: u64,
proxy_connection_write_latency_last_us_max: u64,
proxy_connection_write_latency_ewma_us_max: u64,
body_backpressure_total: u64,
flow_window_blocked_ms: u64,
connection_health_score: u64,
stream_reset_total: u64,
stream_reset_reasons: BTreeMap<String, u64>,
drain_total: u64,
drain_reasons: BTreeMap<String, u64>,
scheduler_selected_conn_total: u64,
proxy_connections_protocol_v1: u64,
proxy_connections_protocol_v2: u64,
proxy_connections_protocol_v3: u64,
}
#[derive(Debug, Serialize)]
struct AcceptanceReport {
required: bool,
passed: bool,
success_rate_bps: u64,
min_success_rate_bps: u64,
congestion_free: bool,
no_queue_full_rejections: bool,
reasons: Vec<String>,
}
#[tokio::main]
@@ -77,8 +121,21 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
async fn run_suite(
config: &GatewayTunnelBaselineConfig,
) -> Result<GatewayTunnelBaselineReport, Box<dyn std::error::Error>> {
let tunnel = TunnelHarness::start(TunnelHarnessConfig::default()).await?;
let peer = connect_protocol_peer(tunnel.base_url()).await?;
let tunnel = TunnelHarness::start(TunnelHarnessConfig {
outbound_queue_capacity: config.outbound_queue_capacity,
..TunnelHarnessConfig::default()
})
.await?;
let mut peers = connect_protocol_peers(tunnel.base_url(), config.tunnel_connections).await?;
let fault_injection = config.close_one_connection_after.map(|delay| {
let peer = peers.pop();
tokio::spawn(async move {
if let Some(peer) = peer {
tokio::time::sleep(delay).await;
peer.abort();
}
})
});
let result = run_http_load_probe(&HttpLoadProbeConfig {
url: format!(
@@ -90,7 +147,7 @@ async fn run_suite(
"content-type".to_string(),
"application/octet-stream".to_string(),
)]),
body: Some(relay_envelope()),
body: Some(relay_envelope(config.request_body_bytes)),
total_requests: config.total_requests,
concurrency: config.concurrency,
timeout: config.timeout,
@@ -100,16 +157,39 @@ async fn run_suite(
.map_err(std::io::Error::other)?;
let tunnel_metrics = capture_tunnel_metrics(tunnel.base_url()).await?;
drop(peer);
if let Some(task) = fault_injection {
let _ = task.await;
}
drop(peers);
let acceptance = evaluate_acceptance(config, &result, &tunnel_metrics);
if config.require_acceptance && !acceptance.passed {
return Err(std::io::Error::other(format!(
"gateway tunnel acceptance failed: {:?}",
acceptance.reasons
))
.into());
}
Ok(GatewayTunnelBaselineReport {
suite: "gateway_tunnel_stream_baseline",
config: GatewayTunnelEffectiveConfig {
total_requests: config.total_requests,
concurrency: config.concurrency,
request_body_bytes: config.request_body_bytes,
tunnel_connections: config.tunnel_connections,
outbound_queue_capacity: config.outbound_queue_capacity,
close_one_connection_after_ms: config
.close_one_connection_after
.map(|duration| duration.as_millis() as u64),
timeout_ms: config.timeout.as_millis() as u64,
},
scenario: result,
tunnel_metrics,
acceptance,
})
}
fn relay_envelope() -> Vec<u8> {
fn relay_envelope(request_body_bytes: usize) -> Vec<u8> {
let meta = protocol::RequestMeta {
method: "POST".to_string(),
url: "https://baseline.example/v1/chat/completions".to_string(),
@@ -129,16 +209,29 @@ fn relay_envelope() -> Vec<u8> {
transport_profile: None,
};
let meta_json = serde_json::to_vec(&meta).expect("tunnel relay metadata should serialize");
let body = br#"{"model":"gpt-5","messages":[{"role":"user","content":"hello"}]}"#;
let body = vec![b'x'; request_body_bytes];
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.extend_from_slice(&body);
envelope
}
async fn connect_protocol_peers(
tunnel_base_url: &str,
count: usize,
) -> Result<Vec<tokio::task::JoinHandle<()>>, Box<dyn std::error::Error>> {
let count = count.max(1);
let mut peers = Vec::with_capacity(count);
for index in 0..count {
peers.push(connect_protocol_peer(tunnel_base_url, index).await?);
}
Ok(peers)
}
async fn connect_protocol_peer(
tunnel_base_url: &str,
index: usize,
) -> Result<tokio::task::JoinHandle<()>, Box<dyn std::error::Error>> {
let ws_url = format!(
"{}{}",
@@ -158,7 +251,7 @@ async fn connect_protocol_peer(
);
request.headers_mut().insert(
"x-node-name",
http::HeaderValue::from_static("proxy-baseline"),
http::HeaderValue::from_str(&format!("proxy-baseline-{index}"))?,
);
request.headers_mut().insert(
"x-tunnel-max-streams",
@@ -167,6 +260,29 @@ async fn connect_protocol_peer(
let (socket, _response) = tokio_tungstenite::connect_async(request).await?;
let (mut sink, mut stream) = socket.split();
sink.send(Message::Binary(
protocol::encode_hello(&protocol::HelloPayload {
protocol_version: aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION,
capabilities: vec![
"flow-control".to_string(),
"reset-stream".to_string(),
"graceful-drain".to_string(),
],
session_id: Some(format!("baseline-session-{index}")),
replica_id: Some(format!("baseline-replica-{index}")),
})
.into(),
))
.await?;
sink.send(Message::Binary(
protocol::encode_settings(&protocol::SettingsPayload {
initial_stream_window_bytes: 4 * 1024 * 1024,
min_window_update_bytes: 1024 * 1024,
drain_deadline_ms: 30_000,
})
.into(),
))
.await?;
Ok(tokio::spawn(async move {
while let Some(message) = stream.next().await {
let Ok(message) = message else {
@@ -250,6 +366,40 @@ async fn capture_tunnel_metrics(
&[],
)
.unwrap_or_default(),
body_backpressure_total: find_metric_value_u64(
&samples,
"tunnel_body_backpressure_total",
&[],
)
.unwrap_or_default(),
flow_window_blocked_ms: find_metric_value_u64(
&samples,
"tunnel_flow_window_blocked_ms",
&[],
)
.unwrap_or_default(),
connection_health_score: find_metric_value_u64(
&samples,
"tunnel_connection_health_score",
&[],
)
.unwrap_or_default(),
stream_reset_total: find_metric_value_u64(
&samples,
"tunnel_stream_reset_total",
&[("reason", "all")],
)
.unwrap_or_default(),
stream_reset_reasons: collect_reason_metrics(&samples, "tunnel_stream_reset_total"),
drain_total: find_metric_value_u64(&samples, "tunnel_drain_total", &[("reason", "all")])
.unwrap_or_default(),
drain_reasons: collect_reason_metrics(&samples, "tunnel_drain_total"),
scheduler_selected_conn_total: find_metric_value_u64(
&samples,
"tunnel_scheduler_selected_conn_total",
&[],
)
.unwrap_or_default(),
proxy_connections_protocol_v1: find_metric_value_u64(
&samples,
"tunnel_proxy_connections_protocol_v1",
@@ -262,9 +412,86 @@ async fn capture_tunnel_metrics(
&[],
)
.unwrap_or_default(),
proxy_connections_protocol_v3: find_metric_value_u64(
&samples,
"tunnel_proxy_connections_protocol_v3",
&[],
)
.unwrap_or_default(),
})
}
fn collect_reason_metrics(
samples: &[PrometheusSample],
metric_name: &str,
) -> BTreeMap<String, u64> {
samples
.iter()
.filter(|sample| {
(sample.name == metric_name || sample.name.ends_with(&format!("_{metric_name}")))
&& sample
.labels
.get("reason")
.is_some_and(|reason| reason != "all")
})
.filter_map(|sample| {
let reason = sample.labels.get("reason")?.clone();
let value = sample.value.parse::<u64>().ok()?;
Some((reason, value))
})
.collect()
}
fn evaluate_acceptance(
config: &GatewayTunnelBaselineConfig,
result: &HttpLoadProbeResult,
metrics: &TunnelMetricsSnapshot,
) -> AcceptanceReport {
let min_success_rate_bps = if config.require_acceptance { 9_950 } else { 1 };
let successful_requests = result
.status_counts
.iter()
.filter(|(status, _)| (200u16..300u16).contains(status))
.map(|(_, count)| *count)
.sum::<usize>();
let success_rate_bps = if result.total_requests == 0 {
0
} else {
((successful_requests as u128) * 10_000 / (result.total_requests as u128)) as u64
};
let congestion_free = metrics.proxy_connection_congested_total == 0;
let no_queue_full_rejections = metrics.outbound_queue_rejected_full_total == 0;
let mut reasons = Vec::new();
if success_rate_bps < min_success_rate_bps {
reasons.push(format!(
"success rate {} bps below required {} bps",
success_rate_bps, min_success_rate_bps
));
}
if !congestion_free {
reasons.push(format!(
"connection congestion total is {}",
metrics.proxy_connection_congested_total
));
}
if !no_queue_full_rejections {
reasons.push(format!(
"outbound queue full rejections total is {}",
metrics.outbound_queue_rejected_full_total
));
}
AcceptanceReport {
required: config.require_acceptance,
passed: reasons.is_empty(),
success_rate_bps,
min_success_rate_bps,
congestion_free,
no_queue_full_rejections,
reasons,
}
}
async fn handle_binary_frame<S>(
sink: &mut S,
data: Vec<u8>,
@@ -285,7 +512,17 @@ where
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 => {
protocol::REQUEST_BODY => {
let payload = protocol::decode_payload(&data, &header).unwrap_or_default();
if !payload.is_empty() {
sink.send(Message::Binary(
protocol::encode_window_update(header.stream_id, payload.len() as u32).into(),
))
.await?;
}
if header.flags & protocol::FLAG_END_STREAM == 0 {
return Ok(());
}
let response_meta = protocol::ResponseMeta {
status: 200,
headers: vec![(
@@ -339,6 +576,25 @@ fn parse_args(
"--concurrency" => {
config.concurrency = next_value(&mut iter, "--concurrency")?.parse()?
}
"--body-bytes" => {
config.request_body_bytes = next_value(&mut iter, "--body-bytes")?.parse()?
}
"--tunnel-connections" => {
config.tunnel_connections =
next_value(&mut iter, "--tunnel-connections")?.parse()?
}
"--outbound-queue-capacity" => {
config.outbound_queue_capacity =
next_value(&mut iter, "--outbound-queue-capacity")?.parse()?
}
"--close-one-connection-after-ms" => {
config.close_one_connection_after = Some(Duration::from_millis(
next_value(&mut iter, "--close-one-connection-after-ms")?.parse()?,
))
}
"--require-acceptance" => {
config.require_acceptance = true;
}
"--timeout-ms" => {
config.timeout =
Duration::from_millis(next_value(&mut iter, "--timeout-ms")?.parse()?)
@@ -359,6 +615,37 @@ fn parse_args(
}
}
}
if config.request_body_bytes == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"--body-bytes must be positive",
)
.into());
}
if config.tunnel_connections == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"--tunnel-connections must be positive",
)
.into());
}
if config.outbound_queue_capacity == 0 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"--outbound-queue-capacity must be positive",
)
.into());
}
if config
.close_one_connection_after
.is_some_and(|duration| duration.is_zero())
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"--close-one-connection-after-ms must be positive",
)
.into());
}
Ok(config)
}
@@ -377,6 +664,6 @@ fn next_value(
fn print_usage() {
eprintln!(
"usage: cargo run -p aether-testkit --bin gateway_tunnel_stream_baseline -- [--requests 200] [--concurrency 20] [--timeout-ms 10000] [--output /tmp/gateway_tunnel_baseline.json]"
"usage: cargo run -p aether-testkit --bin gateway_tunnel_stream_baseline -- [--requests 200] [--concurrency 20] [--body-bytes 6291456] [--tunnel-connections 4] [--outbound-queue-capacity 512] [--close-one-connection-after-ms 1000] [--require-acceptance] [--timeout-ms 10000] [--output /tmp/gateway_tunnel_baseline.json]"
);
}