Worker: acquire the concurrency slot before claiming a job
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 changed file +8 −4
modified
src/aiatlas/services/jobs.py
+8 −4
@@ -80,7 +80,7 @@ async def run_worker(*, concurrency: int | None = None, kinds: list[str] | None | |||||
| 80 | log.info("worker started", extra={"worker": worker, "kinds": kinds or "all"}) | 80 | log.info("worker started", extra={"worker": worker, "kinds": kinds or "all"}) | ||
| 81 | 81 | ||||
| 82 | async def one(job: dict[str, Any]) -> None: | 82 | async def one(job: dict[str, Any]) -> None: | ||
| 83 | − async with sem: | 83 | + try: | ||
| 84 | fn = _HANDLERS.get(job["kind"]) | 84 | fn = _HANDLERS.get(job["kind"]) | ||
| 85 | if fn is None: | 85 | if fn is None: | ||
| 86 | async with transaction() as conn: | 86 | async with transaction() as conn: | ||
@@ -94,15 +94,19 @@ async def run_worker(*, concurrency: int | None = None, kinds: list[str] | None | |||||
| 94 | log.warning("job failed", extra={"job": job["id"], "kind": job["kind"], "error": str(exc)}) | 94 | log.warning("job failed", extra={"job": job["id"], "kind": job["kind"], "error": str(exc)}) | ||
| 95 | async with transaction() as conn: | 95 | async with transaction() as conn: | ||
| 96 | await fail(conn, job, f"{exc.__class__.__name__}: {exc}\n{traceback.format_exc()[-1500:]}") | 96 | await fail(conn, job, f"{exc.__class__.__name__}: {exc}\n{traceback.format_exc()[-1500:]}") | ||
| 97 | + finally: | ||||
| 98 | + sem.release() | ||||
| 97 | 99 | ||||
| 98 | tasks: set[asyncio.Task[None]] = set() | 100 | tasks: set[asyncio.Task[None]] = set() | ||
| 99 | while not stop.is_set(): | 101 | while not stop.is_set(): | ||
| 100 | − if sem.locked(): | 102 | + await sem.acquire() # hold a slot BEFORE claiming, so concurrency bounds claims too | ||
| 101 | − await asyncio.sleep(0.5) | 103 | + if stop.is_set(): | ||
| 102 | − continue | 104 | + sem.release() | ||
| 105 | + break | ||||
| 103 | async with transaction() as conn: | 106 | async with transaction() as conn: | ||
| 104 | job = await claim_next(conn, worker, kinds) | 107 | job = await claim_next(conn, worker, kinds) | ||
| 105 | if job is None: | 108 | if job is None: | ||
| 109 | + sem.release() | ||||
| 106 | try: | 110 | try: | ||
| 107 | await asyncio.wait_for(stop.wait(), timeout=idle_sleep) | 111 | await asyncio.wait_for(stop.wait(), timeout=idle_sleep) | ||
| 108 | except TimeoutError: | 112 | except TimeoutError: | ||
| 109 | 113 | ||||