mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +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:
@@ -39,6 +39,12 @@ pub enum ExecutorServiceError {
|
||||
BodyEncode(serde_json::Error),
|
||||
#[error("failed to build HTTP client: {0}")]
|
||||
ClientBuild(reqwest::Error),
|
||||
#[error("failed to read executor request body: {0}")]
|
||||
RequestRead(String),
|
||||
#[error("executor request body is not valid JSON: {0}")]
|
||||
InvalidRequestJson(serde_json::Error),
|
||||
#[error("executor overloaded: gate {gate} saturated at {limit}")]
|
||||
Overloaded { gate: &'static str, limit: usize },
|
||||
#[error("failed to execute upstream request: {0}")]
|
||||
UpstreamRequest(reqwest::Error),
|
||||
#[error("hub relay request failed: {0}")]
|
||||
|
||||
@@ -4,6 +4,10 @@ use clap::Parser;
|
||||
use tracing::info;
|
||||
|
||||
use aether_executor::server;
|
||||
use aether_runtime::{
|
||||
init_service_runtime, DistributedConcurrencyGate, RedisDistributedConcurrencyConfig,
|
||||
ServiceRuntimeConfig,
|
||||
};
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "aether-executor", about = "Internal Rust executor for Aether")]
|
||||
@@ -20,28 +24,106 @@ struct Args {
|
||||
default_value = "/tmp/aether-executor.sock"
|
||||
)]
|
||||
unix_socket: PathBuf,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_MAX_IN_FLIGHT_REQUESTS")]
|
||||
max_in_flight_requests: Option<usize>,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_LIMIT")]
|
||||
distributed_request_limit: Option<usize>,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_URL")]
|
||||
distributed_request_redis_url: Option<String>,
|
||||
|
||||
#[arg(long, env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_KEY_PREFIX")]
|
||||
distributed_request_redis_key_prefix: Option<String>,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_LEASE_TTL_MS",
|
||||
default_value_t = 30_000
|
||||
)]
|
||||
distributed_request_lease_ttl_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_RENEW_INTERVAL_MS",
|
||||
default_value_t = 10_000
|
||||
)]
|
||||
distributed_request_renew_interval_ms: u64,
|
||||
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_EXECUTOR_DISTRIBUTED_REQUEST_COMMAND_TIMEOUT_MS",
|
||||
default_value_t = 1_000
|
||||
)]
|
||||
distributed_request_command_timeout_ms: u64,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "aether_executor=info".into()),
|
||||
)
|
||||
.init();
|
||||
init_service_runtime(ServiceRuntimeConfig::new(
|
||||
"aether-executor",
|
||||
"aether_executor=info",
|
||||
))?;
|
||||
|
||||
let args = Args::parse();
|
||||
let distributed_request_gate = match args.distributed_request_limit.filter(|limit| *limit > 0) {
|
||||
Some(limit) => {
|
||||
let redis_url = args
|
||||
.distributed_request_redis_url
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"AETHER_EXECUTOR_DISTRIBUTED_REQUEST_REDIS_URL is required when distributed request limit is enabled",
|
||||
)
|
||||
})?;
|
||||
Some(DistributedConcurrencyGate::new_redis(
|
||||
"executor_requests_distributed",
|
||||
limit,
|
||||
RedisDistributedConcurrencyConfig {
|
||||
url: redis_url.to_string(),
|
||||
key_prefix: args
|
||||
.distributed_request_redis_key_prefix
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
lease_ttl_ms: args.distributed_request_lease_ttl_ms.max(1),
|
||||
renew_interval_ms: args.distributed_request_renew_interval_ms.max(1),
|
||||
command_timeout_ms: Some(args.distributed_request_command_timeout_ms.max(1)),
|
||||
},
|
||||
)?)
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
match args.transport.trim().to_ascii_lowercase().as_str() {
|
||||
"unix_socket" | "unix" | "uds" => {
|
||||
info!(socket = %args.unix_socket.display(), "aether-executor started");
|
||||
server::serve_unix(&args.unix_socket).await?;
|
||||
server::serve_unix(
|
||||
&args.unix_socket,
|
||||
args.max_in_flight_requests,
|
||||
distributed_request_gate.clone(),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
"tcp" => {
|
||||
info!(bind = %args.bind, "aether-executor started");
|
||||
server::serve_tcp(&args.bind).await?;
|
||||
info!(
|
||||
bind = %args.bind,
|
||||
max_in_flight_requests = args.max_in_flight_requests.unwrap_or_default(),
|
||||
distributed_request_limit = args.distributed_request_limit.unwrap_or_default(),
|
||||
"aether-executor started"
|
||||
);
|
||||
server::serve_tcp(
|
||||
&args.bind,
|
||||
args.max_in_flight_requests,
|
||||
distributed_request_gate,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
other => {
|
||||
return Err(format!("unsupported executor transport: {other}").into());
|
||||
|
||||
@@ -1,13 +1,20 @@
|
||||
use std::convert::Infallible;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
|
||||
ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType,
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionTelemetry,
|
||||
StreamFrame, StreamFramePayload, StreamFrameType,
|
||||
};
|
||||
use aether_runtime::{
|
||||
maybe_hold_axum_response_permit, prometheus_response, service_up_sample, AdmissionPermit,
|
||||
ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, DistributedConcurrencyError,
|
||||
DistributedConcurrencyGate, DistributedConcurrencySnapshot, MetricKind, MetricLabel,
|
||||
MetricSample,
|
||||
};
|
||||
use async_stream::stream;
|
||||
use axum::body::Body;
|
||||
use axum::extract::State;
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::extract::{Request, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::{get, post};
|
||||
@@ -22,25 +29,132 @@ use crate::{encode_frame, ExecutorServiceError, SyncExecutor};
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AppState {
|
||||
executor: SyncExecutor,
|
||||
request_gate: Option<Arc<ConcurrencyGate>>,
|
||||
distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
fn with_request_concurrency_limit(limit: Option<usize>) -> Self {
|
||||
Self {
|
||||
executor: SyncExecutor::new(),
|
||||
request_gate: limit
|
||||
.filter(|limit| *limit > 0)
|
||||
.map(|limit| Arc::new(ConcurrencyGate::new("executor_requests", limit))),
|
||||
distributed_request_gate: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_distributed_request_gate(mut self, gate: DistributedConcurrencyGate) -> Self {
|
||||
self.distributed_request_gate = Some(Arc::new(gate));
|
||||
self
|
||||
}
|
||||
|
||||
fn request_concurrency_snapshot(&self) -> Option<ConcurrencySnapshot> {
|
||||
self.request_gate.as_ref().map(|gate| gate.snapshot())
|
||||
}
|
||||
|
||||
async fn distributed_request_concurrency_snapshot(
|
||||
&self,
|
||||
) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
|
||||
match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => gate.snapshot().await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn metric_samples(&self) -> Vec<MetricSample> {
|
||||
let mut samples = vec![service_up_sample("aether-executor")];
|
||||
if let Some(snapshot) = self.request_concurrency_snapshot() {
|
||||
samples.extend(snapshot.to_metric_samples("executor_requests"));
|
||||
}
|
||||
if let Some(gate) = self.distributed_request_gate.as_ref() {
|
||||
match gate.snapshot().await {
|
||||
Ok(snapshot) => {
|
||||
samples.extend(snapshot.to_metric_samples("executor_requests_distributed"));
|
||||
}
|
||||
Err(_) => samples.push(
|
||||
MetricSample::new(
|
||||
"concurrency_unavailable",
|
||||
"Whether the distributed concurrency gate is currently unavailable.",
|
||||
MetricKind::Gauge,
|
||||
1,
|
||||
)
|
||||
.with_labels(vec![MetricLabel::new(
|
||||
"gate",
|
||||
"executor_requests_distributed",
|
||||
)]),
|
||||
),
|
||||
}
|
||||
}
|
||||
samples
|
||||
}
|
||||
|
||||
async fn try_acquire_request_permit(
|
||||
&self,
|
||||
) -> Result<Option<AdmissionPermit>, RequestAdmissionError> {
|
||||
let local = self
|
||||
.request_gate
|
||||
.as_ref()
|
||||
.map(|gate| gate.try_acquire())
|
||||
.transpose()
|
||||
.map_err(RequestAdmissionError::Local)?;
|
||||
let distributed = match self.distributed_request_gate.as_ref() {
|
||||
Some(gate) => Some(
|
||||
gate.try_acquire()
|
||||
.await
|
||||
.map_err(RequestAdmissionError::Distributed)?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
Ok(AdmissionPermit::from_parts(local, distributed))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_router() -> Router {
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/v1/execute/sync", post(execute_sync))
|
||||
.route("/v1/execute/stream", post(execute_stream))
|
||||
.with_state(AppState {
|
||||
executor: SyncExecutor::new(),
|
||||
})
|
||||
build_router_with_request_concurrency_limit(None)
|
||||
}
|
||||
|
||||
pub async fn serve_tcp(bind: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
pub fn build_router_with_request_concurrency_limit(limit: Option<usize>) -> Router {
|
||||
build_router_with_request_gates(limit, None)
|
||||
}
|
||||
|
||||
pub fn build_router_with_request_gates(
|
||||
limit: Option<usize>,
|
||||
distributed_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Router {
|
||||
let state = match distributed_gate {
|
||||
Some(gate) => {
|
||||
AppState::with_request_concurrency_limit(limit).with_distributed_request_gate(gate)
|
||||
}
|
||||
None => AppState::with_request_concurrency_limit(limit),
|
||||
};
|
||||
Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/metrics", get(metrics))
|
||||
.route("/v1/execute/sync", post(execute_sync))
|
||||
.route("/v1/execute/stream", post(execute_stream))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
pub async fn serve_tcp(
|
||||
bind: &str,
|
||||
max_in_flight_requests: Option<usize>,
|
||||
distributed_request_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let listener = tokio::net::TcpListener::bind(bind).await?;
|
||||
axum::serve(listener, build_router()).await?;
|
||||
axum::serve(
|
||||
listener,
|
||||
build_router_with_request_gates(max_in_flight_requests, distributed_request_gate),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn serve_unix(socket_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
|
||||
pub async fn serve_unix(
|
||||
socket_path: &Path,
|
||||
max_in_flight_requests: Option<usize>,
|
||||
distributed_request_gate: Option<DistributedConcurrencyGate>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
if let Some(parent) = socket_path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
@@ -49,30 +163,69 @@ pub async fn serve_unix(socket_path: &Path) -> Result<(), Box<dyn std::error::Er
|
||||
}
|
||||
|
||||
let listener = tokio::net::UnixListener::bind(socket_path)?;
|
||||
axum::serve(listener, build_router()).await?;
|
||||
axum::serve(
|
||||
listener,
|
||||
build_router_with_request_gates(max_in_flight_requests, distributed_request_gate),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health() -> impl IntoResponse {
|
||||
Json(json!({"status": "ok"}))
|
||||
async fn health(State(state): State<AppState>) -> impl IntoResponse {
|
||||
let request_concurrency = state.request_concurrency_snapshot().map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
let distributed_request_concurrency = state
|
||||
.distributed_request_concurrency_snapshot()
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|snapshot| {
|
||||
json!({
|
||||
"limit": snapshot.limit,
|
||||
"in_flight": snapshot.in_flight,
|
||||
"available_permits": snapshot.available_permits,
|
||||
"high_watermark": snapshot.high_watermark,
|
||||
"rejected": snapshot.rejected,
|
||||
})
|
||||
});
|
||||
Json(json!({
|
||||
"status": "ok",
|
||||
"component": "aether-executor",
|
||||
"request_concurrency": request_concurrency,
|
||||
"distributed_request_concurrency": distributed_request_concurrency,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn metrics(State(state): State<AppState>) -> Response {
|
||||
prometheus_response(&state.metric_samples().await)
|
||||
}
|
||||
|
||||
async fn execute_sync(
|
||||
State(state): State<AppState>,
|
||||
Json(plan): Json<ExecutionPlan>,
|
||||
) -> Result<Json<ExecutionResult>, AppError> {
|
||||
state
|
||||
.executor
|
||||
.execute_sync(plan)
|
||||
.await
|
||||
.map(Json)
|
||||
.map_err(AppError)
|
||||
request: Request,
|
||||
) -> Result<Response, AppError> {
|
||||
let request_permit = acquire_request_permit(&state).await?;
|
||||
let plan = parse_request_json::<ExecutionPlan>(request).await?;
|
||||
let result = state.executor.execute_sync(plan).await.map_err(AppError)?;
|
||||
Ok(maybe_hold_axum_response_permit(
|
||||
Json(result).into_response(),
|
||||
request_permit,
|
||||
))
|
||||
}
|
||||
|
||||
async fn execute_stream(
|
||||
State(state): State<AppState>,
|
||||
Json(plan): Json<ExecutionPlan>,
|
||||
request: Request,
|
||||
) -> Result<Response, AppError> {
|
||||
let request_permit = acquire_request_permit(&state).await?;
|
||||
let plan = parse_request_json::<ExecutionPlan>(request).await?;
|
||||
let execution = state
|
||||
.executor
|
||||
.execute_stream(plan)
|
||||
@@ -149,7 +302,61 @@ async fn execute_stream(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
axum::http::HeaderValue::from_static("application/x-ndjson"),
|
||||
);
|
||||
Ok(response)
|
||||
Ok(maybe_hold_axum_response_permit(response, request_permit))
|
||||
}
|
||||
|
||||
async fn acquire_request_permit(state: &AppState) -> Result<Option<AdmissionPermit>, AppError> {
|
||||
match state.try_acquire_request_permit().await {
|
||||
Ok(permit) => Ok(permit),
|
||||
Err(RequestAdmissionError::Local(ConcurrencyError::Saturated { gate, limit }))
|
||||
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Saturated {
|
||||
gate,
|
||||
limit,
|
||||
}))
|
||||
| Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Unavailable {
|
||||
gate,
|
||||
limit,
|
||||
..
|
||||
})) => Err(AppError(ExecutorServiceError::Overloaded { gate, limit })),
|
||||
Err(RequestAdmissionError::Local(ConcurrencyError::Closed { gate })) => {
|
||||
Err(AppError(ExecutorServiceError::RequestRead(format!(
|
||||
"executor request concurrency gate {gate} is closed"
|
||||
))))
|
||||
}
|
||||
Err(RequestAdmissionError::Distributed(
|
||||
DistributedConcurrencyError::InvalidConfiguration(message),
|
||||
)) => Err(AppError(ExecutorServiceError::RequestRead(message))),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum RequestAdmissionError {
|
||||
Local(ConcurrencyError),
|
||||
Distributed(DistributedConcurrencyError),
|
||||
}
|
||||
|
||||
async fn parse_request_json<T>(request: Request) -> Result<T, AppError>
|
||||
where
|
||||
T: serde::de::DeserializeOwned,
|
||||
{
|
||||
let body = to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.map_err(|err| AppError(ExecutorServiceError::RequestRead(err.to_string())))?;
|
||||
serde_json::from_slice(&body)
|
||||
.map_err(|err| AppError(ExecutorServiceError::InvalidRequestJson(err)))
|
||||
}
|
||||
|
||||
fn build_overloaded_response(message: &str) -> Response {
|
||||
(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"type": "overloaded",
|
||||
"message": message,
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
@@ -158,6 +365,12 @@ struct AppError(ExecutorServiceError);
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
let status_code = match self.0 {
|
||||
ExecutorServiceError::RequestRead(_) | ExecutorServiceError::InvalidRequestJson(_) => {
|
||||
StatusCode::BAD_REQUEST
|
||||
}
|
||||
ExecutorServiceError::Overloaded { .. } => {
|
||||
return build_overloaded_response(&self.0.to_string());
|
||||
}
|
||||
ExecutorServiceError::StreamUnsupported
|
||||
| ExecutorServiceError::RequestBodyRequired
|
||||
| ExecutorServiceError::BodyDecode(_)
|
||||
@@ -185,3 +398,227 @@ impl IntoResponse for AppError {
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{build_router_with_request_concurrency_limit, build_router_with_request_gates};
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
async fn start_server(app: Router) -> (String, tokio::task::JoinHandle<()>) {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let addr = listener.local_addr().expect("local addr should resolve");
|
||||
let handle = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.expect("server should run");
|
||||
});
|
||||
(format!("http://{addr}"), handle)
|
||||
}
|
||||
|
||||
fn stream_plan(url: String) -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: "req-1".into(),
|
||||
candidate_id: Some("cand-1".into()),
|
||||
provider_name: Some("openai".into()),
|
||||
provider_id: "prov-1".into(),
|
||||
endpoint_id: "ep-1".into(),
|
||||
key_id: "key-1".into(),
|
||||
method: "GET".into(),
|
||||
url,
|
||||
headers: std::collections::BTreeMap::new(),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: true,
|
||||
client_api_format: "openai:chat".into(),
|
||||
provider_api_format: "openai:chat".into(),
|
||||
model_name: Some("gpt-4.1".into()),
|
||||
proxy: None,
|
||||
tls_profile: None,
|
||||
timeouts: Some(ExecutionTimeouts {
|
||||
connect_ms: Some(5_000),
|
||||
total_ms: Some(30_000),
|
||||
..ExecutionTimeouts::default()
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn executor_rejects_second_in_flight_stream_request_with_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/slow",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let executor = build_router_with_request_concurrency_limit(Some(1));
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.post(format!("{executor_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
for _ in 0..50 {
|
||||
if upstream_hits.load(Ordering::SeqCst) == 1 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let second_response = client
|
||||
.post(format!("{executor_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["type"],
|
||||
"overloaded"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
executor_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn executor_rejects_second_in_flight_stream_request_with_distributed_overload() {
|
||||
let upstream_hits = Arc::new(AtomicUsize::new(0));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/slow",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
upstream_hits.fetch_add(1, Ordering::SeqCst);
|
||||
let stream = async_stream::stream! {
|
||||
yield Ok::<_, Infallible>(Bytes::from_static(b"chunk-1"));
|
||||
futures_util::future::pending::<()>().await;
|
||||
};
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from_stream(stream))
|
||||
.expect("response should build")
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let distributed_gate = aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
"executor_requests_distributed",
|
||||
1,
|
||||
);
|
||||
let executor_a = build_router_with_request_gates(None, Some(distributed_gate.clone()));
|
||||
let executor_b = build_router_with_request_gates(None, Some(distributed_gate));
|
||||
let (executor_a_url, executor_a_handle) = start_server(executor_a).await;
|
||||
let (executor_b_url, executor_b_handle) = start_server(executor_b).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let first_response = client
|
||||
.post(format!("{executor_a_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
|
||||
for _ in 0..50 {
|
||||
if upstream_hits.load(Ordering::SeqCst) == 1 {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let second_response = client
|
||||
.post(format!("{executor_b_url}/v1/execute/stream"))
|
||||
.json(&stream_plan(format!("{upstream_url}/slow")))
|
||||
.send()
|
||||
.await
|
||||
.expect("second request should complete");
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(
|
||||
second_response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("json body should decode")["error"]["type"],
|
||||
"overloaded"
|
||||
);
|
||||
assert_eq!(upstream_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
drop(first_response);
|
||||
executor_a_handle.abort();
|
||||
executor_b_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn executor_exposes_request_concurrency_metrics() {
|
||||
let executor = build_router_with_request_gates(
|
||||
Some(4),
|
||||
Some(aether_runtime::DistributedConcurrencyGate::new_in_memory(
|
||||
"executor_requests_distributed",
|
||||
6,
|
||||
)),
|
||||
);
|
||||
let (executor_url, executor_handle) = start_server(executor).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{executor_url}/metrics"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("text/plain; version=0.0.4; charset=utf-8")
|
||||
);
|
||||
let body = response.text().await.expect("body should read");
|
||||
assert!(body.contains("service_up{service=\"aether-executor\"} 1"));
|
||||
assert!(body.contains("concurrency_available_permits{gate=\"executor_requests\"} 4"));
|
||||
assert!(body
|
||||
.contains("concurrency_available_permits{gate=\"executor_requests_distributed\"} 6"));
|
||||
|
||||
executor_handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ use std::time::{Duration, Instant};
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResponseBody,
|
||||
};
|
||||
use aether_http::{apply_http_client_config, HttpClientConfig};
|
||||
use base64::Engine as _;
|
||||
use flate2::write::GzEncoder;
|
||||
use flate2::Compression;
|
||||
@@ -256,10 +257,14 @@ fn gzip_bytes(body_bytes: &[u8]) -> Result<Vec<u8>, ExecutorServiceError> {
|
||||
fn build_relay_client(
|
||||
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
|
||||
) -> Result<reqwest::Client, ExecutorServiceError> {
|
||||
let mut builder = reqwest::Client::builder();
|
||||
if let Some(connect_ms) = timeouts.and_then(|timeouts| timeouts.connect_ms) {
|
||||
builder = builder.connect_timeout(Duration::from_millis(connect_ms));
|
||||
}
|
||||
let builder = apply_http_client_config(
|
||||
reqwest::Client::builder(),
|
||||
&HttpClientConfig {
|
||||
connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
|
||||
use_rustls_tls: false,
|
||||
..HttpClientConfig::default()
|
||||
},
|
||||
);
|
||||
builder.build().map_err(ExecutorServiceError::ClientBuild)
|
||||
}
|
||||
|
||||
@@ -339,10 +344,13 @@ fn build_client(
|
||||
proxy: Option<&ProxySnapshot>,
|
||||
tls_profile: Option<&str>,
|
||||
) -> Result<reqwest::Client, ExecutorServiceError> {
|
||||
let mut builder = reqwest::Client::builder().use_rustls_tls();
|
||||
if let Some(connect_ms) = timeouts.and_then(|timeouts| timeouts.connect_ms) {
|
||||
builder = builder.connect_timeout(Duration::from_millis(connect_ms));
|
||||
}
|
||||
let mut builder = apply_http_client_config(
|
||||
reqwest::Client::builder(),
|
||||
&HttpClientConfig {
|
||||
connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
|
||||
..HttpClientConfig::default()
|
||||
},
|
||||
);
|
||||
builder = apply_tls_profile(builder, tls_profile);
|
||||
if let Some(proxy_url) = resolve_proxy_url(proxy)? {
|
||||
let proxy = reqwest::Proxy::all(&proxy_url).map_err(ExecutorServiceError::InvalidProxy)?;
|
||||
|
||||
Reference in New Issue
Block a user