From d879e8a6c22549bfd0cced4f349629280de9da41 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 11:36:30 +0530 Subject: [PATCH 1/4] =?UTF-8?q?UN-3695=20[FIX]=20PG=20consumer=20=E2=80=94?= =?UTF-8?q?=20supervisor=20shutdown=20grace=20=3D=20VT,=20single=20shared?= =?UTF-8?q?=20join=20deadline?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prefork supervisor SIGKILLed its own children after a hardcoded 30s on shutdown (`_SHUTDOWN_GRACE_SECONDS = 30.0`). But the chart grants the pod `terminationGracePeriodSeconds = VT + buffer` (9120s for worker-pg-file-processing, VT=9060) — so on a graceful SIGTERM (deploy / HPA pod scale-down) the supervisor was knifing an in-flight batch ~300x sooner than the pod grace it lives inside, orphaning the batch (executor finished, but nothing consumed the reply / marked COMPLETE / decremented the barrier / acked the message). - `shutdown_grace_from_env()`: the per-child drain grace now defaults to the consumer VT (`WORKER_PG_QUEUE_CONSUMER_VT_SECONDS`), floored at 30 so it can never be SHORTER than before (non-regressive; OSS/dev with the default VT stays 30, fileproc gets 9060 with no chart change). Explicit `WORKER_PG_QUEUE_CONSUMER_SHUTDOWN_GRACE_SECONDS` overrides. - `_join_children` now uses a SINGLE shared deadline instead of a per-child recomputed one. Children are SIGTERM'd together and drain in parallel, so one shared window still gives each its full grace while bounding total drain to ~grace. A per-child deadline would sum to N x grace and, at grace=VT, blow past the pod grace so k8s SIGKILLs the whole pod, hard-killing siblings that were draining cleanly. PG-worker-component only — the supervisor never runs for Celery workers, so no flag branch is needed and the Celery path is untouched. Tests: shutdown_grace_from_env (VT-tracking / floor / override) + a shared-deadline regression test (38 supervisor tests). Verified e2e with real forked children: grace tracks VT and drains a live child (no premature SIGKILL); 2 wedged children SIGKILL together at the shared deadline, not 2x. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/pg_queue_consumer/supervisor.py | 58 +++++++++++++++++--- workers/tests/test_pg_consumer_supervisor.py | 58 ++++++++++++++++++++ 2 files changed, 107 insertions(+), 9 deletions(-) diff --git a/workers/pg_queue_consumer/supervisor.py b/workers/pg_queue_consumer/supervisor.py index bc2a5564ed..a17663e12b 100644 --- a/workers/pg_queue_consumer/supervisor.py +++ b/workers/pg_queue_consumer/supervisor.py @@ -63,8 +63,12 @@ _MIN_HEALTHY_UPTIME_SECONDS = 10.0 # Consecutive immediate crashes after which the fleet probe is forced unhealthy. _CRASH_LOOP_THRESHOLD = 3 -# How long to wait per child for a graceful drain on shutdown before SIGKILL. -_SHUTDOWN_GRACE_SECONDS = 30.0 +# Fallback per-child graceful-drain budget (s) on shutdown, used only when neither +# an explicit override nor the consumer VT is set — see shutdown_grace_from_env(). +# This is NOT the live value: the live grace defaults to the consumer's visibility +# timeout so a SIGTERM (deploy / HPA scale-down) lets an in-flight batch finish +# instead of a mid-flight SIGKILL that orphans it (UN-3695). +_DEFAULT_SHUTDOWN_GRACE_SECONDS = 30.0 def concurrency_from_env() -> int: @@ -93,6 +97,31 @@ def concurrency_from_env() -> int: return n +def shutdown_grace_from_env() -> float: + """Per-child graceful-drain budget (seconds) on shutdown, before SIGKILL. + + Defaults to the consumer's visibility timeout (``WORKER_PG_QUEUE_CONSUMER_VT_SECONDS``) + — the by-design upper bound on a single task's runtime — so a graceful SIGTERM + (deploy / HPA scale-down) lets an in-flight batch finish rather than being + SIGKILLed mid-flight and orphaned (UN-3695). The chart sets the pod's + ``terminationGracePeriodSeconds`` to VT + a buffer, so even a genuinely-wedged + child is SIGKILLed here just BEFORE k8s reaps the pod. An explicit + ``WORKER_PG_QUEUE_CONSUMER_SHUTDOWN_GRACE_SECONDS`` overrides (and is honoured + as-is, e.g. a short dev drain); otherwise the VT is floored at + ``_DEFAULT_SHUTDOWN_GRACE_SECONDS`` so a tiny/unset VT still gets a sane drain. + + A hardcoded 30s here previously undercut the chart's ~2.5h budget by ~300×, + SIGKILLing in-flight batches on every scale-down/rollout. + """ + from queue_backend.pg_queue.consumer import _DEFAULT_VT_SECONDS, consumer_env + + override = consumer_env("SHUTDOWN_GRACE_SECONDS", None, float) + if override is not None: + return max(0.0, override) + vt = consumer_env("VT_SECONDS", _DEFAULT_VT_SECONDS, int) + return max(float(_DEFAULT_SHUTDOWN_GRACE_SECONDS), float(vt)) + + class _Fleet: """Owns the per-slot child state — pid, last-fork, heartbeat, crash count and pending-restart schedule — keeping them mutually consistent. Slots are @@ -332,6 +361,10 @@ def run_supervised(concurrency: int) -> None: logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO")) fleet = _Fleet(concurrency) + grace_seconds = shutdown_grace_from_env() + logger.info( + "PG-queue consumer supervisor: shutdown drain grace = %.0fs/child", grace_seconds + ) stopping = threading.Event() def _signal_children(sig: int) -> None: @@ -357,7 +390,7 @@ def _on_term(signum: int, _frame: object) -> None: if not _try_fork_child(fleet, slot): stopping.set() _signal_children(signal.SIGTERM) - _join_children(fleet, _SHUTDOWN_GRACE_SECONDS) + _join_children(fleet, grace_seconds) raise RuntimeError( f"PG-queue consumer: os.fork() failed starting child {slot}/" f"{concurrency} — reduce WORKER_PG_QUEUE_CONSUMER_CONCURRENCY or " @@ -373,7 +406,7 @@ def _on_term(signum: int, _frame: object) -> None: finally: stopping.set() _signal_children(signal.SIGTERM) - _join_children(fleet, _SHUTDOWN_GRACE_SECONDS) + _join_children(fleet, grace_seconds) if health is not None: health.stop() logger.info("PG-queue consumer supervisor: stopped") @@ -395,15 +428,22 @@ def _wait_for_exit(pid: int, deadline: float) -> bool: def _join_children(fleet: _Fleet, grace_seconds: float) -> None: - """Wait up to ``grace_seconds`` *per child* for a graceful drain; SIGKILL + - reap any straggler. The budget is per-child (not a single shared deadline) so - one slow-draining child can't starve the others of their grace. + """Wait up to ``grace_seconds`` TOTAL for all children to drain, then SIGKILL + + reap any straggler. A SINGLE shared deadline (not per-child): the children are + all SIGTERM'd together *before* this call and drain in parallel, so one shared + window still gives every child its full grace from that common SIGTERM — while + bounding the total wait to ~``grace_seconds``. A per-child deadline would instead + sum to N×grace and, at grace≈VT (thousands of seconds), blow past the pod's + ``terminationGracePeriodSeconds`` — so k8s SIGKILLs the whole pod, hard-killing + siblings that were still draining cleanly (UN-3695 / debate H2). """ + deadline = time.monotonic() + grace_seconds for slot, pid in fleet.alive_items(): - if _wait_for_exit(pid, time.monotonic() + grace_seconds): + if _wait_for_exit(pid, deadline): continue logger.warning( - "PG-queue consumer: child slot=%s pid=%s did not stop in %ss — SIGKILL", + "PG-queue consumer: child slot=%s pid=%s did not drain within the " + "shared %ss grace — SIGKILL", slot, pid, grace_seconds, diff --git a/workers/tests/test_pg_consumer_supervisor.py b/workers/tests/test_pg_consumer_supervisor.py index c5a1912ca7..64f83b8a7a 100644 --- a/workers/tests/test_pg_consumer_supervisor.py +++ b/workers/tests/test_pg_consumer_supervisor.py @@ -16,6 +16,7 @@ from pg_queue_consumer import supervisor as sup from pg_queue_consumer.supervisor import ( _CRASH_LOOP_THRESHOLD, + _DEFAULT_SHUTDOWN_GRACE_SECONDS, _MAX_CONCURRENCY, _MIN_HEALTHY_UPTIME_SECONDS, _Fleet, @@ -26,12 +27,50 @@ _try_fork_child, _wait_for_exit, concurrency_from_env, + shutdown_grace_from_env, ) _MOD = "pg_queue_consumer.supervisor" _ENV = "WORKER_PG_QUEUE_CONSUMER_CONCURRENCY" +_VT = "WORKER_PG_QUEUE_CONSUMER_VT_SECONDS" +_GRACE = "WORKER_PG_QUEUE_CONSUMER_SHUTDOWN_GRACE_SECONDS" + + +class TestShutdownGraceFromEnv: + """UN-3695: the shutdown drain grace must track the consumer VT (not a hardcoded + 30s), so a SIGTERM drains an in-flight batch instead of SIGKILLing it mid-flight. + """ + + def test_unset_vt_and_override_defaults_to_fallback(self, monkeypatch): + monkeypatch.delenv(_VT, raising=False) + monkeypatch.delenv(_GRACE, raising=False) + assert shutdown_grace_from_env() == _DEFAULT_SHUTDOWN_GRACE_SECONDS + + def test_tracks_vt_when_set(self, monkeypatch): + # The fileproc chart sets VT=9060 → grace must follow it, not stay at 30. + monkeypatch.delenv(_GRACE, raising=False) + monkeypatch.setenv(_VT, "9060") + assert shutdown_grace_from_env() == 9060.0 + + def test_vt_below_fallback_is_floored(self, monkeypatch): + monkeypatch.delenv(_GRACE, raising=False) + monkeypatch.setenv(_VT, "10") + assert shutdown_grace_from_env() == _DEFAULT_SHUTDOWN_GRACE_SECONDS + + def test_explicit_override_wins_over_vt(self, monkeypatch): + monkeypatch.setenv(_VT, "9060") + monkeypatch.setenv(_GRACE, "120") + assert shutdown_grace_from_env() == 120.0 + + def test_override_honoured_as_is_even_below_fallback(self, monkeypatch): + # An explicit short dev drain is respected (not floored) — only the VT path floors. + monkeypatch.delenv(_VT, raising=False) + monkeypatch.setenv(_GRACE, "5") + assert shutdown_grace_from_env() == 5.0 + + class TestConcurrencyFromEnv: def test_unset_defaults_to_one(self, monkeypatch): monkeypatch.delenv(_ENV, raising=False) @@ -258,6 +297,25 @@ def test_straggler_is_sigkilled_and_reaped(self): kill.assert_called_once_with(111, _signal.SIGKILL) waitpid.assert_called_once_with(111, 0) + def test_shared_deadline_not_per_child(self): + # UN-3695 / debate H2: all children must be waited against ONE shared + # deadline. A per-child deadline serializes N wedged children to N×grace, + # which at grace≈VT blows past the pod's terminationGracePeriodSeconds and + # gets siblings hard-killed by k8s. Assert a single shared deadline is used. + f = _Fleet(3) + for slot, pid in enumerate((111, 222, 333)): + f.record_fork(slot, pid) + deadlines: list[float] = [] + + def _capture(_pid, deadline): + deadlines.append(deadline) + return True # each child exits cleanly + + with patch(f"{_MOD}._wait_for_exit", side_effect=_capture): + _join_children(f, grace_seconds=5) + assert len(deadlines) == 3 + assert len(set(deadlines)) == 1 # one shared deadline, not three per-child + class TestSupervisorHealth: def test_no_port_returns_none(self, monkeypatch): From baa0199d430208f2c560aa064554a1b03a2dec03 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 11:58:11 +0530 Subject: [PATCH 2/4] =?UTF-8?q?UN-3695=20[DOC]=20PG=20queue=20=E2=80=94=20?= =?UTF-8?q?reference=20&=20glossary=20(terms=20+=20config=20knobs)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a co-located README as a lookup dictionary for the PG-queue transport: config env vars (consumer / reaper / orchestration / routing) with defaults, a concepts glossary (visibility timeout, claim, redelivery, reaper, barrier, poison/re-park, prefork supervisor, shutdown grace, request-reply/callback, last_progress_at, leader election, fairness), and the table catalogue. Complements the 9e-design.md design doc. Docs only — no code change. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/queue_backend/pg_queue/README.md | 153 +++++++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 workers/queue_backend/pg_queue/README.md diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md new file mode 100644 index 0000000000..23d6675896 --- /dev/null +++ b/workers/queue_backend/pg_queue/README.md @@ -0,0 +1,153 @@ +# PG Queue — Reference & Glossary + +A quick-reference dictionary for the bespoke Postgres-backed work queue that can +stand in for Celery/RabbitMQ (epic UN-3445). For the *design* of a full-execution +migration see [`9e-design.md`](./9e-design.md); this file is the **terminology + config +lookup**. + +The transport is chosen **per org** by the Flipt flag `pg_queue_enabled` +(fail-closed to Celery), so the system runs **with and without** PG unchanged. + +--- + +## 1. Configuration — quick reference + +All knobs are environment variables. Consumer knobs are read via +`consumer_env("", …)` → `WORKER_PG_QUEUE_CONSUMER_`. Authoritative +defaults live in the code constants (`consumer.py`, `pg_queue_consumer/supervisor.py`, +`reaper.py`); cloud overrides live in the Helm chart `values.yaml`. + +**Units:** every duration is in **seconds**, and the **Default column shows the raw +value you set** — a plain number, e.g. `30` means 30 seconds (set `…=30`, not `30s`). +The `*_SECONDS` suffix names the unit; `POLL_INTERVAL` and `BACKOFF_MAX` are also +seconds (floats allowed). Counts (`CONCURRENCY`, `BATCH`, `MAX_ATTEMPTS`) are plain +integers. + +### Consumer (`WORKER_PG_QUEUE_CONSUMER_*`) +| 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) | +| `SHUTDOWN_GRACE_SECONDS` | Per-child graceful-drain budget on SIGTERM before SIGKILL | `= VT` (floored at `30`) | +| `QUEUE` | Queue name(s) this consumer polls (comma-separated) | `default` | +| `BATCH` | Messages claimed per poll | `1` | +| `POLL_INTERVAL` | Time between polls when the queue is empty | `0.1` | +| `BACKOFF_MAX` | Max empty-queue poll backoff | `2.0` | +| `MAX_ATTEMPTS` | Max deliveries before a message is dropped as **poison** | `5` | +| `HEALTH_PORT` | Liveness HTTP port (unset → probe disabled) | unset | +| `HEALTH_STALE_SECONDS` | A poll loop idle beyond this is reported unhealthy | `60` | +| `WORKER_TYPE` | Which source worker's tasks this consumer registers (bootstrap) | — | + +### Reaper (`WORKER_PG_REAPER_*`) +| Env | One-line | Default | +|---|---|---| +| `INTERVAL_SECONDS` | Reaper cycle cadence (stuck-batch / orphan detection) | `5` | +| `SWEEP_SECONDS` | Retention-sweep cadence (expired `pg_task_result`, etc.) | `300` (5 min) | +| `HEALTH_PORT` / `HEALTH_STALE_SECONDS` | Reaper liveness probe | stale `30` | + +### Orchestration / barrier / connection +| Env | One-line | +|---|---| +| `WORKER_PG_BATCH_STUCK_TIMEOUT_SECONDS` | A barrier whose `last_progress_at` hasn't advanced this long is fast-failed by the reaper (UN-3661) | +| `WORKER_PG_ORCHESTRATOR_LEASE_SECONDS` | Leader-election lease duration for the single-orchestrator role | +| `WORKER_PG_DEDUP_RETENTION_SECONDS` | Retention for `pg_batch_dedup` rows | +| `WORKER_PG_QUEUE_CONNECT_RETRIES` / `_BACKOFF` | Reconnect attempts / backoff on a stale or broken DB connection | + +### Routing / flag +| Env / flag | One-line | +|---|---| +| `pg_queue_enabled` (Flipt, `PG_QUEUE_FLAG_KEY`) | Per-org gate that routes dispatch to PG; **fail-closed to Celery** on any Flipt error | +| `WORKER_PG_QUEUE_ENABLED_TASKS` | Comma-separated task names eligible for PG routing (empty → everything stays on Celery) | + +--- + +## 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. + +**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. + +**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. + +**Near-exactly-once** — writes to an *external* system (e.g. a customer destination +DB) can't be transactional with our Postgres markers, so a crash in the tiny +write→mark gap can duplicate. Guards (FileHistory, etc.) reduce this to a narrow +window, but it isn't strictly exactly-once. + +**Poison message / re-park** — a message that fails `MAX_ATTEMPTS` times. Rather than +redeliver forever it is dropped (or re-parked with `POISON_REPARK_VT_SECONDS`, bounded +by a budget) so a permanently-bad row can't wedge the queue. + +**Prefork supervisor / fleet** — `WORKER_PG_QUEUE_CONSUMER_CONCURRENCY > 1` forks N +isolated consumer child processes (the PG analogue of Celery `--pool=prefork +--concurrency=N`). The *supervisor* owns the liveness port, re-forks crashed children +(rate-limited), and drains them on shutdown. The *fleet* is the set of children. + +**Shutdown grace** — on SIGTERM the supervisor waits up to `SHUTDOWN_GRACE_SECONDS` +(= VT) for children to finish their in-flight batch, using a **single shared +deadline**, then SIGKILLs stragglers. Must be ≤ the pod's +`terminationGracePeriodSeconds` (the chart sets `grace = VT + buffer`). + +**Liveness / heartbeat** — each child stamps its last-poll time into a shared array; +the supervisor reports the *oldest* child's staleness on `/health`. Frozen during a +long task, so a wedged child goes stale and trips the probe. + +**Reaper** — a singleton (leader-elected) sweeper that recovers **stranded** work: +fast-fails a barrier whose `last_progress_at` stalled (UN-3661), cascades a terminal +execution to its files, and sweeps expired retention rows. + +**Barrier (`PgBarrierState`)** — the fan-in counter for a batched execution: each +file-batch decrements `remaining`; when it hits 0 the aggregating callback fires. It +carries `last_progress_at` (the reaper's liveness signal for stuck detection). + +**`last_progress_at`** — timestamp bumped whenever a batch makes progress +(UN-3661). Lets the reaper fast-fail a stalled batch in ~stuck-timeout instead of the +full VT/6h horizon. + +**Leader election / orchestrator lock (`pg_orchestrator_lock`)** — a DB lease that +elects a single reaper/orchestrator across replicas; released on graceful shutdown so a +standby takes over promptly. + +**Request-reply (`reply_key`)** — a *blocking* dispatch: the caller enqueues with a +`reply_key`, the consumer stores the result under it in `pg_task_result`, and the +caller polls for it (the PG analogue of Celery's `AsyncResult`). Mutually exclusive +with the callback form. + +**Callback dispatch (`dispatch_with_callback`, `on_success` / `on_error`)** — +fire-and-forget: after the task runs, the consumer **self-chains** the success or error +continuation onto its queue via `_chain_continuation` (the PG analogue of Celery +`link` / `link_error`). No blocking caller. + +**Fairness** — a header carried on a dispatch (`org_id`, workload type, priority) so a +PG-routed run mirrors Celery's fair scheduling. + +--- + +## 3. Tables (`backend/pg_queue/models.py`) + +| Table | Purpose | +|---|---| +| `pg_queue_message` | The queue itself — one row per message; claimed via `SKIP LOCKED` + VT | +| `pg_task_result` | Request-reply results / terminal task status, keyed by reply/task id; TTL'd (`expires_at`) + reaper-swept, `ON CONFLICT DO NOTHING` | +| `pg_barrier_state` | Fan-in barrier counter for batched executions (`remaining`, `last_progress_at`) | +| `pg_batch_dedup` | Batch-level dedup guard (prevents double-dispatch of a batch) | +| `pg_orchestration_claim` | Per-execution orchestration claim (`last_progress_at` for stuck detection) | +| `pg_orchestrator_lock` | Leader-election lease for the singleton reaper/orchestrator | +| `pg_periodic_schedule` | Periodic-schedule store (the Celery-Beat replacement) | + +All PG-queue tables are **droppable side tables** that never touch +`WorkflowExecution` (UN-3533) — the PG transport can be turned off and its tables +dropped without affecting the Celery path. + +--- + +*Defaults here reflect the code constants at time of writing; the code +(`consumer.py` / `supervisor.py` / `reaper.py`) and the chart `values.yaml` are the +source of truth.* From 20462e0448a7ab6601b13a76d374ccf23d87ef4e Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 12:14:27 +0530 Subject: [PATCH 3/4] =?UTF-8?q?UN-3695=20[FIX]=20address=20review=20?= =?UTF-8?q?=E2=80=94=20validate=20grace=20override,=20fix=20shared-deadlin?= =?UTF-8?q?e=20false-alarm,=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review round on the grace=VT / shared-deadline change: - shutdown_grace_from_env: reject a negative / non-finite SHUTDOWN_GRACE_SECONDS override (was silently collapsed by max(0.0, ...) -> 0 = immediate SIGKILL, or inf = shutdown hang) — fail fast like concurrency_from_env; annotate override: float|None; drop the redundant float() on the already-float default. - _wait_for_exit: poll waitpid at least once regardless of the deadline. On the single SHARED deadline a child iterated after the window elapsed was reported 'did not drain' and SIGKILLed — a false hard-kill alarm — even though it exited. - Wording: the grace is a single shared total, not per-child — corrected the docstring, constant comment, startup log ('shared across children'), and the README config row to match the _join_children docstring. - Tests: shared-deadline hardened with an advancing time.monotonic (exact 105/105/105); a 3-child all-wedged straggler case (kill+waitpid x3, one shared deadline); negative / non-finite / malformed override + malformed VT fail-fast; named helpers replace the lambda idiom; float asserts use pytest.approx (Sonar S1244). 43 tests pass; e2e re-verified. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/pg_queue_consumer/supervisor.py | 38 +++++++--- workers/queue_backend/pg_queue/README.md | 2 +- workers/tests/test_pg_consumer_supervisor.py | 80 ++++++++++++++++---- 3 files changed, 96 insertions(+), 24 deletions(-) diff --git a/workers/pg_queue_consumer/supervisor.py b/workers/pg_queue_consumer/supervisor.py index a17663e12b..c406afb6ff 100644 --- a/workers/pg_queue_consumer/supervisor.py +++ b/workers/pg_queue_consumer/supervisor.py @@ -36,6 +36,7 @@ import contextlib import logging +import math import multiprocessing import os import signal @@ -63,8 +64,9 @@ _MIN_HEALTHY_UPTIME_SECONDS = 10.0 # Consecutive immediate crashes after which the fleet probe is forced unhealthy. _CRASH_LOOP_THRESHOLD = 3 -# Fallback per-child graceful-drain budget (s) on shutdown, used only when neither -# an explicit override nor the consumer VT is set — see shutdown_grace_from_env(). +# Fallback graceful-drain budget (s, shared across all children) on shutdown, used +# only when neither an explicit override nor the consumer VT is set — see +# shutdown_grace_from_env(). # This is NOT the live value: the live grace defaults to the consumer's visibility # timeout so a SIGTERM (deploy / HPA scale-down) lets an in-flight batch finish # instead of a mid-flight SIGKILL that orphans it (UN-3695). @@ -98,7 +100,8 @@ def concurrency_from_env() -> int: def shutdown_grace_from_env() -> float: - """Per-child graceful-drain budget (seconds) on shutdown, before SIGKILL. + """Graceful-drain budget (seconds) on shutdown, shared across all children, + before SIGKILL. Defaults to the consumer's visibility timeout (``WORKER_PG_QUEUE_CONSUMER_VT_SECONDS``) — the by-design upper bound on a single task's runtime — so a graceful SIGTERM @@ -115,11 +118,20 @@ def shutdown_grace_from_env() -> float: """ from queue_backend.pg_queue.consumer import _DEFAULT_VT_SECONDS, consumer_env - override = consumer_env("SHUTDOWN_GRACE_SECONDS", None, float) + override: float | None = consumer_env("SHUTDOWN_GRACE_SECONDS", None, float) if override is not None: - return max(0.0, override) + # A negative / non-finite drain budget is never intentional and would + # re-introduce the exact failure this module prevents: <0 or nan → 0 → every + # child SIGKILLed with no drain; inf → shutdown hangs until k8s hard-kills the + # pod. Fail fast at startup, mirroring concurrency_from_env()'s validation. + if not math.isfinite(override) or override < 0: + raise ValueError( + "WORKER_PG_QUEUE_CONSUMER_SHUTDOWN_GRACE_SECONDS must be a finite " + f"number >= 0, got {override!r}" + ) + return override vt = consumer_env("VT_SECONDS", _DEFAULT_VT_SECONDS, int) - return max(float(_DEFAULT_SHUTDOWN_GRACE_SECONDS), float(vt)) + return max(_DEFAULT_SHUTDOWN_GRACE_SECONDS, float(vt)) class _Fleet: @@ -363,7 +375,9 @@ def run_supervised(concurrency: int) -> None: fleet = _Fleet(concurrency) grace_seconds = shutdown_grace_from_env() logger.info( - "PG-queue consumer supervisor: shutdown drain grace = %.0fs/child", grace_seconds + "PG-queue consumer supervisor: shutdown drain grace = %.0fs (shared across " + "children)", + grace_seconds, ) stopping = threading.Event() @@ -415,16 +429,22 @@ def _on_term(signum: int, _frame: object) -> None: def _wait_for_exit(pid: int, deadline: float) -> bool: """Poll ``pid`` until it exits or ``deadline`` (monotonic) passes. True if it exited (or was already reaped). + + Polls ``waitpid`` at least once regardless of the deadline: with the single + SHARED shutdown deadline, a child iterated after the window has already elapsed + would otherwise be reported "did not drain" and SIGKILLed — a false hard-kill + alarm — even though it exited cleanly within the grace. """ - while time.monotonic() < deadline: + while True: try: reaped, _status = os.waitpid(pid, os.WNOHANG) except ChildProcessError: return True if reaped != 0: return True + if time.monotonic() >= deadline: + return False time.sleep(0.1) - return False def _join_children(fleet: _Fleet, grace_seconds: float) -> None: diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md index 23d6675896..fa80e87fbc 100644 --- a/workers/queue_backend/pg_queue/README.md +++ b/workers/queue_backend/pg_queue/README.md @@ -28,7 +28,7 @@ integers. |---|---|---| | `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) | -| `SHUTDOWN_GRACE_SECONDS` | Per-child graceful-drain budget on SIGTERM before SIGKILL | `= VT` (floored at `30`) | +| `SHUTDOWN_GRACE_SECONDS` | Graceful-drain budget (shared across all children) on SIGTERM before SIGKILL | `= VT` (floored at `30`) | | `QUEUE` | Queue name(s) this consumer polls (comma-separated) | `default` | | `BATCH` | Messages claimed per poll | `1` | | `POLL_INTERVAL` | Time between polls when the queue is empty | `0.1` | diff --git a/workers/tests/test_pg_consumer_supervisor.py b/workers/tests/test_pg_consumer_supervisor.py index 64f83b8a7a..bc9145bf7a 100644 --- a/workers/tests/test_pg_consumer_supervisor.py +++ b/workers/tests/test_pg_consumer_supervisor.py @@ -46,29 +46,55 @@ class TestShutdownGraceFromEnv: def test_unset_vt_and_override_defaults_to_fallback(self, monkeypatch): monkeypatch.delenv(_VT, raising=False) monkeypatch.delenv(_GRACE, raising=False) - assert shutdown_grace_from_env() == _DEFAULT_SHUTDOWN_GRACE_SECONDS + assert shutdown_grace_from_env() == pytest.approx(_DEFAULT_SHUTDOWN_GRACE_SECONDS) def test_tracks_vt_when_set(self, monkeypatch): # The fileproc chart sets VT=9060 → grace must follow it, not stay at 30. monkeypatch.delenv(_GRACE, raising=False) monkeypatch.setenv(_VT, "9060") - assert shutdown_grace_from_env() == 9060.0 + assert shutdown_grace_from_env() == pytest.approx(9060.0) def test_vt_below_fallback_is_floored(self, monkeypatch): monkeypatch.delenv(_GRACE, raising=False) monkeypatch.setenv(_VT, "10") - assert shutdown_grace_from_env() == _DEFAULT_SHUTDOWN_GRACE_SECONDS + assert shutdown_grace_from_env() == pytest.approx(_DEFAULT_SHUTDOWN_GRACE_SECONDS) def test_explicit_override_wins_over_vt(self, monkeypatch): monkeypatch.setenv(_VT, "9060") monkeypatch.setenv(_GRACE, "120") - assert shutdown_grace_from_env() == 120.0 + assert shutdown_grace_from_env() == pytest.approx(120.0) def test_override_honoured_as_is_even_below_fallback(self, monkeypatch): # An explicit short dev drain is respected (not floored) — only the VT path floors. monkeypatch.delenv(_VT, raising=False) monkeypatch.setenv(_GRACE, "5") - assert shutdown_grace_from_env() == 5.0 + assert shutdown_grace_from_env() == pytest.approx(5.0) + + def test_negative_override_raises(self, monkeypatch): + # A negative drain budget would collapse to 0 → SIGKILL with no drain (the + # orphan bug this module prevents). Must fail fast, not silently. + monkeypatch.setenv(_GRACE, "-5") + with pytest.raises(ValueError, match="SHUTDOWN_GRACE_SECONDS"): + shutdown_grace_from_env() + + def test_non_finite_override_raises(self, monkeypatch): + # inf would make the join deadline now+inf → shutdown hangs until k8s + # hard-kills the pod; nan collapses to 0. Reject both. + monkeypatch.setenv(_GRACE, "inf") + with pytest.raises(ValueError, match="finite"): + shutdown_grace_from_env() + + def test_malformed_override_raises(self, monkeypatch): + monkeypatch.setenv(_GRACE, "abc") + with pytest.raises(ValueError, match="SHUTDOWN_GRACE_SECONDS"): + shutdown_grace_from_env() + + def test_malformed_vt_raises(self, monkeypatch): + # Fail-fast-at-startup contract: a bad VT surfaces the offending var name. + monkeypatch.delenv(_GRACE, raising=False) + monkeypatch.setenv(_VT, "9060s") + with pytest.raises(ValueError, match="VT_SECONDS"): + shutdown_grace_from_env() class TestConcurrencyFromEnv: @@ -298,23 +324,49 @@ def test_straggler_is_sigkilled_and_reaped(self): waitpid.assert_called_once_with(111, 0) def test_shared_deadline_not_per_child(self): - # UN-3695 / debate H2: all children must be waited against ONE shared - # deadline. A per-child deadline serializes N wedged children to N×grace, - # which at grace≈VT blows past the pod's terminationGracePeriodSeconds and - # gets siblings hard-killed by k8s. Assert a single shared deadline is used. + # UN-3695 / debate H2: all children waited against ONE shared deadline (a + # per-child deadline serializes N wedged children to N×grace and blows past + # the pod terminationGracePeriodSeconds → siblings hard-killed). Advancing + # time.monotonic pins it: shared → the deadline is computed once (105 for + # all three); a per-child recompute would yield 105/205/305. f = _Fleet(3) for slot, pid in enumerate((111, 222, 333)): f.record_fork(slot, pid) deadlines: list[float] = [] - def _capture(_pid, deadline): + def _record_drained(_pid, deadline): deadlines.append(deadline) - return True # each child exits cleanly + return True # child exited within the window + + with ( + patch(f"{_MOD}.time.monotonic", side_effect=[100.0, 200.0, 300.0, 400.0]), + patch(f"{_MOD}._wait_for_exit", side_effect=_record_drained), + ): + _join_children(f, grace_seconds=5) + assert deadlines == pytest.approx([105.0, 105.0, 105.0]) # computed once, shared + + def test_multi_child_stragglers_share_deadline_and_all_sigkilled(self): + # 3 wedged children (none drain) → each SIGKILLed + reaped against the SAME + # shared deadline. Kill-count alone wouldn't catch a per-child regression, + # so also assert the deadline handed to every child was identical. + f = _Fleet(3) + for slot, pid in enumerate((111, 222, 333)): + f.record_fork(slot, pid) + seen: list[float] = [] - with patch(f"{_MOD}._wait_for_exit", side_effect=_capture): + def _record_wedged(_pid, deadline): + seen.append(deadline) + return False # never drains → forces SIGKILL + + with ( + patch(f"{_MOD}._wait_for_exit", side_effect=_record_wedged), + patch(f"{_MOD}.os.kill") as kill, + patch(f"{_MOD}.os.waitpid") as waitpid, + ): _join_children(f, grace_seconds=5) - assert len(deadlines) == 3 - assert len(set(deadlines)) == 1 # one shared deadline, not three per-child + assert kill.call_count == 3 + assert waitpid.call_count == 3 + assert len(set(seen)) == 1 # one shared deadline across all three class TestSupervisorHealth: From 77e0f52218640d6f2dad3fa9af21778381c01b78 Mon Sep 17 00:00:00 2001 From: ali Date: Wed, 8 Jul 2026 12:31:43 +0530 Subject: [PATCH 4/4] UN-3695 [FIX] SIGKILL warning log uses %.0fs to match the startup log (Greptile nit) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The _join_children SIGKILL warning printed grace_seconds with %s (e.g. 9060.0s) while the startup log uses %.0f (9060s). Align to %.0fs. Cosmetic — no behaviour. Co-Authored-By: Claude Opus 4.8 (1M context) --- workers/pg_queue_consumer/supervisor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workers/pg_queue_consumer/supervisor.py b/workers/pg_queue_consumer/supervisor.py index c406afb6ff..e3553fcebd 100644 --- a/workers/pg_queue_consumer/supervisor.py +++ b/workers/pg_queue_consumer/supervisor.py @@ -463,7 +463,7 @@ def _join_children(fleet: _Fleet, grace_seconds: float) -> None: continue logger.warning( "PG-queue consumer: child slot=%s pid=%s did not drain within the " - "shared %ss grace — SIGKILL", + "shared %.0fs grace — SIGKILL", slot, pid, grace_seconds,