Kris
7002557cf4
1. 整理多个文件的导入顺序，移除冗余空行和重复导入 2. 为Nostr订阅添加until参数防止重复拉取 3. 实现长消息分块发送并添加间隔等待 4. 新增Relay健康分数同步任务 5. 新增TLS强制检查配置项并支持从环境变量加载 6. 重构状态恢复逻辑适配异步存储 7. 修复反应发送的参数错误 8. 新增线程模拟自动补全消息上下文 9. 添加解密指标统计和更完善的错误日志
402 lines
16 KiB
Python
402 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import random
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
import websockets
|
|
|
|
from yuxi.channels.infra.circuit_breaker import CircuitBreaker, CircuitBreakerOpenError
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
|
|
class RelayManager:
|
|
def __init__(self, relay_urls: list[str], timeout: float = 30.0, reconnect_interval: float = 5.0):
|
|
self._relay_urls = relay_urls
|
|
self._timeout = timeout
|
|
self._reconnect_interval = reconnect_interval
|
|
self._connections: dict[str, websockets.WebSocketClientProtocol] = {}
|
|
self._event_handlers: list[Callable[[dict], Awaitable[None]]] = []
|
|
self._eose_handlers: list[Callable[[str], Awaitable[None]]] = []
|
|
self._reconnect_tasks: dict[str, asyncio.Task] = {}
|
|
self._subscriptions: list[dict] = []
|
|
self._running = False
|
|
self._circuit_breakers: dict[str, CircuitBreaker] = {}
|
|
self._on_connect_handlers: list[Callable[[str], Awaitable[None]]] = []
|
|
self._on_disconnect_handlers: list[Callable[[str], Awaitable[None]]] = []
|
|
self._relay_scores: dict[str, float] = {}
|
|
self._on_error_handlers: list[Callable[[str, dict], Awaitable[None]]] = []
|
|
self._reconnect_attempts: dict[str, int] = {}
|
|
|
|
for url in relay_urls:
|
|
self._circuit_breakers[url] = CircuitBreaker(
|
|
failure_threshold=5, recovery_timeout=30.0, half_open_max_calls=3
|
|
)
|
|
self._relay_scores[url] = 0.5
|
|
|
|
self._sub_pubkeys: list[str] = []
|
|
|
|
def set_pubkey_filter(self, pubkeys: list[str]) -> None:
|
|
self._sub_pubkeys = pubkeys
|
|
|
|
def _build_subscription_filters(self, filters: list[dict]) -> list[dict]:
|
|
if not self._sub_pubkeys:
|
|
return filters
|
|
return [{**f, "#p": self._sub_pubkeys} for f in filters]
|
|
|
|
def on_connect(self, handler: Callable[[str], Awaitable[None]]) -> None:
|
|
self._on_connect_handlers.append(handler)
|
|
|
|
def on_disconnect(self, handler: Callable[[str], Awaitable[None]]) -> None:
|
|
self._on_disconnect_handlers.append(handler)
|
|
|
|
def on_eose(self, handler: Callable[[str], Awaitable[None]]) -> None:
|
|
self._eose_handlers.append(handler)
|
|
|
|
def on_error(self, handler: Callable[[str, dict], Awaitable[None]]) -> None:
|
|
self._on_error_handlers.append(handler)
|
|
|
|
def _classify_error(self, url: str, error: Exception) -> str:
|
|
error_type = type(error).__name__
|
|
if "timeout" in error_type.lower() or "Timeout" in error_type:
|
|
return "timeout"
|
|
if "connect" in error_type.lower() or "Connection" in error_type:
|
|
return "connection"
|
|
if "circuit" in error_type.lower():
|
|
return "circuit_breaker"
|
|
if "json" in error_type.lower():
|
|
return "parse_error"
|
|
if "websocket" in error_type.lower():
|
|
return "websocket"
|
|
return "unknown"
|
|
|
|
async def _notify_error(self, url: str, category: str, error: Exception) -> None:
|
|
context = {
|
|
"url": url,
|
|
"category": category,
|
|
"error_type": type(error).__name__,
|
|
"error_message": str(error)[:200],
|
|
}
|
|
for handler in self._on_error_handlers:
|
|
try:
|
|
await handler(category, context)
|
|
except Exception:
|
|
logger.debug(f"on_error handler error for {url}", exc_info=True)
|
|
|
|
async def connect_all(self) -> None:
|
|
self._running = True
|
|
tasks = [self._connect_relay(url) for url in self._relay_urls]
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
connected = sum(1 for r in results if r is True)
|
|
logger.info(f"Nostr Relay 连接完成: {connected}/{len(self._relay_urls)} 个 Relay 已连接")
|
|
|
|
async def _connect_relay(self, url: str) -> bool:
|
|
try:
|
|
ws = await asyncio.wait_for(
|
|
websockets.connect(url, ping_interval=20, ping_timeout=10),
|
|
timeout=self._timeout,
|
|
)
|
|
self._connections[url] = ws
|
|
if self._subscriptions:
|
|
sub_id = "forcepilot_nostr_sub"
|
|
effective_filters = self._build_subscription_filters(self._subscriptions)
|
|
req = json.dumps(["REQ", sub_id, *effective_filters])
|
|
await ws.send(req)
|
|
for handler in self._on_connect_handlers:
|
|
try:
|
|
await handler(url)
|
|
except Exception:
|
|
logger.debug(f"on_connect handler error for {url}", exc_info=True)
|
|
logger.info(f"Nostr Relay 已连接: {url}")
|
|
return True
|
|
except Exception as e:
|
|
logger.warning(f"Nostr Relay 连接失败 {url}: {e}")
|
|
self._schedule_reconnect(url)
|
|
return False
|
|
|
|
def _schedule_reconnect(self, url: str) -> None:
|
|
if url in self._reconnect_tasks:
|
|
return
|
|
|
|
attempt = self._reconnect_attempts.get(url, 0) + 1
|
|
self._reconnect_attempts[url] = attempt
|
|
base_delay = min(self._reconnect_interval * (2 ** (attempt - 1)), 300)
|
|
jitter = random.uniform(0, base_delay * 0.3)
|
|
delay = base_delay + jitter
|
|
|
|
async def _reconnect():
|
|
await asyncio.sleep(delay)
|
|
if not self._running:
|
|
self._reconnect_tasks.pop(url, None)
|
|
return
|
|
logger.debug(f"Nostr Relay 尝试重连: {url} (尝试 #{attempt}, 延迟 {delay:.1f}s)")
|
|
success = await self._connect_relay(url)
|
|
if success:
|
|
self._reconnect_attempts.pop(url, None)
|
|
self._reconnect_tasks.pop(url, None)
|
|
|
|
self._reconnect_tasks[url] = asyncio.ensure_future(_reconnect())
|
|
|
|
def list_connections(self) -> dict[str, websockets.WebSocketClientProtocol]:
|
|
return dict(self._connections)
|
|
|
|
def get_circuit_breaker(self, url: str) -> CircuitBreaker | None:
|
|
return self._circuit_breakers.get(url)
|
|
|
|
async def disconnect_all(self) -> None:
|
|
self._running = False
|
|
for url, task in list(self._reconnect_tasks.items()):
|
|
task.cancel()
|
|
self._reconnect_tasks.pop(url, None)
|
|
for url in list(self._connections.keys()):
|
|
ws = self._connections.pop(url, None)
|
|
if ws:
|
|
try:
|
|
await ws.close()
|
|
except Exception:
|
|
pass
|
|
for handler in self._on_disconnect_handlers:
|
|
try:
|
|
await handler(url)
|
|
except Exception:
|
|
logger.debug(f"on_disconnect handler error for {url}", exc_info=True)
|
|
|
|
async def broadcast(self, event: dict) -> int:
|
|
message = json.dumps(["EVENT", event])
|
|
sorted_urls = sorted(
|
|
self._connections.keys(),
|
|
key=lambda url: self._relay_scores.get(url, 0.5),
|
|
reverse=True,
|
|
)
|
|
tasks = []
|
|
for url in sorted_urls:
|
|
cb = self._circuit_breakers.get(url)
|
|
if cb and cb.state == "open":
|
|
continue
|
|
ws = self._connections.get(url)
|
|
if ws and ws.open:
|
|
tasks.append(self._send_with_cb(ws, url, message))
|
|
if not tasks:
|
|
return 0
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
return sum(1 for r in results if r is True)
|
|
|
|
async def broadcast_with_ack(self, event: dict) -> dict:
|
|
message = json.dumps(["EVENT", event])
|
|
sorted_urls = sorted(
|
|
self._connections.keys(),
|
|
key=lambda url: self._relay_scores.get(url, 0.5),
|
|
reverse=True,
|
|
)
|
|
ack_results: dict[str, bool] = {}
|
|
tasks = []
|
|
for url in sorted_urls:
|
|
cb = self._circuit_breakers.get(url)
|
|
if cb and cb.state == "open":
|
|
ack_results[url] = False
|
|
continue
|
|
ws = self._connections.get(url)
|
|
if ws and ws.open:
|
|
tasks.append(self._send_and_wait_ok(ws, url, message, ack_results))
|
|
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
return ack_results
|
|
|
|
async def _send_and_wait_ok(self, ws, url: str, message: str, results: dict[str, bool]) -> None:
|
|
try:
|
|
await ws.send(message)
|
|
raw = await asyncio.wait_for(ws.recv(), timeout=5)
|
|
data = json.loads(raw)
|
|
ok = isinstance(data, list) and len(data) >= 4 and data[0] == "OK" and data[2] is True
|
|
results[url] = ok
|
|
except Exception:
|
|
results[url] = False
|
|
|
|
async def _send_with_cb(
|
|
self, ws: websockets.WebSocketClientProtocol, url: str, message: str, max_retries: int = 3
|
|
) -> bool:
|
|
cb = self._circuit_breakers.get(url)
|
|
retry_delays = [1, 2, 4]
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
if cb:
|
|
await cb.call(lambda: asyncio.wait_for(ws.send(message), timeout=10))
|
|
else:
|
|
await asyncio.wait_for(ws.send(message), timeout=10)
|
|
return True
|
|
except CircuitBreakerOpenError:
|
|
logger.debug(f"Relay {url} circuit breaker open, skipping")
|
|
return False
|
|
except Exception as e:
|
|
if cb:
|
|
await cb.record_failure()
|
|
category = self._classify_error(url, e)
|
|
await self._notify_error(url, category, e)
|
|
if attempt < max_retries - 1:
|
|
await asyncio.sleep(retry_delays[min(attempt, len(retry_delays) - 1)])
|
|
logger.warning(f"Relay 发送失败(已重试{max_retries}次): {url}，断开并调度重连")
|
|
self._connections.pop(url, None)
|
|
self._schedule_reconnect(url)
|
|
return False
|
|
|
|
async def _send_with_retry(
|
|
self, ws: websockets.WebSocketClientProtocol, url: str, message: str, max_retries: int = 3
|
|
) -> bool:
|
|
return await self._send_with_cb(ws, url, message, max_retries)
|
|
|
|
def on_event(self, handler: Callable[[dict], Awaitable[None]]) -> None:
|
|
self._event_handlers.append(handler)
|
|
|
|
async def subscribe(self, filters: list[dict]) -> None:
|
|
self._subscriptions = filters
|
|
sub_id = "forcepilot_nostr_sub"
|
|
effective_filters = self._build_subscription_filters(filters)
|
|
req = json.dumps(["REQ", sub_id, *effective_filters])
|
|
for url in list(self._connections.keys()):
|
|
ws = self._connections.get(url)
|
|
if ws and ws.open:
|
|
try:
|
|
await ws.send(req)
|
|
except Exception:
|
|
pass
|
|
|
|
async def listen(self, url: str) -> None:
|
|
ws = self._connections.get(url)
|
|
if not ws or not ws.open:
|
|
return
|
|
try:
|
|
async for raw in ws:
|
|
if not self._running:
|
|
break
|
|
try:
|
|
data = json.loads(raw)
|
|
if isinstance(data, list) and len(data) >= 3 and data[0] == "EVENT":
|
|
event = data[2]
|
|
if isinstance(event, dict):
|
|
for handler in self._event_handlers:
|
|
await handler(event)
|
|
elif isinstance(data, list) and len(data) >= 2 and data[0] == "EOSE":
|
|
for handler in self._eose_handlers:
|
|
try:
|
|
await handler(url)
|
|
except Exception:
|
|
logger.debug(f"EOSE handler error for {url}", exc_info=True)
|
|
except (json.JSONDecodeError, KeyError, IndexError):
|
|
logger.debug(f"Relay {url} 返回无法解析的消息: {raw[:200]}")
|
|
continue
|
|
except Exception:
|
|
self._connections.pop(url, None)
|
|
for handler in self._on_disconnect_handlers:
|
|
try:
|
|
await handler(url)
|
|
except Exception:
|
|
pass
|
|
self._schedule_reconnect(url)
|
|
|
|
async def query(
|
|
self, filters: list[dict], timeout: float = 10.0, since: int | None = None, until: int | None = None
|
|
) -> list[dict]:
|
|
sub_id = "forcepilot_nostr_query"
|
|
effective_filters: list[dict] = []
|
|
for f in filters:
|
|
f_copy = dict(f)
|
|
if since is not None:
|
|
f_copy["since"] = since
|
|
if until is not None:
|
|
f_copy["until"] = until
|
|
effective_filters.append(f_copy)
|
|
req = json.dumps(["REQ", sub_id, *effective_filters])
|
|
close_req = json.dumps(["CLOSE", sub_id])
|
|
|
|
active_ws = [(url, ws) for url, ws in self._connections.items() if ws and ws.open]
|
|
if not active_ws:
|
|
return []
|
|
|
|
for url, ws in active_ws:
|
|
try:
|
|
await ws.send(req)
|
|
except Exception:
|
|
continue
|
|
|
|
async def _receive_from_relay(_url: str, _ws) -> list[dict]:
|
|
results: list[dict] = []
|
|
try:
|
|
while True:
|
|
raw = await asyncio.wait_for(_ws.recv(), timeout=timeout)
|
|
try:
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
if isinstance(data, list) and len(data) >= 2:
|
|
if data[0] == "EVENT" and len(data) >= 3 and isinstance(data[2], dict):
|
|
results.append(data[2])
|
|
elif data[0] in ("EOSE", "CLOSED"):
|
|
break
|
|
except TimeoutError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
try:
|
|
await _ws.send(close_req)
|
|
except Exception:
|
|
pass
|
|
return results
|
|
|
|
receive_tasks = [_receive_from_relay(url, ws) for url, ws in active_ws]
|
|
all_results = await asyncio.gather(*receive_tasks)
|
|
|
|
seen: set[str] = set()
|
|
events: list[dict] = []
|
|
for batch in all_results:
|
|
for event in batch:
|
|
event_id = event.get("id", "")
|
|
if event_id and event_id not in seen:
|
|
seen.add(event_id)
|
|
events.append(event)
|
|
|
|
return events
|
|
|
|
async def send_auth(self, raw_signed_auth_event: dict, urls: list[str] | None = None) -> dict[str, bool]:
|
|
message = json.dumps(["AUTH", raw_signed_auth_event])
|
|
target_urls = urls or list(self._connections.keys())
|
|
results: dict[str, bool] = {}
|
|
for url in target_urls:
|
|
ws = self._connections.get(url)
|
|
if not ws or not ws.open:
|
|
results[url] = False
|
|
continue
|
|
try:
|
|
await ws.send(message)
|
|
raw = await asyncio.wait_for(ws.recv(), timeout=10)
|
|
data = json.loads(raw)
|
|
results[url] = (isinstance(data, list) and len(data) >= 2 and data[0] == "AUTH" and data[1] == url) or (
|
|
isinstance(data, list) and len(data) >= 2 and data[0] == "OK"
|
|
)
|
|
except Exception:
|
|
results[url] = False
|
|
return results
|
|
|
|
def active_count(self) -> tuple[int, int]:
|
|
active = sum(1 for ws in self._connections.values() if ws.open)
|
|
return active, len(self._relay_urls)
|
|
|
|
def set_relay_scores(self, scores: dict[str, float]) -> None:
|
|
self._relay_scores.update(scores)
|
|
|
|
def sync_scores_from_health_tracker(self, health_tracker, urls: list[str] | None = None) -> None:
|
|
target_urls = urls or list(self._relay_urls)
|
|
for url in target_urls:
|
|
score = health_tracker.get_score(url)
|
|
self._relay_scores[url] = score
|
|
|
|
def get_relay_score(self, url: str) -> float:
|
|
return self._relay_scores.get(url, 0.5)
|
|
|
|
async def query_event(self, event_id: str, timeout: float = 10.0) -> dict | None:
|
|
events = await self.query([{"ids": [event_id], "limit": 1}], timeout=timeout)
|
|
return events[0] if events else None
|