diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md index d2e755251f..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` | 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` | diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index cf799ecfce..278c1a75c9 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,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: diff --git a/workers/queue_backend/pg_queue/connection.py b/workers/queue_backend/pg_queue/connection.py index fa1dc20564..3c15e49b43 100644 --- a/workers/queue_backend/pg_queue/connection.py +++ b/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 diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 50d915979b..fcf16037cc 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: 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 diff --git a/workers/queue_backend/pg_queue/result_backend.py b/workers/queue_backend/pg_queue/result_backend.py index 16f8eefd8c..8af7464b5d 100644 --- a/workers/queue_backend/pg_queue/result_backend.py +++ b/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: diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index fa5f8da603..e9ced9e4d4 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, _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"}) 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 51352c6dd2..773506dd01 100644 --- a/workers/tests/test_pg_result_backend.py +++ b/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