mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
feat(proxy-tunnel): 实现隧道连接池与 TCP 底层优化
Rust 端: - 支持每个 server 多条并行 WebSocket 连接 (tunnel_connections 配置) - 手动控制 TCP 连接: connect/handshake 超时、keepalive、NODELAY (socket2) - 预构建共享 TLS ClientConfig 避免每次重连重新解析根证书 - 增加 stale timeout 检测无数据连接,智能 backoff 按连接存活时长重置 - 每条连接独立 reconnect 计数器,仅主连接 (conn_idx=0) 发送心跳 - dispatcher 的错误帧和 PONG 改用 try_send 避免阻塞读循环 - stream_handler 增加 frame 发送超时保护防止写阻塞 Python 端: - TunnelManager 改为连接池,按 least-loaded 策略分配请求 - handle_incoming_frame 按连接实例路由,unregister 精确移除单条连接 - 心跳和 PONG 回复改为 fire-and-forget 避免阻塞主读循环 - send_frame 增加超时保护防止 TCP 写阻塞级联 - WebSocket 先 accept 再认证,auth 加超时 - 调整 idle timeout (90s) 和 ping 间隔 (15s)
This commit is contained in:
20
aether-proxy/Cargo.lock
generated
20
aether-proxy/Cargo.lock
generated
@@ -27,6 +27,7 @@ dependencies = [
|
|||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"sha2",
|
"sha2",
|
||||||
|
"socket2 0.5.10",
|
||||||
"sysinfo",
|
"sysinfo",
|
||||||
"tar",
|
"tar",
|
||||||
"thiserror 2.0.18",
|
"thiserror 2.0.18",
|
||||||
@@ -36,6 +37,7 @@ dependencies = [
|
|||||||
"tracing",
|
"tracing",
|
||||||
"tracing-subscriber",
|
"tracing-subscriber",
|
||||||
"url",
|
"url",
|
||||||
|
"webpki-roots 0.26.11",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -907,7 +909,7 @@ dependencies = [
|
|||||||
"libc",
|
"libc",
|
||||||
"percent-encoding",
|
"percent-encoding",
|
||||||
"pin-project-lite",
|
"pin-project-lite",
|
||||||
"socket2",
|
"socket2 0.6.2",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tower-service",
|
"tower-service",
|
||||||
"tracing",
|
"tracing",
|
||||||
@@ -1566,7 +1568,7 @@ dependencies = [
|
|||||||
"quinn-udp",
|
"quinn-udp",
|
||||||
"rustc-hash",
|
"rustc-hash",
|
||||||
"rustls",
|
"rustls",
|
||||||
"socket2",
|
"socket2 0.6.2",
|
||||||
"thiserror 2.0.18",
|
"thiserror 2.0.18",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tracing",
|
"tracing",
|
||||||
@@ -1603,7 +1605,7 @@ dependencies = [
|
|||||||
"cfg_aliases",
|
"cfg_aliases",
|
||||||
"libc",
|
"libc",
|
||||||
"once_cell",
|
"once_cell",
|
||||||
"socket2",
|
"socket2 0.6.2",
|
||||||
"tracing",
|
"tracing",
|
||||||
"windows-sys 0.60.2",
|
"windows-sys 0.60.2",
|
||||||
]
|
]
|
||||||
@@ -2149,6 +2151,16 @@ version = "1.15.1"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
|
checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "socket2"
|
||||||
|
version = "0.5.10"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678"
|
||||||
|
dependencies = [
|
||||||
|
"libc",
|
||||||
|
"windows-sys 0.52.0",
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "socket2"
|
name = "socket2"
|
||||||
version = "0.6.2"
|
version = "0.6.2"
|
||||||
@@ -2441,7 +2453,7 @@ dependencies = [
|
|||||||
"parking_lot",
|
"parking_lot",
|
||||||
"pin-project-lite",
|
"pin-project-lite",
|
||||||
"signal-hook-registry",
|
"signal-hook-registry",
|
||||||
"socket2",
|
"socket2 0.6.2",
|
||||||
"tokio-macros",
|
"tokio-macros",
|
||||||
"windows-sys 0.61.2",
|
"windows-sys 0.61.2",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -29,6 +29,8 @@ sysinfo = "0.32"
|
|||||||
libc = "0.2"
|
libc = "0.2"
|
||||||
flate2 = "1"
|
flate2 = "1"
|
||||||
tar = "0.4"
|
tar = "0.4"
|
||||||
|
socket2 = { version = "0.5", features = ["all"] }
|
||||||
|
webpki-roots = "0.26"
|
||||||
|
|
||||||
[profile.release]
|
[profile.release]
|
||||||
lto = true
|
lto = true
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
//! Application lifecycle: initialization, task orchestration, and shutdown.
|
//! Application lifecycle: initialization, task orchestration, and shutdown.
|
||||||
|
|
||||||
use std::sync::atomic::{AtomicU32, AtomicU64};
|
use std::sync::atomic::AtomicU64;
|
||||||
use std::sync::{Arc, RwLock};
|
use std::sync::{Arc, RwLock};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
@@ -65,11 +65,19 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
|||||||
));
|
));
|
||||||
|
|
||||||
// Build reqwest client for tunnel upstream requests (shared).
|
// Build reqwest client for tunnel upstream requests (shared).
|
||||||
let reqwest_client = reqwest::Client::builder()
|
let mut reqwest_builder = reqwest::Client::builder()
|
||||||
.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host)
|
.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host)
|
||||||
.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs))
|
.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs))
|
||||||
.connect_timeout(Duration::from_secs(config.upstream_connect_timeout_secs))
|
.connect_timeout(Duration::from_secs(config.upstream_connect_timeout_secs))
|
||||||
.tcp_nodelay(config.upstream_tcp_nodelay)
|
.tcp_nodelay(config.upstream_tcp_nodelay);
|
||||||
|
|
||||||
|
if config.upstream_tcp_keepalive_secs > 0 {
|
||||||
|
reqwest_builder = reqwest_builder.tcp_keepalive(Some(Duration::from_secs(
|
||||||
|
config.upstream_tcp_keepalive_secs,
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let reqwest_client = reqwest_builder
|
||||||
.build()
|
.build()
|
||||||
.expect("failed to build reqwest client");
|
.expect("failed to build reqwest client");
|
||||||
|
|
||||||
@@ -106,7 +114,6 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
|||||||
dynamic: Arc::new(RwLock::new(DynamicConfig::from_config(&config))),
|
dynamic: Arc::new(RwLock::new(DynamicConfig::from_config(&config))),
|
||||||
active_connections: Arc::new(AtomicU64::new(0)),
|
active_connections: Arc::new(AtomicU64::new(0)),
|
||||||
metrics: Arc::new(ProxyMetrics::new()),
|
metrics: Arc::new(ProxyMetrics::new()),
|
||||||
reconnect_attempts: AtomicU32::new(0),
|
|
||||||
}));
|
}));
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -125,10 +132,12 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Build shared application state
|
// Build shared application state
|
||||||
|
let tunnel_tls_config = Arc::new(crate::tunnel::client::build_tls_config());
|
||||||
let state = Arc::new(AppState {
|
let state = Arc::new(AppState {
|
||||||
config: Arc::new(config),
|
config: Arc::new(config),
|
||||||
dns_cache,
|
dns_cache,
|
||||||
reqwest_client,
|
reqwest_client,
|
||||||
|
tunnel_tls_config,
|
||||||
});
|
});
|
||||||
|
|
||||||
// Shutdown signal channel
|
// Shutdown signal channel
|
||||||
@@ -139,15 +148,18 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
|||||||
"running in tunnel mode"
|
"running in tunnel mode"
|
||||||
);
|
);
|
||||||
|
|
||||||
// Spawn one tunnel task per server
|
// Spawn tunnel connections per server (pool_size connections each)
|
||||||
|
let pool_size = state.config.tunnel_connections.max(1) as usize;
|
||||||
let mut tunnel_handles = Vec::new();
|
let mut tunnel_handles = Vec::new();
|
||||||
for server in &server_contexts {
|
for server in &server_contexts {
|
||||||
let s = Arc::clone(&state);
|
for conn_idx in 0..pool_size {
|
||||||
let srv = Arc::clone(server);
|
let s = Arc::clone(&state);
|
||||||
let rx = shutdown_rx.clone();
|
let srv = Arc::clone(server);
|
||||||
tunnel_handles.push(tokio::spawn(async move {
|
let rx = shutdown_rx.clone();
|
||||||
tunnel::run(&s, &srv, rx).await;
|
tunnel_handles.push(tokio::spawn(async move {
|
||||||
}));
|
tunnel::run(&s, &srv, conn_idx, rx).await;
|
||||||
|
}));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Wait for shutdown signal
|
// Wait for shutdown signal
|
||||||
|
|||||||
@@ -225,6 +225,30 @@ pub struct Config {
|
|||||||
/// Maximum concurrent streams over tunnel (auto-detected from hardware if omitted)
|
/// Maximum concurrent streams over tunnel (auto-detected from hardware if omitted)
|
||||||
#[arg(long, env = "AETHER_PROXY_TUNNEL_MAX_STREAMS")]
|
#[arg(long, env = "AETHER_PROXY_TUNNEL_MAX_STREAMS")]
|
||||||
pub tunnel_max_streams: Option<u32>,
|
pub tunnel_max_streams: Option<u32>,
|
||||||
|
|
||||||
|
/// WebSocket tunnel TCP connect timeout in seconds
|
||||||
|
#[arg(
|
||||||
|
long,
|
||||||
|
env = "AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT",
|
||||||
|
default_value_t = 15
|
||||||
|
)]
|
||||||
|
pub tunnel_connect_timeout_secs: u64,
|
||||||
|
|
||||||
|
/// WebSocket tunnel TCP keepalive in seconds (0 disables)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_TUNNEL_TCP_KEEPALIVE", default_value_t = 30)]
|
||||||
|
pub tunnel_tcp_keepalive_secs: u64,
|
||||||
|
|
||||||
|
/// WebSocket tunnel TCP_NODELAY
|
||||||
|
#[arg(long, env = "AETHER_PROXY_TUNNEL_TCP_NODELAY", default_value_t = true)]
|
||||||
|
pub tunnel_tcp_nodelay: bool,
|
||||||
|
|
||||||
|
/// Tunnel connection staleness timeout in seconds (triggers reconnect if no data received)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_TUNNEL_STALE_TIMEOUT", default_value_t = 45)]
|
||||||
|
pub tunnel_stale_timeout_secs: u64,
|
||||||
|
|
||||||
|
/// Number of parallel WebSocket tunnel connections per server (connection pool)
|
||||||
|
#[arg(long, env = "AETHER_PROXY_TUNNEL_CONNECTIONS", default_value_t = 3)]
|
||||||
|
pub tunnel_connections: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Per-server connection config (used in multi-server TOML `[[servers]]`).
|
/// Per-server connection config (used in multi-server TOML `[[servers]]`).
|
||||||
@@ -306,6 +330,16 @@ pub struct ConfigFile {
|
|||||||
pub tunnel_ping_interval_secs: Option<u64>,
|
pub tunnel_ping_interval_secs: Option<u64>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub tunnel_max_streams: Option<u32>,
|
pub tunnel_max_streams: Option<u32>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tunnel_connect_timeout_secs: Option<u64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tunnel_tcp_keepalive_secs: Option<u64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tunnel_tcp_nodelay: Option<bool>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tunnel_stale_timeout_secs: Option<u64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub tunnel_connections: Option<u32>,
|
||||||
|
|
||||||
/// Multi-server config: each entry connects to a separate Aether instance.
|
/// Multi-server config: each entry connects to a separate Aether instance.
|
||||||
/// When present, top-level aether_url/management_token are ignored for
|
/// When present, top-level aether_url/management_token are ignored for
|
||||||
@@ -548,6 +582,20 @@ impl ConfigFile {
|
|||||||
self.tunnel_ping_interval_secs
|
self.tunnel_ping_interval_secs
|
||||||
);
|
);
|
||||||
set!("AETHER_PROXY_TUNNEL_MAX_STREAMS", self.tunnel_max_streams);
|
set!("AETHER_PROXY_TUNNEL_MAX_STREAMS", self.tunnel_max_streams);
|
||||||
|
set!(
|
||||||
|
"AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT",
|
||||||
|
self.tunnel_connect_timeout_secs
|
||||||
|
);
|
||||||
|
set!(
|
||||||
|
"AETHER_PROXY_TUNNEL_TCP_KEEPALIVE",
|
||||||
|
self.tunnel_tcp_keepalive_secs
|
||||||
|
);
|
||||||
|
set!("AETHER_PROXY_TUNNEL_TCP_NODELAY", self.tunnel_tcp_nodelay);
|
||||||
|
set!(
|
||||||
|
"AETHER_PROXY_TUNNEL_STALE_TIMEOUT",
|
||||||
|
self.tunnel_stale_timeout_secs
|
||||||
|
);
|
||||||
|
set!("AETHER_PROXY_TUNNEL_CONNECTIONS", self.tunnel_connections);
|
||||||
|
|
||||||
// allowed_ports needs special handling (comma-separated)
|
// allowed_ports needs special handling (comma-separated)
|
||||||
if let Some(ref ports) = self.allowed_ports {
|
if let Some(ref ports) = self.allowed_ports {
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
//! Shared application state passed to all subsystems.
|
//! Shared application state passed to all subsystems.
|
||||||
|
|
||||||
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
use std::sync::{Arc, RwLock};
|
use std::sync::{Arc, RwLock};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
@@ -16,6 +16,8 @@ pub struct AppState {
|
|||||||
pub dns_cache: Arc<DnsCache>,
|
pub dns_cache: Arc<DnsCache>,
|
||||||
/// Reqwest client for tunnel upstream requests (shared).
|
/// Reqwest client for tunnel upstream requests (shared).
|
||||||
pub reqwest_client: reqwest::Client,
|
pub reqwest_client: reqwest::Client,
|
||||||
|
/// Shared TLS config for tunnel WebSocket connections (avoids re-parsing root CAs on each reconnect).
|
||||||
|
pub tunnel_tls_config: Arc<rustls::ClientConfig>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Per-server state: one instance per Aether server connection.
|
/// Per-server state: one instance per Aether server connection.
|
||||||
@@ -38,8 +40,6 @@ pub struct ServerContext {
|
|||||||
pub active_connections: Arc<AtomicU64>,
|
pub active_connections: Arc<AtomicU64>,
|
||||||
/// Per-server request/latency metrics.
|
/// Per-server request/latency metrics.
|
||||||
pub metrics: Arc<ProxyMetrics>,
|
pub metrics: Arc<ProxyMetrics>,
|
||||||
/// Reconnect attempt counter (reset on successful connection).
|
|
||||||
pub reconnect_attempts: AtomicU32,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Aggregate metrics for reporting to Aether.
|
/// Aggregate metrics for reporting to Aether.
|
||||||
|
|||||||
@@ -1,13 +1,14 @@
|
|||||||
//! WebSocket tunnel client: connect, authenticate, and run the tunnel.
|
//! WebSocket tunnel client: connect, authenticate, and run the tunnel.
|
||||||
|
|
||||||
use std::sync::atomic::Ordering;
|
use std::sync::atomic::{AtomicU32, Ordering};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use tokio::net::TcpStream;
|
||||||
use tokio::sync::watch;
|
use tokio::sync::watch;
|
||||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||||
use tokio_tungstenite::tungstenite::http;
|
use tokio_tungstenite::tungstenite::http;
|
||||||
use tracing::{debug, info};
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
use crate::state::{AppState, ServerContext};
|
use crate::state::{AppState, ServerContext};
|
||||||
|
|
||||||
@@ -22,16 +23,20 @@ pub enum TunnelOutcome {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Connect to Aether's WebSocket tunnel endpoint and run until disconnected.
|
/// Connect to Aether's WebSocket tunnel endpoint and run until disconnected.
|
||||||
|
///
|
||||||
|
/// `conn_idx` identifies which connection in the pool this is (0-based).
|
||||||
|
/// Only connection 0 sends heartbeats to avoid resetting shared metrics.
|
||||||
pub async fn connect_and_run(
|
pub async fn connect_and_run(
|
||||||
state: &Arc<AppState>,
|
state: &Arc<AppState>,
|
||||||
server: &Arc<ServerContext>,
|
server: &Arc<ServerContext>,
|
||||||
|
conn_idx: usize,
|
||||||
shutdown: &mut watch::Receiver<bool>,
|
shutdown: &mut watch::Receiver<bool>,
|
||||||
) -> Result<TunnelOutcome, anyhow::Error> {
|
) -> Result<TunnelOutcome, anyhow::Error> {
|
||||||
let ws_url = build_tunnel_url(server);
|
let ws_url = build_tunnel_url(server);
|
||||||
info!(url = %ws_url, "connecting tunnel");
|
info!(url = %ws_url, conn = conn_idx, "connecting tunnel");
|
||||||
|
|
||||||
// Build WebSocket request with auth headers
|
// Build WebSocket request with auth headers
|
||||||
let mut request = ws_url.into_client_request()?;
|
let mut request = ws_url.clone().into_client_request()?;
|
||||||
let headers = request.headers_mut();
|
let headers = request.headers_mut();
|
||||||
headers.insert(
|
headers.insert(
|
||||||
"Authorization",
|
"Authorization",
|
||||||
@@ -44,12 +49,59 @@ pub async fn connect_and_run(
|
|||||||
http::HeaderValue::from_str(&server.node_name)?,
|
http::HeaderValue::from_str(&server.node_name)?,
|
||||||
);
|
);
|
||||||
|
|
||||||
// Connect
|
// Parse host:port from URL
|
||||||
let (ws_stream, _response) = tokio_tungstenite::connect_async(request).await?;
|
let uri: http::Uri = ws_url.parse()?;
|
||||||
info!("tunnel connected");
|
let host = uri
|
||||||
|
.host()
|
||||||
|
.ok_or_else(|| anyhow::anyhow!("missing host in tunnel URL"))?;
|
||||||
|
let is_tls = uri.scheme_str() == Some("wss");
|
||||||
|
let port = uri.port_u16().unwrap_or(if is_tls { 443 } else { 80 });
|
||||||
|
|
||||||
// Reset reconnect counter on success
|
// TCP connect with timeout
|
||||||
server.reconnect_attempts.store(0, Ordering::Relaxed);
|
let connect_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
|
||||||
|
let tcp_stream = tokio::time::timeout(connect_timeout, TcpStream::connect((host, port)))
|
||||||
|
.await
|
||||||
|
.map_err(|_| {
|
||||||
|
anyhow::anyhow!(
|
||||||
|
"tunnel TCP connect timeout ({}s)",
|
||||||
|
connect_timeout.as_secs()
|
||||||
|
)
|
||||||
|
})??;
|
||||||
|
|
||||||
|
// Configure TCP parameters via socket2
|
||||||
|
configure_tcp_socket(&tcp_stream, state);
|
||||||
|
|
||||||
|
// WebSocket upgrade (with TLS if wss://)
|
||||||
|
let connector = if is_tls {
|
||||||
|
Some(tokio_tungstenite::Connector::Rustls(Arc::clone(
|
||||||
|
&state.tunnel_tls_config,
|
||||||
|
)))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let handshake_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
|
||||||
|
let (ws_stream, _response) = tokio::time::timeout(
|
||||||
|
handshake_timeout,
|
||||||
|
tokio_tungstenite::client_async_tls_with_config(request, tcp_stream, None, connector),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|_| {
|
||||||
|
anyhow::anyhow!(
|
||||||
|
"tunnel WebSocket handshake timeout ({}s)",
|
||||||
|
handshake_timeout.as_secs()
|
||||||
|
)
|
||||||
|
})??;
|
||||||
|
info!(
|
||||||
|
conn = conn_idx,
|
||||||
|
tcp_keepalive_secs = state.config.tunnel_tcp_keepalive_secs,
|
||||||
|
tcp_nodelay = state.config.tunnel_tcp_nodelay,
|
||||||
|
connect_timeout_secs = state.config.tunnel_connect_timeout_secs,
|
||||||
|
stale_timeout_secs = state.config.tunnel_stale_timeout_secs,
|
||||||
|
"tunnel connected"
|
||||||
|
);
|
||||||
|
|
||||||
|
// NOTE: reconnect_attempts reset is handled by the caller (mod.rs)
|
||||||
|
// based on how long the connection stayed alive.
|
||||||
|
|
||||||
// Split into read/write halves
|
// Split into read/write halves
|
||||||
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
|
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
|
||||||
@@ -58,13 +110,18 @@ pub async fn connect_and_run(
|
|||||||
let ping_interval = Duration::from_secs(state.config.tunnel_ping_interval_secs);
|
let ping_interval = Duration::from_secs(state.config.tunnel_ping_interval_secs);
|
||||||
let (frame_tx, mut writer_handle) = writer::spawn_writer(ws_sink, ping_interval);
|
let (frame_tx, mut writer_handle) = writer::spawn_writer(ws_sink, ping_interval);
|
||||||
|
|
||||||
// Spawn heartbeat task
|
// Spawn heartbeat task (only for primary connection to avoid
|
||||||
let hb_handle = heartbeat::spawn(
|
// resetting shared atomic metrics via swap(0))
|
||||||
Arc::clone(&state.config),
|
let hb_handle = if conn_idx == 0 {
|
||||||
Arc::clone(server),
|
heartbeat::spawn(
|
||||||
frame_tx.clone(),
|
Arc::clone(&state.config),
|
||||||
shutdown.clone(),
|
Arc::clone(server),
|
||||||
);
|
frame_tx.clone(),
|
||||||
|
shutdown.clone(),
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
heartbeat::spawn_noop()
|
||||||
|
};
|
||||||
|
|
||||||
// Run dispatcher (blocks until disconnect or shutdown).
|
// Run dispatcher (blocks until disconnect or shutdown).
|
||||||
// Also watch for writer exit — if the write half dies (e.g. the peer
|
// Also watch for writer exit — if the write half dies (e.g. the peer
|
||||||
@@ -81,7 +138,7 @@ pub async fn connect_and_run(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
_ = &mut writer_handle => {
|
_ = &mut writer_handle => {
|
||||||
debug!("writer task exited, triggering reconnect");
|
warn!("writer task exited, triggering reconnect");
|
||||||
TunnelOutcome::Disconnected
|
TunnelOutcome::Disconnected
|
||||||
}
|
}
|
||||||
_ = shutdown.changed() => {
|
_ = shutdown.changed() => {
|
||||||
@@ -106,9 +163,39 @@ pub async fn connect_and_run(
|
|||||||
Ok(outcome)
|
Ok(outcome)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Configure TCP keepalive and NODELAY on an established socket.
|
||||||
|
fn configure_tcp_socket(stream: &TcpStream, state: &Arc<AppState>) {
|
||||||
|
let sock_ref = socket2::SockRef::from(stream);
|
||||||
|
|
||||||
|
if state.config.tunnel_tcp_keepalive_secs > 0 {
|
||||||
|
let keepalive = socket2::TcpKeepalive::new()
|
||||||
|
.with_time(Duration::from_secs(state.config.tunnel_tcp_keepalive_secs))
|
||||||
|
.with_interval(Duration::from_secs(5))
|
||||||
|
.with_retries(3);
|
||||||
|
if let Err(e) = sock_ref.set_tcp_keepalive(&keepalive) {
|
||||||
|
warn!(error = %e, "failed to set TCP keepalive on tunnel socket");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if state.config.tunnel_tcp_nodelay {
|
||||||
|
if let Err(e) = sock_ref.set_nodelay(true) {
|
||||||
|
warn!(error = %e, "failed to set TCP_NODELAY on tunnel socket");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build rustls ClientConfig with system root certificates.
|
||||||
|
pub fn build_tls_config() -> rustls::ClientConfig {
|
||||||
|
let root_store =
|
||||||
|
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||||
|
rustls::ClientConfig::builder()
|
||||||
|
.with_root_certificates(root_store)
|
||||||
|
.with_no_client_auth()
|
||||||
|
}
|
||||||
|
|
||||||
/// Calculate next reconnect delay with exponential backoff + jitter.
|
/// Calculate next reconnect delay with exponential backoff + jitter.
|
||||||
pub fn next_reconnect_delay(state: &Arc<AppState>, server: &Arc<ServerContext>) -> Duration {
|
pub fn next_reconnect_delay(state: &Arc<AppState>, reconnect_attempts: &AtomicU32) -> Duration {
|
||||||
let attempt = server.reconnect_attempts.fetch_add(1, Ordering::Relaxed);
|
let attempt = reconnect_attempts.fetch_add(1, Ordering::Relaxed);
|
||||||
let base_ms = state.config.tunnel_reconnect_base_ms;
|
let base_ms = state.config.tunnel_reconnect_base_ms;
|
||||||
let max_ms = state.config.tunnel_reconnect_max_ms;
|
let max_ms = state.config.tunnel_reconnect_max_ms;
|
||||||
|
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ use futures_util::StreamExt;
|
|||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
use tokio::task::JoinHandle;
|
use tokio::task::JoinHandle;
|
||||||
use tokio_tungstenite::tungstenite::Message;
|
use tokio_tungstenite::tungstenite::Message;
|
||||||
use tracing::{debug, error, warn};
|
use tracing::{debug, error, info, warn};
|
||||||
|
|
||||||
use crate::state::{AppState, ServerContext};
|
use crate::state::{AppState, ServerContext};
|
||||||
|
|
||||||
@@ -37,11 +37,26 @@ where
|
|||||||
// Track spawned stream handlers so we can wait for them on shutdown
|
// Track spawned stream handlers so we can wait for them on shutdown
|
||||||
let mut handler_handles: Vec<JoinHandle<()>> = Vec::new();
|
let mut handler_handles: Vec<JoinHandle<()>> = Vec::new();
|
||||||
let max_streams = state.config.tunnel_max_streams.unwrap_or(128) as usize;
|
let max_streams = state.config.tunnel_max_streams.unwrap_or(128) as usize;
|
||||||
|
let stale_timeout = Duration::from_secs(state.config.tunnel_stale_timeout_secs);
|
||||||
|
|
||||||
|
// Track last time we received any data to detect stale connections
|
||||||
|
let mut last_data_at = tokio::time::Instant::now();
|
||||||
|
|
||||||
let read_err = loop {
|
let read_err = loop {
|
||||||
let msg_result = match ws_stream.next().await {
|
let msg_result = tokio::select! {
|
||||||
Some(r) => r,
|
msg = ws_stream.next() => {
|
||||||
None => break None, // stream ended
|
match msg {
|
||||||
|
Some(r) => r,
|
||||||
|
None => break None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = tokio::time::sleep_until(last_data_at + stale_timeout) => {
|
||||||
|
warn!(
|
||||||
|
stale_secs = stale_timeout.as_secs(),
|
||||||
|
"tunnel connection stale, no data received"
|
||||||
|
);
|
||||||
|
break None;
|
||||||
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let msg = match msg_result {
|
let msg = match msg_result {
|
||||||
@@ -52,12 +67,15 @@ where
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Any successfully received message proves the connection is alive
|
||||||
|
last_data_at = tokio::time::Instant::now();
|
||||||
|
|
||||||
let data = match msg {
|
let data = match msg {
|
||||||
Message::Binary(data) => Bytes::from(data),
|
Message::Binary(data) => Bytes::from(data),
|
||||||
Message::Ping(_) => continue,
|
Message::Ping(_) => continue,
|
||||||
Message::Pong(_) => continue,
|
Message::Pong(_) => continue,
|
||||||
Message::Close(_) => {
|
Message::Close(_) => {
|
||||||
debug!("received WebSocket close");
|
info!("received WebSocket close");
|
||||||
break None;
|
break None;
|
||||||
}
|
}
|
||||||
_ => continue,
|
_ => continue,
|
||||||
@@ -78,14 +96,13 @@ where
|
|||||||
Ok(m) => m,
|
Ok(m) => m,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
warn!(stream_id = frame.stream_id, error = %e, "invalid request metadata");
|
warn!(stream_id = frame.stream_id, error = %e, "invalid request metadata");
|
||||||
let _ = frame_tx
|
// Use try_send to avoid blocking the read loop
|
||||||
.send(Frame::new(
|
let _ = frame_tx.try_send(Frame::new(
|
||||||
frame.stream_id,
|
frame.stream_id,
|
||||||
MsgType::StreamError,
|
MsgType::StreamError,
|
||||||
0,
|
0,
|
||||||
Bytes::from(format!("invalid request metadata: {e}")),
|
Bytes::from(format!("invalid request metadata: {e}")),
|
||||||
))
|
));
|
||||||
.await;
|
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -95,14 +112,12 @@ where
|
|||||||
stream_id = frame.stream_id,
|
stream_id = frame.stream_id,
|
||||||
"max concurrent streams reached"
|
"max concurrent streams reached"
|
||||||
);
|
);
|
||||||
let _ = frame_tx
|
let _ = frame_tx.try_send(Frame::new(
|
||||||
.send(Frame::new(
|
frame.stream_id,
|
||||||
frame.stream_id,
|
MsgType::StreamError,
|
||||||
MsgType::StreamError,
|
0,
|
||||||
0,
|
Bytes::from("max concurrent streams reached"),
|
||||||
Bytes::from("max concurrent streams reached"),
|
));
|
||||||
))
|
|
||||||
.await;
|
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -147,9 +162,8 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
MsgType::Ping => {
|
MsgType::Ping => {
|
||||||
let _ = frame_tx
|
// Use try_send to avoid blocking the read loop when writer is congested
|
||||||
.send(Frame::control(MsgType::Pong, frame.payload))
|
let _ = frame_tx.try_send(Frame::control(MsgType::Pong, frame.payload));
|
||||||
.await;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
MsgType::HeartbeatAck => {
|
MsgType::HeartbeatAck => {
|
||||||
@@ -157,7 +171,7 @@ where
|
|||||||
}
|
}
|
||||||
|
|
||||||
MsgType::GoAway => {
|
MsgType::GoAway => {
|
||||||
debug!("received GOAWAY");
|
info!("received GOAWAY");
|
||||||
break None;
|
break None;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -28,6 +28,15 @@ 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)`.
|
||||||
|
pub fn spawn_noop() -> HeartbeatHandle {
|
||||||
|
let (ack_tx, _) = tokio::sync::mpsc::channel::<Bytes>(1);
|
||||||
|
// receiver is immediately dropped; on_ack() calls will silently fail
|
||||||
|
HeartbeatHandle { ack_tx }
|
||||||
|
}
|
||||||
|
|
||||||
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
|
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
|
||||||
pub fn spawn(
|
pub fn spawn(
|
||||||
config: Arc<Config>,
|
config: Arc<Config>,
|
||||||
|
|||||||
@@ -5,47 +5,83 @@ pub mod protocol;
|
|||||||
pub mod stream_handler;
|
pub mod stream_handler;
|
||||||
pub mod writer;
|
pub mod writer;
|
||||||
|
|
||||||
|
use std::sync::atomic::{AtomicU32, Ordering};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
use tokio::sync::watch;
|
use tokio::sync::watch;
|
||||||
use tracing::{error, info};
|
use tracing::{error, info};
|
||||||
|
|
||||||
use crate::state::{AppState, ServerContext};
|
use crate::state::{AppState, ServerContext};
|
||||||
|
|
||||||
|
/// Minimum connection duration (seconds) to consider a session "stable".
|
||||||
|
/// If a connection lasts shorter than this, the backoff counter is NOT reset,
|
||||||
|
/// preventing rapid reconnect loops on persistently bad networks.
|
||||||
|
const MIN_STABLE_DURATION: Duration = Duration::from_secs(30);
|
||||||
|
|
||||||
/// Run the tunnel mode main loop (connect, dispatch, reconnect).
|
/// Run the tunnel mode main loop (connect, dispatch, reconnect).
|
||||||
|
///
|
||||||
|
/// `conn_idx` identifies which connection in the pool this is (0-based).
|
||||||
|
/// Only connection 0 sends heartbeats to avoid resetting shared metrics.
|
||||||
pub async fn run(
|
pub async fn run(
|
||||||
state: &Arc<AppState>,
|
state: &Arc<AppState>,
|
||||||
server: &Arc<ServerContext>,
|
server: &Arc<ServerContext>,
|
||||||
|
conn_idx: usize,
|
||||||
mut shutdown: watch::Receiver<bool>,
|
mut shutdown: watch::Receiver<bool>,
|
||||||
) {
|
) {
|
||||||
info!(server = %server.server_label, "starting tunnel");
|
info!(server = %server.server_label, conn = conn_idx, "starting tunnel");
|
||||||
|
|
||||||
|
// Per-connection reconnect counter (avoids N connections interfering
|
||||||
|
// with each other's backoff via the shared ServerContext field).
|
||||||
|
let reconnect_attempts = AtomicU32::new(0);
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
match client::connect_and_run(state, server, &mut shutdown).await {
|
let connect_start = tokio::time::Instant::now();
|
||||||
|
|
||||||
|
match client::connect_and_run(state, server, conn_idx, &mut shutdown).await {
|
||||||
Ok(client::TunnelOutcome::Shutdown) => {
|
Ok(client::TunnelOutcome::Shutdown) => {
|
||||||
info!(server = %server.server_label, "tunnel shut down gracefully");
|
info!(server = %server.server_label, conn = conn_idx, "tunnel shut down gracefully");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
Ok(client::TunnelOutcome::Disconnected) => {
|
Ok(client::TunnelOutcome::Disconnected) => {
|
||||||
info!(server = %server.server_label, "tunnel disconnected, will reconnect");
|
let duration = connect_start.elapsed();
|
||||||
|
if duration >= MIN_STABLE_DURATION {
|
||||||
|
// Stable session -- reset backoff for quick reconnect
|
||||||
|
reconnect_attempts.store(0, Ordering::Relaxed);
|
||||||
|
info!(
|
||||||
|
server = %server.server_label,
|
||||||
|
conn = conn_idx,
|
||||||
|
duration_secs = duration.as_secs(),
|
||||||
|
"tunnel disconnected after stable session"
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
// Short-lived session -- keep backoff increasing
|
||||||
|
info!(
|
||||||
|
server = %server.server_label,
|
||||||
|
conn = conn_idx,
|
||||||
|
duration_secs = duration.as_secs(),
|
||||||
|
"tunnel disconnected quickly, increasing backoff"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
error!(server = %server.server_label, error = %e, "tunnel connection lost");
|
// Connection failed -- keep backoff increasing
|
||||||
|
error!(server = %server.server_label, conn = conn_idx, error = %e, "tunnel connection lost");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if *shutdown.borrow() {
|
if *shutdown.borrow() {
|
||||||
info!(server = %server.server_label, "shutdown requested, not reconnecting");
|
info!(server = %server.server_label, conn = conn_idx, "shutdown requested, not reconnecting");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
let delay = client::next_reconnect_delay(state, server);
|
let delay = client::next_reconnect_delay(state, &reconnect_attempts);
|
||||||
info!(server = %server.server_label, delay_ms = delay.as_millis(), "reconnecting tunnel");
|
info!(server = %server.server_label, conn = conn_idx, delay_ms = delay.as_millis(), "reconnecting tunnel");
|
||||||
|
|
||||||
tokio::select! {
|
tokio::select! {
|
||||||
_ = tokio::time::sleep(delay) => {}
|
_ = tokio::time::sleep(delay) => {}
|
||||||
_ = shutdown.changed() => {
|
_ = shutdown.changed() => {
|
||||||
info!(server = %server.server_label, "shutdown requested during reconnect wait");
|
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during reconnect wait");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -21,6 +21,11 @@ use super::writer::FrameSender;
|
|||||||
/// Maximum response body chunk size per frame (32 KB).
|
/// Maximum response body chunk size per frame (32 KB).
|
||||||
const MAX_CHUNK_SIZE: usize = 32 * 1024;
|
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);
|
||||||
|
|
||||||
/// Handle a single stream: receive body, execute upstream, send response.
|
/// Handle a single stream: receive body, execute upstream, send response.
|
||||||
pub async fn handle_stream(
|
pub async fn handle_stream(
|
||||||
state: Arc<AppState>,
|
state: Arc<AppState>,
|
||||||
@@ -39,6 +44,22 @@ pub async fn handle_stream(
|
|||||||
server.metrics.record_request(start.elapsed());
|
server.metrics.record_request(start.elapsed());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Send a frame to the writer with a timeout. Returns false if send failed.
|
||||||
|
async fn send_frame(tx: &FrameSender, frame: Frame) -> bool {
|
||||||
|
match tokio::time::timeout(FRAME_SEND_TIMEOUT, tx.send(frame)).await {
|
||||||
|
Ok(Ok(())) => true,
|
||||||
|
Ok(Err(_)) => {
|
||||||
|
// Channel closed (writer exited)
|
||||||
|
false
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
// Timeout — writer is congested
|
||||||
|
warn!("frame send timeout (writer congested), abandoning stream");
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async fn handle_stream_inner(
|
async fn handle_stream_inner(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
server: &ServerContext,
|
server: &ServerContext,
|
||||||
@@ -190,14 +211,14 @@ async fn handle_stream_inner(
|
|||||||
headers: resp_headers,
|
headers: resp_headers,
|
||||||
};
|
};
|
||||||
let meta_json = serde_json::to_vec(&resp_meta).unwrap_or_default();
|
let meta_json = serde_json::to_vec(&resp_meta).unwrap_or_default();
|
||||||
let _ = frame_tx
|
if !send_frame(
|
||||||
.send(Frame::new(
|
frame_tx,
|
||||||
stream_id,
|
Frame::new(stream_id, MsgType::ResponseHeaders, 0, meta_json),
|
||||||
MsgType::ResponseHeaders,
|
)
|
||||||
0,
|
.await
|
||||||
meta_json,
|
{
|
||||||
))
|
return;
|
||||||
.await;
|
}
|
||||||
|
|
||||||
// Stream response body
|
// Stream response body
|
||||||
let mut stream = response.bytes_stream();
|
let mut stream = response.bytes_stream();
|
||||||
@@ -205,19 +226,28 @@ async fn handle_stream_inner(
|
|||||||
match chunk_result {
|
match chunk_result {
|
||||||
Ok(chunk) => {
|
Ok(chunk) => {
|
||||||
if chunk.len() <= MAX_CHUNK_SIZE {
|
if chunk.len() <= MAX_CHUNK_SIZE {
|
||||||
// 大多数 chunk 无需分割,直接零拷贝发送
|
if !send_frame(
|
||||||
let _ = frame_tx
|
frame_tx,
|
||||||
.send(Frame::new(stream_id, MsgType::ResponseBody, 0, chunk))
|
Frame::new(stream_id, MsgType::ResponseBody, 0, chunk),
|
||||||
.await;
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return;
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
// 超大 chunk 按 MAX_CHUNK_SIZE 分割(使用 Bytes::slice 避免拷贝)
|
// Split oversized chunks
|
||||||
let mut offset = 0;
|
let mut offset = 0;
|
||||||
while offset < chunk.len() {
|
while offset < chunk.len() {
|
||||||
let end = (offset + MAX_CHUNK_SIZE).min(chunk.len());
|
let end = (offset + MAX_CHUNK_SIZE).min(chunk.len());
|
||||||
let slice = chunk.slice(offset..end);
|
let slice = chunk.slice(offset..end);
|
||||||
let _ = frame_tx
|
if !send_frame(
|
||||||
.send(Frame::new(stream_id, MsgType::ResponseBody, 0, slice))
|
frame_tx,
|
||||||
.await;
|
Frame::new(stream_id, MsgType::ResponseBody, 0, slice),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return;
|
||||||
|
}
|
||||||
offset = end;
|
offset = end;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -231,27 +261,32 @@ async fn handle_stream_inner(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Send STREAM_END
|
// Send STREAM_END
|
||||||
let _ = frame_tx
|
let _ = send_frame(
|
||||||
.send(Frame::new(
|
frame_tx,
|
||||||
|
Frame::new(
|
||||||
stream_id,
|
stream_id,
|
||||||
MsgType::StreamEnd,
|
MsgType::StreamEnd,
|
||||||
flags::END_STREAM,
|
flags::END_STREAM,
|
||||||
Bytes::new(),
|
Bytes::new(),
|
||||||
))
|
),
|
||||||
.await;
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
debug!(stream_id, status, "stream completed");
|
debug!(stream_id, status, "stream completed");
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn send_error(tx: &FrameSender, stream_id: u32, msg: &str) {
|
async fn send_error(tx: &FrameSender, stream_id: u32, msg: &str) {
|
||||||
let _ = tx
|
// Error frames use best-effort delivery — don't block if writer is congested
|
||||||
.send(Frame::new(
|
let _ = send_frame(
|
||||||
|
tx,
|
||||||
|
Frame::new(
|
||||||
stream_id,
|
stream_id,
|
||||||
MsgType::StreamError,
|
MsgType::StreamError,
|
||||||
0,
|
0,
|
||||||
Bytes::from(msg.to_string()),
|
Bytes::from(msg.to_string()),
|
||||||
))
|
),
|
||||||
.await;
|
)
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
fn decompress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
|
fn decompress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
|
||||||
|
|||||||
@@ -23,11 +23,11 @@ router = APIRouter()
|
|||||||
# 单帧最大 64 MB -- AI API 请求体可能包含多张 base64 图片,需要足够余量
|
# 单帧最大 64 MB -- AI API 请求体可能包含多张 base64 图片,需要足够余量
|
||||||
_MAX_FRAME_SIZE = 64 * 1024 * 1024
|
_MAX_FRAME_SIZE = 64 * 1024 * 1024
|
||||||
|
|
||||||
# WebSocket 空闲超时(秒)-- proxy 端 WebSocket ping 间隔 15s + 心跳 30s,180s 提供充足余量
|
# WebSocket 空闲超时(秒)-- 需覆盖客户端 stale_timeout(45s) + 重连延迟(最长30s) 的窗口期
|
||||||
_IDLE_TIMEOUT = 180.0
|
_IDLE_TIMEOUT = 90.0
|
||||||
|
|
||||||
# 服务端应用层 ping 间隔(秒)-- 确保即使 proxy 端心跳延迟,连接也不会因中间代理空闲超时而断开
|
# 服务端应用层 ping 间隔(秒)-- 与客户端 ping 间隔一致,确保高频心跳
|
||||||
_SERVER_PING_INTERVAL = 30.0
|
_SERVER_PING_INTERVAL = 15.0
|
||||||
|
|
||||||
|
|
||||||
async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
|
async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
|
||||||
@@ -78,22 +78,26 @@ async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
|
|||||||
@router.websocket("/api/internal/proxy-tunnel")
|
@router.websocket("/api/internal/proxy-tunnel")
|
||||||
async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
||||||
"""aether-proxy tunnel WebSocket 端点"""
|
"""aether-proxy tunnel WebSocket 端点"""
|
||||||
|
# 先 accept,避免认证(DB/Redis)慢时卡在握手阶段导致网关返回 502。
|
||||||
|
await ws.accept()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
auth = await _authenticate(ws)
|
auth = await asyncio.wait_for(_authenticate(ws), timeout=10.0)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
logger.warning("tunnel auth timeout")
|
||||||
|
await ws.close(code=4002, reason="authentication timeout")
|
||||||
|
return
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning("tunnel auth error: {}", e)
|
logger.warning("tunnel auth error: {}", e)
|
||||||
await ws.accept()
|
|
||||||
await ws.close(code=4002, reason="authentication error")
|
await ws.close(code=4002, reason="authentication error")
|
||||||
return
|
return
|
||||||
|
|
||||||
if not auth:
|
if not auth:
|
||||||
await ws.accept()
|
|
||||||
await ws.close(code=4001, reason="unauthorized")
|
await ws.close(code=4001, reason="unauthorized")
|
||||||
return
|
return
|
||||||
|
|
||||||
node_id: str = auth[0]
|
node_id: str = auth[0]
|
||||||
node_name: str = auth[1]
|
node_name: str = auth[1]
|
||||||
await ws.accept()
|
|
||||||
|
|
||||||
manager = get_tunnel_manager()
|
manager = get_tunnel_manager()
|
||||||
conn = TunnelConnection(node_id, node_name, ws)
|
conn = TunnelConnection(node_id, node_name, ws)
|
||||||
@@ -129,7 +133,7 @@ async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
|||||||
logger.warning("tunnel frame decode error from {}: {}", node_id, e)
|
logger.warning("tunnel frame decode error from {}: {}", node_id, e)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
await manager.handle_incoming_frame(node_id, frame)
|
await manager.handle_incoming_frame(conn, frame)
|
||||||
|
|
||||||
except WebSocketDisconnect:
|
except WebSocketDisconnect:
|
||||||
logger.info("tunnel WebSocket disconnected: node_id={}", node_id)
|
logger.info("tunnel WebSocket disconnected: node_id={}", node_id)
|
||||||
@@ -137,8 +141,11 @@ async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
|||||||
logger.error("tunnel WebSocket error for node_id={}: {}", node_id, e)
|
logger.error("tunnel WebSocket error for node_id={}: {}", node_id, e)
|
||||||
finally:
|
finally:
|
||||||
ping_task.cancel()
|
ping_task.cancel()
|
||||||
manager.unregister(node_id)
|
manager.unregister(conn)
|
||||||
await _update_tunnel_status(node_id, connected=False)
|
if not manager.has_tunnel(node_id):
|
||||||
|
await _update_tunnel_status(node_id, connected=False)
|
||||||
|
else:
|
||||||
|
logger.info("tunnel connection closed but pool still active: node_id={}", node_id)
|
||||||
|
|
||||||
|
|
||||||
async def _ping_loop(conn: TunnelConnection) -> None:
|
async def _ping_loop(conn: TunnelConnection) -> None:
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
WebSocket 隧道管理器
|
WebSocket 隧道管理器
|
||||||
|
|
||||||
管理所有活跃的 aether-proxy tunnel 连接,提供通过隧道发送 HTTP 请求的能力。
|
管理所有活跃的 aether-proxy tunnel 连接,提供通过隧道发送 HTTP 请求的能力。
|
||||||
每个 proxy node 最多一条 tunnel 连接。
|
每个 proxy node 可持有多条 tunnel 连接(连接池),请求按 least-loaded 策略分配。
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -49,9 +49,18 @@ class TunnelConnection:
|
|||||||
def is_alive(self) -> bool:
|
def is_alive(self) -> bool:
|
||||||
return self.ws.client_state == WebSocketState.CONNECTED
|
return self.ws.client_state == WebSocketState.CONNECTED
|
||||||
|
|
||||||
async def send_frame(self, frame: Frame) -> None:
|
async def send_frame(self, frame: Frame, timeout: float = 10.0) -> None:
|
||||||
async with self._write_lock:
|
"""发送帧到 WebSocket,带超时保护防止写阻塞。
|
||||||
await self.ws.send_bytes(frame.encode())
|
|
||||||
|
在高丢包网络下 TCP 写缓冲区可能满,send_bytes 会长时间阻塞。
|
||||||
|
加超时避免所有协程在 _write_lock 上排队导致级联失败。
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
async with asyncio.timeout(timeout):
|
||||||
|
async with self._write_lock:
|
||||||
|
await self.ws.send_bytes(frame.encode())
|
||||||
|
except TimeoutError:
|
||||||
|
raise TunnelStreamError("frame send timeout (writer congested)")
|
||||||
|
|
||||||
def create_stream(self, stream_id: int) -> _StreamState:
|
def create_stream(self, stream_id: int) -> _StreamState:
|
||||||
state = _StreamState(stream_id)
|
state = _StreamState(stream_id)
|
||||||
@@ -163,43 +172,101 @@ class TunnelStreamError(Exception):
|
|||||||
|
|
||||||
|
|
||||||
class TunnelManager:
|
class TunnelManager:
|
||||||
"""管理所有活跃的 tunnel 连接"""
|
"""管理所有活跃的 tunnel 连接(支持每个 node 多条连接的连接池)"""
|
||||||
|
|
||||||
# 单条 tunnel 上允许的最大并发 stream 数(超出时拒绝新请求)
|
# 单条 tunnel 上允许的最大并发 stream 数(超出时拒绝新请求)
|
||||||
MAX_STREAMS_PER_CONN = 2048
|
MAX_STREAMS_PER_CONN = 2048
|
||||||
|
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self._connections: dict[str, TunnelConnection] = {} # node_id -> conn
|
self._connections: dict[str, list[TunnelConnection]] = {} # node_id -> [conn, ...]
|
||||||
|
self._background_tasks: set[asyncio.Task[None]] = set()
|
||||||
|
|
||||||
|
def _background(self, coro: Any) -> None: # noqa: ANN401
|
||||||
|
"""启动 fire-and-forget task,通过 set 持有引用防止 GC 回收,完成后自动清理"""
|
||||||
|
task = asyncio.create_task(coro)
|
||||||
|
self._background_tasks.add(task)
|
||||||
|
task.add_done_callback(self._background_tasks.discard)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def active_count(self) -> int:
|
def active_count(self) -> int:
|
||||||
return len(self._connections)
|
return sum(len(conns) for conns in self._connections.values())
|
||||||
|
|
||||||
def get_connection(self, node_id: str) -> TunnelConnection | None:
|
def get_connection(self, node_id: str) -> TunnelConnection | None:
|
||||||
conn = self._connections.get(node_id)
|
"""获取负载最低的存活连接,同时清理 dead 连接"""
|
||||||
if conn and not conn.is_alive:
|
conns = self._connections.get(node_id)
|
||||||
self._connections.pop(node_id, None)
|
if not conns:
|
||||||
conn.cancel_all_streams()
|
|
||||||
return None
|
return None
|
||||||
return conn
|
|
||||||
|
# 清理 dead 连接
|
||||||
|
alive = [c for c in conns if c.is_alive]
|
||||||
|
dead = [c for c in conns if not c.is_alive]
|
||||||
|
for c in dead:
|
||||||
|
c.cancel_all_streams()
|
||||||
|
|
||||||
|
if not alive:
|
||||||
|
self._connections.pop(node_id, None)
|
||||||
|
return None
|
||||||
|
|
||||||
|
if len(alive) != len(conns):
|
||||||
|
self._connections[node_id] = alive
|
||||||
|
|
||||||
|
# Least-loaded: 选 stream_count 最小的连接
|
||||||
|
return min(alive, key=lambda c: c.stream_count)
|
||||||
|
|
||||||
def register(self, conn: TunnelConnection) -> None:
|
def register(self, conn: TunnelConnection) -> None:
|
||||||
old = self._connections.get(conn.node_id)
|
"""注册一条新连接到连接池"""
|
||||||
if old:
|
conns = self._connections.get(conn.node_id)
|
||||||
old.cancel_all_streams()
|
if conns is None:
|
||||||
self._connections[conn.node_id] = conn
|
conns = []
|
||||||
logger.info("tunnel connected: node_id={}, name={}", conn.node_id, conn.node_name)
|
self._connections[conn.node_id] = conns
|
||||||
|
conns.append(conn)
|
||||||
|
logger.info(
|
||||||
|
"tunnel connected: node_id={}, name={}, pool_size={}",
|
||||||
|
conn.node_id,
|
||||||
|
conn.node_name,
|
||||||
|
len(conns),
|
||||||
|
)
|
||||||
|
|
||||||
def unregister(self, node_id: str) -> None:
|
def unregister(self, conn: TunnelConnection) -> bool:
|
||||||
conn = self._connections.pop(node_id, None)
|
"""
|
||||||
if conn:
|
从连接池中注销指定连接。
|
||||||
conn.cancel_all_streams()
|
|
||||||
logger.info("tunnel disconnected: node_id={}, name={}", node_id, conn.node_name)
|
返回 True 表示成功移除,False 表示该连接已不在池中。
|
||||||
|
"""
|
||||||
|
conns = self._connections.get(conn.node_id)
|
||||||
|
if not conns:
|
||||||
|
return False
|
||||||
|
|
||||||
|
try:
|
||||||
|
conns.remove(conn) # identity comparison via list.remove
|
||||||
|
except ValueError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
conn.cancel_all_streams()
|
||||||
|
|
||||||
|
if not conns:
|
||||||
|
self._connections.pop(conn.node_id, None)
|
||||||
|
|
||||||
|
remaining = len(conns) if conns else 0
|
||||||
|
logger.info(
|
||||||
|
"tunnel disconnected: node_id={}, name={}, remaining={}",
|
||||||
|
conn.node_id,
|
||||||
|
conn.node_name,
|
||||||
|
remaining,
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
def has_tunnel(self, node_id: str) -> bool:
|
def has_tunnel(self, node_id: str) -> bool:
|
||||||
conn = self.get_connection(node_id)
|
conn = self.get_connection(node_id)
|
||||||
return conn is not None
|
return conn is not None
|
||||||
|
|
||||||
|
def connection_count(self, node_id: str) -> int:
|
||||||
|
"""返回指定 node 当前存活的连接数"""
|
||||||
|
conns = self._connections.get(node_id)
|
||||||
|
if not conns:
|
||||||
|
return 0
|
||||||
|
return sum(1 for c in conns if c.is_alive)
|
||||||
|
|
||||||
async def send_request(
|
async def send_request(
|
||||||
self,
|
self,
|
||||||
node_id: str,
|
node_id: str,
|
||||||
@@ -248,10 +315,15 @@ class TunnelManager:
|
|||||||
|
|
||||||
return stream_state
|
return stream_state
|
||||||
|
|
||||||
async def handle_incoming_frame(self, node_id: str, frame: Frame) -> None:
|
async def handle_incoming_frame(self, conn: TunnelConnection, frame: Frame) -> None:
|
||||||
"""处理从 proxy 收到的响应帧"""
|
"""处理从 proxy 收到的响应帧(仅处理当前 active 连接的帧)。
|
||||||
conn = self.get_connection(node_id)
|
|
||||||
if not conn:
|
重要:此方法在 WebSocket 主读循环中被 await 调用,不能长时间阻塞,
|
||||||
|
否则会阻止读取后续帧,导致 proxy 端 TCP 缓冲区满而级联失败。
|
||||||
|
"""
|
||||||
|
# 防止已被移除的连接的帧继续被处理
|
||||||
|
conns = self._connections.get(conn.node_id)
|
||||||
|
if not conns or conn not in conns:
|
||||||
return
|
return
|
||||||
|
|
||||||
stream = conn.get_stream(frame.stream_id)
|
stream = conn.get_stream(frame.stream_id)
|
||||||
@@ -281,10 +353,19 @@ class TunnelManager:
|
|||||||
conn.remove_stream(frame.stream_id)
|
conn.remove_stream(frame.stream_id)
|
||||||
|
|
||||||
elif frame.msg_type == MsgType.HEARTBEAT_DATA:
|
elif frame.msg_type == MsgType.HEARTBEAT_DATA:
|
||||||
await self._handle_heartbeat(conn, frame)
|
# fire-and-forget: 不阻塞主读循环
|
||||||
|
self._background(self._handle_heartbeat(conn, frame))
|
||||||
|
|
||||||
elif frame.msg_type == MsgType.PING:
|
elif frame.msg_type == MsgType.PING:
|
||||||
await conn.send_frame(Frame(0, MsgType.PONG, 0, frame.payload))
|
# fire-and-forget: pong 回复不阻塞读循环
|
||||||
|
self._background(self._send_pong(conn, frame.payload))
|
||||||
|
|
||||||
|
async def _send_pong(self, conn: TunnelConnection, payload: bytes) -> None:
|
||||||
|
"""发送 PONG 回复(fire-and-forget,不阻塞主读循环)"""
|
||||||
|
try:
|
||||||
|
await conn.send_frame(Frame(0, MsgType.PONG, 0, payload))
|
||||||
|
except TunnelStreamError:
|
||||||
|
pass # best-effort pong
|
||||||
|
|
||||||
async def _handle_heartbeat(self, conn: TunnelConnection, frame: Frame) -> None:
|
async def _handle_heartbeat(self, conn: TunnelConnection, frame: Frame) -> None:
|
||||||
"""处理 proxy 上报的心跳数据,更新 DB,返回 ACK"""
|
"""处理 proxy 上报的心跳数据,更新 DB,返回 ACK"""
|
||||||
@@ -320,7 +401,10 @@ class TunnelManager:
|
|||||||
logger.warning("tunnel heartbeat DB update failed: {}", e)
|
logger.warning("tunnel heartbeat DB update failed: {}", e)
|
||||||
ack = {}
|
ack = {}
|
||||||
|
|
||||||
await conn.send_frame(Frame(0, MsgType.HEARTBEAT_ACK, 0, json.dumps(ack).encode()))
|
try:
|
||||||
|
await conn.send_frame(Frame(0, MsgType.HEARTBEAT_ACK, 0, json.dumps(ack).encode()))
|
||||||
|
except TunnelStreamError:
|
||||||
|
logger.debug("heartbeat ACK send failed for node_id={}", conn.node_id)
|
||||||
|
|
||||||
|
|
||||||
# 全局单例
|
# 全局单例
|
||||||
|
|||||||
119
tests/unit/test_tunnel_manager_connection_replacement.py
Normal file
119
tests/unit/test_tunnel_manager_connection_replacement.py
Normal file
@@ -0,0 +1,119 @@
|
|||||||
|
import asyncio
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from starlette.websockets import WebSocketState
|
||||||
|
|
||||||
|
from src.services.proxy_node.tunnel_manager import TunnelConnection, TunnelManager
|
||||||
|
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
|
||||||
|
|
||||||
|
|
||||||
|
class _DummyWebSocket:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.client_state = WebSocketState.CONNECTED
|
||||||
|
self.sent: list[bytes] = []
|
||||||
|
|
||||||
|
async def send_bytes(self, data: bytes) -> None:
|
||||||
|
self.sent.append(data)
|
||||||
|
|
||||||
|
async def close(self, code: int = 1000, reason: str | None = None) -> None: # noqa: ARG002
|
||||||
|
self.client_state = WebSocketState.DISCONNECTED
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_pool_register_and_unregister() -> None:
|
||||||
|
"""register 将连接追加到池中,unregister 按连接实例移除"""
|
||||||
|
manager = TunnelManager()
|
||||||
|
|
||||||
|
ws1 = _DummyWebSocket()
|
||||||
|
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||||
|
manager.register(conn1)
|
||||||
|
assert manager.connection_count("node-1") == 1
|
||||||
|
assert manager.get_connection("node-1") is conn1
|
||||||
|
|
||||||
|
ws2 = _DummyWebSocket()
|
||||||
|
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||||
|
manager.register(conn2)
|
||||||
|
assert manager.connection_count("node-1") == 2
|
||||||
|
|
||||||
|
# unregister conn1 不影响 conn2
|
||||||
|
assert manager.unregister(conn1) is True
|
||||||
|
assert manager.connection_count("node-1") == 1
|
||||||
|
assert manager.get_connection("node-1") is conn2
|
||||||
|
|
||||||
|
# 重复 unregister 返回 False
|
||||||
|
assert manager.unregister(conn1) is False
|
||||||
|
|
||||||
|
# unregister conn2 清空池
|
||||||
|
assert manager.unregister(conn2) is True
|
||||||
|
assert manager.get_connection("node-1") is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_least_loaded_selection() -> None:
|
||||||
|
"""get_connection 返回 stream_count 最小的连接"""
|
||||||
|
manager = TunnelManager()
|
||||||
|
|
||||||
|
ws1 = _DummyWebSocket()
|
||||||
|
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||||
|
ws2 = _DummyWebSocket()
|
||||||
|
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||||
|
manager.register(conn1)
|
||||||
|
manager.register(conn2)
|
||||||
|
|
||||||
|
# 两个都空闲,返回任一(实际返回 min,两者相同时返回第一个)
|
||||||
|
selected = manager.get_connection("node-1")
|
||||||
|
assert selected in (conn1, conn2)
|
||||||
|
|
||||||
|
# 给 conn1 加一个 stream,conn2 应被优先选中
|
||||||
|
conn1.create_stream(2)
|
||||||
|
assert manager.get_connection("node-1") is conn2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_dead_connections_cleaned_on_get() -> None:
|
||||||
|
"""get_connection 自动清理 dead 连接"""
|
||||||
|
manager = TunnelManager()
|
||||||
|
|
||||||
|
ws1 = _DummyWebSocket()
|
||||||
|
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||||
|
ws2 = _DummyWebSocket()
|
||||||
|
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||||
|
manager.register(conn1)
|
||||||
|
manager.register(conn2)
|
||||||
|
|
||||||
|
# 模拟 conn1 断开
|
||||||
|
ws1.client_state = WebSocketState.DISCONNECTED
|
||||||
|
assert manager.get_connection("node-1") is conn2
|
||||||
|
assert manager.connection_count("node-1") == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_removed_connection_frames_ignored() -> None:
|
||||||
|
"""已 unregister 的连接帧不应被处理"""
|
||||||
|
manager = TunnelManager()
|
||||||
|
|
||||||
|
ws1 = _DummyWebSocket()
|
||||||
|
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||||
|
ws2 = _DummyWebSocket()
|
||||||
|
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||||
|
manager.register(conn1)
|
||||||
|
manager.register(conn2)
|
||||||
|
|
||||||
|
# unregister conn1
|
||||||
|
manager.unregister(conn1)
|
||||||
|
|
||||||
|
ping = Frame(0, MsgType.PING, 0, b"hello")
|
||||||
|
|
||||||
|
# conn1 已不在池中,帧应被忽略
|
||||||
|
await manager.handle_incoming_frame(conn1, ping)
|
||||||
|
# 等待 fire-and-forget task 完成
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
assert ws1.sent == []
|
||||||
|
|
||||||
|
# conn2 仍在池中,帧正常处理
|
||||||
|
await manager.handle_incoming_frame(conn2, ping)
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
assert len(ws2.sent) == 1
|
||||||
|
pong = Frame.decode(ws2.sent[0])
|
||||||
|
assert pong.msg_type == MsgType.PONG
|
||||||
|
assert pong.payload == b"hello"
|
||||||
Reference in New Issue
Block a user