mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(proxy-tunnel): 实现隧道连接池与 TCP 底层优化
Rust 端: - 支持每个 server 多条并行 WebSocket 连接 (tunnel_connections 配置) - 手动控制 TCP 连接: connect/handshake 超时、keepalive、NODELAY (socket2) - 预构建共享 TLS ClientConfig 避免每次重连重新解析根证书 - 增加 stale timeout 检测无数据连接，智能 backoff 按连接存活时长重置 - 每条连接独立 reconnect 计数器，仅主连接 (conn_idx=0) 发送心跳 - dispatcher 的错误帧和 PONG 改用 try_send 避免阻塞读循环 - stream_handler 增加 frame 发送超时保护防止写阻塞 Python 端: - TunnelManager 改为连接池，按 least-loaded 策略分配请求 - handle_incoming_frame 按连接实例路由，unregister 精确移除单条连接 - 心跳和 PONG 回复改为 fire-and-forget 避免阻塞主读循环 - send_frame 增加超时保护防止 TCP 写阻塞级联 - WebSocket 先 accept 再认证，auth 加超时 - 调整 idle timeout (90s) 和 ping 间隔 (15s)
This commit is contained in:
fawney19
@@ -23,11 +23,11 @@ router = APIRouter()
|
||||
# 单帧最大 64 MB -- AI API 请求体可能包含多张 base64 图片，需要足够余量
|
||||
_MAX_FRAME_SIZE = 64 * 1024 * 1024
|
||||
|
||||
# WebSocket 空闲超时（秒）-- proxy 端 WebSocket ping 间隔 15s + 心跳 30s，180s 提供充足余量
|
||||
_IDLE_TIMEOUT = 180.0
|
||||
# WebSocket 空闲超时（秒）-- 需覆盖客户端 stale_timeout(45s) + 重连延迟(最长30s) 的窗口期
|
||||
_IDLE_TIMEOUT = 90.0
|
||||
|
||||
# 服务端应用层 ping 间隔（秒）-- 确保即使 proxy 端心跳延迟，连接也不会因中间代理空闲超时而断开
|
||||
_SERVER_PING_INTERVAL = 30.0
|
||||
# 服务端应用层 ping 间隔（秒）-- 与客户端 ping 间隔一致，确保高频心跳
|
||||
_SERVER_PING_INTERVAL = 15.0
|
||||
|
||||
|
||||
async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
|
||||
@@ -78,22 +78,26 @@ async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
|
||||
@router.websocket("/api/internal/proxy-tunnel")
|
||||
async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
||||
"""aether-proxy tunnel WebSocket 端点"""
|
||||
# 先 accept，避免认证（DB/Redis）慢时卡在握手阶段导致网关返回 502。
|
||||
await ws.accept()
|
||||
|
||||
try:
|
||||
auth = await _authenticate(ws)
|
||||
auth = await asyncio.wait_for(_authenticate(ws), timeout=10.0)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("tunnel auth timeout")
|
||||
await ws.close(code=4002, reason="authentication timeout")
|
||||
return
|
||||
except Exception as e:
|
||||
logger.warning("tunnel auth error: {}", e)
|
||||
await ws.accept()
|
||||
await ws.close(code=4002, reason="authentication error")
|
||||
return
|
||||
|
||||
if not auth:
|
||||
await ws.accept()
|
||||
await ws.close(code=4001, reason="unauthorized")
|
||||
return
|
||||
|
||||
node_id: str = auth[0]
|
||||
node_name: str = auth[1]
|
||||
await ws.accept()
|
||||
|
||||
manager = get_tunnel_manager()
|
||||
conn = TunnelConnection(node_id, node_name, ws)
|
||||
@@ -129,7 +133,7 @@ async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
||||
logger.warning("tunnel frame decode error from {}: {}", node_id, e)
|
||||
continue
|
||||
|
||||
await manager.handle_incoming_frame(node_id, frame)
|
||||
await manager.handle_incoming_frame(conn, frame)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info("tunnel WebSocket disconnected: node_id={}", node_id)
|
||||
@@ -137,8 +141,11 @@ async def proxy_tunnel_ws(ws: WebSocket) -> None:
|
||||
logger.error("tunnel WebSocket error for node_id={}: {}", node_id, e)
|
||||
finally:
|
||||
ping_task.cancel()
|
||||
manager.unregister(node_id)
|
||||
await _update_tunnel_status(node_id, connected=False)
|
||||
manager.unregister(conn)
|
||||
if not manager.has_tunnel(node_id):
|
||||
await _update_tunnel_status(node_id, connected=False)
|
||||
else:
|
||||
logger.info("tunnel connection closed but pool still active: node_id={}", node_id)
|
||||
|
||||
|
||||
async def _ping_loop(conn: TunnelConnection) -> None:
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
WebSocket 隧道管理器
|
||||
|
||||
管理所有活跃的 aether-proxy tunnel 连接，提供通过隧道发送 HTTP 请求的能力。
|
||||
每个 proxy node 最多一条 tunnel 连接。
|
||||
每个 proxy node 可持有多条 tunnel 连接（连接池），请求按 least-loaded 策略分配。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -49,9 +49,18 @@ class TunnelConnection:
|
||||
def is_alive(self) -> bool:
|
||||
return self.ws.client_state == WebSocketState.CONNECTED
|
||||
|
||||
async def send_frame(self, frame: Frame) -> None:
|
||||
async with self._write_lock:
|
||||
await self.ws.send_bytes(frame.encode())
|
||||
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)
|
||||
@@ -163,43 +172,101 @@ class TunnelStreamError(Exception):
|
||||
|
||||
|
||||
class TunnelManager:
|
||||
"""管理所有活跃的 tunnel 连接"""
|
||||
"""管理所有活跃的 tunnel 连接（支持每个 node 多条连接的连接池）"""
|
||||
|
||||
# 单条 tunnel 上允许的最大并发 stream 数（超出时拒绝新请求）
|
||||
MAX_STREAMS_PER_CONN = 2048
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._connections: dict[str, TunnelConnection] = {} # node_id -> conn
|
||||
self._connections: dict[str, list[TunnelConnection]] = {} # node_id -> [conn, ...]
|
||||
self._background_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
def _background(self, coro: Any) -> None: # noqa: ANN401
|
||||
"""启动 fire-and-forget task，通过 set 持有引用防止 GC 回收，完成后自动清理"""
|
||||
task = asyncio.create_task(coro)
|
||||
self._background_tasks.add(task)
|
||||
task.add_done_callback(self._background_tasks.discard)
|
||||
|
||||
@property
|
||||
def active_count(self) -> int:
|
||||
return len(self._connections)
|
||||
return sum(len(conns) for conns in self._connections.values())
|
||||
|
||||
def get_connection(self, node_id: str) -> TunnelConnection | None:
|
||||
conn = self._connections.get(node_id)
|
||||
if conn and not conn.is_alive:
|
||||
self._connections.pop(node_id, None)
|
||||
conn.cancel_all_streams()
|
||||
"""获取负载最低的存活连接，同时清理 dead 连接"""
|
||||
conns = self._connections.get(node_id)
|
||||
if not conns:
|
||||
return None
|
||||
return conn
|
||||
|
||||
# 清理 dead 连接
|
||||
alive = [c for c in conns if c.is_alive]
|
||||
dead = [c for c in conns if not c.is_alive]
|
||||
for c in dead:
|
||||
c.cancel_all_streams()
|
||||
|
||||
if not alive:
|
||||
self._connections.pop(node_id, None)
|
||||
return None
|
||||
|
||||
if len(alive) != len(conns):
|
||||
self._connections[node_id] = alive
|
||||
|
||||
# Least-loaded: 选 stream_count 最小的连接
|
||||
return min(alive, key=lambda c: c.stream_count)
|
||||
|
||||
def register(self, conn: TunnelConnection) -> None:
|
||||
old = self._connections.get(conn.node_id)
|
||||
if old:
|
||||
old.cancel_all_streams()
|
||||
self._connections[conn.node_id] = conn
|
||||
logger.info("tunnel connected: node_id={}, name={}", conn.node_id, conn.node_name)
|
||||
"""注册一条新连接到连接池"""
|
||||
conns = self._connections.get(conn.node_id)
|
||||
if conns is None:
|
||||
conns = []
|
||||
self._connections[conn.node_id] = conns
|
||||
conns.append(conn)
|
||||
logger.info(
|
||||
"tunnel connected: node_id={}, name={}, pool_size={}",
|
||||
conn.node_id,
|
||||
conn.node_name,
|
||||
len(conns),
|
||||
)
|
||||
|
||||
def unregister(self, node_id: str) -> None:
|
||||
conn = self._connections.pop(node_id, None)
|
||||
if conn:
|
||||
conn.cancel_all_streams()
|
||||
logger.info("tunnel disconnected: node_id={}, name={}", node_id, conn.node_name)
|
||||
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
|
||||
logger.info(
|
||||
"tunnel disconnected: node_id={}, name={}, remaining={}",
|
||||
conn.node_id,
|
||||
conn.node_name,
|
||||
remaining,
|
||||
)
|
||||
return True
|
||||
|
||||
def has_tunnel(self, node_id: str) -> bool:
|
||||
conn = self.get_connection(node_id)
|
||||
return conn is not None
|
||||
|
||||
def connection_count(self, node_id: str) -> int:
|
||||
"""返回指定 node 当前存活的连接数"""
|
||||
conns = self._connections.get(node_id)
|
||||
if not conns:
|
||||
return 0
|
||||
return sum(1 for c in conns if c.is_alive)
|
||||
|
||||
async def send_request(
|
||||
self,
|
||||
node_id: str,
|
||||
@@ -248,10 +315,15 @@ class TunnelManager:
|
||||
|
||||
return stream_state
|
||||
|
||||
async def handle_incoming_frame(self, node_id: str, frame: Frame) -> None:
|
||||
"""处理从 proxy 收到的响应帧"""
|
||||
conn = self.get_connection(node_id)
|
||||
if not conn:
|
||||
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)
|
||||
@@ -281,10 +353,19 @@ class TunnelManager:
|
||||
conn.remove_stream(frame.stream_id)
|
||||
|
||||
elif frame.msg_type == MsgType.HEARTBEAT_DATA:
|
||||
await self._handle_heartbeat(conn, frame)
|
||||
# fire-and-forget: 不阻塞主读循环
|
||||
self._background(self._handle_heartbeat(conn, frame))
|
||||
|
||||
elif frame.msg_type == MsgType.PING:
|
||||
await conn.send_frame(Frame(0, MsgType.PONG, 0, frame.payload))
|
||||
# fire-and-forget: pong 回复不阻塞读循环
|
||||
self._background(self._send_pong(conn, frame.payload))
|
||||
|
||||
async def _send_pong(self, conn: TunnelConnection, payload: bytes) -> None:
|
||||
"""发送 PONG 回复（fire-and-forget，不阻塞主读循环）"""
|
||||
try:
|
||||
await conn.send_frame(Frame(0, MsgType.PONG, 0, payload))
|
||||
except TunnelStreamError:
|
||||
pass # best-effort pong
|
||||
|
||||
async def _handle_heartbeat(self, conn: TunnelConnection, frame: Frame) -> None:
|
||||
"""处理 proxy 上报的心跳数据，更新 DB，返回 ACK"""
|
||||
@@ -320,7 +401,10 @@ class TunnelManager:
|
||||
logger.warning("tunnel heartbeat DB update failed: {}", e)
|
||||
ack = {}
|
||||
|
||||
await conn.send_frame(Frame(0, MsgType.HEARTBEAT_ACK, 0, json.dumps(ack).encode()))
|
||||
try:
|
||||
await conn.send_frame(Frame(0, MsgType.HEARTBEAT_ACK, 0, json.dumps(ack).encode()))
|
||||
except TunnelStreamError:
|
||||
logger.debug("heartbeat ACK send failed for node_id={}", conn.node_id)
|
||||
|
||||
|
||||
# 全局单例
|
||||
|
||||
Reference in New Issue
Block a user