UN-3695 [FIX] PG consumer — supervisor shutdown grace = VT, single shared join deadline

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,

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("<SUFFIX>", …)` → `WORKER_PG_QUEUE_CONSUMER_<SUFFIX>`. 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.*