perf: claim rows with ordered per-queue scans; add id to _pending_idx

docs/adr/0003-fetch-query-shape.md

@@ -0,0 +1,18 @@
+# The fetch claims rows with ordered per-queue scans, not one filtered sort
+
+The fetch query takes each subscribed queue separately through `CROSS JOIN LATERAL` and runs two
+branches per queue, unleased rows and expired leases. Each branch is an
+`ORDER BY next_attempt_at, id LIMIT n FOR UPDATE SKIP LOCKED` scan, the branches are combined with
+`UNION ALL`, and an outer `ORDER BY ... LIMIT n` picks the rows to claim. The unleased branch reads
+`<table>_pending_idx`, whose key is `(queue, next_attempt_at, id)`, so the scan stops after `n`
+unlocked rows and a fetch costs the same with ten ready rows or a million. The earlier single query,
+`WHERE queue = ANY(:queues) AND (acquired_token IS NULL OR <expired>)`, claimed the same rows in the
+same order, but Postgres could serve neither the OR nor `= ANY` with an ordered index scan, and it
+doesn't push the row limit through `LockRows`, so every fetch read and sorted the whole backlog:
+about 200 ms per fetch at a million ready rows, against 0.1 ms for this shape, on PostgreSQL 13 and
+17 alike. `id` is in the index because rows published in one transaction share `next_attempt_at`,
+and without it a fetch reads the entire tie group. The costs accepted are a statement that is harder
+to read, up to `2 × n` rows per queue locked briefly while only `n` are claimed (other fetchers skip
+them until the claim commits, which affects fairness, not correctness), and a concurrent index
+rebuild when upgrading from 0.14. The expired-lease branch still sorts every expired lease of its
+queue; an index for it was left out until it is measured against a realistic number of hung handlers.

docs/introduction/how-it-works.md

@@ -63,37 +63,43 @@ bulk publish costs one NOTIFY, not one per row.
 
 Per subscriber, two loops run concurrently:
 
-The fetch loop owns a long-lived `AsyncConnection` for the fetch CTE and
-a separate raw asyncpg connection for `LISTEN outbox_<table>`. A single CTE
+The fetch loop owns a long-lived `AsyncConnection` for the fetch query and
+a separate raw asyncpg connection for `LISTEN outbox_<table>`. One statement
 claims rows:
 
 ```sql
-WITH claimed AS (
-    SELECT id FROM outbox
-    WHERE queue = :queue
-      AND next_attempt_at <= now()
-      AND (
-        acquired_token IS NULL
-        OR acquired_at < now() - make_interval(secs => :lease_ttl)
-      )
-    ORDER BY next_attempt_at, id
-    LIMIT :batch
-    FOR UPDATE SKIP LOCKED
+WITH queues AS (SELECT DISTINCT unnest(:queues) AS name),
+pending AS (
+    SELECT p.id, p.next_attempt_at FROM queues JOIN LATERAL (
+        SELECT id, next_attempt_at FROM outbox
+        WHERE queue = queues.name AND next_attempt_at <= now()
+          AND acquired_token IS NULL
+        ORDER BY next_attempt_at, id LIMIT :batch
+        FOR UPDATE SKIP LOCKED
+    ) p ON true
+),
+expired AS (
+    -- the same per-queue scan for leases older than :lease_ttl
+),
+ready AS (
+    SELECT id FROM (SELECT * FROM pending UNION ALL SELECT * FROM expired) c
+    ORDER BY next_attempt_at, id LIMIT :batch
 )
 UPDATE outbox
 SET acquired_token = :uuid, acquired_at = now(),
     deliveries_count = deliveries_count + 1
-WHERE id IN (SELECT id FROM claimed)
+WHERE id IN (SELECT id FROM ready)
 RETURNING *
 ```
 
-This is simplified for illustration. The real query writes each `OR`
-disjunct with its own partial-index predicate spelled out as a conjunct, so
-Postgres can use the `outbox_pending_idx` / `outbox_lease_idx` partial
-indexes instead of a seq-scan. The naive `OR` above is the exact shape the
-code avoids.
+Each queue gets its own ordered scan of `outbox_pending_idx`, which stops
+after `:batch` unlocked rows, so a fetch costs the same with ten ready rows
+or a million. A single `WHERE ... OR ...` over `queue = ANY(:queues)` would
+claim the same rows, but Postgres would have to read and sort the whole
+backlog first. [ADR 0003](https://github.com/modern-python/faststream-outbox/blob/main/docs/adr/0003-fetch-query-shape.md)
+records why the query has this shape.
 
-The CTE reclaims both unleased rows AND rows whose lease has expired
+The query reclaims both unleased rows AND rows whose lease has expired
 (`acquired_at < now() - lease_ttl_seconds`), so there is no separate stuck-row
 reaper. The idle-sleep is short-circuited by NOTIFY via an `asyncio.Event`, so
 idle dispatch latency drops from up to `max_fetch_interval` (default 10s) to

docs/operations/alembic.md

@@ -44,7 +44,7 @@ op.create_index(
 op.create_index(
     "outbox_pending_idx",
     "outbox",
-    ["queue", "next_attempt_at"],
+    ["queue", "next_attempt_at", "id"],
     unique=False,
     postgresql_where=sa.text("acquired_token IS NULL"),
 )
@@ -56,14 +56,15 @@ op.create_index(
 
 The three indexes carry load-bearing partial predicates:
 
-- `outbox_pending_idx`: `(queue, next_attempt_at) WHERE acquired_token IS NULL`.
-  Branch A of the fetch CTE (unleased rows) is served by this index. The
-  WHERE clause is written so Postgres' planner recognizes the implied
-  predicate and uses the partial index; drop the predicate and the planner
-  falls back to a seq-scan as the table grows.
+- `outbox_pending_idx`: `(queue, next_attempt_at, id) WHERE acquired_token IS NULL`.
+  The fetch's unleased branch reads this index in order, one queue at a
+  time, and stops after `limit` unlocked rows. `id` is in the key because
+  rows published in one transaction share `next_attempt_at`; without it a
+  fetch reads the whole tie group. Drop the predicate and the planner falls
+  back to a seq-scan as the table grows.
 - `outbox_lease_idx`: `(queue, acquired_at) WHERE acquired_token IS NOT NULL`.
-  Branch B of the fetch CTE (expired-lease reclaim) is served by this
-  index. Here too, the predicate is load-bearing for fetch performance.
+  The fetch's expired-lease branch is served by this index. Here too, the
+  predicate is load-bearing for fetch performance.
 - `outbox_timer_id_uq`: unique `(queue, timer_id) WHERE timer_id IS NOT NULL`.
   Backs the `timer_id` dedup contract via
   `pg_insert(...).on_conflict_do_nothing(...)`. Without the partial
@@ -92,6 +93,60 @@ indexes are exactly what the broker depends on. When you wire the
 `/health` probe or CI gate, it fails when the live DB drifts from
 this declaration. (It is opt-in; it never runs at `broker.start()`.)
 
+## Upgrading to 0.15: rebuild `outbox_pending_idx` { #upgrading-to-015 }
+
+0.15 claims rows with an ordered scan of `outbox_pending_idx` per queue
+(see [ADR 0003](https://github.com/modern-python/faststream-outbox/blob/main/docs/adr/0003-fetch-query-shape.md)),
+and the index gains `id` as its last column. A table migrated before 0.15
+still has `(queue, next_attempt_at)`, and `validate_schema()` reports it as
+drift.
+
+Rebuild the index without blocking writes, and run the migration **before**
+deploying 0.15. Earlier versions work on the new index, so the order is
+safe; 0.15 on the old index works too, but each fetch then reads every row
+that shares the oldest `next_attempt_at`.
+
+```python
+import sqlalchemy as sa
+from alembic import op
+
+
+def upgrade() -> None:
+    with op.get_context().autocommit_block():
+        op.create_index(
+            "outbox_pending_idx_new",
+            "outbox",
+            ["queue", "next_attempt_at", "id"],
+            postgresql_where=sa.text("acquired_token IS NULL"),
+            postgresql_concurrently=True,
+        )
+        op.drop_index("outbox_pending_idx", table_name="outbox", postgresql_concurrently=True)
+        op.execute("ALTER INDEX outbox_pending_idx_new RENAME TO outbox_pending_idx")
+
+
+def downgrade() -> None:
+    with op.get_context().autocommit_block():
+        op.create_index(
+            "outbox_pending_idx_old",
+            "outbox",
+            ["queue", "next_attempt_at"],
+            postgresql_where=sa.text("acquired_token IS NULL"),
+            postgresql_concurrently=True,
+        )
+        op.drop_index("outbox_pending_idx", table_name="outbox", postgresql_concurrently=True)
+        op.execute("ALTER INDEX outbox_pending_idx_old RENAME TO outbox_pending_idx")
+```
+
+`CONCURRENTLY` can't run inside a transaction, hence `autocommit_block()`.
+For a table in a non-default schema, pass `schema=` to both index calls and
+qualify the name in the `ALTER INDEX`. If a concurrent build fails, it
+leaves an `INVALID` index behind: drop it and run the migration again.
+
+Don't apply what `alembic revision --autogenerate` produces for this
+change. It drops and recreates the index in one transaction without
+`CONCURRENTLY`, which blocks every publish and fetch on the table for the
+whole build.
+
 ## Adding the DLQ after the fact
 
 To opt into [Dead-letter queue](../usage/dlq.md) audit, add
@@ -296,7 +351,7 @@ three indexes and their expected shape:
 
 | Index | Columns | Unique | `postgresql_where` |
 | --- | --- | --- | --- |
-| `outbox_pending_idx` | `queue, next_attempt_at` | no | `acquired_token IS NULL` |
+| `outbox_pending_idx` | `queue, next_attempt_at, id` | no | `acquired_token IS NULL` |
 | `outbox_lease_idx` | `queue, acquired_at` | no | `acquired_token IS NOT NULL` |
 | `outbox_timer_id_uq` | `queue, timer_id` | yes | `timer_id IS NOT NULL` |
 

faststream_outbox/client.py

@@ -23,17 +23,17 @@
 from sqlalchemy import (
     ARRAY,
     Float,
+    Integer,
     String,
-    and_,
-    any_,
     bindparam,
     cast,
     delete,
     func,
     insert,
-    or_,
     select,
     text,
+    true,
+    union_all,
     update,
 )
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -51,7 +51,7 @@
     import typing
     from collections.abc import Mapping, Sequence
 
-    from sqlalchemy import Insert, Table
+    from sqlalchemy import CTE, ColumnElement, Insert, Table, Update
     from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
 
@@ -141,6 +141,50 @@ async def validate_schema(self, *, check_autovacuum: bool = False) -> None: ...
     async def ping(self) -> bool: ...
 
 
+def _build_fetch_stmt(outbox: "Table") -> "Update":
+    """Build the claim statement once; queues, limit, lease TTL and the new token are bind parameters.
+
+    Each queue gets two ``LATERAL`` branches, unleased rows and expired leases, and each branch is an
+    ordered ``LIMIT n FOR UPDATE SKIP LOCKED`` scan, so Postgres stops after ``n`` unlocked rows instead of
+    sorting the whole backlog. See ADR 0003 for why the query has this shape.
+    """
+    t = outbox
+    limit = bindparam("limit", type_=Integer)
+    # ``make_interval(secs => :lease_ttl)`` keeps the cutoff server-side, immune to worker/DB clock skew.
+    lease_cutoff = func.now() - func.make_interval(0, 0, 0, 0, 0, 0, bindparam("lease_ttl", type_=Float))
+    queue_array = ARRAY(String)
+    queue_name = func.unnest(cast(bindparam("queues", type_=queue_array), queue_array)).label("name")
+    queues = select(queue_name).distinct().cte("queues")
+
+    def branch(name: str, *lease_state: "ColumnElement[bool]") -> "CTE":
+        per_queue = (
+            select(t.c.id, t.c.next_attempt_at)
+            .where(t.c.queue == queues.c.name, t.c.next_attempt_at <= func.now(), *lease_state)
+            .order_by(t.c.next_attempt_at, t.c.id)
+            .limit(limit)
+            .with_for_update(skip_locked=True)
+            .lateral(f"{name}_rows")
+        )
+        return select(per_queue.c.id, per_queue.c.next_attempt_at).select_from(queues.join(per_queue, true())).cte(name)
+
+    pending = branch("pending", t.c.acquired_token.is_(None))
+    expired = branch("expired", t.c.acquired_token.is_not(None), t.c.acquired_at < lease_cutoff)
+    candidates = union_all(
+        select(pending.c.id, pending.c.next_attempt_at), select(expired.c.id, expired.c.next_attempt_at)
+    ).subquery("candidates")
+    ready = select(candidates.c.id).order_by(candidates.c.next_attempt_at, candidates.c.id).limit(limit).cte("ready")
+    return (
+        update(t)
+        .where(t.c.id.in_(select(ready.c.id)))
+        .values(
+            acquired_at=func.now(),
+            acquired_token=bindparam("token", type_=t.c.acquired_token.type),
+            deliveries_count=t.c.deliveries_count + 1,
+        )
+        .returning(*t.c)
+    )
+
+
 def _build_dlq_cte(outbox: "Table", dlq: "Table") -> "Insert":
     """Build the DELETE-then-archive CTE once; the row id, lease token and failure context are bind parameters."""
     deleted = (
@@ -171,6 +215,7 @@ def __init__(
         self._engine = engine
         self._table = outbox_table
         self._dlq_table = dlq_table
+        self._fetch_stmt = _build_fetch_stmt(outbox_table)
         self._dlq_cte = None if dlq_table is None else _build_dlq_cte(outbox_table, dlq_table)
 
     @property
@@ -212,50 +257,9 @@ async def fetch(
             raise TypeError(msg)
         if not queues:
             return []
-        token = uuid.uuid4()
-        t = self._table
-
-        # ``make_interval(secs => :lease_ttl)`` keeps the cutoff computation server-side
-        # so lease expiry is immune to clock skew between worker and DB hosts.
-        lease_cutoff = func.now() - func.make_interval(0, 0, 0, 0, 0, 0, bindparam("lease_ttl", type_=Float))
-        ready = (
-            select(t.c.id)
-            .where(
-                t.c.next_attempt_at <= func.now(),
-                # P11: ``queue = ANY(:queues)`` binds a single array param, so the SQL
-                # text is stable regardless of the number of queues — asyncpg can reuse
-                # the prepared statement across ticks instead of recompiling an
-                # OR-of-N-equalities whose text changes with the queue count.
-                t.c.queue == any_(bindparam("queues", list(queues), type_=ARRAY(String))),
-                # The OR is split into two index-implying disjuncts so Postgres'
-                # partial-index inference picks up both `_pending_idx` (Branch A,
-                # `acquired_token IS NULL`) and `_lease_idx` (Branch B,
-                # `acquired_token IS NOT NULL AND acquired_at < lease_cutoff`) and
-                # combines them via BitmapOr. The naive form `acquired_at < lease_cutoff`
-                # alone does not imply `acquired_token IS NOT NULL`, so the planner
-                # cannot prove `_lease_idx` applies and falls back to a seq-scan.
-                or_(
-                    t.c.acquired_token.is_(None),
-                    and_(t.c.acquired_token.is_not(None), t.c.acquired_at < lease_cutoff),
-                ),
-            )
-            .order_by(t.c.next_attempt_at, t.c.id)
-            .limit(limit)
-            .with_for_update(skip_locked=True)
-            .cte("ready")
-        )
-        stmt = (
-            update(t)
-            .where(t.c.id.in_(select(ready.c.id)))
-            .values(
-                acquired_at=func.now(),
-                acquired_token=token,
-                deliveries_count=t.c.deliveries_count + 1,
-            )
-            .returning(*t.c)
-        )
+        params = {"queues": list(queues), "limit": limit, "lease_ttl": lease_ttl_seconds, "token": uuid.uuid4()}
         async with conn.begin():
-            result = await conn.execute(stmt, {"lease_ttl": lease_ttl_seconds})
+            result = await conn.execute(self._fetch_stmt, params)
             rows = result.mappings().all()
         return [_row_to_message(dict(row)) for row in rows]
 

faststream_outbox/schema.py

@@ -121,13 +121,15 @@ def make_outbox_table(metadata: "MetaData", table_name: str = "outbox") -> Table
             name=f"{table_name}{_LEASE_CK_SUFFIX}",
         ),
     )
-    # Partial index that backs the fetch query's hot branch
-    # (`WHERE acquired_token IS NULL AND queue = ? AND next_attempt_at <= now()`).
+    # Partial index that backs the fetch query's unleased branch: an ordered scan per queue on
+    # (next_attempt_at, id) that stops after `limit` unlocked rows. `id` is part of the key because
+    # rows published in one transaction share `next_attempt_at`. See ADR 0003.
     # The expired-lease branch is covered by `_lease_idx` below.
     Index(
         f"{table_name}{_PENDING_IDX_SUFFIX}",
         table.c.queue,
         table.c.next_attempt_at,
+        table.c.id,
         postgresql_where=table.c.acquired_token.is_(None),
     )
     # Partial unique index that backs `publish(..., timer_id=...)`'s ON CONFLICT DO NOTHING

tests/test_client_contract.py

@@ -316,6 +316,25 @@ async def test_fetch_selects_oldest_under_limit(contract: _Harness) -> None:
     assert rid_a not in {m.id for m in msgs}
 
 
+async def test_fetch_selects_oldest_across_queues_under_limit(contract: _Harness) -> None:
+    rid_a_old = await contract.seed(queue="a", next_attempt_offset=-30.0)
+    rid_b_mid = await contract.seed(queue="b", next_attempt_offset=-20.0)
+    await contract.seed(queue="a", next_attempt_offset=-10.0)
+    await contract.seed(queue="b", next_attempt_offset=-5.0)
+    msgs = await contract.fetch(["a", "b"], limit=2, lease_ttl_seconds=60.0)
+    assert {m.id for m in msgs} == {rid_a_old, rid_b_mid}
+
+
+async def test_fetch_selects_oldest_across_pending_and_expired_leases(contract: _Harness) -> None:
+    rid_expired = await contract.seed(
+        queue="orders", next_attempt_offset=-30.0, leased_token=uuid.uuid4(), acquired_age=120.0, deliveries_count=1
+    )
+    rid_pending = await contract.seed(queue="orders", next_attempt_offset=-20.0)
+    await contract.seed(queue="orders", next_attempt_offset=-10.0)
+    msgs = await contract.fetch(["orders"], limit=2, lease_ttl_seconds=60.0)
+    assert {m.id for m in msgs} == {rid_expired, rid_pending}
+
+
 async def test_fetch_respects_limit(contract: _Harness) -> None:
     for _ in range(3):
         await contract.seed(queue="orders")