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

Add tunnel overload protection

Browse Source
This commit is contained in:
fawney19
2026-05-12 18:19:29 +08:00
parent 5509f70ad4
commit 63149fe281
4 changed files with 230 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -23,6 +23,12 @@ 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;
const AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM: u64 = 4;
const AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR: u64 = 40;
#[derive(Debug, Clone, Copy)]
struct TunnelPoolPolicy {
min_connections: usize,
@@ -120,13 +126,25 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
// Collect hardware info (once at startup, sent during registration)
let hw_info = hardware::collect();
if config.max_in_flight_streams.is_none() {
let auto = auto_max_in_flight_streams(&hw_info);
config.max_in_flight_streams = Some(auto);
info!(
max_in_flight_streams = auto,
"auto-detected max_in_flight_streams from hardware"
);
}
// Auto-detect tunnel_max_streams from hardware if not explicitly set
if config.tunnel_max_streams.is_none() {
let auto = (hw_info.estimated_max_concurrency / 10).clamp(64, 1024) as u32;
let stream_limit = config
.max_in_flight_streams
.unwrap_or_else(|| auto_max_in_flight_streams(&hw_info));
let auto = stream_limit.clamp(1, 1024) as u32;
config.tunnel_max_streams = Some(auto);
info!(
tunnel_max_streams = auto,
"auto-detected tunnel_max_streams from hardware"
"auto-detected tunnel_max_streams from stream admission limit"
);
}
let tunnel_pool_sizing = config.resolve_tunnel_pool_sizing(&hw_info)?;
@@ -447,6 +465,22 @@ fn registration_retry_policy(config: &Config) -> HttpRetryConfig {
.normalized()
}
fn auto_max_in_flight_streams(hw_info: &crate::hardware::HardwareInfo) -> usize {
let by_cpu = u64::from(hw_info.cpu_cores.max(1)).saturating_mul(AUTO_STREAM_LIMIT_PER_CPU);
let by_memory = hw_info
.total_memory_mb
.max(AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM)
/ AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM;
let by_estimate = hw_info
.estimated_max_concurrency
.max(AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR)
/ AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR;
let raw = by_cpu.min(by_memory).min(by_estimate).max(1);
usize::try_from(raw)
.unwrap_or(AUTO_STREAM_LIMIT_MAX)
.clamp(AUTO_STREAM_LIMIT_MIN, AUTO_STREAM_LIMIT_MAX)
}
fn build_server_context(
config: &Config,
label: &str,
@@ -845,6 +879,32 @@ mod tests {
assert!(!should_scale_down(200, 1, &policy));
}
#[test]
fn auto_stream_limit_is_conservative_for_tiny_nodes() {
let hw = HardwareInfo {
cpu_cores: 1,
total_memory_mb: 183,
os_info: "test".to_string(),
fd_limit: 65_535,
estimated_max_concurrency: 2_000,
};
assert_eq!(auto_max_in_flight_streams(&hw), 45);
}
#[test]
fn auto_stream_limit_caps_large_nodes() {
let hw = HardwareInfo {
cpu_cores: 64,
total_memory_mb: 262_144,
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);
}
async fn wait_for_registered_server(
server_contexts: &Arc<Mutex<Vec<Arc<ServerContext>>>>,
) -> Arc<ServerContext> {

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

@@ -756,7 +756,11 @@ impl Config {
hw_info: &HardwareInfo,
) -> anyhow::Result<TunnelPoolSizing> {
let per_tunnel_capacity = u64::from(self.tunnel_max_streams.unwrap_or(128).max(1));
let estimated = hw_info.estimated_max_concurrency.max(per_tunnel_capacity);
let estimated = self
.max_in_flight_streams
.and_then(|limit| u64::try_from(limit).ok())
.unwrap_or(hw_info.estimated_max_concurrency)
.max(per_tunnel_capacity);
let cpu_soft_cap = u64::from(hw_info.cpu_cores.max(1))
.saturating_mul(AUTO_TUNNEL_CONNECTIONS_PER_CPU_CAP)
.clamp(
@@ -1613,6 +1617,36 @@ node_name = "proxy-test"
assert_eq!(sizing.max_connections, 4);
}
#[test]
fn auto_tunnel_pool_sizing_respects_stream_admission_limit() {
let config = Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
"--tunnel-max-streams",
"45",
"--max-in-flight-streams",
"45",
]);
let hw = HardwareInfo {
cpu_cores: 1,
total_memory_mb: 183,
os_info: "test".to_string(),
fd_limit: 65_535,
estimated_max_concurrency: 2_000,
};
let sizing = config
.resolve_tunnel_pool_sizing(&hw)
.expect("sizing should resolve");
assert_eq!(sizing.initial_connections, 2);
assert_eq!(sizing.max_connections, 4);
}
#[test]
fn explicit_tunnel_connections_keep_fixed_pool_without_max_override() {
let config = Config::parse_from([

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

@@ -838,20 +838,25 @@ async fn execute_upstream_request(
}
#[allow(clippy::too_many_arguments)]
async fn relay_upstream_response(
async fn relay_upstream_response<B>(
server: &ServerContext,
stream_id: u32,
method: &hyper::Method,
request_url: &url::Url,
frame_tx: &FrameSender,
response: hyper::Response<hyper::body::Incoming>,
response: hyper::Response<B>,
total_dns_ms: u64,
total_elapsed: Duration,
request_timing: upstream_client::RequestTiming,
request_body_size: &AtomicUsize,
redirect_count: usize,
request_body_mode: &'static str,
) -> Option<Duration> {
deadline: Instant,
) -> Option<Duration>
where
B: hyper::body::Body<Data = Bytes> + Send + Unpin + 'static,
B::Error: std::fmt::Display,
{
let status = response.status().as_u16();
let ttfb_ms = total_elapsed.as_millis() as u64;
let mut resp_headers: Vec<(String, String)> = Vec::with_capacity(response.headers().len() + 1);
@@ -911,7 +916,52 @@ async fn relay_upstream_response(
}
let mut stream = response.into_body().into_data_stream();
while let Some(chunk_result) = stream.next().await {
loop {
let Some(remaining) = remaining_timeout(deadline) else {
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
let error_message = "upstream response body timeout".to_string();
log_stream_failure(
stream_log_context(
server,
stream_id,
method,
Some(request_url),
redirect_count,
request_body_size.load(Ordering::Relaxed),
),
&error_message,
total_elapsed,
);
send_error(frame_tx, stream_id, &error_message).await;
return Some(total_elapsed);
};
let chunk_result = match tokio::time::timeout(remaining, stream.next()).await {
Ok(chunk_result) => chunk_result,
Err(_) => {
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
let error_message = "upstream response body timeout".to_string();
log_stream_failure(
stream_log_context(
server,
stream_id,
method,
Some(request_url),
redirect_count,
request_body_size.load(Ordering::Relaxed),
),
&error_message,
total_elapsed,
);
send_error(frame_tx, stream_id, &error_message).await;
return Some(total_elapsed);
}
};
let Some(chunk_result) = chunk_result else {
break;
};
match chunk_result {
Ok(chunk) => {
if chunk.len() <= MAX_CHUNK_SIZE {
@@ -1304,6 +1354,7 @@ async fn handle_stream_inner(
request_body_size.as_ref(),
redirects_followed,
request_body_mode,
deadline,
)
.await;
}
@@ -1341,6 +1392,7 @@ async fn handle_stream_inner(
request_body_size.as_ref(),
redirects_followed,
request_body_mode,
deadline,
)
.await;
}
@@ -1394,6 +1446,7 @@ async fn handle_stream_inner(
request_body_size.as_ref(),
redirects_followed,
request_body_mode,
deadline,
)
.await;
}
@@ -1922,6 +1975,47 @@ mod tests {
&& value.starts_with("text/plain")));
}
#[tokio::test]
async fn response_body_timeout_emits_stream_error() {
let state = sample_state(None, None);
let server = sample_server(&state);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let request_url = url::Url::parse("https://example.com/slow").expect("url");
let request_body_size = AtomicUsize::new(0);
let body = Body::from_stream(futures_util::stream::pending::<
Result<Bytes, std::convert::Infallible>,
>());
let response = Response::builder()
.status(StatusCode::OK)
.body(body)
.expect("response");
relay_upstream_response(
&server,
13,
&hyper::Method::GET,
&request_url,
&frame_tx,
response,
0,
Duration::ZERO,
upstream_client::RequestTiming::default(),
&request_body_size,
0,
"empty",
Instant::now(),
)
.await;
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
assert_eq!(result.response.expect("response metadata").status, 200);
assert_eq!(
result.error.as_deref(),
Some("upstream response body timeout")
);
assert_eq!(server.metrics.stream_errors.load(Ordering::Acquire), 1);
}
#[tokio::test]
async fn follows_redirects_when_explicitly_enabled_for_replayable_post_requests() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")