diff --git a/pgqueuer/adapters/persistence/qb.py b/pgqueuer/adapters/persistence/qb.py index ca85492b..8d58b8db 100644 --- a/pgqueuer/adapters/persistence/qb.py +++ b/pgqueuer/adapters/persistence/qb.py @@ -559,15 +559,12 @@ def build_dequeue_query( """, comment="Entrypoints with free capacity; remaining caps a single dequeue.", ) - # LEAST(batch, NULL) = batch, so unlimited entrypoints keep the full batch. - lateral_limit = f"LEAST({batch}, available.remaining)" else: composer.cte( "available", f"SELECT UNNEST({eps}::text[]) AS entrypoint", comment="No entrypoint carries a concurrency limit; all are available.", ) - lateral_limit = batch # Both gate lines drop out of their block when the budget is unlimited. budget_where = "" @@ -594,9 +591,58 @@ def build_dequeue_query( f"GREATEST(LEAST({batch}, {global_limit} - (SELECT total FROM worker_load)), 0)" ) - composer.cte( - "next_queued", - f""" + if capacity_gated: + # Windowing before the lock keeps the lock node out from under a LIMIT, + # where SKIP LOCKED slides down the backlog and past the cap (#761). + queued_body = f""" + SELECT job.id, job.priority + FROM ( + SELECT job.id, job.priority + FROM available + CROSS JOIN LATERAL ( + SELECT candidate.id, candidate.priority + FROM {queue_table} candidate + WHERE candidate.entrypoint = available.entrypoint + AND candidate.status = 'queued' + AND candidate.execute_after < NOW() + ORDER BY candidate.priority DESC, candidate.id ASC + LIMIT {batch} + FOR UPDATE SKIP LOCKED + ) job + WHERE available.remaining IS NULL + UNION ALL + SELECT job.id, job.priority + FROM available + CROSS JOIN LATERAL ( + SELECT locked.id, locked.priority + FROM ( + SELECT candidate.id + FROM {queue_table} candidate + WHERE candidate.entrypoint = available.entrypoint + AND candidate.status = 'queued' + AND candidate.execute_after < NOW() + ORDER BY candidate.priority DESC, candidate.id ASC + LIMIT LEAST({batch}, available.remaining) + ) capped + CROSS JOIN LATERAL ( + SELECT target.id, target.priority + FROM {queue_table} target + WHERE target.id = capped.id + AND target.status = 'queued' + AND target.execute_after < NOW() + FOR UPDATE OF target SKIP LOCKED + ) locked + ORDER BY locked.priority DESC, locked.id ASC + LIMIT LEAST({batch}, available.remaining) + ) job + WHERE available.remaining IS NOT NULL + ) job + {budget_where} + ORDER BY job.priority DESC, job.id ASC + LIMIT {batch} + """ + else: + queued_body = f""" SELECT job.id, job.priority FROM available CROSS JOIN LATERAL ( @@ -606,13 +652,16 @@ def build_dequeue_query( AND candidate.status = 'queued' AND candidate.execute_after < NOW() ORDER BY candidate.priority DESC, candidate.id ASC - LIMIT {lateral_limit} + LIMIT {batch} FOR UPDATE SKIP LOCKED ) job {budget_where} ORDER BY job.priority DESC, job.id ASC LIMIT {batch} - """, + """ + composer.cte( + "next_queued", + queued_body, comment="New queued jobs; LATERAL hits the (entrypoint, priority, id) index.", ) diff --git a/test/query_shapes/dequeue_both_gates.sql b/test/query_shapes/dequeue_both_gates.sql index 8463f656..403af90b 100644 --- a/test/query_shapes/dequeue_both_gates.sql +++ b/test/query_shapes/dequeue_both_gates.sql @@ -40,16 +40,46 @@ worker_load AS ( -- New queued jobs; LATERAL hits the (entrypoint, priority, id) index. next_queued AS ( SELECT job.id, job.priority - FROM available - CROSS JOIN LATERAL ( - SELECT candidate.id, candidate.priority - FROM pgqueuer candidate - WHERE candidate.entrypoint = available.entrypoint - AND candidate.status = 'queued' - AND candidate.execute_after < NOW() - ORDER BY candidate.priority DESC, candidate.id ASC - LIMIT LEAST($1, available.remaining) - FOR UPDATE SKIP LOCKED + FROM ( + SELECT job.id, job.priority + FROM available + CROSS JOIN LATERAL ( + SELECT candidate.id, candidate.priority + FROM pgqueuer candidate + WHERE candidate.entrypoint = available.entrypoint + AND candidate.status = 'queued' + AND candidate.execute_after < NOW() + ORDER BY candidate.priority DESC, candidate.id ASC + LIMIT $1 + FOR UPDATE SKIP LOCKED + ) job + WHERE available.remaining IS NULL + UNION ALL + SELECT job.id, job.priority + FROM available + CROSS JOIN LATERAL ( + SELECT locked.id, locked.priority + FROM ( + SELECT candidate.id + FROM pgqueuer candidate + WHERE candidate.entrypoint = available.entrypoint + AND candidate.status = 'queued' + AND candidate.execute_after < NOW() + ORDER BY candidate.priority DESC, candidate.id ASC + LIMIT LEAST($1, available.remaining) + ) capped + CROSS JOIN LATERAL ( + SELECT target.id, target.priority + FROM pgqueuer target + WHERE target.id = capped.id + AND target.status = 'queued' + AND target.execute_after < NOW() + FOR UPDATE OF target SKIP LOCKED + ) locked + ORDER BY locked.priority DESC, locked.id ASC + LIMIT LEAST($1, available.remaining) + ) job + WHERE available.remaining IS NOT NULL ) job WHERE (SELECT total FROM worker_load) < $6 ORDER BY job.priority DESC, job.id ASC diff --git a/test/query_shapes/dequeue_entrypoint_gate.sql b/test/query_shapes/dequeue_entrypoint_gate.sql index 3846d391..99dc69c8 100644 --- a/test/query_shapes/dequeue_entrypoint_gate.sql +++ b/test/query_shapes/dequeue_entrypoint_gate.sql @@ -32,16 +32,46 @@ available AS ( -- New queued jobs; LATERAL hits the (entrypoint, priority, id) index. next_queued AS ( SELECT job.id, job.priority - FROM available - CROSS JOIN LATERAL ( - SELECT candidate.id, candidate.priority - FROM pgqueuer candidate - WHERE candidate.entrypoint = available.entrypoint - AND candidate.status = 'queued' - AND candidate.execute_after < NOW() - ORDER BY candidate.priority DESC, candidate.id ASC - LIMIT LEAST($1, available.remaining) - FOR UPDATE SKIP LOCKED + FROM ( + SELECT job.id, job.priority + FROM available + CROSS JOIN LATERAL ( + SELECT candidate.id, candidate.priority + FROM pgqueuer candidate + WHERE candidate.entrypoint = available.entrypoint + AND candidate.status = 'queued' + AND candidate.execute_after < NOW() + ORDER BY candidate.priority DESC, candidate.id ASC + LIMIT $1 + FOR UPDATE SKIP LOCKED + ) job + WHERE available.remaining IS NULL + UNION ALL + SELECT job.id, job.priority + FROM available + CROSS JOIN LATERAL ( + SELECT locked.id, locked.priority + FROM ( + SELECT candidate.id + FROM pgqueuer candidate + WHERE candidate.entrypoint = available.entrypoint + AND candidate.status = 'queued' + AND candidate.execute_after < NOW() + ORDER BY candidate.priority DESC, candidate.id ASC + LIMIT LEAST($1, available.remaining) + ) capped + CROSS JOIN LATERAL ( + SELECT target.id, target.priority + FROM pgqueuer target + WHERE target.id = capped.id + AND target.status = 'queued' + AND target.execute_after < NOW() + FOR UPDATE OF target SKIP LOCKED + ) locked + ORDER BY locked.priority DESC, locked.id ASC + LIMIT LEAST($1, available.remaining) + ) job + WHERE available.remaining IS NOT NULL ) job ORDER BY job.priority DESC, job.id ASC LIMIT $1 diff --git a/test/test_concurrency_limit.py b/test/test_concurrency_limit.py index 48f93639..346fd07a 100644 --- a/test/test_concurrency_limit.py +++ b/test/test_concurrency_limit.py @@ -1,17 +1,21 @@ from __future__ import annotations import asyncio +import contextlib +import uuid from dataclasses import dataclass, field from datetime import timedelta -from typing import Any +from typing import Any, AsyncGenerator, Awaitable, Callable import async_timeout +import asyncpg import pytest +import pytest_asyncio -from pgqueuer.db import Driver +from pgqueuer.db import AsyncpgDriver, Driver from pgqueuer.models import Job from pgqueuer.qm import QueueManager -from pgqueuer.queries import Queries +from pgqueuer.queries import EntrypointExecutionParameter, Queries @dataclass @@ -151,3 +155,99 @@ async def timer() -> None: assert dequeue_batches assert set(dequeue_batches) == {batch_size} + + +@pytest_asyncio.fixture +async def connect(dsn: str) -> AsyncGenerator[Callable[[], Awaitable[asyncpg.Connection]], None]: + """Hand out connections to the per-test database; close them on teardown.""" + async with contextlib.AsyncExitStack() as stack: + + async def _connect() -> asyncpg.Connection: + connection = await asyncpg.connect(dsn=dsn) + stack.push_async_callback(connection.close) + return connection + + yield _connect + + +async def _wait_until_done_or_lock_blocked( + monitor: asyncpg.Connection, + task: asyncio.Task[list[Job]], + backend_pid: int, +) -> None: + """Poll until *task* finished or its backend is waiting on a heavyweight lock. + + Distinguishes "statement completed" (task done) from "statement blocked on + worker A's uncommitted claim" (wait_event_type = 'Lock'), so the test never + relies on a fixed sleep for the interleaving. + """ + async with async_timeout.timeout(30): + while not task.done(): + wait_event_type = await monitor.fetchval( + "SELECT wait_event_type FROM pg_stat_activity WHERE pid = $1", + backend_pid, + ) + if wait_event_type == "Lock": + return + await asyncio.sleep(0.01) + + +@pytest.mark.parametrize("concurrency_limit", (1, 2)) +async def test_concurrency_limit_holds_across_concurrent_dequeues( + connect: Callable[[], Awaitable[asyncpg.Connection]], + concurrency_limit: int, +) -> None: + """Regression test for #761: concurrency_limit must hold across workers. + + Two workers dequeue over separate connections. Worker A claims a full + batch inside an open transaction, freezing the in-flight instant where + its picked rows are neither visible to worker B's snapshot (the capacity + count reads zero) nor released (SKIP LOCKED slides past them onto the + remaining queued rows). Worker B must claim nothing; any claim here + exceeds the entrypoint's global concurrency limit. + + Worker B runs as a task and worker A commits once B has either finished + or provably blocked on A's transaction, so the test stays deterministic + regardless of whether the dequeue implementation skips, waits, or retries. + """ + conn_a = await connect() + conn_b = await connect() + conn_monitor = await connect() + + queries_a = Queries(AsyncpgDriver(conn_a)) + queries_b = Queries(AsyncpgDriver(conn_b)) + + n_jobs = 2 * concurrency_limit + await queries_a.enqueue( + ["fetch"] * n_jobs, + [f"{n}".encode() for n in range(n_jobs)], + [0] * n_jobs, + ) + + def dequeue(q: Queries) -> asyncio.Task[list[Job]]: + return asyncio.ensure_future( + q.dequeue( + batch_size=concurrency_limit, + entrypoints={"fetch": EntrypointExecutionParameter(concurrency_limit)}, + queue_manager_id=uuid.uuid4(), + global_concurrency_limit=None, + heartbeat_timeout=timedelta(minutes=10), + ) + ) + + transaction_a = conn_a.transaction() + await transaction_a.start() + first = await dequeue(queries_a) + assert len(first) == concurrency_limit + + # Worker A's claim is now in flight: uncommitted, row locks held. + task_b = dequeue(queries_b) + await _wait_until_done_or_lock_blocked(conn_monitor, task_b, conn_b.get_server_pid()) + await transaction_a.commit() + + overlap = await asyncio.wait_for(task_b, timeout=30) + assert overlap == [], f"picked {len(first) + len(overlap)} jobs, limit {concurrency_limit}" + + # A's claim is committed and visible; capacity is exhausted either way. + visible = await asyncio.wait_for(dequeue(queries_b), timeout=30) + assert visible == []