254 lines
8.1 KiB
Python
254 lines
8.1 KiB
Python
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])
|
|
|
|
|
|
def _apply_output_policy(payload: dict) -> dict:
|
|
"""PH-020: enforce a single notifier output policy.
|
|
|
|
Per-notifier code keeps its own format-specific shortening (e.g. Telegram
|
|
300 chars), but a hard cap and include-flag apply consistently before any
|
|
delivery. Truncation uses a clear marker so receivers can see it happened.
|
|
"""
|
|
settings = get_settings()
|
|
out = dict(payload)
|
|
text_value = out.get("output")
|
|
if text_value is None:
|
|
return out
|
|
if not settings.notifier_include_output:
|
|
out["output"] = None
|
|
out["output_truncated"] = True
|
|
return out
|
|
max_bytes = max(settings.notifier_max_output_bytes, 0)
|
|
encoded = text_value.encode("utf-8")
|
|
if len(encoded) <= max_bytes:
|
|
return out
|
|
# Reserve room for the marker so the final string never exceeds max_bytes.
|
|
# If max_bytes is smaller than even the marker, drop the body entirely —
|
|
# this keeps the absolute cap a hard guarantee under any setting.
|
|
dropped = len(encoded) - max_bytes
|
|
marker = f"...[truncated {dropped} bytes]"
|
|
marker_bytes = len(marker.encode("utf-8"))
|
|
if marker_bytes >= max_bytes:
|
|
out["output"] = ""
|
|
out["output_truncated"] = True
|
|
return out
|
|
prefix_budget = max_bytes - marker_bytes
|
|
truncated = encoded[:prefix_budget].decode("utf-8", errors="ignore")
|
|
out["output"] = f"{truncated}{marker}"
|
|
out["output_truncated"] = True
|
|
return out
|
|
|
|
|
|
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
|
|
|
|
# PH-020: apply central output policy before handing to any notifier.
|
|
delivery_payload = _apply_output_policy(row["payload"])
|
|
try:
|
|
result = await notifier.deliver(row["event_type"], delivery_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
|
|
|
|
# Bound concurrent DB sessions: each _process_row opens its own session, so an
|
|
# unbounded gather over a large batch can exhaust the connection pool.
|
|
sem = asyncio.Semaphore(settings.notifier_concurrency)
|
|
|
|
async def _guarded(row: dict) -> None:
|
|
async with sem:
|
|
await _process_row(sm, registry, row, settings.notifier_max_attempts)
|
|
|
|
await asyncio.gather(*(_guarded(r) 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()
|