diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md index fa80e87fbc..67ca71633c 100644 --- a/workers/queue_backend/pg_queue/README.md +++ b/workers/queue_backend/pg_queue/README.md @@ -27,7 +27,8 @@ integers. | Env suffix | One-line | Default | |---|---|---| | `CONCURRENCY` | Prefork consumer children per pod (1 = plain single process) | `1` | -| `VT_SECONDS` | **Visibility timeout** — how long a claimed message stays hidden before it can redeliver | `30` (chart: `9060` for file-processing ≈ 2.5h) | +| `VT_SECONDS` | The **drain / max-runtime bound** (drives `SHUTDOWN_GRACE`, `HEALTH_STALE`, the chart guards) — the *claim* window is `LEASE_SECONDS`, not this | `30` (chart: `9060` for file-processing ≈ 2.5h) | +| `LEASE_SECONDS` | **Renewable claim window** — the effective claim window is `min(LEASE, VT)`, renewed every ~that/3 while the task runs; a dead worker's claim expires in ~that → fast redelivery. With the defaults (`LEASE=120`, `VT=30`) it clamps to 30; the chart raises VT so the full 120 applies | `120` | | `SHUTDOWN_GRACE_SECONDS` | Graceful-drain budget (shared across all children) on SIGTERM before SIGKILL | `= VT` (floored at `30`) | | `QUEUE` | Queue name(s) this consumer polls (comma-separated) | `default` | | `BATCH` | Messages claimed per poll | `1` | @@ -64,14 +65,28 @@ integers. ## 2. Concepts — glossary **Visibility timeout (VT)** — when a consumer *claims* a message it becomes invisible -to other consumers for `VT` seconds. If the worker finishes and deletes (acks) it, it's -gone; if the worker dies without acking, the VT expires and the message **redelivers**. -VT is the queue's only "worker died" signal (there's no live broker connection like -RabbitMQ), so recovery latency ≈ VT. +to other consumers until its `vt` passes. If the worker finishes and deletes (acks) it, +it's gone; if the worker dies without acking, the `vt` expires and the message +**redelivers**. There's no live broker connection like RabbitMQ, so the `vt` is the +queue's only "worker died" signal — recovery latency ≈ how far out the `vt` is. + +**Renewable lease (`LEASE_SECONDS`, UN-3695)** — rather than claim for the full `VT` +(up to 2.5h), the consumer claims for a **short** `LEASE` and a background thread +**renews** it (`set_vt`) every ~`LEASE/3` while the task runs. A live-but-slow task +keeps its claim; a **dead** worker's renewal stops, so its `vt` expires in ~`LEASE` → +redelivery in **minutes, not hours**. `VT_SECONDS` is retained as the *drain / +max-runtime bound* (grace, health-stale), and the lease is clamped to it. Note the +lease is **not** a hard cap on runtime — a live-but-hung task keeps renewing forever; +the backstop for that is the liveness probe restarting the pod (process death stops +renewal). The renewal owns its own DB connection (closed on exit) and is best-effort: +a connection death retries within the `~2×` slack the `LEASE/3` interval leaves before +expiry, and escalates to an ERROR log once it keeps failing past `LEASE` (the lease is +then genuinely lost and the message may double-run). Because renewal covers only the +in-flight message, `BATCH_SIZE` is forced to 1 whenever the lease is the short window. **Claim** — an atomic `SELECT … FOR UPDATE SKIP LOCKED` that hides up to `BATCH` -ready rows for `VT` seconds and hands them to one consumer. `SKIP LOCKED` distributes -work across children and replicas without contention. +ready rows for the claim window (`min(LEASE, VT)`) and hands them to one consumer. +`SKIP LOCKED` distributes work across children and replicas without contention. **Redelivery / at-least-once** — a message can be delivered more than once (VT expiry after a crash, or a poison re-park). Handlers must tolerate re-execution. diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index b2055e1520..b90ca57a0e 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -20,12 +20,14 @@ from __future__ import annotations +import contextlib import json import logging import os import signal +import threading import time -from collections.abc import Callable +from collections.abc import Callable, Iterator from enum import Enum from typing import TYPE_CHECKING, TypeVar @@ -36,6 +38,7 @@ from ..barrier import callback_recovery_identity from ..fairness import FAIRNESS_HEADER_NAME from .client import PgQueueClient +from .connection import CONN_DEAD_ERRORS from .liveness import LivenessServer as _BaseLivenessServer from .result_backend import PgResultBackend from .task_payload import to_payload @@ -57,6 +60,15 @@ # raise it, keep vt_seconds > batch_size x worst-case task duration. _DEFAULT_BATCH = 1 _DEFAULT_VT_SECONDS = 30 +# Renewable-lease claim window (UN-3695). A claimed message is hidden for only +# this long; a background thread renews it (~every LEASE/3) while the task runs, so a +# live-but-slow task stays claimed but a DEAD worker's claim expires in ~LEASE → +# redelivery in minutes instead of the full VT. VT_SECONDS is now the drain / +# max-runtime bound (grace, health-stale), NOT the claim window. +_DEFAULT_LEASE_SECONDS = 120 +# Bounded join so a wedged renewal thread can't block the ack — it's a daemon and +# dies with the process regardless. +_LEASE_JOIN_TIMEOUT_SECONDS = 10.0 _DEFAULT_POLL_INTERVAL = 0.1 _DEFAULT_BACKOFF_MAX = 2.0 # A task claimed more than this many times keeps failing — drop it (poison) @@ -201,6 +213,7 @@ def __init__( api_client: InternalAPIClient | None = None, batch_size: int = _DEFAULT_BATCH, vt_seconds: int = _DEFAULT_VT_SECONDS, + lease_seconds: int = _DEFAULT_LEASE_SECONDS, poll_interval: float = _DEFAULT_POLL_INTERVAL, backoff_max: float = _DEFAULT_BACKOFF_MAX, max_attempts: int = _DEFAULT_MAX_ATTEMPTS, @@ -212,6 +225,7 @@ def __init__( for name, value in ( ("batch_size", batch_size), ("vt_seconds", vt_seconds), + ("lease_seconds", lease_seconds), ("poll_interval", poll_interval), ("backoff_max", backoff_max), ("max_attempts", max_attempts), @@ -245,6 +259,33 @@ def __init__( self._api_client = api_client self.batch_size = batch_size self.vt_seconds = vt_seconds + # Renewable lease (UN-3695): the claim is taken for a short LEASE and renewed + # every ~lease/3 while a task runs, so a dead worker's claim expires in ~LEASE + # (fast redelivery) but a live one is never redelivered. VT_SECONDS is the + # drain / max-runtime bound; a lease longer than it is meaningless, so clamp — + # loudly, since a lease > VT is a misconfig, not a silent no-op. + self.lease_seconds = min(lease_seconds, vt_seconds) + if lease_seconds > vt_seconds: + logger.warning( + "WORKER_PG_QUEUE_CONSUMER_LEASE_SECONDS=%s > VT_SECONDS=%s; clamped the " + "claim window to %ss (the lease can't outlast the max-runtime bound)", + lease_seconds, + vt_seconds, + self.lease_seconds, + ) + self._lease_renew_interval = max(1, self.lease_seconds // 3) + # The renewal only extends the IN-FLIGHT message; a batch tail sits claimed but + # unrenewed, so with batch>1 a slow head lets the tail's lease lapse and another + # consumer double-runs it. Force serial claims whenever the lease is the short + # (renewed) claim window; batch>1 stays available when lease == vt (no renewal). + if self.batch_size > 1 and self.lease_seconds < vt_seconds: + logger.warning( + "WORKER_PG_QUEUE_CONSUMER_BATCH_SIZE=%s forced to 1: the renewable lease " + "only covers the in-flight message, so a batch tail could lapse and " + "double-run (UN-3695)", + self.batch_size, + ) + self.batch_size = 1 self.poll_interval = poll_interval self.backoff_max = backoff_max self.max_attempts = max_attempts @@ -276,7 +317,7 @@ def poll_once(self) -> int: for queue_name in self.queue_names: try: messages = self._client.read( - queue_name, vt_seconds=self.vt_seconds, qty=self.batch_size + queue_name, vt_seconds=self.lease_seconds, qty=self.batch_size ) for message in messages: self._handle(message) @@ -289,6 +330,104 @@ def poll_once(self) -> int: ) return total + @contextlib.contextmanager + def _lease_renewal(self, msg_id: int) -> Iterator[None]: + """Keep ``msg_id``'s claim alive while a task runs (UN-3695). + + Starts a daemon thread that renews the short lease (``set_vt``) every + ``lease/3`` — so a live-but-slow task stays claimed — then signals stop and + joins on exit, before the caller acks. If the worker DIES, the thread dies + with it → the lease expires in ~``lease_seconds`` → the message redelivers in + minutes instead of the full VT. The join is bounded by + ``_LEASE_JOIN_TIMEOUT_SECONDS``; a thread still alive after that (wedged in a + stalled ``set_vt``) is logged and abandoned — it owns its own connection (see + the loop), so it can't corrupt the ack path, and a late ``set_vt`` on the + about-to-be-acked row is a benign no-op. + """ + stop = threading.Event() + thread = threading.Thread( + target=self._renew_lease_loop, + args=(msg_id, stop), + name=f"pg-lease-{msg_id}", + daemon=True, + ) + thread.start() + try: + yield + finally: + stop.set() + thread.join(timeout=_LEASE_JOIN_TIMEOUT_SECONDS) + if thread.is_alive(): + logger.error( + "PG-queue consumer: lease-renewal thread for msg_id=%s did not " + "stop within %ss — proceeding to ack; its connection is abandoned " + "until the process exits", + msg_id, + _LEASE_JOIN_TIMEOUT_SECONDS, + ) + + def _make_renew_client(self) -> PgQueueClient: + """Factory for the renewal thread's own connection (patched in tests).""" + return PgQueueClient() + + def _renew_lease_loop(self, msg_id: int, stop: threading.Event) -> None: + """Renew ``msg_id``'s claim every ``_lease_renew_interval`` until ``stop``. + + Waits *then* renews (the initial claim already set the lease), so a task + shorter than the interval never renews and opens no second connection. Owns its + connection start-to-finish (closed on exit) — never shared with the main + thread's claim/ack client. Best-effort: a connection death is retried next tick + (the interval leaves ~2x slack before expiry), but if it keeps failing past + ``lease_seconds`` the lease is genuinely lost (the row can be reclaimed) and it + escalates to ERROR. A ``False`` return means the row was already deleted — and + since the ack runs only after this thread is joined, that means another consumer + reclaimed + acked it, i.e. this task is double-running: logged, then stop. A + non-connection error is left to propagate (fail loud, not swallowed forever). + """ + client: PgQueueClient | None = None + last_ok = time.monotonic() + try: + while not stop.wait(self._lease_renew_interval): + try: + # Built on the FIRST renewal (not on thread start), so a task + # shorter than the interval opens no second connection. + if client is None: + client = self._make_renew_client() + if not client.set_vt(msg_id, self.lease_seconds): + logger.warning( + "PG-queue consumer: lease for msg_id=%s lost (row already " + "gone) — it was reclaimed and this task may double-run", + msg_id, + ) + return + last_ok = time.monotonic() + except CONN_DEAD_ERRORS: + down_for = time.monotonic() - last_ok + if down_for >= self.lease_seconds: + logger.exception( + "PG-queue consumer: lease renewal for msg_id=%s failing for " + "%.0fs (>= lease %ss) — the lease has likely expired and " + "this task may double-run", + msg_id, + down_for, + self.lease_seconds, + ) + else: + logger.warning( + "PG-queue consumer: lease renewal for msg_id=%s failed " + "(retry in %ss; %.0fs of %ss slack used) — a dead " + "connection self-heals on the next tick", + msg_id, + self._lease_renew_interval, + down_for, + self.lease_seconds, + exc_info=True, + ) + finally: + if client is not None: + with contextlib.suppress(Exception): + client.close() + def _handle(self, message: QueueMessage) -> None: payload = message.message task_name = payload.get("task_name") @@ -343,12 +482,15 @@ def _handle(self, message: QueueMessage) -> None: # header so a PG-routed run mirrors the Celery dispatch path. fairness = payload.get("fairness") headers = {FAIRNESS_HEADER_NAME: fairness} if fairness else None - eager = task.apply( - args=payload.get("args") or [], - kwargs=payload.get("kwargs") or {}, - headers=headers, - throw=True, - ) + # Renew the short lease while the (possibly long) task runs, so a dead + # worker's claim expires fast but a live one is never redelivered (UN-3695). + with self._lease_renewal(message.msg_id): + eager = task.apply( + args=payload.get("args") or [], + kwargs=payload.get("kwargs") or {}, + headers=headers, + throw=True, + ) except Exception as exc: if reply_key or on_success or on_error: # Request-reply / async-callback dispatch: surface the failure on @@ -826,10 +968,11 @@ def run(self, *, install_signals: bool = True, require_tasks: bool = True) -> No name for name in self._app.tasks if not name.startswith("celery.") ) logger.info( - "PG-queue consumer started (queues=%r, batch=%s, vt=%ss) — " + "PG-queue consumer started (queues=%r, batch=%s, lease=%ss, vt=%ss) — " "%d application task(s) registered: %s", self.queue_names, self.batch_size, + self.lease_seconds, self.vt_seconds, len(app_tasks), ", ".join(app_tasks) or "(none)", @@ -917,6 +1060,7 @@ def build_consumer_from_env() -> PgQueueConsumer: queue_names=consumer_env("QUEUE", [_DEFAULT_QUEUE], _parse_queue_list), batch_size=consumer_env("BATCH", _DEFAULT_BATCH, int), vt_seconds=consumer_env("VT_SECONDS", _DEFAULT_VT_SECONDS, int), + lease_seconds=consumer_env("LEASE_SECONDS", _DEFAULT_LEASE_SECONDS, int), poll_interval=consumer_env("POLL_INTERVAL", _DEFAULT_POLL_INTERVAL, float), backoff_max=consumer_env("BACKOFF_MAX", _DEFAULT_BACKOFF_MAX, float), max_attempts=consumer_env("MAX_ATTEMPTS", _DEFAULT_MAX_ATTEMPTS, int), diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index f0a992c873..24b505d8f4 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -10,6 +10,8 @@ import logging import os +import threading +import time from unittest.mock import MagicMock, patch import pytest @@ -17,6 +19,7 @@ from queue_backend.fairness import FAIRNESS_HEADER_NAME from queue_backend.pg_queue import to_payload from queue_backend.pg_queue.client import QueueMessage +from queue_backend.pg_queue.connection import CONN_DEAD_ERRORS from queue_backend.pg_queue.consumer import PgQueueConsumer # Registered test tasks (namespaced). apply() runs their bodies in-process. @@ -1055,3 +1058,175 @@ def test_chain_continuation_records_status_even_if_enqueue_fails(self): self._entered_rb(RB).store_result.assert_called_once_with( "tid", result={}, retention_seconds=self._RET ) + + +class TestLeaseRenewal: + """UN-3695: the claim is taken for the short lease and renewed while the task + runs, so a dead worker's claim expires in ~lease (fast redelivery) instead of the + full VT. VT_SECONDS stays the drain/max-runtime bound. + """ + + # --- construction / config --- + + def test_lease_clamped_to_vt_loudly(self, caplog): + with caplog.at_level(logging.WARNING, logger="queue_backend.pg_queue.consumer"): + c = PgQueueConsumer( + ["q"], client=MagicMock(), vt_seconds=100, lease_seconds=200 + ) + assert c.lease_seconds == 100 # clamped — the lease can't outlast VT + assert "clamped the claim window" in caplog.text # loud, not silent + + def test_lease_and_renew_interval(self): + c = PgQueueConsumer(["q"], client=MagicMock(), vt_seconds=9060, lease_seconds=120) + assert c.lease_seconds == 120 + assert c._lease_renew_interval == 40 # lease // 3 → 2 attempts of slack + + def test_tiny_lease_renew_interval_floored_at_1(self): + c = PgQueueConsumer(["q"], client=MagicMock(), vt_seconds=2, lease_seconds=2) + assert c._lease_renew_interval == 1 # max(1, 2 // 3) + + def test_non_positive_lease_rejected(self): + client = MagicMock() + with pytest.raises(ValueError, match="lease_seconds must be positive"): + PgQueueConsumer(["q"], client=client, lease_seconds=0) + + def test_batch_forced_to_1_when_lease_shorter_than_vt(self, caplog): + with caplog.at_level(logging.WARNING, logger="queue_backend.pg_queue.consumer"): + c = PgQueueConsumer( + ["q"], client=MagicMock(), vt_seconds=9060, lease_seconds=120, batch_size=8 + ) + assert c.batch_size == 1 # renewal only covers the in-flight message + assert "forced to 1" in caplog.text + + def test_batch_kept_when_lease_equals_vt(self): + # lease == vt → no short renewal window → batch>1 is safe, left untouched. + c = PgQueueConsumer( + ["q"], client=MagicMock(), vt_seconds=50, lease_seconds=50, batch_size=4 + ) + assert c.lease_seconds == 50 and c.batch_size == 4 + + def test_env_wires_lease_seconds(self, monkeypatch): + from queue_backend.pg_queue import consumer as mod + + monkeypatch.setenv("WORKER_PG_QUEUE_CONSUMER_LEASE_SECONDS", "77") + monkeypatch.setenv("WORKER_PG_QUEUE_CONSUMER_VT_SECONDS", "9060") + with patch.object(mod, "PgQueueClient"): # no real DB connection + c = mod.build_consumer_from_env() + assert c.lease_seconds == 77 + + def test_poll_claims_with_lease_not_vt(self): + client = MagicMock() + client.read.return_value = [] + c = PgQueueConsumer(["q"], client=client, vt_seconds=9060, lease_seconds=120) + c.poll_once() + client.read.assert_called_with("q", vt_seconds=120, qty=c.batch_size) + + # --- the renewal loop (own connection via the _make_renew_client factory) --- + + def _renewing(self, renew_client, lease_seconds=3): + c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=lease_seconds) + c._make_renew_client = MagicMock(return_value=renew_client) + return c + + def test_renews_until_stopped_then_closes_own_client(self): + rc = MagicMock() + rc.set_vt.return_value = True + c = self._renewing(rc) + stop = MagicMock() + stop.wait.side_effect = [False, False, True] # renew twice, then stop + c._renew_lease_loop(99, stop) + assert rc.set_vt.call_count == 2 + rc.set_vt.assert_called_with(99, 3) # renews with the lease + stop.wait.assert_called_with(c._lease_renew_interval) + rc.close.assert_called_once() # own connection released on exit + + def test_no_client_built_for_sub_interval_task(self): + c = self._renewing(MagicMock()) + stop = MagicMock() + stop.wait.side_effect = [True] # stopped before the first renewal tick + c._renew_lease_loop(99, stop) + c._make_renew_client.assert_not_called() # short task opens no 2nd connection + + def test_row_gone_logs_double_run_and_stops(self, caplog): + rc = MagicMock() + rc.set_vt.return_value = False # reclaimed + acked elsewhere (we ack later) + c = self._renewing(rc) + stop = MagicMock() + stop.wait.side_effect = [False, False, True] + with caplog.at_level(logging.WARNING, logger="queue_backend.pg_queue.consumer"): + c._renew_lease_loop(99, stop) + assert rc.set_vt.call_count == 1 # stopped after the first gone + assert "may double-run" in caplog.text # the lost-lease signal is logged + rc.close.assert_called_once() + + def test_conn_death_retried_then_escalates_past_lease(self, caplog): + rc = MagicMock() + rc.set_vt.side_effect = CONN_DEAD_ERRORS[0]("db down") + c = self._renewing(rc, lease_seconds=3) + stop = MagicMock() + stop.wait.side_effect = [False, False, True] + # last_ok=0; 1st failure at t=0 (< lease → WARNING); 2nd at t=10 (≥ lease → ERROR) + with patch( + "queue_backend.pg_queue.consumer.time.monotonic", + side_effect=[0.0, 0.0, 10.0], + ): + with caplog.at_level( + logging.WARNING, logger="queue_backend.pg_queue.consumer" + ): + c._renew_lease_loop(99, stop) # must NOT raise + assert rc.set_vt.call_count == 2 + assert "self-heals" in caplog.text # within-slack retry + assert "may double-run" in caplog.text # escalated once past the lease + + def test_non_connection_error_propagates(self): + # A programming bug is NOT swallowed as a "self-healing blip" forever. + rc = MagicMock() + rc.set_vt.side_effect = AttributeError("bug") + c = self._renewing(rc) + stop = MagicMock() + stop.wait.side_effect = [False, False, True] + with pytest.raises(AttributeError): + c._renew_lease_loop(99, stop) + rc.close.assert_called_once() # still released via finally + + # --- the context manager + _handle wiring --- + + def test_ctx_runs_loop_and_sets_stop_on_exit(self): + c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) + entered = False + with patch.object(c, "_renew_lease_loop") as loop: + with c._lease_renewal(42): + entered = True # CM body runs while the renewal thread is live + assert entered + loop.assert_called_once() + msg_id, stop = loop.call_args.args + assert msg_id == 42 + assert isinstance(stop, threading.Event) and stop.is_set() # stopped on exit + + def test_ctx_logs_when_thread_wedged_past_join_timeout(self, caplog): + c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) + entered = False + # A loop that ignores stop → join times out → the thread is still alive. + with patch.object(c, "_renew_lease_loop", lambda mid, s: time.sleep(0.3)): + with patch( + "queue_backend.pg_queue.consumer._LEASE_JOIN_TIMEOUT_SECONDS", 0.01 + ): + with caplog.at_level( + logging.ERROR, logger="queue_backend.pg_queue.consumer" + ): + with c._lease_renewal(7): + entered = True # CM body runs even though the thread wedges + assert entered + assert "did not stop within" in caplog.text # wedged thread is not silent + + def test_handle_wraps_task_run_in_lease_renewal(self): + # Integration: a real run goes through _lease_renewal (dropping the `with` + # in _handle would keep the unit tests green — this guards that regression). + client = MagicMock() + client.read.return_value = [_msg(1, _ok_payload(3, 4))] + c = PgQueueConsumer(["q"], client=client) + with patch.object(c, "_lease_renewal", wraps=c._lease_renewal) as lease: + c.poll_once() + lease.assert_called_once_with(1) # the in-flight msg_id was leased + assert _calls == [(3, 4)] # task still ran + client.delete.assert_called_once_with(1) # and acked