cleveragents-core/tools/controller/reaper.py

"""Master-side reaper: reset stale-heartbeat workflow_attempts to
pending so another worker can re-acquire.

Per plan v9 simplified: TTL-only reaper. No JOIN-with-workflows
superseded check (deferred — orphan re-attempts no-op cheaply against
terminal workflows when the next worker tries to advance them).

Runs on the master's tick loop every ``CONTROLLER_REAPER_INTERVAL_S``
(default 60s). The reaper:
- Selects ``workflow_attempts WHERE status='in_progress' AND
  lock_heartbeat_at < NOW() - lock_ttl_seconds``
- For each: status → 'pending'; clear locked_by_instance / locked_at /
  lock_heartbeat_at; insert a ``controller_events`` row with
  reason='lock-ttl-expired'.

The reset BUMPS ``pickup_count`` (each reap = one failed-pickup).
The master's separate pickup-guard catches
``pickup_count >= MAX_PICKUPS`` and STUCKs the workflow. Bumping on
reap (rather than on dequeue) means a healthy worker that simply
acquires an attempt doesn't burn a pickup — only crashes /
stale-heartbeat resets do.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone

from sqlalchemy import text
from sqlalchemy.engine import Engine

from .db.session import session_scope

logger = logging.getLogger(__name__)


@dataclass
class ReaperReport:
    """Per-sweep summary; emitted to structured logs + Prometheus."""

    rows_reaped: int = 0
    reaped_attempts: list[tuple[int, str | None]] = field(default_factory=list)
    # Each entry: (attempt_id, locked_by_instance) — useful for
    # operator forensics ("which machine's worker died?")


def reap_stale_attempts(engine: Engine) -> ReaperReport:
    """Reset stale-heartbeat in_progress attempts to pending.

    The TTL is per-row (``lock_ttl_seconds``) — different roles can
    have different TTLs (estimator ~180s; reviewer ~720s; tier-2
    implementer ~2160s). The WHERE clause uses the per-row column.
    """
    report = ReaperReport()
    now = datetime.now(timezone.utc)
    # Two-step approach for portability:
    # 1. SELECT the candidates (so we can log them).
    # 2. UPDATE them.
    # We could do this in one statement on Postgres (RETURNING) and
    # SQLite (RETURNING is supported in 3.35+). Splitting for
    # observability — the reaper isn't a hot path.
    with session_scope(engine) as session:
        dialect = session.bind.dialect.name if session.bind else "sqlite"
        if dialect == "postgresql":
            select_sql = text(
                "SELECT attempt_id, locked_by_instance, lock_heartbeat_at "
                "FROM workflow_attempts "
                "WHERE status = 'in_progress' "
                "  AND lock_heartbeat_at IS NOT NULL "
                "  AND lock_heartbeat_at + (lock_ttl_seconds || ' seconds')::interval < :now"
            )
        else:
            # SQLite: TIMESTAMP arithmetic via julianday().
            select_sql = text(
                "SELECT attempt_id, locked_by_instance, lock_heartbeat_at "
                "FROM workflow_attempts "
                "WHERE status = 'in_progress' "
                "  AND lock_heartbeat_at IS NOT NULL "
                "  AND (julianday(:now) - julianday(lock_heartbeat_at)) "
                "      * 86400 > lock_ttl_seconds"
            )
        rows = session.execute(select_sql, {"now": now}).all()
        if not rows:
            return report

        # 2. UPDATE — use IN-clause with the candidate ids.
        ids = [r.attempt_id for r in rows]
        # SQLAlchemy's expanding bindparam handles IN over a list.
        from sqlalchemy import bindparam
        update_sql = text(
            "UPDATE workflow_attempts SET "
            "  status = 'pending', "
            "  locked_by_instance = NULL, "
            "  locked_at = NULL, "
            "  lock_heartbeat_at = NULL, "
            "  pickup_count = pickup_count + 1 "
            "WHERE attempt_id IN :ids"
        ).bindparams(bindparam("ids", expanding=True))
        result = session.execute(update_sql, {"ids": ids})
        report.rows_reaped = result.rowcount

        # 3. Log + record controller_events.
        events_sql = text(
            "INSERT INTO controller_events "
            "(workflow_id, ts, event_type, payload, forgejo_write_pending, replay_attempts) "
            "VALUES ("
            "  (SELECT workflow_id FROM workflow_attempts WHERE attempt_id = :attempt_id), "
            "  :ts, 'lock-ttl-expired', :payload, 0, 0"
            ")"
        )
        for r in rows:
            report.reaped_attempts.append(
                (r.attempt_id, r.locked_by_instance)
            )
            session.execute(events_sql, {
                "attempt_id": r.attempt_id,
                "ts": now,
                "payload": _payload_for_event(
                    r.attempt_id, r.locked_by_instance,
                    r.lock_heartbeat_at, now,
                ),
            })

    if report.rows_reaped:
        logger.warning(
            "reaper reset %d stale in_progress attempt(s): %s",
            report.rows_reaped,
            [(aid, inst) for (aid, inst) in report.reaped_attempts],
        )
    return report


def _payload_for_event(
    attempt_id: int,
    locked_by_instance: str | None,
    heartbeat_at: datetime | None,
    now: datetime,
) -> str:
    """Build a JSON string payload for the controller_events row."""
    import json
    age_s: float | None = None
    if heartbeat_at is not None:
        try:
            # heartbeat_at may be naive depending on the DB; normalise.
            if heartbeat_at.tzinfo is None:
                from datetime import timezone as _tz
                hb = heartbeat_at.replace(tzinfo=_tz.utc)
            else:
                hb = heartbeat_at
            age_s = (now - hb).total_seconds()
        except Exception:
            age_s = None
    return json.dumps({
        "reaped_attempt_id": attempt_id,
        "previously_locked_by": locked_by_instance,
        "heartbeat_age_s": age_s,
    })


__all__ = ["ReaperReport", "reap_stale_attempts"]