perf: order the expired-lease index like the fetch

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

@@ -14,5 +14,11 @@ about 200 ms per fetch at a million ready rows, against 0.1 ms for this shape, o
 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.
+rebuild when upgrading from 0.14. The expired-lease branch reads `<table>_leased_idx`, keyed
+the same way over leased rows, and filters out leases that haven't expired. Its first index,
+`(queue, acquired_at)`, made the branch read and sort every expired lease of its queue: with 50,000
+of them in a million rows a fetch took about 10 ms on PostgreSQL 17 and 15 ms on 13, against 1.4 ms
+with the ordered index. It replaced that index rather than joining it, because a second index on
+leased rows would add a write to every claim. Leases that are still held sit in the ordered index
+too and the scan steps over them; with 20,000 in flight at the head of the queue a fetch still took
+about 1.5 ms.

docs/operations/alembic.md

@@ -35,9 +35,9 @@ op.create_table(
     sa.PrimaryKeyConstraint("id"),
 )
 op.create_index(
-    "outbox_lease_idx",
+    "outbox_leased_idx",
     "outbox",
-    ["queue", "acquired_at"],
+    ["queue", "next_attempt_at", "id"],
     unique=False,
     postgresql_where=sa.text("acquired_token IS NOT NULL"),
 )
@@ -62,9 +62,11 @@ The three indexes carry load-bearing partial predicates:
   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`.
-  The fetch's expired-lease branch is served by this index. Here too, the
-  predicate is load-bearing for fetch performance.
+- `outbox_leased_idx`: `(queue, next_attempt_at, id) WHERE acquired_token IS NOT NULL`.
+  The fetch's expired-lease branch reads this index in the same order and
+  skips leases that haven't expired yet, so its cost doesn't grow with the
+  number of hung handlers. 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
@@ -147,6 +149,57 @@ 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.
 
+## Upgrading to 0.17: replace `outbox_lease_idx` { #upgrading-to-017 }
+
+0.17 reclaims expired leases with an ordered scan of a new index,
+`outbox_leased_idx` on `(queue, next_attempt_at, id) WHERE acquired_token
+IS NOT NULL`. It replaces `outbox_lease_idx` on `(queue, acquired_at)`, which
+made a fetch read and sort every expired lease on the queue: about 10 ms
+per fetch at 50,000 expired leases in a million rows, against 1.4 ms with
+the new index.
+`validate_schema()` reports a table without `outbox_leased_idx` as drift.
+
+Build the new index and drop the old one without blocking writes, and run
+the migration **before** deploying 0.17. Earlier versions fetch the same
+rows with either index, so the order is safe.
+
+```python
+import sqlalchemy as sa
+from alembic import op
+
+
+def upgrade() -> None:
+    with op.get_context().autocommit_block():
+        op.create_index(
+            "outbox_leased_idx",
+            "outbox",
+            ["queue", "next_attempt_at", "id"],
+            postgresql_where=sa.text("acquired_token IS NOT NULL"),
+            postgresql_concurrently=True,
+        )
+        op.drop_index("outbox_lease_idx", table_name="outbox", postgresql_concurrently=True)
+
+
+def downgrade() -> None:
+    with op.get_context().autocommit_block():
+        op.create_index(
+            "outbox_lease_idx",
+            "outbox",
+            ["queue", "acquired_at"],
+            postgresql_where=sa.text("acquired_token IS NOT NULL"),
+            postgresql_concurrently=True,
+        )
+        op.drop_index("outbox_leased_idx", table_name="outbox", postgresql_concurrently=True)
+```
+
+If a health probe runs `validate_schema()` on 0.16, it reports
+`outbox_lease_idx` missing once the old index is dropped. In that case put
+the `drop_index` in a second migration and run it after 0.17 is deployed;
+0.17 ignores the extra index. As with the 0.15 rebuild, pass `schema=` for a
+table in a non-default schema, drop an `INVALID` index left by a failed
+concurrent build before retrying, and don't apply what
+`alembic revision --autogenerate` produces for this change.
+
 ## Adding the DLQ after the fact
 
 To opt into [Dead-letter queue](../usage/dlq.md) audit, add
@@ -352,7 +405,7 @@ three indexes and their expected shape:
 | Index | Columns | Unique | `postgresql_where` |
 | --- | --- | --- | --- |
 | `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_leased_idx` | `queue, next_attempt_at, id` | no | `acquired_token IS NOT NULL` |
 | `outbox_timer_id_uq` | `queue, timer_id` | yes | `timer_id IS NOT NULL` |
 
 ```python
@@ -367,7 +420,7 @@ op.create_index(
 ```
 
 Substitute the columns / `unique` / predicate from the table above for
-`outbox_pending_idx` and `outbox_lease_idx`.
+`outbox_pending_idx` and `outbox_leased_idx`.
 
 The recipes pass literal names (`'outbox_lease_ck'`, `'outbox_timer_id_uq'`),
 which are the exact names the package emits when no `naming_convention` is set.

docs/tutorials/first-outbox-app.md

@@ -187,8 +187,8 @@ You should see:
  timer_id         | character varying(255)   |           |          |
 Indexes:
     "outbox_pkey" PRIMARY KEY, btree (id)
-    "outbox_lease_idx" btree (queue, acquired_at) WHERE acquired_token IS NOT NULL
-    "outbox_pending_idx" btree (queue, next_attempt_at) WHERE acquired_token IS NULL
+    "outbox_leased_idx" btree (queue, next_attempt_at, id) WHERE acquired_token IS NOT NULL
+    "outbox_pending_idx" btree (queue, next_attempt_at, id) WHERE acquired_token IS NULL
     "outbox_timer_id_uq" UNIQUE, btree (queue, timer_id) WHERE timer_id IS NOT NULL
 Check constraints:
     "outbox_lease_ck" CHECK ((acquired_token IS NULL) = (acquired_at IS NULL))

faststream_outbox/schema.py

@@ -48,7 +48,7 @@
 _CHANNEL_PREFIX = "outbox_"
 _PENDING_IDX_SUFFIX = "_pending_idx"
 _TIMER_ID_UQ_SUFFIX = "_timer_id_uq"
-_LEASE_IDX_SUFFIX = "_lease_idx"
+_LEASED_IDX_SUFFIX = "_leased_idx"
 _LEASE_CK_SUFFIX = "_lease_ck"
 
 
@@ -68,7 +68,7 @@ def validate_table_identifiers(table_name: str) -> None:
             f"{_CHANNEL_PREFIX}{table_name}",
             f"{table_name}{_PENDING_IDX_SUFFIX}",
             f"{table_name}{_TIMER_ID_UQ_SUFFIX}",
-            f"{table_name}{_LEASE_IDX_SUFFIX}",
+            f"{table_name}{_LEASED_IDX_SUFFIX}",
             f"{table_name}{_LEASE_CK_SUFFIX}",
         )
     )
@@ -124,7 +124,6 @@ def make_outbox_table(metadata: "MetaData", table_name: str = "outbox") -> Table
     # 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,
@@ -142,17 +141,12 @@ def make_outbox_table(metadata: "MetaData", table_name: str = "outbox") -> Table
         unique=True,
         postgresql_where=table.c.timer_id.is_not(None),
     )
-    # Partial index that backs the fetch query's expired-lease branch
-    # (`WHERE acquired_token IS NOT NULL AND acquired_at < lease_cutoff`). Without it,
-    # the OR's second arm forces a seq-scan of the whole table on every fetch — invisible
-    # under healthy steady-state (Branch A dominates) but the tail grows linearly with
-    # table size when handlers wedge or `lease_ttl_seconds < P99`. Trade-off: every fetch
-    # UPDATE rewrites (acquired_token, acquired_at), so this index pays write amplification
-    # proportional to the claim rate.
+    # The same ordered scan for the expired-lease branch. See ADR 0003.
     Index(
-        f"{table_name}{_LEASE_IDX_SUFFIX}",
+        f"{table_name}{_LEASED_IDX_SUFFIX}",
         table.c.queue,
-        table.c.acquired_at,
+        table.c.next_attempt_at,
+        table.c.id,
         postgresql_where=table.c.acquired_token.is_not(None),
     )
     return table

faststream_outbox/schema_validation.py

@@ -22,7 +22,7 @@
 from faststream_outbox.autovacuum import _RELOPTIONS_QUERY, autovacuum_findings
 from faststream_outbox.schema import (
     _LEASE_CK_SUFFIX,
-    _LEASE_IDX_SUFFIX,
+    _LEASED_IDX_SUFFIX,
     _PENDING_IDX_SUFFIX,
     _TIMER_ID_UQ_SUFFIX,
     make_dlq_table,
@@ -123,7 +123,7 @@ def _normalize_predicate(predicate: str) -> str:
 _EXPECTED_INDEX_PREDICATES = {
     _PENDING_IDX_SUFFIX: "acquired_token is null",
     _TIMER_ID_UQ_SUFFIX: "timer_id is not null",
-    _LEASE_IDX_SUFFIX: "acquired_token is not null",
+    _LEASED_IDX_SUFFIX: "acquired_token is not null",
 }
 
 # Suffixes whose index must be UNIQUE — the timer_id arbiter relies on it for ON CONFLICT (F2-10).
@@ -251,6 +251,7 @@ def _validate_check_constraints_sync(connection: "Connection", table: "Table", s
 
 
 _PENDING_IDX_UPGRADE_URL = "https://faststream-outbox.modern-python.org/operations/alembic/#upgrading-to-015"
+_LEASED_IDX_UPGRADE_URL = "https://faststream-outbox.modern-python.org/operations/alembic/#upgrading-to-017"
 
 
 def _compose_schema_mismatch_message(errors: list[str], *, has_blind_drift: bool) -> str:
@@ -404,5 +405,9 @@ def _drift_entry_to_error(entry: "tuple[typing.Any, ...]", table_name: str) -> s
         error = f"table '{table_name}' missing index over {cols} (name={idx.name!r}, unique={bool(idx.unique)})"
         if idx.name.endswith(_PENDING_IDX_SUFFIX):
             error += f"; a table created before 0.15 needs this index rebuilt: {_PENDING_IDX_UPGRADE_URL}"
+        elif idx.name.endswith(_LEASED_IDX_SUFFIX):
+            error += (
+                f"; a table created before 0.17 needs it in place of {table_name}_lease_idx: {_LEASED_IDX_UPGRADE_URL}"
+            )
         return error
     return None

tests/test_integration.py

@@ -13,7 +13,7 @@
 import pytest
 from faststream import BaseMiddleware, Context, Depends
 from faststream.kafka import KafkaBroker, TestKafkaBroker
-from sqlalchemy import MetaData, Table, event, func, insert, literal, select, text
+from sqlalchemy import MetaData, Table, case, event, func, insert, literal, select, text
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
 
@@ -424,13 +424,43 @@ async def test_validate_schema_fails_when_pending_index_missing(pg_engine, outbo
         await client.validate_schema()
 
 
-async def test_validate_schema_fails_when_lease_index_missing(pg_engine, outbox_table) -> None:
-    """Missing the expired-lease partial index sends the fetch CTE back to seq-scan; flag it."""
-    drop_sql = f'DROP INDEX "{outbox_table.name}_lease_idx"'
+async def test_validate_schema_reports_the_pre_0_17_lease_index(
+    pg_engine: AsyncEngine,
+    outbox_table: Table,
+) -> None:
+    """An outbox migrated before 0.17 has ``_lease_idx`` on ``(queue, acquired_at)`` and lacks ``_leased_idx``."""
     async with pg_engine.begin() as conn:
-        await conn.exec_driver_sql(drop_sql)
+        await conn.execute(text(f'DROP INDEX "{outbox_table.name}_leased_idx"'))
+        await conn.execute(
+            text(
+                f'CREATE INDEX "{outbox_table.name}_lease_idx" ON {outbox_table.name} (queue, acquired_at) '
+                "WHERE acquired_token IS NOT NULL",
+            ),
+        )
     client = OutboxClient(pg_engine, outbox_table)
-    with pytest.raises(RuntimeError, match="missing index"):
+    with pytest.raises(RuntimeError, match=r"missing index.*leased_idx.*operations/alembic/#upgrading-to-017"):
+        await client.validate_schema()
+
+
+@pytest.mark.parametrize(
+    ("predicate", "error"),
+    [("", "is not a partial index"), ("WHERE acquired_token IS NULL", "has wrong partial predicate")],
+)
+async def test_validate_schema_checks_the_leased_index_predicate(
+    pg_engine: AsyncEngine,
+    outbox_table: Table,
+    predicate: str,
+    error: str,
+) -> None:
+    """A ``_leased_idx`` rebuilt without ``acquired_token IS NOT NULL`` is reported (alembic's diff ignores it)."""
+    idx = f"{outbox_table.name}_leased_idx"
+    async with pg_engine.begin() as conn:
+        await conn.execute(text(f'DROP INDEX "{idx}"'))
+        await conn.execute(
+            text(f'CREATE INDEX "{idx}" ON {outbox_table.name} (queue, next_attempt_at, id) {predicate}')
+        )
+    client = OutboxClient(pg_engine, outbox_table)
+    with pytest.raises(RuntimeError, match=f"'{idx}' {error}"):
         await client.validate_schema()
 
 
@@ -616,6 +646,35 @@ def _scanned_rows(plan: dict[str, Any], table: str) -> list[float]:
     return found
 
 
+async def _explain_fetch(engine: AsyncEngine, table: Table, *, limit: int) -> dict[str, Any]:
+    """Run ``EXPLAIN ANALYZE`` on the statement a real fetch of ``orders`` sends, rolled back."""
+    captured: list[tuple[str, Any]] = []
+
+    def capture(_conn: Any, _cursor: Any, statement: str, parameters: Any, _context: Any, _many: bool) -> None:
+        if "FOR UPDATE SKIP LOCKED" in statement:
+            captured.append((statement, parameters))
+
+    client = OutboxClient(engine, table)
+    event.listen(engine.sync_engine, "before_cursor_execute", capture)
+    try:
+        async with engine.connect() as conn:
+            await client.fetch(conn, ["orders"], limit=limit, lease_ttl_seconds=60.0)
+    finally:
+        event.remove(engine.sync_engine, "before_cursor_execute", capture)
+    statement, parameters = captured[0]
+
+    async with engine.connect() as conn:
+        raw = (await conn.get_raw_connection()).driver_connection
+        assert raw is not None
+        tx = raw.transaction()
+        await tx.start()
+        try:
+            explained = await raw.fetchval(f"EXPLAIN (ANALYZE, FORMAT JSON) {statement}", *parameters)
+        finally:
+            await tx.rollback()
+    return explained[0]["Plan"]
+
+
 async def test_fetch_reads_only_the_rows_it_claims_from_a_backlog(
     pg_engine: AsyncEngine,
     outbox_table: Table,
@@ -633,31 +692,36 @@ async def test_fetch_reads_only_the_rows_it_claims_from_a_backlog(
         await conn.execute(insert(outbox_table).from_select(["queue", "payload"], backlog))
         await conn.execute(text(f"ANALYZE {outbox_table.name}"))
 
-    captured: list[tuple[str, Any]] = []
+    plan = await _explain_fetch(pg_engine, outbox_table, limit=limit)
 
-    def capture(_conn: Any, _cursor: Any, statement: str, parameters: Any, _context: Any, _many: bool) -> None:
-        if "FOR UPDATE SKIP LOCKED" in statement:
-            captured.append((statement, parameters))
+    assert max(_scanned_rows(plan, outbox_table.name)) <= 2 * limit
 
-    client = OutboxClient(pg_engine, outbox_table)
-    event.listen(pg_engine.sync_engine, "before_cursor_execute", capture)
-    try:
-        async with pg_engine.connect() as conn:
-            await client.fetch(conn, ["orders"], limit=limit, lease_ttl_seconds=60.0)
-    finally:
-        event.remove(pg_engine.sync_engine, "before_cursor_execute", capture)
-    statement, parameters = captured[0]
 
-    async with pg_engine.connect() as conn:
-        raw = (await conn.get_raw_connection()).driver_connection
-        assert raw is not None
-        tx = raw.transaction()
-        await tx.start()
-        try:
-            explained = await raw.fetchval(f"EXPLAIN (ANALYZE, FORMAT JSON) {statement}", *parameters)
-        finally:
-            await tx.rollback()
-    plan = explained[0]["Plan"]
+async def test_fetch_reads_only_the_rows_it_claims_from_expired_leases(
+    pg_engine: AsyncEngine,
+    outbox_table: Table,
+) -> None:
+    """INVARIANT: reclaiming expired leases reads about ``limit`` rows, however many handlers hung.
+
+    The expired-lease branch must come off an ordered scan of ``_leased_idx``. An index keyed on
+    ``acquired_at`` makes Postgres read and sort every expired lease on the queue.
+    """
+    limit = 10
+    async with pg_engine.begin() as conn:
+        series = func.generate_series(1, 20000).column_valued("g")
+        leased = series % 4 == 0
+        backlog = select(
+            literal("orders"),
+            literal(b"x"),
+            func.now() - func.make_interval(0, 0, 0, 0, 0, 0, series),
+            case((leased, func.gen_random_uuid())),
+            case((leased, func.now() - func.make_interval(0, 0, 0, 0, 1))),
+        )
+        columns = ["queue", "payload", "next_attempt_at", "acquired_token", "acquired_at"]
+        await conn.execute(insert(outbox_table).from_select(columns, backlog))
+        await conn.execute(text(f"ANALYZE {outbox_table.name}"))
+
+    plan = await _explain_fetch(pg_engine, outbox_table, limit=limit)
 
     assert max(_scanned_rows(plan, outbox_table.name)) <= 2 * limit
 

tests/test_unit.py

@@ -140,16 +140,14 @@ def test_make_outbox_table_declares_timer_unique_index() -> None:
     assert timer_idx.dialect_options["postgresql"]["where"] is not None
 
 
-def test_make_outbox_table_declares_lease_idx() -> None:
+def test_make_outbox_table_declares_leased_idx_in_fetch_order() -> None:
     metadata = MetaData()
     t = make_outbox_table(metadata, table_name="my_outbox")
-    lease_idx = next(idx for idx in t.indexes if idx.name == "my_outbox_lease_idx")
-    assert lease_idx.unique is False
-    assert [c.name for c in lease_idx.columns] == ["queue", "acquired_at"]
-    # Partial-index predicate `acquired_token IS NOT NULL` — the fetch CTE's
-    # Branch B (expired-lease reclaim) relies on this; without it, the OR
-    # disjunct degrades to seq-scan even with `_pending_idx` covering Branch A.
-    assert lease_idx.dialect_options["postgresql"]["where"] is not None
+    leased_idx = next(idx for idx in t.indexes if idx.name == "my_outbox_leased_idx")
+    assert leased_idx.unique is False
+    assert [c.name for c in leased_idx.columns] == ["queue", "next_attempt_at", "id"]
+    assert str(leased_idx.dialect_options["postgresql"]["where"]) == "my_outbox.acquired_token IS NOT NULL"
+    assert "my_outbox_lease_idx" not in {idx.name for idx in t.indexes}
 
 
 def test_make_outbox_table_attaches_to_metadata() -> None: