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
2026-02-27 21:25:42 +08:00
parent a174cf1b02
commit ee356c5e6e
13 changed files with 600 additions and 135 deletions

View File

@@ -0,0 +1,119 @@
import asyncio
import pytest
from starlette.websockets import WebSocketState
from src.services.proxy_node.tunnel_manager import TunnelConnection, TunnelManager
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
class _DummyWebSocket:
def __init__(self) -> None:
self.client_state = WebSocketState.CONNECTED
self.sent: list[bytes] = []
async def send_bytes(self, data: bytes) -> None:
self.sent.append(data)
async def close(self, code: int = 1000, reason: str | None = None) -> None: # noqa: ARG002
self.client_state = WebSocketState.DISCONNECTED
@pytest.mark.asyncio
async def test_pool_register_and_unregister() -> None:
"""register 将连接追加到池中unregister 按连接实例移除"""
manager = TunnelManager()
ws1 = _DummyWebSocket()
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
manager.register(conn1)
assert manager.connection_count("node-1") == 1
assert manager.get_connection("node-1") is conn1
ws2 = _DummyWebSocket()
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
manager.register(conn2)
assert manager.connection_count("node-1") == 2
# unregister conn1 不影响 conn2
assert manager.unregister(conn1) is True
assert manager.connection_count("node-1") == 1
assert manager.get_connection("node-1") is conn2
# 重复 unregister 返回 False
assert manager.unregister(conn1) is False
# unregister conn2 清空池
assert manager.unregister(conn2) is True
assert manager.get_connection("node-1") is None
@pytest.mark.asyncio
async def test_least_loaded_selection() -> None:
"""get_connection 返回 stream_count 最小的连接"""
manager = TunnelManager()
ws1 = _DummyWebSocket()
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
ws2 = _DummyWebSocket()
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
manager.register(conn1)
manager.register(conn2)
# 两个都空闲,返回任一(实际返回 min两者相同时返回第一个
selected = manager.get_connection("node-1")
assert selected in (conn1, conn2)
# 给 conn1 加一个 streamconn2 应被优先选中
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"