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 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."""
    from steward.core.alerts import MetricBatch

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

    async with app.db_sessionmaker() as session:
        async with session.begin():
            batch = MetricBatch()
            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
                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)