SupplyAgent/supply_infra/pipeline/health.py

from __future__ import annotations

from datetime import datetime, timedelta
from typing import Any

from supply_infra.config import InfraSettings, get_infra_settings
from supply_infra.db.repositories.pipeline_lock_repo import PipelineLockRepository
from supply_infra.db.repositories.pipeline_run_repo import PipelineRunRepository
from supply_infra.db.session import get_session
from supply_infra.pipeline.dates import china_now

SCHEDULER_HEARTBEAT_KEY = "pipeline:scheduler"
RECONCILER_HEARTBEAT_KEY = "pipeline:reconciler"
WORKER_HEARTBEAT_PREFIX = "pipeline:worker:"
_ACTIVE_STATUSES = {
    "queued",
    "running",
    "cancelling",
    "blocked_previous_run",
}


def touch_scheduler_heartbeat(
    *,
    owner: str,
    settings: InfraSettings | None = None,
) -> None:
    active = settings or get_infra_settings()
    with get_session() as session:
        PipelineLockRepository(session).touch(
            lock_key=SCHEDULER_HEARTBEAT_KEY,
            owner=owner,
            lease_seconds=max(
                active.pipeline_lease_seconds,
                active.pipeline_heartbeat_seconds * 3,
            ),
            now=china_now(),
        )


def touch_worker_heartbeat(
    *,
    owner: str,
    settings: InfraSettings | None = None,
) -> None:
    active = settings or get_infra_settings()
    with get_session() as session:
        PipelineLockRepository(session).touch(
            lock_key=f"{WORKER_HEARTBEAT_PREFIX}{owner}",
            owner=owner,
            lease_seconds=max(
                active.pipeline_lease_seconds,
                active.pipeline_heartbeat_seconds * 3,
            ),
            now=china_now(),
        )


def touch_reconciler_heartbeat(
    *,
    owner: str,
    settings: InfraSettings | None = None,
) -> None:
    active = settings or get_infra_settings()
    with get_session() as session:
        PipelineLockRepository(session).touch(
            lock_key=RECONCILER_HEARTBEAT_KEY,
            owner=owner,
            lease_seconds=max(
                active.pipeline_lease_seconds,
                active.pipeline_reconcile_seconds * 3,
            ),
            now=china_now(),
        )


def run_alert_level(
    *,
    status: str,
    started_at: datetime | None,
    created_at: datetime,
    deadline_at: datetime | None,
    now: datetime,
    settings: InfraSettings,
) -> str | None:
    if status not in _ACTIVE_STATUSES:
        return None
    if deadline_at is not None:
        remaining = deadline_at - now
        if remaining <= timedelta(minutes=settings.pipeline_final_warn_before_next_minutes):
            return "final_warning"
        if remaining <= timedelta(minutes=settings.pipeline_critical_before_next_minutes):
            return "critical"
    baseline = started_at or created_at
    if now - baseline >= timedelta(hours=settings.pipeline_warn_after_hours):
        return "warning"
    return None


def runtime_components_ready(
    *,
    scheduler_enabled: bool,
    scheduler_running: bool,
    active_worker_count: int,
    expected_worker_count: int,
    reconciler_running: bool,
) -> bool:
    return bool(
        (not scheduler_enabled or scheduler_running)
        and active_worker_count >= expected_worker_count
        and reconciler_running
    )


def pipeline_health_snapshot(
    settings: InfraSettings | None = None,
) -> dict[str, Any]:
    active = settings or get_infra_settings()
    now = china_now()
    with get_session() as session:
        lock_repo = PipelineLockRepository(session)
        scheduler_lock = lock_repo.get(SCHEDULER_HEARTBEAT_KEY)
        reconciler_lock = lock_repo.get(RECONCILER_HEARTBEAT_KEY)
        active_worker_count = lock_repo.count_active_with_prefix(
            WORKER_HEARTBEAT_PREFIX,
            now=now,
        )
        latest = PipelineRunRepository(session).list_recent(limit=1)
        latest_run = latest[0] if latest else None
        scheduler_heartbeat_at = (
            scheduler_lock.heartbeat_at if scheduler_lock is not None else None
        )
        scheduler_lease_until = (
            scheduler_lock.lease_until if scheduler_lock is not None else None
        )
        reconciler_lease_until = (
            reconciler_lock.lease_until if reconciler_lock is not None else None
        )
        latest_snapshot = (
            {
                "run_id": latest_run.run_id,
                "status": latest_run.status,
                "started_at": latest_run.started_at,
                "created_at": latest_run.created_at,
                "deadline_at": latest_run.deadline_at,
            }
            if latest_run is not None
            else None
        )

    scheduler_heartbeat = (
        scheduler_heartbeat_at.isoformat()
        if scheduler_heartbeat_at is not None
        else None
    )
    scheduler_running = bool(
        active.scheduler_enabled
        and scheduler_lease_until is not None
        and scheduler_lease_until > now
    )
    reconciler_running = bool(
        reconciler_lease_until is not None and reconciler_lease_until > now
    )
    workers_running = active_worker_count >= active.pipeline_worker_processes
    alert_level = None
    if latest_snapshot is not None:
        alert_level = run_alert_level(
            status=latest_snapshot["status"],
            started_at=latest_snapshot["started_at"],
            created_at=latest_snapshot["created_at"],
            deadline_at=latest_snapshot["deadline_at"],
            now=now,
            settings=active,
        )
    return {
        "scheduler_enabled": active.scheduler_enabled,
        "scheduler_running": scheduler_running,
        "scheduler_heartbeat_at": scheduler_heartbeat,
        "worker_expected_count": active.pipeline_worker_processes,
        "worker_active_count": active_worker_count,
        "workers_running": workers_running,
        "reconciler_running": reconciler_running,
        "runtime_components_ready": runtime_components_ready(
            scheduler_enabled=active.scheduler_enabled,
            scheduler_running=scheduler_running,
            active_worker_count=active_worker_count,
            expected_worker_count=active.pipeline_worker_processes,
            reconciler_running=reconciler_running,
        ),
        "latest_run_id": latest_snapshot["run_id"] if latest_snapshot else None,
        "latest_run_status": latest_snapshot["status"] if latest_snapshot else None,
        "latest_run_alert_level": alert_level,
        "checked_at": now.isoformat(),
    }