Opt-in batched terminal flush

architecture/drain.md

@@ -16,6 +16,8 @@ The subscriber carries two flags during shutdown: `self.running` (FastStream's e
 - **A timed-out drain is observable.** If `_inflight.join()` doesn't finish within the budget (`move_on_after`'s `cancelled_caught`), `stop()` emits a `WARNING` and a `drain_timeout` recorder metric (`faststream_outbox_drain_timeout_total` / `messaging.outbox.drain_timeout`) before cancelling — abandoned in-flight rows are left to lease-expiry retry, and an operator can now tell a timed-out drain from a clean one.
 - **`graceful_timeout=None`** stays unbounded where FastStream uses it that way (e.g. `ping()`), but the drain wait clamps `None` to a finite fallback (`_DEFAULT_DRAIN_TIMEOUT_SECONDS = 15.0`). `anyio.move_on_after(None)` has deadline `inf`, so without the clamp a single wedged handler would make `_inflight.join()` — and thus `stop()` — never return.
 
+**Batched terminal flush + the drain barrier.** With `terminal_flush_batch_size > 1`, a batchable row's terminal `DELETE` is deferred to a batch flush, so its `task_done()` is deferred with it — called in `_flush_buffer`, not right after `dispatch_one`. This keeps `_inflight.join()` an honest drain barrier: `join()` returns only once every buffered row's `DELETE` has actually landed (or its lease was found lost), never merely because dispatch finished. On a **graceful** `stop()` the buffer is already empty by the time `_inflight.join()` returns — the barrier requires every buffered row's flush to have completed first. On a **timed-out** drain, `_worker_inner`'s `finally` still runs and makes a best-effort, **timeout-bounded** `_flush_buffer` call on loop exit, so completed-but-unflushed rows are deleted if the connection is still healthy. The flush is capped by `_FLUSH_ON_EXIT_TIMEOUT_SECONDS` (2.0s) via `asyncio.timeout`: this block runs during task cancellation (after `_inflight.join()` has already returned or timed out), so an unbounded `await` on a wedged connection here could extend `stop()` past its budget — the cap prevents that. Rows the flush fails to reach (it raises, times out, or the connection is unresponsive) fall back to lease-expiry redelivery, same as an abandoned in-flight row. `_flush_buffer` `task_done()`s every buffered row even when the batch `DELETE` raises (then re-raises), so a flush error can never leave `join()` hanging.
+
 **Why we skip `super().stop()`.** Its `MultiLock.wait_release(graceful_timeout)` would either return instantly (healthy path; `_inflight.join()` already waited a stricter condition) or re-wait the same stuck handlers for another full budget (wedged path; **2× shutdown regression**). The subscriber inlines `TasksMixin.stop`'s cleanup body instead. Per-subscriber shutdown bound: `graceful_timeout`.
 
 ## Broker: parallel-gather subscriber stop

architecture/subscriber.md

@@ -16,6 +16,8 @@ The `WHERE` clause reclaims both unleased rows and expired leases (`acquired_at
 
 **2. `_worker_loop`** — one per `max_workers`. Each pulls from an `asyncio.Queue(maxsize=fetch_batch_size)`, dispatches via `consume()`, and flushes the terminal state. Each worker owns a long-lived `AsyncConnection` (held across reconnect) and routes terminal writes through `delete_with_lease(conn, …)` / `mark_pending_with_lease(conn, …)`, so a drain of N rows costs O(workers) pool checkouts. Flush exceptions propagate (the outer loop rebuilds the connection); the inflight slot still releases in `finally`.
 
+**Batched terminal flush (opt-in, `terminal_flush_batch_size > 1`).** A row is *batchable* iff its terminal write is a plain `DELETE` — success-by-ack or `max_deliveries`, with **no DLQ payload** (`_terminal_has_dlq(row)` is False). `dispatch_one` routes batchable terminal deletes through `_flush_or_buffer`, which appends them to a per-worker buffer instead of flushing inline and returns `True` (the row was buffered). The worker flushes the buffer as one `delete_batch_with_lease(conn, pairs)` — a single `DELETE … RETURNING id` — when it reaches `terminal_flush_batch_size` **or** the inflight queue goes idle (a `get_nowait()` → `QueueEmpty` triggers an opportunistic flush before the blocking `get()`), and once more best-effort in the loop's `finally` on exit. Retry `UPDATE`s and DLQ CTE writes are **never** batched — they stay per-row on the inline path exactly as before. At `terminal_flush_batch_size == 1` (default) and on the test-broker path (`buffer=None`), the `> 1` gate keeps every row inline, so the path is bit-for-bit identical. `task_done()` and the per-row metric are **deferred** to the flush for buffered rows: `_flush_buffer` emits `acked`/`nacked_terminal` for each id present in the `RETURNING` set and `lease_lost(phase=terminal)` for each absent one, then `task_done()`s every buffered row (**including on DB error** — it `task_done`s all, clears, and re-raises so the outer loop rebuilds the connection and the undeleted rows redeliver via lease expiry).
+
 The default ack policy is `AckPolicy.NACK_ON_ERROR`; `REJECT_ON_ERROR` and `MANUAL` are allowed. `AckPolicy.ACK_FIRST` is rejected at registration with `ValueError` — it would delete before the handler runs, defeating the outbox contract. `subscriber/factory.py` raises or warns on other footguns (`lease_ttl_seconds <= max_fetch_interval`, `max_deliveries` without retry, etc.).
 
 `OutboxSubscriber.get_one()` and `__aiter__()` are explicit `NotImplementedError`s — they point operators at `broker.fetch_unprocessed(session=..., queue=...)`. A peek that acquires a lease has surprising `deliveries_count` semantics; lease-free reads belong on `fetch_unprocessed`.
@@ -24,6 +26,8 @@ The default ack policy is `AckPolicy.NACK_ON_ERROR`; `REJECT_ON_ERROR` and `MANU
 
 A subscriber can hold up to `fetch_batch_size + max_workers` leases at once, not `fetch_batch_size` (F1-01). The free-slot computation `free = _inflight.maxsize - qsize()` counts only queued rows, so once `max_workers` rows are checked out for processing, the loop can claim another full batch. Leases are bounded and self-expire via TTL, but when sizing `lease_ttl_seconds` and reasoning about cross-replica contention, reason against `fetch_batch_size + max_workers`, not `fetch_batch_size`.
 
+With batched terminal flush on (`terminal_flush_batch_size > 1`), each worker also holds up to `terminal_flush_batch_size` completed-but-unflushed leases in its buffer, on top of the row it is currently processing — so the ceiling becomes `fetch_batch_size + max_workers × (terminal_flush_batch_size + 1)` (the `+ 1` is the in-hand row each worker holds before it buffers). These leases release only when the batch `DELETE` lands, so size `lease_ttl_seconds` against the extra hold time a full buffer implies.
+
 ## Connection budget
 
 Each subscriber holds `max_workers + 1` SQLAlchemy pool connections steady-state, plus one raw asyncpg connection for `LISTEN`.
@@ -43,7 +47,7 @@ Channel naming is `outbox_<table_name>`. Postgres limits identifiers to 63 bytes
 
 ## The lease-token invariant
 
-Load-bearing. Every terminal write (`delete_with_lease`, `mark_pending_with_lease`) filters on `acquired_token`. If a slow handler's lease expired and a newer fetch reclaimed the row, the slow handler's `DELETE`/`UPDATE` finds `rowcount == 0` and is silently dropped — preventing it from clobbering the new lease holder. Any new fetch or terminal path must preserve this.
+Load-bearing. Every terminal write (`delete_with_lease`, `mark_pending_with_lease`, `delete_batch_with_lease`) filters on `acquired_token`. If a slow handler's lease expired and a newer fetch reclaimed the row, the slow handler's `DELETE`/`UPDATE` finds `rowcount == 0` and is silently dropped — preventing it from clobbering the new lease holder. Any new fetch or terminal path must preserve this. The batched `delete_batch_with_lease` preserves it per-row: its `WHERE (id, acquired_token) IN (…)` matches only pairs that still hold their fetch-time token, and the `RETURNING id` set tells the caller exactly which rows deleted — an id absent from the set is a reclaimed lease and gets `lease_lost(phase=terminal)`, never a false `acked`.
 
 Lease-loss logs at `WARNING` with `extra={"event": "lease_lost", "phase": "terminal"|"retry", "row_id": …, "queue": …, "deliveries_count": …}`. Recurring `event=lease_lost` means `lease_ttl_seconds < handler P99`.
 

benchmarks/__main__.py

@@ -52,6 +52,12 @@ async def _run_sweep(dsn: str, messages: int, repeats: int) -> list[RunResult]:
                 cfg = RunConfig(messages=messages, max_workers=workers, fetch_batch_size=batch)
                 runs = [await run_consumer(engine, cfg) for _ in range(repeats)]
                 results.append(_median_by_wall(runs))
+        # Batched-flush before/after: the tfbs=1 sibling above is consumer/w1/b100; this
+        # tfbs=100 counterpart (consumer/w1/b100/tfbs100) coalesces the terminal DELETEs
+        # so delete_calls drops ~100x. YAGNI: one batched point, not a full cross-product.
+        batched_cfg = RunConfig(messages=messages, max_workers=1, fetch_batch_size=100, terminal_flush_batch_size=100)
+        batched_runs = [await run_consumer(engine, batched_cfg) for _ in range(repeats)]
+        results.append(_median_by_wall(batched_runs))
         producer_cfg = RunConfig(messages=messages)
         producer_runs = [await run_producer(engine, producer_cfg) for _ in range(repeats)]
         results.append(_median_by_wall(producer_runs))

benchmarks/baseline.json

@@ -1,142 +1,165 @@
 {
   "runs": {
     "consumer/w1/b10": {
-      "calls": 14928,
+      "calls": 14880,
       "dead_tup": 10000,
       "delete_calls": 5000,
       "delete_calls_per_msg": 1.0,
       "heap_bytes": 1875968,
       "index_bytes": 344064,
       "insert_calls": 0,
       "messages": 5000,
-      "msgs_per_second": 460.47656165416345,
-      "round_trips_per_msg": 2.9856,
+      "msgs_per_second": 441.21057661925147,
+      "round_trips_per_msg": 2.976,
       "select_calls": 0,
       "tup_del": 5000,
       "tup_hot_upd": 0,
       "tup_ins": 0,
       "tup_upd": 5000,
-      "wal_bytes": 4493088,
-      "wal_bytes_per_msg": 898.6176,
+      "wal_bytes": 4490440,
+      "wal_bytes_per_msg": 898.088,
       "wal_fpi": 242,
-      "wal_records": 36677,
-      "wal_records_per_msg": 7.3354,
-      "wall_seconds": 10.858315963007044
+      "wal_records": 36649,
+      "wal_records_per_msg": 7.3298,
+      "wall_seconds": 11.332457255019108
     },
     "consumer/w1/b100": {
-      "calls": 14598,
+      "calls": 14658,
       "dead_tup": 10000,
       "delete_calls": 5000,
       "delete_calls_per_msg": 1.0,
       "heap_bytes": 2048000,
       "index_bytes": 344064,
       "insert_calls": 0,
       "messages": 5000,
-      "msgs_per_second": 447.63515436129717,
-      "round_trips_per_msg": 2.9196,
+      "msgs_per_second": 450.1358749815796,
+      "round_trips_per_msg": 2.9316,
       "select_calls": 0,
       "tup_del": 5000,
       "tup_hot_upd": 0,
       "tup_ins": 0,
       "tup_upd": 5000,
-      "wal_bytes": 4797960,
-      "wal_bytes_per_msg": 959.592,
-      "wal_fpi": 262,
-      "wal_records": 37645,
-      "wal_records_per_msg": 7.529,
-      "wall_seconds": 11.169810841005528
+      "wal_bytes": 4643216,
+      "wal_bytes_per_msg": 928.6432,
+      "wal_fpi": 245,
+      "wal_records": 37462,
+      "wal_records_per_msg": 7.4924,
+      "wall_seconds": 11.10775718599325
+    },
+    "consumer/w1/b100/tfbs100": {
+      "calls": 204,
+      "dead_tup": 10000,
+      "delete_calls": 50,
+      "delete_calls_per_msg": 0.01,
+      "heap_bytes": 3710976,
+      "index_bytes": 352256,
+      "insert_calls": 0,
+      "messages": 5000,
+      "msgs_per_second": 9673.101687744125,
+      "round_trips_per_msg": 0.0408,
+      "select_calls": 0,
+      "tup_del": 5000,
+      "tup_hot_upd": 0,
+      "tup_ins": 0,
+      "tup_upd": 5000,
+      "wal_bytes": 5580984,
+      "wal_bytes_per_msg": 1116.1968,
+      "wal_fpi": 243,
+      "wal_records": 30330,
+      "wal_records_per_msg": 6.066,
+      "wall_seconds": 0.5168972850078717
     },
     "consumer/w2/b10": {
-      "calls": 8565,
+      "calls": 8586,
       "dead_tup": 10000,
       "delete_calls": 5000,
       "delete_calls_per_msg": 1.0,
-      "heap_bytes": 2097152,
+      "heap_bytes": 2121728,
       "index_bytes": 278528,
       "insert_calls": 0,
       "messages": 5000,
-      "msgs_per_second": 1209.8879048216263,
-      "round_trips_per_msg": 1.713,
+      "msgs_per_second": 1189.9545047874471,
+      "round_trips_per_msg": 1.7172,
       "select_calls": 0,
       "tup_del": 5000,
       "tup_hot_upd": 0,
       "tup_ins": 0,
       "tup_upd": 5000,
-      "wal_bytes": 4738120,
-      "wal_bytes_per_msg": 947.624,
+      "wal_bytes": 4721320,
+      "wal_bytes_per_msg": 944.264,
       "wal_fpi": 249,
-      "wal_records": 34059,
-      "wal_records_per_msg": 6.8118,
-      "wall_seconds": 4.13261425300152
+      "wal_records": 34016,
+      "wal_records_per_msg": 6.8032,
+      "wall_seconds": 4.201841314003104
     },
     "consumer/w2/b100": {
-      "calls": 8904,
+      "calls": 8943,
       "dead_tup": 10000,
       "delete_calls": 5000,
       "delete_calls_per_msg": 1.0,
-      "heap_bytes": 2408448,
+      "heap_bytes": 2424832,
       "index_bytes": 294912,
       "insert_calls": 0,
       "messages": 5000,
-      "msgs_per_second": 978.2605434470029,
-      "round_trips_per_msg": 1.7808,
+      "msgs_per_second": 853.0250182265884,
+      "round_trips_per_msg": 1.7886,
       "select_calls": 0,
       "tup_del": 5000,
       "tup_hot_upd": 0,
       "tup_ins": 0,
       "tup_upd": 5000,
-      "wal_bytes": 4660688,
-      "wal_bytes_per_msg": 932.1376,
+      "wal_bytes": 4668904,
+      "wal_bytes_per_msg": 933.7808,
       "wal_fpi": 243,
-      "wal_records": 34218,
-      "wal_records_per_msg": 6.8436,
-      "wall_seconds": 5.111112814978696
+      "wal_records": 34237,
+      "wal_records_per_msg": 6.8474,
+      "wall_seconds": 5.8614927970047574
     },
     "consumer/w4/b10": {
       "calls": 8007,
       "dead_tup": 10000,
       "delete_calls": 5000,
       "delete_calls_per_msg": 1.0,
-      "heap_bytes": 2097152,
+      "heap_bytes": 2146304,
       "index_bytes": 278528,
       "insert_calls": 0,
       "messages": 5000,
-      "msgs_per_second": 1107.6370562205375,
+      "msgs_per_second": 1161.7567625270935,
       "round_trips_per_msg": 1.6014,
       "select_calls": 0,
       "tup_del": 5000,
       "tup_hot_upd": 0,
       "tup_ins": 0,
       "tup_upd": 5000,
-      "wal_bytes": 4589160,
-      "wal_bytes_per_msg": 917.832,
+      "wal_bytes": 4675072,
+      "wal_bytes_per_msg": 935.0144,
       "wal_fpi": 242,
-      "wal_records": 33485,
-      "wal_records_per_msg": 6.697,
-      "wall_seconds": 4.514114051999059
+      "wal_records": 33711,
+      "wal_records_per_msg": 6.7422,
+      "wall_seconds": 4.30382689498947
     },
     "consumer/w4/b100": {
-      "calls": 6714,
+      "calls": 6729,
       "dead_tup": 10000,
       "delete_calls": 5000,
       "delete_calls_per_msg": 1.0,
-      "heap_bytes": 2760704,
+      "heap_bytes": 2768896,
       "index_bytes": 270336,
       "insert_calls": 0,
       "messages": 5000,
-      "msgs_per_second": 1618.337011594222,
-      "round_trips_per_msg": 1.3428,
+      "msgs_per_second": 1589.6276667449354,
+      "round_trips_per_msg": 1.3458,
       "select_calls": 0,
       "tup_del": 5000,
       "tup_hot_upd": 0,
       "tup_ins": 0,
       "tup_upd": 5000,
-      "wal_bytes": 5176184,
-      "wal_bytes_per_msg": 1035.2368,
-      "wal_fpi": 267,
-      "wal_records": 34364,
-      "wal_records_per_msg": 6.8728,
-      "wall_seconds": 3.0895913299755193
+      "wal_bytes": 4923224,
+      "wal_bytes_per_msg": 984.6448,
+      "wal_fpi": 243,
+      "wal_records": 33866,
+      "wal_records_per_msg": 6.7732,
+      "wall_seconds": 3.1453906500246376
     },
     "producer/w1/b100": {
       "calls": 10002,
@@ -147,19 +170,19 @@
       "index_bytes": 196608,
       "insert_calls": 5000,
       "messages": 5000,
-      "msgs_per_second": 4784.773275886715,
+      "msgs_per_second": 4169.062244595134,
       "round_trips_per_msg": 2.0004,
       "select_calls": 5000,
       "tup_del": 0,
       "tup_hot_upd": 0,
       "tup_ins": 5000,
       "tup_upd": 0,
-      "wal_bytes": 2918536,
-      "wal_bytes_per_msg": 583.7072,
-      "wal_fpi": 0,
-      "wal_records": 15219,
-      "wal_records_per_msg": 3.0438,
-      "wall_seconds": 1.044981593004195
+      "wal_bytes": 2925984,
+      "wal_bytes_per_msg": 585.1968,
+      "wal_fpi": 1,
+      "wal_records": 15220,
+      "wal_records_per_msg": 3.044,
+      "wall_seconds": 1.1993104699940886
     }
   }
 }

benchmarks/config.py

@@ -23,3 +23,6 @@ class RunConfig:
     fetch_batch_size: int = 100
     payload_bytes: int = DEFAULT_PAYLOAD_BYTES
     queue: str = "bench"
+    # 1 = today's per-row terminal DELETE; >1 coalesces plain-delete acks into one
+    # batched ``DELETE ... RETURNING`` per batch (the opt-in batched-flush path).
+    terminal_flush_batch_size: int = 1

benchmarks/report.py

@@ -26,6 +26,21 @@
 # Raw structural totals, gated on EXACT equality. Load-independent across idle/loaded runs.
 # For a consumer run insert/select are 0; for a producer run delete is 0. Exact-0 ==
 # exact-0 passes, so the same key set gates both scenarios.
+#
+# ``delete_calls`` stays EXACT even for the batched (tfbs>1) point. On the per-row path it
+# is ``messages`` (one DELETE per row); on the batched path it is the flush count, and with
+# fetch_batch_size==terminal_flush_batch_size the fetch delivers full buffers, so it is
+# ``messages / tfbs`` exactly (5000/100 == 50). Empty fetch polls under load inflate total
+# ``calls`` but never ``delete_calls`` (a flush fires only on buffered rows), so this count
+# is structural, not timing-dependent -- Task 4's Step 5 measured it at 50 across 5 runs
+# with zero variance. THREE conditions make it exact here: fetch_batch_size == tfbs (full
+# buffers, no partial-buffer flush); messages % tfbs == 0 (no partial tail flush); AND the
+# bench handler does not suspend mid-drain -- the no-op run_consumer handler never yields, so
+# the worker drains a full generation to the size trigger instead of interleaving with the
+# fetch loop. If a future config broke any -- fetch not divisible by tfbs (split buffers),
+# messages not a multiple of tfbs (a partial final flush, e.g. messages=5050/tfbs=100 -> 51
+# delete_calls), or a bench handler that does real I/O (queue empties mid-buffer -> partial
+# idle flush) -- revisit this to a loose upper bound; today all three hold and it is deterministic.
 EXACT_KEYS: tuple[str, ...] = ("delete_calls", "tup_upd", "tup_del", "tup_ins", "insert_calls", "select_calls")
 
 # Near-deterministic: ~5% observed spread, so a 10% upper-bound band leaves headroom.
@@ -74,8 +89,14 @@ def normalize(result: RunResult) -> dict[str, float]:
 
 
 def _key(result: RunResult) -> str:
+    # The tfbs segment is appended ONLY when batching is enabled (>1) so pre-batching
+    # baseline keys (all tfbs=1) stay byte-identical and are not orphaned; a batched
+    # point (e.g. consumer/w1/b100/tfbs100) never collides with its tfbs=1 sibling.
     cfg = result.config
-    return f"{result.scenario}/w{cfg.max_workers}/b{cfg.fetch_batch_size}"
+    base = f"{result.scenario}/w{cfg.max_workers}/b{cfg.fetch_batch_size}"
+    if cfg.terminal_flush_batch_size != 1:
+        return f"{base}/tfbs{cfg.terminal_flush_batch_size}"
+    return base
 
 
 def to_baseline(results: list[RunResult]) -> dict[str, typing.Any]:

benchmarks/workload.py

@@ -129,6 +129,7 @@ async def run_consumer(engine: AsyncEngine, cfg: RunConfig) -> RunResult:
             # max_fetch_interval whenever _inflight drains. min must be <= max.
             min_fetch_interval=0.001,
             max_fetch_interval=0.001,
+            terminal_flush_batch_size=cfg.terminal_flush_batch_size,
         )
         async def _handler(msg: bytes) -> None:
             seen_ids.add(msg)