monlet/server/monlet_server/services/notifier_worker.py

from __future__ import annotations

import asyncio
from datetime import UTC, datetime, timedelta

import httpx
from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker

from .. import metrics
from ..logging_config import get_logger
from ..redaction import redact
from ..settings import Settings, get_settings
from .notifiers import Notifier, build_registry

log = get_logger("monlet.notifier_worker")


# Per ADR-0005: 5s, 15s, 60s, 5m, 30m, 1h, 2h, 4h.
_BACKOFF_SCHEDULE_SEC = (5, 15, 60, 300, 1800, 3600, 7200, 14400)


def _backoff(attempts: int) -> timedelta:
    idx = max(0, min(attempts - 1, len(_BACKOFF_SCHEDULE_SEC) - 1))
    return timedelta(seconds=_BACKOFF_SCHEDULE_SEC[idx])


async def _claim_batch(session, limit: int, lease_sec: int) -> list[dict]:
    res = await session.execute(
        text(
            """
            UPDATE notification_outbox
               SET state = 'sending', updated_at = now()
             WHERE id IN (
                SELECT id FROM notification_outbox
                 WHERE (
                        state IN ('pending','retry')
                        AND (next_attempt_at IS NULL OR next_attempt_at <= now())
                       )
                    OR (
                        state = 'sending'
                        AND updated_at < now() - make_interval(secs => :lease)
                       )
                 ORDER BY next_attempt_at NULLS FIRST
                 LIMIT :lim
                 FOR UPDATE SKIP LOCKED
             )
         RETURNING id, notifier, event_type, attempts, payload
            """
        ),
        {"lim": limit, "lease": lease_sec},
    )
    rows = res.mappings().all()
    return [dict(r) for r in rows]


async def _finalize(
    session,
    row_id,
    *,
    state: str,
    attempts: int,
    next_attempt_at: datetime | None,
    last_error: str | None,
) -> None:
    await session.execute(
        text(
            """
            UPDATE notification_outbox
               SET state = :state,
                   attempts = :attempts,
                   next_attempt_at = :next_attempt_at,
                   last_error = :last_error,
                   updated_at = now()
             WHERE id = :id
            """
        ),
        {
            "state": state,
            "attempts": attempts,
            "next_attempt_at": next_attempt_at,
            "last_error": last_error,
            "id": row_id,
        },
    )


async def _process_row(
    sm: async_sessionmaker,
    registry: dict[str, Notifier],
    row: dict,
    max_attempts: int,
) -> None:
    notifier = registry.get(row["notifier"])
    attempts = int(row["attempts"]) + 1
    async with sm() as session:
        if notifier is None:
            await _finalize(
                session,
                row["id"],
                state="discarded",
                attempts=attempts,
                next_attempt_at=None,
                last_error=f"unknown_notifier:{row['notifier']}",
            )
            await session.commit()
            metrics.notifications_total.labels(notifier=row["notifier"], result="discarded").inc()
            return

        try:
            result = await notifier.deliver(row["event_type"], row["payload"])
        except Exception as exc:  # defensive: notifier should not raise
            result_ok = False
            result_perm = False
            result_err = f"exception:{type(exc).__name__}"
        else:
            result_ok = result.ok
            result_perm = result.permanent
            result_err = redact(result.error) if result.error else None

        if result_ok:
            await _finalize(
                session,
                row["id"],
                state="sent",
                attempts=attempts,
                next_attempt_at=None,
                last_error=None,
            )
            await session.commit()
            metrics.notifications_total.labels(notifier=notifier.name, result="sent").inc()
            return

        if result_perm or attempts >= max_attempts:
            await _finalize(
                session,
                row["id"],
                state="failed",
                attempts=attempts,
                next_attempt_at=None,
                last_error=result_err,
            )
            await session.commit()
            metrics.notifications_total.labels(notifier=notifier.name, result="failed").inc()
            return

        next_at = datetime.now(UTC) + _backoff(attempts)
        await _finalize(
            session,
            row["id"],
            state="retry",
            attempts=attempts,
            next_attempt_at=next_at,
            last_error=result_err,
        )
        await session.commit()
        metrics.notifications_total.labels(notifier=notifier.name, result="retry").inc()


async def tick(
    sm: async_sessionmaker,
    registry: dict[str, Notifier],
    settings: Settings,
) -> int:
    async with sm() as session:
        rows = await _claim_batch(
            session, settings.notifier_batch_size, settings.notifier_lease_sec
        )
        await session.commit()
    if not rows:
        return 0
    await asyncio.gather(
        *(_process_row(sm, registry, r, settings.notifier_max_attempts) for r in rows)
    )
    return len(rows)


async def run_notifier_worker(
    sm: async_sessionmaker,
    stop_event: asyncio.Event,
    client: httpx.AsyncClient | None = None,
) -> None:
    settings = get_settings()
    own_client = False
    if client is None:
        client = httpx.AsyncClient(timeout=settings.notifier_http_timeout_sec)
        own_client = True
    registry = build_registry(settings, client)
    if not registry:
        log.info("notifier_worker_disabled_no_notifiers")
        if own_client:
            await client.aclose()
        return
    log.info("notifier_worker_started", notifiers=list(registry.keys()))
    try:
        while not stop_event.is_set():
            try:
                await tick(sm, registry, settings)
            except Exception as exc:
                log.warning("notifier_tick_failed", error=str(exc))
            try:
                await asyncio.wait_for(stop_event.wait(), timeout=settings.notifier_tick_sec)
            except TimeoutError:
                continue
    finally:
        if own_client:
            await client.aclose()