SupplyAgent/supply_infra/pipeline/reconciler.py

from __future__ import annotations

import logging
import os
import signal
import socket
import threading
from datetime import timedelta
from zoneinfo import ZoneInfo

from supply_infra.config import InfraSettings, get_infra_settings
from supply_infra.db.repositories.pipeline_run_repo import PipelineRunRepository
from supply_infra.db.repositories.pipeline_step_run_repo import PipelineStepRunRepository
from supply_infra.db.session import dispose_engine, get_session
from supply_infra.pipeline.dates import CHINA_TIMEZONE, china_now
from supply_infra.pipeline.enums import RunStatus, StepStatus
from supply_infra.pipeline.health import (
    run_alert_level,
    touch_reconciler_heartbeat,
)
from supply_infra.pipeline.run_service import PIPELINE_KEY, submit_pipeline_run

logger = logging.getLogger(__name__)


def reconcile_once() -> dict[str, int]:
    settings = get_infra_settings()
    now = china_now()
    stats = {
        "retries_promoted": 0,
        "expired_steps": 0,
        "deadline_exceeded": 0,
        "previous_run_unblocked": 0,
        "alerts_emitted": 0,
        "missed_runs_submitted": 0,
    }
    with get_session() as session:
        run_repo = PipelineRunRepository(session)
        step_repo = PipelineStepRunRepository(session)
        stats["retries_promoted"] = step_repo.promote_due_retries(now)

        for step in step_repo.list_expired_running(now):
            run = run_repo.get(step.run_id, for_update=True)
            if run is not None and run.status in {
                RunStatus.CANCELLING.value,
                RunStatus.CANCELLED.value,
            }:
                step.status = StepStatus.CANCELLED.value
                step.finished_at = now
                step.lease_owner = None
                step.lease_until = None
                step.error_code = "cancelled"
                step.error_message = "Cancellation completed after worker lease expired"
                if run.status == RunStatus.CANCELLING.value:
                    run_repo.finish(
                        run,
                        status=RunStatus.CANCELLED.value,
                        now=now,
                    )
                stats["expired_steps"] += 1
                continue
            if run is not None and run.status == RunStatus.DEADLINE_EXCEEDED.value:
                step.status = StepStatus.TIMED_OUT.value
                step.finished_at = now
                step.lease_owner = None
                step.lease_until = None
                step.error_code = "deadline_exceeded"
                step.error_message = "Worker lease expired after pipeline deadline"
                stats["expired_steps"] += 1
                continue
            step.status = StepStatus.TIMED_OUT.value
            step.finished_at = now
            step.lease_owner = None
            step.lease_until = None
            step.error_code = "lease_expired"
            step.error_message = "Worker heartbeat lease expired"
            if step.retryable and step.attempt < step.max_attempts:
                step_repo.create_retry_attempt(step, status=StepStatus.READY.value)
                if run is not None:
                    run.status = RunStatus.QUEUED.value
                    run.current_step = step.step_key
                    run.lease_owner = None
                    run.lease_until = None
            else:
                step_repo.block_unfinished_downstream(
                    run_id=step.run_id,
                    after_order=step.step_order,
                    reason="Worker heartbeat lease expired",
                )
                if run is not None:
                    run_repo.finish(
                        run,
                        status=RunStatus.FAILED.value,
                        now=now,
                        error_code="lease_expired",
                        error_message=f"Lease expired at step {step.step_key}",
                    )
            stats["expired_steps"] += 1

        for run in run_repo.list_past_deadline(now):
            for step in step_repo.list_for_run(run.run_id):
                if step.status == StepStatus.RUNNING.value:
                    step.error_code = "deadline_exceeded"
                    step.error_message = (
                        "Pipeline exceeded next scheduled start; waiting for "
                        "the worker or lease expiry"
                    )
                    continue
                if step.status in {
                    StepStatus.PENDING.value,
                    StepStatus.READY.value,
                    StepStatus.RETRY_WAIT.value,
                    StepStatus.BLOCKED.value,
                }:
                    step.status = StepStatus.BLOCKED.value
                    step.error_code = "deadline_exceeded"
                    step.error_message = "Pipeline exceeded next scheduled start"
            run_repo.finish(
                run,
                status=RunStatus.DEADLINE_EXCEEDED.value,
                now=now,
                error_code="deadline_exceeded",
                error_message="Pipeline exceeded next scheduled start",
            )
            stats["deadline_exceeded"] += 1

        for run in run_repo.list_blocked_previous():
            blocked_by = (run.summary_json or {}).get("blocked_by_run_id")
            previous = run_repo.get(str(blocked_by)) if blocked_by else None
            if previous is not None and previous.status in {
                RunStatus.QUEUED.value,
                RunStatus.RUNNING.value,
                RunStatus.CANCELLING.value,
                RunStatus.BLOCKED_PREVIOUS_RUN.value,
            }:
                continue
            if previous is not None and step_repo.has_running_for_run(previous.run_id):
                continue
            latest = step_repo.latest_attempts(run.run_id)
            first = min(latest.values(), key=lambda item: item.step_order, default=None)
            if first is None:
                continue
            first.status = StepStatus.READY.value
            first.error_code = None
            first.error_message = None
            for step in latest.values():
                if step.step_run_id != first.step_run_id:
                    step.status = StepStatus.PENDING.value
                    step.error_code = None
                    step.error_message = None
            run.status = RunStatus.QUEUED.value
            run.summary_json = None
            stats["previous_run_unblocked"] += 1

        for run in run_repo.list_nonterminal():
            alert_level = run_alert_level(
                status=run.status,
                started_at=run.started_at,
                created_at=run.created_at,
                deadline_at=run.deadline_at,
                now=now,
                settings=settings,
            )
            if alert_level is not None:
                logger.warning(
                    "pipeline_alert level=%s run_id=%s biz_dt=%s status=%s deadline=%s",
                    alert_level,
                    run.run_id,
                    run.biz_dt,
                    run.status,
                    run.deadline_at.isoformat() if run.deadline_at else None,
                )
                stats["alerts_emitted"] += 1

    local_now = now.replace(tzinfo=CHINA_TIMEZONE).astimezone(
        ZoneInfo(settings.scheduler_timezone)
    )
    missed_cutoff = local_now.replace(
        hour=settings.scheduler_cron_hour,
        minute=settings.scheduler_cron_minute,
        second=0,
        microsecond=0,
    ) + timedelta(seconds=settings.pipeline_missed_run_grace_seconds)
    if settings.scheduler_enabled and local_now >= missed_cutoff:
        biz_dt = local_now.strftime("%Y%m%d")
        with get_session() as session:
            existing = PipelineRunRepository(session).get_for_business_day(
                pipeline_key=PIPELINE_KEY,
                biz_dt=biz_dt,
            )
        if existing is None:
            submission = submit_pipeline_run(
                biz_dt=biz_dt,
                trigger_type="reconcile",
                trigger_source="missed_run_reconciler",
                trigger_reason="Scheduled batch was absent after the grace window",
                settings=settings,
            )
            if submission.created:
                logger.error(
                    "pipeline_alert level=critical type=missed_run_recovered "
                    "run_id=%s biz_dt=%s",
                    submission.run_id,
                    biz_dt,
                )
                stats["missed_runs_submitted"] = 1
    return stats


class PipelineReconciler:
    def __init__(self, settings: InfraSettings | None = None) -> None:
        self.settings = settings or get_infra_settings()
        self.owner = f"{socket.gethostname()}:{os.getpid()}"
        self._stop = threading.Event()

    def stop(self, *_args: object) -> None:
        self._stop.set()

    def run_forever(self) -> None:
        signal.signal(signal.SIGTERM, self.stop)
        signal.signal(signal.SIGINT, self.stop)
        logger.info("Pipeline reconciler started")
        try:
            while not self._stop.is_set():
                try:
                    result = reconcile_once()
                    touch_reconciler_heartbeat(
                        owner=self.owner,
                        settings=self.settings,
                    )
                    logger.info("Pipeline reconcile result: %s", result)
                except Exception:
                    logger.exception("Pipeline reconcile failed")
                self._stop.wait(self.settings.pipeline_reconcile_seconds)
        finally:
            dispose_engine()
            logger.info("Pipeline reconciler stopped")