fix(worker): address XBeg9's PR #39 follow-up findings

src/taskq/settings.py

@@ -372,11 +372,13 @@ class WorkerSettings(TaskQSettings):
         "wrapped around each period-1 leader-loop iteration (scheduled_wake, "
         "cron): a stalled PG errors the iteration instead of hanging the "
         "loop past its staleness budget. Checked at load time when the "
-        "watchdog is enabled: timeout + loop period must be < "
-        "max(period x watchdog_tick_grace_factor, watchdog_stale_floor) for "
-        "every bounded loop (leader period-1 loops and the producer), so a "
+        "watchdog is enabled: timeout + the 1.0s leader-loop period must be "
+        "< max(period x watchdog_tick_grace_factor, watchdog_stale_floor) "
+        "for the period-1 leader loops (scheduled_wake, cron), so a "
         "timeout-capped iteration can never false-trip the stale-loop "
-        "detector on a healthy worker.",
+        "detector on a healthy worker. The producer loop is not checked "
+        "(its multi-statement dispatch_batch is not wrapped in a single "
+        "asyncio.timeout).",
     )
     dispatch_oversample: int = Field(
         default=2,
@@ -1004,25 +1006,26 @@ def post_load(self) -> list[ValidationError] | None:
                 )
             )
 
-        # Bounded-loop staleness invariant: every loop whose PG work is
-        # capped by dispatcher_command_timeout ticks once per iteration and
-        # sleeps one period afterwards, so its worst-case tick gap is
-        # timeout + period. That gap must fit the loop's own budget
-        # max(period * watchdog_tick_grace_factor, watchdog_stale_floor) or
-        # detector 2 force-exits a healthy worker mid-degradation
-        # (measured: timeout 10.0 against budget 10.0 produced an 11s tick
-        # gap and a trip at age 10.008s). Only checked when the watchdog is
-        # armed: with watchdog_enabled=False detector 2 is never spawned,
-        # and a stale tick only costs a transient NotReady, which is not
-        # worth blocking boot over.
+        # Bounded-loop staleness invariant: the period-1 leader loops
+        # (scheduled_wake, cron) are wrapped in asyncio.timeout, so their
+        # worst-case tick gap is timeout + period. That gap must fit the
+        # loop's own budget max(period * watchdog_tick_grace_factor,
+        # watchdog_stale_floor) or detector 2 force-exits a healthy worker
+        # mid-degradation (measured: timeout 10.0 against budget 10.0
+        # produced an 11s tick gap and a trip at age 10.008s). Only checked
+        # when the watchdog is armed: with watchdog_enabled=False detector 2
+        # is never spawned, and a stale tick only costs a transient NotReady,
+        # which is not worth blocking boot over.
+        #
+        # The producer loop is deliberately NOT checked here: it is not
+        # wrapped in asyncio.timeout (dispatch_batch is a multi-statement
+        # transaction — BEGIN + resolve_queue_modes + dispatch CTE + INSERTs
+        # + COMMIT, each bounded separately by the pool's command_timeout),
+        # so the timeout + period model does not hold. The actual worst-case
+        # gap is k * timeout + period for k statements, which the invariant
+        # cannot express without knowing k at settings-load time.
         if self.watchdog_enabled:
-            producer_period = (
-                self.notify_poll_interval if self.notify_enabled else self.poll_interval
-            )
-            for loop_label, period in (
-                ("leader loops", 1.0),
-                ("producer loop", producer_period),
-            ):
+            for loop_label, period in (("leader loops", 1.0),):
                 budget = max(period * self.watchdog_tick_grace_factor, self.watchdog_stale_floor)
                 if budget <= period + 1.0:
                     # 1.0 = dispatcher_command_timeout's own ge= minimum: no

src/taskq/worker/_transient.py

@@ -3,8 +3,7 @@
 Every long-lived worker loop that awaits Postgres treats the same set of
 errors as "PG is having a moment; log and retry next tick". That set used
 to be re-enumerated at every call site, and it drifted: heartbeat learned
-``QueryCanceledError`` (the server-side shape of a fired
-``command_timeout``) while the leader loops kept only the OSError
+``QueryCanceledError`` (server-side 57014 cancellation) while the leader loops kept only the OSError
 flavours, so a degraded PG could throw an uncaught error into the worker
 TaskGroup and tear the whole worker down mid-blip. One tuple, one home:
 any shape a site learns, every site learns.
@@ -33,8 +32,10 @@
 #:   ``ConnectionDoesNotExistError`` / ``ConnectionFailureError``).
 #: - ``InterfaceError`` / ``OSError``: the connection is unusable or the
 #:   socket died.
-#: - ``QueryCanceledError``: server-side 57014, asyncpg's OTHER shape for
-#:   a fired ``command_timeout`` (the driver's cancel request landed).
+#: - ``QueryCanceledError``: server-side 57014 — a DBA ran
+#:   pg_cancel_backend, or a server-side ``statement_timeout`` fired. Not
+#:   a client-side ``command_timeout`` (that raises ``TimeoutError``); kept
+#:   in the tuple because a server-side cancel is equally transient.
 #: - ``AdminShutdownError``: 57P01, PG restart/shutdown. An
 #:   OperatorInterventionError, NOT a PostgresConnectionError: notify.py
 #:   learned this one the hard way.

src/taskq/worker/deps.py

@@ -127,7 +127,7 @@ async def open_dedicated_conn(
     hang the caller indefinitely.
     """
     if command_timeout is not None:
-        conn = await asyncpg.connect(dsn, command_timeout=command_timeout)
+        conn = await asyncpg.connect(dsn, command_timeout=command_timeout, timeout=command_timeout)
     else:
         conn = await asyncpg.connect(dsn)
     applied = apply_keepalive_to_conn(conn, label=label) if apply_keepalive else False

src/taskq/worker/leader.py

@@ -292,7 +292,11 @@ async def _election_loop(self, shutdown: asyncio.Event) -> None:
                         # Backstop (see _transient.py): tolerated + logged a
                         # few times, then deliberately fatal; cleanup mirrors
                         # the transient path since conn state is unknown.
-                        guard.unexpected(exc)
+                        # Cleanup runs BEFORE guard.unexpected so the fatal
+                        # iteration still drops the conn and clears
+                        # is_leader — otherwise the dead leader gauge and
+                        # leader_conn reference stay stale until run()'s
+                        # finally.
                         await self._drop_leader_conn(reason="probe_failed")
                         await self._close_leader_owned_conns()
                         log.warning(
@@ -301,6 +305,8 @@ async def _election_loop(self, shutdown: asyncio.Event) -> None:
                             worker_id=str(self._worker_id),
                             error=repr(exc),
                         )
+                        guard.unexpected(exc)
+                        continue
             if self._deps.leader_conn is None or self._deps.leader_conn.is_closed():
                 try:
                     self._deps.leader_conn = await self._open_leader_conn()
@@ -345,7 +351,6 @@ async def _election_loop(self, shutdown: asyncio.Event) -> None:
                 # Backstop (see _transient.py): tolerated + logged a few
                 # times, then deliberately fatal; cleanup mirrors the
                 # transient path since conn state is unknown.
-                guard.unexpected(exc)
                 await self._drop_leader_conn(reason="lock_attempt_failed")
                 await self._close_leader_owned_conns()
                 record_election_attempt(str(self._worker_id), won=False)
@@ -355,6 +360,7 @@ async def _election_loop(self, shutdown: asyncio.Event) -> None:
                     worker_id=str(self._worker_id),
                     error=repr(exc),
                 )
+                guard.unexpected(exc)
                 await asyncio.sleep(self._deps.settings.heartbeat_interval)
                 continue
             if got_lock:
@@ -402,7 +408,6 @@ async def _election_loop(self, shutdown: asyncio.Event) -> None:
                     # Backstop (see _transient.py): tolerated + logged a few
                     # times, then deliberately fatal; cleanup mirrors the
                     # transient path since conn state is unknown.
-                    guard.unexpected(exc)
                     await self._drop_leader_conn(reason="leader_upsert_failed")
                     await self._close_leader_owned_conns()
                     record_election_attempt(str(self._worker_id), won=False)
@@ -412,6 +417,7 @@ async def _election_loop(self, shutdown: asyncio.Event) -> None:
                         worker_id=str(self._worker_id),
                         error=repr(exc),
                     )
+                    guard.unexpected(exc)
                     await asyncio.sleep(self._deps.settings.heartbeat_interval)
                     continue
                 try:
@@ -504,7 +510,6 @@ async def _watchdog_loop(self, shutdown: asyncio.Event) -> None:
                     # Backstop (see _transient.py): tolerated + logged a few
                     # times, then deliberately fatal; cleanup mirrors the
                     # transient path since conn state is unknown.
-                    guard.unexpected(exc)
                     await self._drop_leader_conn(reason="watchdog_probe_failed")
                     await self._close_leader_owned_conns()
                     log.warning(
@@ -513,6 +518,7 @@ async def _watchdog_loop(self, shutdown: asyncio.Event) -> None:
                         worker_id=str(self._worker_id),
                         error=repr(exc),
                     )
+                    guard.unexpected(exc)
                     break
                 await asyncio.sleep(_WATCHDOG_INTERVAL_SECS)
             # Leaving the inner loop means the gate closed: demotion (the
@@ -616,30 +622,29 @@ async def _cron_loop(self, shutdown: asyncio.Event) -> None:
                             self._worker_id,
                         )
                 guard.ok()
-            except Exception as exc:
-                # Why the ordering and the exact-type check: builtin
-                # TimeoutError IS an OSError subclass, so an isinstance
-                # conn-state check would swallow the deadline shapes before
-                # they are seen. asyncio.timeout and asyncpg's local
-                # command_timeout expiry raise the exact builtin; network
-                # death arrives as ConnectionTimeoutError (a subclass) and
-                # correctly falls to the conn-state branch below.
+            except TRANSIENT_PG_ERRORS as exc:
+                # Why TRANSIENT_PG_ERRORS first: the cron loop used to
+                # hand-roll its error classification with isinstance checks
+                # that missed 7 of 12 transient shapes (DeadlockDetectedError,
+                # SerializationError, AdminShutdownError, etc.). Deadlock and
+                # serialization inside a transaction are routine, not
+                # surprises — 5 consecutive killed the worker via the
+                # backstop guard before this fix.
                 if type(exc) is TimeoutError or isinstance(exc, asyncpg.QueryCanceledError):
-                    # Deadline family (iteration deadline or a fired
-                    # command_timeout): the conn is provably responsive,
-                    # because it answered the cancel, or asyncpg has already
-                    # terminated it, which the next tick's transaction()
-                    # surfaces as a conn-state error below. Keep the conn and
-                    # retry: dropping it (and demoting) on every slow tick
-                    # would churn leadership during catch-up bursts.
+                    # Deadline family (iteration deadline or server-side
+                    # cancel): the conn is provably responsive, because it
+                    # answered the cancel, or asyncpg has already terminated
+                    # it, which the next tick's transaction() surfaces as a
+                    # conn-state error below. Keep the conn and retry:
+                    # dropping it (and demoting) on every slow tick would
+                    # churn leadership during catch-up bursts.
                     log.warning(
                         "cron-tick-timeout",
                         kind="cron_tick_timeout",
                         worker_id=str(self._worker_id),
                         error=repr(exc),
                     )
-                    continue
-                if isinstance(
+                elif isinstance(
                     exc, (asyncpg.PostgresConnectionError, asyncpg.InterfaceError, OSError)
                 ):
                     # Conn-state family: the conn is dead or unusable. Drop
@@ -651,7 +656,21 @@ async def _cron_loop(self, shutdown: asyncio.Event) -> None:
                         worker_id=str(self._worker_id),
                         error=repr(exc),
                     )
-                    continue
+                else:
+                    # Other transient (deadlock, serialization, admin
+                    # shutdown, cannot-connect-now, too-many-connections,
+                    # idle-session timeouts): retry next tick. The conn may
+                    # or may not be dead — the next tick's transaction()
+                    # will surface a conn-state error if it is, and the
+                    # transaction rolls back for deadlock/serialization
+                    # leaving the conn usable.
+                    log.warning(
+                        "cron-tick-transient",
+                        kind="cron_tick_transient",
+                        worker_id=str(self._worker_id),
+                        error=repr(exc),
+                    )
+            except Exception as exc:
                 # Backstop for anything outside the transient set (see
                 # _transient.py): tolerated and logged a few times (this
                 # loop's historical blanket catch), then deliberately fatal