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:
119
tests/unit/test_tunnel_manager_connection_replacement.py
Normal file
119
tests/unit/test_tunnel_manager_connection_replacement.py
Normal file
@@ -0,0 +1,119 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from starlette.websockets import WebSocketState
|
||||
|
||||
from src.services.proxy_node.tunnel_manager import TunnelConnection, TunnelManager
|
||||
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
|
||||
|
||||
|
||||
class _DummyWebSocket:
|
||||
def __init__(self) -> None:
|
||||
self.client_state = WebSocketState.CONNECTED
|
||||
self.sent: list[bytes] = []
|
||||
|
||||
async def send_bytes(self, data: bytes) -> None:
|
||||
self.sent.append(data)
|
||||
|
||||
async def close(self, code: int = 1000, reason: str | None = None) -> None: # noqa: ARG002
|
||||
self.client_state = WebSocketState.DISCONNECTED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_register_and_unregister() -> None:
|
||||
"""register 将连接追加到池中,unregister 按连接实例移除"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
assert manager.connection_count("node-1") == 1
|
||||
assert manager.get_connection("node-1") is conn1
|
||||
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn2)
|
||||
assert manager.connection_count("node-1") == 2
|
||||
|
||||
# unregister conn1 不影响 conn2
|
||||
assert manager.unregister(conn1) is True
|
||||
assert manager.connection_count("node-1") == 1
|
||||
assert manager.get_connection("node-1") is conn2
|
||||
|
||||
# 重复 unregister 返回 False
|
||||
assert manager.unregister(conn1) is False
|
||||
|
||||
# unregister conn2 清空池
|
||||
assert manager.unregister(conn2) is True
|
||||
assert manager.get_connection("node-1") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_least_loaded_selection() -> None:
|
||||
"""get_connection 返回 stream_count 最小的连接"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
manager.register(conn2)
|
||||
|
||||
# 两个都空闲,返回任一(实际返回 min,两者相同时返回第一个)
|
||||
selected = manager.get_connection("node-1")
|
||||
assert selected in (conn1, conn2)
|
||||
|
||||
# 给 conn1 加一个 stream,conn2 应被优先选中
|
||||
conn1.create_stream(2)
|
||||
assert manager.get_connection("node-1") is conn2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dead_connections_cleaned_on_get() -> None:
|
||||
"""get_connection 自动清理 dead 连接"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
manager.register(conn2)
|
||||
|
||||
# 模拟 conn1 断开
|
||||
ws1.client_state = WebSocketState.DISCONNECTED
|
||||
assert manager.get_connection("node-1") is conn2
|
||||
assert manager.connection_count("node-1") == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_removed_connection_frames_ignored() -> None:
|
||||
"""已 unregister 的连接帧不应被处理"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
manager.register(conn2)
|
||||
|
||||
# unregister conn1
|
||||
manager.unregister(conn1)
|
||||
|
||||
ping = Frame(0, MsgType.PING, 0, b"hello")
|
||||
|
||||
# conn1 已不在池中,帧应被忽略
|
||||
await manager.handle_incoming_frame(conn1, ping)
|
||||
# 等待 fire-and-forget task 完成
|
||||
await asyncio.sleep(0.05)
|
||||
assert ws1.sent == []
|
||||
|
||||
# conn2 仍在池中,帧正常处理
|
||||
await manager.handle_incoming_frame(conn2, ping)
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(ws2.sent) == 1
|
||||
pong = Frame.decode(ws2.sent[0])
|
||||
assert pong.msg_type == MsgType.PONG
|
||||
assert pong.payload == b"hello"
|
||||
Reference in New Issue
Block a user