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

feat(proxy): 重构 Proxy 节点管理与隧道系统

Browse Source 
- 重构 proxy_nodes 管理端，支持节点注册、心跳、隧道生命周期管理
- 增强 tunnel 嵌入式 hub 和隧道协议
- 重构 aether-proxy 配置、隧道客户端、心跳和调度机制
- 调整 admin OAuth/配额/导入等处理器的参数传递
- 扩展数据迁移模块
- 补充 proxy nodes、OAuth、配额、系统导入等测试
- 更新前端 proxy nodes 视图和 API
This commit is contained in:
fawney19
2026-04-14 22:51:02 +08:00
parent fb31928e44
commit a4e7ac1df6
55 changed files with 3477 additions and 892 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -31,6 +31,7 @@ pub async fn connect_and_run(
server: &Arc<ServerContext>,
conn_idx: usize,
shutdown: &mut watch::Receiver<bool>,
drain: watch::Receiver<bool>,
) -> Result<TunnelOutcome, anyhow::Error> {
let ws_url = build_tunnel_url(server);
info!(url = %ws_url, conn = conn_idx, "connecting tunnel");
@@ -53,8 +54,7 @@ pub async fn connect_and_run(
http::HeaderValue::from_str(&dynamic_node_name)?,
);
// Advertise per-connection max concurrent streams so the backend can
// respect the proxy's capacity limit (backward-compatible: old backends
// ignore this header).
// respect the proxy's capacity limit.
let max_streams = state.config.tunnel_max_streams.unwrap_or(128);
headers.insert("X-Tunnel-Max-Streams", http::HeaderValue::from(max_streams));
@@ -67,13 +67,16 @@ pub async fn connect_and_run(
let port = uri.port_u16().unwrap_or(if is_tls { 443 } else { 80 });
// TCP connect with timeout
let connect_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
let connect_timeout = state
.config
.tunnel_connect_timeout()
.expect("validated config should resolve tunnel connect timeout");
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()
"tunnel TCP connect timeout ({}ms)",
connect_timeout.as_millis()
)
})??;
@@ -96,7 +99,7 @@ pub async fn connect_and_run(
max_message_size: Some(64 << 20),
..Default::default()
};
let handshake_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
let handshake_timeout = connect_timeout;
let (ws_stream, _response) = tokio::time::timeout(
handshake_timeout,
tokio_tungstenite::client_async_tls_with_config(
@@ -109,16 +112,25 @@ pub async fn connect_and_run(
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel WebSocket handshake timeout ({}s)",
handshake_timeout.as_secs()
"tunnel WebSocket handshake timeout ({}ms)",
handshake_timeout.as_millis()
)
})??;
let stale_timeout = state
.config
.tunnel_stale_timeout()
.expect("validated config should resolve tunnel stale timeout");
let ping_interval = state
.config
.tunnel_ping_interval()
.expect("validated config should resolve tunnel ping interval");
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,
connect_timeout_ms = connect_timeout.as_millis(),
stale_timeout_ms = stale_timeout.as_millis(),
ping_interval_ms = ping_interval.as_millis(),
"tunnel connected"
);
@@ -129,8 +141,8 @@ pub async fn connect_and_run(
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
// Spawn writer task (with WebSocket ping keepalive)
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 drain_signal = spawn_drain_signal(conn_idx, frame_tx.clone(), drain.clone());
// Spawn heartbeat task (only for primary connection to avoid
// resetting shared atomic metrics via swap(0))
@@ -153,7 +165,14 @@ pub async fn connect_and_run(
let state_clone = Arc::clone(state);
let server_clone = Arc::clone(server);
let outcome = tokio::select! {
result = dispatcher::run(state_clone, server_clone, ws_read, frame_tx.clone(), hb_handle) => {
result = dispatcher::run(
state_clone,
server_clone,
ws_read,
frame_tx.clone(),
hb_handle,
drain.clone(),
) => {
match result {
Ok(()) => TunnelOutcome::Disconnected,
Err(e) => return Err(e),
@@ -181,6 +200,10 @@ pub async fn connect_and_run(
// Drop our sender; the writer will exit once all stream handler clones
// are also dropped (i.e. after they finish their in-flight work).
drop(frame_tx);
if !drain_signal.is_finished() {
drain_signal.abort();
let _ = drain_signal.await;
}
// Wait for the writer task to finish with a generous timeout — the
// dispatcher already waits up to 30s for stream handlers, so 35s here
@@ -194,6 +217,35 @@ pub async fn connect_and_run(
Ok(outcome)
}
fn spawn_drain_signal(
conn_idx: usize,
frame_tx: writer::FrameSender,
mut drain: watch::Receiver<bool>,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
if !*drain.borrow() {
loop {
if drain.changed().await.is_err() {
return;
}
if *drain.borrow() {
break;
}
}
}
debug!(conn = conn_idx, "sending GOAWAY for tunnel drain");
let _ = tokio::time::timeout(
Duration::from_millis(250),
frame_tx.send(super::protocol::Frame::control(
super::protocol::MsgType::GoAway,
bytes::Bytes::new(),
)),
)
.await;
})
}
/// 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);

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

@@ -7,6 +7,7 @@ use std::time::Duration;
use bytes::Bytes;
use futures_util::StreamExt;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, info, warn};
@@ -25,6 +26,7 @@ pub async fn run<S>(
mut ws_stream: S,
frame_tx: FrameSender,
heartbeat: HeartbeatHandle,
mut drain: watch::Receiver<bool>,
) -> Result<(), anyhow::Error>
where
S: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>>
@@ -38,12 +40,21 @@ where
let mut handler_handles: Vec<JoinHandle<()>> = Vec::new();
let max_streams = state.config.tunnel_max_streams.unwrap_or(128) as usize;
let mut frames_since_cleanup: u32 = 0;
let stale_timeout = Duration::from_secs(state.config.tunnel_stale_timeout_secs);
let stale_timeout = state
.config
.tunnel_stale_timeout()
.expect("validated config should resolve tunnel stale timeout");
// Track last time we received any data to detect stale connections
let mut last_data_at = tokio::time::Instant::now();
let mut draining = *drain.borrow();
let read_err = loop {
if draining && streams.is_empty() {
info!("tunnel drained after in-flight streams completed");
break None;
}
let msg_result = tokio::select! {
msg = ws_stream.next() => {
match msg {
@@ -51,9 +62,19 @@ where
None => break None,
}
}
changed = drain.changed() => {
if changed.is_err() {
continue;
}
if *drain.borrow() {
info!("tunnel drain requested, waiting for in-flight streams");
draining = true;
}
continue;
}
_ = tokio::time::sleep_until(last_data_at + stale_timeout) => {
warn!(
stale_secs = stale_timeout.as_secs(),
stale_ms = stale_timeout.as_millis(),
"tunnel connection stale, no data received"
);
break None;
@@ -92,6 +113,24 @@ where
match frame.msg_type {
MsgType::RequestHeaders => {
if draining {
if frame_tx
.try_send(Frame::new(
frame.stream_id,
MsgType::StreamError,
0,
Bytes::from("tunnel draining"),
))
.is_err()
{
warn!(
stream_id = frame.stream_id,
"writer channel full, StreamError dropped during drain"
);
}
continue;
}
// Decompress if the frame is gzip-compressed, then parse metadata
let payload = match decompress_if_gzip(&frame) {
Ok(p) => p,
@@ -176,6 +215,10 @@ where
let _ = tx.send(frame).await;
if is_end {
streams.remove(&sid);
if draining && streams.is_empty() {
info!("tunnel drained after request body completion");
break None;
}
}
}
}
@@ -184,6 +227,10 @@ where
// Client-side cancellation or end
if let Some(tx) = streams.remove(&frame.stream_id) {
let _ = tx.send(frame).await;
if draining && streams.is_empty() {
info!("tunnel drained after stream termination");
break None;
}
}
}

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

@@ -24,7 +24,7 @@ static NON_ROOT_UPGRADE_WARNED: AtomicBool = AtomicBool::new(false);
enum AckDecision {
Accept {
heartbeat_id: Option<u64>,
heartbeat_id: u64,
upgrade_to: Option<String>,
},
Ignore,
@@ -143,16 +143,8 @@ pub fn spawn(
upgrade_to,
} => {
if let Some((pending_id, _)) = pending {
match ack_id {
Some(id) if id == pending_id => {
pending = None;
}
None => {
// Backward-compatible with servers that don't echo
// heartbeat_id in ACK payload yet.
pending = None;
}
_ => {}
if ack_id == pending_id {
pending = None;
}
}
maybe_trigger_upgrade(upgrade_to);
@@ -266,6 +258,7 @@ async fn build_heartbeat_payload(
"node_id": node_id,
"heartbeat_session_id": heartbeat_session_id,
"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,
@@ -283,10 +276,8 @@ async fn build_heartbeat_payload(
fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
if payload.is_empty() {
return AckDecision::Accept {
heartbeat_id: None,
upgrade_to: None,
};
warn!("received empty heartbeat ACK");
return AckDecision::Ignore;
}
#[derive(serde::Deserialize)]
@@ -295,8 +286,7 @@ fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
remote_config: Option<RemoteConfig>,
#[serde(default)]
config_version: u64,
#[serde(default)]
heartbeat_id: Option<u64>,
heartbeat_id: u64,
#[serde(default)]
upgrade_to: Option<String>,
}
@@ -371,3 +361,74 @@ fn maybe_trigger_upgrade(version: Option<String>) {
}
});
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, RwLock};
use arc_swap::ArcSwap;
use clap::Parser;
use super::{handle_ack, AckDecision};
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::{ProxyMetrics, ServerContext};
fn sample_server() -> Arc<ServerContext> {
let config = Arc::new(crate::config::Config::parse_from([
"aether-proxy",
"--aether-url",
"https://example.com",
"--management-token",
"ae_test",
"--node-name",
"proxy-test",
]));
Arc::new(ServerContext {
server_label: "heartbeat-test".to_string(),
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: config.node_name.clone(),
node_id: Arc::new(RwLock::new("node-123".to_string())),
aether_client: Arc::new(AetherClient::new(
&config,
&config.aether_url,
&config.management_token,
)),
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
})
}
#[test]
fn heartbeat_ack_requires_heartbeat_id() {
let server = sample_server();
let decision = handle_ack(
&server,
br#"{"config_version":1,"remote_config":{"heartbeat_interval":9}}"#,
);
assert!(matches!(decision, AckDecision::Ignore));
assert_eq!(server.dynamic.load().heartbeat_interval, 5);
}
#[test]
fn heartbeat_ack_applies_remote_config_with_heartbeat_id() {
let server = sample_server();
let decision = handle_ack(
&server,
br#"{"heartbeat_id":7,"config_version":1,"remote_config":{"heartbeat_interval":9}}"#,
);
assert!(matches!(
decision,
AckDecision::Accept {
heartbeat_id: 7,
upgrade_to: None
}
));
assert_eq!(server.dynamic.load().heartbeat_interval, 9);
}
}

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

@@ -36,10 +36,16 @@ pub async fn run(
server: &Arc<ServerContext>,
conn_idx: usize,
mut shutdown: watch::Receiver<bool>,
mut drain: watch::Receiver<bool>,
) {
info!(server = %server.server_label, conn = conn_idx, "starting tunnel");
let reconnect_salt = compute_connection_salt(server, conn_idx);
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drain requested before startup");
return;
}
let startup_delay = compute_startup_stagger(conn_idx, reconnect_salt);
if !startup_delay.is_zero() {
info!(
@@ -54,14 +60,24 @@ pub async fn run(
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during startup stagger");
return;
}
_ = drain.changed() => {
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drain requested during startup stagger");
return;
}
}
}
}
let mut consecutive_failures: u32 = 0;
loop {
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drained, exiting slot");
return;
}
let started_at = Instant::now();
match client::connect_and_run(state, server, conn_idx, &mut shutdown).await {
match client::connect_and_run(state, server, conn_idx, &mut shutdown, drain.clone()).await {
Ok(client::TunnelOutcome::Shutdown) => {
info!(server = %server.server_label, conn = conn_idx, "tunnel shut down gracefully");
return;
@@ -78,6 +94,10 @@ pub async fn run(
info!(server = %server.server_label, conn = conn_idx, "shutdown requested, not reconnecting");
return;
}
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drained after disconnect");
return;
}
// Reset backoff after a stable session to keep recovery snappy when
// failures are only occasional.
@@ -108,6 +128,12 @@ pub async fn run(
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during reconnect wait");
return;
}
_ = drain.changed() => {
if *drain.borrow() {
info!(server = %server.server_label, conn = conn_idx, "tunnel drain requested during reconnect wait");
return;
}
}
}
}
}
@@ -266,8 +292,9 @@ mod tests {
let proxy_task = tokio::spawn({
let state = Arc::clone(&state);
let server = Arc::clone(&server);
let (_drain_tx, drain_rx) = watch::channel(false);
async move {
run(&state, &server, 0, shutdown_rx).await;
run(&state, &server, 0, shutdown_rx, drain_rx).await;
}
});
@@ -486,13 +513,18 @@ mod tests {
log_max_files: 30,
tunnel_reconnect_base_ms: 50,
tunnel_reconnect_max_ms: 250,
tunnel_ping_interval_secs: 1,
tunnel_ping_interval_ms: 1_000,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 2,
tunnel_connect_timeout_ms: 2_000,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 5,
tunnel_connections: 1,
tunnel_stale_timeout_ms: 5_000,
tunnel_connections: Some(1),
tunnel_connections_max: Some(1),
tunnel_scale_check_interval_ms: 1_000,
tunnel_scale_up_threshold_percent: 70,
tunnel_scale_down_threshold_percent: 35,
tunnel_scale_down_grace_secs: 15,
}
}

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

@@ -1594,13 +1594,18 @@ mod tests {
log_max_files: 30,
tunnel_reconnect_base_ms: 500,
tunnel_reconnect_max_ms: 30_000,
tunnel_ping_interval_secs: 15,
tunnel_ping_interval_ms: 15_000,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 15,
tunnel_connect_timeout_ms: 15_000,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 45,
tunnel_connections: 1,
tunnel_stale_timeout_ms: 45_000,
tunnel_connections: Some(1),
tunnel_connections_max: Some(1),
tunnel_scale_check_interval_ms: 1_000,
tunnel_scale_up_threshold_percent: 70,
tunnel_scale_down_threshold_percent: 35,
tunnel_scale_down_grace_secs: 15,
}
}