diff --git a/workers/pg_queue_consumer/supervisor.py b/workers/pg_queue_consumer/supervisor.py index bc2a5564ed..e3553fcebd 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,13 @@ _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 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). +_DEFAULT_SHUTDOWN_GRACE_SECONDS = 30.0 def concurrency_from_env() -> int: @@ -93,6 +99,41 @@ def concurrency_from_env() -> int: return n +def shutdown_grace_from_env() -> float: + """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 + (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: float | None = consumer_env("SHUTDOWN_GRACE_SECONDS", None, float) + if override is not None: + # 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(_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 +373,12 @@ 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 (shared across " + "children)", + grace_seconds, + ) stopping = threading.Event() def _signal_children(sig: int) -> None: @@ -357,7 +404,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 +420,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") @@ -382,28 +429,41 @@ 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: - """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 %.0fs grace — SIGKILL", slot, pid, grace_seconds, diff --git a/workers/queue_backend/pg_queue/README.md b/workers/queue_backend/pg_queue/README.md new file mode 100644 index 0000000000..fa80e87fbc --- /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` | 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` | +| `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.* diff --git a/workers/tests/test_pg_consumer_supervisor.py b/workers/tests/test_pg_consumer_supervisor.py index c5a1912ca7..bc9145bf7a 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,76 @@ _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() == 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() == 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() == 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() == 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() == 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: def test_unset_defaults_to_one(self, monkeypatch): monkeypatch.delenv(_ENV, raising=False) @@ -258,6 +323,51 @@ 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 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 _record_drained(_pid, deadline): + deadlines.append(deadline) + 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] = [] + + 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 kill.call_count == 3 + assert waitpid.call_count == 3 + assert len(set(seen)) == 1 # one shared deadline across all three + class TestSupervisorHealth: def test_no_port_returns_none(self, monkeypatch):