UN-4254 [PERF] Run PG-queue ops in autocommit and slow idle polling

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 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` |

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,15 @@ 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 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 = is_connection_dead(exc) or type(exc) is psycopg2.DatabaseError
+            # ``_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()
             except Exception:

workers/queue_backend/pg_queue/connection.py

@@ -82,6 +82,25 @@ 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.
+
+    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()``,
+    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

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: 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
 # A task claimed more than this many times keeps failing — drop it (poison)
 # rather than redeliver forever.
 _DEFAULT_MAX_ATTEMPTS = 5

workers/queue_backend/pg_queue/result_backend.py

@@ -81,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:
@@ -305,6 +309,10 @@ 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
@@ -316,7 +324,7 @@ def _cursor(self) -> Iterator[Any]:
                 yield cur
             conn.commit()
         except Exception as exc:
-            conn_dead = is_connection_dead(exc)
+            conn_dead = is_dead_after_autocommit_error(exc)
             try:
                 conn.rollback()
             except Exception:

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, _dequeue_any_sql
@@ -578,6 +579,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"})

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):

workers/tests/test_pg_result_backend.py

@@ -584,6 +584,49 @@ 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
+
+    def test_bare_database_error_drops_owned_conn(self):
+        # 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
+        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]
     )
@@ -656,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:
@@ -687,6 +726,39 @@ 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_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):
+        # 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
+                # can't race ahead of the termination (PG14+).
+                cur.execute("SELECT pg_terminate_backend(%s, 5000)", (pid,))
+            pg_conn.commit()
+
+        key = _key()
+        assert owned_rb.conn.autocommit is True
+        old_pid = owned_rb.conn.get_backend_pid()
+        terminate(old_pid)
+        with pytest.raises(psycopg2.OperationalError):
+            owned_rb.get_result(key)
+        assert owned_rb._conn is None  # dropped, so the next call reconnects
+
+        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:
     """The early Redis-unavailable fallback in ``_wait_via_redis`` must degrade to