mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
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
263
crates/aether-testkit/src/bin/hub_tunnel_stream_baseline.rs
Normal file
263
crates/aether-testkit/src/bin/hub_tunnel_stream_baseline.rs
Normal 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]"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user