51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
import websockets
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ProbeResult:
|
||
|
|
url: str
|
||
|
|
connected: bool
|
||
|
|
latency_ms: float | None = None
|
||
|
|
error: str | None = None
|
||
|
|
|
||
|
|
|
||
|
|
async def probe_relay(url: str, timeout: float = 10.0) -> ProbeResult:
|
||
|
|
start = asyncio.get_event_loop().time()
|
||
|
|
try:
|
||
|
|
ws = await asyncio.wait_for(
|
||
|
|
websockets.connect(url, ping_interval=20, ping_timeout=10),
|
||
|
|
timeout=timeout,
|
||
|
|
)
|
||
|
|
latency_ms = (asyncio.get_event_loop().time() - start) * 1000
|
||
|
|
sub_id = "forcepilot_probe"
|
||
|
|
req = json.dumps(["REQ", sub_id, {"kinds": [1], "limit": 1}])
|
||
|
|
await asyncio.wait_for(ws.send(req), timeout=5)
|
||
|
|
eose_received = False
|
||
|
|
try:
|
||
|
|
while not eose_received:
|
||
|
|
raw = await asyncio.wait_for(ws.recv(), timeout=5)
|
||
|
|
data = json.loads(raw)
|
||
|
|
if isinstance(data, list) and len(data) >= 2 and data[0] == "EOSE":
|
||
|
|
eose_received = True
|
||
|
|
except TimeoutError:
|
||
|
|
pass
|
||
|
|
close_req = json.dumps(["CLOSE", sub_id])
|
||
|
|
try:
|
||
|
|
await ws.send(close_req)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
await ws.close()
|
||
|
|
return ProbeResult(url=url, connected=True, latency_ms=latency_ms)
|
||
|
|
except Exception as e:
|
||
|
|
latency_ms = (asyncio.get_event_loop().time() - start) * 1000
|
||
|
|
logger.debug(f"Relay 探测失败 {url}: {e}")
|
||
|
|
return ProbeResult(url=url, connected=False, latency_ms=latency_ms, error=str(e))
|