From 2157b4a7c9a738e332abcd60dc4e5087b8bd0b3d Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 13:04:13 +0530 Subject: [PATCH 1/5] =?UTF-8?q?UN-3695=20[FIX]=20PG=20consumer=20=E2=80=94?= =?UTF-8?q?=20renewable=20lease=20so=20a=20dead=20worker=20redelivers=20in?= =?UTF-8?q?=20minutes,=20not=20the=20full=20VT?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The claim was taken for the full visibility timeout (VT=9060s / 2.5h for file-processing), so a HARD-killed worker (OOM / node loss) stranded its in-flight message for ~2.5h before the vt expired and it redelivered. PR A fixed GRACEFUL kills (drain on SIGTERM); this covers the hard-kill path. - New LEASE_SECONDS (default 120): the message is claimed for the short lease, and a daemon thread renews it (client.set_vt) every ~lease/3 while the task runs. A live-but-slow task keeps its claim; a dead worker's renewal stops with the process so its lease expires in ~lease → redelivery in minutes. The lease is capped at vt_seconds; the renewal uses its OWN connection (the main thread owns the claim/ack connection) and is best-effort (a transient set_vt failure retries within the ~2x slack the lease/3 interval leaves before expiry; a False return = row already gone → stop). - VT_SECONDS is retained as the drain / max-runtime bound (it drives PR A's grace and the chart's health-stale / guards) — it is no longer the claim window. The claim now uses LEASE_SECONDS. - The renewal thread is started only around the task run, joined before the ack so the two never race the row, and creates its connection only on the first actual renewal (a task shorter than the interval opens no second connection). Paired chart change (unstract-cloud): sets WORKER_PG_QUEUE_CONSUMER_LEASE_SECONDS + a render guard LEASE <= VT. Tests: lease cap / renew-interval / claim-uses-lease / renew-until-stopped / stop-on-row-gone / swallow-transient-error / ctx-sets-stop-on-exit. 78 consumer tests pass. End-to-end dead-worker redelivery timing validated in the integration env. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/queue_backend/pg_queue/README.md | 20 +++-- workers/queue_backend/pg_queue/consumer.py | 94 ++++++++++++++++++++-- workers/tests/test_pg_queue_consumer.py | 63 +++++++++++++++ 3 files changed, 165 insertions(+), 12 deletions(-) diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md index fa80e87fbc..372b11d689 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` | Since PR C, 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** — a claimed message is hidden only this long and renewed every ~LEASE/3 while the task runs; a dead worker's claim expires in ~LEASE → fast redelivery (capped at VT) | `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,10 +65,19 @@ 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 PR C)** — 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 capped at it. The renewal +uses its own DB connection and is best-effort (a transient failure retries within the +`~2×` slack the `LEASE/3` interval leaves before expiry). **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 diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index b2055e1520..b0ee30f814 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -20,10 +20,12 @@ from __future__ import annotations +import contextlib import json import logging import os import signal +import threading import time from collections.abc import Callable from enum import Enum @@ -57,6 +59,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 PR C). 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 +212,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 +224,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 +258,16 @@ def __init__( self._api_client = api_client self.batch_size = batch_size self.vt_seconds = vt_seconds + # Renewable lease (UN-3695 PR C): claim for lease_seconds (short) and renew + # every ~lease/3 while a task runs. Capped at vt_seconds so the lease is + # genuinely the shorter window; setting lease >= vt degrades to the old + # single-long-claim behaviour (renewal fires rarely). + self.lease_seconds = min(lease_seconds, vt_seconds) + self._lease_renew_interval = max(1, self.lease_seconds // 3) + # Second client (own connection) for lease renewals — the main thread owns + # self._client for the read/ack, so a renewal must not share that connection. + # Lazily created on the first task run. + self._renew_client: PgQueueClient | None = None self.poll_interval = poll_interval self.backoff_max = backoff_max self.max_attempts = max_attempts @@ -276,7 +299,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 +312,59 @@ def poll_once(self) -> int: ) return total + @contextlib.contextmanager + def _lease_renewal(self, msg_id: int): + """Keep ``msg_id``'s claim alive while a task runs (UN-3695 PR C). + + The claim was taken for the short ``lease_seconds``; this starts a daemon + thread that renews it (``set_vt``) every ``lease/3`` — so a live-but-slow task + stays claimed — then stops + joins on exit. 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 renewal uses its OWN + connection (``_renew_client``); the main thread owns ``_client`` for the ack, + and the thread is joined here before that ack so the two never race the row. + """ + 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) + + def _renew_lease_loop(self, msg_id: int, stop: threading.Event) -> None: + """Renew the claim every ``_lease_renew_interval`` until ``stop`` is set. + + Waits *then* renews (the initial claim already set the lease), so a task + shorter than the interval never renews. Best-effort: a transient error is + logged and retried next tick (the interval leaves ~2x slack before the + current lease expires); a ``False`` return means the row is already gone + (acked by us / expired + reclaimed) → stop. + """ + while not stop.wait(self._lease_renew_interval): + try: + # Created on the FIRST actual renewal (not on thread start), so a + # task shorter than the interval opens no second connection. + if self._renew_client is None: + self._renew_client = PgQueueClient() + if not self._renew_client.set_vt(msg_id, self.lease_seconds): + return # row gone — nothing left to renew + except Exception: + logger.warning( + "PG-queue consumer: lease renewal failed for msg_id=%s (retry in " + "%ss; lease %ss) — a dead connection self-heals on the next tick", + msg_id, + self._lease_renew_interval, + self.lease_seconds, + exc_info=True, + ) + def _handle(self, message: QueueMessage) -> None: payload = message.message task_name = payload.get("task_name") @@ -343,12 +419,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 (PR C). + 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 @@ -917,6 +996,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..2f74192aa4 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -10,6 +10,7 @@ import logging import os +import threading from unittest.mock import MagicMock, patch import pytest @@ -1055,3 +1056,65 @@ 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 PR C: 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. + """ + + def test_lease_capped_at_vt(self): + c = PgQueueConsumer(["q"], client=MagicMock(), vt_seconds=100, lease_seconds=200) + assert c.lease_seconds == 100 # capped — the lease is the shorter window + + 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_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) + + def test_renew_loop_renews_until_stopped(self): + c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) + c._renew_client = MagicMock() + c._renew_client.set_vt.return_value = True + stop = MagicMock() + stop.wait.side_effect = [False, False, True] # renew twice, then stop + c._renew_lease_loop(99, stop) + assert c._renew_client.set_vt.call_count == 2 + c._renew_client.set_vt.assert_called_with(99, 3) # renews with the lease + stop.wait.assert_called_with(c._lease_renew_interval) + + def test_renew_loop_stops_when_row_gone(self): + c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) + c._renew_client = MagicMock() + c._renew_client.set_vt.return_value = False # row already acked/expired + stop = MagicMock() + stop.wait.side_effect = [False, False, True] + c._renew_lease_loop(99, stop) + assert c._renew_client.set_vt.call_count == 1 # stopped after the first gone + + def test_renew_loop_swallows_transient_error_and_retries(self): + c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) + c._renew_client = MagicMock() + c._renew_client.set_vt.side_effect = [RuntimeError("db blip"), True] + stop = MagicMock() + stop.wait.side_effect = [False, False, True] + c._renew_lease_loop(99, stop) # must NOT raise + assert c._renew_client.set_vt.call_count == 2 # retried after the error + + def test_lease_renewal_ctx_runs_loop_and_sets_stop_on_exit(self): + c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) + with patch.object(c, "_renew_lease_loop") as loop: + with c._lease_renewal(42): + pass + 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 From f49ebb22a525cedfc69141323c169220a275f88b Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 13:12:53 +0530 Subject: [PATCH 2/5] =?UTF-8?q?UN-3695=20[FIX]=20test=20=E2=80=94=20fill?= =?UTF-8?q?=20empty=20with-block=20in=20lease-renewal=20CM=20test=20(Sonar?= =?UTF-8?q?=20S108)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The context-manager test entered `with c._lease_renewal(42): pass` — an empty nested block (python:S108). Replace the bare pass with an observable statement (entered = True) asserted after, so the block is non-empty and still verifies the CM body runs while the renewal thread is live. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/tests/test_pg_queue_consumer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 2f74192aa4..243728ec2d 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -1111,9 +1111,11 @@ def test_renew_loop_swallows_transient_error_and_retries(self): def test_lease_renewal_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): - pass + 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 From fbbf03ef5efda690be12a3b746f6315bb46bc8b4 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 13:31:48 +0530 Subject: [PATCH 3/5] =?UTF-8?q?UN-3695=20[FIX]=20address=20review=20?= =?UTF-8?q?=E2=80=94=20renewal=20thread=20owns=20its=20connection,=20obser?= =?UTF-8?q?vable=20failures,=20batch=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review round on the renewable-lease change (#2154): - Race: the renewal thread no longer shares self._renew_client. It owns a thread-local connection built lazily on the first renewal and closed in a finally, so a join-timeout wedge can't leave the next message's thread reusing a half-used psycopg2 connection. _make_renew_client() is the (patchable) factory. - Wedge is no longer silent: after the bounded join, a still-alive thread is logged at ERROR (connection abandoned until process exit). - set_vt -> False mid-run is a double-run signal (the ack runs only after this thread is joined, so the row can only be gone if another consumer reclaimed+acked it) — now logged, and the docstring corrected (was wrongly "acked by us"). - Narrowed the renewal except to CONN_DEAD_ERRORS (a non-connection error propagates instead of being swallowed as a "self-healing blip" forever) and added escalation: once renewal keeps failing past lease_seconds the lease is genuinely lost → ERROR. - BATCH_SIZE forced to 1 when lease < vt: renewal only covers the in-flight message, so a batch tail could lapse and double-run. All shipped workers are batch=1; this guards the knob. - Clamp is loud: lease > vt now warns (was a silent min()); comment corrected. - Iterator return annotation on the context manager; "PR C" shorthand → UN-3695. - README: effective claim window is min(LEASE, VT); lease is not a hard runtime cap (liveness probe is the backstop); batch-forced-to-1 note. - Tests: +9 (loud clamp / tiny-interval floor / non-positive reject / batch-forced / batch-kept-at-lease==vt / env wiring / renew+close / no-conn-for-sub-interval / row-gone-logs / conn-death-retry-then-escalate / non-conn-propagates / wedged-log / _handle-wraps-in-lease). 87 consumer tests pass. Deferred: lease metrics counters (ConsumerMetrics isn't wired to the consumer instance) — a separate small observability change. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/queue_backend/pg_queue/README.md | 23 +-- workers/queue_backend/pg_queue/consumer.py | 150 ++++++++++++++------ workers/tests/test_pg_queue_consumer.py | 157 +++++++++++++++++---- 3 files changed, 253 insertions(+), 77 deletions(-) diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md index 372b11d689..67ca71633c 100644 --- a/workers/queue_backend/pg_queue/README.md +++ b/workers/queue_backend/pg_queue/README.md @@ -27,8 +27,8 @@ integers. | Env suffix | One-line | Default | |---|---|---| | `CONCURRENCY` | Prefork consumer children per pod (1 = plain single process) | `1` | -| `VT_SECONDS` | Since PR C, 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** — a claimed message is hidden only this long and renewed every ~LEASE/3 while the task runs; a dead worker's claim expires in ~LEASE → fast redelivery (capped at VT) | `120` | +| `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` | @@ -70,18 +70,23 @@ 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 PR C)** — rather than claim for the full -`VT` (up to 2.5h), the consumer claims for a **short** `LEASE` and a background thread +**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 capped at it. The renewal -uses its own DB connection and is best-effort (a transient failure retries within the -`~2×` slack the `LEASE/3` interval leaves before expiry). +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 b0ee30f814..5914379386 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -27,7 +27,7 @@ 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 @@ -38,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 @@ -59,7 +60,7 @@ # 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 PR C). A claimed message is hidden for only +# 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 / @@ -258,16 +259,33 @@ def __init__( self._api_client = api_client self.batch_size = batch_size self.vt_seconds = vt_seconds - # Renewable lease (UN-3695 PR C): claim for lease_seconds (short) and renew - # every ~lease/3 while a task runs. Capped at vt_seconds so the lease is - # genuinely the shorter window; setting lease >= vt degrades to the old - # single-long-claim behaviour (renewal fires rarely). + # 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) - # Second client (own connection) for lease renewals — the main thread owns - # self._client for the read/ack, so a renewal must not share that connection. - # Lazily created on the first task run. - self._renew_client: PgQueueClient | None = None + # 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 @@ -313,16 +331,18 @@ def poll_once(self) -> int: return total @contextlib.contextmanager - def _lease_renewal(self, msg_id: int): - """Keep ``msg_id``'s claim alive while a task runs (UN-3695 PR C). - - The claim was taken for the short ``lease_seconds``; this starts a daemon - thread that renews it (``set_vt``) every ``lease/3`` — so a live-but-slow task - stays claimed — then stops + joins on exit. 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 renewal uses its OWN - connection (``_renew_client``); the main thread owns ``_client`` for the ack, - and the thread is joined here before that ack so the two never race the row. + 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( @@ -337,33 +357,77 @@ def _lease_renewal(self, msg_id: int): 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 the claim every ``_lease_renew_interval`` until ``stop`` is set. + """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. Best-effort: a transient error is - logged and retried next tick (the interval leaves ~2x slack before the - current lease expires); a ``False`` return means the row is already gone - (acked by us / expired + reclaimed) → stop. + 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). """ - while not stop.wait(self._lease_renew_interval): - try: - # Created on the FIRST actual renewal (not on thread start), so a - # task shorter than the interval opens no second connection. - if self._renew_client is None: - self._renew_client = PgQueueClient() - if not self._renew_client.set_vt(msg_id, self.lease_seconds): - return # row gone — nothing left to renew - except Exception: - logger.warning( - "PG-queue consumer: lease renewal failed for msg_id=%s (retry in " - "%ss; lease %ss) — a dead connection self-heals on the next tick", - msg_id, - self._lease_renew_interval, - self.lease_seconds, - exc_info=True, - ) + 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.error( + "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, + exc_info=True, + ) + 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 @@ -420,7 +484,7 @@ def _handle(self, message: QueueMessage) -> None: fairness = payload.get("fairness") headers = {FAIRNESS_HEADER_NAME: fairness} if fairness else None # 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 (PR C). + # 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 [], diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index 243728ec2d..c8a3d56029 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -11,6 +11,7 @@ import logging import os import threading +import time from unittest.mock import MagicMock, patch import pytest @@ -18,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. @@ -1059,20 +1061,58 @@ def test_chain_continuation_records_status_even_if_enqueue_fails(self): class TestLeaseRenewal: - """UN-3695 PR C: 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. + """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. """ - def test_lease_capped_at_vt(self): - c = PgQueueConsumer(["q"], client=MagicMock(), vt_seconds=100, lease_seconds=200) - assert c.lease_seconds == 100 # capped — the lease is the shorter window + # --- 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): + with pytest.raises(ValueError, match="lease_seconds must be positive"): + PgQueueConsumer(["q"], client=MagicMock(), 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 = [] @@ -1080,36 +1120,77 @@ def test_poll_claims_with_lease_not_vt(self): c.poll_once() client.read.assert_called_with("q", vt_seconds=120, qty=c.batch_size) - def test_renew_loop_renews_until_stopped(self): - c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) - c._renew_client = MagicMock() - c._renew_client.set_vt.return_value = True + # --- 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 c._renew_client.set_vt.call_count == 2 - c._renew_client.set_vt.assert_called_with(99, 3) # renews with the lease + 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_renew_loop_stops_when_row_gone(self): - c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) - c._renew_client = MagicMock() - c._renew_client.set_vt.return_value = False # row already acked/expired + def test_no_client_built_for_sub_interval_task(self): + c = self._renewing(MagicMock()) stop = MagicMock() - stop.wait.side_effect = [False, False, True] + stop.wait.side_effect = [True] # stopped before the first renewal tick c._renew_lease_loop(99, stop) - assert c._renew_client.set_vt.call_count == 1 # stopped after the first gone + c._make_renew_client.assert_not_called() # short task opens no 2nd connection - def test_renew_loop_swallows_transient_error_and_retries(self): - c = PgQueueConsumer(["q"], client=MagicMock(), lease_seconds=3) - c._renew_client = MagicMock() - c._renew_client.set_vt.side_effect = [RuntimeError("db blip"), True] + 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] - c._renew_lease_loop(99, stop) # must NOT raise - assert c._renew_client.set_vt.call_count == 2 # retried after the error + with pytest.raises(AttributeError): + c._renew_lease_loop(99, stop) + rc.close.assert_called_once() # still released via finally - def test_lease_renewal_ctx_runs_loop_and_sets_stop_on_exit(self): + # --- 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: @@ -1120,3 +1201,29 @@ def test_lease_renewal_ctx_runs_loop_and_sets_stop_on_exit(self): 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) + # 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): + pass + 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 From dc232ebbc00c1d1a6c698ec3e0c2d83e4f8aac16 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 13:52:40 +0530 Subject: [PATCH 4/5] =?UTF-8?q?UN-3695=20[FIX]=20resolve=20SonarCloud=20sm?= =?UTF-8?q?ells=20=E2=80=94=20logging.exception,=20single-throw=20pytest.r?= =?UTF-8?q?aises,=20empty=20block?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - consumer.py: escalation log uses logger.exception() (ERROR + traceback) instead of logger.error(..., exc_info=True) in the renewal except block. - tests: one throwing call per pytest.raises block; fill the empty _lease_renewal context-manager with-block (S108). Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/queue_backend/pg_queue/consumer.py | 3 +-- workers/tests/test_pg_queue_consumer.py | 7 +++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index 5914379386..b04c8dd828 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -404,14 +404,13 @@ def _renew_lease_loop(self, msg_id: int, stop: threading.Event) -> None: except CONN_DEAD_ERRORS: down_for = time.monotonic() - last_ok if down_for >= self.lease_seconds: - logger.error( + 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, - exc_info=True, ) else: logger.warning( diff --git a/workers/tests/test_pg_queue_consumer.py b/workers/tests/test_pg_queue_consumer.py index c8a3d56029..24b505d8f4 100644 --- a/workers/tests/test_pg_queue_consumer.py +++ b/workers/tests/test_pg_queue_consumer.py @@ -1086,8 +1086,9 @@ def test_tiny_lease_renew_interval_floored_at_1(self): 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=MagicMock(), lease_seconds=0) + 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"): @@ -1204,6 +1205,7 @@ def test_ctx_runs_loop_and_sets_stop_on_exit(self): 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( @@ -1213,7 +1215,8 @@ def test_ctx_logs_when_thread_wedged_past_join_timeout(self, caplog): logging.ERROR, logger="queue_backend.pg_queue.consumer" ): with c._lease_renewal(7): - pass + 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): From e3d72d6258b878cb7d4bb47457981282dfa4f5e1 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 14:23:15 +0530 Subject: [PATCH 5/5] UN-3695 [FIX] startup log shows lease (the claim window), not just vt (Greptile P2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit vt_seconds is now the drain/max-runtime bound, not the claim window — the startup log showing only vt=%ss misleads an operator into thinking the claim window is 9060s when it's actually lease_seconds (e.g. 120s). Add lease=%ss. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/queue_backend/pg_queue/consumer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/workers/queue_backend/pg_queue/consumer.py b/workers/queue_backend/pg_queue/consumer.py index b04c8dd828..b90ca57a0e 100644 --- a/workers/queue_backend/pg_queue/consumer.py +++ b/workers/queue_backend/pg_queue/consumer.py @@ -968,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)",