diff --git a/gateway/forward.py b/gateway/forward.py index db17b07..f11fde7 100644 --- a/gateway/forward.py +++ b/gateway/forward.py @@ -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: diff --git a/gateway/pool.py b/gateway/pool.py index b538cb0..6ccdd64 100644 --- a/gateway/pool.py +++ b/gateway/pool.py @@ -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() # --------------------------------------------------------------------------- diff --git a/tests/test_gateway_pool.py b/tests/test_gateway_pool.py index a4e0b61..6239055 100644 --- a/tests/test_gateway_pool.py +++ b/tests/test_gateway_pool.py @@ -234,3 +234,118 @@ def test_no_primary_config_all_are_primary(): ]) w = _run_acquire(workers) assert w.url == "http://w1:1" + + +# --------------------------------------------------------------------------- +# 优先级插队（接口5）：全局槽 + worker 池 +# --------------------------------------------------------------------------- + +_PRIO_CFG = { + "dispatch": { + "queue_wait_seconds": 5, # 测试用：等待足够长，避免排队超时 + "max_global_concurrency": 1, + "max_queue_wait_seconds": 5, + } +} + + +def test_global_slot_priority_jumps_queue(): + """普通请求已排队时，优先级请求插队到最前，先获得槽位。""" + pool.init_global_slot(_PRIO_CFG) + order = [] + + async def run(): + assert await pool.acquire_global_slot(_PRIO_CFG) is True # 占住唯一槽 + + async def normal(): + if await pool.acquire_global_slot(_PRIO_CFG): + order.append("normal") + pool.release_global_slot() + + async def prio(): + if await pool.acquire_global_slot(_PRIO_CFG, priority=1): + order.append("prio") + pool.release_global_slot() + + t1 = asyncio.create_task(normal()) + await asyncio.sleep(0.05) # 确保 normal 先进入等待队列 + t2 = asyncio.create_task(prio()) + await asyncio.sleep(0.05) # prio 此时应已插队到队首 + pool.release_global_slot() # 释放槽位 → 应唤醒 prio 而非 normal + await asyncio.wait_for(asyncio.gather(t1, t2), timeout=3) + + asyncio.run(run()) + assert order == ["prio", "normal"], f"优先级请求应先获得槽位，实际顺序={order}" + assert pool.get_pool_status()["global_busy"] == 0 + + +def test_global_slot_same_priority_fifo(): + """同优先级仍按 FIFO，不插队。""" + pool.init_global_slot(_PRIO_CFG) + order = [] + + async def run(): + assert await pool.acquire_global_slot(_PRIO_CFG) is True + + async def waiter(tag): + if await pool.acquire_global_slot(_PRIO_CFG): + order.append(tag) + pool.release_global_slot() + + tasks = [] + for tag in ("a", "b", "c"): + tasks.append(asyncio.create_task(waiter(tag))) + await asyncio.sleep(0.05) # 保证入队顺序 + pool.release_global_slot() + await asyncio.wait_for(asyncio.gather(*tasks), timeout=3) + + asyncio.run(run()) + assert order == ["a", "b", "c"], f"同优先级应保持 FIFO，实际顺序={order}" + + +def _run_worker_prio_scenario(acquire_specs): + """构造「唯一 worker 忙 → 多个请求排队 → 释放」场景，返回获得顺序。 + + acquire_specs: [(tag, priority), ...]，按此顺序发起请求。 + """ + workers = _make_pool([ + ("http://w1:1", True, True, True), # 唯一 worker，初始即忙 + ]) + order = [] + + async def run(): + pool._pool_condition = asyncio.Condition() + pool._worker_waiters.clear() + orig = pool._workers + pool._workers = workers + try: + async def waiter(tag, prio_): + w = await pool.acquire_worker(_PRIO_CFG, priority=prio_) + order.append(tag) + await pool.release_worker(w) + + tasks = [] + for tag, prio_ in acquire_specs: + tasks.append(asyncio.create_task(waiter(tag, prio_))) + await asyncio.sleep(0.05) # 保证入队顺序 + # 释放 worker → 唤醒优先级最高的等待者 + await pool.release_worker(workers["http://w1:1"]) + await asyncio.wait_for(asyncio.gather(*tasks), timeout=3) + finally: + pool._workers = orig + pool._worker_waiters.clear() + + asyncio.run(run()) + return order + + +def test_worker_queue_priority_jumps_queue(): + """普通请求已排队时，优先级请求插队到最前，先获得 worker。""" + order = _run_worker_prio_scenario([("normal", 0), ("prio", 1)]) + assert order == ["prio", "normal"], f"优先级请求应先获得 worker，实际顺序={order}" + + +def test_worker_queue_same_priority_fifo(): + """同优先级仍按 FIFO，不插队。""" + order = _run_worker_prio_scenario([("a", 0), ("b", 0), ("c", 0)]) + assert order == ["a", "b", "c"], f"同优先级应保持 FIFO，实际顺序={order}"