cleveragents-core/tools/controller/worker/heartbeat.py

"""Worker heartbeat thread.

Per plan v9 (simplified): TTL-only heartbeat. The thread updates
``workflow_attempts.lock_heartbeat_at`` every ``HEARTBEAT_INTERVAL_S``
(default 30s). If the UPDATE returns rowcount=0, the lock has been
reaped (TTL exceeded; reaper reset the row); the heartbeat thread
raises ``HeartbeatStopped`` via an event the worker thread polls.

Activity tracking (v6 enhancement) is intentionally deferred per
v9 simplification: TTL alone catches dead workers within
``lock_ttl_seconds`` (default 600s). Re-add activity tracking only
if real hangs are observed.
"""
from __future__ import annotations

import logging
import os
import threading
import time
from datetime import datetime, timezone

from sqlalchemy import text
from sqlalchemy.engine import Engine

logger = logging.getLogger(__name__)


HEARTBEAT_INTERVAL_S = int(os.environ.get("CONTROLLER_HEARTBEAT_INTERVAL_S", "30"))


class HeartbeatStopped(RuntimeError):
    """Raised when the heartbeat thread detects it has lost the lock.

    Worker code polls ``Heartbeat.lost_lock_event`` and aborts cleanly
    when set (no output write — the reaper has already re-pended
    the attempt; another worker will retry).
    """


class Heartbeat:
    """Per-attempt heartbeat thread.

    Usage:
        hb = Heartbeat(engine, attempt_id, instance_id)
        hb.start()
        try:
            # … long-running work; periodically check hb.lost_lock_event
            if hb.lost_lock_event.is_set():
                raise HeartbeatStopped("lock reaped")
        finally:
            hb.stop()
    """

    def __init__(
        self,
        engine: Engine,
        attempt_id: int,
        instance_id: str,
        *,
        interval_s: int = HEARTBEAT_INTERVAL_S,
    ) -> None:
        self.engine = engine
        self.attempt_id = attempt_id
        self.instance_id = instance_id
        self.interval_s = interval_s
        self.lost_lock_event = threading.Event()
        self._stop_event = threading.Event()
        self._thread: threading.Thread | None = None
        self._last_heartbeat_at: datetime | None = None

    def start(self) -> None:
        if self._thread is not None:
            raise RuntimeError("heartbeat already started")
        self._thread = threading.Thread(
            target=self._loop,
            name=f"hb-{self.attempt_id}",
            daemon=True,
        )
        self._thread.start()

    def stop(self, *, timeout: float = 5.0) -> None:
        self._stop_event.set()
        if self._thread is not None:
            self._thread.join(timeout=timeout)

    def last_heartbeat_at(self) -> datetime | None:
        """Last successful heartbeat write (UTC). None until first
        UPDATE succeeds."""
        return self._last_heartbeat_at

    def _loop(self) -> None:
        # Fire immediately so the lock_heartbeat_at column reflects
        # the worker's activity before the first interval lapses.
        if not self._tick():
            return
        while not self._stop_event.wait(self.interval_s):
            if not self._tick():
                return

    def _tick(self) -> bool:
        """Update lock_heartbeat_at. Returns False if lock was lost
        (worker should abort)."""
        now = datetime.now(timezone.utc)
        try:
            with self.engine.connect() as conn:
                result = conn.execute(
                    text(
                        "UPDATE workflow_attempts "
                        "SET lock_heartbeat_at = :now "
                        "WHERE attempt_id = :attempt_id "
                        "  AND locked_by_instance = :instance"
                    ),
                    {
                        "now": now,
                        "attempt_id": self.attempt_id,
                        "instance": self.instance_id,
                    },
                )
                conn.commit()
                rowcount = result.rowcount
        except Exception as exc:  # noqa: BLE001 — best-effort heartbeat
            # Transient DB error. Log + retry next tick. The lock_ttl
            # eventually trips if these stack up.
            logger.warning(
                "heartbeat UPDATE failed for attempt_id=%s: %s",
                self.attempt_id, exc,
            )
            return True

        if rowcount == 0:
            logger.warning(
                "heartbeat lost lock for attempt_id=%s instance=%s — "
                "reaper has re-pended; worker should abort",
                self.attempt_id, self.instance_id,
            )
            self.lost_lock_event.set()
            return False

        self._last_heartbeat_at = now
        return True