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:
fawney19
2026-03-17 20:22:12 +08:00
parent 460eb5434d
commit 73a5325a38
24 changed files with 2253 additions and 3124 deletions

View File

@@ -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)

View File

@@ -0,0 +1,3 @@
from .hub import router
__all__ = ["router"]

97
src/api/internal/hub.py Normal file
View 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

View File

@@ -21,6 +21,7 @@ from src.api.announcements import router as announcement_router
# API路由
from src.api.auth import router as auth_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.payment import router as payment_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(public_router) # 公开API端点用户可查看提供商和模型
app.include_router(monitoring_router) # 监控端点
app.include_router(internal_router) # Hub 本地控制面端点
@app.get("/readyz", include_in_schema=False)

View File

@@ -101,23 +101,13 @@ async def _on_startup() -> None:
else:
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:
logger.info("启动 ProxyNode 心跳检测调度器...")
await proxy_node_health_scheduler.start()
async def _on_shutdown() -> None:
"""优雅关闭 tunnel 连接并停止心跳检测调度器"""
"""停止心跳检测调度器"""
import logging
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")
from src.services.proxy_node.hub_transport import shutdown_hub_connection_manager
await shutdown_hub_connection_manager()
from src.clients import get_redis_client
redis_client = await get_redis_client()

View File

@@ -12,13 +12,10 @@ from __future__ import annotations
import os
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_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)
@@ -26,14 +23,13 @@ class HubConfig:
enabled: bool
url: str
connect_timeout_seconds: float
ping_interval_seconds: float
send_timeout_seconds: float
max_streams: int
max_frame_size: int
@property
def worker_ws_url(self) -> str:
return f"{self.url.rstrip('/')}/worker"
def local_relay_base_url(self) -> str:
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
@@ -56,10 +52,6 @@ def get_hub_config() -> HubConfig:
enabled=docker_runtime,
url=_DOCKER_HUB_URL,
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

View File

@@ -1,43 +1,23 @@
"""
Hub 模式 tunnel transport
Worker 通过单条到 aether-hub 的 WebSocket 长连接转发 tunnel 帧
Worker 通过本机 aether-hub 的 HTTP relay 访问 tunnel 数据面,不再维护 /worker WebSocket
"""
from __future__ import annotations
import asyncio
import gzip
import json
import time as _time
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any
import struct
from typing import TYPE_CHECKING
import aiohttp
import httpx
from aiohttp import WSMsgType
from src.core.logger import logger
from .hub_config import HubConfig, get_hub_config
from .tunnel_manager import TunnelStreamError, _StreamState
from .tunnel_protocol import Frame, FrameFlags, MsgType, normalize_heartbeat_id
from .hub_config import get_hub_config
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(
{
"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)
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")
_RELAY_CONTENT_TYPE = "application/vnd.aether.tunnel-envelope"
_TUNNEL_ERROR_HEADER = "x-aether-tunnel-error"
class HubTunnelTransport(httpx.AsyncBaseTransport):
"""通过 aether-hub 转发请求的 httpx transport。"""
"""通过本机 aether-hub relay 转发请求的 httpx transport。"""
def __init__(self, node_id: str, timeout: float = 60.0) -> None:
self._node_id = node_id
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:
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] = {}
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")
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:
stream_state = await manager.send_request(
self._node_id,
method=request.method,
url=str(request.url),
headers=headers,
body=body,
timeout=self._timeout,
)
await stream_state.wait_headers(timeout=self._timeout)
relay_response = await self._relay_client.send(relay_request, stream=True)
except httpx.ConnectError as exc:
raise httpx.ConnectError(f"hub local relay connect failed: {exc}") from exc
except httpx.TimeoutException as exc:
raise httpx.ConnectError(f"hub local relay timeout: {exc}") from exc
return httpx.Response(
status_code=stream_state.status,
headers=httpx.Headers(stream_state.headers),
stream=HubResponseStream(manager, stream_state, timeout=self._timeout),
)
except TunnelStreamError as e:
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
tunnel_error = relay_response.headers.get(_TUNNEL_ERROR_HEADER)
if tunnel_error:
message = await _read_error_message(relay_response)
if tunnel_error == "timeout":
raise httpx.ReadTimeout(message or "hub relay timed out")
raise httpx.ConnectError(message or f"hub relay error: {tunnel_error}")
def _cleanup_stream(
self,
manager: HubConnectionManager,
stream_state: _StreamState | None,
) -> None:
if stream_state is None:
return
manager.remove_stream(stream_state.stream_id)
return httpx.Response(
status_code=relay_response.status_code,
headers=httpx.Headers(relay_response.headers),
stream=HubRelayResponseStream(relay_response),
request=request,
)
async def aclose(self) -> None:
await self._relay_client.aclose()
class HubResponseStream(httpx.AsyncByteStream):
def __init__(
self,
manager: HubConnectionManager,
stream_state: _StreamState,
timeout: float = 60.0,
) -> None:
self._manager = manager
self._stream_state = stream_state
self._timeout = timeout
class HubRelayResponseStream(httpx.AsyncByteStream):
def __init__(self, response: httpx.Response) -> None:
self._response = response
async def __aiter__(self) -> AsyncGenerator[bytes, None]:
try:
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
async for chunk in self._response.aiter_raw():
yield chunk
finally:
self._manager.remove_stream(self._stream_state.stream_id)
await self._response.aclose()
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]:
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 _encode_relay_envelope(meta: dict[str, object], 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
async def _compress_frame_payload_async(data: bytes) -> tuple[bytes, int]:
if len(data) < _TUNNEL_ASYNC_COMPRESS_THRESHOLD:
return _compress_frame_payload(data)
return await asyncio.to_thread(_compress_frame_payload, data)
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
async def _read_error_message(response: httpx.Response) -> str:
try:
payload = await response.aread()
return payload.decode("utf-8", errors="replace").strip()
finally:
await response.aclose()

View File

@@ -27,7 +27,7 @@ from src.core.logger import logger
# ---------------------------------------------------------------------------
_proxy_node_cache: dict[str, tuple[dict[str, Any] | None, float]] = {}
_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_MAX_SIZE = 256

View File

@@ -18,7 +18,14 @@ from sqlalchemy import func, update
from sqlalchemy.orm import Session
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 (
inject_auth_into_proxy_url,
@@ -117,7 +124,7 @@ def _normalize_proxy_metadata(
def build_heartbeat_ack(node: ProxyNode) -> dict[str, Any]:
"""从心跳后的节点构建 ACK 响应 payload供 hub_transport / tunnel_manager 使用)。"""
"""从心跳后的节点构建 ACK 响应 payload供 hub 控制面回调使用)。"""
result: dict[str, Any] = {}
if not node.remote_config:
return result
@@ -423,6 +430,60 @@ class ProxyNodeService:
raise NotFoundException(f"ProxyNode {node_id} 不存在", "proxy_node")
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
def unregister_node(db: Session, *, node_id: str) -> ProxyNode:
"""注销节点(设置为 OFFLINE"""

View File

@@ -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
# 标记 drainingsend_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