Files
hair/gateway/pool.py
T
xsl 6b24bf0b0d feat(gateway): 被动摘除后进入 60s 隔离期,防探测复活横跳
落地 docs/网关摘除隔离期实施方案.md(9·19 worker 磁盘满事故第 3 层防线):
- pool.py: WorkerState.offline_until + _cooldown_seconds();mark_worker_unhealthy
  摘除时设隔离截止时刻,隔离期内探测成功也不复活(_try_revive 守卫)
- pool.py: 删除从未被调用的死代码 _check_worker_health/_mark_failure
- config.py + config.example.json: health_check.cooldown_seconds(默认 60,0=关闭)
- tests/test_gateway_pool.py: 9 用例(摘除设隔离/隔离内不复活/期满复活/整链路等)

网关机已重启验证:启动日志含"隔离期=60s",/gateway-health 正常。
2026-09-20 23:02:18 +08:00

335 lines
13 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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
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
_workers = {url: WorkerState(url=url) for url in cfg["workers"]}
_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", 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 池已关闭")
# ---------------------------------------------------------------------------
# 全局串行槽
# ---------------------------------------------------------------------------
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。
全忙则排队等待,最多 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 使用)。"""
global_info = {
"global_max": _global_max,
"global_busy": _global_busy,
"global_waiting": _global_waiting,
}
if not _workers:
return {"total": 0, "healthy": 0, "busy": 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)
return {"total": total, "healthy": healthy, "busy": busy, **global_info}