"""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") def _cooldown_seconds() -> float: """被动摘除后的隔离期时长（秒），配置 health_check.cooldown_seconds，默认 60。""" try: from gateway.config import get_config return float(get_config()["health_check"].get("cooldown_seconds", 60)) except Exception: # noqa: BLE001 return 60.0 # --------------------------------------------------------------------------- # 异常 # --------------------------------------------------------------------------- class NoWorkerAvailable(Exception): """无可用 worker（池空或全忙排队超时）。""" pass # --------------------------------------------------------------------------- # 数据结构 # --------------------------------------------------------------------------- @dataclass class WorkerState: """单个 worker 的运行时状态。""" url: str online: bool = False # 初始 offline，等健康检查通过后上线 busy: bool = False primary: bool = True # 主 worker（True 优先派发）；全部主 worker 下线才用备组 consecutive_failures: int = 0 consecutive_successes: int = 0 offline_until: float = 0.0 # 被动摘除后的隔离截止时刻（time.monotonic() 基准） # --------------------------------------------------------------------------- # 全局状态 # --------------------------------------------------------------------------- _workers: Dict[str, WorkerState] = {} _pool_condition: Optional[asyncio.Condition] = None _health_task: Optional[asyncio.Task] = None _shutdown_event: Optional[asyncio.Event] = None # 全局串行槽：同一时间最多 max_global_concurrency 个请求进入 worker 派发，其余排队。 # ⚠️ 依赖单进程 uvicorn 部署（hair-gateway.service 无 --workers）。 # 若改多 worker / gunicorn，进程内信号量会失效，需换成跨进程锁（文件锁/Redis）。 _global_sem: Optional[asyncio.Semaphore] = None _global_max: int = 1 _global_busy: int = 0 _global_waiting: int = 0 # --------------------------------------------------------------------------- # 健康检查后台任务 # --------------------------------------------------------------------------- def _try_revive(w: WorkerState, healthy_threshold: int) -> bool: """探测成功后是否允许将该 worker 复活上线。 条件：未在线 + 连续成功达到阈值 + **不在被动摘除隔离期内** （offline_until 之前即使 /health 探测成功也压着不上线，防横跳）。 """ if w.online or w.consecutive_successes < healthy_threshold: return False if time.monotonic() < w.offline_until: logger.debug("Worker %s 探测成功但处于隔离期，暂不复活", w.url) return False return True 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 | 隔离期=%.0fs | workers=%d", interval, unhealthy_threshold, healthy_threshold, _cooldown_seconds(), 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 _try_revive(w, healthy_threshold): 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 # 主 worker 分组：配置了 primary_workers 则按其标记，否则全部视为主 worker primary_set = set(cfg.get("primary_workers") or cfg["workers"]) _workers = {url: WorkerState(url=url, primary=url in primary_set) for url in cfg["workers"]} n_primary = sum(1 for w in _workers.values() if w.primary) _pool_condition = asyncio.Condition() _shutdown_event = asyncio.Event() # 全局串行槽（独立于 worker busy 标志） init_global_slot(cfg) # 立即做一轮健康检查以快速上线 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 | 主worker=%d", len(_workers), online_count, n_primary) # 启动后台健康检查 _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 池已关闭") # --------------------------------------------------------------------------- # 全局串行槽 # --------------------------------------------------------------------------- def init_global_slot(cfg: dict) -> None: """初始化全局串行信号量。 max_global_concurrency=1 → 同一时间只有一个请求进入 worker 派发，其余在信号量上排队。 依赖单进程 uvicorn 部署；多进程需换跨进程锁。 """ global _global_sem, _global_max, _global_busy, _global_waiting n = int(cfg.get("dispatch", {}).get("max_global_concurrency", 1)) if n < 1: n = 1 _global_sem = asyncio.Semaphore(n) _global_max = n _global_busy = 0 _global_waiting = 0 logger.info("全局并发槽初始化: max=%d（单进程生效）", n) async def acquire_global_slot(cfg: dict) -> bool: """获取一个全局处理槽。True=获得；False=排队超时。 超时由 dispatch.max_queue_wait_seconds 控制（需 < nginx proxy_read_timeout 600s）。 """ global _global_busy, _global_waiting if _global_sem is None: return True # 未初始化，不限流 timeout = float(cfg.get("dispatch", {}).get("max_queue_wait_seconds", 590)) _global_waiting += 1 try: await asyncio.wait_for(_global_sem.acquire(), timeout=timeout) _global_busy += 1 return True except asyncio.TimeoutError: logger.warning("全局排队超时（%.1fs），拒绝请求 | waiting=%d", timeout, _global_waiting) return False finally: _global_waiting -= 1 def release_global_slot() -> None: """释放全局处理槽。""" global _global_busy if _global_sem is not None: _global_busy -= 1 _global_sem.release() # --------------------------------------------------------------------------- # 派发 # --------------------------------------------------------------------------- async def mark_worker_unhealthy(w: WorkerState) -> None: """被动标记：转发失败时立即将该 worker 下线，并进入隔离期。 隔离期内即使 /health 探测成功也不复活（由 _try_revive 守卫）， 避免"摘除→8 秒复活→再失败"的横跳。 """ async with _pool_condition: if w.online: w.online = False w.consecutive_failures += 1 cooldown = _cooldown_seconds() w.offline_until = time.monotonic() + cooldown logger.warning("⚠ Worker 被动下线: %s（转发失败，隔离 %.0fs）", w.url, cooldown) async def acquire_worker(cfg: dict) -> WorkerState: """从在线池中获取一个空闲 worker。 主备分组：只要还有主 worker（primary）在线，就只在主组内挑空闲 （全忙则排队等待，不用备组）；只有主组全部下线时才用备用 worker。 全忙则排队等待，最多 queue_wait_seconds；池空立即抛异常。 """ queue_wait = cfg["dispatch"]["queue_wait_seconds"] deadline = time.monotonic() + queue_wait async with _pool_condition: while True: online = [w for w in _workers.values() if w.online] primary_online = [w for w in online if w.primary] # 主组有在线 → 只在主组内挑；主组全挂 → 才轮到备组 candidates = primary_online if primary_online else online idle = [w for w in candidates if not w.busy] if not idle: if not online: raise NoWorkerAvailable("后端服务暂不可用，请稍后重试") # 全忙，等待 remaining = deadline - time.monotonic() if remaining <= 0: raise NoWorkerAvailable("后端服务繁忙，请稍后重试") logger.info("所有可用 worker 全忙（%d 在线），等待 %.1fs...", len(candidates), remaining) try: await asyncio.wait_for(_pool_condition.wait(), timeout=remaining) except asyncio.TimeoutError: raise NoWorkerAvailable("后端服务繁忙，请稍后重试") continue # 重新检查 # 取第一个空闲 worker w = idle[0] w.busy = True if not w.primary: logger.info("主 worker 全部下线，使用备用 worker: %s", w.url) 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 使用）。""" global_info = { "global_max": _global_max, "global_busy": _global_busy, "global_waiting": _global_waiting, } if not _workers: return {"total": 0, "healthy": 0, "busy": 0, "primary_total": 0, "primary_healthy": 0, **global_info} 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) primary_total = sum(1 for w in _workers.values() if w.primary) primary_healthy = sum(1 for w in _workers.values() if w.primary and w.online) return {"total": total, "healthy": healthy, "busy": busy, "primary_total": primary_total, "primary_healthy": primary_healthy, **global_info}