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

feat(tunnel/usage): proxy writer 双优先级队列、hub 连接压力感知选择、usage 请求记录级别控制及 trace 页面 proxy timing 增强

Browse Source
This commit is contained in:
fawney19
2026-04-16 20:14:50 +08:00
parent 65cd9dc3e5
commit 0dce667019
13 changed files with 1229 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
apps/aether-gateway/src/data/state/integrations.rs
View File

@@ -16,16 +16,37 @@ use aether_data_contracts::repository::settlement::{StoredUsageSettlement, Usage
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
use aether_usage_runtime::{
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRuntimeAccess,
UsageSettlementWriter,
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRequestRecordLevel,
UsageRuntimeAccess, UsageSettlementWriter,
};
use aether_video_tasks_core::StoredVideoTaskReadSide;
use async_trait::async_trait;
use serde_json::Value;
use super::GatewayDataState;
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::provider_transport::ProviderTransportSnapshotSource;
const REQUEST_RECORD_LEVEL_KEY: &str = "request_record_level";
const LEGACY_REQUEST_LOG_LEVEL_KEY: &str = "request_log_level";
fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestRecordLevel {
let Some(value) = value.and_then(Value::as_str).map(str::trim) else {
return UsageRequestRecordLevel::Full;
};
if value.eq_ignore_ascii_case("basic")
|| value.eq_ignore_ascii_case("base")
|| value.eq_ignore_ascii_case("headers")
|| value.eq_ignore_ascii_case("minimal")
|| value.eq_ignore_ascii_case("none")
{
UsageRequestRecordLevel::Basic
} else {
UsageRequestRecordLevel::Full
}
}
#[async_trait]
impl RequestAuditReader for GatewayDataState {
async fn find_request_usage_audit_by_request_id(
@@ -161,6 +182,7 @@ impl UsageBillingEventEnricher for GatewayDataState {
}
}
#[async_trait]
impl UsageRuntimeAccess for GatewayDataState {
fn has_usage_writer(&self) -> bool {
GatewayDataState::has_usage_writer(self)
@@ -173,6 +195,16 @@ impl UsageRuntimeAccess for GatewayDataState {
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
GatewayDataState::usage_worker_runner(self)
}
async fn request_record_level(&self) -> Result<UsageRequestRecordLevel, DataLayerError> {
let value = GatewayDataState::find_system_config_value(self, REQUEST_RECORD_LEVEL_KEY)
.await?
.or(
GatewayDataState::find_system_config_value(self, LEGACY_REQUEST_LOG_LEVEL_KEY)
.await?,
);
Ok(usage_request_record_level_from_value(value.as_ref()))
}
}
#[async_trait]
@@ -188,10 +220,11 @@ impl UsageRecordWriter for GatewayDataState {
#[cfg(test)]
mod tests {
use aether_billing::enrich_usage_event_with_billing;
use serde_json::Value;
use aether_usage_runtime::UsageRuntimeAccess;
use serde_json::{json, Value};
use super::GatewayDataState;
use crate::usage::{UsageEvent, UsageEventData, UsageEventType};
use crate::usage::{UsageEvent, UsageEventData, UsageEventType, UsageRequestRecordLevel};
#[tokio::test]
async fn enriches_completed_usage_event_with_billing_snapshot() {
@@ -255,4 +288,43 @@ mod tests {
Some("complete")
);
}
#[tokio::test]
async fn usage_runtime_access_reads_base_request_record_level_as_basic() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([(
"request_record_level".to_string(),
json!("base"),
)]);
let level = UsageRuntimeAccess::request_record_level(&state)
.await
.expect("request record level should read");
assert_eq!(level, UsageRequestRecordLevel::Basic);
}
#[tokio::test]
async fn usage_runtime_access_falls_back_to_legacy_request_log_level_alias() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([(
"request_log_level".to_string(),
json!("headers"),
)]);
let level = UsageRuntimeAccess::request_record_level(&state)
.await
.expect("legacy request log level should read");
assert_eq!(level, UsageRequestRecordLevel::Basic);
}
#[tokio::test]
async fn usage_runtime_access_defaults_missing_request_record_level_to_full() {
let state = GatewayDataState::disabled();
let level = UsageRuntimeAccess::request_record_level(&state)
.await
.expect("missing request record level should fall back");
assert_eq!(level, UsageRequestRecordLevel::Full);
}
}

85
apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/trace.rs
View File

@@ -1,5 +1,5 @@
use super::super::test_support::{
request_context, sample_candidate, sample_endpoint, sample_key, sample_provider,
request_context, sample_candidate, sample_endpoint, sample_key, sample_provider, sample_usage,
};
use super::local_monitoring_response;
use crate::AppState;
@@ -10,6 +10,7 @@ use std::sync::Arc;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
#[tokio::test]
async fn admin_monitoring_trace_request_returns_local_payload() {
@@ -156,6 +157,88 @@ async fn admin_monitoring_trace_request_keeps_format_conversion_disabled_candida
);
}
#[tokio::test]
async fn admin_monitoring_trace_request_enriches_proxy_timing_from_usage_audit() {
let mut candidate = sample_candidate(
"cand-used",
"request-1",
1,
RequestCandidateStatus::Success,
Some(101),
Some(33),
Some(200),
);
candidate.extra_data = Some(json!({
"proxy": {
"node_id": "proxy-node-1",
"node_name": "edge-1",
"source": "provider"
}
}));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![candidate]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let mut usage = sample_usage(
"request-1",
"provider-1",
"OpenAI",
40,
0.02,
"completed",
Some(200),
100,
);
usage.candidate_id = Some("cand-used".to_string());
usage.response_headers = Some(json!({
"x-proxy-timing": "{\"connection_acquire_ms\":125,\"response_wait_ms\":475,\"ttfb_ms\":600}"
}));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
let data_state =
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidates,
usage_repository,
)
.with_provider_catalog_reader(provider_catalog);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let context = request_context(
http::Method::GET,
"/api/admin/monitoring/trace/request-1?attempted_only=true",
);
let response = local_monitoring_response(&state, &context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(
payload["candidates"][0]["extra_data"]["first_byte_time_ms"],
json!(30)
);
assert_eq!(
payload["candidates"][0]["extra_data"]["proxy"]["ttfb_ms"],
json!(600)
);
assert_eq!(
payload["candidates"][0]["extra_data"]["proxy"]["timing"]["connection_acquire_ms"],
json!(125)
);
assert_eq!(
payload["candidates"][0]["extra_data"]["proxy"]["timing"]["response_wait_ms"],
json!(475)
);
}
#[tokio::test]
async fn admin_monitoring_trace_provider_stats_returns_local_payload() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![

6
apps/aether-gateway/src/handlers/admin/observability/monitoring/trace.rs
View File

@@ -51,9 +51,15 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
attempted_only,
));
};
let usage = state
.data
.read_request_usage_audit(&request_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(build_admin_monitoring_trace_request_payload_response(
&trace,
usage.as_ref(),
))
}

145
apps/aether-gateway/src/tests/usage/local.rs
View File

@@ -371,6 +371,151 @@ async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage()
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let upstream = Router::new().route(
"/api/internal/gateway/report-sync",
any(|_request: Request| async move { Json(json!({"ok": true})) }),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(|_request: Request| async move {
Json(json!({
"request_id": "trace-openai-chat-local-report-sync-base-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-local-report-sync-base-123",
"object": "chat.completion",
"model": "gpt-5-upstream",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "body should not be persisted"
}
}],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
}
},
"telemetry": {
"elapsed_ms": 25
}
}))
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-local-report-sync-base")),
sample_local_openai_auth_snapshot(
"api-key-openai-usage-local-base-1",
"user-openai-usage-local-base-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests([(
"request_record_level".to_string(),
json!("base"),
)]),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-local-report-sync-base",
)
.header(
TRACE_ID_HEADER,
"trace-openai-chat-local-report-sync-base-123",
)
.body(
serde_json::to_string(&json!({
"model": "gpt-5",
"messages": [{
"role": "user",
"content": "request body should not be persisted"
}]
}))
.expect("request should encode"),
)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let body_json: serde_json::Value = response.json().await.expect("body should parse");
assert_eq!(body_json["model"], "gpt-5-upstream");
let stored_usage = wait_for_usage_status(
usage_repository.as_ref(),
"trace-openai-chat-local-report-sync-base-123",
"completed",
)
.await;
assert_eq!(stored_usage.status, "completed");
assert_eq!(stored_usage.total_tokens, 5);
assert_eq!(stored_usage.response_time_ms, Some(25));
assert!(stored_usage.request_body.is_none());
assert!(stored_usage.request_body_ref.is_none());
assert!(stored_usage.provider_request_body.is_none());
assert!(stored_usage.provider_request_body_ref.is_none());
assert!(stored_usage.response_body.is_none());
assert!(stored_usage.response_body_ref.is_none());
assert!(stored_usage.client_response_body.is_none());
assert!(stored_usage.client_response_body_ref.is_none());
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-local-report-sync-base-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exhaust_after_retryable_sync_failure(
) {

346
apps/aether-gateway/src/tunnel/embedded/hub.rs
View File

@@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering}
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_runtime::{BoundedQueueSender, MetricKind, MetricSample, QueueSendError};
use aether_runtime::{BoundedQueueSender, MetricKind, MetricSample, QueueSendError, QueueSnapshot};
use axum::extract::ws::Message;
use bytes::Bytes;
use dashmap::DashMap;
@@ -16,6 +16,8 @@ use super::control_plane::ControlPlaneClient;
use super::protocol;
const MAX_REQUEST_BODY_FRAME_SIZE: usize = 32 * 1024;
const SOFT_AVOID_QUEUE_PRESSURE_PERCENT: u64 = 50;
const SOFT_AVOID_STREAM_PRESSURE_PERCENT: u64 = 85;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SendStatus {
@@ -75,6 +77,10 @@ impl BoundedOutbound {
let _ = self.close_tx.send(true);
true
}
pub fn snapshot(&self) -> QueueSnapshot {
self.tx.snapshot()
}
}
pub struct ProxyConn {
@@ -86,6 +92,7 @@ pub struct ProxyConn {
pub stream_count: AtomicUsize,
pub max_streams: usize,
draining: AtomicBool,
congested_total: AtomicU64,
}
impl ProxyConn {
@@ -106,6 +113,7 @@ impl ProxyConn {
stream_count: AtomicUsize::new(0),
max_streams,
draining: AtomicBool::new(false),
congested_total: AtomicU64::new(0),
}
}
@@ -180,6 +188,7 @@ impl ProxyConn {
let was_closing = self.outbound.is_closing();
let status = self.outbound.send(msg);
if status == SendStatus::Congested && !was_closing {
self.congested_total.fetch_add(1, Ordering::Relaxed);
warn!(
conn_id = self.id,
node_id = %self.node_id,
@@ -190,6 +199,62 @@ impl ProxyConn {
}
status
}
fn snapshot(&self) -> ProxyConnSnapshot {
let outbound = self.outbound.snapshot();
let stream_count = self.stream_count.load(Ordering::Relaxed);
let queue_pressure_percent = percent_u64(outbound.depth, outbound.capacity);
let stream_pressure_percent = percent_u64(stream_count, self.max_streams);
let soft_avoid = queue_pressure_percent >= SOFT_AVOID_QUEUE_PRESSURE_PERCENT
|| stream_pressure_percent >= SOFT_AVOID_STREAM_PRESSURE_PERCENT;
ProxyConnSnapshot {
conn_id: self.id,
available: self.is_available(),
closing: self.outbound.is_closing(),
draining: self.is_draining(),
stream_count,
max_streams: self.max_streams,
stream_pressure_percent,
outbound,
queue_pressure_percent,
soft_avoid,
congested_total: self.congested_total.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, Copy)]
struct ProxyConnSnapshot {
conn_id: u64,
available: bool,
closing: bool,
draining: bool,
stream_count: usize,
max_streams: usize,
stream_pressure_percent: u64,
outbound: QueueSnapshot,
queue_pressure_percent: u64,
soft_avoid: bool,
congested_total: u64,
}
#[derive(Clone)]
struct ProxyConnCandidate {
conn: Arc<ProxyConn>,
snapshot: ProxyConnSnapshot,
}
impl ProxyConnCandidate {
fn rank_key(&self) -> (u8, u64, u64, usize, usize, u64) {
(
u8::from(self.snapshot.soft_avoid),
self.snapshot.queue_pressure_percent,
self.snapshot.stream_pressure_percent,
self.snapshot.outbound.depth,
self.snapshot.stream_count,
self.snapshot.conn_id,
)
}
}
#[derive(Debug, Clone)]
@@ -343,6 +408,9 @@ pub struct HubRouter {
next_local_stream_id: AtomicU64,
control_plane: ControlPlaneClient,
node_status_tx: mpsc::UnboundedSender<NodeStatusEvent>,
soft_avoid_selection_total: AtomicU64,
selection_retry_total: AtomicU64,
selection_unavailable_total: AtomicU64,
}
#[derive(Debug)]
@@ -391,6 +459,9 @@ impl HubRouter {
next_local_stream_id: AtomicU64::new(1),
control_plane,
node_status_tx,
soft_avoid_selection_total: AtomicU64::new(0),
selection_retry_total: AtomicU64::new(0),
selection_unavailable_total: AtomicU64::new(0),
})
}
@@ -477,28 +548,28 @@ impl HubRouter {
}
}
fn get_proxy_conn(&self, node_id: &str) -> Option<Arc<ProxyConn>> {
let map = self.proxy_conns.read();
let conns = map.get(node_id)?;
let result = conns
.iter()
.filter(|c| c.is_available())
.min_by_key(|c| c.stream_count.load(Ordering::Relaxed))
.cloned();
if result.is_none() && !conns.is_empty() {
warn!(
node_id = %node_id,
total_conns = conns.len(),
closing = conns.iter().filter(|c| c.outbound.is_closing()).count(),
draining = conns.iter().filter(|c| c.is_draining()).count(),
"no available proxy connection despite registered connections"
);
}
result
fn ranked_proxy_conn_candidates(&self, node_id: &str) -> Vec<ProxyConnCandidate> {
let conns = {
let map = self.proxy_conns.read();
map.get(node_id)
.map(|entries| entries.to_vec())
.unwrap_or_default()
};
let mut candidates = conns
.into_iter()
.filter_map(|conn| {
let snapshot = conn.snapshot();
snapshot
.available
.then_some(ProxyConnCandidate { conn, snapshot })
})
.collect::<Vec<_>>();
candidates.sort_by_key(|candidate| candidate.rank_key());
candidates
}
pub fn has_local_proxy(&self, node_id: &str) -> bool {
self.get_proxy_conn(node_id).is_some()
!self.ranked_proxy_conn_candidates(node_id).is_empty()
}
pub fn open_local_stream(
@@ -506,12 +577,49 @@ impl HubRouter {
node_id: &str,
meta: &protocol::RequestMeta,
) -> Result<Arc<LocalStream>, String> {
let proxy_conn = self
.get_proxy_conn(node_id)
.ok_or_else(|| format!("no proxy connection for node {node_id}"))?;
let proxy_stream_id = proxy_conn
.alloc_stream_id()
.ok_or_else(|| format!("stream limit reached for node {node_id}"))?;
let candidates = self.ranked_proxy_conn_candidates(node_id);
if candidates.is_empty() {
self.selection_unavailable_total
.fetch_add(1, Ordering::Relaxed);
self.warn_no_available_proxy_connection(node_id);
return Err(format!("no proxy connection for node {node_id}"));
}
let mut skipped_candidates = 0usize;
let mut selected_candidate = None;
let mut proxy_stream_id = None;
for candidate in candidates {
match candidate.conn.alloc_stream_id() {
Some(stream_id) => {
proxy_stream_id = Some(stream_id);
selected_candidate = Some(candidate);
break;
}
None => skipped_candidates = skipped_candidates.saturating_add(1),
}
}
let Some(candidate) = selected_candidate else {
self.selection_unavailable_total
.fetch_add(1, Ordering::Relaxed);
return Err(format!("stream limit reached for node {node_id}"));
};
if skipped_candidates > 0 {
self.selection_retry_total.fetch_add(1, Ordering::Relaxed);
}
if candidate.snapshot.soft_avoid {
self.soft_avoid_selection_total
.fetch_add(1, Ordering::Relaxed);
debug!(
node_id = %node_id,
conn_id = candidate.snapshot.conn_id,
queue_pressure_percent = candidate.snapshot.queue_pressure_percent,
stream_pressure_percent = candidate.snapshot.stream_pressure_percent,
"selected high-pressure proxy connection because no lower-pressure alternative was available"
);
}
let proxy_conn = candidate.conn;
let proxy_stream_id = proxy_stream_id.expect("selected candidate should carry a stream id");
// Encode frames before registering the stream so that encoding failures
// (practically impossible but theoretically possible) don't leak a stream
@@ -556,6 +664,7 @@ impl HubRouter {
proxy_stream_id = proxy_stream_id,
local_stream_id = local_stream_id,
stream_count = proxy_conn.stream_count.load(Ordering::Relaxed),
queue_depth = proxy_conn.outbound.snapshot().depth,
send_status = ?send_status,
"open_local_stream dispatched"
);
@@ -569,6 +678,28 @@ impl HubRouter {
}
}
fn warn_no_available_proxy_connection(&self, node_id: &str) {
let conns = {
let map = self.proxy_conns.read();
map.get(node_id)
.map(|entries| entries.to_vec())
.unwrap_or_default()
};
if conns.is_empty() {
return;
}
let snapshots = conns.iter().map(|conn| conn.snapshot()).collect::<Vec<_>>();
warn!(
node_id = %node_id,
total_conns = snapshots.len(),
available = snapshots.iter().filter(|snapshot| snapshot.available).count(),
closing = snapshots.iter().filter(|snapshot| snapshot.closing).count(),
draining = snapshots.iter().filter(|snapshot| snapshot.draining).count(),
soft_avoid = snapshots.iter().filter(|snapshot| snapshot.soft_avoid).count(),
"no available proxy connection despite registered connections"
);
}
pub fn push_local_request_body(
&self,
local_stream_id: u64,
@@ -873,15 +1004,72 @@ impl HubRouter {
}
pub fn stats(&self) -> HubStats {
let proxy_conns = self.proxy_conns.read();
let total_proxy = proxy_conns.values().map(|v| v.len()).sum();
let nodes = proxy_conns.len();
drop(proxy_conns);
let proxy_conns = self
.proxy_conns_by_id
.iter()
.map(|entry| entry.value().snapshot())
.collect::<Vec<_>>();
let total_proxy = proxy_conns.len();
let nodes = self.proxy_conns.read().len();
let available_proxy_connections = proxy_conns
.iter()
.filter(|snapshot| snapshot.available)
.count();
let closing_proxy_connections = proxy_conns
.iter()
.filter(|snapshot| snapshot.closing)
.count();
let draining_proxy_connections = proxy_conns
.iter()
.filter(|snapshot| snapshot.draining)
.count();
let soft_avoid_proxy_connections = proxy_conns
.iter()
.filter(|snapshot| snapshot.available && snapshot.soft_avoid)
.count();
let outbound_queue_depth_total = proxy_conns
.iter()
.map(|snapshot| snapshot.outbound.depth)
.sum();
let outbound_queue_depth_max = proxy_conns
.iter()
.map(|snapshot| snapshot.outbound.depth)
.max()
.unwrap_or(0);
let outbound_queue_capacity_total = proxy_conns
.iter()
.map(|snapshot| snapshot.outbound.capacity)
.sum();
let outbound_queue_rejected_full_total = proxy_conns
.iter()
.map(|snapshot| snapshot.outbound.rejected_full_total)
.sum();
let outbound_queue_rejected_closed_total = proxy_conns
.iter()
.map(|snapshot| snapshot.outbound.rejected_closed_total)
.sum();
let proxy_connection_congested_total = proxy_conns
.iter()
.map(|snapshot| snapshot.congested_total)
.sum();
HubStats {
proxy_connections: total_proxy,
available_proxy_connections,
closing_proxy_connections,
draining_proxy_connections,
soft_avoid_proxy_connections,
nodes,
active_streams: self.local_streams.len(),
outbound_queue_depth_total,
outbound_queue_depth_max,
outbound_queue_capacity_total,
outbound_queue_rejected_full_total,
outbound_queue_rejected_closed_total,
proxy_connection_congested_total,
soft_avoid_selection_total: self.soft_avoid_selection_total.load(Ordering::Relaxed),
selection_retry_total: self.selection_retry_total.load(Ordering::Relaxed),
selection_unavailable_total: self.selection_unavailable_total.load(Ordering::Relaxed),
}
}
}
@@ -893,11 +1081,31 @@ fn current_unix_secs() -> u64 {
.as_secs()
}
fn percent_u64(value: usize, total: usize) -> u64 {
if total == 0 {
return 0;
}
((value as u128) * 100 / (total as u128)) as u64
}
#[derive(serde::Serialize)]
pub struct HubStats {
pub proxy_connections: usize,
pub available_proxy_connections: usize,
pub closing_proxy_connections: usize,
pub draining_proxy_connections: usize,
pub soft_avoid_proxy_connections: usize,
pub nodes: usize,
pub active_streams: usize,
pub outbound_queue_depth_total: usize,
pub outbound_queue_depth_max: usize,
pub outbound_queue_capacity_total: usize,
pub outbound_queue_rejected_full_total: u64,
pub outbound_queue_rejected_closed_total: u64,
pub proxy_connection_congested_total: u64,
pub soft_avoid_selection_total: u64,
pub selection_retry_total: u64,
pub selection_unavailable_total: u64,
}
impl HubStats {
@@ -909,6 +1117,30 @@ impl HubStats {
MetricKind::Gauge,
self.proxy_connections as u64,
),
MetricSample::new(
"tunnel_proxy_connections_available",
"Current number of proxy connections available for new work.",
MetricKind::Gauge,
self.available_proxy_connections as u64,
),
MetricSample::new(
"tunnel_proxy_connections_closing",
"Current number of proxy connections marked closing.",
MetricKind::Gauge,
self.closing_proxy_connections as u64,
),
MetricSample::new(
"tunnel_proxy_connections_draining",
"Current number of proxy connections marked draining.",
MetricKind::Gauge,
self.draining_proxy_connections as u64,
),
MetricSample::new(
"tunnel_proxy_connections_soft_avoid",
"Current number of available proxy connections currently soft-avoided by the scheduler.",
MetricKind::Gauge,
self.soft_avoid_proxy_connections as u64,
),
MetricSample::new(
"tunnel_nodes",
"Current number of connected logical nodes.",
@@ -921,6 +1153,60 @@ impl HubStats {
MetricKind::Gauge,
self.active_streams as u64,
),
MetricSample::new(
"tunnel_proxy_outbound_queue_depth_total",
"Current aggregate depth across proxy outbound queues.",
MetricKind::Gauge,
self.outbound_queue_depth_total as u64,
),
MetricSample::new(
"tunnel_proxy_outbound_queue_depth_max",
"Current maximum depth observed on a single proxy outbound queue.",
MetricKind::Gauge,
self.outbound_queue_depth_max as u64,
),
MetricSample::new(
"tunnel_proxy_outbound_queue_capacity_total",
"Current aggregate capacity across proxy outbound queues.",
MetricKind::Gauge,
self.outbound_queue_capacity_total as u64,
),
MetricSample::new(
"tunnel_proxy_outbound_queue_rejected_full_total",
"Total proxy outbound queue sends rejected because a queue was full.",
MetricKind::Counter,
self.outbound_queue_rejected_full_total,
),
MetricSample::new(
"tunnel_proxy_outbound_queue_rejected_closed_total",
"Total proxy outbound queue sends rejected because a queue was closed.",
MetricKind::Counter,
self.outbound_queue_rejected_closed_total,
),
MetricSample::new(
"tunnel_proxy_connection_congested_total",
"Total number of times a proxy outbound queue became congested.",
MetricKind::Counter,
self.proxy_connection_congested_total,
),
MetricSample::new(
"tunnel_proxy_soft_avoid_selection_total",
"Total number of times the scheduler had to pick a high-pressure proxy connection.",
MetricKind::Counter,
self.soft_avoid_selection_total,
),
MetricSample::new(
"tunnel_proxy_selection_retry_total",
"Total number of times the scheduler retried a lower-ranked proxy connection after a race on stream allocation.",
MetricKind::Counter,
self.selection_retry_total,
),
MetricSample::new(
"tunnel_proxy_selection_unavailable_total",
"Total number of relay selections that failed because no proxy connection was available.",
MetricKind::Counter,
self.selection_unavailable_total,
),
]
}
}

3
apps/aether-gateway/src/usage/mod.rs
View File

@@ -6,7 +6,8 @@ pub(crate) mod write;
pub(crate) use aether_usage_runtime::UsageRuntime;
pub use aether_usage_runtime::UsageRuntimeConfig;
pub(crate) use aether_usage_runtime::{
now_ms, UsageEvent, UsageEventData, UsageEventType, UsageQueue, USAGE_EVENT_VERSION,
now_ms, UsageEvent, UsageEventData, UsageEventType, UsageQueue, UsageRequestRecordLevel,
USAGE_EVENT_VERSION,
};
pub(crate) use reporting::{
spawn_sync_report, submit_stream_report, submit_sync_report, GatewayStreamReportRequest,

6
apps/aether-proxy/src/tunnel/dispatcher.rs
View File

@@ -404,10 +404,12 @@ mod tests {
#[tokio::test]
async fn try_send_stream_error_emits_stream_error_frame() {
let (frame_tx, mut frame_rx) = bounded_queue::<Frame>(4);
let (high_tx, mut high_rx) = bounded_queue::<Frame>(4);
let (normal_tx, _normal_rx) = bounded_queue::<Frame>(4);
let frame_tx = FrameSender::from_test_queues(high_tx, normal_tx);
try_send_stream_error(&frame_tx, 9, "proxy request body dispatch stalled");
let frame = frame_rx
let frame = high_rx
.recv()
.await
.expect("stream error frame should enqueue");

168
apps/aether-proxy/src/tunnel/stream_handler.rs
View File

@@ -1164,15 +1164,20 @@ fn build_prefixed_request_body(
mod tests {
use std::collections::HashMap;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::atomic::AtomicU64;
use std::sync::Once;
use std::sync::{Mutex, Once};
use std::task::{Context, Poll};
use aether_runtime::{bounded_queue, ConcurrencyGate, DistributedConcurrencyGate};
use aether_runtime::{ConcurrencyGate, DistributedConcurrencyGate};
use arc_swap::ArcSwap;
use axum::body::Body;
use axum::http::{header, Response, StatusCode};
use axum::routing::{get, post};
use axum::Router;
use futures_util::Sink;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::{Error as WebSocketError, Message};
use super::*;
use crate::config::Config;
@@ -1370,14 +1375,22 @@ mod tests {
let state = sample_state_for_port(addr.port());
cache_test_host(&state, host, addr).await;
let server_ctx = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let (_body_tx, body_rx) = mpsc::channel(1);
let mut meta = sample_request_meta();
meta.url = format!("http://{host}:{}/start", addr.port());
handle_stream(Arc::clone(&state), server_ctx, 5, meta, body_rx, frame_tx).await;
let result = collect_stream_result(&mut frame_rx).await;
handle_stream(
Arc::clone(&state),
server_ctx,
5,
meta,
body_rx,
frame_tx.clone(),
)
.await;
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
server.abort();
assert!(
@@ -1423,14 +1436,22 @@ mod tests {
let state = sample_state_for_port(addr.port());
cache_test_host(&state, host, addr).await;
let server_ctx = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let (_body_tx, body_rx) = mpsc::channel(1);
let mut meta = sample_request_meta();
meta.url = format!("http://{host}:{}/ok", addr.port());
handle_stream(Arc::clone(&state), server_ctx, 3, meta, body_rx, frame_tx).await;
let result = collect_stream_result(&mut frame_rx).await;
handle_stream(
Arc::clone(&state),
server_ctx,
3,
meta,
body_rx,
frame_tx.clone(),
)
.await;
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
server.abort();
assert!(
@@ -1486,7 +1507,7 @@ mod tests {
let state = sample_state_for_port(addr.port());
cache_test_host(&state, host, addr).await;
let server_ctx = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let (body_tx, body_rx) = mpsc::channel(4);
body_tx
.send(TunnelFrame::new(
@@ -1504,8 +1525,16 @@ mod tests {
meta.url = format!("http://{host}:{}/start", addr.port());
meta.follow_redirects = Some(true);
handle_stream(Arc::clone(&state), server_ctx, 1, meta, body_rx, frame_tx).await;
let result = collect_stream_result(&mut frame_rx).await;
handle_stream(
Arc::clone(&state),
server_ctx,
1,
meta,
body_rx,
frame_tx.clone(),
)
.await;
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
server.abort();
assert!(
@@ -1553,15 +1582,23 @@ mod tests {
let state = sample_state_for_port(addr.port());
cache_test_host(&state, host, addr).await;
let server_ctx = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let (_body_tx, body_rx) = mpsc::channel(1);
let mut meta = sample_request_meta();
meta.url = format!("http://{host}:{}/start", addr.port());
meta.follow_redirects = Some(false);
handle_stream(Arc::clone(&state), server_ctx, 7, meta, body_rx, frame_tx).await;
let result = collect_stream_result(&mut frame_rx).await;
handle_stream(
Arc::clone(&state),
server_ctx,
7,
meta,
body_rx,
frame_tx.clone(),
)
.await;
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
server.abort();
assert!(
@@ -1618,7 +1655,7 @@ mod tests {
let state = sample_state_for_budget(addr.port(), 0);
cache_test_host(&state, host, addr).await;
let server_ctx = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(16);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let (body_tx, body_rx) = mpsc::channel(4);
body_tx
.send(TunnelFrame::new(
@@ -1636,8 +1673,16 @@ mod tests {
meta.url = format!("http://{host}:{}/start", addr.port());
meta.follow_redirects = Some(true);
handle_stream(Arc::clone(&state), server_ctx, 11, meta, body_rx, frame_tx).await;
let result = collect_stream_result(&mut frame_rx).await;
handle_stream(
Arc::clone(&state),
server_ctx,
11,
meta,
body_rx,
frame_tx.clone(),
)
.await;
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
server.abort();
assert!(
@@ -1663,7 +1708,7 @@ mod tests {
let _permit = gate.try_acquire().expect("first permit");
let state = sample_state(Some(gate), None);
let server = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let (_body_tx, body_rx) = mpsc::channel(1);
handle_stream(
@@ -1672,11 +1717,15 @@ mod tests {
7,
sample_request_meta(),
body_rx,
frame_tx,
frame_tx.clone(),
)
.await;
let frame = frame_rx.recv().await.expect("overload frame");
let frame = collect_emitted_frames(frame_tx, sent, writer_handle)
.await
.into_iter()
.find(|frame| frame.msg_type == MsgType::StreamError)
.expect("overload frame");
assert_eq!(frame.stream_id, 7);
assert_eq!(frame.msg_type, MsgType::StreamError);
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
@@ -1700,7 +1749,7 @@ mod tests {
let _permit = gate.try_acquire().await.expect("first permit");
let state = sample_state(None, Some(gate));
let server = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let (_body_tx, body_rx) = mpsc::channel(1);
handle_stream(
@@ -1709,11 +1758,15 @@ mod tests {
9,
sample_request_meta(),
body_rx,
frame_tx,
frame_tx.clone(),
)
.await;
let frame = frame_rx.recv().await.expect("overload frame");
let frame = collect_emitted_frames(frame_tx, sent, writer_handle)
.await
.into_iter()
.find(|frame| frame.msg_type == MsgType::StreamError)
.expect("overload frame");
assert_eq!(frame.stream_id, 9);
assert_eq!(frame.msg_type, MsgType::StreamError);
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
@@ -1881,20 +1934,85 @@ mod tests {
.await;
}
#[derive(Clone, Default)]
struct VecSink {
sent: Arc<Mutex<Vec<Message>>>,
}
impl Sink<Message> for VecSink {
type Error = WebSocketError;
fn poll_ready(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
self.sent.lock().expect("sink lock").push(item);
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
fn spawn_test_writer() -> (FrameSender, Arc<Mutex<Vec<Message>>>, JoinHandle<()>) {
let sink = VecSink::default();
let sent = Arc::clone(&sink.sent);
let (frame_tx, handle) = crate::tunnel::writer::spawn_writer(sink, Duration::from_secs(60));
(frame_tx, sent, handle)
}
struct StreamResult {
response: Option<ResponseMeta>,
body: Bytes,
error: Option<String>,
}
async fn collect_emitted_frames(
frame_tx: FrameSender,
sent: Arc<Mutex<Vec<Message>>>,
writer_handle: JoinHandle<()>,
) -> Vec<TunnelFrame> {
drop(frame_tx);
writer_handle.await.expect("writer should exit cleanly");
sent.lock()
.expect("sink lock")
.iter()
.filter_map(|message| match message {
Message::Binary(data) => {
Some(TunnelFrame::decode(data.clone().into()).expect("frame should decode"))
}
Message::Ping(_) | Message::Pong(_) | Message::Close(_) => None,
other => panic!("unexpected writer message: {other:?}"),
})
.collect()
}
async fn collect_stream_result(
frame_rx: &mut aether_runtime::BoundedQueueReceiver<TunnelFrame>,
frame_tx: FrameSender,
sent: Arc<Mutex<Vec<Message>>>,
writer_handle: JoinHandle<()>,
) -> StreamResult {
let mut response = None;
let mut body = BytesMut::new();
let mut error = None;
while let Some(frame) = frame_rx.recv().await {
for frame in collect_emitted_frames(frame_tx, sent, writer_handle).await {
match frame.msg_type {
MsgType::ResponseHeaders => {
let payload = decompress_if_gzip(&frame).expect("headers payload");

227
apps/aether-proxy/src/tunnel/writer.rs
View File

@@ -7,7 +7,10 @@
use std::time::Duration;
use aether_runtime::{bounded_queue, BoundedQueueSender};
use aether_contracts::tunnel::MsgType;
#[cfg(test)]
use aether_runtime::QueueSnapshot;
use aether_runtime::{bounded_queue, BoundedQueueSender, QueueSendError};
use futures_util::SinkExt;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message;
@@ -15,8 +18,60 @@ use tracing::{debug, error, trace};
use super::protocol::Frame;
const HIGH_PRIORITY_QUEUE_CAPACITY: usize = 64;
const NORMAL_PRIORITY_QUEUE_CAPACITY: usize = 256;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum FramePriority {
High,
Normal,
}
#[cfg(test)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FrameQueueSnapshots {
pub high: QueueSnapshot,
pub normal: QueueSnapshot,
}
/// Sender half — cloned by stream handlers and heartbeat.
pub type FrameSender = BoundedQueueSender<Frame>;
#[derive(Debug, Clone)]
pub struct FrameSender {
high_tx: BoundedQueueSender<Frame>,
normal_tx: BoundedQueueSender<Frame>,
}
impl FrameSender {
pub async fn send(&self, frame: Frame) -> Result<(), QueueSendError<Frame>> {
match classify_frame_priority(&frame) {
FramePriority::High => self.high_tx.send(frame).await,
FramePriority::Normal => self.normal_tx.send(frame).await,
}
}
pub fn try_send(&self, frame: Frame) -> Result<(), QueueSendError<Frame>> {
match classify_frame_priority(&frame) {
FramePriority::High => self.high_tx.try_send(frame),
FramePriority::Normal => self.normal_tx.try_send(frame),
}
}
#[cfg(test)]
pub fn snapshots(&self) -> FrameQueueSnapshots {
FrameQueueSnapshots {
high: self.high_tx.snapshot(),
normal: self.normal_tx.snapshot(),
}
}
#[cfg(test)]
pub(crate) fn from_test_queues(
high_tx: BoundedQueueSender<Frame>,
normal_tx: BoundedQueueSender<Frame>,
) -> Self {
Self { high_tx, normal_tx }
}
}
/// Spawn the writer task. Returns the sender and a JoinHandle for cleanup.
///
@@ -26,33 +81,56 @@ pub fn spawn_writer<S>(mut sink: S, ping_interval: Duration) -> (FrameSender, Jo
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{
let (tx, mut rx) = bounded_queue::<Frame>(256);
let (high_tx, mut high_rx) = bounded_queue::<Frame>(HIGH_PRIORITY_QUEUE_CAPACITY);
let (normal_tx, mut normal_rx) = bounded_queue::<Frame>(NORMAL_PRIORITY_QUEUE_CAPACITY);
let tx = FrameSender { high_tx, normal_tx };
let handle = tokio::spawn(async move {
let mut ping_ticker = tokio::time::interval(ping_interval);
let mut high_open = true;
let mut normal_open = true;
ping_ticker.tick().await; // skip first immediate tick
loop {
if let Ok(frame) = high_rx.try_recv() {
if !write_frame(&mut sink, frame).await {
break;
}
continue;
}
if !high_open && !normal_open {
break;
}
tokio::select! {
frame = rx.recv() => {
biased;
frame = high_rx.recv(), if high_open => {
match frame {
Some(frame) => {
let data = frame.encode();
if let Err(e) = sink.send(Message::Binary(data.into())).await {
error!(error = %e, "failed to write frame to WebSocket");
if !write_frame(&mut sink, frame).await {
break;
}
}
None => break, // all senders dropped
None => high_open = false,
}
}
_ = ping_ticker.tick() => {
_ = ping_ticker.tick(), if high_open || normal_open => {
if let Err(e) = sink.send(Message::Ping(vec![])).await {
error!(error = %e, "failed to send WebSocket ping");
break;
}
trace!("sent WebSocket ping");
}
frame = normal_rx.recv(), if normal_open => {
match frame {
Some(frame) => {
if !write_frame(&mut sink, frame).await {
break;
}
}
None => normal_open = false,
}
}
}
}
debug!("writer task exiting");
@@ -61,3 +139,134 @@ where
(tx, handle)
}
fn classify_frame_priority(frame: &Frame) -> FramePriority {
match frame.msg_type {
MsgType::ResponseHeaders
| MsgType::StreamError
| MsgType::Ping
| MsgType::Pong
| MsgType::GoAway
| MsgType::HeartbeatData
| MsgType::HeartbeatAck => FramePriority::High,
MsgType::RequestHeaders
| MsgType::RequestBody
| MsgType::ResponseBody
| MsgType::StreamEnd => FramePriority::Normal,
}
}
async fn write_frame<S>(sink: &mut S, frame: Frame) -> bool
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{
let data = frame.encode();
if let Err(e) = sink.send(Message::Binary(data.into())).await {
error!(error = %e, "failed to write frame to WebSocket");
return false;
}
true
}
#[cfg(test)]
mod tests {
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;
use futures_util::Sink;
use tokio_tungstenite::tungstenite::{Error, Message};
use super::spawn_writer;
use crate::tunnel::protocol::Frame;
use aether_contracts::tunnel::MsgType;
#[derive(Clone, Default)]
struct VecSink {
sent: Arc<Mutex<Vec<Message>>>,
}
impl Sink<Message> for VecSink {
type Error = Error;
fn poll_ready(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
self.sent.lock().expect("sink lock").push(item);
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn prioritizes_control_frames_ahead_of_buffered_body_frames() {
let sink = VecSink::default();
let sent = Arc::clone(&sink.sent);
let (sender, handle) = spawn_writer(sink, Duration::from_secs(60));
for idx in 0..8u8 {
sender
.try_send(Frame::new(
7,
MsgType::ResponseBody,
0,
bytes::Bytes::from(vec![idx; 32]),
))
.expect("frame send should succeed");
}
sender
.try_send(Frame::new(
7,
MsgType::StreamError,
0,
bytes::Bytes::from_static(b"boom"),
))
.expect("frame send should succeed");
let snapshots = sender.snapshots();
assert!(snapshots.high.enqueued_total >= 1);
assert!(snapshots.normal.enqueued_total >= 8);
tokio::time::sleep(Duration::from_millis(30)).await;
drop(sender);
handle.await.expect("writer should exit cleanly");
let sent = sent.lock().expect("sink lock");
assert!(
sent.len() >= 2,
"writer should flush both body and control frames"
);
let first = match &sent[0] {
Message::Binary(data) => {
Frame::decode(data.clone().into()).expect("frame should decode")
}
other => panic!("unexpected first message: {other:?}"),
};
let second = match &sent[1] {
Message::Binary(data) => {
Frame::decode(data.clone().into()).expect("frame should decode")
}
other => panic!("unexpected second message: {other:?}"),
};
assert_eq!(first.msg_type, MsgType::StreamError);
assert_eq!(second.msg_type, MsgType::ResponseBody);
}
}

137
crates/aether-admin/src/observability/monitoring.rs
View File

@@ -1,6 +1,7 @@
use aether_data_contracts::repository::{
candidates::{DecisionTrace, DecisionTraceCandidate},
candidates::{DecisionTrace, DecisionTraceCandidate, RequestCandidateStatus},
provider_catalog::StoredProviderCatalogKey,
usage::StoredRequestUsageAudit,
};
use axum::{
body::Body,
@@ -261,11 +262,20 @@ pub fn build_admin_monitoring_trace_provider_stats_payload_response(
pub fn build_admin_monitoring_trace_request_payload_response(
trace: &DecisionTrace,
usage: Option<&StoredRequestUsageAudit>,
) -> Response<Body> {
let usage_candidate_id =
usage.and_then(|item| resolve_admin_monitoring_usage_candidate_id(trace, item));
let candidates = trace
.candidates
.iter()
.map(build_admin_monitoring_trace_request_candidate_payload)
.map(|item| {
let matched_usage = usage_candidate_id
.as_deref()
.filter(|candidate_id| *candidate_id == item.candidate.id.as_str())
.and(usage);
build_admin_monitoring_trace_request_candidate_payload(item, matched_usage)
})
.collect::<Vec<_>>();
Json(json!({
"request_id": trace.request_id,
@@ -279,6 +289,7 @@ pub fn build_admin_monitoring_trace_request_payload_response(
pub fn build_admin_monitoring_trace_request_candidate_payload(
item: &DecisionTraceCandidate,
usage: Option<&StoredRequestUsageAudit>,
) -> Value {
let candidate = &item.candidate;
json!({
@@ -316,13 +327,133 @@ pub fn build_admin_monitoring_trace_request_candidate_payload(
"error_message": candidate.error_message,
"latency_ms": candidate.latency_ms,
"concurrent_requests": candidate.concurrent_requests,
"extra_data": candidate.extra_data,
"extra_data": build_admin_monitoring_trace_candidate_extra_data(candidate.extra_data.as_ref(), usage),
"created_at": unix_ms_to_rfc3339(candidate.created_at_unix_ms),
"started_at": candidate.started_at_unix_ms.and_then(unix_ms_to_rfc3339),
"finished_at": candidate.finished_at_unix_ms.and_then(unix_ms_to_rfc3339),
})
}
fn resolve_admin_monitoring_usage_candidate_id(
trace: &DecisionTrace,
usage: &StoredRequestUsageAudit,
) -> Option<String> {
if usage.request_id.trim() != trace.request_id {
return None;
}
if let Some(candidate_id) = usage
.routing_candidate_id()
.map(str::trim)
.filter(|value| !value.is_empty())
{
return Some(candidate_id.to_string());
}
let candidate_index = usage.routing_candidate_index()?;
trace
.candidates
.iter()
.filter(|item| u64::from(item.candidate.candidate_index) == candidate_index)
.max_by_key(|item| {
(
item.candidate.retry_index,
item.candidate.finished_at_unix_ms.unwrap_or_default(),
item.candidate.started_at_unix_ms.unwrap_or_default(),
admin_monitoring_candidate_status_rank(item.candidate.status),
)
})
.map(|item| item.candidate.id.clone())
}
fn admin_monitoring_candidate_status_rank(status: RequestCandidateStatus) -> u8 {
match status {
RequestCandidateStatus::Success => 7,
RequestCandidateStatus::Streaming => 6,
RequestCandidateStatus::Pending => 5,
RequestCandidateStatus::Failed => 4,
RequestCandidateStatus::Cancelled => 3,
RequestCandidateStatus::Skipped => 2,
RequestCandidateStatus::Unused => 1,
RequestCandidateStatus::Available => 0,
}
}
fn build_admin_monitoring_trace_candidate_extra_data(
existing: Option<&Value>,
usage: Option<&StoredRequestUsageAudit>,
) -> Value {
let mut extra_data = match existing {
Some(Value::Object(object)) => Some(object.clone()),
Some(other) => return other.clone(),
None => None,
};
if let Some(usage) = usage {
let extra_object = extra_data.get_or_insert_with(serde_json::Map::new);
if let Some(first_byte_time_ms) = usage.first_byte_time_ms {
extra_object
.entry("first_byte_time_ms".to_string())
.or_insert_with(|| json!(first_byte_time_ms));
}
if let Some(proxy_value) = extra_object.get_mut("proxy") {
if let Some(proxy_object) = proxy_value.as_object_mut() {
let proxy_timing = parse_admin_monitoring_usage_proxy_timing(usage);
let proxy_ttfb_ms = proxy_timing
.as_ref()
.and_then(|timing| timing.get("ttfb_ms"))
.and_then(Value::as_u64)
.or(usage.first_byte_time_ms);
if let Some(ttfb_ms) = proxy_ttfb_ms {
proxy_object
.entry("ttfb_ms".to_string())
.or_insert_with(|| json!(ttfb_ms));
}
if let Some(timing) = proxy_timing {
proxy_object.entry("timing".to_string()).or_insert(timing);
}
}
}
}
match extra_data {
Some(object) => Value::Object(object),
None => Value::Null,
}
}
fn parse_admin_monitoring_usage_proxy_timing(usage: &StoredRequestUsageAudit) -> Option<Value> {
admin_monitoring_header_value(usage.response_headers.as_ref(), "x-proxy-timing")
.or_else(|| {
admin_monitoring_header_value(usage.client_response_headers.as_ref(), "x-proxy-timing")
})
.and_then(|raw| parse_admin_monitoring_proxy_timing_value(&raw))
}
fn admin_monitoring_header_value(headers: Option<&Value>, name: &str) -> Option<String> {
headers
.and_then(Value::as_object)
.and_then(|object| {
object
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case(name))
.map(|(_, value)| value)
})
.and_then(|value| match value {
Value::String(text) => Some(text.trim().to_string()),
Value::Object(object) => Some(Value::Object(object.clone()).to_string()),
_ => None,
})
.filter(|value| !value.is_empty())
}
fn parse_admin_monitoring_proxy_timing_value(raw: &str) -> Option<Value> {
serde_json::from_str::<Value>(raw)
.ok()
.filter(Value::is_object)
}
#[allow(clippy::too_many_arguments)]
pub fn build_admin_monitoring_system_status_payload_response(
timestamp: chrono::DateTime<chrono::Utc>,

2
crates/aether-admin/src/system.rs
View File

@@ -1201,7 +1201,7 @@ pub fn admin_system_config_default_value(key: &str) -> Option<serde_json::Value>
"site_subtitle" => Some(json!("AI Gateway")),
"default_user_initial_gift_usd" => Some(json!(10.0)),
"password_policy_level" => Some(json!("weak")),
REQUEST_RECORD_LEVEL_KEY => Some(json!("basic")),
REQUEST_RECORD_LEVEL_KEY => Some(json!("full")),
"max_request_body_size" => Some(json!(5_242_880)),
"max_response_body_size" => Some(json!(5_242_880)),
"sensitive_headers" => Some(json!([

4
crates/aether-usage-runtime/src/lib.rs
View File

@@ -30,7 +30,9 @@ pub use report_context::{
build_locally_actionable_report_context_from_request_candidate,
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
};
pub use runtime::{UsageBillingEventEnricher, UsageRuntime, UsageRuntimeAccess};
pub use runtime::{
UsageBillingEventEnricher, UsageRequestRecordLevel, UsageRuntime, UsageRuntimeAccess,
};
pub use settlement::{settle_usage_if_needed, UsageSettlementWriter};
pub use standardized_usage::StandardizedUsage;
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};

97
crates/aether-usage-runtime/src/runtime.rs
View File

@@ -24,12 +24,24 @@ pub trait UsageBillingEventEnricher: Send + Sync {
async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum UsageRequestRecordLevel {
Basic,
#[default]
Full,
}
#[async_trait]
pub trait UsageRuntimeAccess:
UsageRecordWriter + UsageSettlementWriter + UsageBillingEventEnricher + Send + Sync
{
fn has_usage_writer(&self) -> bool;
fn has_usage_worker_runner(&self) -> bool;
fn usage_worker_runner(&self) -> Option<RedisStreamRunner>;
async fn request_record_level(&self) -> Result<UsageRequestRecordLevel, DataLayerError> {
Ok(UsageRequestRecordLevel::Full)
}
}
#[derive(Debug, Clone)]
@@ -197,6 +209,7 @@ impl UsageRuntime {
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_sync_terminal_usage_event_offthread(input).await {
Ok(mut event) => {
apply_request_record_level_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_sync_terminal_billing_enrichment_failed",
@@ -244,6 +257,7 @@ impl UsageRuntime {
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_stream_terminal_usage_event_offthread(input).await {
Ok(mut event) => {
apply_request_record_level_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_stream_terminal_billing_enrichment_failed",
@@ -289,6 +303,7 @@ impl UsageRuntime {
if !self.is_enabled() {
return;
}
apply_request_record_level_from_data(data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_terminal_billing_enrichment_failed",
@@ -431,6 +446,40 @@ fn join_error_to_data_layer(err: tokio::task::JoinError) -> DataLayerError {
DataLayerError::UnexpectedValue(format!("usage builder task join failed: {err}"))
}
async fn apply_request_record_level_from_data<T>(data: &T, event: &mut UsageEvent)
where
T: UsageRuntimeAccess,
{
match data.request_record_level().await {
Ok(level) => apply_request_record_level(level, event),
Err(err) => {
warn!(
event_name = "usage_request_record_level_read_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "full",
error = %err,
"usage runtime failed to read request record level; keeping full capture"
);
}
}
}
fn apply_request_record_level(level: UsageRequestRecordLevel, event: &mut UsageEvent) {
if !matches!(level, UsageRequestRecordLevel::Basic) {
return;
}
event.data.request_body = None;
event.data.request_body_ref = None;
event.data.provider_request_body = None;
event.data.provider_request_body_ref = None;
event.data.response_body = None;
event.data.response_body_ref = None;
event.data.client_response_body = None;
event.data.client_response_body_ref = None;
}
fn boxed_usage_task<F>(task: F) -> Pin<Box<dyn Future<Output = ()> + Send>>
where
F: Future<Output = ()> + Send + 'static,
@@ -444,3 +493,51 @@ fn now_unix_secs() -> u64 {
.unwrap_or_default()
.as_secs()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{apply_request_record_level, UsageRequestRecordLevel};
use crate::{UsageEvent, UsageEventData, UsageEventType};
#[test]
fn basic_request_record_level_strips_body_capture_but_preserves_derived_fields() {
let mut event = UsageEvent::new(
UsageEventType::Failed,
"req-basic-1",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
total_tokens: Some(42),
error_message: Some("upstream failed".to_string()),
request_body: Some(json!({"messages":[{"role":"user","content":"hello"}]})),
request_body_ref: Some("usage://request/req-basic-1/request_body".to_string()),
provider_request_body: Some(json!({"model":"gpt-5"})),
provider_request_body_ref: Some(
"usage://request/req-basic-1/provider_request_body".to_string(),
),
response_body: Some(json!({"error":{"message":"bad gateway"}})),
response_body_ref: Some("usage://request/req-basic-1/response_body".to_string()),
client_response_body: Some(json!({"detail":"bad gateway"})),
client_response_body_ref: Some(
"usage://request/req-basic-1/client_response_body".to_string(),
),
..UsageEventData::default()
},
);
apply_request_record_level(UsageRequestRecordLevel::Basic, &mut event);
assert_eq!(event.data.total_tokens, Some(42));
assert_eq!(event.data.error_message.as_deref(), Some("upstream failed"));
assert!(event.data.request_body.is_none());
assert!(event.data.request_body_ref.is_none());
assert!(event.data.provider_request_body.is_none());
assert!(event.data.provider_request_body_ref.is_none());
assert!(event.data.response_body.is_none());
assert!(event.data.response_body_ref.is_none());
assert!(event.data.client_response_body.is_none());
assert!(event.data.client_response_body_ref.is_none());
}
}