mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
refactor(tunnel): 移除直连tunnel模式,统一使用Hub转发
- 删除 TunnelTransport 及 TunnelManager 相关引用,所有 tunnel 请求统一走 Hub - 简化 health_scheduler,不再依赖进程内 TunnelManager 状态判断节点在线 - 简化 resolver,移除本地 tunnel miss 判断与多 worker 告警逻辑 - 简化 service 中 tunnel 连通性检测,统一以 DB 状态为准 - 启动/关闭流程移除 hub_enabled 分支,始终初始化 Hub 连接 - Rust 端 RequestMeta.timeout 增加浮点数反序列化支持,Python 端确保发送整数
This commit is contained in:
@@ -144,7 +144,7 @@ pub struct RequestMeta {
|
||||
pub method: String,
|
||||
pub url: String,
|
||||
pub headers: std::collections::HashMap<String, String>,
|
||||
#[serde(default = "default_timeout")]
|
||||
#[serde(default = "default_timeout", deserialize_with = "deserialize_timeout")]
|
||||
pub timeout: u64,
|
||||
}
|
||||
|
||||
@@ -152,6 +152,36 @@ 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON payload for RESPONSE_HEADERS frames.
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
pub struct ResponseMeta {
|
||||
@@ -209,3 +239,23 @@ fn compress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
|
||||
let compressed = encoder.finish()?;
|
||||
Ok(Bytes::from(compressed))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RequestMeta;
|
||||
|
||||
#[test]
|
||||
fn request_meta_accepts_integer_timeout() {
|
||||
let raw = br#"{"method":"GET","url":"https://example.com","headers":{},"timeout":15}"#;
|
||||
let meta: RequestMeta = serde_json::from_slice(raw).expect("parse request meta");
|
||||
assert_eq!(meta.timeout, 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_meta_accepts_integer_like_float_timeout() {
|
||||
let raw =
|
||||
br#"{"method":"GET","url":"https://example.com","headers":{},"timeout":15.0}"#;
|
||||
let meta: RequestMeta = serde_json::from_slice(raw).expect("parse request meta");
|
||||
assert_eq!(meta.timeout, 15);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -515,11 +515,6 @@ app.include_router(dashboard_router) # 仪表盘端点
|
||||
app.include_router(public_router) # 公开API端点(用户可查看提供商和模型)
|
||||
app.include_router(monitoring_router) # 监控端点
|
||||
|
||||
# WebSocket 隧道端点(aether-proxy tunnel 模式)
|
||||
from src.api.admin.proxy_tunnel import router as proxy_tunnel_router
|
||||
|
||||
app.include_router(proxy_tunnel_router)
|
||||
|
||||
|
||||
def main() -> Any:
|
||||
# 初始化新日志系统
|
||||
|
||||
@@ -23,9 +23,8 @@ if TYPE_CHECKING:
|
||||
def _reset_tunnel_connected_on_startup() -> None:
|
||||
"""服务端启动时将所有 tunnel_connected=True 的节点重置为 False/OFFLINE。
|
||||
|
||||
服务端重启后 TunnelManager 内存状态丢失,但 DB 中可能残留
|
||||
tunnel_connected=True 的记录。如果不重置,health_scheduler 会错误地
|
||||
将这些节点标记为 ONLINE,而实际上 tunnel 并未连接。
|
||||
服务端重启后 DB 中可能残留 tunnel_connected=True 的记录。
|
||||
如果不重置,节点会在 Hub 状态广播到来前短暂显示 ONLINE。
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
@@ -78,37 +77,27 @@ async def _on_startup() -> None:
|
||||
|
||||
from src.config import config
|
||||
from src.services.proxy_node.health_scheduler import get_proxy_node_health_scheduler
|
||||
from src.services.proxy_node.hub_config import get_hub_config
|
||||
from src.utils.task_coordinator import StartupTaskCoordinator
|
||||
|
||||
logger = logging.getLogger("aether.modules.proxy_nodes")
|
||||
hub_enabled = get_hub_config().enabled
|
||||
|
||||
if config.worker_processes > 1 and not hub_enabled:
|
||||
logger.warning(
|
||||
"检测到 WEB_CONCURRENCY={}。Proxy tunnel 连接是进程内资源,"
|
||||
"多 worker 场景可能出现节点显示 ONLINE 但当前 worker 无可用 tunnel 的情况。"
|
||||
"建议设置 GUNICORN_WORKERS/WEB_CONCURRENCY=1。",
|
||||
config.worker_processes,
|
||||
)
|
||||
elif config.worker_processes > 1 and hub_enabled:
|
||||
if config.worker_processes > 1:
|
||||
logger.info(
|
||||
"检测到 WEB_CONCURRENCY={},Hub 模式已启用,允许多 worker 共享 tunnel。",
|
||||
"检测到 WEB_CONCURRENCY={},Hub 模式允许多 worker 共享 tunnel。",
|
||||
config.worker_processes,
|
||||
)
|
||||
|
||||
# Hub 模式下,worker 启动时主动建立 /worker 长连接:
|
||||
# 启动时主动建立 /worker 长连接:
|
||||
# - 立即接收 NODE_STATUS 广播,避免“proxy 已连但 UI 仍显示离线”的窗口期
|
||||
# - 确保后续 tunnel 请求不需要首请求触发懒连接
|
||||
if hub_enabled:
|
||||
from src.services.proxy_node.hub_transport import get_hub_connection_manager
|
||||
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: {}", e)
|
||||
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: {}", e)
|
||||
|
||||
from src.clients import get_redis_client
|
||||
|
||||
@@ -132,22 +121,13 @@ async def _on_shutdown() -> None:
|
||||
import logging
|
||||
|
||||
from src.services.proxy_node.health_scheduler import get_proxy_node_health_scheduler
|
||||
from src.services.proxy_node.hub_config import get_hub_config
|
||||
from src.utils.task_coordinator import StartupTaskCoordinator
|
||||
|
||||
logger = logging.getLogger("aether.modules.proxy_nodes")
|
||||
|
||||
hub_enabled = get_hub_config().enabled
|
||||
if hub_enabled:
|
||||
from src.services.proxy_node.hub_transport import shutdown_hub_connection_manager
|
||||
from src.services.proxy_node.hub_transport import shutdown_hub_connection_manager
|
||||
|
||||
await shutdown_hub_connection_manager()
|
||||
else:
|
||||
# 先向所有 tunnel 连接发送 GoAway,让 proxy 端立即重连到其他 worker
|
||||
from src.services.proxy_node.tunnel_manager import get_tunnel_manager
|
||||
|
||||
manager = get_tunnel_manager()
|
||||
await manager.shutdown_all()
|
||||
await shutdown_hub_connection_manager()
|
||||
|
||||
from src.clients import get_redis_client
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
ProxyNode 心跳检测调度器
|
||||
|
||||
定期检查 proxy_nodes 的连接健康状态,更新节点状态:
|
||||
- 本地 TunnelManager 观测到连接 -> ONLINE(自愈)
|
||||
- 心跳正常且 tunnel_connected=True -> ONLINE(自愈)
|
||||
- 心跳超时(跨 worker 共享信号) -> OFFLINE
|
||||
"""
|
||||
|
||||
@@ -83,15 +83,11 @@ class ProxyNodeHealthScheduler:
|
||||
await self._cleanup_old_events()
|
||||
|
||||
async def _check_heartbeats(self) -> None:
|
||||
from src.services.proxy_node.tunnel_manager import get_tunnel_manager
|
||||
|
||||
manager = get_tunnel_manager()
|
||||
db = create_session()
|
||||
try:
|
||||
now = datetime.now(timezone.utc)
|
||||
# 检查所有非手动节点(手动节点无心跳,始终保持 ONLINE)
|
||||
# 包括 OFFLINE 节点:tunnel 重连后如果 _update_tunnel_status 失败,
|
||||
# 健康检查需要能根据 TunnelManager 内存状态将其恢复为 ONLINE
|
||||
# 包括 OFFLINE 节点:心跳恢复后可自愈
|
||||
nodes = (
|
||||
db.query(ProxyNode)
|
||||
.filter(
|
||||
@@ -104,23 +100,6 @@ class ProxyNodeHealthScheduler:
|
||||
|
||||
changed = 0
|
||||
for node in nodes:
|
||||
# 注意:TunnelManager 仅是当前 worker 的进程内状态,跨 worker 不共享。
|
||||
# 因此“本地无 tunnel”不能直接判定 OFFLINE(可能连接在其他 worker)。
|
||||
# OFFLINE 统一由心跳超时判定,避免多进程误判。
|
||||
actually_connected_local = manager.has_tunnel(node.id)
|
||||
|
||||
if actually_connected_local:
|
||||
if not node.tunnel_connected:
|
||||
node.tunnel_connected = True
|
||||
node.tunnel_connected_at = now
|
||||
changed += 1
|
||||
if node.status != ProxyNodeStatus.ONLINE:
|
||||
node.status = ProxyNodeStatus.ONLINE
|
||||
node.updated_at = now
|
||||
changed += 1
|
||||
continue
|
||||
|
||||
# 本地无 tunnel:仅在心跳超时时标记 OFFLINE
|
||||
if heartbeat_is_stale(node, now):
|
||||
if node.tunnel_connected:
|
||||
node.tunnel_connected = False
|
||||
@@ -130,6 +109,13 @@ class ProxyNodeHealthScheduler:
|
||||
node.status = ProxyNodeStatus.OFFLINE
|
||||
node.updated_at = now
|
||||
changed += 1
|
||||
continue
|
||||
|
||||
# 心跳正常且连接状态为已连时,确保 ONLINE(自愈状态不一致)
|
||||
if node.tunnel_connected and node.status != ProxyNodeStatus.ONLINE:
|
||||
node.status = ProxyNodeStatus.ONLINE
|
||||
node.updated_at = now
|
||||
changed += 1
|
||||
|
||||
if changed:
|
||||
db.commit()
|
||||
|
||||
@@ -465,7 +465,7 @@ class HubConnectionManager:
|
||||
"method": method,
|
||||
"url": url,
|
||||
"headers": headers,
|
||||
"timeout": timeout,
|
||||
"timeout": int(timeout),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
|
||||
@@ -10,7 +10,6 @@ from __future__ import annotations
|
||||
import gzip
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
from urllib.parse import quote, urlparse
|
||||
@@ -27,55 +26,7 @@ from src.core.logger import logger
|
||||
_proxy_node_cache: dict[str, tuple[dict[str, Any] | None, float]] = {}
|
||||
_PROXY_NODE_CACHE_TTL_SECONDS = 15.0
|
||||
_PROXY_NODE_CACHE_NEGATIVE_TTL_SECONDS = 5.0 # 不可用节点使用更短的 TTL,加速恢复感知
|
||||
_PROXY_NODE_CACHE_TUNNEL_LOCAL_MISS_TTL_SECONDS = 0.5 # 本地 worker 无 tunnel 时,快速重试
|
||||
_PROXY_NODE_CACHE_MAX_SIZE = 256
|
||||
_TUNNEL_LOCAL_MISS_LOG_COOLDOWN_SECONDS = 30.0
|
||||
_tunnel_local_miss_log_next_at: dict[str, float] = {}
|
||||
|
||||
|
||||
def _is_hub_mode_enabled() -> bool:
|
||||
from src.services.proxy_node.hub_config import get_hub_config
|
||||
|
||||
return get_hub_config().enabled
|
||||
|
||||
|
||||
def _build_tunnel_local_miss_message(node_id: str) -> str:
|
||||
"""构建“当前 worker 无本地 tunnel”的用户可读错误信息。"""
|
||||
return (
|
||||
f"代理节点 {node_id} 当前 worker 无 tunnel 连接(pid={os.getpid()})。"
|
||||
"这通常发生在多 worker 部署(WEB_CONCURRENCY/GUNICORN_WORKERS > 1)时。"
|
||||
"请设置 GUNICORN_WORKERS=1(或 WEB_CONCURRENCY=1),"
|
||||
"或确保每个 worker 都建立该节点的 tunnel 连接。"
|
||||
)
|
||||
|
||||
|
||||
def _is_tunnel_local_miss(node_id: str) -> bool:
|
||||
"""判断节点是否“全局在线但当前 worker 无本地 tunnel 连接”。
|
||||
|
||||
仅在错误路径调用(build_proxy_url 解析失败时),用于提供更准确的报错信息。
|
||||
"""
|
||||
if _is_hub_mode_enabled():
|
||||
return False
|
||||
|
||||
from src.database import create_session
|
||||
from src.models.database import ProxyNode, ProxyNodeStatus
|
||||
from src.services.proxy_node.tunnel_manager import get_tunnel_manager
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
node = db.query(ProxyNode).filter(ProxyNode.id == node_id).first()
|
||||
if not node:
|
||||
return False
|
||||
if node.is_manual or not node.tunnel_mode:
|
||||
return False
|
||||
if node.status != ProxyNodeStatus.ONLINE:
|
||||
return False
|
||||
manager = get_tunnel_manager()
|
||||
return not manager.has_tunnel(node_id)
|
||||
except Exception:
|
||||
return False
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _get_proxy_node_info(node_id: str) -> dict[str, Any] | None:
|
||||
@@ -116,41 +67,14 @@ def _get_proxy_node_info(node_id: str) -> dict[str, Any] | None:
|
||||
_proxy_node_cache[node_id] = (None, now + _PROXY_NODE_CACHE_NEGATIVE_TTL_SECONDS)
|
||||
return None
|
||||
|
||||
# tunnel 模式节点:
|
||||
# - Hub 模式:信任 DB 的 tunnel_connected/status(由 Hub 广播统一维护)
|
||||
# - 非 Hub 模式:以当前 worker 本地 TunnelManager 状态为准
|
||||
# tunnel 模式节点:统一以 DB 的 tunnel_connected/status 为准(由 Hub 广播维护)
|
||||
if node.tunnel_mode and not node.is_manual:
|
||||
if _is_hub_mode_enabled():
|
||||
if node.status != ProxyNodeStatus.ONLINE or not bool(node.tunnel_connected):
|
||||
_proxy_node_cache[node_id] = (
|
||||
None,
|
||||
now + _PROXY_NODE_CACHE_NEGATIVE_TTL_SECONDS,
|
||||
)
|
||||
return None
|
||||
else:
|
||||
from src.services.proxy_node.tunnel_manager import get_tunnel_manager
|
||||
|
||||
manager = get_tunnel_manager()
|
||||
if not manager.has_tunnel(node_id):
|
||||
# 多 worker 部署时,DB 可能显示 ONLINE(其他 worker 有 tunnel),
|
||||
# 但当前 worker 无本地连接,请求仍不可用。记录限频告警便于定位。
|
||||
if now >= _tunnel_local_miss_log_next_at.get(node_id, 0.0):
|
||||
_tunnel_local_miss_log_next_at[node_id] = (
|
||||
now + _TUNNEL_LOCAL_MISS_LOG_COOLDOWN_SECONDS
|
||||
)
|
||||
logger.warning(
|
||||
"tunnel node {} has no local connection on pid={} "
|
||||
"(db_status={}, db_tunnel_connected={}), request may fail on this worker",
|
||||
node_id,
|
||||
os.getpid(),
|
||||
str(getattr(node, "status", "unknown")),
|
||||
bool(getattr(node, "tunnel_connected", False)),
|
||||
)
|
||||
_proxy_node_cache[node_id] = (
|
||||
None,
|
||||
now + _PROXY_NODE_CACHE_TUNNEL_LOCAL_MISS_TTL_SECONDS,
|
||||
)
|
||||
return None
|
||||
if node.status != ProxyNodeStatus.ONLINE or not bool(node.tunnel_connected):
|
||||
_proxy_node_cache[node_id] = (
|
||||
None,
|
||||
now + _PROXY_NODE_CACHE_NEGATIVE_TTL_SECONDS,
|
||||
)
|
||||
return None
|
||||
value: dict[str, Any] = {
|
||||
"name": node.name,
|
||||
"ip": node.ip,
|
||||
@@ -199,7 +123,6 @@ _SYSTEM_PROXY_CACHE_TTL = 60.0
|
||||
def invalidate_proxy_node_cache(node_id: str) -> None:
|
||||
"""主动清除指定节点的信息缓存(tunnel 断开时调用,避免使用过期的连接状态)"""
|
||||
_proxy_node_cache.pop(node_id, None)
|
||||
_tunnel_local_miss_log_next_at.pop(node_id, None)
|
||||
|
||||
|
||||
def invalidate_system_proxy_cache() -> None:
|
||||
@@ -509,11 +432,7 @@ def build_proxy_url(proxy_config: dict[str, Any]) -> str | None:
|
||||
node_info = _get_proxy_node_info(node_id)
|
||||
if not node_info:
|
||||
logger.warning("代理节点不可用(离线或不存在): node_id={}", node_id)
|
||||
message = (
|
||||
_build_tunnel_local_miss_message(node_id)
|
||||
if _is_tunnel_local_miss(node_id)
|
||||
else f"代理节点 {node_id} 不可用"
|
||||
)
|
||||
message = f"代理节点 {node_id} 不可用"
|
||||
raise ProxyNodeUnavailableError(message, node_id=node_id)
|
||||
|
||||
# 手动节点:直接使用存储的代理 URL(含认证信息)
|
||||
|
||||
@@ -556,38 +556,14 @@ class ProxyNodeService:
|
||||
|
||||
# tunnel 节点:通过 WebSocket tunnel 测试
|
||||
if not node.is_manual:
|
||||
from src.services.proxy_node.hub_config import get_hub_config
|
||||
|
||||
hub_enabled = get_hub_config().enabled
|
||||
connected = False
|
||||
if hub_enabled:
|
||||
connected = bool(node.tunnel_connected) and node.status == ProxyNodeStatus.ONLINE
|
||||
else:
|
||||
# 非 Hub 模式:以当前 worker 本地 TunnelManager 状态为准,
|
||||
# 避免多 worker 场景的跨进程状态误判。
|
||||
from src.services.proxy_node.tunnel_manager import get_tunnel_manager
|
||||
|
||||
manager = get_tunnel_manager()
|
||||
connected = manager.has_tunnel(node.id)
|
||||
connected = bool(node.tunnel_connected) and node.status == ProxyNodeStatus.ONLINE
|
||||
|
||||
if not connected:
|
||||
hint = ""
|
||||
try:
|
||||
from src.config import config
|
||||
|
||||
if config.worker_processes > 1 and not hub_enabled:
|
||||
hint = (
|
||||
"(当前 worker 无 tunnel 连接;检测到多 worker 部署,"
|
||||
"建议设置 GUNICORN_WORKERS/WEB_CONCURRENCY=1)"
|
||||
)
|
||||
except Exception:
|
||||
# 配置读取失败时保持原始错误,避免影响主流程
|
||||
hint = ""
|
||||
return {
|
||||
"success": False,
|
||||
"latency_ms": None,
|
||||
"exit_ip": None,
|
||||
"error": f"tunnel 未连接{hint}",
|
||||
"error": "tunnel 未连接",
|
||||
}
|
||||
result = await _test_tunnel_connectivity(node.id)
|
||||
|
||||
|
||||
@@ -1,135 +1,15 @@
|
||||
"""
|
||||
Tunnel httpx Transport
|
||||
|
||||
自定义 httpx AsyncBaseTransport,将 HTTP 请求通过 WebSocket tunnel 发送到 aether-proxy。
|
||||
对 handler 层完全透明 -- 只需在创建 httpx.AsyncClient 时使用此 transport。
|
||||
统一通过 aether-hub 转发 tunnel 请求。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from .tunnel_manager import TunnelManager, TunnelStreamError, _StreamState, get_tunnel_manager
|
||||
|
||||
_HOP_BY_HOP_HEADERS = frozenset(
|
||||
{
|
||||
"host",
|
||||
"transfer-encoding",
|
||||
"content-length",
|
||||
"connection",
|
||||
"upgrade",
|
||||
"keep-alive",
|
||||
"proxy-authorization",
|
||||
"proxy-connection",
|
||||
"te",
|
||||
"trailer",
|
||||
}
|
||||
)
|
||||
|
||||
# bytes 版本,用于直接比较 httpx raw headers(key 已经是小写 bytes)
|
||||
_HOP_BY_HOP_HEADERS_BYTES = frozenset(h.encode("ascii") for h in _HOP_BY_HOP_HEADERS)
|
||||
|
||||
|
||||
class TunnelTransport(httpx.AsyncBaseTransport):
|
||||
"""通过 WebSocket tunnel 发送请求的 httpx transport"""
|
||||
|
||||
def __init__(self, node_id: str, timeout: float = 60.0) -> None:
|
||||
self._node_id = node_id
|
||||
self._timeout = timeout
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
manager = get_tunnel_manager()
|
||||
|
||||
# 构建 headers dict(跳过 hop-by-hop 和 httpx 内部 headers)
|
||||
# request.headers.raw 返回 (bytes, bytes) 元组,key 已经是小写
|
||||
headers: dict[str, str] = {}
|
||||
for key, value in request.headers.raw:
|
||||
if key not in _HOP_BY_HOP_HEADERS_BYTES:
|
||||
headers[key.decode("latin-1")] = value.decode("latin-1")
|
||||
|
||||
# 读取 body -- request.content 在 json= 传参时已由 httpx 序列化好;
|
||||
# 对 stream 类型的 request 需要先 read() 才能拿到完整 content。
|
||||
body = request.content or await request.aread() or None
|
||||
|
||||
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)
|
||||
|
||||
# 构建 httpx.Response(流式 body)
|
||||
resp_headers = httpx.Headers(stream_state.headers)
|
||||
|
||||
return httpx.Response(
|
||||
status_code=stream_state.status,
|
||||
headers=resp_headers,
|
||||
stream=TunnelResponseStream(
|
||||
manager, self._node_id, stream_state, timeout=self._timeout
|
||||
),
|
||||
)
|
||||
|
||||
except TunnelStreamError as e:
|
||||
self._cleanup_stream(manager, stream_state)
|
||||
# 区分连接阶段和响应阶段的错误
|
||||
if stream_state and stream_state.status > 0:
|
||||
raise httpx.ReadError(str(e)) from e
|
||||
raise httpx.ConnectError(str(e)) from e
|
||||
except asyncio.TimeoutError:
|
||||
self._cleanup_stream(manager, stream_state)
|
||||
raise httpx.ReadTimeout("tunnel request timeout") from None
|
||||
|
||||
def _cleanup_stream(self, manager: TunnelManager, stream_state: _StreamState | None) -> None:
|
||||
if stream_state is None:
|
||||
return
|
||||
# 优先从 stream 记住的原始连接上移除,避免连接池竞态
|
||||
conn = stream_state._conn
|
||||
if conn is None:
|
||||
conn = manager.get_connection(self._node_id)
|
||||
if conn:
|
||||
conn.remove_stream(stream_state.stream_id)
|
||||
|
||||
|
||||
class TunnelResponseStream(httpx.AsyncByteStream):
|
||||
"""将 tunnel stream 的 body chunks 包装为 httpx AsyncByteStream"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
manager: TunnelManager,
|
||||
node_id: str,
|
||||
stream_state: _StreamState,
|
||||
timeout: float = 60.0,
|
||||
) -> None:
|
||||
self._manager = manager
|
||||
self._node_id = node_id
|
||||
self._stream_state = stream_state
|
||||
self._timeout = timeout
|
||||
|
||||
async def __aiter__(self) -> AsyncGenerator[bytes, None]:
|
||||
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
|
||||
yield chunk
|
||||
|
||||
async def aclose(self) -> None:
|
||||
# 从 stream 记住的原始连接上精确移除,避免连接池竞态
|
||||
conn = self._stream_state._conn
|
||||
if conn is None:
|
||||
conn = self._manager.get_connection(self._node_id)
|
||||
if conn:
|
||||
conn.remove_stream(self._stream_state.stream_id)
|
||||
|
||||
|
||||
def is_tunnel_node(node_info: dict[str, Any] | None) -> bool:
|
||||
"""检查节点是否为 tunnel 模式且已连接"""
|
||||
@@ -139,12 +19,7 @@ def is_tunnel_node(node_info: dict[str, Any] | None) -> bool:
|
||||
|
||||
|
||||
def create_tunnel_transport(node_id: str, timeout: float = 60.0) -> httpx.AsyncBaseTransport:
|
||||
"""根据配置创建 tunnel transport(Hub 模式或直连 tunnel 模式)。"""
|
||||
from .hub_config import get_hub_config
|
||||
"""创建统一的 Hub tunnel transport。"""
|
||||
from .hub_transport import HubTunnelTransport
|
||||
|
||||
hub_cfg = get_hub_config()
|
||||
if hub_cfg.enabled:
|
||||
from .hub_transport import HubTunnelTransport
|
||||
|
||||
return HubTunnelTransport(node_id, timeout=timeout)
|
||||
return TunnelTransport(node_id, timeout=timeout)
|
||||
return HubTunnelTransport(node_id, timeout=timeout)
|
||||
|
||||
@@ -1,29 +1,9 @@
|
||||
import pytest
|
||||
|
||||
from src.services.proxy_node.hub_config import reset_hub_config_cache
|
||||
from src.services.proxy_node.hub_transport import HubTunnelTransport
|
||||
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
|
||||
from src.services.proxy_node.tunnel_transport import TunnelTransport, create_tunnel_transport
|
||||
from src.services.proxy_node.tunnel_transport import create_tunnel_transport
|
||||
|
||||
|
||||
def test_create_tunnel_transport_uses_legacy_transport_when_hub_disabled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("DOCKER_CONTAINER", "false")
|
||||
monkeypatch.setattr("src.services.proxy_node.hub_config.os.path.exists", lambda _: False)
|
||||
reset_hub_config_cache()
|
||||
|
||||
transport = create_tunnel_transport("node-1", timeout=12.0)
|
||||
assert isinstance(transport, TunnelTransport)
|
||||
|
||||
|
||||
def test_create_tunnel_transport_uses_hub_transport_when_hub_enabled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("DOCKER_CONTAINER", "true")
|
||||
monkeypatch.setattr("src.services.proxy_node.hub_config.os.path.exists", lambda _: False)
|
||||
reset_hub_config_cache()
|
||||
|
||||
def test_create_tunnel_transport_always_uses_hub_transport() -> None:
|
||||
transport = create_tunnel_transport("node-1", timeout=12.0)
|
||||
assert isinstance(transport, HubTunnelTransport)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user