mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
refactor(hub): 用本地 HTTP relay 替代 Worker WebSocket 长连接
Hub 数据面改为 /local/relay/{node_id} HTTP 端点,Worker 通过本机
HTTP 请求转发 tunnel 帧,不再维护 /worker WebSocket 长连接。
Hub 侧:
- 新增 control_plane.rs: Hub 通过 HTTP 回调 Aether app 处理心跳 ACK 和节点状态变更
- 新增 local_relay.rs: 接收本地 HTTP 请求,在 Hub 内部打开 LocalStream 并透传到 proxy
- 移除 worker_conn.rs 及 Worker WebSocket 处理逻辑
- 简化 protocol.rs: 移除 NODE_STATUS 帧类型,抽取通用 encode_frame/decode_payload
Python 侧:
- 删除 tunnel_manager.py 及其 WebSocket 连接管理器 (HubConnectionManager)
- 简化 hub_transport.py 为 HTTP relay 调用
- 新增 src/api/internal/hub.py 接收 Hub 控制面回调 (heartbeat/node-status)
- hub_config.py 移除 WebSocket 相关配置,改为 HTTP relay URL
- service.py 新增 update_tunnel_status 方法
- 删除 src/api/admin/proxy_tunnel.py (旧管理接口)
- proxy_node 缓存 TTL 从 15s 降至 3s 加速状态感知
This commit is contained in:
804
aether-hub/Cargo.lock
generated
804
aether-hub/Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -16,6 +16,9 @@ dashmap = "6"
|
|||||||
parking_lot = "0.12"
|
parking_lot = "0.12"
|
||||||
flate2 = "1"
|
flate2 = "1"
|
||||||
futures-util = "0.3"
|
futures-util = "0.3"
|
||||||
|
bytes = "1"
|
||||||
|
async-stream = "0.3"
|
||||||
|
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
||||||
|
|
||||||
[profile.release]
|
[profile.release]
|
||||||
lto = true
|
lto = true
|
||||||
|
|||||||
85
aether-hub/src/control_plane.rs
Normal file
85
aether-hub/src/control_plane.rs
Normal file
@@ -0,0 +1,85 @@
|
|||||||
|
use reqwest::Client;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ControlPlaneClient {
|
||||||
|
client: Option<Client>,
|
||||||
|
base_url: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ControlPlaneClient {
|
||||||
|
pub fn new(base_url: String) -> Self {
|
||||||
|
let client = Client::builder()
|
||||||
|
.timeout(std::time::Duration::from_secs(10))
|
||||||
|
.build()
|
||||||
|
.ok();
|
||||||
|
Self { client, base_url }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn disabled() -> Self {
|
||||||
|
Self {
|
||||||
|
client: None,
|
||||||
|
base_url: String::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn heartbeat_ack(&self, payload: &[u8]) -> Result<Vec<u8>, String> {
|
||||||
|
let Some(client) = &self.client else {
|
||||||
|
return Ok(b"{}".to_vec());
|
||||||
|
};
|
||||||
|
let url = format!(
|
||||||
|
"{}/api/internal/hub/heartbeat",
|
||||||
|
self.base_url.trim_end_matches('/')
|
||||||
|
);
|
||||||
|
let response = client
|
||||||
|
.post(&url)
|
||||||
|
.header("content-type", "application/json")
|
||||||
|
.body(payload.to_vec())
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("heartbeat callback request failed: {e}"))?;
|
||||||
|
if !response.status().is_success() {
|
||||||
|
return Err(format!(
|
||||||
|
"heartbeat callback failed with status {}",
|
||||||
|
response.status()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
response
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map(|bytes| bytes.to_vec())
|
||||||
|
.map_err(|e| format!("heartbeat callback body read failed: {e}"))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn push_node_status(
|
||||||
|
&self,
|
||||||
|
node_id: &str,
|
||||||
|
connected: bool,
|
||||||
|
conn_count: usize,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let Some(client) = &self.client else {
|
||||||
|
return Ok(());
|
||||||
|
};
|
||||||
|
let url = format!(
|
||||||
|
"{}/api/internal/hub/node-status",
|
||||||
|
self.base_url.trim_end_matches('/')
|
||||||
|
);
|
||||||
|
let response = client
|
||||||
|
.post(&url)
|
||||||
|
.json(&serde_json::json!({
|
||||||
|
"node_id": node_id,
|
||||||
|
"connected": connected,
|
||||||
|
"conn_count": conn_count,
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| format!("node-status callback request failed: {e}"))?;
|
||||||
|
if response.status().is_success() {
|
||||||
|
Ok(())
|
||||||
|
} else {
|
||||||
|
Err(format!(
|
||||||
|
"node-status callback failed with status {}",
|
||||||
|
response.status()
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
File diff suppressed because it is too large
Load Diff
169
aether-hub/src/local_relay.rs
Normal file
169
aether-hub/src/local_relay.rs
Normal file
@@ -0,0 +1,169 @@
|
|||||||
|
use std::io;
|
||||||
|
use std::net::SocketAddr;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use async_stream::stream;
|
||||||
|
use axum::body::{Body, Bytes};
|
||||||
|
use axum::extract::{ConnectInfo, Path, State};
|
||||||
|
use axum::http::{HeaderMap, HeaderName, HeaderValue, Response, StatusCode};
|
||||||
|
use axum::response::IntoResponse;
|
||||||
|
use tracing::warn;
|
||||||
|
|
||||||
|
use crate::hub::LocalBodyEvent;
|
||||||
|
use crate::protocol;
|
||||||
|
use crate::AppState;
|
||||||
|
|
||||||
|
pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error";
|
||||||
|
|
||||||
|
struct StreamGuard {
|
||||||
|
hub: std::sync::Arc<crate::hub::HubRouter>,
|
||||||
|
stream_id: u64,
|
||||||
|
finished: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for StreamGuard {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if !self.finished {
|
||||||
|
self.hub
|
||||||
|
.cancel_local_stream(self.stream_id, "local relay client dropped");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn relay_request(
|
||||||
|
Path(node_id): Path<String>,
|
||||||
|
State(state): State<AppState>,
|
||||||
|
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||||
|
body: Bytes,
|
||||||
|
) -> impl IntoResponse {
|
||||||
|
if !addr.ip().is_loopback() {
|
||||||
|
return tunnel_error_response(
|
||||||
|
StatusCode::FORBIDDEN,
|
||||||
|
"forbidden",
|
||||||
|
"local relay only accepts loopback requests",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let (meta, request_body) = match decode_envelope(body) {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(error) => {
|
||||||
|
return tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let stream = match state.hub.open_local_stream(&node_id, &meta, request_body) {
|
||||||
|
Ok(stream) => stream,
|
||||||
|
Err(error) => {
|
||||||
|
return tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let request_guard = StreamGuard {
|
||||||
|
hub: state.hub.clone(),
|
||||||
|
stream_id: stream.id,
|
||||||
|
finished: false,
|
||||||
|
};
|
||||||
|
|
||||||
|
let wait_timeout = Duration::from_secs(meta.timeout.clamp(5, 300));
|
||||||
|
let response_head = match stream.wait_headers(wait_timeout).await {
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(error) => {
|
||||||
|
state.hub.cancel_local_stream(stream.id, &error);
|
||||||
|
return tunnel_error_response(StatusCode::GATEWAY_TIMEOUT, "timeout", &error);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let Some(mut body_rx) = stream.take_body_receiver() else {
|
||||||
|
state
|
||||||
|
.hub
|
||||||
|
.cancel_local_stream(stream.id, "missing relay response body receiver");
|
||||||
|
return tunnel_error_response(
|
||||||
|
StatusCode::BAD_GATEWAY,
|
||||||
|
"relay",
|
||||||
|
"missing relay response body receiver",
|
||||||
|
);
|
||||||
|
};
|
||||||
|
|
||||||
|
let hub = state.hub.clone();
|
||||||
|
let stream_id = stream.id;
|
||||||
|
let body_stream = stream! {
|
||||||
|
let mut guard = request_guard;
|
||||||
|
guard.hub = hub;
|
||||||
|
guard.stream_id = stream_id;
|
||||||
|
while let Some(event) = body_rx.recv().await {
|
||||||
|
match event {
|
||||||
|
LocalBodyEvent::Chunk(chunk) => yield Ok::<Bytes, io::Error>(chunk),
|
||||||
|
LocalBodyEvent::End => {
|
||||||
|
guard.finished = true;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
LocalBodyEvent::Error(error) => {
|
||||||
|
guard.finished = true;
|
||||||
|
yield Err(io::Error::other(error));
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
guard.finished = true;
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut builder = Response::builder().status(response_head.status);
|
||||||
|
if let Some(headers) = builder.headers_mut() {
|
||||||
|
append_headers(headers, &response_head.headers);
|
||||||
|
}
|
||||||
|
match builder.body(Body::from_stream(body_stream)) {
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(error) => {
|
||||||
|
warn!(error = %error, "failed to build relay response");
|
||||||
|
tunnel_error_response(
|
||||||
|
StatusCode::BAD_GATEWAY,
|
||||||
|
"relay",
|
||||||
|
"failed to build relay response",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode_envelope(body: Bytes) -> Result<(protocol::RequestMeta, Bytes), String> {
|
||||||
|
if body.len() < 4 {
|
||||||
|
return Err("relay envelope too short".to_string());
|
||||||
|
}
|
||||||
|
let meta_len = u32::from_be_bytes([body[0], body[1], body[2], body[3]]) as usize;
|
||||||
|
let meta_end = 4usize
|
||||||
|
.checked_add(meta_len)
|
||||||
|
.ok_or_else(|| "relay envelope length overflow".to_string())?;
|
||||||
|
if body.len() < meta_end {
|
||||||
|
return Err("relay envelope metadata truncated".to_string());
|
||||||
|
}
|
||||||
|
let meta = serde_json::from_slice::<protocol::RequestMeta>(&body[4..meta_end])
|
||||||
|
.map_err(|e| format!("invalid relay metadata: {e}"))?;
|
||||||
|
Ok((meta, body.slice(meta_end..)))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) {
|
||||||
|
for (name, value) in headers {
|
||||||
|
let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let Ok(value) = HeaderValue::from_str(value) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
target.append(name, value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tunnel_error_response(status: StatusCode, kind: &str, message: &str) -> Response<Body> {
|
||||||
|
let mut builder = Response::builder().status(status);
|
||||||
|
if let Some(headers) = builder.headers_mut() {
|
||||||
|
headers.insert(
|
||||||
|
HeaderName::from_static(TUNNEL_ERROR_HEADER),
|
||||||
|
HeaderValue::from_str(kind).unwrap_or_else(|_| HeaderValue::from_static("relay")),
|
||||||
|
);
|
||||||
|
headers.insert(
|
||||||
|
axum::http::header::CONTENT_TYPE,
|
||||||
|
HeaderValue::from_static("text/plain; charset=utf-8"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.body(Body::from(message.to_string()))
|
||||||
|
.unwrap_or_else(|_| Response::new(Body::from("relay error")))
|
||||||
|
}
|
||||||
@@ -1,20 +1,23 @@
|
|||||||
|
mod control_plane;
|
||||||
mod hub;
|
mod hub;
|
||||||
|
mod local_relay;
|
||||||
mod protocol;
|
mod protocol;
|
||||||
mod proxy_conn;
|
mod proxy_conn;
|
||||||
mod worker_conn;
|
|
||||||
|
|
||||||
use std::sync::Arc;
|
use std::net::SocketAddr;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use axum::extract::ws::WebSocketUpgrade;
|
use axum::extract::ws::WebSocketUpgrade;
|
||||||
use axum::extract::State;
|
use axum::extract::State;
|
||||||
use axum::response::{IntoResponse, Json};
|
use axum::response::{IntoResponse, Json};
|
||||||
use axum::routing::get;
|
use axum::routing::{get, post};
|
||||||
use axum::Router;
|
use axum::Router;
|
||||||
use clap::Parser;
|
use clap::Parser;
|
||||||
use tracing::{info, warn};
|
use tracing::{info, warn};
|
||||||
|
|
||||||
|
use crate::control_plane::ControlPlaneClient;
|
||||||
use crate::hub::{ConnConfig, HubRouter};
|
use crate::hub::{ConnConfig, HubRouter};
|
||||||
|
use crate::local_relay::relay_request;
|
||||||
|
|
||||||
#[derive(Parser, Debug)]
|
#[derive(Parser, Debug)]
|
||||||
#[command(name = "aether-hub", about = "Tunnel Hub for Aether")]
|
#[command(name = "aether-hub", about = "Tunnel Hub for Aether")]
|
||||||
@@ -27,10 +30,6 @@ struct Args {
|
|||||||
#[arg(long, default_value_t = 0, env = "TUNNEL_HUB_PROXY_IDLE_TIMEOUT")]
|
#[arg(long, default_value_t = 0, env = "TUNNEL_HUB_PROXY_IDLE_TIMEOUT")]
|
||||||
proxy_idle_timeout: u64,
|
proxy_idle_timeout: u64,
|
||||||
|
|
||||||
/// Worker-side idle timeout in seconds (0 to disable)
|
|
||||||
#[arg(long, default_value_t = 60, env = "TUNNEL_HUB_WORKER_IDLE_TIMEOUT")]
|
|
||||||
worker_idle_timeout: u64,
|
|
||||||
|
|
||||||
/// Ping interval in seconds (for both sides)
|
/// Ping interval in seconds (for both sides)
|
||||||
#[arg(long, default_value_t = 15, env = "TUNNEL_HUB_PING_INTERVAL")]
|
#[arg(long, default_value_t = 15, env = "TUNNEL_HUB_PING_INTERVAL")]
|
||||||
ping_interval: u64,
|
ping_interval: u64,
|
||||||
@@ -46,14 +45,21 @@ struct Args {
|
|||||||
env = "TUNNEL_HUB_OUTBOUND_QUEUE_CAPACITY"
|
env = "TUNNEL_HUB_OUTBOUND_QUEUE_CAPACITY"
|
||||||
)]
|
)]
|
||||||
outbound_queue_capacity: usize,
|
outbound_queue_capacity: usize,
|
||||||
|
|
||||||
|
/// Local Aether app base URL for control-plane callbacks
|
||||||
|
#[arg(
|
||||||
|
long,
|
||||||
|
default_value = "http://127.0.0.1:8084",
|
||||||
|
env = "TUNNEL_HUB_APP_BASE_URL"
|
||||||
|
)]
|
||||||
|
app_base_url: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
struct AppState {
|
pub struct AppState {
|
||||||
hub: Arc<HubRouter>,
|
pub hub: std::sync::Arc<HubRouter>,
|
||||||
proxy_conn_cfg: ConnConfig,
|
pub proxy_conn_cfg: ConnConfig,
|
||||||
worker_conn_cfg: ConnConfig,
|
pub max_streams: usize,
|
||||||
max_streams: usize,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
@@ -68,7 +74,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
|
|
||||||
let args = Args::parse();
|
let args = Args::parse();
|
||||||
|
|
||||||
let hub = HubRouter::new();
|
let hub = HubRouter::new(ControlPlaneClient::new(args.app_base_url));
|
||||||
let outbound_queue_capacity = args.outbound_queue_capacity.clamp(8, 4096);
|
let outbound_queue_capacity = args.outbound_queue_capacity.clamp(8, 4096);
|
||||||
let ping_interval = Duration::from_secs(args.ping_interval);
|
let ping_interval = Duration::from_secs(args.ping_interval);
|
||||||
let state = AppState {
|
let state = AppState {
|
||||||
@@ -78,11 +84,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
idle_timeout: Duration::from_secs(args.proxy_idle_timeout),
|
idle_timeout: Duration::from_secs(args.proxy_idle_timeout),
|
||||||
outbound_queue_capacity,
|
outbound_queue_capacity,
|
||||||
},
|
},
|
||||||
worker_conn_cfg: ConnConfig {
|
|
||||||
ping_interval,
|
|
||||||
idle_timeout: Duration::from_secs(args.worker_idle_timeout),
|
|
||||||
outbound_queue_capacity,
|
|
||||||
},
|
|
||||||
max_streams: args.max_streams,
|
max_streams: args.max_streams,
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -90,13 +91,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
.route("/health", get(health))
|
.route("/health", get(health))
|
||||||
.route("/stats", get(stats))
|
.route("/stats", get(stats))
|
||||||
.route("/proxy", get(ws_proxy))
|
.route("/proxy", get(ws_proxy))
|
||||||
.route("/worker", get(ws_worker))
|
.route("/local/relay/{node_id}", post(relay_request))
|
||||||
.with_state(state);
|
.with_state(state);
|
||||||
|
|
||||||
let listener = tokio::net::TcpListener::bind(&args.bind).await?;
|
let listener = tokio::net::TcpListener::bind(&args.bind).await?;
|
||||||
info!(bind = %args.bind, "aether-hub started");
|
info!(bind = %args.bind, "aether-hub started");
|
||||||
|
|
||||||
axum::serve(listener, app).await?;
|
axum::serve(
|
||||||
|
listener,
|
||||||
|
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -160,10 +165,3 @@ async fn ws_proxy(
|
|||||||
})
|
})
|
||||||
.into_response()
|
.into_response()
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn ws_worker(ws: WebSocketUpgrade, State(state): State<AppState>) -> impl IntoResponse {
|
|
||||||
ws.max_frame_size(64 * 1024 * 1024)
|
|
||||||
.on_upgrade(move |socket| {
|
|
||||||
worker_conn::handle_worker_connection(socket, state.hub, state.worker_conn_cfg)
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -22,8 +22,6 @@ pub const PONG: u8 = 0x11;
|
|||||||
pub const GOAWAY: u8 = 0x12;
|
pub const GOAWAY: u8 = 0x12;
|
||||||
pub const HEARTBEAT_DATA: u8 = 0x13;
|
pub const HEARTBEAT_DATA: u8 = 0x13;
|
||||||
pub const HEARTBEAT_ACK: u8 = 0x14;
|
pub const HEARTBEAT_ACK: u8 = 0x14;
|
||||||
pub const NODE_STATUS: u8 = 0x15;
|
|
||||||
|
|
||||||
// Flags
|
// Flags
|
||||||
pub const FLAG_END_STREAM: u8 = 0x01;
|
pub const FLAG_END_STREAM: u8 = 0x01;
|
||||||
pub const FLAG_GZIP_COMPRESSED: u8 = 0x02;
|
pub const FLAG_GZIP_COMPRESSED: u8 = 0x02;
|
||||||
@@ -50,162 +48,89 @@ impl FrameHeader {
|
|||||||
payload_len: u32::from_be_bytes([data[6], data[7], data[8], data[9]]),
|
payload_len: u32::from_be_bytes([data[6], data[7], data[8], data[9]]),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Check if this is a stream-terminating frame
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||||
#[inline]
|
pub struct RequestMeta {
|
||||||
pub fn is_stream_terminal(&self) -> bool {
|
pub method: String,
|
||||||
self.msg_type == STREAM_END || self.msg_type == STREAM_ERROR
|
pub url: String,
|
||||||
|
pub headers: std::collections::HashMap<String, String>,
|
||||||
|
#[serde(default = "default_timeout", deserialize_with = "deserialize_timeout")]
|
||||||
|
pub timeout: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_timeout() -> u64 {
|
||||||
|
60
|
||||||
|
}
|
||||||
|
|
||||||
|
fn deserialize_timeout<'de, D>(deserializer: D) -> Result<u64, D::Error>
|
||||||
|
where
|
||||||
|
D: serde::Deserializer<'de>,
|
||||||
|
{
|
||||||
|
#[derive(serde::Deserialize)]
|
||||||
|
#[serde(untagged)]
|
||||||
|
enum TimeoutValue {
|
||||||
|
Int(u64),
|
||||||
|
Float(f64),
|
||||||
|
}
|
||||||
|
|
||||||
|
match <TimeoutValue as serde::Deserialize>::deserialize(deserializer)? {
|
||||||
|
TimeoutValue::Int(v) => Ok(v),
|
||||||
|
TimeoutValue::Float(v) => {
|
||||||
|
if !v.is_finite() || v < 0.0 {
|
||||||
|
return Err(serde::de::Error::custom(
|
||||||
|
"timeout must be a non-negative finite number",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if v.fract() != 0.0 {
|
||||||
|
return Err(serde::de::Error::custom("timeout must be integer seconds"));
|
||||||
|
}
|
||||||
|
if v > (u64::MAX as f64) {
|
||||||
|
return Err(serde::de::Error::custom("timeout is too large"));
|
||||||
|
}
|
||||||
|
Ok(v as u64)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||||
pub struct RequestHeadersExtracted {
|
pub struct ResponseMeta {
|
||||||
pub node_id: String,
|
pub status: u16,
|
||||||
pub rebuilt_frame: Vec<u8>,
|
pub headers: Vec<(String, String)>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn encode_frame(stream_id: u32, msg_type: u8, flags: u8, payload: &[u8]) -> Vec<u8> {
|
||||||
|
let mut buf = Vec::with_capacity(HEADER_SIZE + payload.len());
|
||||||
|
buf.extend_from_slice(&stream_id.to_be_bytes());
|
||||||
|
buf.push(msg_type);
|
||||||
|
buf.push(flags);
|
||||||
|
buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
|
||||||
|
buf.extend_from_slice(payload);
|
||||||
|
buf
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Encode a STREAM_ERROR frame for a given stream_id with an error message
|
/// Encode a STREAM_ERROR frame for a given stream_id with an error message
|
||||||
pub fn encode_stream_error(stream_id: u32, msg: &str) -> Vec<u8> {
|
pub fn encode_stream_error(stream_id: u32, msg: &str) -> Vec<u8> {
|
||||||
let payload = msg.as_bytes();
|
encode_frame(stream_id, STREAM_ERROR, 0, msg.as_bytes())
|
||||||
let mut buf = Vec::with_capacity(HEADER_SIZE + payload.len());
|
|
||||||
buf.extend_from_slice(&stream_id.to_be_bytes());
|
|
||||||
buf.push(STREAM_ERROR);
|
|
||||||
buf.push(0); // flags
|
|
||||||
buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
|
|
||||||
buf.extend_from_slice(payload);
|
|
||||||
buf
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Encode a NODE_STATUS frame (stream_id=0, Hub-generated)
|
|
||||||
pub fn encode_node_status(node_id: &str, connected: bool, conn_count: usize) -> Vec<u8> {
|
|
||||||
let payload = serde_json::json!({
|
|
||||||
"node_id": node_id,
|
|
||||||
"connected": connected,
|
|
||||||
"conn_count": conn_count,
|
|
||||||
});
|
|
||||||
let payload_bytes = payload.to_string().into_bytes();
|
|
||||||
let mut buf = Vec::with_capacity(HEADER_SIZE + payload_bytes.len());
|
|
||||||
buf.extend_from_slice(&0u32.to_be_bytes()); // stream_id = 0
|
|
||||||
buf.push(NODE_STATUS);
|
|
||||||
buf.push(0); // flags
|
|
||||||
buf.extend_from_slice(&(payload_bytes.len() as u32).to_be_bytes());
|
|
||||||
buf.extend_from_slice(&payload_bytes);
|
|
||||||
buf
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Encode a PING frame (stream_id=0)
|
/// Encode a PING frame (stream_id=0)
|
||||||
pub fn encode_ping() -> Vec<u8> {
|
pub fn encode_ping() -> Vec<u8> {
|
||||||
let mut buf = Vec::with_capacity(HEADER_SIZE);
|
encode_frame(0, PING, 0, &[])
|
||||||
buf.extend_from_slice(&0u32.to_be_bytes());
|
|
||||||
buf.push(PING);
|
|
||||||
buf.push(0);
|
|
||||||
buf.extend_from_slice(&0u32.to_be_bytes());
|
|
||||||
buf
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Encode a PONG frame (stream_id=0, echo payload)
|
/// Encode a PONG frame (stream_id=0, echo payload)
|
||||||
pub fn encode_pong(payload: &[u8]) -> Vec<u8> {
|
pub fn encode_pong(payload: &[u8]) -> Vec<u8> {
|
||||||
let mut buf = Vec::with_capacity(HEADER_SIZE + payload.len());
|
encode_frame(0, PONG, 0, payload)
|
||||||
buf.extend_from_slice(&0u32.to_be_bytes());
|
|
||||||
buf.push(PONG);
|
|
||||||
buf.push(0);
|
|
||||||
buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
|
|
||||||
buf.extend_from_slice(payload);
|
|
||||||
buf
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Encode a GOAWAY frame (stream_id=0)
|
/// Encode a GOAWAY frame (stream_id=0)
|
||||||
pub fn encode_goaway() -> Vec<u8> {
|
pub fn encode_goaway() -> Vec<u8> {
|
||||||
let mut buf = Vec::with_capacity(HEADER_SIZE);
|
encode_frame(0, GOAWAY, 0, &[])
|
||||||
buf.extend_from_slice(&0u32.to_be_bytes());
|
|
||||||
buf.push(GOAWAY);
|
|
||||||
buf.push(0);
|
|
||||||
buf.extend_from_slice(&0u32.to_be_bytes());
|
|
||||||
buf
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Rewrite the stream_id in raw frame bytes (first 4 bytes) -- near zero-copy
|
|
||||||
#[inline]
|
|
||||||
pub fn rewrite_stream_id(data: &mut [u8], new_stream_id: u32) {
|
|
||||||
let bytes = new_stream_id.to_be_bytes();
|
|
||||||
data[0] = bytes[0];
|
|
||||||
data[1] = bytes[1];
|
|
||||||
data[2] = bytes[2];
|
|
||||||
data[3] = bytes[3];
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Get the payload portion of a raw frame (after the 10-byte header)
|
|
||||||
#[inline]
|
|
||||||
pub fn frame_payload(data: &[u8]) -> &[u8] {
|
|
||||||
if data.len() > HEADER_SIZE {
|
|
||||||
&data[HEADER_SIZE..]
|
|
||||||
} else {
|
|
||||||
&[]
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Parse REQUEST_HEADERS payload, extract `node_id`, strip it from JSON,
|
|
||||||
/// and rebuild a new REQUEST_HEADERS frame with `new_stream_id`.
|
|
||||||
///
|
|
||||||
/// If the source frame is gzip-compressed, this function will decode it first,
|
|
||||||
/// then try to re-encode with gzip (only keeps compression when payload shrinks).
|
|
||||||
pub fn rebuild_request_headers_without_node_id(
|
|
||||||
data: &[u8],
|
|
||||||
new_stream_id: u32,
|
|
||||||
) -> Result<RequestHeadersExtracted, String> {
|
|
||||||
let header = FrameHeader::parse(data).ok_or_else(|| "invalid frame header".to_string())?;
|
|
||||||
if header.msg_type != REQUEST_HEADERS {
|
|
||||||
return Err("frame is not REQUEST_HEADERS".to_string());
|
|
||||||
}
|
|
||||||
|
|
||||||
let payload = frame_payload_by_header(data, &header)
|
|
||||||
.ok_or_else(|| "incomplete REQUEST_HEADERS payload".to_string())?;
|
|
||||||
|
|
||||||
let decoded_payload = if header.flags & FLAG_GZIP_COMPRESSED != 0 {
|
|
||||||
let mut decoder = GzDecoder::new(payload);
|
|
||||||
let mut decoded = Vec::new();
|
|
||||||
decoder
|
|
||||||
.read_to_end(&mut decoded)
|
|
||||||
.map_err(|e| format!("failed to decompress REQUEST_HEADERS: {e}"))?;
|
|
||||||
decoded
|
|
||||||
} else {
|
|
||||||
payload.to_vec()
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut meta: serde_json::Value = serde_json::from_slice(&decoded_payload)
|
|
||||||
.map_err(|e| format!("invalid REQUEST_HEADERS JSON: {e}"))?;
|
|
||||||
let obj = meta
|
|
||||||
.as_object_mut()
|
|
||||||
.ok_or_else(|| "REQUEST_HEADERS payload must be a JSON object".to_string())?;
|
|
||||||
|
|
||||||
let node_id = obj
|
|
||||||
.remove("node_id")
|
|
||||||
.and_then(|v| v.as_str().map(|s| s.to_string()))
|
|
||||||
.map(|s| s.trim().to_string())
|
|
||||||
.filter(|s| !s.is_empty())
|
|
||||||
.ok_or_else(|| "missing node_id in REQUEST_HEADERS".to_string())?;
|
|
||||||
|
|
||||||
let stripped_payload = serde_json::to_vec(&meta)
|
|
||||||
.map_err(|e| format!("failed to encode REQUEST_HEADERS payload: {e}"))?;
|
|
||||||
let (final_payload, flags) =
|
|
||||||
maybe_recompress_payload(&stripped_payload, header.flags & FLAG_GZIP_COMPRESSED != 0)
|
|
||||||
.map_err(|e| format!("failed to recompress REQUEST_HEADERS payload: {e}"))?;
|
|
||||||
|
|
||||||
let mut rebuilt = Vec::with_capacity(HEADER_SIZE + final_payload.len());
|
|
||||||
rebuilt.extend_from_slice(&new_stream_id.to_be_bytes());
|
|
||||||
rebuilt.push(REQUEST_HEADERS);
|
|
||||||
rebuilt.push(flags);
|
|
||||||
rebuilt.extend_from_slice(&(final_payload.len() as u32).to_be_bytes());
|
|
||||||
rebuilt.extend_from_slice(&final_payload);
|
|
||||||
|
|
||||||
Ok(RequestHeadersExtracted {
|
|
||||||
node_id,
|
|
||||||
rebuilt_frame: rebuilt,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
fn frame_payload_by_header<'a>(data: &'a [u8], header: &FrameHeader) -> Option<&'a [u8]> {
|
pub fn frame_payload_by_header<'a>(data: &'a [u8], header: &FrameHeader) -> Option<&'a [u8]> {
|
||||||
let payload_len = header.payload_len as usize;
|
let payload_len = header.payload_len as usize;
|
||||||
let end = HEADER_SIZE.checked_add(payload_len)?;
|
let end = HEADER_SIZE.checked_add(payload_len)?;
|
||||||
if data.len() < end {
|
if data.len() < end {
|
||||||
@@ -214,6 +139,25 @@ fn frame_payload_by_header<'a>(data: &'a [u8], header: &FrameHeader) -> Option<&
|
|||||||
Some(&data[HEADER_SIZE..end])
|
Some(&data[HEADER_SIZE..end])
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn decode_payload(data: &[u8], header: &FrameHeader) -> Result<Vec<u8>, String> {
|
||||||
|
let payload = frame_payload_by_header(data, header)
|
||||||
|
.ok_or_else(|| "incomplete frame payload".to_string())?;
|
||||||
|
if header.flags & FLAG_GZIP_COMPRESSED != 0 {
|
||||||
|
let mut decoder = GzDecoder::new(payload);
|
||||||
|
let mut decoded = Vec::new();
|
||||||
|
decoder
|
||||||
|
.read_to_end(&mut decoded)
|
||||||
|
.map_err(|e| format!("failed to decompress payload: {e}"))?;
|
||||||
|
Ok(decoded)
|
||||||
|
} else {
|
||||||
|
Ok(payload.to_vec())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn compress_payload(payload: &[u8]) -> Result<(Vec<u8>, u8), std::io::Error> {
|
||||||
|
maybe_recompress_payload(payload, true)
|
||||||
|
}
|
||||||
|
|
||||||
fn maybe_recompress_payload(
|
fn maybe_recompress_payload(
|
||||||
payload: &[u8],
|
payload: &[u8],
|
||||||
prefer_gzip: bool,
|
prefer_gzip: bool,
|
||||||
|
|||||||
@@ -140,7 +140,7 @@ async fn run_proxy_reader(
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
hub.handle_proxy_frame(conn.id, &mut data);
|
hub.handle_proxy_frame(conn.id, &mut data).await;
|
||||||
}
|
}
|
||||||
Some(Ok(Message::Close(_))) | None => {
|
Some(Ok(Message::Close(_))) | None => {
|
||||||
info!(conn_id = conn.id, node_id = %conn.node_id, "proxy WebSocket closed");
|
info!(conn_id = conn.id, node_id = %conn.node_id, "proxy WebSocket closed");
|
||||||
|
|||||||
@@ -1,204 +0,0 @@
|
|||||||
/// Worker-side WebSocket connection handler
|
|
||||||
///
|
|
||||||
/// Handles the lifecycle of a single Gunicorn worker connection:
|
|
||||||
/// accept -> read loop (route frames via Hub) -> cleanup
|
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::time::{Duration, Instant};
|
|
||||||
|
|
||||||
use axum::extract::ws::{Message, WebSocket};
|
|
||||||
use futures_util::{SinkExt, StreamExt};
|
|
||||||
use tokio::sync::{mpsc, watch};
|
|
||||||
use tracing::{debug, info, warn};
|
|
||||||
|
|
||||||
use crate::hub::{ConnConfig, HubRouter, SendStatus, WorkerConn};
|
|
||||||
use crate::protocol;
|
|
||||||
|
|
||||||
pub async fn handle_worker_connection(ws: WebSocket, hub: Arc<HubRouter>, cfg: ConnConfig) {
|
|
||||||
let conn_id = hub.alloc_conn_id();
|
|
||||||
let (mut ws_tx, ws_rx) = ws.split();
|
|
||||||
|
|
||||||
let (tx, mut rx) = mpsc::channel::<Message>(cfg.outbound_queue_capacity);
|
|
||||||
let (close_tx, mut close_rx) = watch::channel(false);
|
|
||||||
|
|
||||||
let conn = Arc::new(WorkerConn::new(conn_id, tx, close_tx));
|
|
||||||
hub.register_worker(conn.clone());
|
|
||||||
|
|
||||||
let writer = tokio::spawn(async move {
|
|
||||||
loop {
|
|
||||||
tokio::select! {
|
|
||||||
msg = rx.recv() => match msg {
|
|
||||||
Some(msg) => {
|
|
||||||
if ws_tx.send(msg).await.is_err() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
None => break,
|
|
||||||
},
|
|
||||||
changed = close_rx.changed() => {
|
|
||||||
if changed.is_err() || *close_rx.borrow() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
let _ = ws_tx.close().await;
|
|
||||||
});
|
|
||||||
|
|
||||||
let liveness_clock = Instant::now();
|
|
||||||
let last_seen_ms = Arc::new(AtomicU64::new(0));
|
|
||||||
|
|
||||||
let reader_hub = hub.clone();
|
|
||||||
let reader_conn = conn.clone();
|
|
||||||
let reader_last_seen_ms = last_seen_ms.clone();
|
|
||||||
let liveness_conn = conn.clone();
|
|
||||||
let mut reader = tokio::spawn(async move {
|
|
||||||
run_worker_reader(
|
|
||||||
ws_rx,
|
|
||||||
reader_hub,
|
|
||||||
conn_id,
|
|
||||||
reader_conn,
|
|
||||||
reader_last_seen_ms,
|
|
||||||
liveness_clock,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
});
|
|
||||||
|
|
||||||
let liveness_last_seen_ms = last_seen_ms.clone();
|
|
||||||
let mut liveness = tokio::spawn(async move {
|
|
||||||
run_worker_liveness(
|
|
||||||
conn_id,
|
|
||||||
liveness_conn,
|
|
||||||
cfg.ping_interval,
|
|
||||||
cfg.idle_timeout,
|
|
||||||
liveness_last_seen_ms,
|
|
||||||
liveness_clock,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
});
|
|
||||||
|
|
||||||
let reader_finished = tokio::select! {
|
|
||||||
res = &mut reader => {
|
|
||||||
if let Err(err) = res {
|
|
||||||
warn!(worker_id = conn_id, error = %err, "worker reader task failed");
|
|
||||||
}
|
|
||||||
true
|
|
||||||
}
|
|
||||||
res = &mut liveness => {
|
|
||||||
if let Err(err) = res {
|
|
||||||
warn!(worker_id = conn_id, error = %err, "worker liveness task failed");
|
|
||||||
}
|
|
||||||
false
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
conn.request_close();
|
|
||||||
if !reader_finished {
|
|
||||||
reader.abort();
|
|
||||||
let _ = reader.await;
|
|
||||||
}
|
|
||||||
if reader_finished {
|
|
||||||
liveness.abort();
|
|
||||||
let _ = liveness.await;
|
|
||||||
}
|
|
||||||
|
|
||||||
hub.unregister_worker(conn_id);
|
|
||||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
||||||
writer.abort();
|
|
||||||
let _ = writer.await;
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn run_worker_reader(
|
|
||||||
mut ws_rx: futures_util::stream::SplitStream<WebSocket>,
|
|
||||||
hub: Arc<HubRouter>,
|
|
||||||
conn_id: u64,
|
|
||||||
conn: Arc<WorkerConn>,
|
|
||||||
last_seen_ms: Arc<AtomicU64>,
|
|
||||||
liveness_clock: Instant,
|
|
||||||
) {
|
|
||||||
loop {
|
|
||||||
match ws_rx.next().await {
|
|
||||||
Some(Ok(Message::Binary(data))) => {
|
|
||||||
last_seen_ms.store(elapsed_millis(liveness_clock), Ordering::Relaxed);
|
|
||||||
|
|
||||||
let mut data = data.to_vec();
|
|
||||||
if data.len() < protocol::HEADER_SIZE {
|
|
||||||
debug!(worker_id = conn_id, "frame too small, skipping");
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let header = match protocol::FrameHeader::parse(&data) {
|
|
||||||
Some(h) => h,
|
|
||||||
None => continue,
|
|
||||||
};
|
|
||||||
|
|
||||||
if header.msg_type == protocol::HEARTBEAT_ACK {
|
|
||||||
hub.handle_worker_heartbeat_ack(&mut data);
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Some(err_msg) = hub.handle_worker_frame(conn_id, &mut data) {
|
|
||||||
let err_frame = protocol::encode_stream_error(header.stream_id, &err_msg);
|
|
||||||
let _ = conn.send(Message::Binary(err_frame.into()));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Some(Ok(Message::Close(_))) | None => {
|
|
||||||
info!(worker_id = conn_id, "worker WebSocket closed");
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
Some(Err(e)) => {
|
|
||||||
warn!(worker_id = conn_id, error = %e, "worker WebSocket error");
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
_ => {}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn run_worker_liveness(
|
|
||||||
conn_id: u64,
|
|
||||||
conn: Arc<WorkerConn>,
|
|
||||||
ping_interval: Duration,
|
|
||||||
idle_timeout: Duration,
|
|
||||||
last_seen_ms: Arc<AtomicU64>,
|
|
||||||
liveness_clock: Instant,
|
|
||||||
) {
|
|
||||||
let ping_interval_ms = ping_interval.as_millis().max(1) as u64;
|
|
||||||
let idle_timeout_ms = idle_timeout.as_millis() as u64;
|
|
||||||
|
|
||||||
loop {
|
|
||||||
tokio::time::sleep(ping_interval).await;
|
|
||||||
|
|
||||||
let ping = protocol::encode_ping();
|
|
||||||
if !matches!(conn.send(Message::Binary(ping.into())), SendStatus::Queued) {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
if idle_timeout.is_zero() {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let now_ms = elapsed_millis(liveness_clock);
|
|
||||||
let last_seen = last_seen_ms.load(Ordering::Relaxed);
|
|
||||||
let silent_for_ms = now_ms.saturating_sub(last_seen);
|
|
||||||
if silent_for_ms < idle_timeout_ms {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
|
|
||||||
let missed_heartbeats = (silent_for_ms / ping_interval_ms).max(1);
|
|
||||||
warn!(
|
|
||||||
worker_id = conn_id,
|
|
||||||
idle_timeout_secs = idle_timeout.as_secs(),
|
|
||||||
silent_for_ms = silent_for_ms,
|
|
||||||
missed_heartbeats = missed_heartbeats,
|
|
||||||
"worker heartbeat timeout"
|
|
||||||
);
|
|
||||||
let _ = conn.send(Message::Binary(protocol::encode_goaway().into()));
|
|
||||||
conn.request_close();
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn elapsed_millis(started_at: Instant) -> u64 {
|
|
||||||
started_at.elapsed().as_millis().min(u64::MAX as u128) as u64
|
|
||||||
}
|
|
||||||
@@ -1,352 +0,0 @@
|
|||||||
"""
|
|
||||||
WebSocket 隧道端点
|
|
||||||
|
|
||||||
aether-proxy 通过此端点建立 tunnel 连接。
|
|
||||||
路径: /api/internal/proxy-tunnel
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import os
|
|
||||||
from datetime import datetime, timezone
|
|
||||||
|
|
||||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
||||||
|
|
||||||
from src.core.logger import logger
|
|
||||||
from src.services.proxy_node.health_scheduler import heartbeat_is_stale
|
|
||||||
from src.services.proxy_node.tunnel_manager import (
|
|
||||||
TunnelConnection,
|
|
||||||
get_tunnel_manager,
|
|
||||||
)
|
|
||||||
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
|
|
||||||
|
|
||||||
router = APIRouter()
|
|
||||||
|
|
||||||
# Per-node 锁: 防止并发的 connect/disconnect 写入 DB 时出现竞态(后断连覆盖先连接)
|
|
||||||
_node_status_locks: dict[str, asyncio.Lock] = {}
|
|
||||||
|
|
||||||
# Per-node idle timeout 连续计数: 用于抑制重复日志
|
|
||||||
_idle_timeout_counts: dict[str, int] = {}
|
|
||||||
|
|
||||||
|
|
||||||
def _get_node_lock(node_id: str) -> asyncio.Lock:
|
|
||||||
lock = _node_status_locks.get(node_id)
|
|
||||||
if lock is None:
|
|
||||||
lock = asyncio.Lock()
|
|
||||||
_node_status_locks[node_id] = lock
|
|
||||||
return lock
|
|
||||||
|
|
||||||
|
|
||||||
# 单帧最大 64 MB -- AI API 请求体可能包含多张 base64 图片,需要足够余量
|
|
||||||
_MAX_FRAME_SIZE = 64 * 1024 * 1024
|
|
||||||
|
|
||||||
# 默认 WebSocket 空闲超时(秒)-- 0 表示禁用(依赖 PING/PONG 心跳检测连接存活)
|
|
||||||
_DEFAULT_IDLE_TIMEOUT = 0.0
|
|
||||||
|
|
||||||
# 默认服务端应用层 ping 间隔(秒)-- 与客户端 ping 间隔一致,确保高频心跳
|
|
||||||
_DEFAULT_SERVER_PING_INTERVAL = 15.0
|
|
||||||
|
|
||||||
|
|
||||||
def _env_float(name: str, default: float, *, min_value: float, max_value: float) -> float:
|
|
||||||
"""读取并校验浮点环境变量,非法时回退默认值。"""
|
|
||||||
raw = os.getenv(name, "").strip()
|
|
||||||
if not raw:
|
|
||||||
return default
|
|
||||||
try:
|
|
||||||
value = float(raw)
|
|
||||||
except ValueError:
|
|
||||||
logger.warning("invalid {}={}, fallback to {}", name, raw, default)
|
|
||||||
return default
|
|
||||||
if value < min_value or value > max_value:
|
|
||||||
logger.warning(
|
|
||||||
"{}={} out of range [{}, {}], fallback to {}",
|
|
||||||
name,
|
|
||||||
value,
|
|
||||||
min_value,
|
|
||||||
max_value,
|
|
||||||
default,
|
|
||||||
)
|
|
||||||
return default
|
|
||||||
return value
|
|
||||||
|
|
||||||
|
|
||||||
# 允许通过环境变量在弱网环境下调大容忍窗口(无需改代码)
|
|
||||||
_SERVER_PING_INTERVAL = _env_float(
|
|
||||||
"AETHER_PROXY_TUNNEL_SERVER_PING_INTERVAL",
|
|
||||||
_DEFAULT_SERVER_PING_INTERVAL,
|
|
||||||
min_value=5.0,
|
|
||||||
max_value=120.0,
|
|
||||||
)
|
|
||||||
_IDLE_TIMEOUT = _env_float(
|
|
||||||
"AETHER_PROXY_TUNNEL_SERVER_IDLE_TIMEOUT",
|
|
||||||
_DEFAULT_IDLE_TIMEOUT,
|
|
||||||
min_value=0.0,
|
|
||||||
max_value=600.0,
|
|
||||||
)
|
|
||||||
|
|
||||||
# 避免 idle timeout 过小导致 ping 尚未生效就被服务端断开(0 表示禁用,跳过校验)
|
|
||||||
if _IDLE_TIMEOUT > 0 and _IDLE_TIMEOUT <= _SERVER_PING_INTERVAL * 2:
|
|
||||||
adjusted_idle = max(_SERVER_PING_INTERVAL * 3, 30.0)
|
|
||||||
logger.warning(
|
|
||||||
"AETHER_PROXY_TUNNEL_SERVER_IDLE_TIMEOUT too low for ping interval, auto-adjust to {}",
|
|
||||||
adjusted_idle,
|
|
||||||
)
|
|
||||||
_IDLE_TIMEOUT = adjusted_idle
|
|
||||||
|
|
||||||
|
|
||||||
async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
|
|
||||||
"""验证 WebSocket 连接的认证信息,返回 (node_id, node_name) 或 None
|
|
||||||
|
|
||||||
认证方式:Bearer <management_token>,通过 Management Token 系统验证。
|
|
||||||
authenticate_management_token 是 async 方法(内部有 Redis 速率限制),
|
|
||||||
因此直接 await 调用。节点存在性检查复用同一 session。
|
|
||||||
"""
|
|
||||||
auth = ws.headers.get("authorization", "")
|
|
||||||
if not auth.startswith("Bearer "):
|
|
||||||
return None
|
|
||||||
|
|
||||||
token = auth[7:]
|
|
||||||
if not token or not token.startswith("ae_"):
|
|
||||||
return None
|
|
||||||
|
|
||||||
client_ip = getattr(ws.client, "host", "unknown") if ws.client else "unknown"
|
|
||||||
node_id_header = ws.headers.get("x-node-id", "").strip()
|
|
||||||
node_name_header = ws.headers.get("x-node-name", "").strip()
|
|
||||||
|
|
||||||
if not node_id_header:
|
|
||||||
return None
|
|
||||||
|
|
||||||
from src.database import create_session
|
|
||||||
from src.models.database import ProxyNode
|
|
||||||
from src.services.auth.service import AuthService
|
|
||||||
|
|
||||||
db = create_session()
|
|
||||||
try:
|
|
||||||
result = await AuthService.authenticate_management_token(db, token, client_ip)
|
|
||||||
if not result:
|
|
||||||
return None
|
|
||||||
|
|
||||||
# 节点存在性检查(复用同一 session,避免额外连接开销)
|
|
||||||
exists = db.query(
|
|
||||||
db.query(ProxyNode).filter(ProxyNode.id == node_id_header).exists()
|
|
||||||
).scalar()
|
|
||||||
if not exists:
|
|
||||||
logger.warning("tunnel auth: node_id={} not found in DB", node_id_header)
|
|
||||||
return None
|
|
||||||
finally:
|
|
||||||
db.close()
|
|
||||||
|
|
||||||
return node_id_header, node_name_header or node_id_header
|
|
||||||
|
|
||||||
|
|
||||||
@router.websocket("/api/internal/proxy-tunnel")
|
|
||||||
async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
|
||||||
"""aether-proxy tunnel WebSocket 端点"""
|
|
||||||
# 先 accept,避免认证(DB/Redis)慢时卡在握手阶段导致网关返回 502。
|
|
||||||
await ws.accept()
|
|
||||||
|
|
||||||
try:
|
|
||||||
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:
|
|
||||||
logger.warning("tunnel auth error: {}", e)
|
|
||||||
await ws.close(code=4002, reason="authentication error")
|
|
||||||
return
|
|
||||||
|
|
||||||
if not auth:
|
|
||||||
await ws.close(code=4001, reason="unauthorized")
|
|
||||||
return
|
|
||||||
|
|
||||||
node_id: str = auth[0]
|
|
||||||
node_name: str = auth[1]
|
|
||||||
|
|
||||||
# Read proxy-advertised max concurrent streams (backward-compatible:
|
|
||||||
# old proxies don't send this header, we fall back to the default).
|
|
||||||
max_streams_raw = ws.headers.get("x-tunnel-max-streams", "").strip()
|
|
||||||
max_streams: int | None = None
|
|
||||||
if max_streams_raw:
|
|
||||||
try:
|
|
||||||
max_streams = int(max_streams_raw)
|
|
||||||
except ValueError:
|
|
||||||
pass
|
|
||||||
|
|
||||||
manager = get_tunnel_manager()
|
|
||||||
conn = TunnelConnection(node_id, node_name, ws, max_streams=max_streams)
|
|
||||||
node_lock = _get_node_lock(node_id)
|
|
||||||
|
|
||||||
manager.register(conn)
|
|
||||||
|
|
||||||
# 在 per-node 锁保护下更新 DB,防止并发的 connect/disconnect 写入竞态
|
|
||||||
async with node_lock:
|
|
||||||
await _update_tunnel_status(
|
|
||||||
node_id,
|
|
||||||
connected=True,
|
|
||||||
observed_at=datetime.now(timezone.utc),
|
|
||||||
)
|
|
||||||
|
|
||||||
# 启动服务端 ping 任务,防止中间代理因空闲超时关闭连接
|
|
||||||
ping_task = asyncio.create_task(_ping_loop(conn))
|
|
||||||
|
|
||||||
disconnect_reason: str | None = None
|
|
||||||
try:
|
|
||||||
oversized_count = 0
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
if _IDLE_TIMEOUT > 0:
|
|
||||||
data = await asyncio.wait_for(ws.receive_bytes(), timeout=_IDLE_TIMEOUT)
|
|
||||||
else:
|
|
||||||
data = await ws.receive_bytes()
|
|
||||||
except asyncio.TimeoutError:
|
|
||||||
count = _idle_timeout_counts.get(node_id, 0) + 1
|
|
||||||
_idle_timeout_counts[node_id] = count
|
|
||||||
# 首次 warning,后续每 10 次打印一条 info,其余 debug
|
|
||||||
if count == 1:
|
|
||||||
logger.warning("tunnel idle timeout for node_id={}", node_id)
|
|
||||||
elif count % 10 == 0:
|
|
||||||
logger.info(
|
|
||||||
"tunnel idle timeout for node_id={} (repeated {} times)",
|
|
||||||
node_id,
|
|
||||||
count,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logger.debug("tunnel idle timeout for node_id={}", node_id)
|
|
||||||
disconnect_reason = "idle timeout"
|
|
||||||
await ws.close(code=4004, reason="idle timeout")
|
|
||||||
break
|
|
||||||
if len(data) > _MAX_FRAME_SIZE:
|
|
||||||
oversized_count += 1
|
|
||||||
logger.warning("tunnel frame too large from {}: {} bytes", node_id, len(data))
|
|
||||||
if oversized_count >= 5:
|
|
||||||
logger.warning("too many oversized frames from {}, closing", node_id)
|
|
||||||
disconnect_reason = "too many oversized frames"
|
|
||||||
await ws.close(code=4003, reason="too many oversized frames")
|
|
||||||
break
|
|
||||||
continue
|
|
||||||
oversized_count = 0 # 正常帧重置计数
|
|
||||||
_idle_timeout_counts.pop(node_id, None) # 收到正常帧,重置 idle timeout 计数
|
|
||||||
manager.reset_reconnect_count(node_id) # 连接正常活跃,重置 reconnect 计数
|
|
||||||
try:
|
|
||||||
frame = Frame.decode(data)
|
|
||||||
except ValueError as e:
|
|
||||||
logger.warning("tunnel frame decode error from {}: {}", node_id, e)
|
|
||||||
continue
|
|
||||||
|
|
||||||
await manager.handle_incoming_frame(conn, frame)
|
|
||||||
|
|
||||||
except WebSocketDisconnect:
|
|
||||||
disconnect_reason = "WebSocket disconnected"
|
|
||||||
logger.info("tunnel WebSocket disconnected: node_id={}", node_id)
|
|
||||||
except Exception as e:
|
|
||||||
disconnect_reason = f"error: {e}"
|
|
||||||
logger.error("tunnel WebSocket error for node_id={}: {}", node_id, e)
|
|
||||||
finally:
|
|
||||||
ping_task.cancel()
|
|
||||||
# 在 per-node 锁保护下执行 unregister + 连接池计数检查 + DB 更新,
|
|
||||||
# 确保整个序列是原子的,避免"断连写 OFFLINE 覆盖新连接写 ONLINE"的竞态
|
|
||||||
async with node_lock:
|
|
||||||
manager.unregister(conn)
|
|
||||||
if manager.connection_count(node_id) == 0:
|
|
||||||
await _update_tunnel_status(
|
|
||||||
node_id,
|
|
||||||
connected=False,
|
|
||||||
detail=disconnect_reason,
|
|
||||||
observed_at=datetime.now(timezone.utc),
|
|
||||||
)
|
|
||||||
# 不清理锁: asyncio.Lock 极轻量,清理可能导致并发新连接拿到不同锁实例
|
|
||||||
else:
|
|
||||||
logger.info("tunnel connection closed but pool still active: node_id={}", node_id)
|
|
||||||
|
|
||||||
|
|
||||||
async def _ping_loop(conn: TunnelConnection) -> None:
|
|
||||||
"""定期发送应用层 PING 帧,保持连接活跃"""
|
|
||||||
try:
|
|
||||||
while True:
|
|
||||||
await asyncio.sleep(_SERVER_PING_INTERVAL)
|
|
||||||
if not conn.is_alive:
|
|
||||||
break
|
|
||||||
try:
|
|
||||||
await conn.send_frame(Frame(0, MsgType.PING, 0, b""))
|
|
||||||
except Exception as e:
|
|
||||||
logger.debug("ping loop send failed for node_id={}: {}", conn.node_id, e)
|
|
||||||
break
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
async def _update_tunnel_status(
|
|
||||||
node_id: str,
|
|
||||||
*,
|
|
||||||
connected: bool,
|
|
||||||
detail: str | None = None,
|
|
||||||
observed_at: datetime | None = None,
|
|
||||||
) -> None:
|
|
||||||
"""更新 ProxyNode 的 tunnel 连接状态并记录事件(在线程池中执行)"""
|
|
||||||
|
|
||||||
def _sync_update() -> None:
|
|
||||||
from src.database import create_session
|
|
||||||
from src.models.database import ProxyNode, ProxyNodeEvent, ProxyNodeStatus
|
|
||||||
|
|
||||||
db = create_session()
|
|
||||||
try:
|
|
||||||
node = db.query(ProxyNode).filter(ProxyNode.id == node_id).first()
|
|
||||||
if node:
|
|
||||||
event_time = observed_at or datetime.now(timezone.utc)
|
|
||||||
last_transition = node.tunnel_connected_at
|
|
||||||
if last_transition and last_transition.tzinfo is None:
|
|
||||||
last_transition = last_transition.replace(tzinfo=timezone.utc)
|
|
||||||
|
|
||||||
# 忽略乱序的旧事件,避免快速重连时旧状态覆盖新状态
|
|
||||||
stale_event = bool(last_transition and event_time < last_transition)
|
|
||||||
if stale_event:
|
|
||||||
detail_text = f"[stale_ignored] {detail}" if detail else "[stale_ignored]"
|
|
||||||
db.add(
|
|
||||||
ProxyNodeEvent(
|
|
||||||
node_id=node_id,
|
|
||||||
event_type="connected" if connected else "disconnected",
|
|
||||||
detail=detail_text,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
db.commit()
|
|
||||||
return
|
|
||||||
|
|
||||||
event_detail = detail
|
|
||||||
if connected:
|
|
||||||
node.tunnel_connected = True
|
|
||||||
node.tunnel_connected_at = event_time
|
|
||||||
node.status = ProxyNodeStatus.ONLINE
|
|
||||||
else:
|
|
||||||
# 断连不立即强制 OFFLINE。若心跳仍新鲜,可能仍有其他连接存活
|
|
||||||
# (连接池或跨 worker),避免误判写回 OFFLINE。
|
|
||||||
if heartbeat_is_stale(node, event_time):
|
|
||||||
node.tunnel_connected = False
|
|
||||||
node.tunnel_connected_at = event_time
|
|
||||||
node.status = ProxyNodeStatus.OFFLINE
|
|
||||||
else:
|
|
||||||
event_detail = (
|
|
||||||
f"[heartbeat_fresh] {detail}" if detail else "[heartbeat_fresh]"
|
|
||||||
)
|
|
||||||
|
|
||||||
# 记录连接事件
|
|
||||||
event = ProxyNodeEvent(
|
|
||||||
node_id=node_id,
|
|
||||||
event_type="connected" if connected else "disconnected",
|
|
||||||
detail=event_detail,
|
|
||||||
)
|
|
||||||
db.add(event)
|
|
||||||
db.commit()
|
|
||||||
finally:
|
|
||||||
db.close()
|
|
||||||
|
|
||||||
try:
|
|
||||||
await asyncio.to_thread(_sync_update)
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("failed to update tunnel status for {}: {}", node_id, e)
|
|
||||||
|
|
||||||
# 清除节点信息缓存,确保后续请求能立即感知连接状态变化
|
|
||||||
from src.services.proxy_node.resolver import invalidate_proxy_node_cache
|
|
||||||
|
|
||||||
invalidate_proxy_node_cache(node_id)
|
|
||||||
3
src/api/internal/__init__.py
Normal file
3
src/api/internal/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
|||||||
|
from .hub import router
|
||||||
|
|
||||||
|
__all__ = ["router"]
|
||||||
97
src/api/internal/hub.py
Normal file
97
src/api/internal/hub.py
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import ipaddress
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from fastapi import APIRouter, HTTPException, Request
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
from src.services.proxy_node.service import ProxyNodeService, build_heartbeat_ack
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/api/internal/hub", tags=["Internal - Hub"], include_in_schema=False)
|
||||||
|
|
||||||
|
|
||||||
|
class HubHeartbeatRequest(BaseModel):
|
||||||
|
node_id: str = Field(..., min_length=1, max_length=36)
|
||||||
|
heartbeat_interval: int | None = Field(None, ge=5, le=600)
|
||||||
|
active_connections: int | None = Field(None, ge=0)
|
||||||
|
total_requests: int | None = Field(None, ge=0)
|
||||||
|
avg_latency_ms: float | None = Field(None, ge=0)
|
||||||
|
failed_requests: int | None = Field(None, ge=0)
|
||||||
|
dns_failures: int | None = Field(None, ge=0)
|
||||||
|
stream_errors: int | None = Field(None, ge=0)
|
||||||
|
proxy_metadata: dict[str, Any] | None = None
|
||||||
|
proxy_version: str | None = Field(None, max_length=20)
|
||||||
|
|
||||||
|
|
||||||
|
class HubNodeStatusRequest(BaseModel):
|
||||||
|
node_id: str = Field(..., min_length=1, max_length=36)
|
||||||
|
connected: bool
|
||||||
|
conn_count: int = Field(0, ge=0)
|
||||||
|
|
||||||
|
|
||||||
|
def _ensure_loopback(request: Request) -> None:
|
||||||
|
host = request.client.host if request.client else ""
|
||||||
|
try:
|
||||||
|
if not ipaddress.ip_address(host).is_loopback:
|
||||||
|
raise ValueError(host)
|
||||||
|
except ValueError as exc:
|
||||||
|
raise HTTPException(status_code=403, detail="loopback access only") from exc
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/heartbeat")
|
||||||
|
async def hub_heartbeat(request: Request, payload: HubHeartbeatRequest) -> dict[str, Any]:
|
||||||
|
_ensure_loopback(request)
|
||||||
|
|
||||||
|
def _sync_apply() -> dict[str, Any]:
|
||||||
|
from src.database import create_session
|
||||||
|
|
||||||
|
db = create_session()
|
||||||
|
try:
|
||||||
|
node = ProxyNodeService.heartbeat(
|
||||||
|
db,
|
||||||
|
node_id=payload.node_id,
|
||||||
|
heartbeat_interval=payload.heartbeat_interval,
|
||||||
|
active_connections=payload.active_connections,
|
||||||
|
total_requests=payload.total_requests,
|
||||||
|
avg_latency_ms=payload.avg_latency_ms,
|
||||||
|
failed_requests=payload.failed_requests,
|
||||||
|
dns_failures=payload.dns_failures,
|
||||||
|
stream_errors=payload.stream_errors,
|
||||||
|
proxy_metadata=payload.proxy_metadata,
|
||||||
|
proxy_version=payload.proxy_version,
|
||||||
|
)
|
||||||
|
return build_heartbeat_ack(node)
|
||||||
|
finally:
|
||||||
|
db.close()
|
||||||
|
|
||||||
|
try:
|
||||||
|
return await asyncio.to_thread(_sync_apply)
|
||||||
|
except Exception as exc:
|
||||||
|
raise HTTPException(status_code=500, detail=f"heartbeat sync failed: {exc}") from exc
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/node-status")
|
||||||
|
async def hub_node_status(request: Request, payload: HubNodeStatusRequest) -> dict[str, Any]:
|
||||||
|
_ensure_loopback(request)
|
||||||
|
|
||||||
|
def _sync_apply() -> dict[str, Any]:
|
||||||
|
from src.database import create_session
|
||||||
|
|
||||||
|
db = create_session()
|
||||||
|
try:
|
||||||
|
node = ProxyNodeService.update_tunnel_status(
|
||||||
|
db,
|
||||||
|
node_id=payload.node_id,
|
||||||
|
connected=payload.connected,
|
||||||
|
conn_count=payload.conn_count,
|
||||||
|
)
|
||||||
|
return {"updated": node is not None}
|
||||||
|
finally:
|
||||||
|
db.close()
|
||||||
|
|
||||||
|
try:
|
||||||
|
return await asyncio.to_thread(_sync_apply)
|
||||||
|
except Exception as exc:
|
||||||
|
raise HTTPException(status_code=500, detail=f"node status sync failed: {exc}") from exc
|
||||||
@@ -21,6 +21,7 @@ from src.api.announcements import router as announcement_router
|
|||||||
# API路由
|
# API路由
|
||||||
from src.api.auth import router as auth_router
|
from src.api.auth import router as auth_router
|
||||||
from src.api.dashboard import router as dashboard_router
|
from src.api.dashboard import router as dashboard_router
|
||||||
|
from src.api.internal import router as internal_router
|
||||||
from src.api.monitoring import router as monitoring_router
|
from src.api.monitoring import router as monitoring_router
|
||||||
from src.api.payment import router as payment_router
|
from src.api.payment import router as payment_router
|
||||||
from src.api.public import router as public_router
|
from src.api.public import router as public_router
|
||||||
@@ -715,6 +716,7 @@ app.include_router(announcement_router) # 公告系统
|
|||||||
app.include_router(dashboard_router) # 仪表盘端点
|
app.include_router(dashboard_router) # 仪表盘端点
|
||||||
app.include_router(public_router) # 公开API端点(用户可查看提供商和模型)
|
app.include_router(public_router) # 公开API端点(用户可查看提供商和模型)
|
||||||
app.include_router(monitoring_router) # 监控端点
|
app.include_router(monitoring_router) # 监控端点
|
||||||
|
app.include_router(internal_router) # Hub 本地控制面端点
|
||||||
|
|
||||||
|
|
||||||
@app.get("/readyz", include_in_schema=False)
|
@app.get("/readyz", include_in_schema=False)
|
||||||
|
|||||||
@@ -101,23 +101,13 @@ async def _on_startup() -> None:
|
|||||||
else:
|
else:
|
||||||
logger.info("检测到其他 worker 已运行 ProxyNode 心跳检测,本实例跳过")
|
logger.info("检测到其他 worker 已运行 ProxyNode 心跳检测,本实例跳过")
|
||||||
|
|
||||||
# 在可能的状态重置之后再建立 /worker 长连接,避免“先同步在线,再被重置离线”的竞态。
|
|
||||||
from src.services.proxy_node.hub_transport import get_hub_connection_manager
|
|
||||||
|
|
||||||
try:
|
|
||||||
await get_hub_connection_manager().ensure_connected()
|
|
||||||
logger.info("Hub worker channel initialized on startup")
|
|
||||||
except Exception as e:
|
|
||||||
# ensure_connected 失败时内部会启动重连循环,这里仅记录告警不阻塞启动
|
|
||||||
logger.warning("Hub worker channel init failed, reconnecting in background: %s", e)
|
|
||||||
|
|
||||||
if active:
|
if active:
|
||||||
logger.info("启动 ProxyNode 心跳检测调度器...")
|
logger.info("启动 ProxyNode 心跳检测调度器...")
|
||||||
await proxy_node_health_scheduler.start()
|
await proxy_node_health_scheduler.start()
|
||||||
|
|
||||||
|
|
||||||
async def _on_shutdown() -> None:
|
async def _on_shutdown() -> None:
|
||||||
"""优雅关闭 tunnel 连接并停止心跳检测调度器"""
|
"""停止心跳检测调度器"""
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
from src.services.proxy_node.health_scheduler import get_proxy_node_health_scheduler
|
from src.services.proxy_node.health_scheduler import get_proxy_node_health_scheduler
|
||||||
@@ -125,10 +115,6 @@ async def _on_shutdown() -> None:
|
|||||||
|
|
||||||
logger = logging.getLogger("aether.modules.proxy_nodes")
|
logger = logging.getLogger("aether.modules.proxy_nodes")
|
||||||
|
|
||||||
from src.services.proxy_node.hub_transport import shutdown_hub_connection_manager
|
|
||||||
|
|
||||||
await shutdown_hub_connection_manager()
|
|
||||||
|
|
||||||
from src.clients import get_redis_client
|
from src.clients import get_redis_client
|
||||||
|
|
||||||
redis_client = await get_redis_client()
|
redis_client = await get_redis_client()
|
||||||
|
|||||||
@@ -12,13 +12,10 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import os
|
import os
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
from urllib.parse import quote
|
||||||
|
|
||||||
_DOCKER_HUB_URL = "ws://127.0.0.1:8085"
|
_DOCKER_HUB_URL = "http://127.0.0.1:8085"
|
||||||
_DOCKER_HUB_CONNECT_TIMEOUT_SECONDS = 5.0
|
_DOCKER_HUB_CONNECT_TIMEOUT_SECONDS = 5.0
|
||||||
_DOCKER_HUB_PING_INTERVAL_SECONDS = 15.0
|
|
||||||
_DOCKER_HUB_SEND_TIMEOUT_SECONDS = 10.0
|
|
||||||
_DOCKER_HUB_MAX_STREAMS = 2048
|
|
||||||
_DOCKER_HUB_MAX_FRAME_SIZE = 64 * 1024 * 1024
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
@@ -26,14 +23,13 @@ class HubConfig:
|
|||||||
enabled: bool
|
enabled: bool
|
||||||
url: str
|
url: str
|
||||||
connect_timeout_seconds: float
|
connect_timeout_seconds: float
|
||||||
ping_interval_seconds: float
|
|
||||||
send_timeout_seconds: float
|
|
||||||
max_streams: int
|
|
||||||
max_frame_size: int
|
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def worker_ws_url(self) -> str:
|
def local_relay_base_url(self) -> str:
|
||||||
return f"{self.url.rstrip('/')}/worker"
|
return f"{self.url.rstrip('/')}/local/relay"
|
||||||
|
|
||||||
|
def local_relay_url(self, node_id: str) -> str:
|
||||||
|
return f"{self.local_relay_base_url}/{quote(node_id, safe='')}"
|
||||||
|
|
||||||
|
|
||||||
_hub_config: HubConfig | None = None
|
_hub_config: HubConfig | None = None
|
||||||
@@ -56,10 +52,6 @@ def get_hub_config() -> HubConfig:
|
|||||||
enabled=docker_runtime,
|
enabled=docker_runtime,
|
||||||
url=_DOCKER_HUB_URL,
|
url=_DOCKER_HUB_URL,
|
||||||
connect_timeout_seconds=_DOCKER_HUB_CONNECT_TIMEOUT_SECONDS,
|
connect_timeout_seconds=_DOCKER_HUB_CONNECT_TIMEOUT_SECONDS,
|
||||||
ping_interval_seconds=_DOCKER_HUB_PING_INTERVAL_SECONDS,
|
|
||||||
send_timeout_seconds=_DOCKER_HUB_SEND_TIMEOUT_SECONDS,
|
|
||||||
max_streams=_DOCKER_HUB_MAX_STREAMS,
|
|
||||||
max_frame_size=_DOCKER_HUB_MAX_FRAME_SIZE,
|
|
||||||
)
|
)
|
||||||
return _hub_config
|
return _hub_config
|
||||||
|
|
||||||
|
|||||||
@@ -1,43 +1,23 @@
|
|||||||
"""
|
"""
|
||||||
Hub 模式 tunnel transport
|
Hub 模式 tunnel transport
|
||||||
|
|
||||||
Worker 通过单条到 aether-hub 的 WebSocket 长连接转发 tunnel 帧。
|
Worker 通过本机 aether-hub 的 HTTP relay 访问 tunnel 数据面,不再维护 /worker WebSocket。
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import gzip
|
|
||||||
import json
|
import json
|
||||||
import time as _time
|
import struct
|
||||||
from datetime import datetime, timezone
|
from typing import TYPE_CHECKING
|
||||||
from typing import TYPE_CHECKING, Any
|
|
||||||
|
|
||||||
import aiohttp
|
|
||||||
import httpx
|
import httpx
|
||||||
from aiohttp import WSMsgType
|
|
||||||
|
|
||||||
from src.core.logger import logger
|
from .hub_config import get_hub_config
|
||||||
|
|
||||||
from .hub_config import HubConfig, get_hub_config
|
|
||||||
from .tunnel_manager import TunnelStreamError, _StreamState
|
|
||||||
from .tunnel_protocol import Frame, FrameFlags, MsgType, normalize_heartbeat_id
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncGenerator, Coroutine
|
from collections.abc import AsyncGenerator
|
||||||
|
|
||||||
|
|
||||||
_TUNNEL_COMPRESS_MIN_SIZE = 512
|
|
||||||
_TUNNEL_ASYNC_COMPRESS_THRESHOLD = 64 * 1024
|
|
||||||
_RECONNECT_DELAYS_SECONDS: tuple[float, ...] = (0.0, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0)
|
|
||||||
_HEARTBEAT_DEDUP_TTL_SECONDS = 600
|
|
||||||
_LOOP_WATCHDOG_INTERVAL_SECONDS = 1.0
|
|
||||||
_LOOP_LAG_WARNING_SECONDS = 1.0
|
|
||||||
_LOOP_LAG_DEGRADE_SECONDS = 3.0
|
|
||||||
_LOOP_LAG_DEGRADE_MIN_COOLDOWN_SECONDS = 2.0
|
|
||||||
_LOOP_LAG_DEGRADE_MAX_COOLDOWN_SECONDS = 5.0
|
|
||||||
_LOOP_LAG_WARNING_LOG_INTERVAL_SECONDS = 10.0
|
|
||||||
|
|
||||||
_HOP_BY_HOP_HEADERS = frozenset(
|
_HOP_BY_HOP_HEADERS = frozenset(
|
||||||
{
|
{
|
||||||
"host",
|
"host",
|
||||||
@@ -53,762 +33,105 @@ _HOP_BY_HOP_HEADERS = frozenset(
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
_HOP_BY_HOP_HEADERS_BYTES = frozenset(h.encode("ascii") for h in _HOP_BY_HOP_HEADERS)
|
_HOP_BY_HOP_HEADERS_BYTES = frozenset(h.encode("ascii") for h in _HOP_BY_HOP_HEADERS)
|
||||||
|
_RELAY_CONTENT_TYPE = "application/vnd.aether.tunnel-envelope"
|
||||||
|
_TUNNEL_ERROR_HEADER = "x-aether-tunnel-error"
|
||||||
class HubConnectionManager:
|
|
||||||
"""Worker 进程级 Hub 连接管理器(单例)。"""
|
|
||||||
|
|
||||||
def __init__(self, config: HubConfig | None = None) -> None:
|
|
||||||
self._config = config or get_hub_config()
|
|
||||||
self._session: aiohttp.ClientSession | None = None
|
|
||||||
self._ws: aiohttp.ClientWebSocketResponse | None = None
|
|
||||||
|
|
||||||
self._connect_lock = asyncio.Lock()
|
|
||||||
self._write_lock = asyncio.Lock()
|
|
||||||
|
|
||||||
self._next_stream_id = 2
|
|
||||||
self._pending_streams: dict[int, _StreamState] = {}
|
|
||||||
|
|
||||||
self._reader_task: asyncio.Task[None] | None = None
|
|
||||||
self._ping_task: asyncio.Task[None] | None = None
|
|
||||||
self._reconnect_task: asyncio.Task[None] | None = None
|
|
||||||
self._watchdog_task: asyncio.Task[None] | None = None
|
|
||||||
self._background_tasks: set[asyncio.Task[None]] = set()
|
|
||||||
|
|
||||||
self._closing = False
|
|
||||||
self._degraded_until: float = 0.0
|
|
||||||
self._last_loop_lag_warning_ts: float = 0.0
|
|
||||||
|
|
||||||
self._disconnect_count = 0 # 连续断开计数,用于抑制重复日志
|
|
||||||
# 连续快速断开退避:防止 Hub 端持续发送 GOAWAY 时产生重连风暴
|
|
||||||
self._last_disconnect_ts: float = 0.0
|
|
||||||
self._rapid_disconnect_count: int = 0
|
|
||||||
|
|
||||||
@property
|
|
||||||
def is_connected(self) -> bool:
|
|
||||||
ws = self._ws
|
|
||||||
return ws is not None and not ws.closed
|
|
||||||
|
|
||||||
def _background(self, coro: Coroutine[Any, Any, None]) -> None:
|
|
||||||
task = asyncio.create_task(coro)
|
|
||||||
self._background_tasks.add(task)
|
|
||||||
task.add_done_callback(self._background_tasks.discard)
|
|
||||||
|
|
||||||
async def ensure_connected(self) -> None:
|
|
||||||
self._ensure_watchdog_running()
|
|
||||||
if self._closing:
|
|
||||||
raise TunnelStreamError("hub connection manager is shutting down")
|
|
||||||
if self.is_connected:
|
|
||||||
return
|
|
||||||
|
|
||||||
async with self._connect_lock:
|
|
||||||
if self._closing:
|
|
||||||
raise TunnelStreamError("hub connection manager is shutting down")
|
|
||||||
if self.is_connected:
|
|
||||||
return
|
|
||||||
try:
|
|
||||||
await self._connect_once()
|
|
||||||
except Exception as e:
|
|
||||||
self._start_reconnect_loop()
|
|
||||||
raise TunnelStreamError(f"failed to connect hub worker channel: {e}") from e
|
|
||||||
|
|
||||||
async def _ensure_session(self) -> aiohttp.ClientSession:
|
|
||||||
if self._session is None or self._session.closed:
|
|
||||||
self._session = aiohttp.ClientSession()
|
|
||||||
return self._session
|
|
||||||
|
|
||||||
async def _connect_once(self) -> None:
|
|
||||||
session = await self._ensure_session()
|
|
||||||
|
|
||||||
ws = await session.ws_connect(
|
|
||||||
self._config.worker_ws_url,
|
|
||||||
timeout=self._config.connect_timeout_seconds,
|
|
||||||
autoping=False,
|
|
||||||
heartbeat=None,
|
|
||||||
max_msg_size=self._config.max_frame_size,
|
|
||||||
)
|
|
||||||
|
|
||||||
old_ws = self._ws
|
|
||||||
self._ws = ws
|
|
||||||
if old_ws is not None and not old_ws.closed:
|
|
||||||
try:
|
|
||||||
await old_ws.close()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
|
|
||||||
if self._reader_task is not None:
|
|
||||||
self._reader_task.cancel()
|
|
||||||
if self._ping_task is not None:
|
|
||||||
self._ping_task.cancel()
|
|
||||||
|
|
||||||
self._reader_task = asyncio.create_task(self._reader_loop(ws))
|
|
||||||
self._ping_task = asyncio.create_task(self._ping_loop(ws))
|
|
||||||
if self._disconnect_count == 0:
|
|
||||||
logger.info("Hub worker channel connected: {}", self._config.worker_ws_url)
|
|
||||||
else:
|
|
||||||
logger.debug(
|
|
||||||
"Hub worker channel connected: {} (after {} disconnects)",
|
|
||||||
self._config.worker_ws_url,
|
|
||||||
self._disconnect_count,
|
|
||||||
)
|
|
||||||
self._disconnect_count = 0
|
|
||||||
|
|
||||||
def _ensure_watchdog_running(self) -> None:
|
|
||||||
if self._closing:
|
|
||||||
return
|
|
||||||
if self._watchdog_task is not None and not self._watchdog_task.done():
|
|
||||||
return
|
|
||||||
self._watchdog_task = asyncio.create_task(self._loop_watchdog())
|
|
||||||
|
|
||||||
def _record_loop_lag(self, lag_seconds: float) -> None:
|
|
||||||
if lag_seconds < _LOOP_LAG_WARNING_SECONDS:
|
|
||||||
return
|
|
||||||
|
|
||||||
now = _time.monotonic()
|
|
||||||
if lag_seconds >= _LOOP_LAG_DEGRADE_SECONDS:
|
|
||||||
cooldown = min(
|
|
||||||
_LOOP_LAG_DEGRADE_MAX_COOLDOWN_SECONDS,
|
|
||||||
max(_LOOP_LAG_DEGRADE_MIN_COOLDOWN_SECONDS, lag_seconds * 1.5),
|
|
||||||
)
|
|
||||||
degraded_until = now + cooldown
|
|
||||||
self._degraded_until = max(self._degraded_until, degraded_until)
|
|
||||||
logger.warning(
|
|
||||||
"Hub worker event loop lag detected: lag={:.2f}s, pausing new streams for {:.1f}s",
|
|
||||||
lag_seconds,
|
|
||||||
cooldown,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
|
|
||||||
if now - self._last_loop_lag_warning_ts >= _LOOP_LAG_WARNING_LOG_INTERVAL_SECONDS:
|
|
||||||
self._last_loop_lag_warning_ts = now
|
|
||||||
logger.warning("Hub worker event loop lag observed: lag={:.2f}s", lag_seconds)
|
|
||||||
|
|
||||||
def _raise_if_degraded(self) -> None:
|
|
||||||
remaining = self._degraded_until - _time.monotonic()
|
|
||||||
if remaining <= 0:
|
|
||||||
return
|
|
||||||
raise TunnelStreamError(f"hub worker event loop degraded, retry in {remaining:.1f}s")
|
|
||||||
|
|
||||||
async def _loop_watchdog(self) -> None:
|
|
||||||
interval = _LOOP_WATCHDOG_INTERVAL_SECONDS
|
|
||||||
expected_at = _time.monotonic() + interval
|
|
||||||
try:
|
|
||||||
while not self._closing:
|
|
||||||
await asyncio.sleep(interval)
|
|
||||||
now = _time.monotonic()
|
|
||||||
lag_seconds = max(0.0, now - expected_at)
|
|
||||||
expected_at = now + interval
|
|
||||||
self._record_loop_lag(lag_seconds)
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
return
|
|
||||||
|
|
||||||
def _start_reconnect_loop(self) -> None:
|
|
||||||
if self._closing:
|
|
||||||
return
|
|
||||||
if not self._config.enabled:
|
|
||||||
return
|
|
||||||
if self._reconnect_task is not None and not self._reconnect_task.done():
|
|
||||||
return
|
|
||||||
self._reconnect_task = asyncio.create_task(self._reconnect_loop())
|
|
||||||
|
|
||||||
async def _reconnect_loop(self) -> None:
|
|
||||||
# 如果连续快速断开(GOAWAY 风暴),初始 attempt 跳过 0-delay 阶段
|
|
||||||
rapid = self._rapid_disconnect_count
|
|
||||||
attempt = min(rapid, len(_RECONNECT_DELAYS_SECONDS) - 1)
|
|
||||||
while not self._closing and not self.is_connected:
|
|
||||||
delay = _RECONNECT_DELAYS_SECONDS[min(attempt, len(_RECONNECT_DELAYS_SECONDS) - 1)]
|
|
||||||
if delay > 0:
|
|
||||||
await asyncio.sleep(delay)
|
|
||||||
try:
|
|
||||||
async with self._connect_lock:
|
|
||||||
if self._closing or self.is_connected:
|
|
||||||
break
|
|
||||||
await self._connect_once()
|
|
||||||
if self.is_connected:
|
|
||||||
logger.debug("Hub worker channel reconnected (attempt {})", attempt + 1)
|
|
||||||
break
|
|
||||||
except Exception as e:
|
|
||||||
attempt += 1
|
|
||||||
if attempt <= 3 or attempt % 10 == 0 or attempt in (20, 50, 100):
|
|
||||||
logger.debug("Hub reconnect attempt {} failed: {}", attempt, e)
|
|
||||||
|
|
||||||
async def _handle_disconnect(
|
|
||||||
self,
|
|
||||||
reason: str,
|
|
||||||
*,
|
|
||||||
ws: aiohttp.ClientWebSocketResponse | None = None,
|
|
||||||
) -> None:
|
|
||||||
current: aiohttp.ClientWebSocketResponse | None = None
|
|
||||||
async with self._connect_lock:
|
|
||||||
if self._ws is None:
|
|
||||||
return
|
|
||||||
if ws is not None and self._ws is not ws:
|
|
||||||
return
|
|
||||||
current = self._ws
|
|
||||||
self._ws = None
|
|
||||||
|
|
||||||
if current is not None and not current.closed:
|
|
||||||
try:
|
|
||||||
await current.close()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
|
|
||||||
if self._pending_streams:
|
|
||||||
affected_count = len(self._pending_streams)
|
|
||||||
affected_ids = list(self._pending_streams.keys())[:10] # 最多记录 10 个
|
|
||||||
logger.warning(
|
|
||||||
"Hub disconnect affecting {} in-flight streams: reason={}, stream_ids={}{}",
|
|
||||||
affected_count,
|
|
||||||
reason,
|
|
||||||
affected_ids,
|
|
||||||
"..." if affected_count > 10 else "",
|
|
||||||
)
|
|
||||||
for state in self._pending_streams.values():
|
|
||||||
state.set_error("hub disconnected")
|
|
||||||
self._pending_streams.clear()
|
|
||||||
|
|
||||||
if not self._closing:
|
|
||||||
self._disconnect_count += 1
|
|
||||||
|
|
||||||
# 追踪连续快速断开:如果距上次断开不足 2 秒,累加计数;否则重置
|
|
||||||
now_mono = _time.monotonic()
|
|
||||||
if now_mono - self._last_disconnect_ts < 2.0:
|
|
||||||
self._rapid_disconnect_count += 1
|
|
||||||
else:
|
|
||||||
self._rapid_disconnect_count = 0
|
|
||||||
self._last_disconnect_ts = now_mono
|
|
||||||
|
|
||||||
if self._disconnect_count <= 1:
|
|
||||||
logger.warning("Hub worker channel disconnected: {}", reason)
|
|
||||||
elif self._disconnect_count % 10 == 0:
|
|
||||||
logger.info(
|
|
||||||
"Hub worker channel disconnected: {} (repeated {} times)",
|
|
||||||
reason,
|
|
||||||
self._disconnect_count,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logger.debug("Hub worker channel disconnected: {}", reason)
|
|
||||||
self._start_reconnect_loop()
|
|
||||||
|
|
||||||
async def _send_frame(self, frame: Frame) -> None:
|
|
||||||
ws = self._ws
|
|
||||||
if ws is None or ws.closed:
|
|
||||||
raise TunnelStreamError("hub not connected")
|
|
||||||
|
|
||||||
try:
|
|
||||||
async with asyncio.timeout(self._config.send_timeout_seconds):
|
|
||||||
async with self._write_lock:
|
|
||||||
await ws.send_bytes(frame.encode())
|
|
||||||
except TimeoutError as e:
|
|
||||||
await self._handle_disconnect("send timeout", ws=ws)
|
|
||||||
raise TunnelStreamError("hub frame send timeout") from e
|
|
||||||
except Exception as e:
|
|
||||||
await self._handle_disconnect(f"send failed: {e}", ws=ws)
|
|
||||||
raise TunnelStreamError(f"hub frame send failed: {e}") from e
|
|
||||||
|
|
||||||
async def _reader_loop(self, ws: aiohttp.ClientWebSocketResponse) -> None:
|
|
||||||
try:
|
|
||||||
while not self._closing:
|
|
||||||
msg = await ws.receive()
|
|
||||||
|
|
||||||
if msg.type == WSMsgType.BINARY:
|
|
||||||
raw = msg.data
|
|
||||||
if isinstance(raw, memoryview):
|
|
||||||
raw = raw.tobytes()
|
|
||||||
elif isinstance(raw, bytearray):
|
|
||||||
raw = bytes(raw)
|
|
||||||
if not isinstance(raw, bytes):
|
|
||||||
continue
|
|
||||||
try:
|
|
||||||
frame = Frame.decode(raw)
|
|
||||||
except Exception as e:
|
|
||||||
logger.debug("invalid frame from hub: {}", e)
|
|
||||||
continue
|
|
||||||
await self._handle_incoming_frame(frame)
|
|
||||||
continue
|
|
||||||
|
|
||||||
if msg.type == WSMsgType.CLOSE or msg.type == WSMsgType.CLOSED:
|
|
||||||
break
|
|
||||||
|
|
||||||
if msg.type == WSMsgType.ERROR:
|
|
||||||
logger.debug("hub ws reader error: {}", ws.exception())
|
|
||||||
break
|
|
||||||
|
|
||||||
if msg.type == WSMsgType.PING:
|
|
||||||
payload = msg.data if isinstance(msg.data, bytes) else b""
|
|
||||||
self._background(self._send_pong(payload))
|
|
||||||
continue
|
|
||||||
|
|
||||||
# TEXT / PONG / 其他类型直接忽略
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
return
|
|
||||||
except Exception as e:
|
|
||||||
logger.debug("hub reader loop aborted: {}", e)
|
|
||||||
finally:
|
|
||||||
await self._handle_disconnect("reader ended", ws=ws)
|
|
||||||
|
|
||||||
async def _ping_loop(self, ws: aiohttp.ClientWebSocketResponse) -> None:
|
|
||||||
try:
|
|
||||||
while not self._closing:
|
|
||||||
await asyncio.sleep(self._config.ping_interval_seconds)
|
|
||||||
if self._ws is not ws or ws.closed:
|
|
||||||
break
|
|
||||||
try:
|
|
||||||
await self._send_frame(Frame(0, MsgType.PING, 0, b""))
|
|
||||||
except TunnelStreamError:
|
|
||||||
break
|
|
||||||
except asyncio.CancelledError:
|
|
||||||
return
|
|
||||||
|
|
||||||
async def _send_pong(self, payload: bytes) -> None:
|
|
||||||
try:
|
|
||||||
await self._send_frame(Frame(0, MsgType.PONG, 0, payload))
|
|
||||||
except TunnelStreamError:
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def _handle_incoming_frame(self, frame: Frame) -> None:
|
|
||||||
match frame.msg_type:
|
|
||||||
# -- stream-level frames --
|
|
||||||
case MsgType.RESPONSE_HEADERS:
|
|
||||||
stream = self._pending_streams.get(frame.stream_id)
|
|
||||||
if not stream:
|
|
||||||
return
|
|
||||||
try:
|
|
||||||
payload = _decompress_frame_payload(frame)
|
|
||||||
meta = json.loads(payload)
|
|
||||||
stream.set_response_headers(meta["status"], meta.get("headers", []))
|
|
||||||
except Exception as e:
|
|
||||||
stream.set_error(f"invalid response headers: {e}")
|
|
||||||
self._pending_streams.pop(frame.stream_id, None)
|
|
||||||
|
|
||||||
case MsgType.RESPONSE_BODY:
|
|
||||||
stream = self._pending_streams.get(frame.stream_id)
|
|
||||||
if stream:
|
|
||||||
stream.push_body_chunk(_decompress_frame_payload(frame))
|
|
||||||
|
|
||||||
case MsgType.STREAM_END:
|
|
||||||
stream = self._pending_streams.pop(frame.stream_id, None)
|
|
||||||
if stream:
|
|
||||||
stream.set_done()
|
|
||||||
|
|
||||||
case MsgType.STREAM_ERROR:
|
|
||||||
stream = self._pending_streams.pop(frame.stream_id, None)
|
|
||||||
if stream:
|
|
||||||
message = (
|
|
||||||
frame.payload.decode(errors="replace") if frame.payload else "stream error"
|
|
||||||
)
|
|
||||||
logger.warning(
|
|
||||||
"Hub received STREAM_ERROR: stream_id={}, message={}",
|
|
||||||
frame.stream_id,
|
|
||||||
message[:500],
|
|
||||||
)
|
|
||||||
stream.set_error(message)
|
|
||||||
|
|
||||||
# -- connection-level frames --
|
|
||||||
case MsgType.PING:
|
|
||||||
self._background(self._send_pong(frame.payload))
|
|
||||||
|
|
||||||
case MsgType.PONG:
|
|
||||||
pass
|
|
||||||
|
|
||||||
case MsgType.GOAWAY:
|
|
||||||
await self._handle_disconnect("received GOAWAY")
|
|
||||||
|
|
||||||
case MsgType.HEARTBEAT_DATA:
|
|
||||||
self._background(self._handle_heartbeat(frame))
|
|
||||||
|
|
||||||
case MsgType.HEARTBEAT_ACK:
|
|
||||||
pass
|
|
||||||
|
|
||||||
case MsgType.NODE_STATUS:
|
|
||||||
self._background(self._handle_node_status(frame.payload))
|
|
||||||
|
|
||||||
async def _handle_heartbeat(self, frame: Frame) -> None:
|
|
||||||
try:
|
|
||||||
data = json.loads(frame.payload) if frame.payload else {}
|
|
||||||
except Exception:
|
|
||||||
data = {}
|
|
||||||
|
|
||||||
node_id = str(data.get("node_id") or "").strip()
|
|
||||||
heartbeat_session_id = str(data.get("heartbeat_session_id") or "").strip()
|
|
||||||
if len(heartbeat_session_id) > 128:
|
|
||||||
heartbeat_session_id = heartbeat_session_id[:128]
|
|
||||||
heartbeat_id = normalize_heartbeat_id(data.get("heartbeat_id"))
|
|
||||||
ack: dict[str, object] = {}
|
|
||||||
if heartbeat_id is not None:
|
|
||||||
ack["heartbeat_id"] = heartbeat_id
|
|
||||||
|
|
||||||
should_process = True
|
|
||||||
if node_id and heartbeat_id is not None:
|
|
||||||
if heartbeat_session_id:
|
|
||||||
dedup_key = f"hub:heartbeat:{node_id}:{heartbeat_session_id}:{heartbeat_id}"
|
|
||||||
else:
|
|
||||||
dedup_key = f"hub:heartbeat:{node_id}:{heartbeat_id}"
|
|
||||||
try:
|
|
||||||
from src.clients import get_redis_client
|
|
||||||
|
|
||||||
redis = await get_redis_client()
|
|
||||||
if redis:
|
|
||||||
acquired = await redis.set(
|
|
||||||
dedup_key,
|
|
||||||
"1",
|
|
||||||
ex=_HEARTBEAT_DEDUP_TTL_SECONDS,
|
|
||||||
nx=True,
|
|
||||||
)
|
|
||||||
if not acquired:
|
|
||||||
should_process = False
|
|
||||||
except Exception:
|
|
||||||
# Redis 不可用时降级为不去重,避免心跳链路阻塞
|
|
||||||
pass
|
|
||||||
|
|
||||||
def _sync_heartbeat() -> dict[str, object]:
|
|
||||||
from src.database import create_session
|
|
||||||
from src.services.proxy_node.service import ProxyNodeService, build_heartbeat_ack
|
|
||||||
|
|
||||||
if not node_id:
|
|
||||||
return {}
|
|
||||||
|
|
||||||
db = create_session()
|
|
||||||
try:
|
|
||||||
node = ProxyNodeService.heartbeat(
|
|
||||||
db,
|
|
||||||
node_id=node_id,
|
|
||||||
active_connections=data.get("active_connections"),
|
|
||||||
total_requests=data.get("total_requests"),
|
|
||||||
avg_latency_ms=data.get("avg_latency_ms"),
|
|
||||||
failed_requests=data.get("failed_requests"),
|
|
||||||
dns_failures=data.get("dns_failures"),
|
|
||||||
stream_errors=data.get("stream_errors"),
|
|
||||||
proxy_metadata=data.get("proxy_metadata"),
|
|
||||||
proxy_version=data.get("proxy_version"),
|
|
||||||
)
|
|
||||||
return build_heartbeat_ack(node)
|
|
||||||
finally:
|
|
||||||
db.close()
|
|
||||||
|
|
||||||
if should_process:
|
|
||||||
try:
|
|
||||||
ack.update(await asyncio.to_thread(_sync_heartbeat))
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("hub heartbeat DB update failed: {}", e)
|
|
||||||
|
|
||||||
try:
|
|
||||||
await self._send_frame(
|
|
||||||
Frame(
|
|
||||||
frame.stream_id,
|
|
||||||
MsgType.HEARTBEAT_ACK,
|
|
||||||
0,
|
|
||||||
json.dumps(ack, ensure_ascii=False).encode("utf-8"),
|
|
||||||
)
|
|
||||||
)
|
|
||||||
except TunnelStreamError:
|
|
||||||
logger.debug("hub heartbeat ACK send failed")
|
|
||||||
|
|
||||||
async def _handle_node_status(self, payload: bytes) -> None:
|
|
||||||
try:
|
|
||||||
data = json.loads(payload) if payload else {}
|
|
||||||
except Exception:
|
|
||||||
return
|
|
||||||
|
|
||||||
node_id = str(data.get("node_id") or "").strip()
|
|
||||||
if not node_id:
|
|
||||||
return
|
|
||||||
|
|
||||||
connected = bool(data.get("connected"))
|
|
||||||
conn_count = int(data.get("conn_count") or 0)
|
|
||||||
|
|
||||||
# 所有 worker 都需要立即失效本地缓存,保证请求路由正确
|
|
||||||
try:
|
|
||||||
from src.services.proxy_node.resolver import invalidate_proxy_node_cache
|
|
||||||
|
|
||||||
invalidate_proxy_node_cache(node_id)
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
|
|
||||||
# 使用 Redis SETNX 去重:同一次 NODE_STATUS 广播只有一个 worker 执行 DB 写入,
|
|
||||||
# 避免 N 个 worker 并发写同一行并产生 N 条重复事件记录。
|
|
||||||
dedup_key = f"hub:node_status:{node_id}:{connected}:{conn_count}"
|
|
||||||
try:
|
|
||||||
from src.clients import get_redis_client
|
|
||||||
|
|
||||||
redis = await get_redis_client()
|
|
||||||
if redis:
|
|
||||||
acquired = await redis.set(dedup_key, "1", ex=10, nx=True)
|
|
||||||
if not acquired:
|
|
||||||
return
|
|
||||||
except Exception:
|
|
||||||
# Redis 不可用时不去重,允许重复写入(幂等)
|
|
||||||
pass
|
|
||||||
|
|
||||||
def _sync_update() -> None:
|
|
||||||
from src.database import create_session
|
|
||||||
from src.models.database import ProxyNode, ProxyNodeEvent, ProxyNodeStatus
|
|
||||||
|
|
||||||
db = create_session()
|
|
||||||
try:
|
|
||||||
node = db.query(ProxyNode).filter(ProxyNode.id == node_id).first()
|
|
||||||
if not node:
|
|
||||||
return
|
|
||||||
|
|
||||||
now = datetime.now(timezone.utc)
|
|
||||||
|
|
||||||
node.tunnel_connected = connected
|
|
||||||
if connected:
|
|
||||||
node.tunnel_connected_at = now
|
|
||||||
node.status = ProxyNodeStatus.ONLINE if connected else ProxyNodeStatus.OFFLINE
|
|
||||||
node.updated_at = now
|
|
||||||
|
|
||||||
event = ProxyNodeEvent(
|
|
||||||
node_id=node_id,
|
|
||||||
event_type="connected" if connected else "disconnected",
|
|
||||||
detail=f"[hub_node_status] conn_count={conn_count}",
|
|
||||||
)
|
|
||||||
db.add(event)
|
|
||||||
db.commit()
|
|
||||||
except Exception:
|
|
||||||
try:
|
|
||||||
db.rollback()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
raise
|
|
||||||
finally:
|
|
||||||
db.close()
|
|
||||||
|
|
||||||
try:
|
|
||||||
await asyncio.to_thread(_sync_update)
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("hub NODE_STATUS DB update failed: node_id={}, error={}", node_id, e)
|
|
||||||
|
|
||||||
async def send_request(
|
|
||||||
self,
|
|
||||||
node_id: str,
|
|
||||||
*,
|
|
||||||
method: str,
|
|
||||||
url: str,
|
|
||||||
headers: dict[str, str],
|
|
||||||
body: bytes | None = None,
|
|
||||||
timeout: float = 60.0,
|
|
||||||
) -> _StreamState:
|
|
||||||
await self.ensure_connected()
|
|
||||||
self._raise_if_degraded()
|
|
||||||
|
|
||||||
if len(self._pending_streams) >= self._config.max_streams:
|
|
||||||
raise TunnelStreamError(
|
|
||||||
f"hub stream limit reached ({self._config.max_streams}) for node {node_id}"
|
|
||||||
)
|
|
||||||
stream_id = self._alloc_stream_id()
|
|
||||||
stream_state = _StreamState(stream_id)
|
|
||||||
self._pending_streams[stream_id] = stream_state
|
|
||||||
|
|
||||||
try:
|
|
||||||
meta = json.dumps(
|
|
||||||
{
|
|
||||||
"node_id": node_id,
|
|
||||||
"method": method,
|
|
||||||
"url": url,
|
|
||||||
"headers": headers,
|
|
||||||
"timeout": int(timeout),
|
|
||||||
},
|
|
||||||
ensure_ascii=False,
|
|
||||||
separators=(",", ":"),
|
|
||||||
).encode("utf-8")
|
|
||||||
|
|
||||||
meta_payload, meta_flags = _compress_frame_payload(meta)
|
|
||||||
await self._send_frame(
|
|
||||||
Frame(stream_id, MsgType.REQUEST_HEADERS, meta_flags, meta_payload)
|
|
||||||
)
|
|
||||||
|
|
||||||
body_data = body or b""
|
|
||||||
if body_data:
|
|
||||||
body_payload, body_flags = await _compress_frame_payload_async(body_data)
|
|
||||||
else:
|
|
||||||
body_payload, body_flags = body_data, 0
|
|
||||||
body_flags |= FrameFlags.END_STREAM
|
|
||||||
await self._send_frame(Frame(stream_id, MsgType.REQUEST_BODY, body_flags, body_payload))
|
|
||||||
except Exception:
|
|
||||||
self._pending_streams.pop(stream_id, None)
|
|
||||||
raise
|
|
||||||
|
|
||||||
return stream_state
|
|
||||||
|
|
||||||
def remove_stream(self, stream_id: int) -> None:
|
|
||||||
self._pending_streams.pop(stream_id, None)
|
|
||||||
|
|
||||||
def _alloc_stream_id(self) -> int:
|
|
||||||
sid = self._next_stream_id
|
|
||||||
self._next_stream_id = sid + 2 if sid < 0xFFFF_FFFE else 2
|
|
||||||
return sid
|
|
||||||
|
|
||||||
async def shutdown(self) -> None:
|
|
||||||
self._closing = True
|
|
||||||
|
|
||||||
if self._reconnect_task is not None:
|
|
||||||
self._reconnect_task.cancel()
|
|
||||||
if self._reader_task is not None:
|
|
||||||
self._reader_task.cancel()
|
|
||||||
if self._ping_task is not None:
|
|
||||||
self._ping_task.cancel()
|
|
||||||
if self._watchdog_task is not None:
|
|
||||||
self._watchdog_task.cancel()
|
|
||||||
|
|
||||||
tasks = list(self._background_tasks)
|
|
||||||
for task in tasks:
|
|
||||||
task.cancel()
|
|
||||||
if tasks:
|
|
||||||
await asyncio.gather(*tasks, return_exceptions=True)
|
|
||||||
|
|
||||||
if self._ws is not None and not self._ws.closed:
|
|
||||||
try:
|
|
||||||
await self._ws.close()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
self._ws = None
|
|
||||||
|
|
||||||
if self._session is not None and not self._session.closed:
|
|
||||||
try:
|
|
||||||
await self._session.close()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
self._session = None
|
|
||||||
|
|
||||||
if self._pending_streams:
|
|
||||||
for state in self._pending_streams.values():
|
|
||||||
state.set_error("hub connection manager shutdown")
|
|
||||||
self._pending_streams.clear()
|
|
||||||
|
|
||||||
logger.info("Hub connection manager shutdown completed")
|
|
||||||
|
|
||||||
|
|
||||||
class HubTunnelTransport(httpx.AsyncBaseTransport):
|
class HubTunnelTransport(httpx.AsyncBaseTransport):
|
||||||
"""通过 aether-hub 转发请求的 httpx transport。"""
|
"""通过本机 aether-hub relay 转发请求的 httpx transport。"""
|
||||||
|
|
||||||
def __init__(self, node_id: str, timeout: float = 60.0) -> None:
|
def __init__(self, node_id: str, timeout: float = 60.0) -> None:
|
||||||
self._node_id = node_id
|
self._node_id = node_id
|
||||||
self._timeout = timeout
|
self._timeout = timeout
|
||||||
|
|
||||||
|
config = get_hub_config()
|
||||||
|
relay_timeout = max(timeout + 5.0, config.connect_timeout_seconds)
|
||||||
|
self._relay_client = httpx.AsyncClient(
|
||||||
|
transport=httpx.AsyncHTTPTransport(retries=0),
|
||||||
|
timeout=httpx.Timeout(
|
||||||
|
connect=config.connect_timeout_seconds,
|
||||||
|
read=relay_timeout,
|
||||||
|
write=relay_timeout,
|
||||||
|
pool=relay_timeout,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||||
manager = get_hub_connection_manager()
|
config = get_hub_config()
|
||||||
|
if not config.enabled:
|
||||||
|
raise httpx.ConnectError("hub local relay is unavailable outside docker runtime")
|
||||||
|
|
||||||
headers: dict[str, str] = {}
|
headers: dict[str, str] = {}
|
||||||
for key, value in request.headers.raw:
|
for key, value in request.headers.raw:
|
||||||
if key not in _HOP_BY_HOP_HEADERS_BYTES:
|
if key.lower() not in _HOP_BY_HOP_HEADERS_BYTES:
|
||||||
headers[key.decode("latin-1")] = value.decode("latin-1")
|
headers[key.decode("latin-1")] = value.decode("latin-1")
|
||||||
|
|
||||||
body = request.content or await request.aread() or None
|
body = request.content or await request.aread() or b""
|
||||||
|
envelope = _encode_relay_envelope(
|
||||||
|
{
|
||||||
|
"method": request.method,
|
||||||
|
"url": str(request.url),
|
||||||
|
"headers": headers,
|
||||||
|
"timeout": int(self._timeout),
|
||||||
|
},
|
||||||
|
body,
|
||||||
|
)
|
||||||
|
|
||||||
|
relay_request = self._relay_client.build_request(
|
||||||
|
"POST",
|
||||||
|
config.local_relay_url(self._node_id),
|
||||||
|
headers={"content-type": _RELAY_CONTENT_TYPE},
|
||||||
|
content=envelope,
|
||||||
|
)
|
||||||
|
|
||||||
stream_state: _StreamState | None = None
|
|
||||||
try:
|
try:
|
||||||
stream_state = await manager.send_request(
|
relay_response = await self._relay_client.send(relay_request, stream=True)
|
||||||
self._node_id,
|
except httpx.ConnectError as exc:
|
||||||
method=request.method,
|
raise httpx.ConnectError(f"hub local relay connect failed: {exc}") from exc
|
||||||
url=str(request.url),
|
except httpx.TimeoutException as exc:
|
||||||
headers=headers,
|
raise httpx.ConnectError(f"hub local relay timeout: {exc}") from exc
|
||||||
body=body,
|
|
||||||
timeout=self._timeout,
|
|
||||||
)
|
|
||||||
await stream_state.wait_headers(timeout=self._timeout)
|
|
||||||
|
|
||||||
return httpx.Response(
|
tunnel_error = relay_response.headers.get(_TUNNEL_ERROR_HEADER)
|
||||||
status_code=stream_state.status,
|
if tunnel_error:
|
||||||
headers=httpx.Headers(stream_state.headers),
|
message = await _read_error_message(relay_response)
|
||||||
stream=HubResponseStream(manager, stream_state, timeout=self._timeout),
|
if tunnel_error == "timeout":
|
||||||
)
|
raise httpx.ReadTimeout(message or "hub relay timed out")
|
||||||
except TunnelStreamError as e:
|
raise httpx.ConnectError(message or f"hub relay error: {tunnel_error}")
|
||||||
stream_id = stream_state.stream_id if stream_state else None
|
|
||||||
has_headers = bool(stream_state and stream_state.status > 0)
|
|
||||||
logger.warning(
|
|
||||||
"HubTunnelTransport error: node_id={}, url={}, stream_id={}, "
|
|
||||||
"has_headers={}, error={}",
|
|
||||||
self._node_id,
|
|
||||||
str(request.url),
|
|
||||||
stream_id,
|
|
||||||
has_headers,
|
|
||||||
e,
|
|
||||||
)
|
|
||||||
self._cleanup_stream(manager, stream_state)
|
|
||||||
if has_headers:
|
|
||||||
raise httpx.ReadError(str(e)) from e
|
|
||||||
raise httpx.ConnectError(str(e)) from e
|
|
||||||
except asyncio.TimeoutError:
|
|
||||||
stream_id = stream_state.stream_id if stream_state else None
|
|
||||||
logger.warning(
|
|
||||||
"HubTunnelTransport timeout: node_id={}, url={}, stream_id={}, timeout={:.0f}s",
|
|
||||||
self._node_id,
|
|
||||||
str(request.url),
|
|
||||||
stream_id,
|
|
||||||
self._timeout,
|
|
||||||
)
|
|
||||||
self._cleanup_stream(manager, stream_state)
|
|
||||||
raise httpx.ReadTimeout("hub tunnel request timeout") from None
|
|
||||||
except Exception:
|
|
||||||
self._cleanup_stream(manager, stream_state)
|
|
||||||
raise
|
|
||||||
|
|
||||||
def _cleanup_stream(
|
return httpx.Response(
|
||||||
self,
|
status_code=relay_response.status_code,
|
||||||
manager: HubConnectionManager,
|
headers=httpx.Headers(relay_response.headers),
|
||||||
stream_state: _StreamState | None,
|
stream=HubRelayResponseStream(relay_response),
|
||||||
) -> None:
|
request=request,
|
||||||
if stream_state is None:
|
)
|
||||||
return
|
|
||||||
manager.remove_stream(stream_state.stream_id)
|
async def aclose(self) -> None:
|
||||||
|
await self._relay_client.aclose()
|
||||||
|
|
||||||
|
|
||||||
class HubResponseStream(httpx.AsyncByteStream):
|
class HubRelayResponseStream(httpx.AsyncByteStream):
|
||||||
def __init__(
|
def __init__(self, response: httpx.Response) -> None:
|
||||||
self,
|
self._response = response
|
||||||
manager: HubConnectionManager,
|
|
||||||
stream_state: _StreamState,
|
|
||||||
timeout: float = 60.0,
|
|
||||||
) -> None:
|
|
||||||
self._manager = manager
|
|
||||||
self._stream_state = stream_state
|
|
||||||
self._timeout = timeout
|
|
||||||
|
|
||||||
async def __aiter__(self) -> AsyncGenerator[bytes, None]:
|
async def __aiter__(self) -> AsyncGenerator[bytes, None]:
|
||||||
try:
|
try:
|
||||||
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
|
async for chunk in self._response.aiter_raw():
|
||||||
yield chunk
|
yield chunk
|
||||||
finally:
|
finally:
|
||||||
self._manager.remove_stream(self._stream_state.stream_id)
|
await self._response.aclose()
|
||||||
|
|
||||||
async def aclose(self) -> None:
|
async def aclose(self) -> None:
|
||||||
self._manager.remove_stream(self._stream_state.stream_id)
|
await self._response.aclose()
|
||||||
|
|
||||||
|
|
||||||
def _compress_frame_payload(data: bytes) -> tuple[bytes, int]:
|
def _encode_relay_envelope(meta: dict[str, object], body: bytes) -> bytes:
|
||||||
if len(data) >= _TUNNEL_COMPRESS_MIN_SIZE:
|
meta_json = json.dumps(meta, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||||
compressed = gzip.compress(data, compresslevel=6)
|
return struct.pack("!I", len(meta_json)) + meta_json + body
|
||||||
if len(compressed) < len(data):
|
|
||||||
return compressed, FrameFlags.GZIP_COMPRESSED
|
|
||||||
return data, 0
|
|
||||||
|
|
||||||
|
|
||||||
async def _compress_frame_payload_async(data: bytes) -> tuple[bytes, int]:
|
async def _read_error_message(response: httpx.Response) -> str:
|
||||||
if len(data) < _TUNNEL_ASYNC_COMPRESS_THRESHOLD:
|
try:
|
||||||
return _compress_frame_payload(data)
|
payload = await response.aread()
|
||||||
return await asyncio.to_thread(_compress_frame_payload, data)
|
return payload.decode("utf-8", errors="replace").strip()
|
||||||
|
finally:
|
||||||
|
await response.aclose()
|
||||||
def _decompress_frame_payload(frame: Frame) -> bytes:
|
|
||||||
if frame.is_gzip:
|
|
||||||
return gzip.decompress(frame.payload)
|
|
||||||
return frame.payload
|
|
||||||
|
|
||||||
|
|
||||||
_hub_connection_manager: HubConnectionManager | None = None
|
|
||||||
|
|
||||||
|
|
||||||
def get_hub_connection_manager() -> HubConnectionManager:
|
|
||||||
global _hub_connection_manager
|
|
||||||
if _hub_connection_manager is None:
|
|
||||||
_hub_connection_manager = HubConnectionManager()
|
|
||||||
return _hub_connection_manager
|
|
||||||
|
|
||||||
|
|
||||||
async def shutdown_hub_connection_manager() -> None:
|
|
||||||
global _hub_connection_manager
|
|
||||||
if _hub_connection_manager is None:
|
|
||||||
return
|
|
||||||
await _hub_connection_manager.shutdown()
|
|
||||||
_hub_connection_manager = None
|
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ from src.core.logger import logger
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
_proxy_node_cache: dict[str, tuple[dict[str, Any] | None, float]] = {}
|
_proxy_node_cache: dict[str, tuple[dict[str, Any] | None, float]] = {}
|
||||||
_proxy_node_cache_lock = threading.Lock()
|
_proxy_node_cache_lock = threading.Lock()
|
||||||
_PROXY_NODE_CACHE_TTL_SECONDS = 15.0
|
_PROXY_NODE_CACHE_TTL_SECONDS = 3.0
|
||||||
_PROXY_NODE_CACHE_NEGATIVE_TTL_SECONDS = 5.0 # 不可用节点使用更短的 TTL,加速恢复感知
|
_PROXY_NODE_CACHE_NEGATIVE_TTL_SECONDS = 5.0 # 不可用节点使用更短的 TTL,加速恢复感知
|
||||||
_PROXY_NODE_CACHE_MAX_SIZE = 256
|
_PROXY_NODE_CACHE_MAX_SIZE = 256
|
||||||
|
|
||||||
|
|||||||
@@ -18,7 +18,14 @@ from sqlalchemy import func, update
|
|||||||
from sqlalchemy.orm import Session
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||||
from src.models.database import Provider, ProviderEndpoint, ProxyNode, ProxyNodeStatus, SystemConfig
|
from src.models.database import (
|
||||||
|
Provider,
|
||||||
|
ProviderEndpoint,
|
||||||
|
ProxyNode,
|
||||||
|
ProxyNodeEvent,
|
||||||
|
ProxyNodeStatus,
|
||||||
|
SystemConfig,
|
||||||
|
)
|
||||||
|
|
||||||
from .resolver import (
|
from .resolver import (
|
||||||
inject_auth_into_proxy_url,
|
inject_auth_into_proxy_url,
|
||||||
@@ -117,7 +124,7 @@ def _normalize_proxy_metadata(
|
|||||||
|
|
||||||
|
|
||||||
def build_heartbeat_ack(node: ProxyNode) -> dict[str, Any]:
|
def build_heartbeat_ack(node: ProxyNode) -> dict[str, Any]:
|
||||||
"""从心跳后的节点构建 ACK 响应 payload(供 hub_transport / tunnel_manager 使用)。"""
|
"""从心跳后的节点构建 ACK 响应 payload(供 hub 控制面回调使用)。"""
|
||||||
result: dict[str, Any] = {}
|
result: dict[str, Any] = {}
|
||||||
if not node.remote_config:
|
if not node.remote_config:
|
||||||
return result
|
return result
|
||||||
@@ -423,6 +430,60 @@ class ProxyNodeService:
|
|||||||
raise NotFoundException(f"ProxyNode {node_id} 不存在", "proxy_node")
|
raise NotFoundException(f"ProxyNode {node_id} 不存在", "proxy_node")
|
||||||
return refreshed
|
return refreshed
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def update_tunnel_status(
|
||||||
|
db: Session,
|
||||||
|
*,
|
||||||
|
node_id: str,
|
||||||
|
connected: bool,
|
||||||
|
conn_count: int = 0,
|
||||||
|
detail: str | None = None,
|
||||||
|
observed_at: datetime | None = None,
|
||||||
|
) -> ProxyNode | None:
|
||||||
|
"""根据 Hub 连接池状态更新 tunnel 连接状态并记录事件。"""
|
||||||
|
node = db.query(ProxyNode).filter(ProxyNode.id == node_id).first()
|
||||||
|
if not node:
|
||||||
|
return None
|
||||||
|
|
||||||
|
event_time = observed_at or datetime.now(timezone.utc)
|
||||||
|
last_transition = node.tunnel_connected_at
|
||||||
|
if last_transition and last_transition.tzinfo is None:
|
||||||
|
last_transition = last_transition.replace(tzinfo=timezone.utc)
|
||||||
|
|
||||||
|
event_type = "connected" if connected else "disconnected"
|
||||||
|
event_detail = detail or f"[hub_node_status] conn_count={max(int(conn_count), 0)}"
|
||||||
|
|
||||||
|
if last_transition and event_time < last_transition:
|
||||||
|
db.add(
|
||||||
|
ProxyNodeEvent(
|
||||||
|
node_id=node_id,
|
||||||
|
event_type=event_type,
|
||||||
|
detail=f"[stale_ignored] {event_detail}",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
db.commit()
|
||||||
|
return node
|
||||||
|
|
||||||
|
node.tunnel_connected = connected
|
||||||
|
node.tunnel_connected_at = event_time
|
||||||
|
node.status = ProxyNodeStatus.ONLINE if connected else ProxyNodeStatus.OFFLINE
|
||||||
|
node.updated_at = event_time
|
||||||
|
|
||||||
|
db.add(
|
||||||
|
ProxyNodeEvent(
|
||||||
|
node_id=node_id,
|
||||||
|
event_type=event_type,
|
||||||
|
detail=event_detail,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
db.commit()
|
||||||
|
db.refresh(node)
|
||||||
|
|
||||||
|
from .resolver import invalidate_proxy_node_cache
|
||||||
|
|
||||||
|
invalidate_proxy_node_cache(node_id)
|
||||||
|
return node
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def unregister_node(db: Session, *, node_id: str) -> ProxyNode:
|
def unregister_node(db: Session, *, node_id: str) -> ProxyNode:
|
||||||
"""注销节点(设置为 OFFLINE)"""
|
"""注销节点(设置为 OFFLINE)"""
|
||||||
|
|||||||
@@ -1,606 +0,0 @@
|
|||||||
"""
|
|
||||||
WebSocket 隧道管理器
|
|
||||||
|
|
||||||
管理所有活跃的 aether-proxy tunnel 连接,提供通过隧道发送 HTTP 请求的能力。
|
|
||||||
每个 proxy node 可持有多条 tunnel 连接(连接池),请求按 least-loaded 策略分配。
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import gzip
|
|
||||||
import json
|
|
||||||
import time
|
|
||||||
from typing import TYPE_CHECKING, Any
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from collections.abc import AsyncGenerator
|
|
||||||
|
|
||||||
from starlette.websockets import WebSocket, WebSocketState
|
|
||||||
|
|
||||||
from src.core.logger import logger
|
|
||||||
|
|
||||||
from .tunnel_protocol import Frame, FrameFlags, MsgType, normalize_heartbeat_id
|
|
||||||
|
|
||||||
# 隧道帧压缩的最小 payload 大小(字节)
|
|
||||||
# 小于此值的帧压缩收益不大,反而增加 CPU 开销
|
|
||||||
_TUNNEL_COMPRESS_MIN_SIZE = 512
|
|
||||||
|
|
||||||
|
|
||||||
class TunnelConnection:
|
|
||||||
"""单条 tunnel 连接"""
|
|
||||||
|
|
||||||
__slots__ = (
|
|
||||||
"node_id",
|
|
||||||
"node_name",
|
|
||||||
"ws",
|
|
||||||
"connected_at",
|
|
||||||
"max_streams",
|
|
||||||
"_pending_streams",
|
|
||||||
"_write_lock",
|
|
||||||
"_next_stream_id",
|
|
||||||
)
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
node_id: str,
|
|
||||||
node_name: str,
|
|
||||||
ws: WebSocket,
|
|
||||||
max_streams: int | None = None,
|
|
||||||
) -> None:
|
|
||||||
self.node_id = node_id
|
|
||||||
self.node_name = node_name
|
|
||||||
self.ws = ws
|
|
||||||
self.connected_at = time.time()
|
|
||||||
# Per-connection max concurrent streams: use proxy-advertised value
|
|
||||||
# (from X-Tunnel-Max-Streams header), clamped to [64, 2048].
|
|
||||||
# Falls back to TunnelManager.MAX_STREAMS_PER_CONN if not provided.
|
|
||||||
if max_streams is not None:
|
|
||||||
self.max_streams = max(64, min(max_streams, 2048))
|
|
||||||
else:
|
|
||||||
self.max_streams = TunnelManager.MAX_STREAMS_PER_CONN
|
|
||||||
self._pending_streams: dict[int, _StreamState] = {}
|
|
||||||
self._write_lock = asyncio.Lock()
|
|
||||||
# Per-connection stream ID 分配器(Aether 端使用偶数,从 2 开始)
|
|
||||||
self._next_stream_id: int = 2
|
|
||||||
|
|
||||||
@property
|
|
||||||
def is_alive(self) -> bool:
|
|
||||||
return self.ws.client_state == WebSocketState.CONNECTED
|
|
||||||
|
|
||||||
async def send_frame(self, frame: Frame, timeout: float = 10.0) -> None:
|
|
||||||
"""发送帧到 WebSocket,带超时保护防止写阻塞。
|
|
||||||
|
|
||||||
在高丢包网络下 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:
|
|
||||||
state = _StreamState(stream_id, conn=self)
|
|
||||||
self._pending_streams[stream_id] = state
|
|
||||||
return state
|
|
||||||
|
|
||||||
def get_stream(self, stream_id: int) -> _StreamState | None:
|
|
||||||
return self._pending_streams.get(stream_id)
|
|
||||||
|
|
||||||
def remove_stream(self, stream_id: int) -> None:
|
|
||||||
self._pending_streams.pop(stream_id, None)
|
|
||||||
|
|
||||||
@property
|
|
||||||
def stream_count(self) -> int:
|
|
||||||
return len(self._pending_streams)
|
|
||||||
|
|
||||||
def has_stream(self, stream_id: int) -> bool:
|
|
||||||
return stream_id in self._pending_streams
|
|
||||||
|
|
||||||
def alloc_stream_id(self, max_streams: int) -> int:
|
|
||||||
"""分配一个未被占用的偶数 stream_id,回绕时跳过飞行中的 ID"""
|
|
||||||
# 最多尝试 max_streams + 16 次(飞行中的 stream 数量不超过 max_streams)
|
|
||||||
for _ in range(max_streams + 16):
|
|
||||||
sid = self._next_stream_id
|
|
||||||
self._next_stream_id += 2
|
|
||||||
if self._next_stream_id > 0xFFFF_FFFE:
|
|
||||||
self._next_stream_id = 2
|
|
||||||
if sid not in self._pending_streams:
|
|
||||||
return sid
|
|
||||||
raise TunnelStreamError("stream ID space exhausted")
|
|
||||||
|
|
||||||
def cancel_all_streams(self) -> None:
|
|
||||||
if self._pending_streams:
|
|
||||||
logger.warning(
|
|
||||||
"tunnel cancel_all_streams: node_id={}, name={}, count={}",
|
|
||||||
self.node_id,
|
|
||||||
self.node_name,
|
|
||||||
len(self._pending_streams),
|
|
||||||
)
|
|
||||||
for state in self._pending_streams.values():
|
|
||||||
state.set_error("tunnel disconnected")
|
|
||||||
self._pending_streams.clear()
|
|
||||||
|
|
||||||
|
|
||||||
class _StreamState:
|
|
||||||
"""跟踪单个 stream 的响应状态"""
|
|
||||||
|
|
||||||
__slots__ = (
|
|
||||||
"stream_id",
|
|
||||||
"status",
|
|
||||||
"headers",
|
|
||||||
"_header_event",
|
|
||||||
"_body_chunks",
|
|
||||||
"_done_event",
|
|
||||||
"_error",
|
|
||||||
"_conn",
|
|
||||||
)
|
|
||||||
|
|
||||||
def __init__(self, stream_id: int, conn: TunnelConnection | None = None) -> None:
|
|
||||||
self.stream_id = stream_id
|
|
||||||
self.status: int = 0
|
|
||||||
self.headers: list[list[str]] = []
|
|
||||||
self._header_event = asyncio.Event()
|
|
||||||
self._body_chunks: asyncio.Queue[bytes | None] = asyncio.Queue()
|
|
||||||
self._done_event = asyncio.Event()
|
|
||||||
self._error: str | None = None
|
|
||||||
self._conn = conn
|
|
||||||
|
|
||||||
def set_response_headers(self, status: int, headers: list[list[str]] | dict[str, str]) -> None:
|
|
||||||
self.status = status
|
|
||||||
# headers 可能是 [[k, v], ...] (多值) 或 {k: v} (旧格式兼容)
|
|
||||||
if isinstance(headers, list):
|
|
||||||
self.headers = headers # type: ignore[assignment]
|
|
||||||
else:
|
|
||||||
self.headers = list(headers.items()) # type: ignore[assignment]
|
|
||||||
self._header_event.set()
|
|
||||||
|
|
||||||
def push_body_chunk(self, data: bytes) -> None:
|
|
||||||
self._body_chunks.put_nowait(data)
|
|
||||||
|
|
||||||
def set_done(self) -> None:
|
|
||||||
self._body_chunks.put_nowait(None) # sentinel
|
|
||||||
self._done_event.set()
|
|
||||||
|
|
||||||
def set_error(self, msg: str) -> None:
|
|
||||||
self._error = msg
|
|
||||||
self._header_event.set()
|
|
||||||
self._body_chunks.put_nowait(None)
|
|
||||||
self._done_event.set()
|
|
||||||
|
|
||||||
async def wait_headers(self, timeout: float = 60.0) -> None:
|
|
||||||
await asyncio.wait_for(self._header_event.wait(), timeout=timeout)
|
|
||||||
if self._error:
|
|
||||||
raise TunnelStreamError(self._error)
|
|
||||||
|
|
||||||
async def iter_body(self, chunk_timeout: float = 60.0) -> AsyncGenerator[bytes, None]:
|
|
||||||
chunks_received = 0
|
|
||||||
total_bytes = 0
|
|
||||||
while True:
|
|
||||||
try:
|
|
||||||
chunk = await asyncio.wait_for(self._body_chunks.get(), timeout=chunk_timeout)
|
|
||||||
except asyncio.TimeoutError:
|
|
||||||
self._error = "body chunk timeout"
|
|
||||||
self._done_event.set()
|
|
||||||
logger.warning(
|
|
||||||
"tunnel stream body chunk timeout: stream_id={}, "
|
|
||||||
"chunk_timeout={:.0f}s, chunks_received={}, total_bytes={}",
|
|
||||||
self.stream_id,
|
|
||||||
chunk_timeout,
|
|
||||||
chunks_received,
|
|
||||||
total_bytes,
|
|
||||||
)
|
|
||||||
raise TunnelStreamError("body chunk timeout")
|
|
||||||
if chunk is None:
|
|
||||||
if self._error:
|
|
||||||
logger.warning(
|
|
||||||
"tunnel stream ended with error: stream_id={}, error={}, "
|
|
||||||
"chunks_received={}, total_bytes={}",
|
|
||||||
self.stream_id,
|
|
||||||
self._error,
|
|
||||||
chunks_received,
|
|
||||||
total_bytes,
|
|
||||||
)
|
|
||||||
raise TunnelStreamError(self._error)
|
|
||||||
return
|
|
||||||
chunks_received += 1
|
|
||||||
total_bytes += len(chunk)
|
|
||||||
yield chunk
|
|
||||||
|
|
||||||
|
|
||||||
class TunnelStreamError(Exception):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 全局 TunnelManager 单例
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
class TunnelManager:
|
|
||||||
"""管理所有活跃的 tunnel 连接(支持每个 node 多条连接的连接池)"""
|
|
||||||
|
|
||||||
# 单条 tunnel 上允许的最大并发 stream 数(超出时拒绝新请求)
|
|
||||||
MAX_STREAMS_PER_CONN = 2048
|
|
||||||
|
|
||||||
def __init__(self) -> None:
|
|
||||||
self._connections: dict[str, list[TunnelConnection]] = {} # node_id -> [conn, ...]
|
|
||||||
self._background_tasks: set[asyncio.Task[None]] = set()
|
|
||||||
self._draining: bool = False
|
|
||||||
# 每个 node 的连续 reconnect 计数,用于抑制高频 connect/disconnect 日志
|
|
||||||
self._reconnect_counts: dict[str, int] = {}
|
|
||||||
|
|
||||||
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
|
|
||||||
def active_count(self) -> int:
|
|
||||||
return sum(len(conns) for conns in self._connections.values())
|
|
||||||
|
|
||||||
def get_connection(self, node_id: str) -> TunnelConnection | None:
|
|
||||||
"""获取负载最低的存活连接,同时清理 dead 连接"""
|
|
||||||
conns = self._connections.get(node_id)
|
|
||||||
if not conns:
|
|
||||||
return None
|
|
||||||
|
|
||||||
# 清理 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:
|
|
||||||
"""注册一条新连接到连接池"""
|
|
||||||
conns = self._connections.get(conn.node_id)
|
|
||||||
if conns is None:
|
|
||||||
conns = []
|
|
||||||
self._connections[conn.node_id] = conns
|
|
||||||
conns.append(conn)
|
|
||||||
reconn = self._reconnect_counts.get(conn.node_id, 0)
|
|
||||||
if reconn == 0:
|
|
||||||
logger.info(
|
|
||||||
"tunnel connected: node_id={}, name={}, pool_size={}",
|
|
||||||
conn.node_id,
|
|
||||||
conn.node_name,
|
|
||||||
len(conns),
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logger.debug(
|
|
||||||
"tunnel reconnected: node_id={}, name={}, pool_size={}, reconnect_count={}",
|
|
||||||
conn.node_id,
|
|
||||||
conn.node_name,
|
|
||||||
len(conns),
|
|
||||||
reconn,
|
|
||||||
)
|
|
||||||
|
|
||||||
def unregister(self, conn: TunnelConnection) -> bool:
|
|
||||||
"""
|
|
||||||
从连接池中注销指定连接。
|
|
||||||
|
|
||||||
返回 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
|
|
||||||
reconn = self._reconnect_counts.get(conn.node_id, 0) + 1
|
|
||||||
self._reconnect_counts[conn.node_id] = reconn
|
|
||||||
# 首次断开 info,后续每 10 次 info 汇总,其余 debug
|
|
||||||
if reconn == 1:
|
|
||||||
logger.info(
|
|
||||||
"tunnel disconnected: node_id={}, name={}, remaining={}",
|
|
||||||
conn.node_id,
|
|
||||||
conn.node_name,
|
|
||||||
remaining,
|
|
||||||
)
|
|
||||||
elif reconn % 10 == 0:
|
|
||||||
logger.info(
|
|
||||||
"tunnel disconnected: node_id={}, name={}, remaining={} (repeated {} times)",
|
|
||||||
conn.node_id,
|
|
||||||
conn.node_name,
|
|
||||||
remaining,
|
|
||||||
reconn,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
logger.debug(
|
|
||||||
"tunnel disconnected: node_id={}, name={}, remaining={}",
|
|
||||||
conn.node_id,
|
|
||||||
conn.node_name,
|
|
||||||
remaining,
|
|
||||||
)
|
|
||||||
return True
|
|
||||||
|
|
||||||
def reset_reconnect_count(self, node_id: str) -> None:
|
|
||||||
"""重置 reconnect 计数(当连接恢复正常活动时调用)"""
|
|
||||||
self._reconnect_counts.pop(node_id, None)
|
|
||||||
|
|
||||||
async def shutdown_all(self, drain_timeout: float = 60.0) -> None:
|
|
||||||
"""优雅关闭所有 tunnel 连接:drain 飞行中请求 -> GoAway -> 关闭 WebSocket。
|
|
||||||
|
|
||||||
在 worker 即将退出时调用。先标记 draining 阻止新请求进入,
|
|
||||||
等待飞行中的 stream 完成(最多 drain_timeout 秒),
|
|
||||||
然后发送 GoAway 让 proxy 端重连到其他 worker。
|
|
||||||
"""
|
|
||||||
all_conns = [c for conns in self._connections.values() for c in conns]
|
|
||||||
if not all_conns:
|
|
||||||
return
|
|
||||||
|
|
||||||
# 标记 draining,send_request 将拒绝新请求
|
|
||||||
self._draining = True
|
|
||||||
|
|
||||||
total_streams = sum(c.stream_count for c in all_conns)
|
|
||||||
if total_streams > 0:
|
|
||||||
logger.info(
|
|
||||||
"draining {} in-flight streams on {} connections (timeout={}s)",
|
|
||||||
total_streams,
|
|
||||||
len(all_conns),
|
|
||||||
drain_timeout,
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
await asyncio.wait_for(self._wait_streams_drain(all_conns), timeout=drain_timeout)
|
|
||||||
except asyncio.TimeoutError:
|
|
||||||
remaining = sum(c.stream_count for c in all_conns)
|
|
||||||
logger.warning("drain timeout, {} streams still in-flight", remaining)
|
|
||||||
|
|
||||||
logger.info("sending GoAway to {} tunnel connections", len(all_conns))
|
|
||||||
|
|
||||||
async def _close_conn(conn: TunnelConnection) -> None:
|
|
||||||
try:
|
|
||||||
await asyncio.wait_for(
|
|
||||||
conn.send_frame(Frame(0, MsgType.GOAWAY, 0, b"")),
|
|
||||||
timeout=2.0,
|
|
||||||
)
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
try:
|
|
||||||
await conn.ws.close(code=1001, reason="server shutting down")
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
|
|
||||||
await asyncio.gather(*(_close_conn(c) for c in all_conns), return_exceptions=True)
|
|
||||||
|
|
||||||
async def _wait_streams_drain(self, conns: list[TunnelConnection]) -> None:
|
|
||||||
"""轮询等待所有连接的 pending_streams 清空"""
|
|
||||||
while any(c.stream_count > 0 for c in conns):
|
|
||||||
await asyncio.sleep(0.5)
|
|
||||||
|
|
||||||
def has_tunnel(self, node_id: str) -> bool:
|
|
||||||
"""检查指定 node 是否有存活的 tunnel 连接(纯检查,无副作用)
|
|
||||||
|
|
||||||
与 get_connection 不同,此方法不会清理 dead 连接,
|
|
||||||
避免在 finally 块或 health_scheduler 中误清理刚注册的连接。
|
|
||||||
"""
|
|
||||||
conns = self._connections.get(node_id)
|
|
||||||
if not conns:
|
|
||||||
return False
|
|
||||||
return any(c.is_alive for c in conns)
|
|
||||||
|
|
||||||
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(
|
|
||||||
self,
|
|
||||||
node_id: str,
|
|
||||||
*,
|
|
||||||
method: str,
|
|
||||||
url: str,
|
|
||||||
headers: dict[str, str],
|
|
||||||
body: bytes | None = None,
|
|
||||||
timeout: float = 60.0,
|
|
||||||
) -> _StreamState:
|
|
||||||
"""
|
|
||||||
通过 tunnel 发送 HTTP 请求,返回 StreamState 用于读取响应。
|
|
||||||
"""
|
|
||||||
if self._draining:
|
|
||||||
raise TunnelStreamError("tunnel manager is draining, rejecting new requests")
|
|
||||||
|
|
||||||
conn = self.get_connection(node_id)
|
|
||||||
if not conn:
|
|
||||||
raise TunnelStreamError(f"tunnel not connected for node {node_id}")
|
|
||||||
|
|
||||||
if conn.stream_count >= conn.max_streams:
|
|
||||||
raise TunnelStreamError(
|
|
||||||
f"tunnel stream limit reached ({conn.max_streams}) for node {node_id}"
|
|
||||||
)
|
|
||||||
|
|
||||||
stream_id = conn.alloc_stream_id(conn.max_streams)
|
|
||||||
stream_state = conn.create_stream(stream_id)
|
|
||||||
|
|
||||||
try:
|
|
||||||
# 发送 REQUEST_HEADERS(大元数据帧压缩)
|
|
||||||
meta = json.dumps(
|
|
||||||
{
|
|
||||||
"method": method,
|
|
||||||
"url": url,
|
|
||||||
"headers": headers,
|
|
||||||
"timeout": int(timeout),
|
|
||||||
}
|
|
||||||
).encode()
|
|
||||||
meta_payload, meta_flags = _compress_frame_payload(meta)
|
|
||||||
await conn.send_frame(
|
|
||||||
Frame(stream_id, MsgType.REQUEST_HEADERS, meta_flags, meta_payload)
|
|
||||||
)
|
|
||||||
|
|
||||||
# 发送 REQUEST_BODY + END_STREAM(大请求体帧压缩)
|
|
||||||
body_data = body or b""
|
|
||||||
if body_data:
|
|
||||||
body_payload, body_flags = _compress_frame_payload(body_data)
|
|
||||||
else:
|
|
||||||
body_payload, body_flags = body_data, 0
|
|
||||||
body_flags |= FrameFlags.END_STREAM
|
|
||||||
await conn.send_frame(Frame(stream_id, MsgType.REQUEST_BODY, body_flags, body_payload))
|
|
||||||
except Exception:
|
|
||||||
conn.remove_stream(stream_id)
|
|
||||||
raise
|
|
||||||
|
|
||||||
return stream_state
|
|
||||||
|
|
||||||
async def handle_incoming_frame(self, conn: TunnelConnection, frame: Frame) -> None:
|
|
||||||
"""处理从 proxy 收到的响应帧(仅处理当前 active 连接的帧)。
|
|
||||||
|
|
||||||
重要:此方法在 WebSocket 主读循环中被 await 调用,不能长时间阻塞,
|
|
||||||
否则会阻止读取后续帧,导致 proxy 端 TCP 缓冲区满而级联失败。
|
|
||||||
"""
|
|
||||||
# 防止已被移除的连接的帧继续被处理
|
|
||||||
conns = self._connections.get(conn.node_id)
|
|
||||||
if not conns or conn not in conns:
|
|
||||||
return
|
|
||||||
|
|
||||||
stream = conn.get_stream(frame.stream_id)
|
|
||||||
|
|
||||||
if frame.msg_type == MsgType.RESPONSE_HEADERS:
|
|
||||||
if not stream:
|
|
||||||
return
|
|
||||||
try:
|
|
||||||
payload = _decompress_frame_payload(frame)
|
|
||||||
meta = json.loads(payload)
|
|
||||||
stream.set_response_headers(meta["status"], meta.get("headers", []))
|
|
||||||
except Exception as e:
|
|
||||||
stream.set_error(f"invalid response headers: {e}")
|
|
||||||
|
|
||||||
elif frame.msg_type == MsgType.RESPONSE_BODY:
|
|
||||||
if stream:
|
|
||||||
payload = _decompress_frame_payload(frame)
|
|
||||||
stream.push_body_chunk(payload)
|
|
||||||
|
|
||||||
elif frame.msg_type == MsgType.STREAM_END:
|
|
||||||
if stream:
|
|
||||||
stream.set_done()
|
|
||||||
conn.remove_stream(frame.stream_id)
|
|
||||||
|
|
||||||
elif frame.msg_type == MsgType.STREAM_ERROR:
|
|
||||||
if stream:
|
|
||||||
msg = frame.payload.decode(errors="replace") if frame.payload else "stream error"
|
|
||||||
logger.warning(
|
|
||||||
"tunnel received STREAM_ERROR: node={}, stream_id={}, message={}",
|
|
||||||
conn.node_name,
|
|
||||||
frame.stream_id,
|
|
||||||
msg[:500],
|
|
||||||
)
|
|
||||||
stream.set_error(msg)
|
|
||||||
conn.remove_stream(frame.stream_id)
|
|
||||||
|
|
||||||
elif frame.msg_type == MsgType.HEARTBEAT_DATA:
|
|
||||||
# fire-and-forget: 不阻塞主读循环
|
|
||||||
self._background(self._handle_heartbeat(conn, frame))
|
|
||||||
|
|
||||||
elif frame.msg_type == MsgType.PING:
|
|
||||||
# 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:
|
|
||||||
"""处理 proxy 上报的心跳数据,更新 DB,返回 ACK"""
|
|
||||||
try:
|
|
||||||
data = json.loads(frame.payload) if frame.payload else {}
|
|
||||||
except Exception:
|
|
||||||
data = {}
|
|
||||||
heartbeat_id = normalize_heartbeat_id(data.get("heartbeat_id"))
|
|
||||||
ack: dict[str, Any] = {}
|
|
||||||
if heartbeat_id is not None:
|
|
||||||
ack["heartbeat_id"] = heartbeat_id
|
|
||||||
|
|
||||||
def _sync_heartbeat() -> dict[str, Any]:
|
|
||||||
from src.database import create_session
|
|
||||||
from src.services.proxy_node.service import ProxyNodeService, build_heartbeat_ack
|
|
||||||
|
|
||||||
db = create_session()
|
|
||||||
try:
|
|
||||||
node = ProxyNodeService.heartbeat(
|
|
||||||
db,
|
|
||||||
node_id=conn.node_id,
|
|
||||||
active_connections=data.get("active_connections"),
|
|
||||||
total_requests=data.get("total_requests"),
|
|
||||||
avg_latency_ms=data.get("avg_latency_ms"),
|
|
||||||
failed_requests=data.get("failed_requests"),
|
|
||||||
dns_failures=data.get("dns_failures"),
|
|
||||||
stream_errors=data.get("stream_errors"),
|
|
||||||
proxy_metadata=data.get("proxy_metadata"),
|
|
||||||
proxy_version=data.get("proxy_version"),
|
|
||||||
)
|
|
||||||
return build_heartbeat_ack(node)
|
|
||||||
finally:
|
|
||||||
db.close()
|
|
||||||
|
|
||||||
try:
|
|
||||||
ack.update(await asyncio.to_thread(_sync_heartbeat))
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("tunnel heartbeat DB update failed: {}", e)
|
|
||||||
|
|
||||||
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)
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# 隧道帧压缩 / 解压
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def _compress_frame_payload(data: bytes) -> tuple[bytes, int]:
|
|
||||||
"""按配置对帧 payload 进行 gzip 压缩。
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
(payload, flags) — 若压缩则 flags 含 GZIP_COMPRESSED,否则 flags=0。
|
|
||||||
"""
|
|
||||||
if len(data) >= _TUNNEL_COMPRESS_MIN_SIZE:
|
|
||||||
compressed = gzip.compress(data, compresslevel=6)
|
|
||||||
# 仅在压缩确实缩小时使用
|
|
||||||
if len(compressed) < len(data):
|
|
||||||
return compressed, FrameFlags.GZIP_COMPRESSED
|
|
||||||
return data, 0
|
|
||||||
|
|
||||||
|
|
||||||
def _decompress_frame_payload(frame: Frame) -> bytes:
|
|
||||||
"""如果帧设置了 GZIP_COMPRESSED 标志则解压,否则原样返回。"""
|
|
||||||
if frame.is_gzip:
|
|
||||||
return gzip.decompress(frame.payload)
|
|
||||||
return frame.payload
|
|
||||||
|
|
||||||
|
|
||||||
# 全局单例
|
|
||||||
_tunnel_manager: TunnelManager | None = None
|
|
||||||
|
|
||||||
|
|
||||||
def get_tunnel_manager() -> TunnelManager:
|
|
||||||
global _tunnel_manager
|
|
||||||
if _tunnel_manager is None:
|
|
||||||
_tunnel_manager = TunnelManager()
|
|
||||||
return _tunnel_manager
|
|
||||||
323
tests/e2e_hub_relay.py
Normal file
323
tests/e2e_hub_relay.py
Normal file
@@ -0,0 +1,323 @@
|
|||||||
|
"""
|
||||||
|
aether-hub local relay 端到端测试
|
||||||
|
|
||||||
|
测试流程:
|
||||||
|
1. 启动 aether-hub(绑定随机端口)
|
||||||
|
2. 用 websockets 库模拟一个 aether-proxy client 连接到 Hub
|
||||||
|
3. Mock proxy 在收到请求帧后返回固定响应帧
|
||||||
|
4. 通过 Hub 的 /local/relay/{node_id} HTTP API 发送请求
|
||||||
|
5. 验证完整链路: HTTP request -> Hub -> WS frame -> mock proxy -> WS frame -> Hub -> HTTP response
|
||||||
|
|
||||||
|
运行: uv run python tests/e2e_hub_relay.py
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import gzip
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import signal
|
||||||
|
import struct
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Protocol constants (mirror aether-hub/src/protocol.rs)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
HEADER_SIZE = 10
|
||||||
|
|
||||||
|
REQUEST_HEADERS = 0x01
|
||||||
|
REQUEST_BODY = 0x02
|
||||||
|
RESPONSE_HEADERS = 0x03
|
||||||
|
RESPONSE_BODY = 0x04
|
||||||
|
STREAM_END = 0x05
|
||||||
|
STREAM_ERROR = 0x06
|
||||||
|
PING = 0x10
|
||||||
|
PONG = 0x11
|
||||||
|
GOAWAY = 0x12
|
||||||
|
|
||||||
|
FLAG_END_STREAM = 0x01
|
||||||
|
FLAG_GZIP_COMPRESSED = 0x02
|
||||||
|
|
||||||
|
|
||||||
|
def encode_frame(stream_id: int, msg_type: int, flags: int, payload: bytes) -> bytes:
|
||||||
|
header = struct.pack(">I", stream_id) + bytes([msg_type, flags]) + struct.pack(">I", len(payload))
|
||||||
|
return header + payload
|
||||||
|
|
||||||
|
|
||||||
|
def parse_frame(data: bytes) -> tuple[int, int, int, bytes] | None:
|
||||||
|
if len(data) < HEADER_SIZE:
|
||||||
|
return None
|
||||||
|
stream_id = struct.unpack(">I", data[0:4])[0]
|
||||||
|
msg_type = data[4]
|
||||||
|
flags = data[5]
|
||||||
|
payload_len = struct.unpack(">I", data[6:10])[0]
|
||||||
|
if len(data) < HEADER_SIZE + payload_len:
|
||||||
|
return None
|
||||||
|
payload = data[HEADER_SIZE : HEADER_SIZE + payload_len]
|
||||||
|
if flags & FLAG_GZIP_COMPRESSED:
|
||||||
|
payload = gzip.decompress(payload)
|
||||||
|
return stream_id, msg_type, flags, payload
|
||||||
|
|
||||||
|
|
||||||
|
def encode_relay_envelope(meta: dict, body: bytes) -> bytes:
|
||||||
|
meta_json = json.dumps(meta, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||||
|
return struct.pack("!I", len(meta_json)) + meta_json + body
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Mock aether-proxy: connects to Hub via WebSocket, handles request frames
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
async def mock_proxy(hub_ws_url: str, node_id: str, ready_event: asyncio.Event) -> None:
|
||||||
|
"""Simulate an aether-proxy node that echoes requests as fixed responses."""
|
||||||
|
try:
|
||||||
|
import websockets
|
||||||
|
except ImportError:
|
||||||
|
print("SKIP: websockets package not installed (uv pip install websockets)")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
headers = {
|
||||||
|
"X-Node-ID": node_id,
|
||||||
|
"X-Node-Name": f"test-{node_id}",
|
||||||
|
}
|
||||||
|
|
||||||
|
async with websockets.connect(
|
||||||
|
hub_ws_url,
|
||||||
|
additional_headers=headers,
|
||||||
|
max_size=64 * 1024 * 1024,
|
||||||
|
) as ws:
|
||||||
|
ready_event.set()
|
||||||
|
print(f" [mock-proxy] connected to hub as node_id={node_id}")
|
||||||
|
|
||||||
|
request_meta: dict | None = None
|
||||||
|
request_body: bytes = b""
|
||||||
|
|
||||||
|
async for raw_msg in ws:
|
||||||
|
if not isinstance(raw_msg, bytes):
|
||||||
|
continue
|
||||||
|
|
||||||
|
parsed = parse_frame(raw_msg)
|
||||||
|
if parsed is None:
|
||||||
|
continue
|
||||||
|
|
||||||
|
stream_id, msg_type, flags, payload = parsed
|
||||||
|
|
||||||
|
if msg_type == PING:
|
||||||
|
await ws.send(encode_frame(0, PONG, 0, payload))
|
||||||
|
continue
|
||||||
|
|
||||||
|
if msg_type == REQUEST_HEADERS:
|
||||||
|
request_meta = json.loads(payload)
|
||||||
|
print(f" [mock-proxy] stream={stream_id} got REQUEST_HEADERS: {request_meta.get('method')} {request_meta.get('url')}")
|
||||||
|
|
||||||
|
elif msg_type == REQUEST_BODY:
|
||||||
|
request_body = payload
|
||||||
|
is_end = bool(flags & FLAG_END_STREAM)
|
||||||
|
print(f" [mock-proxy] stream={stream_id} got REQUEST_BODY ({len(payload)} bytes, end={is_end})")
|
||||||
|
|
||||||
|
if is_end and request_meta:
|
||||||
|
# Send response: 200 OK with echoed body
|
||||||
|
resp_meta = {
|
||||||
|
"status": 200,
|
||||||
|
"headers": [
|
||||||
|
["content-type", "application/json"],
|
||||||
|
["x-test-echo", "true"],
|
||||||
|
],
|
||||||
|
}
|
||||||
|
resp_meta_json = json.dumps(resp_meta, separators=(",", ":")).encode("utf-8")
|
||||||
|
await ws.send(encode_frame(stream_id, RESPONSE_HEADERS, 0, resp_meta_json))
|
||||||
|
|
||||||
|
echo_body = json.dumps({
|
||||||
|
"echo": True,
|
||||||
|
"received_method": request_meta.get("method"),
|
||||||
|
"received_url": request_meta.get("url"),
|
||||||
|
"received_body_len": len(request_body),
|
||||||
|
}, separators=(",", ":")).encode("utf-8")
|
||||||
|
await ws.send(encode_frame(stream_id, RESPONSE_BODY, 0, echo_body))
|
||||||
|
await ws.send(encode_frame(stream_id, STREAM_END, 0, b""))
|
||||||
|
print(f" [mock-proxy] stream={stream_id} sent response (200, {len(echo_body)} bytes)")
|
||||||
|
|
||||||
|
request_meta = None
|
||||||
|
request_body = b""
|
||||||
|
|
||||||
|
elif msg_type == STREAM_ERROR:
|
||||||
|
error_msg = payload.decode("utf-8", errors="replace")
|
||||||
|
print(f" [mock-proxy] stream={stream_id} got STREAM_ERROR: {error_msg}")
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Test runner
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
async def run_test() -> bool:
|
||||||
|
hub_port = 18085
|
||||||
|
hub_bind = f"127.0.0.1:{hub_port}"
|
||||||
|
hub_binary = os.path.join(
|
||||||
|
os.path.dirname(__file__),
|
||||||
|
"..",
|
||||||
|
"aether-hub",
|
||||||
|
"target",
|
||||||
|
"release",
|
||||||
|
"aether-hub",
|
||||||
|
)
|
||||||
|
hub_binary = os.path.normpath(hub_binary)
|
||||||
|
|
||||||
|
if not os.path.isfile(hub_binary):
|
||||||
|
print(f"FAIL: hub binary not found at {hub_binary}")
|
||||||
|
print(" run: cd aether-hub && cargo build --release")
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Start aether-hub (with control plane disabled since we don't have the app running)
|
||||||
|
print(f"[1/5] Starting aether-hub on {hub_bind} ...")
|
||||||
|
hub_proc = subprocess.Popen(
|
||||||
|
[
|
||||||
|
hub_binary,
|
||||||
|
"--bind", hub_bind,
|
||||||
|
"--proxy-idle-timeout", "0",
|
||||||
|
"--ping-interval", "30",
|
||||||
|
# Use a non-existent app URL -- control plane callbacks will fail silently
|
||||||
|
"--app-base-url", "http://127.0.0.1:19999",
|
||||||
|
],
|
||||||
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=subprocess.STDOUT,
|
||||||
|
text=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Wait for hub to be ready
|
||||||
|
for _ in range(30):
|
||||||
|
await asyncio.sleep(0.2)
|
||||||
|
try:
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
resp = await client.get(f"http://{hub_bind}/health", timeout=1.0)
|
||||||
|
if resp.status_code == 200:
|
||||||
|
print(" hub is healthy")
|
||||||
|
break
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
print("FAIL: hub did not start in time")
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Check initial stats
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
stats = (await client.get(f"http://{hub_bind}/stats")).json()
|
||||||
|
print(f" initial stats: {stats}")
|
||||||
|
assert stats["proxy_connections"] == 0
|
||||||
|
assert stats["nodes"] == 0
|
||||||
|
|
||||||
|
# Start mock proxy
|
||||||
|
node_id = "test-node-e2e"
|
||||||
|
proxy_ready = asyncio.Event()
|
||||||
|
print(f"\n[2/5] Connecting mock proxy (node_id={node_id}) ...")
|
||||||
|
proxy_task = asyncio.create_task(
|
||||||
|
mock_proxy(f"ws://{hub_bind}/proxy", node_id, proxy_ready)
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(proxy_ready.wait(), timeout=5.0)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
print("FAIL: mock proxy did not connect in time")
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Give hub a moment to register
|
||||||
|
await asyncio.sleep(0.3)
|
||||||
|
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
stats = (await client.get(f"http://{hub_bind}/stats")).json()
|
||||||
|
print(f" stats after connect: {stats}")
|
||||||
|
assert stats["proxy_connections"] == 1, f"expected 1 proxy connection, got {stats['proxy_connections']}"
|
||||||
|
assert stats["nodes"] == 1
|
||||||
|
|
||||||
|
# Send request through local relay
|
||||||
|
print(f"\n[3/5] Sending request via local relay ...")
|
||||||
|
request_body = b'{"model":"test","messages":[]}'
|
||||||
|
envelope = encode_relay_envelope(
|
||||||
|
{
|
||||||
|
"method": "POST",
|
||||||
|
"url": "https://api.example.com/v1/chat/completions",
|
||||||
|
"headers": {
|
||||||
|
"content-type": "application/json",
|
||||||
|
"authorization": "Bearer sk-test-123",
|
||||||
|
},
|
||||||
|
"timeout": 30,
|
||||||
|
},
|
||||||
|
request_body,
|
||||||
|
)
|
||||||
|
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
relay_url = f"http://{hub_bind}/local/relay/{node_id}"
|
||||||
|
resp = await client.post(
|
||||||
|
relay_url,
|
||||||
|
content=envelope,
|
||||||
|
headers={"content-type": "application/vnd.aether.tunnel-envelope"},
|
||||||
|
timeout=10.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
print(f" relay response: status={resp.status_code}")
|
||||||
|
assert resp.status_code == 200, f"expected 200, got {resp.status_code}: {resp.text}"
|
||||||
|
|
||||||
|
echo = resp.json()
|
||||||
|
print(f" echo body: {echo}")
|
||||||
|
assert echo["echo"] is True
|
||||||
|
assert echo["received_method"] == "POST"
|
||||||
|
assert echo["received_url"] == "https://api.example.com/v1/chat/completions"
|
||||||
|
assert echo["received_body_len"] == len(request_body)
|
||||||
|
|
||||||
|
assert resp.headers.get("x-test-echo") == "true"
|
||||||
|
|
||||||
|
# Verify active streams cleaned up
|
||||||
|
print(f"\n[4/5] Verifying stream cleanup ...")
|
||||||
|
await asyncio.sleep(0.2)
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
stats = (await client.get(f"http://{hub_bind}/stats")).json()
|
||||||
|
print(f" stats after request: {stats}")
|
||||||
|
assert stats["active_streams"] == 0, f"expected 0 active streams, got {stats['active_streams']}"
|
||||||
|
|
||||||
|
# Test error case: request to non-existent node
|
||||||
|
print(f"\n[5/5] Testing error cases ...")
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
resp = await client.post(
|
||||||
|
f"http://{hub_bind}/local/relay/non-existent-node",
|
||||||
|
content=encode_relay_envelope(
|
||||||
|
{"method": "GET", "url": "https://example.com", "headers": {}, "timeout": 5},
|
||||||
|
b"",
|
||||||
|
),
|
||||||
|
headers={"content-type": "application/vnd.aether.tunnel-envelope"},
|
||||||
|
timeout=5.0,
|
||||||
|
)
|
||||||
|
assert resp.status_code == 503, f"expected 503 for missing node, got {resp.status_code}"
|
||||||
|
assert resp.headers.get("x-aether-tunnel-error") == "connect"
|
||||||
|
print(f" missing node: status={resp.status_code}, error={resp.text}")
|
||||||
|
|
||||||
|
# Test: request from non-loopback should be rejected
|
||||||
|
# (can't easily test from non-loopback, but verify header is present for valid errors)
|
||||||
|
|
||||||
|
# Cleanup: cancel proxy
|
||||||
|
proxy_task.cancel()
|
||||||
|
try:
|
||||||
|
await proxy_task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
print("\n" + "=" * 50)
|
||||||
|
print("ALL TESTS PASSED")
|
||||||
|
print("=" * 50)
|
||||||
|
return True
|
||||||
|
|
||||||
|
finally:
|
||||||
|
hub_proc.send_signal(signal.SIGTERM)
|
||||||
|
try:
|
||||||
|
hub_proc.wait(timeout=5)
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
hub_proc.kill()
|
||||||
|
hub_proc.wait()
|
||||||
|
print("\n[cleanup] hub process stopped")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
success = asyncio.run(run_test())
|
||||||
|
sys.exit(0 if success else 1)
|
||||||
@@ -1,205 +1,90 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
from types import SimpleNamespace
|
import struct
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
import httpx
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from src.services.proxy_node.hub_config import HubConfig
|
from src.services.proxy_node.hub_config import HubConfig
|
||||||
from src.services.proxy_node.hub_transport import HubConnectionManager
|
from src.services.proxy_node.hub_transport import HubTunnelTransport
|
||||||
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
|
|
||||||
|
|
||||||
|
|
||||||
class _StubRedis:
|
class _FakeRelayClient:
|
||||||
def __init__(self, *, set_result: bool = True, set_exc: Exception | None = None) -> None:
|
def __init__(self, response: httpx.Response) -> None:
|
||||||
self.set_result = set_result
|
self.response = response
|
||||||
self.set_exc = set_exc
|
self.sent_request: httpx.Request | None = None
|
||||||
self.calls: list[tuple[str, str, int | None, bool | None]] = []
|
|
||||||
|
|
||||||
async def set(
|
def build_request(self, method: str, url: str, **kwargs: Any) -> httpx.Request:
|
||||||
self,
|
return httpx.Request(method, url, **kwargs)
|
||||||
key: str,
|
|
||||||
value: str,
|
async def send(self, request: httpx.Request, *, stream: bool = False) -> httpx.Response:
|
||||||
ex: int | None = None,
|
_ = stream
|
||||||
nx: bool | None = None,
|
self.sent_request = request
|
||||||
) -> bool:
|
self.response.request = request
|
||||||
self.calls.append((key, value, ex, nx))
|
return self.response
|
||||||
if self.set_exc is not None:
|
|
||||||
raise self.set_exc
|
async def aclose(self) -> None:
|
||||||
return self.set_result
|
return None
|
||||||
|
|
||||||
|
|
||||||
def _build_manager() -> HubConnectionManager:
|
def _relay_config() -> HubConfig:
|
||||||
return HubConnectionManager(
|
return HubConfig(
|
||||||
HubConfig(
|
enabled=True,
|
||||||
enabled=True,
|
url="http://127.0.0.1:8085",
|
||||||
url="ws://127.0.0.1:8085",
|
connect_timeout_seconds=1.0,
|
||||||
connect_timeout_seconds=1.0,
|
|
||||||
ping_interval_seconds=1.0,
|
|
||||||
send_timeout_seconds=1.0,
|
|
||||||
max_streams=16,
|
|
||||||
max_frame_size=1024 * 1024,
|
|
||||||
)
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _heartbeat_frame(payload: dict[str, Any]) -> Frame:
|
@pytest.mark.asyncio
|
||||||
return Frame(0, MsgType.HEARTBEAT_DATA, 0, json.dumps(payload).encode("utf-8"))
|
async def test_transport_encodes_local_relay_envelope(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
transport = HubTunnelTransport("node-1", timeout=12.0)
|
||||||
|
fake_client = _FakeRelayClient(httpx.Response(200, content=b"ok"))
|
||||||
|
monkeypatch.setattr("src.services.proxy_node.hub_transport.get_hub_config", _relay_config)
|
||||||
|
monkeypatch.setattr(transport, "_relay_client", fake_client)
|
||||||
|
|
||||||
|
request = httpx.Request(
|
||||||
|
"POST",
|
||||||
|
"https://example.com/v1/chat/completions",
|
||||||
|
headers={"content-type": "application/json", "connection": "keep-alive"},
|
||||||
|
content=b'{"hello":"world"}',
|
||||||
|
)
|
||||||
|
|
||||||
def _decode_ack(frame: Frame) -> dict[str, Any]:
|
response = await transport.handle_async_request(request)
|
||||||
assert frame.msg_type == MsgType.HEARTBEAT_ACK
|
assert response.status_code == 200
|
||||||
if not frame.payload:
|
await response.aclose()
|
||||||
return {}
|
|
||||||
return json.loads(frame.payload.decode("utf-8"))
|
assert fake_client.sent_request is not None
|
||||||
|
payload = fake_client.sent_request.content
|
||||||
|
assert payload is not None
|
||||||
|
meta_len = struct.unpack("!I", payload[:4])[0]
|
||||||
|
meta = json.loads(payload[4 : 4 + meta_len].decode("utf-8"))
|
||||||
|
assert meta == {
|
||||||
|
"method": "POST",
|
||||||
|
"url": "https://example.com/v1/chat/completions",
|
||||||
|
"headers": {"content-type": "application/json"},
|
||||||
|
"timeout": 12,
|
||||||
|
}
|
||||||
|
assert payload[4 + meta_len :] == b'{"hello":"world"}'
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_handle_heartbeat_normalizes_id_and_updates_db(
|
async def test_transport_maps_relay_timeout_to_read_timeout(
|
||||||
monkeypatch: pytest.MonkeyPatch,
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
) -> None:
|
) -> None:
|
||||||
manager = _build_manager()
|
transport = HubTunnelTransport("node-1", timeout=12.0)
|
||||||
captured_frames: list[Frame] = []
|
fake_client = _FakeRelayClient(
|
||||||
heartbeat_calls: list[dict[str, Any]] = []
|
httpx.Response(
|
||||||
|
504,
|
||||||
async def _fake_send_frame(frame: Frame) -> None:
|
headers={"x-aether-tunnel-error": "timeout"},
|
||||||
captured_frames.append(frame)
|
content=b"relay timed out",
|
||||||
|
|
||||||
class _FakeSession:
|
|
||||||
def close(self) -> None:
|
|
||||||
return
|
|
||||||
|
|
||||||
async def _fake_get_redis_client(*, require_redis: bool = False) -> _StubRedis:
|
|
||||||
_ = require_redis
|
|
||||||
return redis
|
|
||||||
|
|
||||||
def _fake_heartbeat(db: Any, **kwargs: Any) -> Any:
|
|
||||||
heartbeat_calls.append(kwargs)
|
|
||||||
return SimpleNamespace(
|
|
||||||
remote_config={"heartbeat_interval": 8, "upgrade_to": "0.2.3"},
|
|
||||||
config_version=5,
|
|
||||||
)
|
|
||||||
|
|
||||||
redis = _StubRedis(set_result=True)
|
|
||||||
monkeypatch.setattr(manager, "_send_frame", _fake_send_frame)
|
|
||||||
monkeypatch.setattr("src.database.create_session", lambda: _FakeSession())
|
|
||||||
monkeypatch.setattr("src.clients.get_redis_client", _fake_get_redis_client)
|
|
||||||
monkeypatch.setattr(
|
|
||||||
"src.services.proxy_node.service.ProxyNodeService.heartbeat", _fake_heartbeat
|
|
||||||
)
|
|
||||||
|
|
||||||
await manager._handle_heartbeat(
|
|
||||||
_heartbeat_frame(
|
|
||||||
{
|
|
||||||
"node_id": "node-1",
|
|
||||||
"heartbeat_session_id": "sess-1",
|
|
||||||
"heartbeat_id": 15.0,
|
|
||||||
"active_connections": 3,
|
|
||||||
"total_requests": 10,
|
|
||||||
"failed_requests": 1,
|
|
||||||
"dns_failures": 2,
|
|
||||||
"stream_errors": 0,
|
|
||||||
"proxy_metadata": {"version": "0.2.1"},
|
|
||||||
}
|
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
monkeypatch.setattr("src.services.proxy_node.hub_transport.get_hub_config", _relay_config)
|
||||||
|
monkeypatch.setattr(transport, "_relay_client", fake_client)
|
||||||
|
|
||||||
assert len(redis.calls) == 1
|
request = httpx.Request("GET", "https://example.com")
|
||||||
assert redis.calls[0][0] == "hub:heartbeat:node-1:sess-1:15"
|
|
||||||
assert heartbeat_calls and heartbeat_calls[0]["node_id"] == "node-1"
|
|
||||||
assert heartbeat_calls[0]["proxy_metadata"] == {"version": "0.2.1"}
|
|
||||||
assert len(captured_frames) == 1
|
|
||||||
|
|
||||||
ack = _decode_ack(captured_frames[0])
|
with pytest.raises(httpx.ReadTimeout, match="relay timed out"):
|
||||||
assert ack["heartbeat_id"] == 15
|
await transport.handle_async_request(request)
|
||||||
assert isinstance(ack["heartbeat_id"], int)
|
|
||||||
assert ack["remote_config"] == {"heartbeat_interval": 8, "upgrade_to": "0.2.3"}
|
|
||||||
assert ack["config_version"] == 5
|
|
||||||
assert ack["upgrade_to"] == "0.2.3"
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_handle_heartbeat_duplicate_skips_db_update(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
||||||
manager = _build_manager()
|
|
||||||
captured_frames: list[Frame] = []
|
|
||||||
heartbeat_called = False
|
|
||||||
|
|
||||||
async def _fake_send_frame(frame: Frame) -> None:
|
|
||||||
captured_frames.append(frame)
|
|
||||||
|
|
||||||
async def _fake_get_redis_client(*, require_redis: bool = False) -> _StubRedis:
|
|
||||||
_ = require_redis
|
|
||||||
return redis
|
|
||||||
|
|
||||||
def _fake_heartbeat(db: Any, **kwargs: Any) -> Any:
|
|
||||||
_ = db, kwargs
|
|
||||||
nonlocal heartbeat_called
|
|
||||||
heartbeat_called = True
|
|
||||||
return SimpleNamespace(remote_config={"heartbeat_interval": 8}, config_version=5)
|
|
||||||
|
|
||||||
redis = _StubRedis(set_result=False)
|
|
||||||
monkeypatch.setattr(manager, "_send_frame", _fake_send_frame)
|
|
||||||
monkeypatch.setattr("src.clients.get_redis_client", _fake_get_redis_client)
|
|
||||||
monkeypatch.setattr(
|
|
||||||
"src.services.proxy_node.service.ProxyNodeService.heartbeat", _fake_heartbeat
|
|
||||||
)
|
|
||||||
|
|
||||||
await manager._handle_heartbeat(
|
|
||||||
_heartbeat_frame({"node_id": "node-1", "heartbeat_id": 77, "total_requests": 20})
|
|
||||||
)
|
|
||||||
|
|
||||||
assert len(redis.calls) == 1
|
|
||||||
assert heartbeat_called is False
|
|
||||||
assert len(captured_frames) == 1
|
|
||||||
ack = _decode_ack(captured_frames[0])
|
|
||||||
assert ack == {"heartbeat_id": 77}
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_handle_heartbeat_redis_error_falls_back_to_db_update(
|
|
||||||
monkeypatch: pytest.MonkeyPatch,
|
|
||||||
) -> None:
|
|
||||||
manager = _build_manager()
|
|
||||||
captured_frames: list[Frame] = []
|
|
||||||
heartbeat_calls: list[dict[str, Any]] = []
|
|
||||||
|
|
||||||
async def _fake_send_frame(frame: Frame) -> None:
|
|
||||||
captured_frames.append(frame)
|
|
||||||
|
|
||||||
class _FakeSession:
|
|
||||||
def close(self) -> None:
|
|
||||||
return
|
|
||||||
|
|
||||||
async def _fake_get_redis_client(*, require_redis: bool = False) -> _StubRedis:
|
|
||||||
_ = require_redis
|
|
||||||
return redis
|
|
||||||
|
|
||||||
def _fake_heartbeat(db: Any, **kwargs: Any) -> Any:
|
|
||||||
heartbeat_calls.append(kwargs)
|
|
||||||
return SimpleNamespace(remote_config=None, config_version=0)
|
|
||||||
|
|
||||||
redis = _StubRedis(set_exc=RuntimeError("redis unavailable"))
|
|
||||||
monkeypatch.setattr(manager, "_send_frame", _fake_send_frame)
|
|
||||||
monkeypatch.setattr("src.database.create_session", lambda: _FakeSession())
|
|
||||||
monkeypatch.setattr("src.clients.get_redis_client", _fake_get_redis_client)
|
|
||||||
monkeypatch.setattr(
|
|
||||||
"src.services.proxy_node.service.ProxyNodeService.heartbeat", _fake_heartbeat
|
|
||||||
)
|
|
||||||
|
|
||||||
await manager._handle_heartbeat(
|
|
||||||
_heartbeat_frame(
|
|
||||||
{
|
|
||||||
"node_id": "node-1",
|
|
||||||
"heartbeat_id": 99,
|
|
||||||
"total_requests": 1,
|
|
||||||
"proxy_metadata": {"version": "0.2.2"},
|
|
||||||
}
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
assert heartbeat_calls and heartbeat_calls[0]["total_requests"] == 1
|
|
||||||
assert heartbeat_calls[0]["proxy_metadata"] == {"version": "0.2.2"}
|
|
||||||
assert len(captured_frames) == 1
|
|
||||||
ack = _decode_ack(captured_frames[0])
|
|
||||||
assert ack == {"heartbeat_id": 99}
|
|
||||||
|
|||||||
@@ -2,46 +2,40 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from src.services.proxy_node.hub_transport import HubResponseStream
|
from src.services.proxy_node.hub_transport import HubRelayResponseStream
|
||||||
from src.services.proxy_node.tunnel_manager import _StreamState
|
|
||||||
|
|
||||||
|
|
||||||
class _Manager:
|
class _FakeResponse:
|
||||||
def __init__(self) -> None:
|
def __init__(self, chunks: list[bytes]) -> None:
|
||||||
self.removed: list[int] = []
|
self._chunks = chunks
|
||||||
|
self.closed = False
|
||||||
|
|
||||||
def remove_stream(self, stream_id: int) -> None:
|
async def aiter_raw(self): # type: ignore[override]
|
||||||
self.removed.append(stream_id)
|
for chunk in self._chunks:
|
||||||
|
yield chunk
|
||||||
|
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
self.closed = True
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_hub_response_stream_removes_stream_after_normal_iteration() -> None:
|
async def test_hub_relay_response_stream_closes_after_iteration() -> None:
|
||||||
manager = _Manager()
|
response = _FakeResponse([b"hello", b"world"])
|
||||||
state = _StreamState(7)
|
stream = HubRelayResponseStream(response) # type: ignore[arg-type]
|
||||||
state.set_response_headers(200, {})
|
|
||||||
state.push_body_chunk(b"hello")
|
|
||||||
state.set_done()
|
|
||||||
|
|
||||||
stream = HubResponseStream(manager, state, timeout=0.1)
|
|
||||||
chunks = []
|
chunks = []
|
||||||
async for chunk in stream:
|
async for chunk in stream:
|
||||||
chunks.append(chunk)
|
chunks.append(chunk)
|
||||||
|
|
||||||
assert chunks == [b"hello"]
|
assert chunks == [b"hello", b"world"]
|
||||||
assert manager.removed == [7]
|
assert response.closed is True
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_hub_response_stream_removes_stream_after_body_error() -> None:
|
async def test_hub_relay_response_stream_aclose_closes_response() -> None:
|
||||||
manager = _Manager()
|
response = _FakeResponse([])
|
||||||
state = _StreamState(9)
|
stream = HubRelayResponseStream(response) # type: ignore[arg-type]
|
||||||
state.set_response_headers(200, {})
|
|
||||||
state.set_error("boom")
|
|
||||||
|
|
||||||
stream = HubResponseStream(manager, state, timeout=0.1)
|
await stream.aclose()
|
||||||
|
|
||||||
with pytest.raises(Exception):
|
assert response.closed is True
|
||||||
async for _chunk in stream:
|
|
||||||
pass
|
|
||||||
|
|
||||||
assert manager.removed == [9]
|
|
||||||
|
|||||||
@@ -1,66 +1,14 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import time
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
from src.services.proxy_node.hub_config import HubConfig
|
from src.services.proxy_node.hub_config import HubConfig
|
||||||
from src.services.proxy_node.hub_transport import HubConnectionManager
|
|
||||||
from src.services.proxy_node.tunnel_manager import TunnelStreamError
|
|
||||||
|
|
||||||
|
|
||||||
def _build_manager() -> HubConnectionManager:
|
def test_local_relay_url_uses_http_path() -> None:
|
||||||
return HubConnectionManager(
|
config = HubConfig(
|
||||||
HubConfig(
|
enabled=True,
|
||||||
enabled=True,
|
url="http://127.0.0.1:8085",
|
||||||
url="ws://127.0.0.1:8085",
|
connect_timeout_seconds=1.0,
|
||||||
connect_timeout_seconds=1.0,
|
|
||||||
ping_interval_seconds=1.0,
|
|
||||||
send_timeout_seconds=1.0,
|
|
||||||
max_streams=16,
|
|
||||||
max_frame_size=1024 * 1024,
|
|
||||||
)
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
assert (
|
||||||
def test_record_loop_lag_warning_does_not_degrade() -> None:
|
config.local_relay_url("node a/1")
|
||||||
manager = _build_manager()
|
== "http://127.0.0.1:8085/local/relay/node%20a%2F1"
|
||||||
|
)
|
||||||
manager._record_loop_lag(1.5)
|
|
||||||
|
|
||||||
assert manager._degraded_until == 0.0
|
|
||||||
|
|
||||||
|
|
||||||
def test_record_loop_lag_degrades_manager(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
||||||
manager = _build_manager()
|
|
||||||
now = 1234.0
|
|
||||||
monkeypatch.setattr("src.services.proxy_node.hub_transport._time.monotonic", lambda: now)
|
|
||||||
|
|
||||||
manager._record_loop_lag(4.0)
|
|
||||||
|
|
||||||
assert manager._degraded_until == pytest.approx(now + 12.0)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_send_request_rejects_while_manager_degraded(
|
|
||||||
monkeypatch: pytest.MonkeyPatch,
|
|
||||||
) -> None:
|
|
||||||
manager = _build_manager()
|
|
||||||
manager._degraded_until = time.monotonic() + 5.0
|
|
||||||
|
|
||||||
async def _fake_ensure_connected() -> None:
|
|
||||||
return None
|
|
||||||
|
|
||||||
monkeypatch.setattr(manager, "ensure_connected", _fake_ensure_connected)
|
|
||||||
|
|
||||||
with pytest.raises(TunnelStreamError, match="event loop degraded"):
|
|
||||||
await manager.send_request(
|
|
||||||
"node-1",
|
|
||||||
method="POST",
|
|
||||||
url="https://example.com/v1/chat/completions",
|
|
||||||
headers={"content-type": "application/json"},
|
|
||||||
body=b"{}",
|
|
||||||
timeout=5.0,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert manager._pending_streams == {}
|
|
||||||
|
|||||||
@@ -1,165 +0,0 @@
|
|||||||
import asyncio
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
from starlette.websockets import WebSocketState
|
|
||||||
|
|
||||||
from src.services.proxy_node.tunnel_manager import (
|
|
||||||
TunnelConnection,
|
|
||||||
TunnelManager,
|
|
||||||
TunnelStreamError,
|
|
||||||
)
|
|
||||||
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"
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_max_streams_from_header() -> None:
|
|
||||||
"""TunnelConnection respects proxy-advertised max_streams (clamped)"""
|
|
||||||
ws = _DummyWebSocket()
|
|
||||||
|
|
||||||
# Explicit value within range
|
|
||||||
conn = TunnelConnection("n", "n", ws, max_streams=256) # type: ignore[arg-type]
|
|
||||||
assert conn.max_streams == 256
|
|
||||||
|
|
||||||
# Clamped to minimum 64
|
|
||||||
conn_low = TunnelConnection("n", "n", ws, max_streams=10) # type: ignore[arg-type]
|
|
||||||
assert conn_low.max_streams == 64
|
|
||||||
|
|
||||||
# Clamped to maximum 2048
|
|
||||||
conn_high = TunnelConnection("n", "n", ws, max_streams=9999) # type: ignore[arg-type]
|
|
||||||
assert conn_high.max_streams == 2048
|
|
||||||
|
|
||||||
# None falls back to TunnelManager.MAX_STREAMS_PER_CONN
|
|
||||||
conn_default = TunnelConnection("n", "n", ws) # type: ignore[arg-type]
|
|
||||||
assert conn_default.max_streams == TunnelManager.MAX_STREAMS_PER_CONN
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_send_request_respects_per_conn_max_streams() -> None:
|
|
||||||
"""send_request raises TunnelStreamError when per-connection limit is reached"""
|
|
||||||
manager = TunnelManager()
|
|
||||||
ws = _DummyWebSocket()
|
|
||||||
# Set a very low max_streams (clamped to minimum 64)
|
|
||||||
conn = TunnelConnection("node-1", "node-1", ws, max_streams=64) # type: ignore[arg-type]
|
|
||||||
manager.register(conn)
|
|
||||||
|
|
||||||
# Fill up to max_streams
|
|
||||||
for i in range(64):
|
|
||||||
conn.create_stream(i * 2 + 2)
|
|
||||||
|
|
||||||
assert conn.stream_count == 64
|
|
||||||
|
|
||||||
# Next send_request should fail
|
|
||||||
with pytest.raises(TunnelStreamError, match="stream limit reached"):
|
|
||||||
await manager.send_request("node-1", method="GET", url="https://example.com", headers={})
|
|
||||||
Reference in New Issue
Block a user