FabledSteward/plugins/prometheus/scheduler.py

# plugins/prometheus/scheduler.py
from __future__ import annotations

import asyncio
import logging
import time
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from sqlalchemy import select

from steward.core.counter_rates import CounterRates
from steward.core.scheduler import ScheduledTask

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

# How often the scheduler looks for due targets. A target is scraped within
# one tick of its interval, so this is the cadence's worst-case lateness.
TICK_SECONDS = 5

# A scrape past this is given up, whatever the client's own timeouts did:
# a slow drip of bytes can keep each read inside its timeout for ever.
SCRAPE_DEADLINE_SECONDS = 30

# Counters' previous samples, for their rates.
RATES = CounterRates()


def make_scrape_task(app: "Quart") -> ScheduledTask:
    async def tick() -> None:
        await scrape_due(app)

    return ScheduledTask(name="prometheus_scrape", coro_factory=tick,
                         interval_seconds=TICK_SECONDS, run_on_startup=True)


def is_due(target, now: datetime) -> bool:
    """Whether a target's interval has elapsed since its last scrape, with
    half a tick of slack so tick jitter doesn't push it back a whole tick."""
    if not target.enabled:
        return False
    if target.last_scraped_at is None:
        return True
    elapsed = (now - target.last_scraped_at).total_seconds()
    return elapsed >= target.interval_seconds - TICK_SECONDS / 2


async def scrape_due(app: "Quart", scrape=None) -> None:
    """Scrape every due target at once, then write each one's outcome to its
    row and its series to the metric tables. `scrape` stands in for
    scraper.scrape in tests."""
    from steward.core.crypto import decrypt_secret

    from .models import PrometheusTarget
    from .scraper import ScrapeError
    from .scraper import scrape as fetch

    scrape = scrape or fetch
    now = datetime.now(timezone.utc)
    async with app.db_sessionmaker() as session:
        targets = [t for t in (await session.execute(select(PrometheusTarget))).scalars().all()
                   if is_due(t, now)]
    if not targets:
        return

    async def one(t):
        started = time.monotonic()
        token = decrypt_secret(t.token, context=f"prometheus target {t.name} token")
        try:
            families = await asyncio.wait_for(scrape(t.url, token, t.verify_tls),
                                              SCRAPE_DEADLINE_SECONDS)
            return families, (time.monotonic() - started) * 1000
        except TimeoutError:
            return ScrapeError(f"Gave up after {SCRAPE_DEADLINE_SECONDS} s."), \
                (time.monotonic() - started) * 1000
        except Exception as exc:  # noqa: BLE001 — every failure is the target's outcome
            return exc, (time.monotonic() - started) * 1000

    results = await asyncio.gather(*(one(t) for t in targets))
    await record(app, targets, results, datetime.now(timezone.utc))


async def forget_gone(session, target: str, kept: set[tuple[str, str]]) -> None:
    """Drop the newest value of each series of `target` that its last scrape
    didn't keep: the endpoint stopped reporting it (an old build_info after an
    upgrade, a deleted job), or a filter or the cap now leaves it out. Pages,
    widgets and query_metrics read the newest values, so they show what the
    endpoint says now. History stays for charts and ages out with retention.
    Alert rules keep their own last-sample time, so one on a gone series
    still goes no_data."""
    from sqlalchemy import delete, tuple_

    from steward.models.metrics import PluginMetricLatest

    from .service import SOURCE

    L = PluginMetricLatest
    stored = (await session.execute(select(L.resource_name, L.metric_name).where(
        L.source_module == SOURCE,
        (L.resource_name == target) | L.resource_name.startswith(target + ":", autoescape=True)))).all()
    gone = [(r, m) for r, m in stored if (r, m) not in kept]
    if gone:
        await session.execute(delete(L).where(
            L.source_module == SOURCE, tuple_(L.resource_name, L.metric_name).in_(gone)))


async def record(app: "Quart", targets, results, now: datetime) -> None:
    """Each target's outcome on its row, and its kept series (counters as
    rates) plus its own is_up / scrape_ms / series counts as metrics. After a
    good scrape, the series it no longer has are forgotten (forget_gone); a
    failed scrape says nothing about them, so they stay."""
    from steward.core.alerts import MetricBatch

    from .models import PrometheusTarget
    from .scraper import ScrapeError
    from .service import (IS_UP, OWN_METRICS, SCRAPE_MS, SERIES_CAPPED, SERIES_KEPT, SOURCE,
                          map_families)

    async with app.db_sessionmaker() as session:
        async with session.begin():
            batch = MetricBatch()
            answered: list[tuple[str, set[tuple[str, str]]]] = []
            for target, (result, ms) in zip(targets, results):
                row = await session.get(PrometheusTarget, target.id)
                if row is None:          # deleted while it was being scraped
                    continue
                row.last_scraped_at = now
                if isinstance(result, BaseException):
                    if not isinstance(result, ScrapeError):
                        logger.warning("Prometheus scrape of %s failed: %r", target.name, result)
                    row.last_error = str(result) or "The scrape failed."
                    batch.add(source_module=SOURCE, resource_name=row.name, metric_name=IS_UP,
                              value=0.0, recorded_at=now)
                    continue
                mapped = map_families(row.name, result, list(row.keep or []), list(row.drop or []),
                                      row.series_cap)
                row.last_error = None
                row.last_ok_at = now
                row.last_kept = len(mapped.series)
                row.last_dropped = mapped.dropped
                row.last_capped = mapped.capped
                # A counter's first scrape gives no rate yet; it is still kept.
                answered.append((row.name, {(x.resource, x.metric) for x in mapped.series}
                                 | {(row.name, m) for m in OWN_METRICS}))
                for series in mapped.series:
                    value = series.value
                    if series.counter:
                        value = RATES.rate((row.id, series.resource, series.metric), value, now)
                        if value is None:
                            continue
                    batch.add(source_module=SOURCE, resource_name=series.resource,
                              metric_name=series.metric, value=value, recorded_at=now)
                for metric_name, value in [(IS_UP, 1.0), (SCRAPE_MS, ms),
                                           (SERIES_KEPT, len(mapped.series)),
                                           (SERIES_CAPPED, mapped.capped)]:
                    batch.add(source_module=SOURCE, resource_name=row.name,
                              metric_name=metric_name, value=value, recorded_at=now)
            await batch.record(session)
            for name, kept in answered:
                await forget_gone(session, name, kept)