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

refactor(proxy): improve tunnel throughput and observability

Browse Source
This commit is contained in:
fawney19
2026-05-19 23:49:36 +08:00
parent 57655bdb25
commit f5deed8709
30 changed files with 1306 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -5097,6 +5097,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let plan = ExecutionPlan {
@@ -5226,6 +5227,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let plan = ExecutionPlan {

2
apps/aether-gateway/src/execution_runtime/stream_pump.rs
View File

@@ -1291,6 +1291,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let plan = ExecutionPlan {
@@ -1435,6 +1436,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let plan = ExecutionPlan {

1
apps/aether-gateway/src/execution_runtime/transport.rs
View File

@@ -2326,6 +2326,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let plan = ExecutionPlan {

1
apps/aether-gateway/src/maintenance/runtime/tests.rs
View File

@@ -475,6 +475,7 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
proxy_tx,
proxy_close_tx,
16,
2,
)));
let responder_hub = tunnel_state.hub.clone();

1
apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs
View File

@@ -1224,6 +1224,7 @@ async fn gateway_tests_connected_tunnel_proxy_nodes_with_active_probe() {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let gateway = build_router_with_state(state);

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

@@ -1,5 +1,5 @@
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicU8, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
@@ -91,8 +91,11 @@ pub struct ProxyConn {
next_stream_id: AtomicU32,
pub stream_count: AtomicUsize,
pub max_streams: usize,
pub protocol_version: AtomicU8,
draining: AtomicBool,
congested_total: AtomicU64,
write_latency_last_us: AtomicU64,
write_latency_ewma_us: AtomicU64,
}
impl ProxyConn {
@@ -103,6 +106,7 @@ impl ProxyConn {
tx: BoundedQueueSender<Message>,
close_tx: watch::Sender<bool>,
max_streams: usize,
protocol_version: u8,
) -> Self {
Self {
id,
@@ -112,11 +116,28 @@ impl ProxyConn {
next_stream_id: AtomicU32::new(2),
stream_count: AtomicUsize::new(0),
max_streams,
protocol_version: AtomicU8::new(protocol_version.max(1)),
draining: AtomicBool::new(false),
congested_total: AtomicU64::new(0),
write_latency_last_us: AtomicU64::new(0),
write_latency_ewma_us: AtomicU64::new(0),
}
}
pub fn record_write_latency(&self, elapsed: std::time::Duration) {
let micros = u64::try_from(elapsed.as_micros()).unwrap_or(u64::MAX);
self.write_latency_last_us.store(micros, Ordering::Relaxed);
let current = self.write_latency_ewma_us.load(Ordering::Relaxed);
let next = if current == 0 {
micros
} else {
let delta = micros as i128 - current as i128;
(current as i128 + (delta / 8)).max(1) as u64
};
self.write_latency_ewma_us
.store(next.max(1), Ordering::Relaxed);
}
pub fn alloc_stream_id(&self) -> Option<u32> {
let mut current = self.stream_count.load(Ordering::Relaxed);
loop {
@@ -214,11 +235,14 @@ impl ProxyConn {
draining: self.is_draining(),
stream_count,
max_streams: self.max_streams,
protocol_version: self.protocol_version.load(Ordering::Relaxed),
stream_pressure_percent,
outbound,
queue_pressure_percent,
soft_avoid,
congested_total: self.congested_total.load(Ordering::Relaxed),
write_latency_last_us: self.write_latency_last_us.load(Ordering::Relaxed),
write_latency_ewma_us: self.write_latency_ewma_us.load(Ordering::Relaxed),
}
}
}
@@ -231,11 +255,14 @@ struct ProxyConnSnapshot {
draining: bool,
stream_count: usize,
max_streams: usize,
protocol_version: u8,
stream_pressure_percent: u64,
outbound: QueueSnapshot,
queue_pressure_percent: u64,
soft_avoid: bool,
congested_total: u64,
write_latency_last_us: u64,
write_latency_ewma_us: u64,
}
#[derive(Clone)]
@@ -245,11 +272,12 @@ struct ProxyConnCandidate {
}
impl ProxyConnCandidate {
fn rank_key(&self) -> (u8, u64, u64, usize, usize, u64) {
fn rank_key(&self) -> (u8, u64, u64, u64, usize, usize, u64) {
(
u8::from(self.snapshot.soft_avoid),
self.snapshot.queue_pressure_percent,
self.snapshot.stream_pressure_percent,
self.snapshot.write_latency_ewma_us,
self.snapshot.outbound.depth,
self.snapshot.stream_count,
self.snapshot.conn_id,
@@ -744,8 +772,7 @@ impl HubRouter {
payload: &[u8],
end_stream: bool,
) -> Result<(), String> {
let (body_payload, body_flags) = protocol::compress_payload(payload)
.map_err(|e| format!("failed to compress request body: {e}"))?;
let (body_payload, body_flags) = protocol::raw_payload(payload);
let body_frame = protocol::encode_frame(
proxy_stream_id,
protocol::REQUEST_BODY,
@@ -1052,6 +1079,24 @@ impl HubRouter {
.iter()
.map(|snapshot| snapshot.congested_total)
.sum();
let protocol_v1_proxy_connections = proxy_conns
.iter()
.filter(|snapshot| snapshot.protocol_version == 1)
.count();
let protocol_v2_proxy_connections = proxy_conns
.iter()
.filter(|snapshot| snapshot.protocol_version >= 2)
.count();
let write_latency_last_us_max = proxy_conns
.iter()
.map(|snapshot| snapshot.write_latency_last_us)
.max()
.unwrap_or(0);
let write_latency_ewma_us_max = proxy_conns
.iter()
.map(|snapshot| snapshot.write_latency_ewma_us)
.max()
.unwrap_or(0);
HubStats {
proxy_connections: total_proxy,
@@ -1059,6 +1104,8 @@ impl HubRouter {
closing_proxy_connections,
draining_proxy_connections,
soft_avoid_proxy_connections,
protocol_v1_proxy_connections,
protocol_v2_proxy_connections,
nodes,
active_streams: self.local_streams.len(),
outbound_queue_depth_total,
@@ -1067,6 +1114,8 @@ impl HubRouter {
outbound_queue_rejected_full_total,
outbound_queue_rejected_closed_total,
proxy_connection_congested_total,
proxy_connection_write_latency_last_us_max: write_latency_last_us_max,
proxy_connection_write_latency_ewma_us_max: write_latency_ewma_us_max,
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),
@@ -1095,6 +1144,8 @@ pub struct HubStats {
pub closing_proxy_connections: usize,
pub draining_proxy_connections: usize,
pub soft_avoid_proxy_connections: usize,
pub protocol_v1_proxy_connections: usize,
pub protocol_v2_proxy_connections: usize,
pub nodes: usize,
pub active_streams: usize,
pub outbound_queue_depth_total: usize,
@@ -1103,6 +1154,8 @@ pub struct HubStats {
pub outbound_queue_rejected_full_total: u64,
pub outbound_queue_rejected_closed_total: u64,
pub proxy_connection_congested_total: u64,
pub proxy_connection_write_latency_last_us_max: u64,
pub proxy_connection_write_latency_ewma_us_max: u64,
pub soft_avoid_selection_total: u64,
pub selection_retry_total: u64,
pub selection_unavailable_total: u64,
@@ -1141,6 +1194,18 @@ impl HubStats {
MetricKind::Gauge,
self.soft_avoid_proxy_connections as u64,
),
MetricSample::new(
"tunnel_proxy_connections_protocol_v1",
"Current number of connected proxy sockets still using tunnel protocol v1.",
MetricKind::Gauge,
self.protocol_v1_proxy_connections as u64,
),
MetricSample::new(
"tunnel_proxy_connections_protocol_v2",
"Current number of connected proxy sockets using tunnel protocol v2.",
MetricKind::Gauge,
self.protocol_v2_proxy_connections as u64,
),
MetricSample::new(
"tunnel_nodes",
"Current number of connected logical nodes.",
@@ -1189,6 +1254,18 @@ impl HubStats {
MetricKind::Counter,
self.proxy_connection_congested_total,
),
MetricSample::new(
"tunnel_proxy_connection_write_latency_last_us_max",
"Maximum observed last write latency across proxy connections in microseconds.",
MetricKind::Gauge,
self.proxy_connection_write_latency_last_us_max,
),
MetricSample::new(
"tunnel_proxy_connection_write_latency_ewma_us_max",
"Maximum observed write latency EWMA across proxy connections in microseconds.",
MetricKind::Gauge,
self.proxy_connection_write_latency_ewma_us_max,
),
MetricSample::new(
"tunnel_proxy_soft_avoid_selection_total",
"Total number of times the scheduler had to pick a high-pressure proxy connection.",
@@ -1250,6 +1327,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
));
hub.register_proxy(proxy);
@@ -1285,6 +1363,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
));
hub.register_proxy(proxy);
@@ -1303,6 +1382,7 @@ mod tests {
};
let first_header = protocol::FrameHeader::parse(&first).expect("first body header");
assert_eq!(first_header.msg_type, protocol::REQUEST_BODY);
assert_eq!(first_header.flags & protocol::FLAG_GZIP_COMPRESSED, 0);
assert_eq!(first_header.flags & protocol::FLAG_END_STREAM, 0);
let second = match proxy_rx.try_recv().expect("second body frame") {
@@ -1311,6 +1391,7 @@ mod tests {
};
let second_header = protocol::FrameHeader::parse(&second).expect("second body header");
assert_eq!(second_header.msg_type, protocol::REQUEST_BODY);
assert_eq!(second_header.flags & protocol::FLAG_GZIP_COMPRESSED, 0);
assert_ne!(second_header.flags & protocol::FLAG_END_STREAM, 0);
}
@@ -1327,6 +1408,7 @@ mod tests {
proxy_one_tx,
proxy_one_close_tx,
16,
2,
));
hub.register_proxy(Arc::clone(&proxy_one));
@@ -1339,6 +1421,7 @@ mod tests {
proxy_two_tx,
proxy_two_close_tx,
16,
2,
));
hub.register_proxy(Arc::clone(&proxy_two));
@@ -1387,6 +1470,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
));
hub.register_proxy(proxy);
@@ -1414,6 +1498,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
));
hub.register_proxy(Arc::clone(&proxy));

2
apps/aether-gateway/src/tunnel/embedded/local_relay.rs
View File

@@ -629,6 +629,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let meta = protocol::RequestMeta {
@@ -743,6 +744,7 @@ mod tests {
proxy_tx,
proxy_close_tx,
16,
2,
)));
let meta = protocol::RequestMeta {

30
apps/aether-gateway/src/tunnel/embedded/mod.rs
View File

@@ -211,6 +211,7 @@ pub async fn ws_proxy(
.to_string();
let max_streams = resolve_proxy_max_streams(&headers, state.max_streams);
let protocol_version = resolve_proxy_protocol_version(&headers);
if node_id.is_empty() {
warn!("proxy connection rejected: missing X-Node-ID header");
@@ -251,6 +252,7 @@ pub async fn ws_proxy(
node_id,
node_name,
max_streams,
protocol_version,
state.proxy_conn_cfg,
)
.await
@@ -268,11 +270,20 @@ fn resolve_proxy_max_streams(headers: &HeaderMap, fallback: usize) -> usize {
.clamp(1, 2048)
}
fn resolve_proxy_protocol_version(headers: &HeaderMap) -> u8 {
headers
.get(aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u8>().ok())
.filter(|value| *value >= 1)
.unwrap_or(1)
}
#[cfg(test)]
mod tests {
use axum::http::{HeaderMap, HeaderValue};
use super::resolve_proxy_max_streams;
use super::{resolve_proxy_max_streams, resolve_proxy_protocol_version};
#[test]
fn proxy_max_streams_honors_small_advertised_capacity() {
@@ -289,4 +300,21 @@ mod tests {
assert_eq!(resolve_proxy_max_streams(&headers, 128), 2048);
}
#[test]
fn proxy_protocol_version_defaults_to_v1_when_header_missing() {
let headers = HeaderMap::new();
assert_eq!(resolve_proxy_protocol_version(&headers), 1);
}
#[test]
fn proxy_protocol_version_reads_advertised_version() {
let mut headers = HeaderMap::new();
headers.insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
HeaderValue::from_static("2"),
);
assert_eq!(resolve_proxy_protocol_version(&headers), 2);
}
}

5
apps/aether-gateway/src/tunnel/embedded/protocol.rs
View File

@@ -12,3 +12,8 @@ pub fn compress_payload(payload: &[u8]) -> Result<(Vec<u8>, u8), std::io::Error>
aether_contracts::tunnel::compress_payload(Bytes::copy_from_slice(payload));
Ok((compressed.to_vec(), flags))
}
pub fn raw_payload(payload: &[u8]) -> (Vec<u8>, u8) {
let (payload, flags) = aether_contracts::tunnel::raw_payload(Bytes::copy_from_slice(payload));
(payload.to_vec(), flags)
}

7
apps/aether-gateway/src/tunnel/embedded/proxy_conn.rs
View File

@@ -23,6 +23,7 @@ pub async fn handle_proxy_connection(
node_id: String,
node_name: String,
max_streams: usize,
protocol_version: u8,
cfg: ConnConfig,
) {
let conn_id = hub.alloc_conn_id();
@@ -38,6 +39,7 @@ pub async fn handle_proxy_connection(
tx,
close_tx,
max_streams,
protocol_version,
));
hub.register_proxy(conn.clone());
@@ -55,12 +57,15 @@ pub async fn handle_proxy_connection(
Message::Binary(b) => b.len(),
_ => 0,
};
let send_started_at = std::time::Instant::now();
let send_result = tokio::time::timeout(
Duration::from_secs(15),
ws_tx.send(msg),
).await;
match send_result {
Ok(Ok(())) => {}
Ok(Ok(())) => {
writer_conn.record_write_latency(send_started_at.elapsed());
}
Ok(Err(e)) => {
let snapshot = writer_conn.outbound.snapshot();
warn!(

4
apps/aether-proxy/Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "aether-proxy"
version = "0.3.11"
version = "0.3.12"
edition = "2021"
description = "Tunnel proxy for Aether"
@@ -9,6 +9,7 @@ aether-contracts.workspace = true
aether-http.workspace = true
aether-runtime.workspace = true
aether-runtime-state.workspace = true
axum.workspace = true
tokio = { version = "1", features = ["full"] }
reqwest.workspace = true
hyper = { version = "1", features = ["client", "http1", "http2"] }
@@ -43,4 +44,3 @@ webpki-roots = "0.26"
[dev-dependencies]
aether-gateway.workspace = true
axum.workspace = true

302
apps/aether-proxy/src/app.rs
View File

@@ -6,9 +6,14 @@ use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use aether_http::{jittered_delay_for_retry, HttpRetryConfig};
use aether_runtime::{init_reloadable_service_tracing, wait_for_shutdown_signal, ConcurrencyGate};
use aether_runtime::{
init_reloadable_service_tracing, prometheus_response, wait_for_shutdown_signal, ConcurrencyGate,
};
use aether_runtime_state::{RedisClientConfig, RuntimeSemaphoreConfig, RuntimeState};
use arc_swap::ArcSwap;
use axum::extract::State as AxumState;
use axum::routing::get;
use axum::{Json, Router};
use tokio::sync::{watch, Mutex};
use tokio::task::JoinHandle;
use tracing::{error, info, warn};
@@ -24,10 +29,14 @@ use crate::{hardware, target_filter, tunnel};
type TaskHandles = Arc<Mutex<Vec<JoinHandle<()>>>>;
const AUTO_STREAM_LIMIT_MIN: usize = 16;
const AUTO_STREAM_LIMIT_MAX: usize = 512;
const AUTO_STREAM_LIMIT_PER_CPU: u64 = 64;
// Keep the automatic fallback large enough for real load while still
// protecting tiny nodes from overcommitting by default.
const AUTO_STREAM_LIMIT_MAX: usize = 2048;
// Bias toward throughput: let a 16-core class box auto-land near 2k streams,
// then rely on FD / memory estimates and the hard cap to keep smaller hosts safe.
const AUTO_STREAM_LIMIT_PER_CPU: u64 = 128;
const AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM: u64 = 4;
const AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR: u64 = 40;
const AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR: u64 = 12;
#[derive(Debug, Clone, Copy)]
struct TunnelPoolPolicy {
@@ -75,6 +84,12 @@ struct ManagedTunnel {
draining: bool,
}
#[derive(Clone)]
struct DiagnosticsState {
state: Arc<AppState>,
server_contexts: Arc<Mutex<Vec<Arc<ServerContext>>>>,
}
/// Run the full application lifecycle after config has been parsed.
pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Result<()> {
config.validate()?;
@@ -275,6 +290,20 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
// Shutdown signal channel
let (shutdown_tx, shutdown_rx) = watch::channel(false);
let diagnostics_handle = if let Some(bind_addr) = state.config.diagnostics_bind {
let listener = tokio::net::TcpListener::bind(bind_addr).await?;
Some(spawn_diagnostics_server(
listener,
DiagnosticsState {
state: Arc::clone(&state),
server_contexts: Arc::clone(&server_contexts),
},
shutdown_rx.clone(),
)?)
} else {
None
};
info!(
active_servers = server_contexts.lock().await.len(),
"running in tunnel mode"
@@ -314,6 +343,9 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
wait_for_shutdown().await;
info!("shutdown signal received, cleaning up...");
let _ = shutdown_tx.send(true);
if let Some(handle) = diagnostics_handle {
let _ = handle.await;
}
await_all_handles(&retry_handles).await;
// Graceful unregister from all servers (including retry-registered ones)
@@ -335,6 +367,160 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
Ok(())
}
fn spawn_diagnostics_server(
listener: tokio::net::TcpListener,
diagnostics_state: DiagnosticsState,
mut shutdown: watch::Receiver<bool>,
) -> std::io::Result<JoinHandle<()>> {
let bind_addr = listener.local_addr()?;
let app = Router::new()
.route("/health", get(diagnostics_health))
.route("/metrics", get(diagnostics_metrics))
.route("/stats", get(diagnostics_stats))
.with_state(diagnostics_state);
info!(bind = %bind_addr, "proxy diagnostics server listening");
Ok(tokio::spawn(async move {
let graceful_shutdown = async move {
while !*shutdown.borrow() {
if shutdown.changed().await.is_err() {
break;
}
}
};
if let Err(error) = axum::serve(listener, app)
.with_graceful_shutdown(graceful_shutdown)
.await
{
error!(error = %error, "proxy diagnostics server exited with error");
}
}))
}
async fn diagnostics_health(
AxumState(diagnostics): AxumState<DiagnosticsState>,
) -> Json<serde_json::Value> {
let servers = diagnostics.server_contexts.lock().await.clone();
let active_connections = servers
.iter()
.map(|server| server.active_connections.load(Ordering::Acquire))
.sum::<u64>();
let stream_concurrency = diagnostics
.state
.stream_concurrency_snapshot()
.map(concurrency_snapshot_json);
let distributed_stream_concurrency =
distributed_stream_concurrency_json(&diagnostics.state).await;
Json(serde_json::json!({
"status": "ok",
"service": "aether-proxy",
"version": env!("CARGO_PKG_VERSION"),
"protocol_version": aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION,
"server_count": servers.len(),
"active_connections": active_connections,
"stream_concurrency": stream_concurrency,
"distributed_stream_concurrency": distributed_stream_concurrency,
}))
}
async fn diagnostics_metrics(
AxumState(diagnostics): AxumState<DiagnosticsState>,
) -> impl axum::response::IntoResponse {
let mut samples = diagnostics.state.metric_samples().await;
let servers = diagnostics.server_contexts.lock().await.clone();
for server in servers {
samples.extend(server.metric_samples());
}
prometheus_response(&samples)
}
async fn diagnostics_stats(
AxumState(diagnostics): AxumState<DiagnosticsState>,
) -> Json<serde_json::Value> {
let servers = diagnostics.server_contexts.lock().await.clone();
let active_connections = servers
.iter()
.map(|server| server.active_connections.load(Ordering::Acquire))
.sum::<u64>();
let server_stats = servers
.iter()
.map(|server| diagnostics_server_stats(server))
.collect::<Vec<_>>();
let stream_concurrency = diagnostics
.state
.stream_concurrency_snapshot()
.map(concurrency_snapshot_json);
let distributed_stream_concurrency =
distributed_stream_concurrency_json(&diagnostics.state).await;
Json(serde_json::json!({
"status": "ok",
"service": "aether-proxy",
"version": env!("CARGO_PKG_VERSION"),
"protocol_version": aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION,
"capacities": {
"max_concurrent_connections": diagnostics.state.config.max_concurrent_connections,
"max_in_flight_streams": diagnostics.state.config.max_in_flight_streams,
"distributed_stream_limit": diagnostics.state.config.distributed_stream_limit,
"tunnel_max_streams": diagnostics.state.config.tunnel_max_streams,
"tunnel_connections": diagnostics.state.config.tunnel_connections,
"tunnel_connections_max": diagnostics.state.config.tunnel_connections_max,
"diagnostics_bind": diagnostics.state.config.diagnostics_bind.map(|addr| addr.to_string()),
},
"server_count": servers.len(),
"active_connections": active_connections,
"stream_concurrency": stream_concurrency,
"distributed_stream_concurrency": distributed_stream_concurrency,
"resource_usage": diagnostics.state.resource_monitor.snapshot(),
"servers": server_stats,
}))
}
fn diagnostics_server_stats(server: &ServerContext) -> serde_json::Value {
let node_id = server.node_id.read().unwrap().clone();
let dynamic = server.dynamic.load();
serde_json::json!({
"server": server.server_label.clone(),
"node_id": node_id,
"node_name": dynamic.node_name.clone(),
"active_connections": server.active_connections.load(Ordering::Acquire),
"proxy_metrics": server.metrics.snapshot(),
"tunnel_metrics": server.tunnel_metrics.snapshot(),
"recent_tunnel_errors": server.tunnel_metrics.recent_errors(16),
})
}
fn concurrency_snapshot_json(snapshot: aether_runtime::ConcurrencySnapshot) -> serde_json::Value {
serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected_total": snapshot.rejected,
})
}
fn runtime_semaphore_snapshot_json(
snapshot: aether_runtime_state::RuntimeSemaphoreSnapshot,
) -> serde_json::Value {
serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected_total": snapshot.rejected,
})
}
async fn distributed_stream_concurrency_json(state: &AppState) -> Option<serde_json::Value> {
match state.distributed_stream_concurrency_snapshot().await {
Ok(Some(snapshot)) => Some(runtime_semaphore_snapshot_json(snapshot)),
Ok(None) => None,
Err(error) => Some(serde_json::json!({ "error": error.to_string() })),
}
}
#[allow(clippy::too_many_arguments)]
async fn spawn_registration_recovery_tasks(
state: Arc<AppState>,
@@ -771,7 +957,7 @@ mod tests {
use axum::extract::State as AxumState;
use axum::http::StatusCode as AxumStatusCode;
use axum::routing::post;
use axum::routing::{get, post};
use axum::Router;
use serde_json::json;
@@ -845,6 +1031,72 @@ mod tests {
gateway_handle.abort();
}
#[tokio::test]
async fn diagnostics_routes_report_health_metrics_and_stats() {
ensure_rustls_provider();
let state = sample_state(sample_config("https://aether.example.com"));
let server = sample_registered_server(&state, "server", "node-diagnostics");
let server_contexts = Arc::new(Mutex::new(vec![server]));
let router = Router::new()
.route("/health", get(diagnostics_health))
.route("/metrics", get(diagnostics_metrics))
.route("/stats", get(diagnostics_stats))
.with_state(DiagnosticsState {
state: Arc::clone(&state),
server_contexts,
});
let port = reserve_local_port().expect("diagnostics port should reserve");
let handle = spawn_router_on_port(port, router)
.await
.expect("diagnostics test server should start");
let client = reqwest::Client::new();
let base_url = format!("http://127.0.0.1:{port}");
let health: serde_json::Value = client
.get(format!("{base_url}/health"))
.send()
.await
.expect("health request should send")
.error_for_status()
.expect("health response should be success")
.json()
.await
.expect("health response should parse");
assert_eq!(health["status"], "ok");
assert_eq!(health["service"], "aether-proxy");
assert_eq!(health["server_count"], 1);
let metrics = client
.get(format!("{base_url}/metrics"))
.send()
.await
.expect("metrics request should send")
.error_for_status()
.expect("metrics response should be success")
.text()
.await
.expect("metrics response should read");
assert!(metrics.contains("service_up{service=\"aether-proxy\"} 1"));
assert!(metrics.contains("proxy_active_connections{server=\"server\"} 0"));
let stats: serde_json::Value = client
.get(format!("{base_url}/stats"))
.send()
.await
.expect("stats request should send")
.error_for_status()
.expect("stats response should be success")
.json()
.await
.expect("stats response should parse");
assert_eq!(stats["status"], "ok");
assert_eq!(stats["protocol_version"], 2);
assert_eq!(stats["servers"][0]["node_id"], "node-diagnostics");
handle.abort();
}
#[test]
fn desired_tunnel_connections_expands_when_load_crosses_high_water() {
let policy = TunnelPoolPolicy {
@@ -892,6 +1144,19 @@ mod tests {
assert_eq!(auto_max_in_flight_streams(&hw), 45);
}
#[test]
fn auto_stream_limit_scales_to_high_band_on_mid_size_nodes() {
let hw = HardwareInfo {
cpu_cores: 16,
total_memory_mb: 65_536,
os_info: "test".to_string(),
fd_limit: 1_048_576,
estimated_max_concurrency: 500_000,
};
assert_eq!(auto_max_in_flight_streams(&hw), AUTO_STREAM_LIMIT_MAX);
}
#[test]
fn auto_stream_limit_caps_large_nodes() {
let hw = HardwareInfo {
@@ -1015,6 +1280,31 @@ mod tests {
})
}
fn sample_registered_server(
state: &Arc<ProxyAppState>,
label: &str,
node_id: &str,
) -> Arc<ServerContext> {
let entry = ServerEntry {
aether_url: state.config.aether_url.clone(),
management_token: state.config.management_token.clone(),
node_name: Some(state.config.node_name.clone()),
};
let client = Arc::new(AetherClient::new(
&state.config,
&state.config.aether_url,
&state.config.management_token,
));
build_server_context(
&state.config,
label,
&entry,
client,
&state.config.node_name,
node_id.to_string(),
)
}
fn sample_config(aether_url: &str) -> Config {
Config {
aether_url: aether_url.to_string(),
@@ -1036,6 +1326,7 @@ mod tests {
aether_retry_max_attempts: 1,
aether_retry_base_delay_ms: 50,
aether_retry_max_delay_ms: 100,
diagnostics_bind: None,
max_concurrent_connections: None,
max_in_flight_streams: None,
distributed_stream_limit: None,
@@ -1053,6 +1344,7 @@ mod tests {
upstream_tcp_nodelay: true,
upstream_proxy_url: None,
redirect_replay_budget_bytes: DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES,
emit_proxy_timing_header: true,
log_level: "info".to_string(),
log_destination: ProxyLogDestinationArg::Stdout,
log_dir: None,

28
apps/aether-proxy/src/config.rs
View File

@@ -1,4 +1,5 @@
use std::fmt;
use std::net::SocketAddr;
use std::path::Path;
use std::time::Duration;
@@ -372,6 +373,11 @@ pub struct Config {
)]
pub aether_retry_max_delay_ms: u64,
/// Optional local diagnostics listener for /health, /metrics, and /stats.
/// Bind only to loopback addresses, for example 127.0.0.1:9311.
#[arg(long, env = "AETHER_PROXY_DIAGNOSTICS_BIND")]
pub diagnostics_bind: Option<SocketAddr>,
/// Maximum concurrent TCP connections (defaults to hardware estimate)
#[arg(long, env = "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS")]
pub max_concurrent_connections: Option<u64>,
@@ -479,6 +485,14 @@ pub struct Config {
)]
pub redirect_replay_budget_bytes: usize,
/// Emit detailed x-proxy-timing headers on tunneled upstream responses.
#[arg(
long,
env = "AETHER_PROXY_EMIT_PROXY_TIMING_HEADER",
default_value_t = true
)]
pub emit_proxy_timing_header: bool,
/// Log level (trace, debug, info, warn, error)
#[arg(long, env = "AETHER_PROXY_LOG_LEVEL", default_value = "info")]
pub log_level: String,
@@ -686,6 +700,11 @@ impl Config {
if self.aether_retry_max_attempts == 0 {
anyhow::bail!("aether_retry_max_attempts must be >= 1");
}
if let Some(addr) = self.diagnostics_bind {
if !addr.ip().is_loopback() {
anyhow::bail!("diagnostics_bind must use a loopback address");
}
}
if self.upstream_connect_timeout_secs == 0 {
anyhow::bail!("upstream_connect_timeout_secs must be > 0");
}
@@ -886,6 +905,8 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_retry_max_delay_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub diagnostics_bind: Option<SocketAddr>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_concurrent_connections: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub dns_cache_ttl_secs: Option<u64>,
@@ -910,6 +931,8 @@ pub struct ConfigFile {
)]
pub redirect_replay_budget_bytes: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub emit_proxy_timing_header: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub log_level: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub log_destination: Option<ProxyLogDestinationArg>,
@@ -1048,6 +1071,7 @@ impl ConfigFile {
"AETHER_PROXY_AETHER_RETRY_MAX_DELAY_MS",
self.aether_retry_max_delay_ms
);
set!("AETHER_PROXY_DIAGNOSTICS_BIND", self.diagnostics_bind);
set!(
"AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS",
self.max_concurrent_connections
@@ -1079,6 +1103,10 @@ impl ConfigFile {
"AETHER_PROXY_REDIRECT_REPLAY_BUDGET_BYTES",
self.redirect_replay_budget_bytes
);
set!(
"AETHER_PROXY_EMIT_PROXY_TIMING_HEADER",
self.emit_proxy_timing_header
);
set!("AETHER_PROXY_LOG_LEVEL", self.log_level);
set!(
"AETHER_PROXY_LOG_DESTINATION",

255
apps/aether-proxy/src/state.rs
View File

@@ -6,7 +6,10 @@ use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_runtime::{AdmissionPermit, ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot};
use aether_runtime::{
service_up_sample, AdmissionPermit, ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot,
MetricKind, MetricLabel, MetricSample,
};
use aether_runtime_state::{RuntimeSemaphore, RuntimeSemaphoreError, RuntimeSemaphoreSnapshot};
use crate::config::Config;
@@ -59,6 +62,23 @@ pub struct ServerContext {
pub tunnel_metrics: Arc<TunnelMetrics>,
}
impl ServerContext {
pub fn metric_samples(&self) -> Vec<MetricSample> {
let mut samples = self.metrics.to_metric_samples(&self.server_label);
samples.extend(self.tunnel_metrics.to_metric_samples(&self.server_label));
samples.push(
MetricSample::new(
"proxy_active_connections",
"Current number of active tunneled streams handled by this proxy server context.",
MetricKind::Gauge,
self.active_connections.load(Ordering::Acquire),
)
.with_labels(vec![MetricLabel::new("server", self.server_label.clone())]),
);
samples
}
}
/// Aggregate metrics for reporting to Aether.
pub struct ProxyMetrics {
pub total_requests: AtomicU64,
@@ -68,6 +88,43 @@ pub struct ProxyMetrics {
pub failed_requests: AtomicU64,
pub dns_failures: AtomicU64,
pub stream_errors: AtomicU64,
pub slow_requests: AtomicU64,
}
#[derive(Debug, Clone, Copy, Default, serde::Serialize)]
pub struct ProxyMetricsSnapshot {
pub total_requests: u64,
pub total_latency_ns: u64,
pub failed_requests: u64,
pub dns_failures: u64,
pub stream_errors: u64,
pub slow_requests: u64,
}
impl ProxyMetricsSnapshot {
pub fn average_latency_ns(self) -> Option<u64> {
self.total_latency_ns.checked_div(self.total_requests)
}
pub fn average_latency_ms(self) -> Option<f64> {
self.average_latency_ns()
.map(|value| value as f64 / 1_000_000.0)
}
pub fn delta_since(self, baseline: Self) -> Self {
Self {
total_requests: self.total_requests.saturating_sub(baseline.total_requests),
total_latency_ns: self
.total_latency_ns
.saturating_sub(baseline.total_latency_ns),
failed_requests: self
.failed_requests
.saturating_sub(baseline.failed_requests),
dns_failures: self.dns_failures.saturating_sub(baseline.dns_failures),
stream_errors: self.stream_errors.saturating_sub(baseline.stream_errors),
slow_requests: self.slow_requests.saturating_sub(baseline.slow_requests),
}
}
}
impl ProxyMetrics {
@@ -78,6 +135,7 @@ impl ProxyMetrics {
failed_requests: AtomicU64::new(0),
dns_failures: AtomicU64::new(0),
stream_errors: AtomicU64::new(0),
slow_requests: AtomicU64::new(0),
}
}
@@ -88,6 +146,77 @@ impl ProxyMetrics {
self.total_requests.fetch_add(1, Ordering::Release);
self.total_latency_ns.fetch_add(nanos, Ordering::Release);
}
pub fn record_slow_request(&self) {
self.slow_requests.fetch_add(1, Ordering::Release);
}
pub fn snapshot(&self) -> ProxyMetricsSnapshot {
ProxyMetricsSnapshot {
total_requests: self.total_requests.load(Ordering::Acquire),
total_latency_ns: self.total_latency_ns.load(Ordering::Acquire),
failed_requests: self.failed_requests.load(Ordering::Acquire),
dns_failures: self.dns_failures.load(Ordering::Acquire),
stream_errors: self.stream_errors.load(Ordering::Acquire),
slow_requests: self.slow_requests.load(Ordering::Acquire),
}
}
pub fn to_metric_samples(&self, server_label: &str) -> Vec<MetricSample> {
let snapshot = self.snapshot();
let labels = vec![MetricLabel::new("server", server_label)];
vec![
MetricSample::new(
"proxy_requests_total",
"Total number of tunneled upstream requests completed by the proxy.",
MetricKind::Counter,
snapshot.total_requests,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_request_latency_total_ns",
"Cumulative proxy request latency in nanoseconds through upstream response headers.",
MetricKind::Counter,
snapshot.total_latency_ns,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_request_latency_avg_ns",
"Average proxy request latency in nanoseconds through upstream response headers.",
MetricKind::Gauge,
snapshot.average_latency_ns().unwrap_or(0),
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_failed_requests_total",
"Total number of tunneled upstream requests that failed before response headers.",
MetricKind::Counter,
snapshot.failed_requests,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_dns_failures_total",
"Total number of tunneled upstream requests rejected or failed during target validation or DNS.",
MetricKind::Counter,
snapshot.dns_failures,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_stream_errors_total",
"Total number of tunneled response body stream errors.",
MetricKind::Counter,
snapshot.stream_errors,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_slow_requests_total",
"Total number of tunneled requests crossing the proxy slow-request threshold.",
MetricKind::Counter,
snapshot.slow_requests,
)
.with_labels(labels),
]
}
}
const RECENT_TUNNEL_ERROR_CAPACITY: usize = 64;
@@ -106,7 +235,7 @@ pub struct TunnelErrorEvent {
pub operator_action: String,
}
#[derive(Debug, Clone, Copy, Default)]
#[derive(Debug, Clone, Copy, Default, serde::Serialize)]
pub struct TunnelMetricsSnapshot {
pub connect_attempts: u64,
pub connect_successes: u64,
@@ -298,6 +427,104 @@ impl TunnelMetrics {
error_events_total: self.error_events_total.load(Ordering::Acquire),
}
}
pub fn to_metric_samples(&self, server_label: &str) -> Vec<MetricSample> {
let snapshot = self.snapshot();
let labels = vec![MetricLabel::new("server", server_label)];
vec![
MetricSample::new(
"proxy_tunnel_connect_attempts_total",
"Total number of WebSocket tunnel connection attempts.",
MetricKind::Counter,
snapshot.connect_attempts,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_connect_successes_total",
"Total number of successful WebSocket tunnel connections.",
MetricKind::Counter,
snapshot.connect_successes,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_connect_errors_total",
"Total number of WebSocket tunnel connection errors.",
MetricKind::Counter,
snapshot.connect_errors,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_disconnects_total",
"Total number of WebSocket tunnel disconnects.",
MetricKind::Counter,
snapshot.disconnects,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_heartbeat_sent_total",
"Total number of tunnel heartbeats sent.",
MetricKind::Counter,
snapshot.heartbeat_sent,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_heartbeat_ack_total",
"Total number of tunnel heartbeat acknowledgements received.",
MetricKind::Counter,
snapshot.heartbeat_ack,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_heartbeat_rtt_last_ms",
"Last observed tunnel heartbeat round-trip time in milliseconds.",
MetricKind::Gauge,
snapshot.heartbeat_rtt_last_ms,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_heartbeat_rtt_avg_ms",
"Average observed tunnel heartbeat round-trip time in milliseconds.",
MetricKind::Gauge,
snapshot.heartbeat_rtt_avg_ms().unwrap_or(0.0) as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_ws_in_frames_total",
"Total number of WebSocket frames received by the proxy tunnel.",
MetricKind::Counter,
snapshot.ws_in_frames,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_ws_in_bytes_total",
"Total number of WebSocket bytes received by the proxy tunnel.",
MetricKind::Counter,
snapshot.ws_in_bytes,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_ws_out_frames_total",
"Total number of WebSocket frames sent by the proxy tunnel.",
MetricKind::Counter,
snapshot.ws_out_frames,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_ws_out_bytes_total",
"Total number of WebSocket bytes sent by the proxy tunnel.",
MetricKind::Counter,
snapshot.ws_out_bytes,
)
.with_labels(labels.clone()),
MetricSample::new(
"proxy_tunnel_error_events_total",
"Total number of classified tunnel error events recorded by the proxy.",
MetricKind::Counter,
snapshot.error_events_total,
)
.with_labels(labels),
]
}
}
fn now_unix_secs() -> u64 {
@@ -425,6 +652,30 @@ pub enum ProxyAdmissionError {
}
impl AppState {
pub async fn metric_samples(&self) -> Vec<MetricSample> {
let mut samples = vec![service_up_sample("aether-proxy")];
if let Some(snapshot) = self.stream_concurrency_snapshot() {
samples.extend(snapshot.to_metric_samples("proxy_streams"));
}
if let Some(gate) = self.distributed_stream_gate.as_ref() {
match gate.snapshot().await {
Ok(snapshot) => {
samples.extend(snapshot.to_metric_samples("proxy_streams_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", "proxy_streams_distributed")]),
),
}
}
samples
}
pub fn with_stream_concurrency_gate(mut self, gate: Arc<ConcurrencyGate>) -> Self {
self.stream_gate = Some(gate);
self

5
apps/aether-proxy/src/tunnel/client.rs
View File

@@ -12,6 +12,7 @@ use tracing::{debug, info, warn};
use crate::egress_proxy::{connect_target_via_proxy, ProxyConnectOptions, UpstreamProxyConfig};
use crate::state::{AppState, ServerContext};
use aether_contracts::tunnel::{CURRENT_TUNNEL_PROTOCOL_VERSION, TUNNEL_PROTOCOL_VERSION_HEADER};
use super::{dispatcher, heartbeat, writer};
@@ -44,6 +45,10 @@ pub async fn connect_and_run(
"Authorization",
http::HeaderValue::from_str(&format!("Bearer {}", server.management_token))?,
);
headers.insert(
TUNNEL_PROTOCOL_VERSION_HEADER,
http::HeaderValue::from_str(&CURRENT_TUNNEL_PROTOCOL_VERSION.to_string())?,
);
let node_id = server.node_id.read().unwrap().clone();
headers.insert("X-Node-Id", http::HeaderValue::from_str(&node_id)?);
// Use dynamic node_name (may be updated by remote config) instead of

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

@@ -343,7 +343,7 @@ fn stream_frame_dispatch_timeout() -> Duration {
#[cfg(not(test))]
{
Duration::from_secs(5)
Duration::from_millis(500)
}
}

122
apps/aether-proxy/src/tunnel/heartbeat.rs
View File

@@ -13,8 +13,7 @@ use tracing::{debug, info, warn};
use crate::registration::client::RemoteConfig;
use crate::runtime;
use crate::state::AppState;
use crate::state::ServerContext;
use crate::state::{AppState, ProxyMetricsSnapshot, ServerContext};
use super::protocol::{Frame, MsgType};
use super::writer::FrameSender;
@@ -45,7 +44,7 @@ impl HeartbeatHandle {
/// Create a no-op heartbeat handle that silently discards ACKs.
/// Used for non-primary tunnel connections (conn_idx > 0) to avoid
/// resetting shared atomic metrics via `swap(0)`.
/// duplicating heartbeat ACK processing.
pub fn spawn_noop() -> HeartbeatHandle {
let (ack_tx, _) = tokio::sync::mpsc::channel::<Bytes>(1);
// receiver is immediately dropped; on_ack() calls will silently fail
@@ -54,17 +53,15 @@ pub fn spawn_noop() -> HeartbeatHandle {
#[derive(Debug, Clone, Copy, Default)]
struct HeartbeatSnapshot {
requests: u64,
latency_ns: u64,
failed: u64,
dns_failures: u64,
stream_errors: u64,
cumulative: ProxyMetricsSnapshot,
window: ProxyMetricsSnapshot,
}
#[derive(Debug, Clone, Copy)]
struct PendingHeartbeat {
heartbeat_id: u64,
snapshot: HeartbeatSnapshot,
cumulative: ProxyMetricsSnapshot,
sent_at: Option<Instant>,
}
@@ -82,9 +79,10 @@ pub fn spawn(
let initial_interval = Duration::from_secs(server.dynamic.load().heartbeat_interval);
let mut current_interval = initial_interval;
// At most one in-flight heartbeat snapshot is tracked at a time.
// Snapshot is only cleared after receiving an ACK, which avoids losing
// interval counters when ACK/frame delivery is temporarily unstable.
// We keep the last ACKed cumulative snapshot so each payload can
// report both monotonic totals and the delta since the previous ACK.
let mut pending: Option<PendingHeartbeat> = None;
let mut last_acked_snapshot = ProxyMetricsSnapshot::default();
let mut next_heartbeat_id: u64 = 1;
let heartbeat_session_id = format!(
"{}-{}",
@@ -104,15 +102,17 @@ pub fn spawn(
let pending_entry = if let Some(entry) = pending {
entry
} else {
let snap = collect_snapshot(&server);
let cumulative = server.metrics.snapshot();
let id = next_heartbeat_id;
next_heartbeat_id = next_heartbeat_id.wrapping_add(1);
if next_heartbeat_id == 0 {
next_heartbeat_id = 1;
}
let window = cumulative.delta_since(last_acked_snapshot);
let entry = PendingHeartbeat {
heartbeat_id: id,
snapshot: snap,
snapshot: HeartbeatSnapshot { cumulative, window },
cumulative,
sent_at: None,
};
pending = Some(entry);
@@ -128,9 +128,6 @@ pub fn spawn(
).await;
let frame = Frame::control(MsgType::HeartbeatData, payload);
if frame_tx.send(frame).await.is_err() {
if let Some(entry) = pending.take() {
restore_snapshot(&server, entry.snapshot);
}
break; // Writer closed
}
server.tunnel_metrics.record_heartbeat_sent();
@@ -165,6 +162,7 @@ pub fn spawn(
if let Some(sent_at) = entry.sent_at {
server.tunnel_metrics.record_heartbeat_ack(sent_at.elapsed());
}
last_acked_snapshot = entry.cumulative;
pending = None;
}
}
@@ -175,9 +173,6 @@ pub fn spawn(
}
_ = shutdown.changed() => {
debug!("heartbeat task shutting down");
if let Some(entry) = pending.take() {
restore_snapshot(&server, entry.snapshot);
}
break;
}
}
@@ -187,49 +182,6 @@ pub fn spawn(
HeartbeatHandle { ack_tx }
}
fn collect_snapshot(server: &ServerContext) -> HeartbeatSnapshot {
HeartbeatSnapshot {
requests: server.metrics.total_requests.swap(0, Ordering::AcqRel),
latency_ns: server.metrics.total_latency_ns.swap(0, Ordering::AcqRel),
failed: server.metrics.failed_requests.swap(0, Ordering::AcqRel),
dns_failures: server.metrics.dns_failures.swap(0, Ordering::AcqRel),
stream_errors: server.metrics.stream_errors.swap(0, Ordering::AcqRel),
}
}
fn restore_snapshot(server: &ServerContext, snap: HeartbeatSnapshot) {
if snap.requests > 0 {
server
.metrics
.total_requests
.fetch_add(snap.requests, Ordering::Release);
}
if snap.latency_ns > 0 {
server
.metrics
.total_latency_ns
.fetch_add(snap.latency_ns, Ordering::Release);
}
if snap.failed > 0 {
server
.metrics
.failed_requests
.fetch_add(snap.failed, Ordering::Release);
}
if snap.dns_failures > 0 {
server
.metrics
.dns_failures
.fetch_add(snap.dns_failures, Ordering::Release);
}
if snap.stream_errors > 0 {
server
.metrics
.stream_errors
.fetch_add(snap.stream_errors, Ordering::Release);
}
}
async fn build_heartbeat_payload(
state: &AppState,
server: &ServerContext,
@@ -242,12 +194,26 @@ async fn build_heartbeat_payload(
let recent_errors = server.tunnel_metrics.recent_errors(8);
let resource_usage = state.resource_monitor.snapshot();
let avg_latency_ms = if snapshot.requests > 0 {
Some(snapshot.latency_ns as f64 / snapshot.requests as f64 / 1_000_000.0)
} else {
None
};
let cumulative = snapshot.cumulative;
let window = snapshot.window;
let cumulative_metrics = serde_json::json!({
"total_requests": cumulative.total_requests,
"total_latency_ns": cumulative.total_latency_ns,
"avg_latency_ms": cumulative.average_latency_ms(),
"failed_requests": cumulative.failed_requests,
"dns_failures": cumulative.dns_failures,
"stream_errors": cumulative.stream_errors,
"slow_requests": cumulative.slow_requests,
});
let window_metrics = serde_json::json!({
"total_requests": window.total_requests,
"total_latency_ns": window.total_latency_ns,
"avg_latency_ms": window.average_latency_ms(),
"failed_requests": window.failed_requests,
"dns_failures": window.dns_failures,
"stream_errors": window.stream_errors,
"slow_requests": window.slow_requests,
});
let local_admission = state.stream_concurrency_snapshot().map(|snapshot| {
serde_json::json!({
"limit": snapshot.limit,
@@ -284,11 +250,23 @@ async fn build_heartbeat_payload(
"heartbeat_id": heartbeat_id,
"heartbeat_interval": server.dynamic.load().heartbeat_interval,
"active_connections": server.active_connections.load(Ordering::Acquire),
"total_requests": snapshot.requests,
"avg_latency_ms": avg_latency_ms,
"failed_requests": snapshot.failed,
"dns_failures": snapshot.dns_failures,
"stream_errors": snapshot.stream_errors,
"total_requests": cumulative.total_requests,
"avg_latency_ms": cumulative.average_latency_ms(),
"failed_requests": cumulative.failed_requests,
"dns_failures": cumulative.dns_failures,
"stream_errors": cumulative.stream_errors,
"slow_requests": cumulative.slow_requests,
"window_total_requests": window.total_requests,
"window_total_latency_ns": window.total_latency_ns,
"window_avg_latency_ms": window.average_latency_ms(),
"window_failed_requests": window.failed_requests,
"window_dns_failures": window.dns_failures,
"window_stream_errors": window.stream_errors,
"window_slow_requests": window.slow_requests,
"proxy_metrics": {
"cumulative": cumulative_metrics,
"window": window_metrics,
},
"proxy_metadata": {
"version": CURRENT_VERSION,
"admission": admission,

2
apps/aether-proxy/src/tunnel/mod.rs
View File

@@ -509,6 +509,7 @@ mod tests {
aether_retry_max_attempts: 1,
aether_retry_base_delay_ms: 50,
aether_retry_max_delay_ms: 100,
diagnostics_bind: None,
max_concurrent_connections: None,
max_in_flight_streams: None,
distributed_stream_limit: None,
@@ -526,6 +527,7 @@ mod tests {
upstream_tcp_nodelay: true,
upstream_proxy_url: None,
redirect_replay_budget_bytes: crate::config::DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES,
emit_proxy_timing_header: true,
log_level: "info".to_string(),
log_destination: crate::config::ProxyLogDestinationArg::Stdout,
log_dir: None,

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

@@ -10,7 +10,7 @@ use std::sync::Arc;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use aether_runtime::hold_admission_permit_until;
use aether_runtime::{AdmissionPermit, QueueSendError};
use bytes::{Bytes, BytesMut};
use futures_util::stream;
use futures_util::StreamExt;
@@ -24,8 +24,8 @@ use crate::target_filter;
use crate::upstream_client;
use super::protocol::{
compress_payload, decompress_if_gzip, flags, Frame as TunnelFrame, MsgType, RequestMeta,
ResponseMeta,
compress_payload, decompress_if_gzip, flags, raw_payload, Frame as TunnelFrame, MsgType,
RequestMeta, ResponseMeta,
};
use super::writer::FrameSender;
@@ -33,9 +33,11 @@ use super::writer::FrameSender;
const MAX_CHUNK_SIZE: usize = 32 * 1024;
/// Timeout for sending a single frame to the writer channel.
/// If the writer is congested (TCP backpressure), we abandon the stream
/// rather than blocking indefinitely and exhausting the stream pool.
const FRAME_SEND_TIMEOUT: Duration = Duration::from_secs(30);
/// Control frames are allowed a short wait; body frames fail fast.
const CONTROL_FRAME_SEND_TIMEOUT: Duration = Duration::from_millis(250);
const SLOW_STREAM_LOG_THRESHOLD: Duration = Duration::from_secs(2);
const SUCCESS_LOG_SAMPLE_MODULO: u32 = 256;
const REQUEST_BODY_SPOOL_QUEUE_CAPACITY: usize = 64;
/// Minimum allowed upstream request timeout (seconds).
const MIN_TIMEOUT_SECS: u64 = 5;
@@ -203,21 +205,51 @@ fn log_stream_success(ctx: StreamLogContext<'_>, status: u16, duration: Duration
let url = ctx
.url
.expect("successful requests should always have a URL");
info!(
server = %ctx.server.server_label,
stream_id = ctx.stream_id,
method = %ctx.method,
scheme = url.scheme(),
host = request_log_host(url),
port = request_log_port(url),
path = request_log_path(url),
query_present = url.query().is_some(),
status,
duration_ms = duration.as_millis() as u64,
redirect_count = ctx.redirect_count,
request_body_bytes = ctx.request_body_size,
"proxy request completed"
);
let slow = duration >= SLOW_STREAM_LOG_THRESHOLD;
if slow {
ctx.server.metrics.record_slow_request();
}
let sampled = slow
|| ctx.redirect_count > 0
|| ctx.request_body_size >= 1_048_576
|| ctx.stream_id.is_multiple_of(SUCCESS_LOG_SAMPLE_MODULO);
if sampled {
info!(
server = %ctx.server.server_label,
stream_id = ctx.stream_id,
method = %ctx.method,
scheme = url.scheme(),
host = request_log_host(url),
port = request_log_port(url),
path = request_log_path(url),
query_present = url.query().is_some(),
status,
duration_ms = duration.as_millis() as u64,
redirect_count = ctx.redirect_count,
request_body_bytes = ctx.request_body_size,
slow,
sampled,
"proxy request completed"
);
} else {
debug!(
server = %ctx.server.server_label,
stream_id = ctx.stream_id,
method = %ctx.method,
scheme = url.scheme(),
host = request_log_host(url),
port = request_log_port(url),
path = request_log_path(url),
query_present = url.query().is_some(),
status,
duration_ms = duration.as_millis() as u64,
redirect_count = ctx.redirect_count,
request_body_bytes = ctx.request_body_size,
slow,
sampled,
"proxy request completed"
);
}
}
fn log_stream_failure(ctx: StreamLogContext<'_>, error: &str, duration: Duration) {
@@ -458,7 +490,7 @@ fn prepare_request_body(
deadline: Instant,
replay_budget_bytes: usize,
) -> PreparedRequestBody {
let (spool_tx, spool_rx) = mpsc::unbounded_channel();
let (spool_tx, spool_rx) = mpsc::channel(REQUEST_BODY_SPOOL_QUEUE_CAPACITY);
let replay_state = if replay_budget_bytes == 0 {
None
} else {
@@ -558,12 +590,11 @@ fn remaining_timeout(deadline: Instant) -> Option<Duration> {
async fn spool_request_body(
mut body_rx: mpsc::Receiver<TunnelFrame>,
spool_tx: mpsc::UnboundedSender<SpoolBodyEvent>,
mut spool_tx: mpsc::Sender<SpoolBodyEvent>,
replay_state: Option<Arc<RequestBodyReplayState>>,
body_size: Arc<AtomicUsize>,
deadline: Instant,
) {
let mut spool_tx = Some(spool_tx);
loop {
let frame = match recv_body_frame_with_deadline(&mut body_rx, deadline).await {
Ok(frame) => frame,
@@ -571,7 +602,7 @@ async fn spool_request_body(
if let Some(state) = &replay_state {
state.fail(message.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message));
let _ = send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message)).await;
return;
}
};
@@ -580,7 +611,7 @@ async fn spool_request_body(
if let Some(state) = &replay_state {
state.finish();
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::End);
let _ = send_spool_event(&mut spool_tx, SpoolBodyEvent::End).await;
return;
};
@@ -594,7 +625,8 @@ async fn spool_request_body(
if let Some(state) = &replay_state {
state.fail(message.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message));
let _ =
send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message)).await;
return;
}
};
@@ -604,14 +636,22 @@ async fn spool_request_body(
if let Some(state) = &replay_state {
state.push_chunk(payload.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Data(payload));
if send_spool_event(&mut spool_tx, SpoolBodyEvent::Data(payload))
.await
.is_err()
{
if let Some(state) = &replay_state {
state.fail("request body replay channel closed".to_string());
}
return;
}
}
if end_stream {
if let Some(state) = &replay_state {
state.finish();
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::End);
let _ = send_spool_event(&mut spool_tx, SpoolBodyEvent::End).await;
return;
}
}
@@ -621,14 +661,14 @@ async fn spool_request_body(
if let Some(state) = &replay_state {
state.fail(message.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message));
let _ = send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message)).await;
return;
}
MsgType::StreamEnd => {
if let Some(state) = &replay_state {
state.finish();
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::End);
let _ = send_spool_event(&mut spool_tx, SpoolBodyEvent::End).await;
return;
}
_ => continue,
@@ -636,15 +676,11 @@ async fn spool_request_body(
}
}
fn send_spool_event(
spool_tx: &mut Option<mpsc::UnboundedSender<SpoolBodyEvent>>,
async fn send_spool_event(
spool_tx: &mut mpsc::Sender<SpoolBodyEvent>,
event: SpoolBodyEvent,
) {
if let Some(sender) = spool_tx.as_ref() {
if sender.send(event).is_err() {
*spool_tx = None;
}
}
) -> Result<(), ()> {
spool_tx.send(event).await.map_err(|_| ())
}
fn remove_headers_case_insensitive(headers: &mut Vec<(String, String)>, blocked: &[&str]) {
@@ -851,6 +887,7 @@ async fn relay_upstream_response<B>(
request_body_size: &AtomicUsize,
redirect_count: usize,
request_body_mode: &'static str,
emit_proxy_timing_header: bool,
deadline: Instant,
) -> Option<Duration>
where
@@ -882,7 +919,9 @@ where
"mode": "tunnel",
"redirect_count": redirect_count,
});
resp_headers.push(("x-proxy-timing".to_string(), timing.to_string()));
if emit_proxy_timing_header {
resp_headers.push(("x-proxy-timing".to_string(), timing.to_string()));
}
let resp_meta = ResponseMeta {
status,
headers: resp_headers,
@@ -965,7 +1004,7 @@ where
match chunk_result {
Ok(chunk) => {
if chunk.len() <= MAX_CHUNK_SIZE {
let (payload, extra_flags) = compress_payload(chunk);
let (payload, extra_flags) = raw_payload(chunk);
if !send_frame(
frame_tx,
TunnelFrame::new(stream_id, MsgType::ResponseBody, extra_flags, payload),
@@ -991,7 +1030,7 @@ where
while offset < chunk.len() {
let end = (offset + MAX_CHUNK_SIZE).min(chunk.len());
let slice = chunk.slice(offset..end);
let (payload, extra_flags) = compress_payload(slice);
let (payload, extra_flags) = raw_payload(slice);
if !send_frame(
frame_tx,
TunnelFrame::new(
@@ -1142,10 +1181,8 @@ pub async fn handle_stream(
server.active_connections.fetch_add(1, Ordering::Release);
let connect_elapsed = hold_admission_permit_until(permit, async {
handle_stream_inner(&state, &server, stream_id, meta, body_rx, &frame_tx).await
})
.await;
let connect_elapsed =
handle_stream_inner(&state, &server, stream_id, meta, body_rx, &frame_tx, permit).await;
server.active_connections.fetch_sub(1, Ordering::Release);
if let Some(d) = connect_elapsed {
@@ -1155,16 +1192,41 @@ pub async fn handle_stream(
/// Send a frame to the writer with a timeout. Returns false if send failed.
async fn send_frame(tx: &FrameSender, frame: TunnelFrame) -> bool {
match tokio::time::timeout(FRAME_SEND_TIMEOUT, tx.send(frame)).await {
Ok(Ok(())) => true,
Ok(Err(_)) => {
// Channel closed (writer exited)
false
let stream_id = frame.stream_id;
let msg_type = frame.msg_type;
let flags = frame.flags;
let is_body_frame = matches!(
msg_type,
MsgType::RequestBody | MsgType::ResponseBody | MsgType::StreamEnd
);
if is_body_frame {
match tx.try_send(frame) {
Ok(()) => true,
Err(QueueSendError::Full(_)) => {
warn!(
stream_id,
msg_type = ?msg_type,
flags = flags,
"writer channel full for body frame, abandoning stream"
);
false
}
Err(QueueSendError::Closed(_)) => false,
}
Err(_) => {
// Timeout — writer is congested
warn!("frame send timeout (writer congested), abandoning stream");
false
} else {
match tokio::time::timeout(CONTROL_FRAME_SEND_TIMEOUT, tx.send(frame)).await {
Ok(Ok(())) => true,
Ok(Err(_)) => false,
Err(_) => {
warn!(
stream_id,
msg_type = ?msg_type,
flags = flags,
"control frame send timeout (writer congested), abandoning stream"
);
false
}
}
}
}
@@ -1179,6 +1241,7 @@ async fn handle_stream_inner(
meta: RequestMeta,
body_rx: mpsc::Receiver<TunnelFrame>,
frame_tx: &FrameSender,
mut admission_permit: Option<AdmissionPermit>,
) -> Option<Duration> {
let mut current_method: hyper::Method = parse_request_method(&meta.method);
let mut current_url = match url::Url::parse(&meta.url) {
@@ -1341,6 +1404,7 @@ async fn handle_stream_inner(
redirects_followed,
) {
RedirectDecision::Stop => {
drop(admission_permit.take());
return relay_upstream_response(
server,
stream_id,
@@ -1354,6 +1418,7 @@ async fn handle_stream_inner(
request_body_size.as_ref(),
redirects_followed,
request_body_mode,
state.config.emit_proxy_timing_header,
deadline,
)
.await;
@@ -1379,6 +1444,7 @@ async fn handle_stream_inner(
continue;
}
Ok(None) => {
drop(admission_permit.take());
return relay_upstream_response(
server,
stream_id,
@@ -1392,6 +1458,7 @@ async fn handle_stream_inner(
request_body_size.as_ref(),
redirects_followed,
request_body_mode,
state.config.emit_proxy_timing_header,
deadline,
)
.await;
@@ -1433,6 +1500,7 @@ async fn handle_stream_inner(
}
}
drop(admission_permit.take());
return relay_upstream_response(
server,
stream_id,
@@ -1446,6 +1514,7 @@ async fn handle_stream_inner(
request_body_size.as_ref(),
redirects_followed,
request_body_mode,
state.config.emit_proxy_timing_header,
deadline,
)
.await;
@@ -1474,7 +1543,7 @@ fn build_streaming_request_body(
}
fn build_spooled_request_body(
spool_rx: mpsc::UnboundedReceiver<SpoolBodyEvent>,
spool_rx: mpsc::Receiver<SpoolBodyEvent>,
) -> upstream_client::UpstreamRequestBody {
let body_stream = stream::unfold((spool_rx, false), |(mut spool_rx, finished)| async move {
if finished {
@@ -2003,6 +2072,7 @@ mod tests {
&request_body_size,
0,
"empty",
true,
Instant::now(),
)
.await;
@@ -2456,6 +2526,7 @@ mod tests {
aether_retry_max_attempts: 3,
aether_retry_base_delay_ms: 200,
aether_retry_max_delay_ms: 2_000,
diagnostics_bind: None,
max_concurrent_connections: None,
max_in_flight_streams: None,
distributed_stream_limit: None,
@@ -2473,6 +2544,7 @@ mod tests {
upstream_tcp_nodelay: true,
upstream_proxy_url: None,
redirect_replay_budget_bytes: crate::config::DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES,
emit_proxy_timing_header: true,
log_level: "info".to_string(),
log_destination: crate::config::ProxyLogDestinationArg::Stdout,
log_dir: None,