Recycle worker connections when the fetch loop reconnects (#263)

After a silent primary loss the fetch loop hit query_timeout_seconds and
reconnected, but each worker then wrote on its own idle connection to the
dead primary and waited out a second timeout. Workers now track a
generation bumped on every fetch-loop error; a worker whose connection
predates it requeues the row, drops the connection and reconnects before
running the handler.

Author: lesnik512

docs/operations/checklist.md

@@ -28,8 +28,8 @@ story.
 - [ ] 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).
+  that went silent takes about that value plus `max_fetch_interval`. 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

@@ -383,9 +383,10 @@ How fast the subscriber recovers depends on how the old primary went away:
 - **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.
+  the loop reconnects. When the fetch loop reconnects, the workers drop their
+  connections too before writing again. Recovery takes about
+  `query_timeout_seconds` plus up to `max_fetch_interval` for an idle subscriber
+  to run its next fetch, which is about 70 s with the defaults.
   A subscriber with a long lease for slow handlers should set a shorter
   `query_timeout_seconds`.
 

faststream_outbox/subscriber/usecase.py

@@ -70,6 +70,10 @@
 _FLUSH_ON_EXIT_TIMEOUT_SECONDS = 2.0
 
 
+class _StaleWriterError(Exception):
+    """The fetch loop reconnected since this worker opened its writer connection."""
+
+
 @dataclasses.dataclass(slots=True)
 class _PendingFlush:
     """A batchable terminal delete awaiting its batch flush, with the metric to emit once it lands."""
@@ -216,6 +220,8 @@ def __init__(
         # from self.running: running must stay True during drain so FastStream's
         # SubscriberUsecase.consume() doesn't early-exit on rows already in _inflight.
         self._stopping: bool = False
+        # Bumped on every fetch-loop error; a worker whose writer predates the bump reconnects.
+        self._fetch_generation: int = 0
 
     @property
     def _client(self) -> "AbstractOutboxClient":
@@ -366,6 +372,7 @@ async def _fetch_loop(self) -> None:
             open_resources=self._open_fetch_resources,
             inner=self._fetch_inner,
             halt_on_drain=True,
+            recycle_writers_on_error=True,
         )
 
     @asynccontextmanager
@@ -586,6 +593,7 @@ async def _run_with_reconnect(
         open_resources: Callable[["AsyncEngine | None"], AbstractAsyncContextManager[Mapping[str, object]]],
         inner: Callable[..., Awaitable[None]],
         halt_on_drain: bool = False,
+        recycle_writers_on_error: bool = False,
     ) -> None:
         """Reconnect-with-backoff scaffold shared by ``_fetch_loop`` and ``_worker_loop``.
 
@@ -614,7 +622,14 @@ async def _run_with_reconnect(
                 async with open_resources(client.engine) as kwargs:
                     started = time.monotonic()
                     await inner(**kwargs)
+            except _StaleWriterError:
+                self._log(
+                    log_level=logging.INFO,
+                    message=f"Outbox {name} loop reconnecting: the fetch loop lost its connection",
+                )
             except Exception as e:  # noqa: BLE001
+                if recycle_writers_on_error:
+                    self._fetch_generation += 1
                 # Reset the backoff counter when the connection was healthy for a
                 # sustained window before failing — otherwise the lifetime error count
                 # accrues and every transient blip after the first handful costs the
@@ -638,6 +653,7 @@ async def _worker_inner(self, *, writer_conn: "AsyncConnection | None") -> None:
         actual delete, not merely for dispatch.
         """
         buffer: list[_PendingFlush] = []
+        generation = self._fetch_generation
         batch_size = self._config.terminal_flush_batch_size
         max_buffer_age = self._config.lease_ttl_seconds / 2
         try:
@@ -650,6 +666,12 @@ async def _worker_inner(self, *, writer_conn: "AsyncConnection | None") -> None:
                         row = await self._inflight.get()
                 else:
                     row = await self._inflight.get()
+                if writer_conn is not None and self._fetch_generation != generation:
+                    self._inflight.put_nowait(row)
+                    self._inflight.task_done()
+                    await self._flush_buffer(buffer, writer_conn=writer_conn)
+                    await writer_conn.invalidate()
+                    raise _StaleWriterError
                 if self._inflight.maxsize - self._inflight.qsize() >= self._refill_at:
                     self._notify_event.set()
                 buffered = False

tests/test_integration.py

@@ -890,6 +890,73 @@ async def handle(body: dict) -> None:
     assert received == [{"order_id": 1}]
 
 
+async def test_worker_recycles_writer_connection_after_fetch_loop_reconnects(pg_engine, outbox_table) -> None:
+    """After the fetch loop loses its connection, the worker must not write on its own stale one.
+
+    A silently dropped primary leaves the idle writer connection dead too; reusing it costs a
+    second ``query_timeout_seconds`` before the first post-failover row can be acknowledged.
+    """
+    received: list[dict] = []
+    delete_pids: list[int] = []
+    hang_next_fetch: list[bool] = []
+    broker = OutboxBroker(pg_engine, outbox_table=outbox_table, graceful_timeout=2.0)
+    original_fetch = OutboxClient.fetch
+    original_delete = OutboxClient.delete_with_lease
+
+    async def hang_on_demand(
+        self: OutboxClient,
+        conn: AsyncConnection | None,
+        queues: Sequence[str],
+        *,
+        limit: int,
+        lease_ttl_seconds: float,
+    ) -> list[OutboxInnerMessage]:
+        if hang_next_fetch:
+            hang_next_fetch.clear()
+            await asyncio.Event().wait()
+        return await original_fetch(self, conn, queues, limit=limit, lease_ttl_seconds=lease_ttl_seconds)
+
+    async def record_pid(
+        self: OutboxClient,
+        conn: AsyncConnection | None,
+        row_id: int,
+        token: uuid.UUID,
+        *,
+        dlq_payload: Mapping[str, Any] | None = None,
+    ) -> bool:
+        assert conn is not None
+        delete_pids.append((await conn.exec_driver_sql("SELECT pg_backend_pid()")).scalar_one())
+        return await original_delete(self, conn, row_id, token, dlq_payload=dlq_payload)
+
+    @broker.subscriber(
+        "orders",
+        min_fetch_interval=0.05,
+        max_fetch_interval=0.2,
+        lease_ttl_seconds=30.0,
+        query_timeout_seconds=0.5,
+    )
+    async def handle(body: dict) -> None:
+        received.append(body)
+
+    session_factory = async_sessionmaker(pg_engine, expire_on_commit=False)
+    with (
+        mock.patch.object(OutboxClient, "fetch", hang_on_demand),
+        mock.patch.object(OutboxClient, "delete_with_lease", record_pid),
+    ):
+        async with broker:
+            async with session_factory() as session, session.begin():
+                await broker.publish({"order_id": 1}, queue="orders", session=session)
+            await _wait_until(lambda: len(delete_pids) == 1, timeout=5.0)
+            hang_next_fetch.append(True)
+            await _wait_until(lambda: not hang_next_fetch, timeout=5.0)
+            async with session_factory() as session, session.begin():
+                await broker.publish({"order_id": 2}, queue="orders", session=session)
+            await _wait_until(lambda: len(delete_pids) == 2, timeout=10.0)
+
+    assert received == [{"order_id": 1}, {"order_id": 2}]
+    assert delete_pids[0] != delete_pids[1]
+
+
 async def test_hung_terminal_write_times_out_and_worker_reconnects(pg_engine, outbox_table) -> None:
     """A terminal DELETE stuck on a silently dropped socket must not wedge the worker forever."""
     received: list[dict] = []