Bound subscriber statements so loops reconnect after a silent primary loss (#262)

* Bound subscriber statements so loops reconnect after a silent primary loss

A fetch or terminal write on a connection whose primary vanished without
closing it (power loss, network partition) hung until the kernel's TCP
timeout, so the subscriber stopped consuming after a failover. Each fetch,
delete, batch delete and retry update is now capped at the new
query_timeout_seconds subscriber option (default: lease_ttl_seconds); on
timeout the connection is invalidated and the loop reconnects.

Adds failover drills (just test-dr): a primary/streaming-replica compose
cluster where the primary is either killed or cut off from the network and
the replica is promoted. Not part of CI.

* Skip failover drills unless DR_CLUSTER_DSN is set

The floors job runs pytest . and sets POSTGRES_DSN, so the drills were
collected and tried to connect to a cluster that does not exist.

Author: lesnik512

AGENTS.md

@@ -19,6 +19,9 @@ the vocabulary. Read it before naming a concept in code, a test name, or an issu
   pytest`; `tests/test_integration.py` skips if `POSTGRES_DSN` is unreachable. The compose Postgres
   publishes no host port, so run the DB tests through `just test`. The coverage gate is on by
   default, so a partial run trips it; pass `--no-cov` while iterating.
+- `just test-dr` runs the failover drills in `tests_dr/` (primary killed, primary cut off from the
+  network), each on a fresh primary/replica cluster from `docker-compose.dr.yml`. It needs the Docker
+  socket, takes a few minutes, and is not in CI.
 - Nothing validates Markdown links outside `docs/`. `just docs-build` runs `mkdocs --strict` over
   the site only; root Markdown, `.github/`, and `docs/agents/` are unchecked.
 

docker-compose.dr.yml

@@ -0,0 +1,52 @@
+name: outbox-dr
+
+services:
+  dr:
+    build:
+      context: .
+      dockerfile: ./Dockerfile
+    user: root
+    volumes:
+      - .:/code
+      - /code/.venv
+      - /var/run/docker.sock:/var/run/docker.sock
+    depends_on:
+      pg-replica:
+        condition: service_healthy
+    environment:
+      - DR_CLUSTER_DSN=postgresql+asyncpg://outbox:outbox@/outbox?host=pg-primary:5432&host=pg-replica:5432&target_session_attrs=primary
+      - PRIMARY_DSN=postgresql://outbox:outbox@pg-primary:5432/outbox
+      - REPLICA_DSN=postgresql://outbox:outbox@pg-replica:5432/outbox
+
+  pg-primary:
+    image: postgres:17
+    container_name: outbox-dr-primary
+    environment:
+      - POSTGRES_USER=outbox
+      - POSTGRES_PASSWORD=outbox
+      - POSTGRES_DB=outbox
+    volumes:
+      - ./tests_dr/primary-init.sh:/docker-entrypoint-initdb.d/primary-init.sh:ro
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U outbox -d outbox"]
+      interval: 1s
+      timeout: 3s
+      retries: 30
+
+  pg-replica:
+    image: postgres:17
+    container_name: outbox-dr-replica
+    user: postgres
+    entrypoint: ["bash", "/replica-entrypoint.sh"]
+    environment:
+      - PGPASSWORD=outbox
+    volumes:
+      - ./tests_dr/replica-entrypoint.sh:/replica-entrypoint.sh:ro
+    depends_on:
+      pg-primary:
+        condition: service_healthy
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U outbox -d outbox"]
+      interval: 1s
+      timeout: 3s
+      retries: 60

docs/operations/checklist.md

@@ -25,6 +25,11 @@ story.
   cutoff is server-side `make_interval(...)`, immune to clock skew.
   Tuning: [Subscriber § Slow handlers: dedicated
   queue](../usage/subscriber.md#slow-handlers-dedicated-queue).
+- [ ] Engine URL follows the primary across a failover (a proxy, or a
+  multi-host URL with `target_session_attrs=primary`). Subscribers with a
+  long lease set a shorter `query_timeout_seconds`; recovery from a primary
+  that went silent takes about twice that value. See [Subscriber § Database
+  failover](../usage/subscriber.md#database-failover).
 - [ ] Slow handlers segregated onto their own subscriber with a
   taller `lease_ttl_seconds`. Don't raise it globally; that delays
   reclaim of *actually* stuck rows everywhere.

docs/usage/subscriber.md

@@ -90,6 +90,7 @@ Per-subscriber knobs, passed to `@broker.subscriber("…", …)`:
 | `min_fetch_interval` | `1.0` s | Base for the adaptive idle backoff (jittered ±50%, so an actual wait can land below it); no sleep at all while fetches keep returning rows |
 | `max_fetch_interval` | `10.0` s | Ceiling for the adaptive idle backoff (with jitter) |
 | `lease_ttl_seconds` | `60.0` s | How long a claim is valid before another fetch may reclaim it. Must exceed your handler's P99 with margin. |
+| `query_timeout_seconds` | `None` (uses `lease_ttl_seconds`) | Client-side cap on each fetch and terminal write. A statement still running when it expires is abandoned, its connection is closed, and the loop reconnects. Set it below a long lease so a lost primary is noticed sooner. Must be `> 0` and at most `lease_ttl_seconds`. See [Database failover](#database-failover). |
 | `max_deliveries` | `None` (unbounded) | Total claims (including lease-expiry re-claims) after which the row is dropped without invoking the handler. Defends against handlers that consistently wedge. Must be at least `1`; use `None`, not `0`, for no cap. |
 | `terminal_flush_batch_size` | `1` (off) | Coalesce completed terminal `DELETE`s into one `DELETE … RETURNING` per N rows. `1` is one round-trip per message (unchanged). Higher trades a wider crash-redelivery window for far fewer round-trips. See [Batching terminal deletes](#batching-terminal-deletes). |
 | `ack_policy` | `AckPolicy.NACK_ON_ERROR` | See [Ack policy](#ack-policy) |
@@ -360,6 +361,39 @@ startup with `FATAL: too many connections`.
 
 *Operator-side: [Production checklist § Sizing](../operations/checklist.md#sizing).*
 
+## Database failover
+
+Point the engine at whichever node is currently primary, either through a proxy
+or virtual IP, or with a multi-host URL that asks asyncpg for the writable node:
+
+```python
+engine = create_async_engine(
+    "postgresql+asyncpg://user:pass@/app?host=pg-1:5432&host=pg-2:5432&target_session_attrs=primary",
+)
+```
+
+Without `target_session_attrs=primary`, a reconnect can land on a standby, where
+the fetch fails with a read-only error.
+
+How fast the subscriber recovers depends on how the old primary went away:
+
+- **It closed its connections** (the process crashed or was killed). The next
+  statement fails at once, and the fetch and worker loops back off and
+  reconnect within seconds.
+- **It went silent** (the host lost power, or the network was partitioned).
+  Statements hang because no error ever arrives. Each fetch and terminal write
+  is capped at `query_timeout_seconds`, after which the connection is dropped and
+  the loop reconnects. The fetch loop hits that cap first; the worker hits it
+  again on its own stale connection once it gets a row. Recovery takes about
+  `2 × query_timeout_seconds`, which is two minutes with the default 60 s lease.
+  A subscriber with a long lease for slow handlers should set a shorter
+  `query_timeout_seconds`.
+
+With asynchronous replication the new primary can miss the last commits on the
+old one. A lost insert disappears together with the business rows from the same
+transaction. A lost terminal `DELETE` brings a handled row back, so it is
+delivered again; handlers are idempotent for that reason anyway.
+
 ## Read-only inspection
 
 `subscriber.get_one()` and `async for msg in subscriber:` are not

faststream_outbox/fastapi/router.py

@@ -167,6 +167,7 @@ def subscriber(  # ty: ignore[invalid-method-override]  # noqa: PLR0913
         min_fetch_interval: float = 1.0,
         max_fetch_interval: float = 10.0,
         lease_ttl_seconds: float = 60.0,
+        query_timeout_seconds: float | None = None,
         max_deliveries: int | None = None,
         terminal_flush_batch_size: int = 1,
         ack_policy: AckPolicy | None = None,
@@ -200,6 +201,7 @@ def subscriber(  # ty: ignore[invalid-method-override]  # noqa: PLR0913
                 min_fetch_interval=min_fetch_interval,
                 max_fetch_interval=max_fetch_interval,
                 lease_ttl_seconds=lease_ttl_seconds,
+                query_timeout_seconds=query_timeout_seconds,
                 max_deliveries=max_deliveries,
                 terminal_flush_batch_size=terminal_flush_batch_size,
                 ack_policy=ack_policy,

faststream_outbox/registrator.py

@@ -33,6 +33,7 @@ def subscriber(  # ty: ignore[invalid-method-override]
         min_fetch_interval: float = 1.0,
         max_fetch_interval: float = 10.0,
         lease_ttl_seconds: float = 60.0,
+        query_timeout_seconds: float | None = None,
         max_deliveries: int | None = None,
         terminal_flush_batch_size: int = 1,
         ack_policy: AckPolicy | None = None,
@@ -56,6 +57,7 @@ def subscriber(  # ty: ignore[invalid-method-override]
             min_fetch_interval=min_fetch_interval,
             max_fetch_interval=max_fetch_interval,
             lease_ttl_seconds=lease_ttl_seconds,
+            query_timeout_seconds=query_timeout_seconds,
             max_deliveries=max_deliveries,
             terminal_flush_batch_size=terminal_flush_batch_size,
             ack_policy=ack_policy,

faststream_outbox/router.py

@@ -31,6 +31,7 @@ def __init__(  # noqa: PLR0913
         min_fetch_interval: float = 1.0,
         max_fetch_interval: float = 10.0,
         lease_ttl_seconds: float = 60.0,
+        query_timeout_seconds: float | None = None,
         max_deliveries: int | None = None,
         terminal_flush_batch_size: int = 1,
         ack_policy: AckPolicy | None = None,
@@ -51,6 +52,7 @@ def __init__(  # noqa: PLR0913
             min_fetch_interval=min_fetch_interval,
             max_fetch_interval=max_fetch_interval,
             lease_ttl_seconds=lease_ttl_seconds,
+            query_timeout_seconds=query_timeout_seconds,
             max_deliveries=max_deliveries,
             terminal_flush_batch_size=terminal_flush_batch_size,
             ack_policy=ack_policy,

faststream_outbox/subscriber/config.py

@@ -74,6 +74,7 @@ class OutboxSubscriberConfig(SubscriberUsecaseConfig):
     min_fetch_interval: float
     max_fetch_interval: float
     lease_ttl_seconds: float
+    query_timeout_seconds: float | None = None
     max_deliveries: int | None
     terminal_flush_batch_size: int
     propagate_inbound_headers: bool
@@ -127,6 +128,7 @@ def _validate(self) -> None:  # noqa: C901  # flat sequence of independent knob
         if self.lease_ttl_seconds <= 0:
             msg = f"lease_ttl_seconds must be > 0, got {self.lease_ttl_seconds}"
             raise ValueError(msg)
+        self._validate_query_timeout()
         if self.min_fetch_interval > self.max_fetch_interval:
             msg = (
                 f"min_fetch_interval ({self.min_fetch_interval}) must be <= max_fetch_interval "
@@ -169,6 +171,20 @@ def _validate(self) -> None:  # noqa: C901  # flat sequence of independent knob
                 f"2 * max_fetch_interval + P99(handler).",
             )
 
+    def _validate_query_timeout(self) -> None:
+        if self.query_timeout_seconds is None:
+            return
+        if self.query_timeout_seconds <= 0:
+            msg = f"query_timeout_seconds must be > 0, got {self.query_timeout_seconds}"
+            raise ValueError(msg)
+        if self.query_timeout_seconds > self.lease_ttl_seconds:
+            msg = (
+                f"query_timeout_seconds ({self.query_timeout_seconds}) must be <= lease_ttl_seconds "
+                f"({self.lease_ttl_seconds}): a fetch or terminal write still running past the lease "
+                f"has already lost the row to reclaim."
+            )
+            raise ValueError(msg)
+
     def _validate_numbers(self) -> None:
         """Reject bools passed as counts, NaN or infinite intervals, and a ``max_deliveries`` below 1."""
         for name in ("max_workers", "fetch_batch_size", "terminal_flush_batch_size", "max_deliveries"):
@@ -178,14 +194,19 @@ def _validate_numbers(self) -> None:
             if isinstance(value, bool) or not isinstance(value, int):
                 msg = f"{name} must be an int, got {value!r}"
                 raise TypeError(msg)
-        for name in ("min_fetch_interval", "max_fetch_interval", "lease_ttl_seconds"):
-            if not math.isfinite(getattr(self, name)):
+        for name in ("min_fetch_interval", "max_fetch_interval", "lease_ttl_seconds", "query_timeout_seconds"):
+            if getattr(self, name) is not None and not math.isfinite(getattr(self, name)):
                 msg = f"{name} must be a finite number, got {getattr(self, name)}"
                 raise ValueError(msg)
         if self.max_deliveries is not None and self.max_deliveries < 1:
             msg = f"max_deliveries must be >= 1 or None (unlimited), got {self.max_deliveries}"
             raise ValueError(msg)
 
+    @property
+    def query_timeout(self) -> float:
+        """Client-side cap on each fetch and terminal write; ``lease_ttl_seconds`` unless set."""
+        return self.lease_ttl_seconds if self.query_timeout_seconds is None else self.query_timeout_seconds
+
     @property
     def ack_policy(self) -> AckPolicy:
         if self._ack_policy is EMPTY:

faststream_outbox/subscriber/factory.py

@@ -13,7 +13,7 @@
     from faststream_outbox.retry import RetryStrategyProto
 
 
-def create_subscriber(
+def create_subscriber(  # noqa: PLR0913
     *,
     queues: list[str],
     max_workers: int,
@@ -25,6 +25,7 @@ def create_subscriber(
     max_deliveries: int | None,
     terminal_flush_batch_size: int,
     config: "OutboxBrokerConfig",
+    query_timeout_seconds: float | None = None,
     ack_policy: AckPolicy | None = None,
     propagate_inbound_headers: bool = False,
     title_: str | None = None,
@@ -43,6 +44,7 @@ def create_subscriber(
         min_fetch_interval=min_fetch_interval,
         max_fetch_interval=max_fetch_interval,
         lease_ttl_seconds=lease_ttl_seconds,
+        query_timeout_seconds=query_timeout_seconds,
         max_deliveries=max_deliveries,
         terminal_flush_batch_size=terminal_flush_batch_size,
         propagate_inbound_headers=propagate_inbound_headers,

faststream_outbox/subscriber/usecase.py

@@ -103,6 +103,8 @@ class _PendingFlush:
 # the kernel keepalive; cap it and fall back to an immediate terminate.
 _LISTEN_CLOSE_TIMEOUT = 5.0
 
+_T = typing.TypeVar("_T")
+
 
 def _compute_backoff(attempt: int, ceiling: float, *, base: float = 1.0) -> float:
     """Exponential backoff with ±50% jitter, capped at *ceiling*.
@@ -427,11 +429,14 @@ async def _fetch_inner(
                 # fetch_conn is None only in the test-broker path against FakeOutboxClient
                 # (which ignores conn); the real OutboxClient raises if conn is None. The
                 # AbstractOutboxClient surface admits None so both implementations type-check.
-                rows = await self._client.fetch(
+                rows = await self._bounded(
                     fetch_conn,
-                    self._queues,
-                    limit=free,
-                    lease_ttl_seconds=self._config.lease_ttl_seconds,
+                    self._client.fetch(
+                        fetch_conn,
+                        self._queues,
+                        limit=free,
+                        lease_ttl_seconds=self._config.lease_ttl_seconds,
+                    ),
                 )
                 for row in rows:
                     self._inflight.put_nowait(row)
@@ -528,6 +533,20 @@ def _on_notify(self, *args: object) -> None:
         if payload is None or payload in self._served_queues:
             self._notify_event.set()
 
+    async def _bounded(self, conn: "AsyncConnection | None", statement: Awaitable[_T]) -> _T:
+        """Cap a fetch/terminal write at ``query_timeout`` so a silently dropped socket raises.
+
+        On timeout the connection is invalidated (closed without a round-trip on the dead
+        socket) and the error propagates, so the outer loop reconnects.
+        """
+        try:
+            async with asyncio.timeout(self._config.query_timeout):
+                return await statement
+        except TimeoutError:
+            if conn is not None:
+                await conn.invalidate()
+            raise
+
     async def _worker_loop(self) -> None:
         """Thin wrapper around :meth:`_run_with_reconnect` for the worker path."""
         await self._run_with_reconnect(
@@ -862,7 +881,7 @@ async def _flush_buffer(
             return
         pairs = [(p.row.id, p.row.acquired_token) for p in buffer if p.row.acquired_token is not None]
         try:
-            deleted = await self._client.delete_batch_with_lease(writer_conn, pairs)
+            deleted = await self._bounded(writer_conn, self._client.delete_batch_with_lease(writer_conn, pairs))
         except BaseException:  # incl. CancelledError from the exit-flush timeout: balance task_done before re-raising
             for pending in buffer:
                 self._emit_dispatch_abandoned(pending.row)
@@ -944,11 +963,9 @@ async def _flush_terminal(
                     self._outer_config.last_exception_renderer,
                 ),
             }
-        deleted = await self._client.delete_with_lease(
+        deleted = await self._bounded(
             writer_conn,
-            row.id,
-            row.acquired_token,
-            dlq_payload=dlq_payload,
+            self._client.delete_with_lease(writer_conn, row.id, row.acquired_token, dlq_payload=dlq_payload),
         )
         if not deleted:
             self._emit_lease_lost(row, phase="terminal")
@@ -977,14 +994,17 @@ async def _flush_retry(
         # P19: first_attempt_at / last_attempt_at are the worker's clock; next_attempt_at
         # is computed server-side (now() + delay) inside mark_pending_with_lease, so the
         # reschedule time stays clock-skew-immune even though these audit timestamps don't.
-        updated = await self._client.mark_pending_with_lease(
+        updated = await self._bounded(
             writer_conn,
-            row.id,
-            row.acquired_token,
-            delay_seconds=row.pending_delay_seconds,
-            attempts_count=row.attempts_count,
-            first_attempt_at=row.first_attempt_at,  # ty: ignore[invalid-argument-type]
-            last_attempt_at=row.last_attempt_at,  # ty: ignore[invalid-argument-type]
+            self._client.mark_pending_with_lease(
+                writer_conn,
+                row.id,
+                row.acquired_token,
+                delay_seconds=row.pending_delay_seconds,
+                attempts_count=row.attempts_count,
+                first_attempt_at=row.first_attempt_at,  # ty: ignore[invalid-argument-type]
+                last_attempt_at=row.last_attempt_at,  # ty: ignore[invalid-argument-type]
+            ),
         )
         if not updated:
             self._emit_lease_lost(row, phase="retry")