svcforge/scripts/redis_budget.py

#!/usr/bin/env python3
"""Project month-end Redis command burn from the live counter. Exit 1 if it blows the budget.

    $ python3 scripts/redis_budget.py
    $ python3 scripts/redis_budget.py --url http://localhost:8000/metrics --budget 500000
    $ python3 scripts/redis_budget.py --url http://api:8000/metrics \
                                      --url http://worker:9100/metrics \
                                      --url http://reconciler:9100/metrics

The budget is per Upstash database; the counters are per process. api, worker and
reconciler each keep their own registry (see obs.py — one process per pod, no multiproc
directory), so scraping one endpoint measures one third of the burn. Repeat `--url` to sum
them; a run that covers fewer than three sources says so in its output rather than printing
a reassuring number derived from one process.

Upstash's free tier is 500,000 commands/month, which is far smaller than it sounds:

    500,000 / month = 16,129 / day = 11 / minute = 0.19 / second, sustained

0.19 commands per second is the entire engineering constraint. One worker polling Redis
every five seconds spends 518,400/month — the whole budget, to learn nothing. That is why
Redis is only ever on the request path here, and why this script exists: the rule is easy
to state and invisible to violate. A `cache.get()` added inside the reconciler's
per-instance loop is one line in review and 2,160,000 commands/month in production.

**This projects the burn rather than alarming on the counter.** Exhausting the budget is a slow,
silent failure with a cliff at the end: nothing degrades, nothing pages, every call
succeeds, and then the month rolls over and every Redis call starts erroring at once. By
then the fix is a bill or an outage. A burn rate extrapolated from the counter is visible
on day two, which is the only time it is cheap to fix.

Stdlib only, on purpose — this is meant to run from CI, from a laptop, or from inside a
pod that has nothing but python3, without an environment to activate first.
"""

from __future__ import annotations

import argparse
import sys
import time
import urllib.request
from urllib.parse import urlparse

# A 30-day month. Upstash bills on a calendar month; 30 days is the honest rounding and
# errs slightly pessimistic on the long ones, which is the correct direction for a budget.
_MONTH_S = 30 * 24 * 60 * 60
_FREE_TIER_BUDGET = 500_000

_COMMANDS_METRIC = "svcforge_redis_commands_total"
_START_TIME_METRIC = "process_start_time_seconds"


class BudgetError(RuntimeError):
    """The metrics endpoint did not give us enough to project from."""


def scrape(url: str, timeout_s: float = 5.0) -> str:
    """GET the Prometheus text exposition. http/https only."""
    if urlparse(url).scheme not in ("http", "https"):
        raise BudgetError(f"refusing to fetch a non-http(s) url: {url}")
    # S310 is satisfied by the scheme check above: this cannot open file:// or ftp://.
    with urllib.request.urlopen(url, timeout=timeout_s) as resp:  # noqa: S310
        body: str = resp.read().decode("utf-8")
    return body


def _parse_sample(line: str) -> tuple[str, dict[str, str], float] | None:
    """One exposition line -> (name, labels, value). None for comments and blanks.

    A deliberately small parser rather than prometheus_client's: importing the library
    would mean this script only runs where the app's venv is already active, which is
    exactly where you least need to check the budget.
    """
    line = line.strip()
    if not line or line.startswith("#"):
        return None
    head, _, raw_value = line.rpartition(" ")
    if not head:
        return None
    try:
        value = float(raw_value)
    except ValueError:
        return None

    name, brace, rest = head.partition("{")
    labels: dict[str, str] = {}
    if brace:
        for pair in rest.rstrip("}").split(","):
            key, eq, val = pair.partition("=")
            if eq:
                labels[key.strip()] = val.strip().strip('"')
    return name.strip(), labels, value


def collect(text: str) -> tuple[dict[str, float], float]:
    """Extract per-op command totals and the process start time from a scrape."""
    per_op: dict[str, float] = {}
    started_at: float | None = None

    for line in text.splitlines():
        parsed = _parse_sample(line)
        if parsed is None:
            continue
        name, labels, value = parsed
        # prometheus_client exposes counters with a `_total` suffix already; tolerate both
        # spellings so this keeps working if the client library changes its mind.
        if name in (_COMMANDS_METRIC, _COMMANDS_METRIC.removesuffix("_total")):
            per_op[labels.get("op", "unknown")] = value
        elif name == _START_TIME_METRIC:
            started_at = value

    if not per_op:
        raise BudgetError(
            f"{_COMMANDS_METRIC} is not exposed. Either the API never imported "
            f"svcforge_core.adapters.redis, or you are scraping the wrong process."
        )
    if started_at is None:
        raise BudgetError(
            f"{_START_TIME_METRIC} is missing, so there is no window to project over. "
            f"It comes from prometheus_client's default collector."
        )
    return per_op, started_at


def merge(scrapes: list[tuple[dict[str, float], float]]) -> tuple[dict[str, float], float]:
    """Fold several processes' scrapes into one budget view.

    The budget is per-Upstash-database, but the counters are per-process: api, worker and
    reconciler each hold their own prometheus_client registry, so scraping one of them
    projects a third of the truth. Command totals sum across processes; the window is the
    EARLIEST start time, because a counter that has been running longest bounds how far back
    the summed total can be attributed — using the latest would inflate the rate.
    """
    per_op: dict[str, float] = {}
    for scraped, _ in scrapes:
        for op, value in scraped.items():
            per_op[op] = per_op.get(op, 0.0) + value
    return per_op, min(started for _, started in scrapes)


def report(per_op: dict[str, float], started_at: float, budget: int, now: float, sources: int = 1) -> int:
    """Print the projection. Returns the process exit code."""
    elapsed_s = max(1.0, now - started_at)
    total = sum(per_op.values())
    rate = total / elapsed_s
    projected = rate * _MONTH_S

    print(f"window          {elapsed_s / 3600:.2f} h since process start")
    print(f"commands        {total:,.0f}")
    for op, value in sorted(per_op.items(), key=lambda kv: -kv[1]):
        share = (value / total * 100) if total else 0.0
        print(f"  {op:<14}{value:>12,.0f}  ({share:.1f}%)")
    print(f"rate            {rate:.4f} /s   (budget allows {budget / _MONTH_S:.4f} /s sustained)")
    print(f"projected       {projected:,.0f} / month")
    print(f"budget          {budget:,} / month")
    print(f"sources         {sources} process(es) scraped")
    if sources < 3:
        # Be honest about what was measured. api/worker/reconciler each keep their own
        # in-process registry (obs.py: one process per pod, no multiproc dir), so a
        # single-endpoint run undercounts the shared Upstash budget by however many
        # processes were left out. An optimistic verdict here is worse than no verdict.
        print(
            "                UNDERCOUNT: svcforge runs api + worker + reconciler, each with "
            "its own\n                registry. Pass --url once per process for the real "
            "total; the numbers\n                above cover only what was scraped."
        )

    if total < 100:
        # Extrapolating a month from a handful of commands is astrology. Say so rather than
        # printing a confident number derived from six samples.
        print("verdict         INCONCLUSIVE — fewer than 100 commands; let it run longer")
        return 0
    if projected >= budget:
        headroom = projected / budget
        print(f"verdict         OVER BUDGET — {headroom:.1f}x. Find the Redis call in a loop.")
        return 1
    share = projected / budget * 100
    print(f"verdict         OK — {share:.1f}% of budget, {budget / projected:.1f}x headroom")
    return 0


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--url",
        action="append",
        dest="urls",
        metavar="URL",
        help="Prometheus endpoint; repeat once per process (api, worker, reconciler)",
    )
    parser.add_argument("--budget", type=int, default=_FREE_TIER_BUDGET, help="commands per month")
    args = parser.parse_args(argv)

    # action="append" cannot carry a default (argparse appends to it), so apply it here.
    urls: list[str] = args.urls or ["http://localhost:8000/metrics"]

    scrapes: list[tuple[dict[str, float], float]] = []
    for url in urls:
        try:
            scrapes.append(collect(scrape(url)))
        except BudgetError as exc:
            print(f"error: {url}: {exc}", file=sys.stderr)
            return 2
        except OSError as exc:
            print(f"error: cannot scrape {url}: {exc}", file=sys.stderr)
            return 2

    per_op, started_at = merge(scrapes)
    return report(per_op, started_at, args.budget, time.time(), sources=len(scrapes))


if __name__ == "__main__":
    raise SystemExit(main())