Skip to content

Commit fa14dd9

Browse files
authored
perf: order the expired-lease index like the fetch (#240)
Replace <table>_lease_idx (queue, acquired_at) with <table>_leased_idx (queue, next_attempt_at, id) WHERE acquired_token IS NOT NULL, so the fetch's expired-lease branch stops after `limit` rows instead of sorting every expired lease on the queue. validate_schema() reports a table without the new index and points at the upgrade section. Closes #220
1 parent 358c4fd commit fa14dd9

7 files changed

Lines changed: 181 additions & 61 deletions

File tree

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,11 @@ about 200 ms per fetch at a million ready rows, against 0.1 ms for this shape, o
1414
and without it a fetch reads the entire tie group. The costs accepted are a statement that is harder
1515
to read, up to `2 × n` rows per queue locked briefly while only `n` are claimed (other fetchers skip
1616
them until the claim commits, which affects fairness, not correctness), and a concurrent index
17-
rebuild when upgrading from 0.14. The expired-lease branch still sorts every expired lease of its
18-
queue; an index for it was left out until it is measured against a realistic number of hung handlers.
17+
rebuild when upgrading from 0.14. The expired-lease branch reads `<table>_leased_idx`, keyed
18+
the same way over leased rows, and filters out leases that haven't expired. Its first index,
19+
`(queue, acquired_at)`, made the branch read and sort every expired lease of its queue: with 50,000
20+
of them in a million rows a fetch took about 10 ms on PostgreSQL 17 and 15 ms on 13, against 1.4 ms
21+
with the ordered index. It replaced that index rather than joining it, because a second index on
22+
leased rows would add a write to every claim. Leases that are still held sit in the ordered index
23+
too and the scan steps over them; with 20,000 in flight at the head of the queue a fetch still took
24+
about 1.5 ms.

‎docs/operations/alembic.md‎

Lines changed: 60 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,9 @@ op.create_table(
3535
sa.PrimaryKeyConstraint("id"),
3636
)
3737
op.create_index(
38-
"outbox_lease_idx",
38+
"outbox_leased_idx",
3939
"outbox",
40-
["queue", "acquired_at"],
40+
["queue", "next_attempt_at", "id"],
4141
unique=False,
4242
postgresql_where=sa.text("acquired_token IS NOT NULL"),
4343
)
@@ -62,9 +62,11 @@ The three indexes carry load-bearing partial predicates:
6262
rows published in one transaction share `next_attempt_at`; without it a
6363
fetch reads the whole tie group. Drop the predicate and the planner falls
6464
back to a seq-scan as the table grows.
65-
- `outbox_lease_idx`: `(queue, acquired_at) WHERE acquired_token IS NOT NULL`.
66-
The fetch's expired-lease branch is served by this index. Here too, the
67-
predicate is load-bearing for fetch performance.
65+
- `outbox_leased_idx`: `(queue, next_attempt_at, id) WHERE acquired_token IS NOT NULL`.
66+
The fetch's expired-lease branch reads this index in the same order and
67+
skips leases that haven't expired yet, so its cost doesn't grow with the
68+
number of hung handlers. Here too, the predicate is load-bearing for
69+
fetch performance.
6870
- `outbox_timer_id_uq`: unique `(queue, timer_id) WHERE timer_id IS NOT NULL`.
6971
Backs the `timer_id` dedup contract via
7072
`pg_insert(...).on_conflict_do_nothing(...)`. Without the partial
@@ -147,6 +149,57 @@ change. It drops and recreates the index in one transaction without
147149
`CONCURRENTLY`, which blocks every publish and fetch on the table for the
148150
whole build.
149151

152+
## Upgrading to 0.17: replace `outbox_lease_idx` { #upgrading-to-017 }
153+
154+
0.17 reclaims expired leases with an ordered scan of a new index,
155+
`outbox_leased_idx` on `(queue, next_attempt_at, id) WHERE acquired_token
156+
IS NOT NULL`. It replaces `outbox_lease_idx` on `(queue, acquired_at)`, which
157+
made a fetch read and sort every expired lease on the queue: about 10 ms
158+
per fetch at 50,000 expired leases in a million rows, against 1.4 ms with
159+
the new index.
160+
`validate_schema()` reports a table without `outbox_leased_idx` as drift.
161+
162+
Build the new index and drop the old one without blocking writes, and run
163+
the migration **before** deploying 0.17. Earlier versions fetch the same
164+
rows with either index, so the order is safe.
165+
166+
```python
167+
import sqlalchemy as sa
168+
from alembic import op
169+
170+
171+
def upgrade() -> None:
172+
with op.get_context().autocommit_block():
173+
op.create_index(
174+
"outbox_leased_idx",
175+
"outbox",
176+
["queue", "next_attempt_at", "id"],
177+
postgresql_where=sa.text("acquired_token IS NOT NULL"),
178+
postgresql_concurrently=True,
179+
)
180+
op.drop_index("outbox_lease_idx", table_name="outbox", postgresql_concurrently=True)
181+
182+
183+
def downgrade() -> None:
184+
with op.get_context().autocommit_block():
185+
op.create_index(
186+
"outbox_lease_idx",
187+
"outbox",
188+
["queue", "acquired_at"],
189+
postgresql_where=sa.text("acquired_token IS NOT NULL"),
190+
postgresql_concurrently=True,
191+
)
192+
op.drop_index("outbox_leased_idx", table_name="outbox", postgresql_concurrently=True)
193+
```
194+
195+
If a health probe runs `validate_schema()` on 0.16, it reports
196+
`outbox_lease_idx` missing once the old index is dropped. In that case put
197+
the `drop_index` in a second migration and run it after 0.17 is deployed;
198+
0.17 ignores the extra index. As with the 0.15 rebuild, pass `schema=` for a
199+
table in a non-default schema, drop an `INVALID` index left by a failed
200+
concurrent build before retrying, and don't apply what
201+
`alembic revision --autogenerate` produces for this change.
202+
150203
## Adding the DLQ after the fact
151204

152205
To opt into [Dead-letter queue](../usage/dlq.md) audit, add
@@ -352,7 +405,7 @@ three indexes and their expected shape:
352405
| Index | Columns | Unique | `postgresql_where` |
353406
| --- | --- | --- | --- |
354407
| `outbox_pending_idx` | `queue, next_attempt_at, id` | no | `acquired_token IS NULL` |
355-
| `outbox_lease_idx` | `queue, acquired_at` | no | `acquired_token IS NOT NULL` |
408+
| `outbox_leased_idx` | `queue, next_attempt_at, id` | no | `acquired_token IS NOT NULL` |
356409
| `outbox_timer_id_uq` | `queue, timer_id` | yes | `timer_id IS NOT NULL` |
357410

358411
```python
@@ -367,7 +420,7 @@ op.create_index(
367420
```
368421

369422
Substitute the columns / `unique` / predicate from the table above for
370-
`outbox_pending_idx` and `outbox_lease_idx`.
423+
`outbox_pending_idx` and `outbox_leased_idx`.
371424

372425
The recipes pass literal names (`'outbox_lease_ck'`, `'outbox_timer_id_uq'`),
373426
which are the exact names the package emits when no `naming_convention` is set.

‎docs/tutorials/first-outbox-app.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -187,8 +187,8 @@ You should see:
187187
timer_id | character varying(255) | | |
188188
Indexes:
189189
"outbox_pkey" PRIMARY KEY, btree (id)
190-
"outbox_lease_idx" btree (queue, acquired_at) WHERE acquired_token IS NOT NULL
191-
"outbox_pending_idx" btree (queue, next_attempt_at) WHERE acquired_token IS NULL
190+
"outbox_leased_idx" btree (queue, next_attempt_at, id) WHERE acquired_token IS NOT NULL
191+
"outbox_pending_idx" btree (queue, next_attempt_at, id) WHERE acquired_token IS NULL
192192
"outbox_timer_id_uq" UNIQUE, btree (queue, timer_id) WHERE timer_id IS NOT NULL
193193
Check constraints:
194194
"outbox_lease_ck" CHECK ((acquired_token IS NULL) = (acquired_at IS NULL))

‎faststream_outbox/schema.py‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@
4848
_CHANNEL_PREFIX = "outbox_"
4949
_PENDING_IDX_SUFFIX = "_pending_idx"
5050
_TIMER_ID_UQ_SUFFIX = "_timer_id_uq"
51-
_LEASE_IDX_SUFFIX = "_lease_idx"
51+
_LEASED_IDX_SUFFIX = "_leased_idx"
5252
_LEASE_CK_SUFFIX = "_lease_ck"
5353

5454

@@ -68,7 +68,7 @@ def validate_table_identifiers(table_name: str) -> None:
6868
f"{_CHANNEL_PREFIX}{table_name}",
6969
f"{table_name}{_PENDING_IDX_SUFFIX}",
7070
f"{table_name}{_TIMER_ID_UQ_SUFFIX}",
71-
f"{table_name}{_LEASE_IDX_SUFFIX}",
71+
f"{table_name}{_LEASED_IDX_SUFFIX}",
7272
f"{table_name}{_LEASE_CK_SUFFIX}",
7373
)
7474
)
@@ -124,7 +124,6 @@ def make_outbox_table(metadata: "MetaData", table_name: str = "outbox") -> Table
124124
# Partial index that backs the fetch query's unleased branch: an ordered scan per queue on
125125
# (next_attempt_at, id) that stops after `limit` unlocked rows. `id` is part of the key because
126126
# rows published in one transaction share `next_attempt_at`. See ADR 0003.
127-
# The expired-lease branch is covered by `_lease_idx` below.
128127
Index(
129128
f"{table_name}{_PENDING_IDX_SUFFIX}",
130129
table.c.queue,
@@ -142,17 +141,12 @@ def make_outbox_table(metadata: "MetaData", table_name: str = "outbox") -> Table
142141
unique=True,
143142
postgresql_where=table.c.timer_id.is_not(None),
144143
)
145-
# Partial index that backs the fetch query's expired-lease branch
146-
# (`WHERE acquired_token IS NOT NULL AND acquired_at < lease_cutoff`). Without it,
147-
# the OR's second arm forces a seq-scan of the whole table on every fetch — invisible
148-
# under healthy steady-state (Branch A dominates) but the tail grows linearly with
149-
# table size when handlers wedge or `lease_ttl_seconds < P99`. Trade-off: every fetch
150-
# UPDATE rewrites (acquired_token, acquired_at), so this index pays write amplification
151-
# proportional to the claim rate.
144+
# The same ordered scan for the expired-lease branch. See ADR 0003.
152145
Index(
153-
f"{table_name}{_LEASE_IDX_SUFFIX}",
146+
f"{table_name}{_LEASED_IDX_SUFFIX}",
154147
table.c.queue,
155-
table.c.acquired_at,
148+
table.c.next_attempt_at,
149+
table.c.id,
156150
postgresql_where=table.c.acquired_token.is_not(None),
157151
)
158152
return table

‎faststream_outbox/schema_validation.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
from faststream_outbox.autovacuum import _RELOPTIONS_QUERY, autovacuum_findings
2323
from faststream_outbox.schema import (
2424
_LEASE_CK_SUFFIX,
25-
_LEASE_IDX_SUFFIX,
25+
_LEASED_IDX_SUFFIX,
2626
_PENDING_IDX_SUFFIX,
2727
_TIMER_ID_UQ_SUFFIX,
2828
make_dlq_table,
@@ -123,7 +123,7 @@ def _normalize_predicate(predicate: str) -> str:
123123
_EXPECTED_INDEX_PREDICATES = {
124124
_PENDING_IDX_SUFFIX: "acquired_token is null",
125125
_TIMER_ID_UQ_SUFFIX: "timer_id is not null",
126-
_LEASE_IDX_SUFFIX: "acquired_token is not null",
126+
_LEASED_IDX_SUFFIX: "acquired_token is not null",
127127
}
128128

129129
# 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
251251

252252

253253
_PENDING_IDX_UPGRADE_URL = "https://faststream-outbox.modern-python.org/operations/alembic/#upgrading-to-015"
254+
_LEASED_IDX_UPGRADE_URL = "https://faststream-outbox.modern-python.org/operations/alembic/#upgrading-to-017"
254255

255256

256257
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
404405
error = f"table '{table_name}' missing index over {cols} (name={idx.name!r}, unique={bool(idx.unique)})"
405406
if idx.name.endswith(_PENDING_IDX_SUFFIX):
406407
error += f"; a table created before 0.15 needs this index rebuilt: {_PENDING_IDX_UPGRADE_URL}"
408+
elif idx.name.endswith(_LEASED_IDX_SUFFIX):
409+
error += (
410+
f"; a table created before 0.17 needs it in place of {table_name}_lease_idx: {_LEASED_IDX_UPGRADE_URL}"
411+
)
407412
return error
408413
return None

‎tests/test_integration.py‎

Lines changed: 92 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import pytest
1414
from faststream import BaseMiddleware, Context, Depends
1515
from faststream.kafka import KafkaBroker, TestKafkaBroker
16-
from sqlalchemy import MetaData, Table, event, func, insert, literal, select, text
16+
from sqlalchemy import MetaData, Table, case, event, func, insert, literal, select, text
1717
from sqlalchemy.exc import IntegrityError
1818
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
1919

@@ -424,13 +424,43 @@ async def test_validate_schema_fails_when_pending_index_missing(pg_engine, outbo
424424
await client.validate_schema()
425425

426426

427-
async def test_validate_schema_fails_when_lease_index_missing(pg_engine, outbox_table) -> None:
428-
"""Missing the expired-lease partial index sends the fetch CTE back to seq-scan; flag it."""
429-
drop_sql = f'DROP INDEX "{outbox_table.name}_lease_idx"'
427+
async def test_validate_schema_reports_the_pre_0_17_lease_index(
428+
pg_engine: AsyncEngine,
429+
outbox_table: Table,
430+
) -> None:
431+
"""An outbox migrated before 0.17 has ``_lease_idx`` on ``(queue, acquired_at)`` and lacks ``_leased_idx``."""
430432
async with pg_engine.begin() as conn:
431-
await conn.exec_driver_sql(drop_sql)
433+
await conn.execute(text(f'DROP INDEX "{outbox_table.name}_leased_idx"'))
434+
await conn.execute(
435+
text(
436+
f'CREATE INDEX "{outbox_table.name}_lease_idx" ON {outbox_table.name} (queue, acquired_at) '
437+
"WHERE acquired_token IS NOT NULL",
438+
),
439+
)
432440
client = OutboxClient(pg_engine, outbox_table)
433-
with pytest.raises(RuntimeError, match="missing index"):
441+
with pytest.raises(RuntimeError, match=r"missing index.*leased_idx.*operations/alembic/#upgrading-to-017"):
442+
await client.validate_schema()
443+
444+
445+
@pytest.mark.parametrize(
446+
("predicate", "error"),
447+
[("", "is not a partial index"), ("WHERE acquired_token IS NULL", "has wrong partial predicate")],
448+
)
449+
async def test_validate_schema_checks_the_leased_index_predicate(
450+
pg_engine: AsyncEngine,
451+
outbox_table: Table,
452+
predicate: str,
453+
error: str,
454+
) -> None:
455+
"""A ``_leased_idx`` rebuilt without ``acquired_token IS NOT NULL`` is reported (alembic's diff ignores it)."""
456+
idx = f"{outbox_table.name}_leased_idx"
457+
async with pg_engine.begin() as conn:
458+
await conn.execute(text(f'DROP INDEX "{idx}"'))
459+
await conn.execute(
460+
text(f'CREATE INDEX "{idx}" ON {outbox_table.name} (queue, next_attempt_at, id) {predicate}')
461+
)
462+
client = OutboxClient(pg_engine, outbox_table)
463+
with pytest.raises(RuntimeError, match=f"'{idx}' {error}"):
434464
await client.validate_schema()
435465

436466

@@ -616,6 +646,35 @@ def _scanned_rows(plan: dict[str, Any], table: str) -> list[float]:
616646
return found
617647

618648

649+
async def _explain_fetch(engine: AsyncEngine, table: Table, *, limit: int) -> dict[str, Any]:
650+
"""Run ``EXPLAIN ANALYZE`` on the statement a real fetch of ``orders`` sends, rolled back."""
651+
captured: list[tuple[str, Any]] = []
652+
653+
def capture(_conn: Any, _cursor: Any, statement: str, parameters: Any, _context: Any, _many: bool) -> None:
654+
if "FOR UPDATE SKIP LOCKED" in statement:
655+
captured.append((statement, parameters))
656+
657+
client = OutboxClient(engine, table)
658+
event.listen(engine.sync_engine, "before_cursor_execute", capture)
659+
try:
660+
async with engine.connect() as conn:
661+
await client.fetch(conn, ["orders"], limit=limit, lease_ttl_seconds=60.0)
662+
finally:
663+
event.remove(engine.sync_engine, "before_cursor_execute", capture)
664+
statement, parameters = captured[0]
665+
666+
async with engine.connect() as conn:
667+
raw = (await conn.get_raw_connection()).driver_connection
668+
assert raw is not None
669+
tx = raw.transaction()
670+
await tx.start()
671+
try:
672+
explained = await raw.fetchval(f"EXPLAIN (ANALYZE, FORMAT JSON) {statement}", *parameters)
673+
finally:
674+
await tx.rollback()
675+
return explained[0]["Plan"]
676+
677+
619678
async def test_fetch_reads_only_the_rows_it_claims_from_a_backlog(
620679
pg_engine: AsyncEngine,
621680
outbox_table: Table,
@@ -633,31 +692,36 @@ async def test_fetch_reads_only_the_rows_it_claims_from_a_backlog(
633692
await conn.execute(insert(outbox_table).from_select(["queue", "payload"], backlog))
634693
await conn.execute(text(f"ANALYZE {outbox_table.name}"))
635694

636-
captured: list[tuple[str, Any]] = []
695+
plan = await _explain_fetch(pg_engine, outbox_table, limit=limit)
637696

638-
def capture(_conn: Any, _cursor: Any, statement: str, parameters: Any, _context: Any, _many: bool) -> None:
639-
if "FOR UPDATE SKIP LOCKED" in statement:
640-
captured.append((statement, parameters))
697+
assert max(_scanned_rows(plan, outbox_table.name)) <= 2 * limit
641698

642-
client = OutboxClient(pg_engine, outbox_table)
643-
event.listen(pg_engine.sync_engine, "before_cursor_execute", capture)
644-
try:
645-
async with pg_engine.connect() as conn:
646-
await client.fetch(conn, ["orders"], limit=limit, lease_ttl_seconds=60.0)
647-
finally:
648-
event.remove(pg_engine.sync_engine, "before_cursor_execute", capture)
649-
statement, parameters = captured[0]
650699

651-
async with pg_engine.connect() as conn:
652-
raw = (await conn.get_raw_connection()).driver_connection
653-
assert raw is not None
654-
tx = raw.transaction()
655-
await tx.start()
656-
try:
657-
explained = await raw.fetchval(f"EXPLAIN (ANALYZE, FORMAT JSON) {statement}", *parameters)
658-
finally:
659-
await tx.rollback()
660-
plan = explained[0]["Plan"]
700+
async def test_fetch_reads_only_the_rows_it_claims_from_expired_leases(
701+
pg_engine: AsyncEngine,
702+
outbox_table: Table,
703+
) -> None:
704+
"""INVARIANT: reclaiming expired leases reads about ``limit`` rows, however many handlers hung.
705+
706+
The expired-lease branch must come off an ordered scan of ``_leased_idx``. An index keyed on
707+
``acquired_at`` makes Postgres read and sort every expired lease on the queue.
708+
"""
709+
limit = 10
710+
async with pg_engine.begin() as conn:
711+
series = func.generate_series(1, 20000).column_valued("g")
712+
leased = series % 4 == 0
713+
backlog = select(
714+
literal("orders"),
715+
literal(b"x"),
716+
func.now() - func.make_interval(0, 0, 0, 0, 0, 0, series),
717+
case((leased, func.gen_random_uuid())),
718+
case((leased, func.now() - func.make_interval(0, 0, 0, 0, 1))),
719+
)
720+
columns = ["queue", "payload", "next_attempt_at", "acquired_token", "acquired_at"]
721+
await conn.execute(insert(outbox_table).from_select(columns, backlog))
722+
await conn.execute(text(f"ANALYZE {outbox_table.name}"))
723+
724+
plan = await _explain_fetch(pg_engine, outbox_table, limit=limit)
661725

662726
assert max(_scanned_rows(plan, outbox_table.name)) <= 2 * limit
663727

‎tests/test_unit.py‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -140,16 +140,14 @@ def test_make_outbox_table_declares_timer_unique_index() -> None:
140140
assert timer_idx.dialect_options["postgresql"]["where"] is not None
141141

142142

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

154152

155153
def test_make_outbox_table_attaches_to_metadata() -> None:

0 commit comments

Comments
 (0)