diff --git a/workers/pg_queue_consumer/supervisor.py b/workers/pg_queue_consumer/supervisor.py index b852a3369d..e88014791b 100644 --- a/workers/pg_queue_consumer/supervisor.py +++ b/workers/pg_queue_consumer/supervisor.py @@ -370,7 +370,7 @@ def _run_child(slot: int, heartbeats, loaded) -> None: # noqa: ANN001 (ctypes a def _publish_heartbeat() -> None: # last-poll wall-time = now − (seconds since last poll). Frozen while a - # task runs (the consumer stamps its heartbeat before each queue read), + # task runs (the consumer stamps its heartbeat before each claim), # so a child stuck on a too-long task goes stale exactly as the single # consumer does. Guarded so a transient error (e.g. teardown during # shutdown) logs loudly and the loop continues instead of dying silently diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index c87c1b06e2..cf799ecfce 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -22,10 +22,12 @@ place; the schema (``pg_queue_message`` table + partial claim index) is a plain Django migration with no DB-side function. -The cached connection is kept usable across calls: every operation rolls -back on error, and a connection that goes bad (dropped socket / PgBouncer -recycle) is discarded so the next call reconnects — so one blip can't -permanently wedge the consumer. +The cached connection is kept usable across calls: a connection that goes bad +(dropped socket / PgBouncer recycle) is discarded so the next call reconnects — +so one blip can't permanently wedge the consumer. A connection the client opens +itself runs in autocommit, so every operation must stay ONE statement (each +commits on its own); an injected connection keeps the caller's transaction +semantics and is committed or rolled back per operation. """ from __future__ import annotations @@ -37,6 +39,8 @@ 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 @@ -136,6 +140,46 @@ def _dequeue_sql() -> str: """ +def _dequeue_any_sql() -> str: + """Multi-queue variant of :func:`_dequeue_sql` (see :meth:`PgQueueClient.read_any`). + + ``WITH ORDINALITY`` numbers the queues 1..n; the outer ORDER BY rotates + that numbering by the caller's ``start`` so the queue served first changes + each cycle (every queue is still probed on every call). Within a queue the + order is the same ``priority DESC, msg_id`` top-N the single-queue claim + uses, on the same partial index. + """ + msg = qualified("pg_queue_message") + return f""" +WITH locked AS ( + SELECT c.msg_id, q.ord, c.priority + FROM unnest(%s::text[]) WITH ORDINALITY AS q(name, ord) + CROSS JOIN LATERAL ( + SELECT m.msg_id, m.priority + FROM {msg} m + WHERE m.queue_name = q.name + AND m.state = '{_READY}' + ORDER BY m.priority DESC, m.msg_id + FOR UPDATE SKIP LOCKED + LIMIT %s + ) c + ORDER BY (q.ord - 1 - %s + %s) %% %s, c.priority DESC, c.msg_id + LIMIT %s +), claimed AS ( + UPDATE {msg} q + SET state = '{_CLAIMED}', + vt = now() + make_interval(secs => %s), + read_ct = read_ct + 1 + FROM locked + WHERE q.msg_id = locked.msg_id + RETURNING q.msg_id, q.message, q.read_ct, q.queue_name, locked.ord, q.priority +) +SELECT msg_id, message, read_ct, queue_name + FROM claimed + ORDER BY (ord - 1 - %s + %s) %% %s, priority DESC, msg_id +""" + + @dataclass(frozen=True, slots=True) class QueueMessage: """A claimed queue message. @@ -147,11 +191,14 @@ class QueueMessage: from the dequeue (``read_ct = read_ct + 1 RETURNING``), so consumers can cap redelivery of a poison message. No default: a ``0`` ("never claimed") can't come from the dequeue and would silently bypass the poison guard. + ``queue_name`` is the queue the message was claimed from (also no default); + the multi-queue consumer uses it to rotate past the queue it just served. """ msg_id: int message: dict[str, Any] read_ct: int + queue_name: str # The whole enqueue contract (columns + their defaults) in one place. `send()` @@ -200,17 +247,29 @@ def __init__(self, conn: PgConnection | None = None) -> None: @property def conn(self) -> PgConnection: if self._conn is None: - self._conn = create_pg_connection() + conn = create_pg_connection() + # Every operation here is ONE statement, so on a connection this + # client owns an explicit transaction only adds a BEGIN + COMMIT + # round trip to each (3 statements per claim instead of 1). + # Autocommit is set only on a connection this client creates: an + # injected one may carry the caller's open transaction (unstract- + # cloud's agentic fan-out claim relies on send() committing its + # claim INSERT together with the enqueue), and create_pg_connection + # itself stays non-autocommit for pg_barrier, which relies on that. + conn.autocommit = True + self._conn = conn 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. On an injected connection a failed + statement leaves an aborted transaction, so we always roll back. On an + owned connection (autocommit) each statement already committed or failed + on its own, so keep every operation to ONE statement; commit/rollback + send nothing there. A dead connection can't be reused, so (when we own it) we + drop the cached handle and the next call reconnects. """ conn = self.conn try: @@ -221,13 +280,17 @@ def _cursor(self) -> Iterator[Any]: # 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. # NOTE: ``send()``'s retry deliberately catches only # ``_CONN_DEAD_ERRORS``; a server death that surfaces as a *bare* - # ``psycopg2.DatabaseError`` is still treated as dead here (via the - # failed-rollback branch below, which drops the handle) but is - # intentionally NOT retried by ``send()`` — it's left to the next - # call's reconnect. - conn_dead = is_connection_dead(exc) + # ``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 try: conn.rollback() except Exception: @@ -280,7 +343,10 @@ def send( targets) is indistinguishable here from an idle reap, so the retry would re-insert an already-committed row. 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 + (idle reap) but cannot remove the *rare* one (post-commit death). On a + client-owned (autocommit) connection the INSERT commits at statement end, + so that ambiguous window is the INSERT's own reply rather than a separate + COMMIT's — the same class of case, still rare. That residual duplicate leans on the module's at-least-once / idempotent- consumer contract (see the module docstring header) as the GENERAL backstop — every consumer here must already tolerate redelivery. @@ -388,7 +454,78 @@ def read( cur.execute(_dequeue_sql(), (queue_name, qty, vt_seconds)) rows = cur.fetchall() return [ - QueueMessage(msg_id=int(r[0]), message=r[1], read_ct=int(r[2])) for r in rows + QueueMessage( + msg_id=int(r[0]), message=r[1], read_ct=int(r[2]), queue_name=queue_name + ) + for r in rows + ] + + def read_any( + self, + queue_names: list[str], + *, + vt_seconds: int = 30, + qty: int = 1, + start: int = 0, + ) -> list[QueueMessage]: + """Claim up to ``qty`` ready messages from several queues in ONE statement. + + Same claim contract as :meth:`read`, but one statement covers every + queue — on a client-owned (autocommit) connection an idle multi-queue + consumer costs one statement per cycle. Each queue is probed with its + own indexed top-N (a ``LATERAL`` per queue on the partial claim index), + so cost does not grow with the backlog. + + Queue order is round-robin from ``start`` (an index into + ``queue_names``), NOT priority: messages come from the first non-empty + queue at or after ``start``, then the next, and so on. The caller should + pass ``start`` = one past the queue it served last (each message's + ``queue_name``), so every busy queue gets an equal turn. ``priority`` + orders messages within a queue. + + Accepted trade-off: the probe locks up to ``qty`` ready rows in EVERY + queue (``FOR UPDATE SKIP LOCKED``), and only the chosen ones are claimed; + the rest stay ``ready`` and are unlocked when the claim commits (at + statement end on a client-owned autocommit connection). A + reader racing that statement — another multi-queue consumer, or a + single-queue consumer of one of these queues — skips the locked rows, so + if the only ready rows it wants are locked it gets ``[]`` and takes the + consumer's empty-poll back-off (starting at ``poll_interval``) instead of + a message. This only happens while consumers contend for the last rows; + idle polls lock nothing. + + Accepted trade-off (autocommit): on a client-owned connection the claim + commits at statement end. If the connection dies after that but before + the rows are read back, the claim stands with nobody running it — the + messages stay ``claimed`` until their lease lapses (``vt_seconds``) and + the reaper re-arms them, and that delivery counts toward ``read_ct``. The + caller sees the connection error (the consumer logs it as a failed + claim). + """ + if vt_seconds <= 0: + raise ValueError(f"vt_seconds must be positive, got {vt_seconds}") + if qty <= 0: + raise ValueError(f"qty must be positive, got {qty}") + if not queue_names: + raise ValueError("queue_names must be non-empty") + # A repeated name would be probed twice and its rows would take batch + # slots twice, crowding other queues out (fewer than qty returned). + queues = list(dict.fromkeys(queue_names)) + n = len(queues) + rotate = (start % n, n, n) + with self._cursor() as cur: + # Param order matches the %s positions in _dequeue_any_sql(): + # queues, per-queue LIMIT, rotation, LIMIT, vt, rotation again. + cur.execute( + _dequeue_any_sql(), + (queues, qty, *rotate, qty, vt_seconds, *rotate), + ) + rows = cur.fetchall() + return [ + QueueMessage( + msg_id=int(r[0]), message=r[1], read_ct=int(r[2]), queue_name=r[3] + ) + for r in rows ] def set_vt(self, msg_id: int, vt_seconds: int) -> bool: diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index c31b427957..50d915979b 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -96,9 +96,10 @@ # response, and the row is TTL'd. _TASK_STATUS_RETENTION_SECONDS = 86400 # Liveness: a poll loop that hasn't cycled in this many seconds is reported -# unhealthy. The heartbeat is stamped before each queue read and frozen during -# task execution, so this threshold doubles as an UPPER BOUND on a single task's -# wall-clock: a task running longer than it trips the probe → pod restart (or, +# unhealthy. The heartbeat is stamped before each claim and frozen while the +# claimed batch runs, so this threshold doubles as an UPPER BOUND on a batch's +# wall-clock (one task at the default batch_size=1): a batch running longer than +# it trips the probe → pod restart (or, # under the prefork supervisor with CHILD_WATCHDOG on, a SIGKILL of just that # child) → the in-flight task is killed and (at-least-once) redelivered. 60s suits the # current sub-second leaf (send_webhook_notification); for longer-running tasks, @@ -311,14 +312,15 @@ def __init__( if not queue_names: raise ValueError("queue_names must be a non-empty list") # One process can drain several queues (9f) — e.g. a file_processing - # consumer drains both file_processing and api_file_processing. Each is - # read once per cycle in list order (poll_once): this prevents starvation - # (every queue gets a read each cycle) but does NOT equalize throughput — - # an always-full queue_names[0] claims its full batch and runs it before - # later queues. Copy + de-dup (order-preserving): a duplicate would - # double-read a queue per cycle, and storing the caller's list by + # consumer drains both file_processing and api_file_processing. All are + # claimed from in one statement per cycle, in a rotation that starts past + # the queue served last (poll_once): busy queues take equal turns. Copy + + # de-dup (order-preserving): a duplicate would probe a queue twice per + # claim and give it two rotation slots, and storing the caller's list by # reference would let a later mutation bypass the non-empty validation. self.queue_names = list(dict.fromkeys(queue_names)) + # Round-robin start for the multi-queue claim (see poll_once). + self._next_queue = 0 self._client = client if client is not None else PgQueueClient() self._app = app if app is not None else current_app # Lazily built the first time a poison drop needs to mark an execution @@ -366,40 +368,74 @@ def __init__( self._result_backend: PgResultBackend | None = None # Heartbeat for the liveness probe: monotonic timestamp of the most # recent poll attempt. Seeded at construction so a just-started consumer - # reads healthy. Updated before EACH queue read in poll_once, so a loop + # reads healthy. Updated before each claim in poll_once, so a loop # wedged on a long-running task goes stale and is caught — # something pgrep-based --status and the launch-time check cannot see. self._last_poll_monotonic = time.monotonic() def poll_once(self) -> int: - """Claim + process one batch per queue (read once each, in list order); - returns the total number of messages claimed across all queues this cycle. - - Each queue is isolated: a read/handle failure on one queue is logged and - skipped so the others still get their turn, and the work already done this - cycle still counts (so run() doesn't take the empty-queue backoff path - after a partial failure). + """Claim + process one batch; returns the number of messages processed. + + A single-queue consumer reads its queue. A multi-queue consumer claims + across all its queues in ONE statement (:meth:`PgQueueClient.read_any`), + so an idle cycle costs one statement however many queues it drains. The + next claim starts one past the queue just served (or one further along + after an empty cycle), so busy queues take equal turns and none starves. + + A failed claim is logged and returns 0, so run() backs off. A failure + while handling stops the batch: it is logged with the msg_ids left + unstarted, and the messages already handled are still counted (so run() + does not back off after real work). Accepted trade-off, ``batch_size > 1`` only + (it is forced to 1 unless the lease equals the max runtime): the + unstarted messages stay claimed until their lease lapses and the reaper + re-arms them, costing each a ``read_ct`` attempt toward the poison cap. + Running them instead would, during a DB outage, run work whose ack + cannot land, so it would run twice. """ - total = 0 - for queue_name in self.queue_names: - # Per queue, not per cycle: a cycle runs one task from EACH queue, and - # HEALTH_STALE is sized for one task, so a per-cycle stamp would let - # two legitimate long tasks add up to a false "stale". - self._last_poll_monotonic = time.monotonic() + # HEALTH_STALE is sized for one batch (see the constant). + self._last_poll_monotonic = time.monotonic() + try: + messages = self._claim() + except Exception: + logger.exception( + "PG-queue consumer: claim failed for queues %r", self.queue_names + ) + return 0 + for handled, message in enumerate(messages): try: - messages = self._client.read( - queue_name, vt_seconds=self.lease_seconds, qty=self.batch_size - ) - for message in messages: - self._handle(message) - total += len(messages) + self._handle(message) except Exception: + unstarted = [m.msg_id for m in messages[handled + 1 :]] logger.exception( - "PG-queue consumer: poll failed for queue %r; " - "continuing with the other queues", - queue_name, + "PG-queue consumer: handling msg_id=%s (queue=%r) failed; it " + "redelivers when its lease lapses%s", + message.msg_id, + message.queue_name, + f", as do the unstarted msg_ids {unstarted}" if unstarted else "", ) - return total + return handled + return len(messages) + + def _claim(self) -> list[QueueMessage]: + """One claim for this cycle, advancing the multi-queue rotation.""" + if len(self.queue_names) == 1: + return self._client.read( + self.queue_names[0], vt_seconds=self.lease_seconds, qty=self.batch_size + ) + messages = self._client.read_any( + self.queue_names, + vt_seconds=self.lease_seconds, + qty=self.batch_size, + start=self._next_queue, + ) + # One past the queue served last, so the queues after it go first next + # time; an empty cycle just moves on by one. + if messages: + served = self.queue_names.index(messages[-1].queue_name) + else: + served = self._next_queue + self._next_queue = (served + 1) % len(self.queue_names) + return messages @contextlib.contextmanager def _lease_renewal(self, msg_id: int) -> Iterator[None]: @@ -545,6 +581,17 @@ def _resolve_runnable_task( def _handle(self, message: QueueMessage) -> None: payload = message.message + if not isinstance(payload, dict): + # A raw writer can store any JSON; a non-object would raise below on + # every delivery and stop its batch each time, so drop it as malformed. + logger.error( + "PG-queue consumer: payload is not a JSON object (msg_id=%s) — " + "dropping malformed message: %r", + message.msg_id, + payload, + ) + self._client.delete(message.msg_id) + return task_name = payload.get("task_name") # Request-reply (executor RPC): a unique key the dispatching caller is # blocking on. Read up front so the drop branches below can store a diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index 6a0b52bb66..fa5f8da603 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -24,7 +24,7 @@ import pytest 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 +from queue_backend.pg_queue.client import _SEND_RETRY_BACKOFF_SECONDS, _dequeue_any_sql from queue_backend.pg_queue.connection import create_pg_connection from queue_backend.pg_queue.reaper import promote_due_scheduled, rearm_expired_claims from queue_backend.pg_queue.schema import qualified @@ -119,7 +119,9 @@ def test_read_runs_skip_locked_dequeue(self): assert "ORDER BY priority DESC" in sql # fairness L3 claim order # Param order follows the %s positions: queue_name, qty, vt_seconds. assert params == ("q1", 3, 15) - assert msgs == [QueueMessage(msg_id=7, message={"k": "v"}, read_ct=1)] + assert msgs == [ + QueueMessage(msg_id=7, message={"k": "v"}, read_ct=1, queue_name="q1") + ] conn.commit.assert_called_once() def test_delete_returns_true_when_row_removed(self): @@ -214,6 +216,54 @@ def test_owned_conn_recovered_when_rollback_fails(self, monkeypatch): conn.close.assert_called_once() assert client._conn is None + def test_owned_conn_dropped_on_bare_database_error(self, monkeypatch): + # Autocommit rollback() sends nothing, so it can't prove the connection + # dead; a bare DatabaseError is treated as dead (a heuristic). + client, conn, factory = self._owned_client( + monkeypatch, execute_raises=psycopg2.DatabaseError("server gone") + ) + with pytest.raises(psycopg2.DatabaseError): + client.read("q") + conn.close.assert_called_once() + assert client._conn is None + + def test_owned_conn_kept_on_statement_error(self, monkeypatch): + # A statement error mapped to a SQLSTATE subclass keeps the connection. + client, conn, factory = self._owned_client( + monkeypatch, execute_raises=psycopg2.IntegrityError("dup") + ) + with pytest.raises(psycopg2.IntegrityError): + client.read("q") + conn.close.assert_not_called() + assert client._conn is conn + + def test_injected_conn_kept_on_bare_database_error(self): + cur = MagicMock() + cur.execute.side_effect = psycopg2.DatabaseError("server gone") + conn = MagicMock(closed=0) + conn.cursor.return_value = _CursorCtx(cur) + client = PgQueueClient(conn=conn) + with pytest.raises(psycopg2.DatabaseError): + client.read("q") + conn.close.assert_not_called() + + def test_owned_conn_is_autocommit(self, monkeypatch): + # One statement per operation: an explicit transaction would only add a + # BEGIN + COMMIT round trip to each claim (UN-4176). + conn = MagicMock(closed=0, autocommit=False) + monkeypatch.setattr( + "queue_backend.pg_queue.client.create_pg_connection", lambda: conn + ) + PgQueueClient().conn + assert conn.autocommit is True + + def test_injected_conn_autocommit_untouched(self): + # An injected conn may be the caller's open transaction (the agentic + # fan-out outbox) — flipping it would commit that transaction early. + conn = MagicMock(closed=0, autocommit=False) + PgQueueClient(conn=conn).conn + assert conn.autocommit is False + def test_injected_conn_never_closed_on_error(self): cur = MagicMock() cur.execute.side_effect = psycopg2.OperationalError("dead") @@ -534,6 +584,7 @@ def test_send_read_delete_roundtrip(self, pg_conn, queue_name): msgs = client.read(queue_name, vt_seconds=30, qty=10) assert [m.msg_id for m in msgs] == [msg_id] assert msgs[0].message == {"hello": "world"} + assert msgs[0].queue_name == queue_name assert client.delete(msg_id) is True assert client.read(queue_name, vt_seconds=30, qty=10) == [] # gone @@ -867,5 +918,234 @@ def test_no_double_delivery_across_readers(self, pg_conn, queue_name): conn_b.close() +@pytest.fixture +def queue_names(pg_conn, queue_name): + """Three sibling queues for the multi-queue claim; cleaned up by name.""" + names = [f"{queue_name}_{s}" for s in ("a", "b", "c")] + yield names + pg_conn.rollback() + with pg_conn.cursor() as cur: + cur.execute("DELETE FROM pg_queue_message WHERE queue_name = ANY(%s)", (names,)) + pg_conn.commit() + + +class TestReadAnyIntegration: + """``read_any``: one claim across several queues (UN-4176).""" + + def test_empty_queues_claim_nothing(self, pg_conn, queue_names): + assert PgQueueClient(conn=pg_conn).read_any(queue_names, vt_seconds=30) == [] + + def test_claims_from_whichever_queue_has_work(self, pg_conn, queue_names): + client = PgQueueClient(conn=pg_conn) + mid = client.send(queue_names[2], {"n": "c"}) + msgs = client.read_any(queue_names, vt_seconds=30, start=0) + assert [m.msg_id for m in msgs] == [mid] + assert msgs[0].read_ct == 1 + assert msgs[0].queue_name == queue_names[2] + with pg_conn.cursor() as cur: + # The lease must hold: vt pushed ~vt_seconds out, not left at now(). + cur.execute( + "SELECT state, vt > now() + interval '20 seconds' " + "FROM pg_queue_message WHERE msg_id = %s", + (mid,), + ) + assert cur.fetchone() == ("claimed", True) + pg_conn.commit() + assert client.read_any(queue_names, vt_seconds=30) == [] + + def test_start_picks_the_first_queue_probed(self, pg_conn, queue_names): + # Every queue has work; the rotating start decides which is served, so + # no queue starves behind an always-full one. + client = PgQueueClient(conn=pg_conn) + for name in queue_names: + client.send(name, {"q": name}) + served = [ + client.read_any(queue_names, vt_seconds=30, start=s)[0].message["q"] + for s in (1, 2, 0) + ] + assert served == [queue_names[1], queue_names[2], queue_names[0]] + + def test_start_wraps_and_skips_empty_queues(self, pg_conn, queue_names): + client = PgQueueClient(conn=pg_conn) + client.send(queue_names[0], {"q": "a"}) + client.send(queue_names[1], {"q": "b"}) + # start=2 → c (empty) is skipped, wraps to a. + msgs = client.read_any(queue_names, vt_seconds=30, start=2) + assert [m.message["q"] for m in msgs] == ["a"] + + def test_priority_orders_within_a_queue(self, pg_conn, queue_names): + client = PgQueueClient(conn=pg_conn) + client.send(queue_names[0], {"n": "low"}, priority=1) + client.send(queue_names[0], {"n": "high"}, priority=9) + msgs = client.read_any(queue_names, vt_seconds=30) + assert [m.message["n"] for m in msgs] == ["high"] + + def test_batch_orders_by_priority_within_a_queue(self, pg_conn, queue_names): + client = PgQueueClient(conn=pg_conn) + client.send(queue_names[0], {"n": "low"}, priority=1) + client.send(queue_names[0], {"n": "high"}, priority=9) + msgs = client.read_any(queue_names, vt_seconds=30, qty=2) + assert [m.message["n"] for m in msgs] == ["high", "low"] + + def test_duplicate_queue_names_still_fill_the_batch(self, pg_conn, queue_names): + # Probed twice, a repeated queue could hand both probes the same row and + # return fewer than qty messages. + # A repeated queue's second probe returns the same rows again and takes + # batch slots, crowding out the next queue's message. + client = PgQueueClient(conn=pg_conn) + client.send(queue_names[0], {"n": 1}) + client.send(queue_names[0], {"n": 2}) + client.send(queue_names[1], {"n": 3}) + queues = [queue_names[0], queue_names[0], queue_names[1]] + msgs = client.read_any(queues, vt_seconds=30, qty=3) + assert [m.message["n"] for m in msgs] == [1, 2, 3] + + def test_batch_spans_queues_in_rotation_order(self, pg_conn, queue_names): + client = PgQueueClient(conn=pg_conn) + client.send(queue_names[0], {"n": "a1"}) + client.send(queue_names[1], {"n": "b1"}) + client.send(queue_names[1], {"n": "b2"}) + msgs = client.read_any(queue_names, vt_seconds=30, qty=3, start=1) + assert [m.message["n"] for m in msgs] == ["b1", "b2", "a1"] + + def test_unchosen_candidates_stay_ready(self, pg_conn, queue_names): + # The per-queue probe locks a candidate in every queue; only the chosen + # one may be claimed — the others must still be claimable afterwards. + client = PgQueueClient(conn=pg_conn) + ids = [client.send(name, {"q": name}) for name in queue_names] + claimed = client.read_any(queue_names, vt_seconds=30, qty=1) + assert len(claimed) == 1 + with pg_conn.cursor() as cur: + cur.execute( + "SELECT msg_id, state FROM pg_queue_message WHERE msg_id = ANY(%s)", + (ids,), + ) + states = dict(cur.fetchall()) + pg_conn.commit() + assert sorted(states.values()) == ["claimed", "ready", "ready"] + + def test_held_claim_is_skipped_not_waited_on(self, pg_conn, queue_names): + # Reader A's claim transaction is still open — its probe locks are held — + # when reader B claims. B must neither block on those rows (lock_timeout + # turns a wait into an error) nor return a row reader A has locked: it + # takes the next message. + client = PgQueueClient(conn=pg_conn) + ids = [client.send(name, {"q": name}) for name in queue_names for _ in range(2)] + n = len(queue_names) + conn_a = create_pg_connection(env_prefix="TEST_DB_") + conn_b = create_pg_connection(env_prefix="TEST_DB_") + try: + with conn_a.cursor() as cur: + # Same params read_any(start=0, qty=1, vt_seconds=30) sends, + # but left uncommitted so A's locks stay held. + cur.execute( + _dequeue_any_sql(), + (list(queue_names), 1, 0, n, n, 1, 30, 0, n, n), + ) + a_ids = [r[0] for r in cur.fetchall()] + with conn_b.cursor() as cur: + cur.execute("SET lock_timeout = '2s'") + conn_b.commit() + b_ids = [ + m.msg_id + for m in PgQueueClient(conn=conn_b).read_any( + queue_names, vt_seconds=30, start=0 + ) + ] + conn_a.commit() + finally: + conn_a.close() + conn_b.close() + # Both start at queue 0: A holds its head, so B gets that queue's second. + assert a_ids == [ids[0]] + assert b_ids == [ids[1]] + + def test_concurrent_readers_never_double_claim(self, pg_conn, queue_names): + # Released together by a barrier; results come back through futures so a + # reader's exception fails the test instead of dying in its thread. The + # interleaving itself is not guaranteed — the test above pins the + # contended claim deterministically. + import threading + from concurrent.futures import ThreadPoolExecutor + + client = PgQueueClient(conn=pg_conn) + for i in range(60): + client.send(queue_names[i % 3], {"i": i}) + conns = [create_pg_connection(env_prefix="TEST_DB_") for _ in range(2)] + barrier = threading.Barrier(2) + + def drain(conn, start): + reader = PgQueueClient(conn=conn) + claimed: list[list[int]] = [] + barrier.wait(timeout=10) + while msgs := reader.read_any(queue_names, vt_seconds=30, start=start): + claimed.append([m.msg_id for m in msgs]) + start += 1 + return claimed + + try: + with ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(drain, conn, s) for s, conn in enumerate(conns)] + batches = [f.result(timeout=60) for f in futures] + finally: + for conn in conns: + conn.close() + assert all(len(b) == 1 for reader in batches for b in reader) # qty=1 + seen = [{mid for b in reader for mid in b} for reader in batches] + assert seen[0].isdisjoint(seen[1]) + assert len(seen[0]) + len(seen[1]) == 60 + + def test_owned_connection_claims_without_a_transaction( + self, pg_conn, queue_names, monkeypatch + ): + # The autocommit claim must still persist (no COMMIT is sent). + monkeypatch.setattr( + "queue_backend.pg_queue.client.create_pg_connection", + lambda: create_pg_connection(env_prefix="TEST_DB_"), + ) + with PgQueueClient() as owned: + mid = owned.send(queue_names[0], {"n": 1}) + assert [m.msg_id for m in owned.read_any(queue_names, vt_seconds=30)] == [mid] + assert owned.conn.autocommit is True + with pg_conn.cursor() as cur: + cur.execute("SELECT state FROM pg_queue_message WHERE msg_id = %s", (mid,)) + assert cur.fetchone()[0] == "claimed" + pg_conn.commit() + + def test_owned_connection_reconnects_after_backend_termination( + self, pg_conn, queue_names, monkeypatch + ): + # On an owned autocommit connection rollback() sends nothing, so it can't + # prove the connection dead; here the termination surfaces as a + # connection-dead error: the next call raises, drops the handle, and the + # call after reconnects. + monkeypatch.setattr( + "queue_backend.pg_queue.client.create_pg_connection", + lambda: create_pg_connection(env_prefix="TEST_DB_"), + ) + with PgQueueClient() as owned: + owned.read_any(queue_names, vt_seconds=30) + old_pid = owned.conn.get_backend_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)", (old_pid,)) + pg_conn.commit() + with pytest.raises(psycopg2.OperationalError): + owned.read_any(queue_names, vt_seconds=30) + assert owned.read_any(queue_names, vt_seconds=30) == [] + assert owned.conn.get_backend_pid() != old_pid + + def test_rejects_bad_args(self): + # Validation runs before any SQL, so no database is needed. + client = PgQueueClient(conn=MagicMock()) + with pytest.raises(ValueError): + client.read_any([], vt_seconds=30) + with pytest.raises(ValueError): + client.read_any(["q"], vt_seconds=0) + with pytest.raises(ValueError): + client.read_any(["q"], vt_seconds=30, qty=0) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 80555b3a80..93f22892d2 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -51,8 +51,10 @@ def _clear_calls(): _calls.clear() -def _msg(msg_id, payload, *, read_ct=1): - return QueueMessage(msg_id=msg_id, message=payload, read_ct=read_ct) +def _msg(msg_id, payload, *, read_ct=1, queue_name="qa"): + return QueueMessage( + msg_id=msg_id, message=payload, read_ct=read_ct, queue_name=queue_name + ) def _ok_payload(x, y=0): @@ -666,7 +668,7 @@ def test_poll_once_refreshes_heartbeat(self): assert consumer.seconds_since_last_poll() < 1.0 def test_heartbeat_stamped_before_read(self): - # Pins the headline design: the stamp lands before each queue read (not + # Pins the headline design: the stamp lands before each claim (not # after it), so a task running longer than the threshold still trips # the probe. A bottom-of-poll stamp would pass test_poll_once_refreshes # but fail here. @@ -681,23 +683,19 @@ def test_heartbeat_stamped_before_read(self): consumer.poll_once() assert seen["during"] > before # refreshed BEFORE read ran, not after - def test_heartbeat_restamped_before_each_queue_read(self): - # UN-4223: a cycle runs one task per queue and HEALTH_STALE is sized for - # one task, so each queue's read must start from a fresh stamp — a single - # per-cycle stamp would add two long tasks into a false "stale". + def test_heartbeat_stamped_before_multi_queue_claim(self): + # The multi-queue path claims once per cycle (read_any); the stamp must + # still land before that claim, not after the task it returns. client = MagicMock() consumer = PgQueueConsumer(["q1", "q2"], client=client) - seen: list[float] = [] - - def _read(*_a, **_k): - seen.append(consumer._last_poll_monotonic) - consumer._last_poll_monotonic -= 5000 # the first queue's long task - return [] - - client.read.side_effect = _read + before = consumer._last_poll_monotonic + seen: dict[str, float] = {} + client.read_any.side_effect = lambda *a, **k: ( + seen.setdefault("during", consumer._last_poll_monotonic), + [], + )[1] consumer.poll_once() - assert len(seen) == 2 - assert seen[1] > seen[0] - 1 # re-stamped, not inherited from queue 1 + assert seen["during"] > before def test_health_server_disabled_without_port(self): # No port configured → no server bound (opt-in). @@ -873,49 +871,190 @@ def test_enqueue_poll_execute_ack(self, pg_client): ) pg_client.conn.commit() + def test_multi_queue_poll_executes_and_acks_on_postgres(self, pg_client): + # The multi-queue claim end to end: real read_any rows feed the real + # dispatch and the delete() ack, which the stubbed tests can't see. + base = f"test_consumer_mq_{os.getpid()}" + queues = [f"{base}_a", f"{base}_empty", f"{base}_b"] + try: + pg_client.send(queues[0], to_payload("test_pg_consumer.ok", args=[11])) + pg_client.send(queues[2], to_payload("test_pg_consumer.ok", args=[22])) + consumer = PgQueueConsumer(queues, client=pg_client) + claimed = [consumer.poll_once() for _ in range(3)] + assert claimed == [1, 1, 0] + assert (11, 0) in _calls and (22, 0) in _calls + assert pg_client.read_any(queues, vt_seconds=30, qty=10) == [] # acked + finally: + with pg_client.conn.cursor() as cur: + cur.execute( + "DELETE FROM pg_queue_message WHERE queue_name = ANY(%s)", (queues,) + ) + pg_client.conn.commit() + + def test_busy_queues_take_equal_turns_around_an_empty_one(self, pg_client): + # An empty queue between two busy ones must not give the queue after it + # an extra turn: claims alternate. + base = f"test_consumer_rr_{os.getpid()}" + queues = [f"{base}_a", f"{base}_empty", f"{base}_b"] + try: + for _ in range(3): + pg_client.send(queues[0], to_payload("test_pg_consumer.ok", args=[1])) + pg_client.send(queues[2], to_payload("test_pg_consumer.ok", args=[1])) + consumer = PgQueueConsumer(queues, client=pg_client) + served: list[str] = [] + consumer._handle = lambda m: served.append(m.queue_name) + for _ in range(6): + consumer.poll_once() + assert served == [queues[0], queues[2]] * 3 + finally: + with pg_client.conn.cursor() as cur: + cur.execute( + "DELETE FROM pg_queue_message WHERE queue_name = ANY(%s)", (queues,) + ) + pg_client.conn.commit() + class TestMultiQueue: - """9f: one consumer drains several queues round-robin.""" + """9f: one consumer drains several queues — one claim per cycle (UN-4176).""" - def test_round_robin_reads_each_queue_and_aggregates(self): + def test_multi_queue_claims_once_per_cycle(self): client = MagicMock() - # qa → 2 msgs, qb → 1 msg; poll_once returns the total across both. - client.read.side_effect = [ - [_msg(1, _ok_payload(1)), _msg(2, _ok_payload(2))], - [_msg(3, _ok_payload(3))], - ] + client.read_any.return_value = [_msg(1, _ok_payload(1))] claimed = PgQueueConsumer(["qa", "qb"], client=client).poll_once() - assert claimed == 3 - # Read once per queue, in order. - assert [c.args[0] for c in client.read.call_args_list] == ["qa", "qb"] + assert claimed == 1 + client.read_any.assert_called_once() + client.read.assert_not_called() # never one statement per queue + assert client.read_any.call_args.args[0] == ["qa", "qb"] + + def test_single_queue_keeps_the_single_queue_claim(self): + client = MagicMock() + client.read.return_value = [] + PgQueueConsumer(["q"], client=client).poll_once() + client.read.assert_called_once() + client.read_any.assert_not_called() - def test_empty_queue_does_not_starve_the_others(self): + def test_next_claim_starts_past_the_queue_served(self): + # Rotating by one regardless of what was served would hand a busy queue + # an extra turn for every empty queue before it; starting one past the + # queue actually served keeps busy queues on equal turns. client = MagicMock() - client.read.side_effect = [[], [_msg(1, _ok_payload(1))]] - assert PgQueueConsumer(["empty", "busy"], client=client).poll_once() == 1 + client.read_any.side_effect = [ + [_msg(1, _ok_payload(1), queue_name="qc")], # start 0 fell through to qc + [_msg(2, _ok_payload(2), queue_name="qa")], + [], + ] + consumer = PgQueueConsumer(["qa", "qb", "qc"], client=client) + for _ in range(3): + consumer.poll_once() + assert [c.kwargs["start"] for c in client.read_any.call_args_list] == [0, 0, 1] + assert consumer._next_queue == 2 # empty cycle moves on by one + + @pytest.mark.parametrize( + ("failing", "expected_done", "unstarted"), + [(1, 0, [2, 3]), (2, 1, [3]), (3, 2, [])], + ) + def test_handler_failure_stops_the_batch_and_counts_the_work_done( + self, caplog, failing, expected_done, unstarted + ): + # Running the rest after a failure could, in a DB outage, run work whose + # ack cannot land; so the batch stops, but work done still counts (0 on + # a first-message failure, so run() backs off). + client = MagicMock() + client.read_any.return_value = [ + _msg(1, _ok_payload(1), queue_name="qa"), + _msg(2, _ok_payload(2), queue_name="qa"), + _msg(3, _ok_payload(3), queue_name="qb"), + ] + consumer = PgQueueConsumer( + ["qa", "qb"], client=client, batch_size=3, lease_seconds=60, vt_seconds=60 + ) + handled: list[int] = [] + + def _handle(message): + if message.msg_id == failing: + raise RuntimeError("boom") + handled.append(message.msg_id) + + consumer._handle = _handle + with caplog.at_level(logging.ERROR): + assert consumer.poll_once() == expected_done + assert handled == list(range(1, failing)) + assert f"msg_id={failing} " in caplog.text + if unstarted: + assert f"unstarted msg_ids {unstarted}" in caplog.text + else: + assert "unstarted" not in caplog.text + # Rotation follows the claim, not what ran: the batch ended on qb. + assert consumer._next_queue == 0 + + def test_non_object_payload_is_dropped_not_raised(self): + # A non-dict payload must not raise out of _handle (it would stop its + # batch on every delivery); it is dropped and acked like any malformed one. + client = MagicMock() + client.read_any.return_value = [ + _msg(1, ["not", "an", "object"]), + _msg(2, _ok_payload(7)), + ] + consumer = PgQueueConsumer( + ["qa", "qb"], client=client, batch_size=2, lease_seconds=60, vt_seconds=60 + ) + assert consumer.poll_once() == 2 + assert (7, 0) in _calls # the batch-mate still ran + assert [c.args[0] for c in client.delete.call_args_list] == [1, 2] + + def test_rotation_follows_the_last_message_of_a_batch(self): + client = MagicMock() + client.read_any.side_effect = [ + [ + _msg(1, _ok_payload(1), queue_name="qa"), + _msg(2, _ok_payload(2), queue_name="qc"), + ], + [], + ] + consumer = PgQueueConsumer( + ["qa", "qb", "qc"], + client=client, + batch_size=3, + lease_seconds=60, + vt_seconds=60, + ) + consumer._handle = lambda m: None + consumer.poll_once() + consumer.poll_once() + # Batch ended on qc (index 2), so the next claim starts at qa (0). + assert [c.kwargs["start"] for c in client.read_any.call_args_list] == [0, 0] + + def test_start_queue_rotates_every_cycle(self): + # With nothing claimed the start still moves on, so every queue gets to + # be probed first in turn. + client = MagicMock() + client.read_any.return_value = [] + consumer = PgQueueConsumer(["qa", "qb", "qc"], client=client) + for _ in range(4): + consumer.poll_once() + assert [c.kwargs["start"] for c in client.read_any.call_args_list] == [ + 0, + 1, + 2, + 0, + ] def test_rejects_empty_queue_list(self): client = MagicMock() with pytest.raises(ValueError, match="queue_names"): PgQueueConsumer([], client=client) - def test_one_queue_failing_does_not_abort_the_others(self): - """A read failure on one queue is isolated: the already-acked work this - cycle still counts (so run() doesn't take the empty backoff path), and - the failure doesn't propagate out of poll_once. - """ + def test_claim_failure_counts_as_an_empty_cycle(self): + """A failed claim is logged, not raised, and returns 0 so run() backs off.""" client = MagicMock() - client.read.side_effect = [[_msg(1, _ok_payload(1))], RuntimeError("boom")] - claimed = PgQueueConsumer(["qa", "qb"], client=client).poll_once() - assert claimed == 1 # qa's message counted despite qb raising - assert _calls == [(1, 0)] # qa task body ran - client.delete.assert_called_once_with(1) # qa acked + client.read_any.side_effect = RuntimeError("boom") + assert PgQueueConsumer(["qa", "qb"], client=client).poll_once() == 0 - def test_read_uses_batch_size_qty_per_queue(self): + def test_claim_uses_batch_size_qty(self): client = MagicMock() - client.read.return_value = [] + client.read_any.return_value = [] PgQueueConsumer(["qa", "qb"], client=client, batch_size=2).poll_once() - assert [c.kwargs["qty"] for c in client.read.call_args_list] == [2, 2] + assert client.read_any.call_args.kwargs["qty"] == 2 def test_duplicate_queues_are_deduped(self): """A duplicate (e.g. env "q,q") must not double-read the same queue."""