wehub-resource-sync
111 lines
3.6 KiB
Python
111 lines
3.6 KiB
Python
"""
|
|
work_queue.py - bounded background-job execution with per-principal quotas.
|
|
|
|
/crawl/job and /llm/job used FastAPI BackgroundTasks with no bound: a client
|
|
could enqueue unlimited background jobs and exhaust memory / browser slots, and
|
|
one caller could starve others.
|
|
|
|
This replaces that with a fixed worker pool draining an asyncio.Queue, plus an
|
|
optional per-principal concurrency cap. Everything is configurable, and any
|
|
limit set to 0 (or null) means "unbounded" - i.e. the previous behavior is fully
|
|
recoverable:
|
|
|
|
limits.queue.maxsize 0 => unbounded queue (never 503)
|
|
limits.queue.workers worker pool size (>=1)
|
|
limits.queue.per_principal 0 => no per-caller cap (never 429)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Awaitable, Callable, Dict, Optional
|
|
|
|
logger = logging.getLogger("crawl4ai.workqueue")
|
|
|
|
JobFactory = Callable[[], Awaitable[None]]
|
|
|
|
|
|
class QueueFull(Exception):
|
|
"""The bounded job queue is full -> 503 Retry-After."""
|
|
|
|
|
|
class QuotaExceeded(Exception):
|
|
"""The principal has too many concurrent jobs -> 429."""
|
|
|
|
|
|
class WorkQueue:
|
|
def __init__(self, maxsize: int = 0, workers: int = 4, per_principal: int = 0):
|
|
self.maxsize = max(0, int(maxsize)) # 0 = unbounded
|
|
self.workers = max(1, int(workers))
|
|
self.per_principal = max(0, int(per_principal)) # 0 = unlimited
|
|
self._q: Optional[asyncio.Queue] = None
|
|
self._tasks: list = []
|
|
self._counts: Dict[str, int] = {}
|
|
|
|
@property
|
|
def started(self) -> bool:
|
|
return self._q is not None
|
|
|
|
async def start(self) -> None:
|
|
self._q = asyncio.Queue(maxsize=self.maxsize)
|
|
self._tasks = [asyncio.create_task(self._worker()) for _ in range(self.workers)]
|
|
logger.info(
|
|
"work queue started (maxsize=%s, workers=%s, per_principal=%s)",
|
|
self.maxsize or "unbounded", self.workers, self.per_principal or "unlimited",
|
|
)
|
|
|
|
async def stop(self) -> None:
|
|
for t in self._tasks:
|
|
t.cancel()
|
|
self._tasks = []
|
|
self._q = None
|
|
|
|
async def _worker(self) -> None:
|
|
assert self._q is not None
|
|
while True:
|
|
principal, factory = await self._q.get()
|
|
try:
|
|
await factory()
|
|
except Exception:
|
|
logger.exception("background job failed")
|
|
finally:
|
|
self._release(principal)
|
|
self._q.task_done()
|
|
|
|
def _release(self, principal: Optional[str]) -> None:
|
|
if not principal:
|
|
return
|
|
n = self._counts.get(principal, 0) - 1
|
|
if n <= 0:
|
|
self._counts.pop(principal, None)
|
|
else:
|
|
self._counts[principal] = n
|
|
|
|
def submit(self, factory: JobFactory, principal: Optional[str] = None) -> None:
|
|
"""Enqueue a job. Raises QuotaExceeded / QueueFull (mapped to 429 / 503)."""
|
|
if self._q is None:
|
|
raise RuntimeError("work queue not started")
|
|
if self.per_principal and principal:
|
|
if self._counts.get(principal, 0) >= self.per_principal:
|
|
raise QuotaExceeded()
|
|
try:
|
|
self._q.put_nowait((principal, factory))
|
|
except asyncio.QueueFull:
|
|
raise QueueFull()
|
|
if principal:
|
|
self._counts[principal] = self._counts.get(principal, 0) + 1
|
|
|
|
|
|
# Process-wide singleton, set at server boot.
|
|
_JOB_QUEUE: Optional[WorkQueue] = None
|
|
|
|
|
|
def set_job_queue(q: Optional[WorkQueue]) -> None:
|
|
global _JOB_QUEUE
|
|
_JOB_QUEUE = q
|
|
|
|
|
|
def get_job_queue() -> Optional[WorkQueue]:
|
|
return _JOB_QUEUE
|