From 55a31ddda9b8f1923f8256999615648064a59248 Mon Sep 17 00:00:00 2001 From: JeyBee Date: Sun, 30 Aug 2026 09:17:20 +0200 Subject: [PATCH] fix(qb): window candidates before locking concurrency_limit could be exceeded by workers dequeuing concurrently. The capacity count reads its own statement snapshot, so a concurrent uncommitted claim is invisible, and LIMIT ... FOR UPDATE SKIP LOCKED slides past the rows another worker holds onto further queued rows. Limited entrypoints now fix their candidate window before any locking. The window's LIMIT stops the planner pulling the subquery up, which would put the lock node back under a LIMIT and restore the slide-down. Each window row is re-fetched by primary key in its own LATERAL and locked there, so SKIP LOCKED can only skip rows inside the window: a row a concurrent worker holds shrinks this pick rather than pushing it over the limit. That closes the reported race, because both workers derive the same window and collide on the same row locks. It stops holding as soon as the ordering changes under them: a higher-priority job arriving inside another worker's in-flight claim gives the second worker a window the first never locked. Closing that needs a capacity slot on the row, which needs a migration, so it lands separately once #751 is in. Unlimited entrypoints keep the direct scan -- the slide is pure throughput there. The ungated shapes render byte-identical SQL, which is what proves this touches only entrypoints that carry a limit. Refs #761 --- pgqueuer/adapters/persistence/qb.py | 65 +++++++++-- test/query_shapes/dequeue_both_gates.sql | 50 +++++++-- test/query_shapes/dequeue_entrypoint_gate.sql | 50 +++++++-- test/test_concurrency_limit.py | 106 +++++++++++++++++- 4 files changed, 240 insertions(+), 31 deletions(-) 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 == []