feat(gateway): 接口5排队插队优先级，两级队列均支持
- 全局槽换为 _PrioritySemaphore：priority 大者排到等待队列最前，同级 FIFO， 无等待者时才允许直接获取（收紧 asyncio.Semaphore 的 barging）； 授予槽位瞬间请求被取消时归还槽位，防泄漏 - worker 池 Condition FIFO 等待改为显式优先级等待队列，release 唤醒队首； 被唤醒但 worker 被抢时按原序号重排；授予/取消竞态时转手唤醒下一位 - forward.py 新增 PRIORITY_PATHS（现仅接口5 /api/v1/hairline/generate） - 新增 4 个单测：两级队列的插队 + 同级 FIFO
This commit is contained in:
+17
-6
@@ -29,6 +29,15 @@ from gateway.reqlog import extract_form_params
|
||||
|
||||
logger = logging.getLogger("gateway.forward")
|
||||
|
||||
# 插队优先级接口：排队（全局槽 / worker 池）时排到等待队列最前，同优先级 FIFO。
|
||||
PRIORITY_PATHS = frozenset({
|
||||
"/api/v1/hairline/generate", # 接口5：发际线PNG生成
|
||||
})
|
||||
|
||||
|
||||
def _priority_of(path: str) -> int:
|
||||
return 1 if path in PRIORITY_PATHS else 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 全局 httpx client（连接复用）
|
||||
@@ -137,7 +146,8 @@ async def proxy_request(request: Request, path: str) -> JSONResponse:
|
||||
"""代理一次请求到 worker，处理重试与 base64 改写。
|
||||
|
||||
流程：
|
||||
1. acquire worker（排队等待空闲）
|
||||
0. 获取全局串行槽（PRIORITY_PATHS 中的接口插队到最前）
|
||||
1. acquire worker（排队等待空闲，同样支持优先级插队）
|
||||
2. 重构 multipart/form 请求，加 X-Internal-Token
|
||||
3. 转发到 worker
|
||||
4. 成功 → release worker → 改写 base64 → 返回 JSONResponse
|
||||
@@ -161,8 +171,9 @@ async def proxy_request(request: Request, path: str) -> JSONResponse:
|
||||
except Exception as ex: # noqa: BLE001
|
||||
request.state.log_request_params = {"_parse_error": str(ex)}
|
||||
|
||||
# --- 全局串行槽（GPU 请求在此排队；超时拒绝） ---
|
||||
if not await acquire_global_slot(cfg):
|
||||
# --- 全局串行槽（GPU 请求在此排队；超时拒绝；优先级接口插队到最前） ---
|
||||
priority = _priority_of(path)
|
||||
if not await acquire_global_slot(cfg, priority):
|
||||
request.state.worker_url = ""
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
@@ -175,13 +186,13 @@ async def proxy_request(request: Request, path: str) -> JSONResponse:
|
||||
)
|
||||
|
||||
try:
|
||||
return await _dispatch_with_retries(request, path, cfg, body)
|
||||
return await _dispatch_with_retries(request, path, cfg, body, priority)
|
||||
finally:
|
||||
release_global_slot()
|
||||
|
||||
|
||||
async def _dispatch_with_retries(
|
||||
request: Request, path: str, cfg: dict, body: bytes
|
||||
request: Request, path: str, cfg: dict, body: bytes, priority: int = 0
|
||||
) -> JSONResponse:
|
||||
"""在已持有全局串行槽的前提下：选 worker → 转发 → 重试 → 改写 base64 → 返回。
|
||||
|
||||
@@ -203,7 +214,7 @@ async def _dispatch_with_retries(
|
||||
for attempt in range(attempts):
|
||||
worker = None
|
||||
try:
|
||||
worker = await acquire_worker(cfg)
|
||||
worker = await acquire_worker(cfg, priority)
|
||||
# 记录当前使用的 worker，供日志中间件读取
|
||||
request.state.worker_url = worker.url
|
||||
except NoWorkerAvailable:
|
||||
|
||||
+139
-31
@@ -5,6 +5,8 @@
|
||||
- 每个 worker 一个 busy 标志（per_worker_concurrency=1）
|
||||
- acquire_worker() 从在线池挑空闲 worker；全忙排队；池空→NoWorkerAvailable
|
||||
- release_worker() 释放 worker
|
||||
- 两级排队（全局槽 / worker 池）均支持优先级：priority 大者插到队列最前，
|
||||
同优先级 FIFO；已有等待者时新请求不得越过（严格按优先级排序）
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -61,15 +63,74 @@ _pool_condition: Optional[asyncio.Condition] = None
|
||||
_health_task: Optional[asyncio.Task] = None
|
||||
_shutdown_event: Optional[asyncio.Event] = None
|
||||
|
||||
# worker 池等待者：[(-priority, 入队序号, future)]，按 (-priority, 序号) 升序排列。
|
||||
# 等待者在持有 _pool_condition 锁时增删；release_worker() 唤醒队首（优先级最高者）。
|
||||
_worker_waiters: List[tuple] = []
|
||||
_waiter_seq: int = 0
|
||||
|
||||
# 全局串行槽：同一时间最多 max_global_concurrency 个请求进入 worker 派发，其余排队。
|
||||
# ⚠️ 依赖单进程 uvicorn 部署（hair-gateway.service 无 --workers）。
|
||||
# 若改多 worker / gunicorn，进程内信号量会失效，需换成跨进程锁（文件锁/Redis）。
|
||||
_global_sem: Optional[asyncio.Semaphore] = None
|
||||
_global_sem: Optional["_PrioritySemaphore"] = None
|
||||
_global_max: int = 1
|
||||
_global_busy: int = 0
|
||||
_global_waiting: int = 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 优先级信号量
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _PrioritySemaphore:
|
||||
"""带插队优先级的信号量（替代 asyncio.Semaphore，后者不支持优先级）。
|
||||
|
||||
等待者按 (-priority, 入队序号) 排序：priority 大者先获得槽位（插队到最前），
|
||||
同优先级 FIFO。仅当无等待者时才允许直接获取——防止新请求越过排队者
|
||||
（asyncio.Semaphore 本身允许这种插队，这里收紧以保证优先级语义）。
|
||||
"""
|
||||
|
||||
def __init__(self, value: int) -> None:
|
||||
self._value = value
|
||||
self._waiters: List[tuple] = [] # [(-priority, seq, future)]
|
||||
self._seq = 0
|
||||
|
||||
async def acquire(self, priority: int = 0) -> None:
|
||||
if self._value > 0 and not self._waiters:
|
||||
self._value -= 1
|
||||
return
|
||||
loop = asyncio.get_running_loop()
|
||||
fut = loop.create_future()
|
||||
self._seq += 1
|
||||
entry = (-priority, self._seq, fut)
|
||||
idx = 0
|
||||
while idx < len(self._waiters) and self._waiters[idx][:2] <= entry[:2]:
|
||||
idx += 1
|
||||
self._waiters.insert(idx, entry)
|
||||
if priority > 0 and idx < len(self._waiters) - 1:
|
||||
logger.info("优先级请求插队到全局队列第 %d 位（超越 %d 个等待者）",
|
||||
idx + 1, len(self._waiters) - 1 - idx)
|
||||
try:
|
||||
await fut
|
||||
except asyncio.CancelledError:
|
||||
if fut.done() and not fut.cancelled():
|
||||
# 已被 release 授予槽位、但本任务随后被取消 → 归还槽位
|
||||
self.release()
|
||||
else:
|
||||
try:
|
||||
self._waiters.remove(entry)
|
||||
except ValueError:
|
||||
pass
|
||||
raise
|
||||
|
||||
def release(self) -> None:
|
||||
while self._waiters:
|
||||
_, _, fut = self._waiters.pop(0)
|
||||
if not fut.done():
|
||||
fut.set_result(None)
|
||||
return
|
||||
self._value += 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 健康检查后台任务
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -214,23 +275,25 @@ async def shutdown_pool() -> None:
|
||||
def init_global_slot(cfg: dict) -> None:
|
||||
"""初始化全局串行信号量。
|
||||
|
||||
max_global_concurrency=1 → 同一时间只有一个请求进入 worker 派发，其余在信号量上排队。
|
||||
max_global_concurrency=N → 同一时间最多 N 个请求进入 worker 派发，其余排队。
|
||||
priority>0 的请求可插队到等待队列最前（见 _PrioritySemaphore）。
|
||||
依赖单进程 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_sem = _PrioritySemaphore(n)
|
||||
_global_max = n
|
||||
_global_busy = 0
|
||||
_global_waiting = 0
|
||||
logger.info("全局并发槽初始化: max=%d（单进程生效）", n)
|
||||
|
||||
|
||||
async def acquire_global_slot(cfg: dict) -> bool:
|
||||
async def acquire_global_slot(cfg: dict, priority: int = 0) -> bool:
|
||||
"""获取一个全局处理槽。True=获得；False=排队超时。
|
||||
|
||||
priority>0 的请求插队到等待队列最前，同优先级 FIFO。
|
||||
超时由 dispatch.max_queue_wait_seconds 控制（需 < nginx proxy_read_timeout 600s）。
|
||||
"""
|
||||
global _global_busy, _global_waiting
|
||||
@@ -239,7 +302,7 @@ async def acquire_global_slot(cfg: dict) -> bool:
|
||||
timeout = float(cfg.get("dispatch", {}).get("max_queue_wait_seconds", 590))
|
||||
_global_waiting += 1
|
||||
try:
|
||||
await asyncio.wait_for(_global_sem.acquire(), timeout=timeout)
|
||||
await asyncio.wait_for(_global_sem.acquire(priority), timeout=timeout)
|
||||
_global_busy += 1
|
||||
return True
|
||||
except asyncio.TimeoutError:
|
||||
@@ -277,56 +340,101 @@ async def mark_worker_unhealthy(w: WorkerState) -> None:
|
||||
w.url, cooldown)
|
||||
|
||||
|
||||
async def acquire_worker(cfg: dict) -> WorkerState:
|
||||
def _insert_worker_waiter(key: tuple, fut: asyncio.Future) -> tuple:
|
||||
"""按 (-priority, 序号) 升序插入等待者（须在持有 _pool_condition 锁时调用）。"""
|
||||
entry = (key[0], key[1], fut)
|
||||
idx = 0
|
||||
while idx < len(_worker_waiters) and _worker_waiters[idx][:2] <= entry[:2]:
|
||||
idx += 1
|
||||
_worker_waiters.insert(idx, entry)
|
||||
return entry
|
||||
|
||||
|
||||
def _wake_best_worker_waiter() -> None:
|
||||
"""唤醒优先级最高的等待者（须在持有 _pool_condition 锁时调用）。"""
|
||||
while _worker_waiters:
|
||||
_, _, fut = _worker_waiters.pop(0)
|
||||
if not fut.done():
|
||||
fut.set_result(None)
|
||||
return
|
||||
|
||||
|
||||
async def acquire_worker(cfg: dict, priority: int = 0) -> WorkerState:
|
||||
"""从在线池中获取一个空闲 worker。
|
||||
|
||||
主备分组：只要还有主 worker（primary）在线，就只在主组内挑空闲
|
||||
（全忙则排队等待，不用备组）；只有主组全部下线时才用备用 worker。
|
||||
全忙则排队等待，最多 queue_wait_seconds；池空立即抛异常。
|
||||
排队按 (-priority, 入队序号) 排序：priority 大者插队到最前，同优先级 FIFO；
|
||||
已有等待者时新请求不得越过（保证优先级语义，防饿死）。
|
||||
"""
|
||||
queue_wait = cfg["dispatch"]["queue_wait_seconds"]
|
||||
deadline = time.monotonic() + queue_wait
|
||||
my_key: Optional[tuple] = None # (-priority, seq)，被唤醒后重排仍用原序号
|
||||
woken = False # 是否刚被 release_worker 唤醒（被唤醒者直接参与挑选）
|
||||
|
||||
async with _pool_condition:
|
||||
while True:
|
||||
while True:
|
||||
async with _pool_condition:
|
||||
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("后端服务暂不可用，请稍后重试")
|
||||
if idle and (woken or not _worker_waiters):
|
||||
# 取第一个空闲 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
|
||||
|
||||
# 全忙，等待
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
raise NoWorkerAvailable("后端服务繁忙，请稍后重试")
|
||||
if not online:
|
||||
raise NoWorkerAvailable("后端服务暂不可用，请稍后重试")
|
||||
|
||||
logger.info("所有可用 worker 全忙（%d 在线），等待 %.1fs...",
|
||||
len(candidates), remaining)
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
raise NoWorkerAvailable("后端服务繁忙，请稍后重试")
|
||||
|
||||
if my_key is None:
|
||||
global _waiter_seq
|
||||
_waiter_seq += 1
|
||||
my_key = (-priority, _waiter_seq)
|
||||
first_enqueue = True
|
||||
else:
|
||||
first_enqueue = False
|
||||
fut = asyncio.get_running_loop().create_future()
|
||||
entry = _insert_worker_waiter(my_key, fut)
|
||||
if first_enqueue:
|
||||
pos = _worker_waiters.index(entry) + 1
|
||||
if not idle:
|
||||
logger.info("所有可用 worker 全忙（%d 在线），等待 %.1fs...",
|
||||
len(candidates), remaining)
|
||||
if priority > 0 and pos < len(_worker_waiters):
|
||||
logger.info("优先级请求插队到 worker 队列第 %d 位（超越 %d 个等待者）",
|
||||
pos, len(_worker_waiters) - pos)
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(fut, timeout=remaining)
|
||||
woken = True
|
||||
except asyncio.TimeoutError:
|
||||
async with _pool_condition:
|
||||
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
|
||||
_worker_waiters.remove(entry)
|
||||
except ValueError:
|
||||
# 授予与取消竞态：已被弹出授予（worker 空闲），转手唤醒下一位
|
||||
if fut.done() and not fut.cancelled():
|
||||
_wake_best_worker_waiter()
|
||||
raise NoWorkerAvailable("后端服务繁忙，请稍后重试")
|
||||
|
||||
|
||||
async def release_worker(w: WorkerState) -> None:
|
||||
"""释放 worker，标记为空闲并通知等待者。"""
|
||||
"""释放 worker，标记为空闲并唤醒优先级最高的等待者。"""
|
||||
async with _pool_condition:
|
||||
w.busy = False
|
||||
logger.debug("Worker 释放: %s", w.url)
|
||||
_pool_condition.notify(1)
|
||||
_wake_best_worker_waiter()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user