"""Worker 健康池 + 空闲派发。 - 后台 asyncio 任务周期性探测每个 worker 的 /health - 连续失败 N 次 → 下线；连续成功 N 次 → 上线 - 每个 worker 一个 busy 标志（per_worker_concurrency=1） - acquire_worker() 从在线池挑空闲 worker；全忙排队；池空→NoWorkerAvailable - release_worker() 释放 worker """ import asyncio import logging import time from dataclasses import dataclass, field from typing import Dict, List, Optional import httpx logger = logging.getLogger("gateway.pool") # --------------------------------------------------------------------------- # 异常 # --------------------------------------------------------------------------- class NoWorkerAvailable(Exception): """无可用 worker（池空或全忙排队超时）。""" pass # --------------------------------------------------------------------------- # 数据结构 # --------------------------------------------------------------------------- @dataclass class WorkerState: """单个 worker 的运行时状态。""" url: str online: bool = False # 初始 offline，等健康检查通过后上线 busy: bool = False consecutive_failures: int = 0 consecutive_successes: int = 0 # --------------------------------------------------------------------------- # 全局状态 # --------------------------------------------------------------------------- _workers: Dict[str, WorkerState] = {} _pool_condition: Optional[asyncio.Condition] = None _health_task: Optional[asyncio.Task] = None _shutdown_event: Optional[asyncio.Event] = None # --------------------------------------------------------------------------- # 健康检查后台任务 # --------------------------------------------------------------------------- async def _check_worker_health( client: httpx.AsyncClient, w: WorkerState, cfg: dict, ) -> None: """探测单个 worker 的 /health，更新上下线状态。""" hc_cfg = cfg["health_check"] url = f"{w.url}{hc_cfg['path']}" token = cfg["shared_password"] try: resp = await client.get( url, headers={"X-Internal-Token": token}, timeout=hc_cfg["timeout_seconds"], ) if resp.status_code == 200: w.consecutive_failures = 0 w.consecutive_successes += 1 if w.consecutive_successes >= hc_cfg["healthy_threshold"]: if not w.online: w.online = True logger.info("Worker 上线: %s（连续成功 %d 次）", w.url, w.consecutive_successes) else: _mark_failure(w, f"HTTP {resp.status_code}") except Exception as exc: _mark_failure(w, str(exc)) def _mark_failure(w: WorkerState, reason: str) -> None: """记录一次失败，达到阈值后下线。""" w.consecutive_successes = 0 w.consecutive_failures += 1 threshold = w.consecutive_failures # used in log hc_threshold = 2 # default, will be overridden if w.consecutive_failures >= hc_threshold: # 实际 threshold 从配置读取，这里先做基本判断 pass logger.debug("Worker %s 健康检查失败 (%d/%d): %s", w.url, w.consecutive_failures, 99, reason) async def _health_check_loop(cfg: dict) -> None: """后台循环：周期性探测所有 worker 健康状态。""" hc_cfg = cfg["health_check"] interval = hc_cfg["interval_seconds"] unhealthy_threshold = hc_cfg["unhealthy_threshold"] healthy_threshold = hc_cfg["healthy_threshold"] token = cfg["shared_password"] logger.info( "健康检查循环启动 | 间隔=%ds | 下线阈值=%d | 上线阈值=%d | workers=%d", interval, unhealthy_threshold, healthy_threshold, len(_workers), ) async with httpx.AsyncClient() as client: while not _shutdown_event.is_set(): for w in _workers.values(): url = f"{w.url}{hc_cfg['path']}" try: resp = await client.get( url, headers={"X-Internal-Token": token}, timeout=hc_cfg["timeout_seconds"], ) if resp.status_code == 200: w.consecutive_failures = 0 w.consecutive_successes += 1 if w.consecutive_successes >= healthy_threshold and not w.online: w.online = True logger.info("✅ Worker 上线: %s", w.url) else: w.consecutive_successes = 0 w.consecutive_failures += 1 if w.consecutive_failures >= unhealthy_threshold and w.online: w.online = False logger.warning("⚠ Worker 下线: %s（HTTP %d，连续失败 %d 次）", w.url, resp.status_code, w.consecutive_failures) except Exception as exc: w.consecutive_successes = 0 w.consecutive_failures += 1 if w.consecutive_failures >= unhealthy_threshold and w.online: w.online = False logger.warning("⚠ Worker 下线: %s（%s，连续失败 %d 次）", w.url, exc, w.consecutive_failures) # 等待下一次探测（支持快速关闭） try: await asyncio.wait_for(_shutdown_event.wait(), timeout=interval) break # shutdown signaled except asyncio.TimeoutError: pass # 正常的 interval 到期 logger.info("健康检查循环已停止") # --------------------------------------------------------------------------- # 初始化 / 关闭 # --------------------------------------------------------------------------- async def init_pool(cfg: dict) -> None: """初始化 worker 池并启动健康检查后台任务。""" global _workers, _pool_condition, _health_task, _shutdown_event _workers = {url: WorkerState(url=url) for url in cfg["workers"]} _pool_condition = asyncio.Condition() _shutdown_event = asyncio.Event() # 立即做一轮健康检查以快速上线 hc_cfg = cfg["health_check"] token = cfg["shared_password"] async with httpx.AsyncClient() as client: for w in _workers.values(): try: resp = await client.get( f"{w.url}{hc_cfg['path']}", headers={"X-Internal-Token": token}, timeout=hc_cfg["timeout_seconds"], ) if resp.status_code == 200: w.consecutive_successes = 1 if hc_cfg["healthy_threshold"] <= 1: w.online = True else: w.consecutive_failures = 1 except Exception: w.consecutive_failures = 1 # 对于 healthy_threshold == 1 的情况，已在上面的检查中上线 # 标记初始在线状态 online_count = sum(1 for w in _workers.values() if w.online) logger.info("Worker 池初始化完成 | 总数=%d | 在线=%d", len(_workers), online_count) # 启动后台健康检查 _health_task = asyncio.create_task(_health_check_loop(cfg)) async def shutdown_pool() -> None: """关闭健康检查任务，释放资源。""" global _shutdown_event, _health_task if _shutdown_event: _shutdown_event.set() if _health_task: _health_task.cancel() try: await _health_task except asyncio.CancelledError: pass _health_task = None logger.info("Worker 池已关闭") # --------------------------------------------------------------------------- # 派发 # --------------------------------------------------------------------------- async def mark_worker_unhealthy(w: WorkerState) -> None: """被动标记：转发失败时立即将该 worker 下线。""" async with _pool_condition: if w.online: w.online = False w.consecutive_failures += 1 logger.warning("⚠ Worker 被动下线: %s（转发失败）", w.url) async def acquire_worker(cfg: dict) -> WorkerState: """从在线池中获取一个空闲 worker。 全忙则排队等待，最多 queue_wait_seconds；池空立即抛异常。 """ queue_wait = cfg["dispatch"]["queue_wait_seconds"] deadline = time.monotonic() + queue_wait async with _pool_condition: while True: # 在线 + 空闲 idle = [w for w in _workers.values() if w.online and not w.busy] if not idle: online = [w for w in _workers.values() if w.online] if not online: raise NoWorkerAvailable("后端服务暂不可用，请稍后重试") # 全忙，等待 remaining = deadline - time.monotonic() if remaining <= 0: raise NoWorkerAvailable("后端服务繁忙，请稍后重试") logger.info("所有 worker 全忙（%d 在线），等待 %.1fs...", len(online), remaining) try: await asyncio.wait_for(_pool_condition.wait(), timeout=remaining) except asyncio.TimeoutError: raise NoWorkerAvailable("后端服务繁忙，请稍后重试") continue # 重新检查 # 取第一个空闲 worker w = idle[0] w.busy = True logger.debug("Worker 分配: %s", w.url) return w async def release_worker(w: WorkerState) -> None: """释放 worker，标记为空闲并通知等待者。""" async with _pool_condition: w.busy = False logger.debug("Worker 释放: %s", w.url) _pool_condition.notify(1) # --------------------------------------------------------------------------- # 状态查询 # --------------------------------------------------------------------------- def get_pool_status() -> dict: """返回当前池状态（供 /gateway-health 使用）。""" if not _workers: return {"total": 0, "healthy": 0, "busy": 0} total = len(_workers) healthy = sum(1 for w in _workers.values() if w.online) busy = sum(1 for w in _workers.values() if w.busy) return {"total": total, "healthy": healthy, "busy": busy}