UN-4223 [FIX] Kill and restart only a hung PG consumer child instead of losing the whole worker pod

workers/pg_queue_consumer/supervisor.py

@@ -23,6 +23,17 @@
 a row, never reaching a real poll) forces the probe to 503 so k8s restarts the
 pod rather than the supervisor masking a wedged fleet with fresh-looking re-forks.
 
+**Child watchdog** (UN-4223, opt-in via ``CHILD_WATCHDOG``): a loaded child whose
+heartbeat goes stale past ``HEALTH_STALE_SECONDS``, or a child still not loaded
+that long after its fork, is SIGKILLed and re-forked on its own. Its lease-renewal
+thread dies with it, so the claim lapses and the reaper redelivers the message
+(the poison cap bounds a task that hangs on every attempt). The siblings keep
+running, so a pod without a liveness probe on ``/health`` loses one slot for one
+stale window instead of the whole fleet. ``/health`` itself still reports 503 from
+the stale crossing until the replacement loads (~20s), because ``oldest_age``
+counts the killed slot's heartbeat. SIGKILL rather than SIGTERM: a graceful stop
+waits for the in-flight task, which is the thing that is hung.
+
 **Readiness** (UN-4136): the same port serves ``/ready``, which answers 200 only
 once EVERY child has finished its ``import worker`` bootstrap and built its
 consumer. ``/health`` cannot say this — the heartbeats are seeded fresh at
@@ -72,6 +83,10 @@
 _MIN_HEALTHY_UPTIME_SECONDS = 10.0
 # Consecutive immediate crashes after which the fleet probe is forced unhealthy.
 _CRASH_LOOP_THRESHOLD = 3
+# Least time a child gets to finish ``import worker`` before the watchdog treats it
+# as hung. HEALTH_STALE is sized for task runtime and can be short (ide-callback:
+# 180s), while a bootstrap under a CPU cap can take minutes.
+_MIN_BOOTSTRAP_BUDGET_SECONDS = 600.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().
@@ -107,6 +122,47 @@ def concurrency_from_env() -> int:
     return n
 
 
+def _parse_bool(raw: str) -> bool:
+    value = raw.strip().lower()
+    if value in ("1", "true", "yes", "on"):
+        return True
+    if value in ("0", "false", "no", "off"):
+        return False
+    raise ValueError(f"expected a boolean, got {raw!r}")
+
+
+def child_watchdog_from_env() -> float | None:
+    """Heartbeat age (seconds) past which the supervisor SIGKILLs a single child, or
+    ``None`` when the watchdog is off.
+
+    Off unless ``WORKER_PG_QUEUE_CONSUMER_CHILD_WATCHDOG`` is true. Opt-in because
+    it turns ``HEALTH_STALE_SECONDS`` into a hard per-task wall-clock cap: on k8s a
+    liveness probe already killed at that threshold (the whole pod), but a
+    docker-compose install has no healthcheck, so there it has never been enforced
+    and may sit below a legitimately long batch. Enable it only where
+    ``HEALTH_STALE_SECONDS`` exceeds the longest legitimate task. Enabling it
+    without that knob is a misconfiguration (its 60s code default would kill any
+    task longer than a minute), so it raises.
+    """
+    from queue_backend.pg_queue.consumer import consumer_env
+
+    if not consumer_env("CHILD_WATCHDOG", False, _parse_bool):
+        return None
+    stale: float | None = consumer_env("HEALTH_STALE_SECONDS", None, float)
+    if stale is None:
+        raise ValueError(
+            "WORKER_PG_QUEUE_CONSUMER_CHILD_WATCHDOG is enabled but "
+            "WORKER_PG_QUEUE_CONSUMER_HEALTH_STALE_SECONDS is unset — set it above "
+            "the longest legitimate task"
+        )
+    if not math.isfinite(stale) or stale <= 0:
+        raise ValueError(
+            "WORKER_PG_QUEUE_CONSUMER_HEALTH_STALE_SECONDS must be a finite number "
+            f"> 0, got {stale!r}"
+        )
+    return stale
+
+
 def shutdown_grace_from_env() -> float:
     """Graceful-drain budget (seconds) on shutdown, shared across all children,
     before SIGKILL.
@@ -169,6 +225,9 @@ def __init__(self, concurrency: int) -> None:
         self._last_fork: dict[int, float] = {}
         self._consecutive_crashes: dict[int, int] = {}
         self._restart_due: dict[int, float] = {}  # slot -> monotonic not-before
+        # Written by the main thread, read by the liveness thread for /metrics; an
+        # int read is atomic under the GIL.
+        self.watchdog_kills = 0
 
     @property
     def concurrency(self) -> int:
@@ -264,6 +323,20 @@ def all_loaded(self) -> bool:
         """Readiness verdict source: True once every slot's child has loaded."""
         return self.loaded_count() == self._n
 
+    def is_loaded(self, slot: int) -> bool:
+        self._validate(slot)
+        return bool(self._loaded[slot])
+
+    def fork_age(self, slot: int) -> float:
+        """Seconds since ``slot``'s current child was forked."""
+        self._validate(slot)
+        return time.monotonic() - self._last_fork.get(slot, time.monotonic())
+
+    def slot_age(self, slot: int) -> float:
+        """Seconds since ``slot``'s child last polled, per its published heartbeat."""
+        self._validate(slot)
+        return time.time() - self._heartbeats[slot]
+
     def oldest_age(self) -> float:
         now = time.time()
         return max((now - hb for hb in self._heartbeats), default=0.0)
@@ -297,7 +370,7 @@ def _run_child(slot: int, heartbeats, loaded) -> None:  # noqa: ANN001 (ctypes a
 
     def _publish_heartbeat() -> None:
         # last-poll wall-time = now − (seconds since last poll). Frozen while a
-        # task runs (the consumer stamps its heartbeat at the top of poll_once),
+        # task runs (the consumer stamps its heartbeat before each queue read),
         # so a child stuck on a too-long task goes stale exactly as the single
         # consumer does. Guarded so a transient error (e.g. teardown during
         # shutdown) logs loudly and the loop continues instead of dying silently
@@ -311,6 +384,9 @@ def _publish_heartbeat() -> None:
                 )
             time.sleep(_REPORT_INTERVAL_SECONDS)
 
+    # Publish once before ``loaded`` so the watchdog never pairs a loaded slot with
+    # the previous child's stale heartbeat.
+    heartbeats[slot] = time.time() - consumer.seconds_since_last_poll()
     threading.Thread(target=_publish_heartbeat, daemon=True, name=f"pg-hb-{slot}").start()
     loaded[slot] = 1
     logger.info(
@@ -406,6 +482,64 @@ def _restart_due_children(fleet: _Fleet, stopping: threading.Event) -> None:
         _try_fork_child(fleet, slot)
 
 
+def _kill_stale_children(fleet: _Fleet, stale_after: float, killed: set[int]) -> None:
+    """SIGKILL every child silent for longer than ``stale_after``.
+
+    A *loaded* child is judged by its heartbeat. A child not yet loaded is judged
+    by time since its fork, against the larger of ``stale_after`` and
+    ``_MIN_BOOTSTRAP_BUDGET_SECONDS``: a re-forked slot keeps its predecessor's old
+    heartbeat until the new child bootstraps (record_fork does not reseed it), so
+    the heartbeat would kill every replacement during import, while a child that
+    hangs in ``import worker`` would otherwise never be judged at all. ``killed``
+    holds pids already signalled and not yet reaped, so a child is killed and
+    logged once; the next ``_reap_dead`` reaps it and schedules the re-fork like
+    any other exit.
+    """
+    bootstrap_budget = max(stale_after, _MIN_BOOTSTRAP_BUDGET_SECONDS)
+    for slot, pid in fleet.alive_items():
+        if pid in killed:
+            continue
+        loaded = fleet.is_loaded(slot)
+        age = fleet.slot_age(slot) if loaded else fleet.fork_age(slot)
+        limit = stale_after if loaded else bootstrap_budget
+        if age <= limit:
+            continue
+        logger.error(
+            "PG-queue consumer: child slot=%s pid=%s %s for %.0fs (> %.0fs) — "
+            "presumed hung; SIGKILL so its message redelivers and the slot is "
+            "re-forked",
+            slot,
+            pid,
+            "has not polled" if loaded else "has not finished loading",
+            age,
+            limit,
+        )
+        killed.add(pid)
+        try:
+            os.kill(pid, signal.SIGKILL)
+        except ProcessLookupError:
+            continue  # exited on its own first; the reap handles it, not a kill
+        fleet.watchdog_kills += 1
+
+
+def _monitor_tick(
+    fleet: _Fleet,
+    stopping: threading.Event,
+    watchdog_after: float | None,
+    killed: set[int],
+) -> None:
+    """One supervisor iteration: reap exits, re-fork due slots, then the watchdog.
+
+    ``killed`` is trimmed to live pids right after the reap, so a recycled pid is
+    never mistaken for one already signalled.
+    """
+    _reap_dead(fleet, stopping)
+    killed.intersection_update(pid for _slot, pid in fleet.alive_items())
+    _restart_due_children(fleet, stopping)
+    if watchdog_after is not None and not stopping.is_set():
+        _kill_stale_children(fleet, watchdog_after, killed)
+
+
 def run_supervised(concurrency: int) -> None:
     """Fork ``concurrency`` consumer children and supervise them until SIGTERM."""
     logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
@@ -417,6 +551,16 @@ def run_supervised(concurrency: int) -> None:
         "children)",
         grace_seconds,
     )
+    watchdog_after = child_watchdog_from_env()
+    if watchdog_after is None:
+        logger.info("PG-queue consumer supervisor: child watchdog off")
+    else:
+        logger.info(
+            "PG-queue consumer supervisor: child watchdog kills a child silent for "
+            "> %.0fs",
+            watchdog_after,
+        )
+    killed: set[int] = set()
     stopping = threading.Event()
 
     def _signal_children(sig: int) -> None:
@@ -452,8 +596,7 @@ def _on_term(signum: int, _frame: object) -> None:
     health = _maybe_start_supervisor_health(fleet)
     try:
         while not stopping.is_set():
-            _reap_dead(fleet, stopping)
-            _restart_due_children(fleet, stopping)
+            _monitor_tick(fleet, stopping, watchdog_after, killed)
             stopping.wait(_MONITOR_INTERVAL_SECONDS)  # responsive to SIGTERM
     finally:
         stopping.set()
@@ -558,6 +701,7 @@ def _extra_status() -> dict[str, object]:
         freshness_fn=fleet.freshness,
         alive_children_fn=lambda: float(fleet.alive_count()),
         concurrency_fn=lambda: float(fleet.concurrency),
+        watchdog_kills_fn=lambda: float(fleet.watchdog_kills),
     )
     server = LivenessServer(
         freshness_fn=fleet.freshness,

workers/queue_backend/pg_queue/README.md

@@ -36,6 +36,7 @@ integers.
 | `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` |
+| `CHILD_WATCHDOG` | With `CONCURRENCY > 1`, SIGKILL and re-fork a child silent past `HEALTH_STALE_SECONDS` (requires it set above the longest legitimate task) | `false` |
 | `WORKER_TYPE` | Which source worker's tasks this consumer registers (bootstrap) | — |
 
 ### Reaper (`WORKER_PG_REAPER_*`)
@@ -76,8 +77,9 @@ keeps its claim; a **dead** worker's renewal stops, so its `vt` expires in ~`LEA
 redelivery in **minutes, not hours**. `VT_SECONDS` is retained as the *drain /
 max-runtime bound* (grace, health-stale), and the lease is clamped to it. Note the
 lease is **not** a hard cap on runtime — a live-but-hung task keeps renewing forever;
-the backstop for that is the liveness probe restarting the pod (process death stops
-renewal). The renewal owns its own DB connection (closed on exit) and is best-effort:
+the backstop is process death, which stops renewal: the supervisor's child watchdog
+(`CHILD_WATCHDOG`) kills a child silent past `HEALTH_STALE_SECONDS`; otherwise a
+liveness probe restarting the pod, where one exists. 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

workers/queue_backend/pg_queue/consumer.py

@@ -96,10 +96,11 @@
 # response, and the row is TTL'd.
 _TASK_STATUS_RETENTION_SECONDS = 86400
 # Liveness: a poll loop that hasn't cycled in this many seconds is reported
-# unhealthy. The heartbeat is stamped at the top of each poll_once and frozen
-# during task execution, so this threshold doubles as an UPPER BOUND on a single
-# task's wall-clock: a task running longer than it trips the probe → pod restart
-# → the in-flight task is killed and (at-least-once) redelivered. 60s suits the
+# unhealthy. The heartbeat is stamped before each queue read and frozen during
+# task execution, so this threshold doubles as an UPPER BOUND on a single task's
+# wall-clock: a task running longer than it trips the probe → pod restart (or,
+# under the prefork supervisor with CHILD_WATCHDOG on, a SIGKILL of just that
+# child) → the in-flight task is killed and (at-least-once) redelivered. 60s suits the
 # current sub-second leaf (send_webhook_notification); for longer-running tasks,
 # raise WORKER_PG_QUEUE_CONSUMER_HEALTH_STALE_SECONDS above
 # max(batch_size x worst_case_task_seconds, backoff_max).
@@ -365,8 +366,8 @@ def __init__(
         self._result_backend: PgResultBackend | None = None
         # Heartbeat for the liveness probe: monotonic timestamp of the most
         # recent poll attempt. Seeded at construction so a just-started consumer
-        # reads healthy. Updated at the TOP of poll_once, so a loop wedged on a
-        # long-running task (poll_once not returning) goes stale and is caught —
+        # reads healthy. Updated before EACH queue read in poll_once, so a loop
+        # wedged on a long-running task goes stale and is caught —
         # something pgrep-based --status and the launch-time check cannot see.
         self._last_poll_monotonic = time.monotonic()
 
@@ -379,9 +380,12 @@ def poll_once(self) -> int:
         cycle still counts (so run() doesn't take the empty-queue backoff path
         after a partial failure).
         """
-        self._last_poll_monotonic = time.monotonic()
         total = 0
         for queue_name in self.queue_names:
+            # Per queue, not per cycle: a cycle runs one task from EACH queue, and
+            # HEALTH_STALE is sized for one task, so a per-cycle stamp would let
+            # two legitimate long tasks add up to a false "stale".
+            self._last_poll_monotonic = time.monotonic()
             try:
                 messages = self._client.read(
                     queue_name, vt_seconds=self.lease_seconds, qty=self.batch_size

workers/queue_backend/pg_queue/metrics.py

@@ -87,6 +87,7 @@ def __init__(
         freshness_fn: Callable[[], float],
         alive_children_fn: Callable[[], float] | None = None,
         concurrency_fn: Callable[[], float] | None = None,
+        watchdog_kills_fn: Callable[[], float] | None = None,
     ) -> None:
         super().__init__()
         self._function_gauge(
@@ -107,6 +108,36 @@ def __init__(
                 "Configured child-process concurrency of the supervisor fleet",
                 concurrency_fn,
             )
+        if watchdog_kills_fn is not None:
+            self.registry.register(
+                _FunctionCounter(
+                    "pg_consumer_child_watchdog_kills",
+                    "Children the supervisor SIGKILLed for going silent past "
+                    "HEALTH_STALE_SECONDS",
+                    watchdog_kills_fn,
+                )
+            )
+
+
+class _FunctionCounter:
+    """A counter read from a callable at scrape time (``Counter`` has no
+    ``set_function``). Exposed as ``<name>_total``.
+    """
+
+    def __init__(self, name: str, doc: str, fn: Callable[[], float]) -> None:
+        self._name = name
+        self._doc = doc
+        self._fn = fn
+
+    def describe(self) -> Iterable[Metric]:
+        from prometheus_client.core import CounterMetricFamily
+
+        return [CounterMetricFamily(self._name, self._doc)]
+
+    def collect(self) -> Iterable[Metric]:
+        from prometheus_client.core import CounterMetricFamily
+
+        yield CounterMetricFamily(self._name, self._doc, value=self._fn())
 
 
 @dataclass(frozen=True)