fix(qb): window candidates before locking

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.",
         )
 

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

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

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 == []