From 82b43fe63777f324841a41817ffd919eb5249f75 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 03:25:27 +0000 Subject: [PATCH 1/9] UN-4176 [FIX] Cut idle PG-queue polling to one autocommit claim per cycle Idle PG-queue consumers were the largest steady source of DB statements: every poll cycle sent one claim per queue, each wrapped in BEGIN/COMMIT. - PgQueueClient: connections the client opens itself run in autocommit, so each single-statement operation (claim, ack, lease renewal, send) is one statement instead of three. Injected connections are untouched (they may be the caller's open transaction), and create_pg_connection stays non-autocommit for pg_barrier. - PgQueueClient.read_any: claims across several queues in one statement, with an indexed per-queue probe (LATERAL on the partial claim index) so cost does not grow with the backlog. - PgQueueConsumer.poll_once: multi-queue consumers make one read_any claim per cycle. The first queue probed rotates every cycle so no queue starves; priority still orders messages within a queue. Single-queue consumers keep the existing claim SQL. No queue lists change. Measured locally with all PG workers idle: queue polling 89.0 -> 9.2 statements/s; API deployment still completes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/queue_backend/pg_queue/client.py | 96 ++++++++++++- workers/queue_backend/pg_queue/consumer.py | 68 +++++---- workers/tests/test_pg_queue_client.py | 160 +++++++++++++++++++++ workers/tests/test_pg_queue_consumer.py | 84 +++++------ 4 files changed, 337 insertions(+), 71 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index c87c1b06e2..dfba29f277 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -136,6 +136,45 @@ 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 first queue probed changes + each cycle. 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, locked.ord, q.priority +) +SELECT msg_id, message, read_ct + FROM claimed + ORDER BY (ord - 1 - %s + %s) %% %s, priority DESC, msg_id +""" + + @dataclass(frozen=True, slots=True) class QueueMessage: """A claimed queue message. @@ -200,7 +239,16 @@ 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 an explicit transaction + # only adds a BEGIN + COMMIT round trip to each (3 statements per + # claim instead of 1 — most of the idle-poll load on the DB). + # Autocommit is set only on a connection this client creates: an + # injected one may be the caller's open transaction (the agentic + # fan-out outbox), 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 @@ -391,6 +439,52 @@ def read( QueueMessage(msg_id=int(r[0]), message=r[1], read_ct=int(r[2])) 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`; replaces one :meth:`read` per queue + per poll cycle, so an idle multi-queue consumer costs one statement per + cycle instead of one per queue. 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: the caller advances ``start`` each cycle + so every non-empty queue keeps getting a turn, as with the per-queue + loop this replaces. ``priority`` still orders messages within a queue. + + The per-queue probe locks up to ``qty`` rows in each queue, and only the + chosen ones are claimed; the rest stay ``ready`` and are unlocked when + the statement commits, so a concurrent reader skips them for those few + milliseconds at most. + """ + 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") + n = len(queue_names) + 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(), + (list(queue_names), qty, *rotate, qty, vt_seconds, *rotate), + ) + rows = cur.fetchall() + return [ + QueueMessage(msg_id=int(r[0]), message=r[1], read_ct=int(r[2])) for r in rows + ] + def set_vt(self, msg_id: int, vt_seconds: int) -> bool: """Re-park a claimed message: extend its lease by another ``vt_seconds``. diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index c31b427957..00dfc52dde 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -311,14 +311,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 + # consumer drains both file_processing and api_file_processing. All are + # claimed from in one statement per cycle, starting from a queue that + # rotates each cycle (poll_once): every non-empty queue gets a turn, so + # none starves. Copy + de-dup (order-preserving): a duplicate would # double-read a queue per cycle, 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 +367,47 @@ 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 claimed. + + A single-queue consumer reads its queue. A multi-queue consumer claims + across all its queues in ONE statement (:meth:`PgQueueClient.read_any`) + rather than one read per queue, so an idle cycle costs one statement + however many queues it drains. The queue probed first rotates every + cycle, so each non-empty queue still gets a turn (no starvation), as it + did when every queue was read once per cycle. A read/handle failure is + logged and the cycle counts as empty, as a failed queue did before. """ - 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() - try: + # HEALTH_STALE is sized for one task, and a cycle now runs one batch. + self._last_poll_monotonic = time.monotonic() + try: + if len(self.queue_names) == 1: messages = self._client.read( - queue_name, vt_seconds=self.lease_seconds, qty=self.batch_size + self.queue_names[0], + vt_seconds=self.lease_seconds, + qty=self.batch_size, ) - for message in messages: - self._handle(message) - total += len(messages) - except Exception: - logger.exception( - "PG-queue consumer: poll failed for queue %r; " - "continuing with the other queues", - queue_name, + else: + messages = self._client.read_any( + self.queue_names, + vt_seconds=self.lease_seconds, + qty=self.batch_size, + start=self._next_queue, ) - return total + self._next_queue = (self._next_queue + 1) % len(self.queue_names) + for message in messages: + self._handle(message) + except Exception: + logger.exception( + "PG-queue consumer: poll failed for queues %r", self.queue_names + ) + return 0 + return len(messages) @contextlib.contextmanager def _lease_renewal(self, msg_id: int) -> Iterator[None]: diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index 6a0b52bb66..6e33c94c6f 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -214,6 +214,23 @@ def test_owned_conn_recovered_when_rollback_fails(self, monkeypatch): conn.close.assert_called_once() assert client._conn is None + 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") @@ -867,5 +884,148 @@ 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 prefix.""" + 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 + 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_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_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_concurrent_readers_never_double_claim(self, pg_conn, queue_names): + import threading + + client_a = PgQueueClient(conn=pg_conn) + for i in range(60): + client_a.send(queue_names[i % 3], {"i": i}) + conn_b = create_pg_connection(env_prefix="TEST_DB_") + seen: list[list[int]] = [[], []] + violations: list[int] = [] + + def drain(client, out, start): + while True: + msgs = client.read_any(queue_names, vt_seconds=30, start=start) + if not msgs: + return + if len(msgs) > 1: + violations.append(len(msgs)) + out.extend(m.msg_id for m in msgs) + start += 1 + + try: + threads = [ + threading.Thread(target=drain, args=(client_a, seen[0], 0)), + threading.Thread( + target=drain, args=(PgQueueClient(conn=conn_b), seen[1], 1) + ), + ] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + finally: + conn_b.close() + assert violations == [] + assert set(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_rejects_bad_args(self, pg_conn): + client = PgQueueClient(conn=pg_conn) + 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..bd32e5ccd3 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -681,23 +681,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). @@ -875,47 +871,55 @@ def test_enqueue_poll_execute_ack(self, pg_client): 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_empty_queue_does_not_starve_the_others(self): + def test_single_queue_keeps_the_single_queue_claim(self): client = MagicMock() - client.read.side_effect = [[], [_msg(1, _ok_payload(1))]] - assert PgQueueConsumer(["empty", "busy"], client=client).poll_once() == 1 + client.read.return_value = [] + PgQueueConsumer(["q"], client=client).poll_once() + client.read.assert_called_once() + client.read_any.assert_not_called() + + def test_start_queue_rotates_every_cycle(self): + # The rotating start is what keeps an always-full first queue from + # starving the rest. + 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.""" From 36b4596dd07822895002878fa650a8d151f92b8d Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 04:08:51 +0000 Subject: [PATCH 2/9] UN-4176 Pin the contended read_any claim deterministically in tests The concurrent-readers test could pass with one reader draining every message before the other read, and a reader's exception died in its thread. Add a test that holds reader A's claim transaction open while reader B claims (lock_timeout turns a wait into an error), and release the parallel drain from a barrier with results collected through futures. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/tests/test_pg_queue_client.py | 89 ++++++++++++++++++--------- 1 file changed, 61 insertions(+), 28 deletions(-) diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index 6e33c94c6f..27b27de78c 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 @@ -963,41 +963,74 @@ def test_unchosen_candidates_stay_ready(self, pg_conn, queue_names): 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 one: 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_a = PgQueueClient(conn=pg_conn) + client = PgQueueClient(conn=pg_conn) for i in range(60): - client_a.send(queue_names[i % 3], {"i": i}) - conn_b = create_pg_connection(env_prefix="TEST_DB_") - seen: list[list[int]] = [[], []] - violations: list[int] = [] - - def drain(client, out, start): - while True: - msgs = client.read_any(queue_names, vt_seconds=30, start=start) - if not msgs: - return - if len(msgs) > 1: - violations.append(len(msgs)) - out.extend(m.msg_id for m in msgs) + 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: - threads = [ - threading.Thread(target=drain, args=(client_a, seen[0], 0)), - threading.Thread( - target=drain, args=(PgQueueClient(conn=conn_b), seen[1], 1) - ), - ] - for t in threads: - t.start() - for t in threads: - t.join(timeout=60) + 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: - conn_b.close() - assert violations == [] - assert set(seen[0]).isdisjoint(seen[1]) + 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( From 351b04c93e9ff12da9a75496cf5a54ce92468784 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 05:03:04 +0000 Subject: [PATCH 3/9] UN-4176 Address round-1 self-review: fair rotation, partial counts, docs - Rotate the multi-queue claim from one past the queue actually served (QueueMessage.queue_name), so busy queues take equal turns; advancing by one per cycle gave a busy queue an extra turn per empty queue before it. - poll_once separates the claim from handling: a handler failure logs the msg_ids left unstarted and still counts the messages processed. - Document the accepted read_any trade-off (the probe locks every queue's head; a racing reader can see [] and back off) and the autocommit one-statement rule; fix stale "each queue read" comments. - Tests: read_any lease (vt), priority order within a qty>1 batch, queue_name on read/read_any, reconnect after pg_terminate_backend on an owned autocommit connection, equal turns around an empty queue on real Postgres; bad-args test no longer needs a database. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/pg_queue_consumer/supervisor.py | 2 +- workers/queue_backend/pg_queue/client.py | 93 ++++++++++++++-------- workers/queue_backend/pg_queue/consumer.py | 80 ++++++++++++------- workers/tests/test_pg_queue_client.py | 51 ++++++++++-- workers/tests/test_pg_queue_consumer.py | 71 +++++++++++++++-- 5 files changed, 221 insertions(+), 76 deletions(-) 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 dfba29f277..dd053b8112 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 @@ -140,9 +142,10 @@ 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 first queue probed changes - each cycle. Within a queue the order is the same ``priority DESC, msg_id`` - top-N the single-queue claim uses, on the same partial index. + 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""" @@ -167,9 +170,9 @@ def _dequeue_any_sql() -> str: read_ct = read_ct + 1 FROM locked WHERE q.msg_id = locked.msg_id - RETURNING q.msg_id, q.message, q.read_ct, locked.ord, q.priority + RETURNING q.msg_id, q.message, q.read_ct, q.queue_name, locked.ord, q.priority ) -SELECT msg_id, message, read_ct +SELECT msg_id, message, read_ct, queue_name FROM claimed ORDER BY (ord - 1 - %s + %s) %% %s, priority DESC, msg_id """ @@ -186,11 +189,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; 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()` @@ -240,13 +246,14 @@ def __init__(self, conn: PgConnection | None = None) -> None: def conn(self) -> PgConnection: if self._conn is None: conn = create_pg_connection() - # Every operation here is ONE statement, so an explicit transaction - # only adds a BEGIN + COMMIT round trip to each (3 statements per - # claim instead of 1 — most of the idle-poll load on the DB). + # 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 be the caller's open transaction (the agentic - # fan-out outbox), and create_pg_connection itself stays - # non-autocommit for pg_barrier, which relies on that. + # 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 @@ -255,10 +262,12 @@ def conn(self) -> PgConnection: 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 are + no-ops 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: @@ -269,6 +278,10 @@ 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 and + # cannot fail; there a server termination surfaces as + # OperationalError (is_connection_dead) and sets conn.closed, which + # the discard below checks too. # 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 @@ -436,7 +449,10 @@ 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( @@ -449,21 +465,27 @@ def read_any( ) -> list[QueueMessage]: """Claim up to ``qty`` ready messages from several queues in ONE statement. - Same claim contract as :meth:`read`; replaces one :meth:`read` per queue - per poll cycle, so an idle multi-queue consumer costs one statement per - cycle instead of one per queue. 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. + 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: the caller advances ``start`` each cycle - so every non-empty queue keeps getting a turn, as with the per-queue - loop this replaces. ``priority`` still orders messages within a queue. - - The per-queue probe locks up to ``qty`` rows in each queue, and only the - chosen ones are claimed; the rest stay ``ready`` and are unlocked when - the statement commits, so a concurrent reader skips them for those few - milliseconds at most. + ``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 statement commits. A + reader racing that statement skips the locked rows, so if each queue's + only ready row is locked it gets ``[]`` and takes the consumer's empty- + poll back-off (starting at ``poll_interval``) instead of a message. This + only happens while several consumers contend for the last rows; idle + polls lock nothing. """ if vt_seconds <= 0: raise ValueError(f"vt_seconds must be positive, got {vt_seconds}") @@ -482,7 +504,10 @@ def read_any( ) 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=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 00dfc52dde..1e672b1859 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -96,7 +96,7 @@ # 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 +# unhealthy. The heartbeat is stamped before each claim 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, # under the prefork supervisor with CHILD_WATCHDOG on, a SIGKILL of just that @@ -312,10 +312,10 @@ def __init__( 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. All are - # claimed from in one statement per cycle, starting from a queue that - # rotates each cycle (poll_once): every non-empty queue gets a turn, so - # none starves. Copy + de-dup (order-preserving): a duplicate would - # double-read a queue per cycle, and storing the caller's list by + # 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). @@ -373,42 +373,62 @@ def __init__( self._last_poll_monotonic = time.monotonic() def poll_once(self) -> int: - """Claim + process one batch; returns the number of messages claimed. + """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`) - rather than one read per queue, so an idle cycle costs one statement - however many queues it drains. The queue probed first rotates every - cycle, so each non-empty queue still gets a turn (no starvation), as it - did when every queue was read once per cycle. A read/handle failure is - logged and the cycle counts as empty, as a failed queue did before. + 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 (they stay claimed until their lease lapses, then the reaper + re-arms them) and the messages already processed are still counted. """ - # HEALTH_STALE is sized for one task, and a cycle now runs one batch. + # HEALTH_STALE is sized for one batch (see the constant). self._last_poll_monotonic = time.monotonic() try: - if len(self.queue_names) == 1: - messages = self._client.read( - self.queue_names[0], - vt_seconds=self.lease_seconds, - qty=self.batch_size, - ) - else: - messages = self._client.read_any( - self.queue_names, - vt_seconds=self.lease_seconds, - qty=self.batch_size, - start=self._next_queue, - ) - self._next_queue = (self._next_queue + 1) % len(self.queue_names) - for message in messages: - self._handle(message) + messages = self._claim() except Exception: logger.exception( - "PG-queue consumer: poll failed for queues %r", self.queue_names + "PG-queue consumer: claim failed for queues %r", self.queue_names ) return 0 + for done, message in enumerate(messages): + try: + self._handle(message) + except Exception: + logger.exception( + "PG-queue consumer: handling msg_id=%s failed; unstarted in " + "this batch: %s", + message.msg_id, + [m.msg_id for m in messages[done + 1 :]], + ) + return done 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 and messages[-1].queue_name in self.queue_names: + 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]: """Keep ``msg_id``'s claim alive while a task runs. diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index 27b27de78c..12aa712902 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -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): @@ -551,6 +553,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 @@ -886,7 +889,7 @@ def test_no_double_delivery_across_readers(self, pg_conn, queue_name): @pytest.fixture def queue_names(pg_conn, queue_name): - """Three sibling queues for the multi-queue claim; cleaned up by prefix.""" + """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() @@ -907,10 +910,17 @@ def test_claims_from_whichever_queue_has_work(self, pg_conn, queue_names): 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: - cur.execute("SELECT state FROM pg_queue_message WHERE msg_id = %s", (mid,)) - assert cur.fetchone()[0] == "claimed" + # 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 @@ -939,6 +949,13 @@ def test_priority_orders_within_a_queue(self, pg_conn, queue_names): 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_batch_spans_queues_in_rotation_order(self, pg_conn, queue_names): client = PgQueueClient(conn=pg_conn) client.send(queue_names[0], {"n": "a1"}) @@ -1050,8 +1067,30 @@ def test_owned_connection_claims_without_a_transaction( assert cur.fetchone()[0] == "claimed" pg_conn.commit() - def test_rejects_bad_args(self, pg_conn): - client = PgQueueClient(conn=pg_conn) + def test_owned_connection_reconnects_after_backend_termination( + self, pg_conn, queue_names, monkeypatch + ): + # Autocommit rollback() cannot fail, so recovery rests on the + # termination surfacing 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: + cur.execute("SELECT pg_terminate_backend(%s)", (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): diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index bd32e5ccd3..359bf7533e 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=""): + 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. @@ -869,6 +871,28 @@ def test_enqueue_poll_execute_ack(self, pg_client): ) 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 used to hand the queue after it a + # second turn (2:1). Claims must now 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 — one claim per cycle (UN-4176).""" @@ -889,9 +913,46 @@ def test_single_queue_keeps_the_single_queue_claim(self): client.read.assert_called_once() client.read_any.assert_not_called() + 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_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 + + def test_handler_failure_counts_the_messages_already_processed(self): + client = MagicMock() + client.read_any.return_value = [ + _msg(1, _ok_payload(1)), + _msg(2, _ok_payload(2)), + _msg(3, _ok_payload(3)), + ] + 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 == 2: + raise RuntimeError("boom") + handled.append(message.msg_id) + + consumer._handle = _handle + assert consumer.poll_once() == 1 # msg 1 counted, so run() won't back off + assert handled == [1] # msg 3 left claimed, not started + def test_start_queue_rotates_every_cycle(self): - # The rotating start is what keeps an always-full first queue from - # starving the rest. + # 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) From 7202b32175b7a56e84f857c30bed81777c83a390 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 05:16:33 +0000 Subject: [PATCH 4/9] UN-4176 Address round-2 self-review: batch failures, accepted trade-offs - poll_once skips a message whose handling fails and runs the rest of the batch, so a deterministic failure cannot strand its batch-mates (each redelivery costs them a read_ct attempt); a connection-dead failure still stops the batch, since the rest could not be acked either. - Warn when a claimed message's queue_name is unknown instead of silently rotating by one. - Document the autocommit trade-offs: a claim lost in flight stays claimed until its lease lapses and counts toward read_ct; send()'s ambiguous post-commit window is the INSERT's own reply. Fix the _cursor NOTE and the HEALTH_STALE comment (a batch, not a task). - Tests: rotation follows the last message of a batch, the unknown-name fallback, skip-and-continue, connection loss stopping the batch. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/queue_backend/pg_queue/client.py | 21 +++++-- workers/queue_backend/pg_queue/consumer.py | 56 +++++++++++++----- workers/tests/test_pg_queue_client.py | 9 +-- workers/tests/test_pg_queue_consumer.py | 69 ++++++++++++++++++++-- 4 files changed, 125 insertions(+), 30 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index dd053b8112..19a5746cef 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -284,10 +284,10 @@ def _cursor(self) -> Iterator[Any]: # the discard below checks too. # 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. + # ``psycopg2.DatabaseError`` is intentionally NOT retried by + # ``send()``. On an owned (autocommit) connection it is dropped here + # only if ``conn.closed`` is set; otherwise the next call reuses it + # and, if it is dead, fails with a classified error and is dropped. conn_dead = is_connection_dead(exc) try: conn.rollback() @@ -341,7 +341,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. @@ -486,6 +489,14 @@ def read_any( poll back-off (starting at ``poll_interval``) instead of a message. This only happens while several 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}") diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 1e672b1859..65c2284f8f 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -42,7 +42,7 @@ from ..barrier import callback_recovery_identity from ..fairness import FAIRNESS_HEADER_NAME from .client import PgQueueClient -from .connection import CONN_DEAD_ERRORS +from .connection import CONN_DEAD_ERRORS, is_connection_dead from .liveness import LivenessServer as _BaseLivenessServer from .result_backend import PgResultBackend from .task_payload import to_payload @@ -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 claim 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, @@ -381,10 +382,14 @@ def poll_once(self) -> int: 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 (they stay claimed until their lease lapses, then the reaper - re-arms them) and the messages already processed are still counted. + A failed claim is logged and returns 0, so run() backs off. A message + whose handling fails is logged and skipped; the rest of the batch still + runs, so a deterministic failure cannot strand its batch-mates (each + redelivery costs them a ``read_ct`` attempt toward the poison cap). The + exception: a connection-dead failure stops the batch — the remaining + messages could not be acked either, so running them would only run them + twice; they stay claimed until their lease lapses and the reaper + re-arms them. Returns the number handled successfully. """ # HEALTH_STALE is sized for one batch (see the constant). self._last_poll_monotonic = time.monotonic() @@ -395,18 +400,27 @@ def poll_once(self) -> int: "PG-queue consumer: claim failed for queues %r", self.queue_names ) return 0 - for done, message in enumerate(messages): + handled = 0 + for i, message in enumerate(messages): try: self._handle(message) - except Exception: + except Exception as exc: + if is_connection_dead(exc): + logger.exception( + "PG-queue consumer: connection lost handling msg_id=%s; " + "left claimed until their lease lapses: %s", + message.msg_id, + [m.msg_id for m in messages[i + 1 :]], + ) + return handled logger.exception( - "PG-queue consumer: handling msg_id=%s failed; unstarted in " - "this batch: %s", + "PG-queue consumer: handling msg_id=%s failed; continuing " + "with the rest of the batch", message.msg_id, - [m.msg_id for m in messages[done + 1 :]], ) - return done - return len(messages) + continue + handled += 1 + return handled def _claim(self) -> list[QueueMessage]: """One claim for this cycle, advancing the multi-queue rotation.""" @@ -422,9 +436,19 @@ def _claim(self) -> list[QueueMessage]: ) # 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 and messages[-1].queue_name in self.queue_names: + if not messages: + served = self._next_queue + elif messages[-1].queue_name in self.queue_names: served = self.queue_names.index(messages[-1].queue_name) else: + # read_any always fills queue_name; reaching here means a client that + # doesn't, which would quietly bring back unequal turns. + logger.warning( + "PG-queue consumer: claimed msg_id=%s has unknown queue_name %r; " + "rotating by one", + messages[-1].msg_id, + messages[-1].queue_name, + ) served = self._next_queue self._next_queue = (served + 1) % len(self.queue_names) return messages diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index 12aa712902..ea8884a5fb 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -983,7 +983,8 @@ def test_unchosen_candidates_stay_ready(self, pg_conn, queue_names): 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 one: it takes the next message. + # 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) @@ -1070,9 +1071,9 @@ def test_owned_connection_claims_without_a_transaction( def test_owned_connection_reconnects_after_backend_termination( self, pg_conn, queue_names, monkeypatch ): - # Autocommit rollback() cannot fail, so recovery rests on the - # termination surfacing as a connection-dead error: the next call - # raises, drops the handle, and the call after reconnects. + # On an owned autocommit connection rollback() sends nothing, so recovery + # rests on the termination surfacing 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_"), diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 359bf7533e..bcb5243398 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -14,6 +14,7 @@ import time from unittest.mock import MagicMock, patch +import psycopg2 import pytest from celery import shared_task from queue_backend.fairness import FAIRNESS_HEADER_NAME @@ -872,8 +873,8 @@ def test_enqueue_poll_execute_ack(self, pg_client): 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 used to hand the queue after it a - # second turn (2:1). Claims must now alternate. + # 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: @@ -929,7 +930,9 @@ def test_next_claim_starts_past_the_queue_served(self): 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 - def test_handler_failure_counts_the_messages_already_processed(self): + def test_handler_failure_skips_only_the_failed_message(self): + # A deterministic failure must not strand its batch-mates: they would + # otherwise pay a read_ct attempt on every redelivery. client = MagicMock() client.read_any.return_value = [ _msg(1, _ok_payload(1)), @@ -947,8 +950,64 @@ def _handle(message): handled.append(message.msg_id) consumer._handle = _handle - assert consumer.poll_once() == 1 # msg 1 counted, so run() won't back off - assert handled == [1] # msg 3 left claimed, not started + assert consumer.poll_once() == 2 # msgs 1 and 3 handled + assert handled == [1, 3] + + def test_connection_loss_stops_the_batch(self): + # Once the connection is gone the rest could not be acked either, so + # running them would only run them twice. + client = MagicMock() + client.read_any.return_value = [ + _msg(1, _ok_payload(1)), + _msg(2, _ok_payload(2)), + _msg(3, _ok_payload(3)), + ] + 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 == 2: + raise psycopg2.OperationalError("server closed the connection") + handled.append(message.msg_id) + + consumer._handle = _handle + assert consumer.poll_once() == 1 + assert handled == [1] + + 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_unknown_queue_name_rotates_by_one_and_warns(self, caplog): + client = MagicMock() + client.read_any.side_effect = [[], [_msg(1, _ok_payload(1), queue_name="??")]] + consumer = PgQueueConsumer(["qa", "qb", "qc"], client=client) + consumer._handle = lambda m: None + consumer.poll_once() # empty: start 0 -> 1 + with caplog.at_level(logging.WARNING): + consumer.poll_once() # start 1, unknown name -> 2 + assert consumer._next_queue == 2 + assert "unknown queue_name" in caplog.text def test_start_queue_rotates_every_cycle(self): # With nothing claimed the start still moves on, so every queue gets to From 5c25907e6bd3ae0fd5bb075c95f8d68be3e9a107 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 05:35:37 +0000 Subject: [PATCH 5/9] UN-4176 Address round-3 self-review: drop batch abort, bare DatabaseError - poll_once no longer stops a batch on a connection-classified handler error: the client drops the dead handle and reconnects for the next ack (and the classifier also matches cancellations and deadlocks), so every failure is now logged and skipped. - On a client-owned (autocommit) connection, rollback() can no longer prove a connection dead, so a bare psycopg2.DatabaseError is treated as dead and the handle dropped, keeping what the failed-rollback path did. - read_any trade-off text names single-queue consumers too; comments no longer claim rollback() cannot fail. - Tests: owned conn dropped on a bare DatabaseError, kept on a statement error; injected conn kept; a connection error does not stop the batch. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/queue_backend/pg_queue/client.py | 34 ++++++++++--------- workers/queue_backend/pg_queue/consumer.py | 22 ++++--------- workers/tests/test_pg_queue_client.py | 38 ++++++++++++++++++++-- workers/tests/test_pg_queue_consumer.py | 10 +++--- 4 files changed, 64 insertions(+), 40 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index 19a5746cef..0cd61ce1c9 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -39,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 @@ -265,8 +267,8 @@ def _cursor(self) -> Iterator[Any]: 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 are - no-ops there. A dead connection can't be reused, so (when we own it) we + 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 @@ -278,17 +280,16 @@ 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 and - # cannot fail; there a server termination surfaces as - # OperationalError (is_connection_dead) and sets conn.closed, which - # the discard below checks too. + # On an owned autocommit connection rollback() sends nothing, so it + # proves nothing there: a bare DatabaseError (no SQLSTATE subclass, + # so not a statement error) is treated as dead directly instead. # NOTE: ``send()``'s retry deliberately catches only # ``_CONN_DEAD_ERRORS``; a server death that surfaces as a *bare* - # ``psycopg2.DatabaseError`` is intentionally NOT retried by - # ``send()``. On an owned (autocommit) connection it is dropped here - # only if ``conn.closed`` is set; otherwise the next call reuses it - # and, if it is dead, fails with a classified error and is dropped. - conn_dead = is_connection_dead(exc) + # ``psycopg2.DatabaseError`` is still dropped here (above, or via + # the failed-rollback branch on an injected connection) 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: @@ -484,11 +485,12 @@ def read_any( 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 statement commits. A - reader racing that statement skips the locked rows, so if each queue's - only ready row is locked it gets ``[]`` and takes the consumer's empty- - poll back-off (starting at ``poll_interval``) instead of a message. This - only happens while several consumers contend for the last rows; idle - polls lock nothing. + 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 diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 65c2284f8f..31066f4dac 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -42,7 +42,7 @@ from ..barrier import callback_recovery_identity from ..fairness import FAIRNESS_HEADER_NAME from .client import PgQueueClient -from .connection import CONN_DEAD_ERRORS, is_connection_dead +from .connection import CONN_DEAD_ERRORS from .liveness import LivenessServer as _BaseLivenessServer from .result_backend import PgResultBackend from .task_payload import to_payload @@ -385,11 +385,9 @@ def poll_once(self) -> int: A failed claim is logged and returns 0, so run() backs off. A message whose handling fails is logged and skipped; the rest of the batch still runs, so a deterministic failure cannot strand its batch-mates (each - redelivery costs them a ``read_ct`` attempt toward the poison cap). The - exception: a connection-dead failure stops the batch — the remaining - messages could not be acked either, so running them would only run them - twice; they stay claimed until their lease lapses and the reaper - re-arms them. Returns the number handled successfully. + redelivery costs them a ``read_ct`` attempt toward the poison cap), and + the client reconnects for their acks if the failure was a dropped + connection. Returns the number handled successfully. """ # HEALTH_STALE is sized for one batch (see the constant). self._last_poll_monotonic = time.monotonic() @@ -401,18 +399,10 @@ def poll_once(self) -> int: ) return 0 handled = 0 - for i, message in enumerate(messages): + for message in messages: try: self._handle(message) - except Exception as exc: - if is_connection_dead(exc): - logger.exception( - "PG-queue consumer: connection lost handling msg_id=%s; " - "left claimed until their lease lapses: %s", - message.msg_id, - [m.msg_id for m in messages[i + 1 :]], - ) - return handled + except Exception: logger.exception( "PG-queue consumer: handling msg_id=%s failed; continuing " "with the rest of the batch", diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index ea8884a5fb..76d856e9cd 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -216,6 +216,37 @@ 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 (no SQLSTATE subclass) is treated as dead. + 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 (SQLSTATE subclass) leaves the connection usable. + 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). @@ -1071,9 +1102,10 @@ def test_owned_connection_claims_without_a_transaction( def test_owned_connection_reconnects_after_backend_termination( self, pg_conn, queue_names, monkeypatch ): - # On an owned autocommit connection rollback() sends nothing, so recovery - # rests on the termination surfacing as a connection-dead error: the - # next call raises, drops the handle, and the call after reconnects. + # 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_"), diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index bcb5243398..cdae88dc31 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -953,9 +953,9 @@ def _handle(message): assert consumer.poll_once() == 2 # msgs 1 and 3 handled assert handled == [1, 3] - def test_connection_loss_stops_the_batch(self): - # Once the connection is gone the rest could not be acked either, so - # running them would only run them twice. + def test_connection_error_does_not_stop_the_batch(self): + # The client drops a dead handle and reconnects for the next ack, so a + # connection error is no reason to strand the rest of the batch. client = MagicMock() client.read_any.return_value = [ _msg(1, _ok_payload(1)), @@ -973,8 +973,8 @@ def _handle(message): handled.append(message.msg_id) consumer._handle = _handle - assert consumer.poll_once() == 1 - assert handled == [1] + assert consumer.poll_once() == 2 + assert handled == [1, 3] def test_rotation_follows_the_last_message_of_a_batch(self): client = MagicMock() From 233e66b27599a921a32bf569940dd050f3164f15 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 05:48:37 +0000 Subject: [PATCH 6/9] UN-4176 Address round-4 self-review: keep main's batch-stop semantics A handler failure stops the batch again, as each queue's batch did on main: rounds 2-4 showed both alternatives cost something (stopping burns a read_ct attempt for the unstarted tail; continuing runs work whose ack cannot land during a DB outage, so it runs twice). Matching main keeps this PR to its scope. Still counts the messages handled and logs the msg_ids left unstarted; the trade-off is written in poll_once. Reword the bare-DatabaseError comment as the heuristic it is. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/queue_backend/pg_queue/client.py | 6 ++-- workers/queue_backend/pg_queue/consumer.py | 29 +++++++++--------- workers/tests/test_pg_queue_client.py | 4 +-- workers/tests/test_pg_queue_consumer.py | 34 ++++------------------ 4 files changed, 27 insertions(+), 46 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index 0cd61ce1c9..8d51a7c3b0 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -281,8 +281,10 @@ def _cursor(self) -> Iterator[Any]: # 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 - # proves nothing there: a bare DatabaseError (no SQLSTATE subclass, - # so not a statement error) is treated as dead directly instead. + # proves nothing 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 dropped here (above, or via diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 31066f4dac..22878573b0 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -382,12 +382,16 @@ def poll_once(self) -> int: 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 message - whose handling fails is logged and skipped; the rest of the batch still - runs, so a deterministic failure cannot strand its batch-mates (each - redelivery costs them a ``read_ct`` attempt toward the poison cap), and - the client reconnects for their acks if the failure was a dropped - connection. Returns the number handled successfully. + A failed claim is logged and returns 0, so run() backs off. A failure + while handling stops the batch, as each queue's batch did before this + loop claimed across queues: 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. """ # HEALTH_STALE is sized for one batch (see the constant). self._last_poll_monotonic = time.monotonic() @@ -398,19 +402,18 @@ def poll_once(self) -> int: "PG-queue consumer: claim failed for queues %r", self.queue_names ) return 0 - handled = 0 - for message in messages: + for handled, message in enumerate(messages): try: self._handle(message) except Exception: logger.exception( - "PG-queue consumer: handling msg_id=%s failed; continuing " - "with the rest of the batch", + "PG-queue consumer: handling msg_id=%s failed; left claimed " + "until their lease lapses: %s", message.msg_id, + [m.msg_id for m in messages[handled + 1 :]], ) - continue - handled += 1 - return handled + return handled + return len(messages) def _claim(self) -> list[QueueMessage]: """One claim for this cycle, advancing the multi-queue rotation.""" diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index 76d856e9cd..b40254c8a2 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -218,7 +218,7 @@ def test_owned_conn_recovered_when_rollback_fails(self, monkeypatch): 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 (no SQLSTATE subclass) is treated as dead. + # dead; a bare DatabaseError is treated as dead (a heuristic). client, conn, factory = self._owned_client( monkeypatch, execute_raises=psycopg2.DatabaseError("server gone") ) @@ -228,7 +228,7 @@ def test_owned_conn_dropped_on_bare_database_error(self, monkeypatch): assert client._conn is None def test_owned_conn_kept_on_statement_error(self, monkeypatch): - # A statement error (SQLSTATE subclass) leaves the connection usable. + # A statement error mapped to a SQLSTATE subclass keeps the connection. client, conn, factory = self._owned_client( monkeypatch, execute_raises=psycopg2.IntegrityError("dup") ) diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index cdae88dc31..d38a8d81cc 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -14,7 +14,6 @@ import time from unittest.mock import MagicMock, patch -import psycopg2 import pytest from celery import shared_task from queue_backend.fairness import FAIRNESS_HEADER_NAME @@ -930,9 +929,9 @@ def test_next_claim_starts_past_the_queue_served(self): 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 - def test_handler_failure_skips_only_the_failed_message(self): - # A deterministic failure must not strand its batch-mates: they would - # otherwise pay a read_ct attempt on every redelivery. + def test_handler_failure_stops_the_batch_and_counts_the_work_done(self): + # 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. client = MagicMock() client.read_any.return_value = [ _msg(1, _ok_payload(1)), @@ -950,31 +949,8 @@ def _handle(message): handled.append(message.msg_id) consumer._handle = _handle - assert consumer.poll_once() == 2 # msgs 1 and 3 handled - assert handled == [1, 3] - - def test_connection_error_does_not_stop_the_batch(self): - # The client drops a dead handle and reconnects for the next ack, so a - # connection error is no reason to strand the rest of the batch. - client = MagicMock() - client.read_any.return_value = [ - _msg(1, _ok_payload(1)), - _msg(2, _ok_payload(2)), - _msg(3, _ok_payload(3)), - ] - 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 == 2: - raise psycopg2.OperationalError("server closed the connection") - handled.append(message.msg_id) - - consumer._handle = _handle - assert consumer.poll_once() == 2 - assert handled == [1, 3] + assert consumer.poll_once() == 1 # msg 1 counted, so run() won't back off + assert handled == [1] # msg 3 left claimed, not started def test_rotation_follows_the_last_message_of_a_batch(self): client = MagicMock() From 8c78d0fee76bbb77e4f1c99882b773b1645a4473 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 06:03:16 +0000 Subject: [PATCH 7/9] UN-4176 Address round-5 self-review: non-object payloads, failure log - _handle drops and acks a payload that is not a JSON object instead of raising on payload.get(); with the batch stopping on a handler failure, such a row would otherwise stop its batch on every delivery. - The handling-failure log names the queue and lists unstarted msg_ids only when there are any. - Fix the _cursor NOTE: an injected connection is never dropped. - Tests: the failure position parametrised over the batch (return value, work done, log text, rotation), and the non-object drop. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/queue_backend/pg_queue/client.py | 7 ++-- workers/queue_backend/pg_queue/consumer.py | 19 +++++++-- workers/tests/test_pg_queue_consumer.py | 46 ++++++++++++++++++---- 3 files changed, 57 insertions(+), 15 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index 8d51a7c3b0..01836a1f8d 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -287,10 +287,9 @@ def _cursor(self) -> Iterator[Any]: # needless reconnect. # NOTE: ``send()``'s retry deliberately catches only # ``_CONN_DEAD_ERRORS``; a server death that surfaces as a *bare* - # ``psycopg2.DatabaseError`` is still dropped here (above, or via - # the failed-rollback branch on an injected connection) but is - # intentionally NOT retried by ``send()`` — it's left to the next - # call's reconnect. + # ``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() diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 22878573b0..26d5720a10 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -406,11 +406,13 @@ def poll_once(self) -> int: try: self._handle(message) except Exception: + unstarted = [m.msg_id for m in messages[handled + 1 :]] logger.exception( - "PG-queue consumer: handling msg_id=%s failed; left claimed " - "until their lease lapses: %s", + "PG-queue consumer: handling msg_id=%s (queue=%r) failed; it " + "redelivers when its lease lapses%s", message.msg_id, - [m.msg_id for m in messages[handled + 1 :]], + message.queue_name, + f", as do the unstarted msg_ids {unstarted}" if unstarted else "", ) return handled return len(messages) @@ -590,6 +592,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_consumer.py b/workers/tests/test_pg_queue_consumer.py index d38a8d81cc..39e3c45b52 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -929,14 +929,21 @@ def test_next_claim_starts_past_the_queue_served(self): 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 - def test_handler_failure_stops_the_batch_and_counts_the_work_done(self): + @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. + # 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)), - _msg(2, _ok_payload(2)), - _msg(3, _ok_payload(3)), + _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 @@ -944,13 +951,36 @@ def test_handler_failure_stops_the_batch_and_counts_the_work_done(self): handled: list[int] = [] def _handle(message): - if message.msg_id == 2: + if message.msg_id == failing: raise RuntimeError("boom") handled.append(message.msg_id) consumer._handle = _handle - assert consumer.poll_once() == 1 # msg 1 counted, so run() won't back off - assert handled == [1] # msg 3 left claimed, not started + 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() From 2b84b049ab6b8f9e7595cba31641ec03965b64a2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 09:19:33 +0000 Subject: [PATCH 8/9] UN-4176 Address review: required queue_name, read_any dedupe, simpler rotation - QueueMessage.queue_name is required (both constructors always set it), so the consumer's unknown-name fallback and its test are gone; message is annotated Any, since a raw writer can store any JSON. - read_any de-duplicates queue names: a repeated queue's second probe returned the same rows and crowded another queue out of the batch. - Rotation is done in Python by binding the queue list already rotated; the SQL orders by ord and takes 4 parameters instead of 10. - Comment accuracy: unlocked when the claim commits; rollback() can't detect a dead server connection; drop the "before this loop" wording. - Tests: end-to-end multi-queue claim, execute and ack on real Postgres; duplicate queue names still fill the batch; pg_terminate_backend waits for the backend to exit. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/queue_backend/pg_queue/client.py | 41 ++++++++++++---------- workers/queue_backend/pg_queue/consumer.py | 19 +++------- workers/tests/test_pg_queue_client.py | 23 ++++++++---- workers/tests/test_pg_queue_consumer.py | 33 ++++++++++------- 4 files changed, 65 insertions(+), 51 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index 01836a1f8d..b1e15c2549 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -143,9 +143,9 @@ 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 + The caller binds the queue list already rotated, so ``WITH ORDINALITY`` + numbers the queues in serving order and ``ord`` alone decides which queue is + served first (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. """ @@ -163,7 +163,7 @@ def _dequeue_any_sql() -> str: FOR UPDATE SKIP LOCKED LIMIT %s ) c - ORDER BY (q.ord - 1 - %s + %s) %% %s, c.priority DESC, c.msg_id + ORDER BY q.ord, c.priority DESC, c.msg_id LIMIT %s ), claimed AS ( UPDATE {msg} q @@ -176,7 +176,7 @@ def _dequeue_any_sql() -> str: ) SELECT msg_id, message, read_ct, queue_name FROM claimed - ORDER BY (ord - 1 - %s + %s) %% %s, priority DESC, msg_id + ORDER BY ord, priority DESC, msg_id """ @@ -184,21 +184,22 @@ def _dequeue_any_sql() -> str: class QueueMessage: """A claimed queue message. - ``message`` is the already-decoded JSONB payload (psycopg2 parses - ``jsonb`` to a Python ``dict``). ``frozen`` freezes the binding only — - the ``dict`` itself is mutable, so treat the payload as read-only by + ``message`` is the already-decoded JSONB payload: a ``dict`` for every + message ``send()`` writes, but any JSON value if a raw writer inserted it + (the consumer drops non-objects as malformed). ``frozen`` freezes the + binding only — the payload itself is mutable, so treat it as read-only by convention. ``read_ct`` is the post-claim delivery count — always ``>= 1`` 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; the multi-queue - consumer uses it to rotate past the queue it just served. + ``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] + message: Any read_ct: int - queue_name: str = "" + queue_name: str # The whole enqueue contract (columns + their defaults) in one place. `send()` @@ -281,7 +282,7 @@ def _cursor(self) -> Iterator[Any]: # 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 - # proves nothing there. As a heuristic, a bare DatabaseError is + # 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. @@ -485,7 +486,8 @@ def read_any( 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 statement commits. A + 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 @@ -507,14 +509,17 @@ def read_any( raise ValueError(f"qty must be positive, got {qty}") if not queue_names: raise ValueError("queue_names must be non-empty") - n = len(queue_names) - rotate = (start % n, n, n) + # A repeated name would be probed twice, and both probes can pick the + # same row (SKIP LOCKED doesn't skip this statement's own locks), which + # would return fewer than qty messages. + queues = list(dict.fromkeys(queue_names)) + k = start % len(queues) with self._cursor() as cur: # Param order matches the %s positions in _dequeue_any_sql(): - # queues, per-queue LIMIT, rotation, LIMIT, vt, rotation again. + # queues in serving order, per-queue LIMIT, LIMIT, vt. cur.execute( _dequeue_any_sql(), - (list(queue_names), qty, *rotate, qty, vt_seconds, *rotate), + (queues[k:] + queues[:k], qty, qty, vt_seconds), ) rows = cur.fetchall() return [ diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 26d5720a10..50d915979b 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -383,10 +383,9 @@ def poll_once(self) -> int: 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, as each queue's batch did before this - loop claimed across queues: 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 + 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. @@ -431,19 +430,9 @@ def _claim(self) -> list[QueueMessage]: ) # One past the queue served last, so the queues after it go first next # time; an empty cycle just moves on by one. - if not messages: - served = self._next_queue - elif messages[-1].queue_name in self.queue_names: + if messages: served = self.queue_names.index(messages[-1].queue_name) else: - # read_any always fills queue_name; reaching here means a client that - # doesn't, which would quietly bring back unequal turns. - logger.warning( - "PG-queue consumer: claimed msg_id=%s has unknown queue_name %r; " - "rotating by one", - messages[-1].msg_id, - messages[-1].queue_name, - ) served = self._next_queue self._next_queue = (served + 1) % len(self.queue_names) return messages diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index b40254c8a2..dfeba80523 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -987,6 +987,19 @@ def test_batch_orders_by_priority_within_a_queue(self, pg_conn, queue_names): 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"}) @@ -1018,17 +1031,13 @@ def test_held_claim_is_skipped_not_waited_on(self, pg_conn, queue_names): # 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), - ) + cur.execute(_dequeue_any_sql(), (list(queue_names), 1, 1, 30)) a_ids = [r[0] for r in cur.fetchall()] with conn_b.cursor() as cur: cur.execute("SET lock_timeout = '2s'") @@ -1114,7 +1123,9 @@ def test_owned_connection_reconnects_after_backend_termination( owned.read_any(queue_names, vt_seconds=30) old_pid = owned.conn.get_backend_pid() with pg_conn.cursor() as cur: - cur.execute("SELECT pg_terminate_backend(%s)", (old_pid,)) + # 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) diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 39e3c45b52..93f22892d2 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -51,7 +51,7 @@ def _clear_calls(): _calls.clear() -def _msg(msg_id, payload, *, read_ct=1, queue_name=""): +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 ) @@ -871,6 +871,26 @@ 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. @@ -1004,17 +1024,6 @@ def test_rotation_follows_the_last_message_of_a_batch(self): # 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_unknown_queue_name_rotates_by_one_and_warns(self, caplog): - client = MagicMock() - client.read_any.side_effect = [[], [_msg(1, _ok_payload(1), queue_name="??")]] - consumer = PgQueueConsumer(["qa", "qb", "qc"], client=client) - consumer._handle = lambda m: None - consumer.poll_once() # empty: start 0 -> 1 - with caplog.at_level(logging.WARNING): - consumer.poll_once() # start 1, unknown name -> 2 - assert consumer._next_queue == 2 - assert "unknown queue_name" in caplog.text - 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. From 364d5a8cda47d6237c1ed5da3935730568845f00 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 16 Oct 2026 09:27:32 +0000 Subject: [PATCH 9/9] UN-4176 Revert out-of-scope review suggestions Keep the PR to its scope and to required changes only: - Restore the SQL-side rotation (cosmetic simplification, not a defect; the de-duplication of queue names stays). - Restore QueueMessage.message's dict annotation and docstring, a pre-existing public field this PR does not need to change. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017Ymcknu7CYVmAt7QoRTTax --- workers/queue_backend/pg_queue/client.py | 39 ++++++++++++------------ workers/tests/test_pg_queue_client.py | 6 +++- 2 files changed, 24 insertions(+), 21 deletions(-) diff --git a/workers/queue_backend/pg_queue/client.py b/workers/queue_backend/pg_queue/client.py index b1e15c2549..cf799ecfce 100644 --- a/workers/queue_backend/pg_queue/client.py +++ b/workers/queue_backend/pg_queue/client.py @@ -143,9 +143,9 @@ def _dequeue_sql() -> str: def _dequeue_any_sql() -> str: """Multi-queue variant of :func:`_dequeue_sql` (see :meth:`PgQueueClient.read_any`). - The caller binds the queue list already rotated, so ``WITH ORDINALITY`` - numbers the queues in serving order and ``ord`` alone decides which queue is - served first (every queue is still probed on every call). Within a queue the + ``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. """ @@ -163,7 +163,7 @@ def _dequeue_any_sql() -> str: FOR UPDATE SKIP LOCKED LIMIT %s ) c - ORDER BY q.ord, c.priority DESC, c.msg_id + ORDER BY (q.ord - 1 - %s + %s) %% %s, c.priority DESC, c.msg_id LIMIT %s ), claimed AS ( UPDATE {msg} q @@ -176,7 +176,7 @@ def _dequeue_any_sql() -> str: ) SELECT msg_id, message, read_ct, queue_name FROM claimed - ORDER BY ord, priority DESC, msg_id + ORDER BY (ord - 1 - %s + %s) %% %s, priority DESC, msg_id """ @@ -184,10 +184,9 @@ def _dequeue_any_sql() -> str: class QueueMessage: """A claimed queue message. - ``message`` is the already-decoded JSONB payload: a ``dict`` for every - message ``send()`` writes, but any JSON value if a raw writer inserted it - (the consumer drops non-objects as malformed). ``frozen`` freezes the - binding only — the payload itself is mutable, so treat it as read-only by + ``message`` is the already-decoded JSONB payload (psycopg2 parses + ``jsonb`` to a Python ``dict``). ``frozen`` freezes the binding only — + the ``dict`` itself is mutable, so treat the payload as read-only by convention. ``read_ct`` is the post-claim delivery count — always ``>= 1`` from the dequeue (``read_ct = read_ct + 1 RETURNING``), so consumers can cap redelivery of a poison message. No default: a ``0`` ("never claimed") @@ -197,7 +196,7 @@ class QueueMessage: """ msg_id: int - message: Any + message: dict[str, Any] read_ct: int queue_name: str @@ -282,10 +281,10 @@ def _cursor(self) -> Iterator[Any]: # 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. + # 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 dropped here when we own the @@ -509,17 +508,17 @@ def read_any( 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 both probes can pick the - # same row (SKIP LOCKED doesn't skip this statement's own locks), which - # would return fewer than qty messages. + # 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)) - k = start % len(queues) + n = len(queues) + rotate = (start % n, n, n) with self._cursor() as cur: # Param order matches the %s positions in _dequeue_any_sql(): - # queues in serving order, per-queue LIMIT, LIMIT, vt. + # queues, per-queue LIMIT, rotation, LIMIT, vt, rotation again. cur.execute( _dequeue_any_sql(), - (queues[k:] + queues[:k], qty, qty, vt_seconds), + (queues, qty, *rotate, qty, vt_seconds, *rotate), ) rows = cur.fetchall() return [ diff --git a/workers/tests/test_pg_queue_client.py b/workers/tests/test_pg_queue_client.py index dfeba80523..fa5f8da603 100644 --- a/workers/tests/test_pg_queue_client.py +++ b/workers/tests/test_pg_queue_client.py @@ -1031,13 +1031,17 @@ def test_held_claim_is_skipped_not_waited_on(self, pg_conn, queue_names): # 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, 1, 30)) + 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'")