drew
0bc734c020
Applies `ruff format` to the accumulated formatting debt on this branch. Formatting-only — no behavioral changes. Required for CI/lint's format gate (`nox -s format -- --check`), which the branch was failing on 288 tracked files that drifted from ruff's canonical style. In-progress WIP files are intentionally excluded so this commit stays a clean formatting-only diff. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
144 lines
4.6 KiB
Python
144 lines
4.6 KiB
Python
"""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
|