From 1250fdf36fa6ffbeb1a63c6e232053277ce4c3d3 Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Thu, 22 Oct 2026 16:15:48 +0530 Subject: [PATCH 1/4] UN-4254 [PERF] Run PG-queue ops in autocommit and slow idle polling The PG-queue client and result backend ran every single-statement operation inside psycopg2's implicit BEGIN ... COMMIT, so each claim (including empty polls), lease renewal, ack, enqueue and result read/write cost three round trips. Their owned connections now run in autocommit and send only the statement. create_pg_connection() itself is unchanged: pg_barrier's reconnect-retry relies on it staying non-autocommit, and the reaper, scheduler and leader election run multi-statement transactions on their own connections. Injected connections keep the caller's mode. Idle consumers polled every 0.1s after a task and backed off only to 2s, so idle DB load scaled with slots x queues. Defaults are now POLL_INTERVAL=0.5 and BACKOFF_MAX=5.0 (still env-tunable); BACKOFF_MAX stays well under HEALTH_STALE_SECONDS. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01KGgUm6WuXrgzYq1WhrFcX7 --- workers/queue_backend/pg_queue/README.md | 4 +- workers/queue_backend/pg_queue/client.py | 26 +++++++++-- workers/queue_backend/pg_queue/consumer.py | 12 ++++- .../queue_backend/pg_queue/result_backend.py | 11 ++++- workers/tests/test_pg_queue_client.py | 46 +++++++++++++++++++ workers/tests/test_pg_result_backend.py | 22 +++++++++ 6 files changed, 111 insertions(+), 10 deletions(-) diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md index d2e755251f..6b6a8b252d 100644 --- a/workers/queue_backend/pg_queue/README.md +++ b/workers/queue_backend/pg_queue/README.md @@ -31,8 +31,8 @@ integers. | `SHUTDOWN_GRACE_SECONDS` | Graceful-drain budget (shared across all children) on SIGTERM before SIGKILL | `= VT` (floored at `30`) | | `QUEUE` | Queue name(s) this consumer polls (comma-separated) | `default` | | `BATCH` | Messages claimed per poll | `1` | -| `POLL_INTERVAL` | Time between polls when the queue is empty | `0.1` | -| `BACKOFF_MAX` | Max empty-queue poll backoff | `2.0` | +| `POLL_INTERVAL` | First sleep after an empty poll (doubles up to `BACKOFF_MAX`; a claim resets it) | `0.5` | +| `BACKOFF_MAX` | Max empty-queue poll backoff — also the worst-case pickup latency on an idle queue with a single poller. Keep well under `HEALTH_STALE_SECONDS` | `5.0` | | `MAX_ATTEMPTS` | Max deliveries before a message is dropped as **poison** | `5` | | `HEALTH_PORT` | Liveness HTTP port (unset → probe disabled) | unset | | `HEALTH_STALE_SECONDS` | A poll loop idle beyond this is reported unhealthy | `60` | diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index c87c1b06e2..5d0bab360f 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -201,16 +201,29 @@ def __init__(self, conn: PgConnection | None = None) -> None: def conn(self) -> PgConnection: if self._conn is None: self._conn = create_pg_connection() + # Every operation here is ONE statement (the claim CTE, set_vt, the + # ack DELETE, the enqueue INSERT) and already atomic on its own. In + # psycopg2's default mode each is wrapped in BEGIN … COMMIT: three + # round trips per operation, empty polls included — the largest + # share of statements on the DB under load (UN-4254). Autocommit + # sends the bare statement. Owned connections only: an injected one + # keeps the caller's mode, and create_pg_connection() itself stays + # non-autocommit because pg_barrier's retry safety depends on it. + self._conn.autocommit = True return self._conn @contextlib.contextmanager def _cursor(self) -> Iterator[Any]: """Yield a cursor; commit on success, roll back + recover on error. - Keeps the cached connection usable: a failed statement leaves the - connection in an aborted transaction, so we always roll back; a - dead connection can't be reused, so (when we own it) we drop the - cached handle and the next call reconnects. + Keeps the cached connection usable: a failed statement on a + non-autocommit (injected) connection leaves it in an aborted + transaction, so we always roll back; a dead connection can't be + reused, so (when we own it) we drop the cached handle and the next + call reconnects. On the owned autocommit connection ``commit()`` / + ``rollback()`` send nothing to the server (no transaction is open), + but ``rollback()`` still raises on a closed connection, which is + what flags it as dead below. """ conn = self.conn try: @@ -278,7 +291,10 @@ def send( committed the row but BEFORE psycopg2 read back ``RETURNING msg_id`` (the commit-loss / PgBouncer server-recycle case this very feature targets) is indistinguishable here from an idle reap, so the retry would - re-insert an already-committed row. The enqueue is therefore + re-insert an already-committed row. On the owned autocommit connection + there is no separate COMMIT: the row is durable as soon as the INSERT + executes, so this window covers the statement's own execution rather + than only a trailing COMMIT. The enqueue is therefore **at-least-once**: the ``reused`` gate removes the *common* duplicate (idle reap) but cannot remove the *rare* one (post-commit death). That residual duplicate leans on the module's at-least-once / idempotent- diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index c31b427957..e84378eef1 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -73,8 +73,16 @@ # Bounded join so a wedged renewal thread can't block the ack — it's a daemon and # dies with the process regardless. _LEASE_JOIN_TIMEOUT_SECONDS = 10.0 -_DEFAULT_POLL_INTERVAL = 0.1 -_DEFAULT_BACKOFF_MAX = 2.0 +# Empty-queue poll backoff: sleep POLL_INTERVAL after an empty claim, doubling to +# BACKOFF_MAX (a claim resets it). Every slot polls every queue it serves, so idle +# DB load is slots x queues / BACKOFF_MAX and tracks worker count, not work. At +# 0.1/2.0 that was the second-largest statement source on the DB and kept a shared +# instance busy with nothing queued (UN-4254). The cost of a higher cap is pickup +# latency on an idle queue: up to BACKOFF_MAX with one poller, much less with +# many (their sleeps are out of phase). Keep BACKOFF_MAX well under +# HEALTH_STALE_SECONDS, which the poll loop must beat between heartbeats. +_DEFAULT_POLL_INTERVAL = 0.5 +_DEFAULT_BACKOFF_MAX = 5.0 # A task claimed more than this many times keeps failing — drop it (poison) # rather than redeliver forever. _DEFAULT_MAX_ATTEMPTS = 5 diff --git a/workers/queue_backend/pg_queue/result_backend.py b/workers/queue_backend/pg_queue/result_backend.py index 16f8eefd8c..bbe78febb4 100644 --- a/workers/queue_backend/pg_queue/result_backend.py +++ b/workers/queue_backend/pg_queue/result_backend.py @@ -305,11 +305,20 @@ def __init__(self, conn: PgConnection | None = None) -> None: def conn(self) -> PgConnection: if self._conn is None: self._conn = create_pg_connection() + # Every write and read here is ONE idempotent statement, so run in + # autocommit and skip psycopg2's BEGIN … COMMIT round trips (UN-4254; + # same reasoning as PgQueueClient.conn). Owned connections only. + self._conn.autocommit = True return self._conn @contextlib.contextmanager def _cursor(self) -> Iterator[Any]: - """Yield a cursor; commit on success, roll back + recover on error.""" + """Yield a cursor; commit on success, roll back + recover on error. + + On the owned autocommit connection ``commit()`` / ``rollback()`` send + nothing to the server; ``rollback()`` still raises on a closed + connection, which is what flags it as dead below. + """ conn = self.conn try: with conn.cursor() as cur: diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index 6a0b52bb66..be8943df71 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -22,6 +22,7 @@ import psycopg2 import pytest +from psycopg2.extensions import TRANSACTION_STATUS_IDLE from queue_backend.fairness import DEFAULT_PRIORITY, MAX_PRIORITY, MIN_PRIORITY from queue_backend.pg_queue import PgQueueClient, QueueMessage from queue_backend.pg_queue.client import _SEND_RETRY_BACKOFF_SECONDS @@ -238,6 +239,22 @@ def test_close_closes_owned_not_injected(self, monkeypatch): PgQueueClient(conn=injected).close() injected.close.assert_not_called() + def test_owned_conn_runs_in_autocommit(self, monkeypatch): + # Single-statement operations must not pay psycopg2's BEGIN/COMMIT + # round trips (UN-4254) — including the reconnect after a dead conn. + client, conn, factory = self._owned_client(monkeypatch) + conn.autocommit = False + assert client.conn.autocommit is True + client._conn = None + conn.autocommit = False + assert client.conn.autocommit is True + assert factory.call_count == 2 + + def test_injected_conn_keeps_callers_mode(self): + injected = MagicMock() + injected.autocommit = False + assert PgQueueClient(conn=injected).conn.autocommit is False + class TestSendReconnectRetry: """``send()``'s one-shot reconnect-retry for a reused, idle-reaped connection. @@ -528,6 +545,35 @@ def queue_name(pg_conn): class TestPgQueueClientIntegration: + def test_owned_autocommit_conn_roundtrip(self, pg_conn, queue_name, monkeypatch): + # The production posture: an owned connection in autocommit. Each + # operation must persist on its own (visible to another session) and + # leave no transaction open between calls. + monkeypatch.setattr( + "queue_backend.pg_queue.client.create_pg_connection", + lambda *a, **k: create_pg_connection(env_prefix="TEST_DB_"), + ) + client = PgQueueClient() + try: + msg_id = client.send(queue_name, {"hello": "autocommit"}) + assert client.conn.autocommit is True + assert client.conn.get_transaction_status() == TRANSACTION_STATUS_IDLE + pg_conn.rollback() + with pg_conn.cursor() as cur: + cur.execute( + "SELECT state FROM pg_queue_message WHERE msg_id = %s", (msg_id,) + ) + assert cur.fetchone() == (QueueMessageState.READY.value,) + pg_conn.rollback() + msgs = client.read(queue_name, vt_seconds=30, qty=10) + assert [m.msg_id for m in msgs] == [msg_id] + assert client.conn.get_transaction_status() == TRANSACTION_STATUS_IDLE + assert client.set_vt(msg_id, 60) is True + assert client.delete(msg_id) is True + assert client.read(queue_name, vt_seconds=30, qty=10) == [] + finally: + client.close() + def test_send_read_delete_roundtrip(self, pg_conn, queue_name): client = PgQueueClient(conn=pg_conn) msg_id = client.send(queue_name, {"hello": "world"}) diff --git a/workers/tests/test_pg_result_backend.py b/workers/tests/test_pg_result_backend.py index 51352c6dd2..1759b4173d 100644 --- a/workers/tests/test_pg_result_backend.py +++ b/workers/tests/test_pg_result_backend.py @@ -584,6 +584,28 @@ def test_injected_connection_not_retried(self, monkeypatch): sleep.assert_not_called() conn.close.assert_not_called() # caller's connection untouched + def test_owned_conn_runs_in_autocommit(self, monkeypatch): + # Single-statement writes/reads skip psycopg2's BEGIN/COMMIT (UN-4254), + # including on the connection the retry reconnects to. + dead, _ = self._conn(execute_side_effect=psycopg2.OperationalError("reap")) + fresh, _ = self._conn() + fresh.autocommit = False + monkeypatch.setattr( + "queue_backend.pg_queue.result_backend.create_pg_connection", + MagicMock(return_value=fresh), + ) + self._no_sleep(monkeypatch) + rb = PgResultBackend() + rb._conn = dead + + rb.store_result("k", result={"ok": True}) + assert fresh.autocommit is True + + def test_injected_conn_keeps_callers_mode(self): + conn, _ = self._conn() + conn.autocommit = False + assert PgResultBackend(conn=conn).conn.autocommit is False + @pytest.mark.parametrize( "exc_type", [psycopg2.OperationalError, psycopg2.InterfaceError] ) From 8df3d4380cc3c291273315d2962e0a58f468db4e Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Thu, 22 Oct 2026 16:15:48 +0530 Subject: [PATCH 2/4] UN-4254 Restore dead-connection detection on the autocommit result backend Remediation of review findings F1, F2, F4, F5, F6, F7, F8: - F1: PgResultBackend._cursor treats a bare DatabaseError as a dead connection, as PgQueueClient._cursor does, since rollback() sends nothing under autocommit; docstring corrected. - F2: pin server-side backend termination on the result backend (live DB). - F4, F5, F8: backoff comment and README rows no longer claim extra pollers shorten pickup, and say a handled message (not a claim) resets the backoff. - F6: drop the PgQueueClient autocommit unit test that duplicated main's. - F7: pin the default backoff cap below the default health-stale threshold. Unstract-Remediation: remediation-deep iteration 1 --- workers/queue_backend/pg_queue/README.md | 4 +- workers/queue_backend/pg_queue/consumer.py | 15 ++--- .../queue_backend/pg_queue/result_backend.py | 11 +++- workers/tests/test_pg_queue_client.py | 11 ---- workers/tests/test_pg_queue_consumer.py | 12 ++++ workers/tests/test_pg_result_backend.py | 66 +++++++++++++++++++ 6 files changed, 95 insertions(+), 24 deletions(-) diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md index 6b6a8b252d..365230e470 100644 --- a/workers/queue_backend/pg_queue/README.md +++ b/workers/queue_backend/pg_queue/README.md @@ -31,8 +31,8 @@ integers. | `SHUTDOWN_GRACE_SECONDS` | Graceful-drain budget (shared across all children) on SIGTERM before SIGKILL | `= VT` (floored at `30`) | | `QUEUE` | Queue name(s) this consumer polls (comma-separated) | `default` | | `BATCH` | Messages claimed per poll | `1` | -| `POLL_INTERVAL` | First sleep after an empty poll (doubles up to `BACKOFF_MAX`; a claim resets it) | `0.5` | -| `BACKOFF_MAX` | Max empty-queue poll backoff — also the worst-case pickup latency on an idle queue with a single poller. Keep well under `HEALTH_STALE_SECONDS` | `5.0` | +| `POLL_INTERVAL` | First sleep after a poll cycle that handles nothing (doubles up to `BACKOFF_MAX`; a cycle that handles a message resets it) | `0.5` | +| `BACKOFF_MAX` | Max empty-queue poll backoff — also the worst-case pickup latency on an idle queue. Keep well under `HEALTH_STALE_SECONDS` | `5.0` | | `MAX_ATTEMPTS` | Max deliveries before a message is dropped as **poison** | `5` | | `HEALTH_PORT` | Liveness HTTP port (unset → probe disabled) | unset | | `HEALTH_STALE_SECONDS` | A poll loop idle beyond this is reported unhealthy | `60` | diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 2fca208ba3..fcf16037cc 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -73,14 +73,13 @@ # Bounded join so a wedged renewal thread can't block the ack — it's a daemon and # dies with the process regardless. _LEASE_JOIN_TIMEOUT_SECONDS = 10.0 -# Empty-queue poll backoff: sleep POLL_INTERVAL after an empty claim, doubling to -# BACKOFF_MAX (a claim resets it). Each slot makes one claim per cycle however many -# queues it serves (read_any), so idle DB load is slots / BACKOFF_MAX and tracks -# worker count, not work. At -# 0.1/2.0 that was the second-largest statement source on the DB and kept a shared -# instance busy with nothing queued (UN-4254). The cost of a higher cap is pickup -# latency on an idle queue: up to BACKOFF_MAX with one poller, much less with -# many (their sleeps are out of phase). Keep BACKOFF_MAX well under +# Empty-queue poll backoff: after a cycle that handles nothing, sleep POLL_INTERVAL, +# doubling up to BACKOFF_MAX; a cycle that handles a message resets it. Each slot +# makes one claim per cycle however many queues it serves (read_any), so idle DB +# load is slots / BACKOFF_MAX and tracks worker count, not work. At 0.1/2.0 that +# was the second-largest statement source on the DB and kept a shared instance +# busy with nothing queued (UN-4254). The cost of a higher cap is pickup latency +# on an idle queue: up to BACKOFF_MAX. Keep BACKOFF_MAX well under # HEALTH_STALE_SECONDS, which the poll loop must beat between heartbeats. _DEFAULT_POLL_INTERVAL = 0.5 _DEFAULT_BACKOFF_MAX = 5.0 diff --git a/workers/queue_backend/pg_queue/result_backend.py b/workers/queue_backend/pg_queue/result_backend.py index bbe78febb4..1e2a1db635 100644 --- a/workers/queue_backend/pg_queue/result_backend.py +++ b/workers/queue_backend/pg_queue/result_backend.py @@ -72,6 +72,7 @@ from collections.abc import Callable, Iterator from typing import TYPE_CHECKING, Any, Final, Self +import psycopg2 import redis from unstract.core.cache.redis_client import create_redis_client @@ -316,8 +317,8 @@ def _cursor(self) -> Iterator[Any]: """Yield a cursor; commit on success, roll back + recover on error. On the owned autocommit connection ``commit()`` / ``rollback()`` send - nothing to the server; ``rollback()`` still raises on a closed - connection, which is what flags it as dead below. + nothing to the server, so ``rollback()`` cannot detect a dead server + connection there; see the heuristic below. """ conn = self.conn try: @@ -325,7 +326,11 @@ def _cursor(self) -> Iterator[Any]: yield cur conn.commit() except Exception as exc: - conn_dead = is_connection_dead(exc) + # As in PgQueueClient._cursor: a server-side termination can surface + # as a bare DatabaseError, which the autocommit rollback() no longer + # catches, so treat it as dead directly. Statement errors map to a + # SQLSTATE subclass; a rare unmapped one costs only a reconnect. + conn_dead = is_connection_dead(exc) or type(exc) is psycopg2.DatabaseError try: conn.rollback() except Exception: diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index bef9ab0ac9..e9ced9e4d4 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -289,17 +289,6 @@ def test_close_closes_owned_not_injected(self, monkeypatch): PgQueueClient(conn=injected).close() injected.close.assert_not_called() - def test_owned_conn_runs_in_autocommit(self, monkeypatch): - # Single-statement operations must not pay psycopg2's BEGIN/COMMIT - # round trips (UN-4254) — including the reconnect after a dead conn. - client, conn, factory = self._owned_client(monkeypatch) - conn.autocommit = False - assert client.conn.autocommit is True - client._conn = None - conn.autocommit = False - assert client.conn.autocommit is True - assert factory.call_count == 2 - class TestSendReconnectRetry: """``send()``'s one-shot reconnect-retry for a reused, idle-reaped connection. diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 93f22892d2..56fab28be8 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -534,6 +534,18 @@ def test_rejects_backoff_max_below_poll_interval(self): with pytest.raises(ValueError, match="backoff_max"): PgQueueConsumer(["q"], client=client, poll_interval=0.5, backoff_max=0.1) + def test_default_backoff_max_stays_under_default_health_stale(self): + # An idle loop sleeps up to backoff_max between heartbeats, so the + # default cap must stay below the default staleness threshold or idle + # consumers fail their liveness probe. + from queue_backend.pg_queue import consumer as consumer_mod + + assert ( + consumer_mod._DEFAULT_POLL_INTERVAL + <= consumer_mod._DEFAULT_BACKOFF_MAX + < consumer_mod._DEFAULT_HEALTH_STALE_SECONDS + ) + class TestRunLoop: def test_run_stops_gracefully(self, monkeypatch): diff --git a/workers/tests/test_pg_result_backend.py b/workers/tests/test_pg_result_backend.py index 1759b4173d..417ec7502a 100644 --- a/workers/tests/test_pg_result_backend.py +++ b/workers/tests/test_pg_result_backend.py @@ -606,6 +606,27 @@ def test_injected_conn_keeps_callers_mode(self): conn.autocommit = False assert PgResultBackend(conn=conn).conn.autocommit is False + def test_bare_database_error_drops_owned_conn(self): + # Under autocommit rollback() sends nothing, so it can't flag a dead + # server connection; a bare DatabaseError (how a server-side termination + # can surface) must drop the owned handle itself, as PgQueueClient does. + dead, _ = self._conn(execute_side_effect=psycopg2.DatabaseError("gone")) + rb = PgResultBackend() + rb._conn = dead + with pytest.raises(psycopg2.DatabaseError): + rb.get_result("k") + assert rb._conn is None + dead.close.assert_called_once() + + def test_statement_error_subclass_keeps_owned_conn(self): + # A SQLSTATE-mapped statement error is not a dead connection. + conn, _ = self._conn(execute_side_effect=psycopg2.DataError("bad input")) + rb = PgResultBackend() + rb._conn = conn + with pytest.raises(psycopg2.DataError): + rb.get_result("k") + assert rb._conn is conn + @pytest.mark.parametrize( "exc_type", [psycopg2.OperationalError, psycopg2.InterfaceError] ) @@ -709,6 +730,51 @@ def test_real_reconnect_heals_closed_connection(self, pg_conn, monkeypatch): pg_conn.commit() assert row is not None and row[0] == "true" + def test_real_reconnect_after_backend_termination(self, pg_conn, monkeypatch): + # The owned connection runs in autocommit, where rollback() sends nothing + # and so can't prove a server-terminated connection dead. Pin that a + # terminated backend still heals: a read drops the handle, and a write + # self-heals through _store_with_reconnect. + os.environ.setdefault("TEST_DB_HOST", "127.0.0.1") + monkeypatch.setattr( + "queue_backend.pg_queue.result_backend.create_pg_connection", + lambda *a, **k: create_pg_connection(env_prefix="TEST_DB_"), + ) + + def terminate(pid): + with pg_conn.cursor() as cur: + # Wait (up to 5 s) for the backend to exit, so the next call + # can't race ahead of the termination (PG14+). + cur.execute("SELECT pg_terminate_backend(%s, 5000)", (pid,)) + pg_conn.commit() + + key = _key() + rb = PgResultBackend() + try: + assert rb.conn.autocommit is True + old_pid = rb.conn.get_backend_pid() + terminate(old_pid) + with pytest.raises(psycopg2.DatabaseError): + rb.get_result(key) + assert rb._conn is None # dropped, so the next call reconnects + + terminate(rb.conn.get_backend_pid()) + rb.store_result(key, result={"healed": True}) # must self-heal + assert rb.conn.get_backend_pid() != old_pid + finally: + rb.close() + + pg_conn.rollback() + with pg_conn.cursor() as cur: + cur.execute( + "SELECT result->>'healed' FROM pg_task_result WHERE task_id = %s", + (key,), + ) + row = cur.fetchone() + cur.execute("DELETE FROM pg_task_result WHERE task_id = %s", (key,)) + pg_conn.commit() + assert row is not None and row[0] == "true" + class TestRedisDownFallbackBudget: """The early Redis-unavailable fallback in ``_wait_via_redis`` must degrade to From f937a3fe6a90be732fab35e49ff6b6631822e382 Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Thu, 22 Oct 2026 16:15:48 +0530 Subject: [PATCH 3/4] UN-4254 Share the autocommit dead-connection check between both PG clients Behaviour-preserving cleanup (simplify pass): - connection.is_dead_after_autocommit_error holds the bare-DatabaseError rule and its rationale once; PgQueueClient._cursor and PgResultBackend._cursor both call it. is_connection_dead is unchanged. - The live-DB reconnect tests share an owned-backend fixture and a read-back helper; the redundant TEST_DB_HOST setdefault is dropped (pg_conn already sets it). Unstract-Remediation: remediation-deep iteration 2 --- workers/queue_backend/pg_queue/client.py | 19 +++--- workers/queue_backend/pg_queue/connection.py | 17 +++++ .../queue_backend/pg_queue/result_backend.py | 20 ++---- workers/tests/test_pg_result_backend.py | 65 +++++++------------ 4 files changed, 56 insertions(+), 65 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index cf799ecfce..307e984223 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -39,14 +39,16 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Self -import psycopg2 - from unstract.core.data_models import QueueMessageState from unstract.core.jsonb import dumps_for_jsonb from ..fairness import DEFAULT_PRIORITY, MAX_PRIORITY, MIN_PRIORITY from .connection import CONN_DEAD_ERRORS as _CONN_DEAD_ERRORS -from .connection import create_pg_connection, is_connection_dead +from .connection import ( + create_pg_connection, + is_connection_dead, + is_dead_after_autocommit_error, +) from .schema import qualified if TYPE_CHECKING: @@ -278,19 +280,14 @@ def _cursor(self) -> Iterator[Any]: conn.commit() except Exception as exc: # A failed rollback proves the connection is unusable, regardless - # of which psycopg2 error subclass was raised (a server-side - # termination mid-statement can surface as a bare DatabaseError). - # On an owned autocommit connection rollback() sends nothing, so it - # can't detect a dead server connection there. As a heuristic, a - # bare DatabaseError is treated as dead directly instead: statement - # errors normally map to a SQLSTATE subclass, and a rare unmapped - # one costs only a needless reconnect. + # of which psycopg2 error subclass was raised; on the owned + # autocommit connection it proves nothing, hence the helper. # NOTE: ``send()``'s retry deliberately catches only # ``_CONN_DEAD_ERRORS``; a server death that surfaces as a *bare* # ``psycopg2.DatabaseError`` is still dropped here when we own the # connection (an injected one is never dropped) but is intentionally # NOT retried by ``send()`` — it's left to the next call's reconnect. - conn_dead = is_connection_dead(exc) or type(exc) is psycopg2.DatabaseError + conn_dead = is_dead_after_autocommit_error(exc) try: conn.rollback() except Exception: diff --git a/workers/queue_backend/pg_queue/connection.py b/workers/queue_backend/pg_queue/connection.py index fa1dc20564..a5c9979de2 100644 --- a/workers/queue_backend/pg_queue/connection.py +++ b/workers/queue_backend/pg_queue/connection.py @@ -82,6 +82,23 @@ def is_connection_dead(exc: BaseException) -> bool: return isinstance(exc, CONN_DEAD_ERRORS) +def is_dead_after_autocommit_error(exc: BaseException) -> bool: + """Whether a statement error on an owned AUTOCOMMIT connection means it is dead. + + On a transactional connection a failed ``rollback()`` proves a dead server. + Under autocommit ``rollback()`` sends nothing, so it proves nothing; and a + server-side termination mid-statement can surface as a *bare* + ``DatabaseError``. Treat that as dead too: statement errors normally map to + a SQLSTATE subclass, and a rare unmapped one costs only a needless reconnect. + + Deliberately separate from :func:`is_connection_dead`: transactional + connections (``pg_barrier``) get this detection from a real ``rollback()``, + and the retry sites must not start treating a bare ``DatabaseError`` as + retryable. + """ + return is_connection_dead(exc) or type(exc) is psycopg2.DatabaseError + + # Bounded retry for *transient* connect failures (DB restart, PgBouncer pool # wait, brief network partition, "too many clients" spikes) — the common cloud # blips. Opening a connection is side-effect-free, so retrying it is safe and diff --git a/workers/queue_backend/pg_queue/result_backend.py b/workers/queue_backend/pg_queue/result_backend.py index 1e2a1db635..8af7464b5d 100644 --- a/workers/queue_backend/pg_queue/result_backend.py +++ b/workers/queue_backend/pg_queue/result_backend.py @@ -72,7 +72,6 @@ from collections.abc import Callable, Iterator from typing import TYPE_CHECKING, Any, Final, Self -import psycopg2 import redis from unstract.core.cache.redis_client import create_redis_client @@ -82,7 +81,11 @@ from .connection import CONN_DEAD_ERRORS as _CONN_DEAD_ERRORS from .connection import PAYLOAD_REJECTED_ERRORS as _PAYLOAD_REJECTED_ERRORS -from .connection import create_pg_connection, is_connection_dead +from .connection import ( + create_pg_connection, + is_connection_dead, + is_dead_after_autocommit_error, +) from .schema import qualified if TYPE_CHECKING: @@ -314,23 +317,14 @@ def conn(self) -> PgConnection: @contextlib.contextmanager def _cursor(self) -> Iterator[Any]: - """Yield a cursor; commit on success, roll back + recover on error. - - On the owned autocommit connection ``commit()`` / ``rollback()`` send - nothing to the server, so ``rollback()`` cannot detect a dead server - connection there; see the heuristic below. - """ + """Yield a cursor; commit on success, roll back + recover on error.""" conn = self.conn try: with conn.cursor() as cur: yield cur conn.commit() except Exception as exc: - # As in PgQueueClient._cursor: a server-side termination can surface - # as a bare DatabaseError, which the autocommit rollback() no longer - # catches, so treat it as dead directly. Statement errors map to a - # SQLSTATE subclass; a rare unmapped one costs only a reconnect. - conn_dead = is_connection_dead(exc) or type(exc) is psycopg2.DatabaseError + conn_dead = is_dead_after_autocommit_error(exc) try: conn.rollback() except Exception: diff --git a/workers/tests/test_pg_result_backend.py b/workers/tests/test_pg_result_backend.py index 417ec7502a..9244a1a50f 100644 --- a/workers/tests/test_pg_result_backend.py +++ b/workers/tests/test_pg_result_backend.py @@ -699,25 +699,21 @@ class TestStoreResultRealReconnect: were unusable — the mock tests can't see that. """ - def test_real_reconnect_heals_closed_connection(self, pg_conn, monkeypatch): + @pytest.fixture + def owned_rb(self, pg_conn, monkeypatch): # The owned backend reconnects via result_backend.create_pg_connection; # point that at the integration DB (TEST_DB_*), the same one pg_conn uses # (the bare DB_* env is the suite's unit-isolation placeholder). - os.environ.setdefault("TEST_DB_HOST", "127.0.0.1") monkeypatch.setattr( "queue_backend.pg_queue.result_backend.create_pg_connection", lambda *a, **k: create_pg_connection(env_prefix="TEST_DB_"), ) - - key = _key() rb = PgResultBackend() # owned, real connection to the test DB - try: - _ = rb.conn # materialise the connection - rb._conn.close() # client-side close -> InterfaceError on next use - rb.store_result(key, result={"healed": True}) # must self-heal - finally: - rb.close() + yield rb + rb.close() + @staticmethod + def _assert_healed_and_delete(pg_conn, key): # The row really landed — read it back on the fixture's own connection. pg_conn.rollback() # clear any aborted txn before reading with pg_conn.cursor() as cur: @@ -730,17 +726,18 @@ def test_real_reconnect_heals_closed_connection(self, pg_conn, monkeypatch): pg_conn.commit() assert row is not None and row[0] == "true" - def test_real_reconnect_after_backend_termination(self, pg_conn, monkeypatch): + def test_real_reconnect_heals_closed_connection(self, pg_conn, owned_rb): + key = _key() + _ = owned_rb.conn # materialise the connection + owned_rb._conn.close() # client-side close -> InterfaceError on next use + owned_rb.store_result(key, result={"healed": True}) # must self-heal + self._assert_healed_and_delete(pg_conn, key) + + def test_real_reconnect_after_backend_termination(self, pg_conn, owned_rb): # The owned connection runs in autocommit, where rollback() sends nothing # and so can't prove a server-terminated connection dead. Pin that a # terminated backend still heals: a read drops the handle, and a write # self-heals through _store_with_reconnect. - os.environ.setdefault("TEST_DB_HOST", "127.0.0.1") - monkeypatch.setattr( - "queue_backend.pg_queue.result_backend.create_pg_connection", - lambda *a, **k: create_pg_connection(env_prefix="TEST_DB_"), - ) - def terminate(pid): with pg_conn.cursor() as cur: # Wait (up to 5 s) for the backend to exit, so the next call @@ -749,31 +746,17 @@ def terminate(pid): pg_conn.commit() key = _key() - rb = PgResultBackend() - try: - assert rb.conn.autocommit is True - old_pid = rb.conn.get_backend_pid() - terminate(old_pid) - with pytest.raises(psycopg2.DatabaseError): - rb.get_result(key) - assert rb._conn is None # dropped, so the next call reconnects - - terminate(rb.conn.get_backend_pid()) - rb.store_result(key, result={"healed": True}) # must self-heal - assert rb.conn.get_backend_pid() != old_pid - finally: - rb.close() + assert owned_rb.conn.autocommit is True + old_pid = owned_rb.conn.get_backend_pid() + terminate(old_pid) + with pytest.raises(psycopg2.DatabaseError): + owned_rb.get_result(key) + assert owned_rb._conn is None # dropped, so the next call reconnects - pg_conn.rollback() - with pg_conn.cursor() as cur: - cur.execute( - "SELECT result->>'healed' FROM pg_task_result WHERE task_id = %s", - (key,), - ) - row = cur.fetchone() - cur.execute("DELETE FROM pg_task_result WHERE task_id = %s", (key,)) - pg_conn.commit() - assert row is not None and row[0] == "true" + terminate(owned_rb.conn.get_backend_pid()) + owned_rb.store_result(key, result={"healed": True}) # must self-heal + assert owned_rb.conn.get_backend_pid() != old_pid + self._assert_healed_and_delete(pg_conn, key) class TestRedisDownFallbackBudget: From bea677d1d23d78c113eb86b2bcac5ba9a852de02 Mon Sep 17 00:00:00 2001 From: Rahul Johny Date: Thu, 22 Oct 2026 16:15:48 +0530 Subject: [PATCH 4/4] UN-4254 Narrow the autocommit dead-connection comments to what the code does The confirming review showed a server-terminated connection raises OperationalError and is marked closed, so it was never only the new helper that caught it. Reword is_dead_after_autocommit_error's docstring, the PgQueueClient._cursor comment and both result-backend test comments to the narrow case the helper actually adds (a bare DatabaseError on a connection psycopg2 still reports open), and have the termination test assert the OperationalError it really sees. Unstract-Remediation: remediation-deep iteration 2 --- workers/queue_backend/pg_queue/client.py | 13 +++++++------ workers/queue_backend/pg_queue/connection.py | 12 +++++++----- workers/tests/test_pg_result_backend.py | 17 +++++++++-------- 3 files changed, 23 insertions(+), 19 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index 307e984223..278c1a75c9 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -280,13 +280,14 @@ def _cursor(self) -> Iterator[Any]: conn.commit() except Exception as exc: # A failed rollback proves the connection is unusable, regardless - # of which psycopg2 error subclass was raised; on the owned - # autocommit connection it proves nothing, hence the helper. + # of which psycopg2 error subclass was raised. On the owned + # autocommit connection a rollback that succeeds sends nothing, so + # the helper also treats a bare DatabaseError as dead. # NOTE: ``send()``'s retry deliberately catches only - # ``_CONN_DEAD_ERRORS``; a server death that surfaces as a *bare* - # ``psycopg2.DatabaseError`` is still dropped here when we own the - # connection (an injected one is never dropped) but is intentionally - # NOT retried by ``send()`` — it's left to the next call's reconnect. + # ``_CONN_DEAD_ERRORS``; a *bare* ``psycopg2.DatabaseError`` is still + # dropped here when we own the connection (an injected one is never + # dropped) but is intentionally NOT retried by ``send()`` — it's + # left to the next call's reconnect. conn_dead = is_dead_after_autocommit_error(exc) try: conn.rollback() diff --git a/workers/queue_backend/pg_queue/connection.py b/workers/queue_backend/pg_queue/connection.py index a5c9979de2..3c15e49b43 100644 --- a/workers/queue_backend/pg_queue/connection.py +++ b/workers/queue_backend/pg_queue/connection.py @@ -85,11 +85,13 @@ def is_connection_dead(exc: BaseException) -> bool: def is_dead_after_autocommit_error(exc: BaseException) -> bool: """Whether a statement error on an owned AUTOCOMMIT connection means it is dead. - On a transactional connection a failed ``rollback()`` proves a dead server. - Under autocommit ``rollback()`` sends nothing, so it proves nothing; and a - server-side termination mid-statement can surface as a *bare* - ``DatabaseError``. Treat that as dead too: statement errors normally map to - a SQLSTATE subclass, and a rare unmapped one costs only a needless reconnect. + A connection the server terminated already reads as dead without this: the + error is an ``OperationalError`` and psycopg2 marks the connection closed. + What autocommit loses is narrower: ``rollback()`` sends nothing, so it no + longer exposes a dead connection that psycopg2 still reports as open. So a + bare ``DatabaseError`` (one psycopg2 mapped to no SQLSTATE subclass) is + treated as dead too; if it was only an unmapped statement error, the cost + is one needless reconnect. Deliberately separate from :func:`is_connection_dead`: transactional connections (``pg_barrier``) get this detection from a real ``rollback()``, diff --git a/workers/tests/test_pg_result_backend.py b/workers/tests/test_pg_result_backend.py index 9244a1a50f..773506dd01 100644 --- a/workers/tests/test_pg_result_backend.py +++ b/workers/tests/test_pg_result_backend.py @@ -607,9 +607,9 @@ def test_injected_conn_keeps_callers_mode(self): assert PgResultBackend(conn=conn).conn.autocommit is False def test_bare_database_error_drops_owned_conn(self): - # Under autocommit rollback() sends nothing, so it can't flag a dead - # server connection; a bare DatabaseError (how a server-side termination - # can surface) must drop the owned handle itself, as PgQueueClient does. + # Under autocommit rollback() sends nothing, so a bare DatabaseError on + # a connection psycopg2 still reports open must drop the owned handle + # through is_dead_after_autocommit_error, as PgQueueClient does. dead, _ = self._conn(execute_side_effect=psycopg2.DatabaseError("gone")) rb = PgResultBackend() rb._conn = dead @@ -734,10 +734,11 @@ def test_real_reconnect_heals_closed_connection(self, pg_conn, owned_rb): self._assert_healed_and_delete(pg_conn, key) def test_real_reconnect_after_backend_termination(self, pg_conn, owned_rb): - # The owned connection runs in autocommit, where rollback() sends nothing - # and so can't prove a server-terminated connection dead. Pin that a - # terminated backend still heals: a read drops the handle, and a write - # self-heals through _store_with_reconnect. + # Pin that a terminated backend heals on the owned autocommit + # connection: a read drops the handle, and a write self-heals through + # _store_with_reconnect. Termination surfaces as OperationalError with + # the connection marked closed, so this path rests on is_connection_dead + # and conn.closed; the bare-DatabaseError branch is the unit test's. def terminate(pid): with pg_conn.cursor() as cur: # Wait (up to 5 s) for the backend to exit, so the next call @@ -749,7 +750,7 @@ def terminate(pid): assert owned_rb.conn.autocommit is True old_pid = owned_rb.conn.get_backend_pid() terminate(old_pid) - with pytest.raises(psycopg2.DatabaseError): + with pytest.raises(psycopg2.OperationalError): owned_rb.get_result(key) assert owned_rb._conn is None # dropped, so the next call reconnects