svcforge/scripts/load.py

"""Throwaway load generator. Enqueue N instances, watch the queue drain, print three numbers.

The point is to find the ceiling on purpose, in a place where finding it is free, so that
the number in RUNBOOK.md comes from an observation instead of a guess.

The ceiling you are looking for is arithmetic:

    total connections = (api_replicas + worker_replicas) x pool_max_size

Supabase's free-tier pooler has a small connection budget. Cross it and the failure does not
look like "too many connections" — it looks like slow claims, then PoolTimeout, then a queue
that grows while every worker looks idle. Once you have watched it once, you recognise it in
two seconds instead of an hour.

Usage:
    python -m scripts.load --count 200 --watch
    python -m scripts.load --count 200 --direct        # skip the API, enqueue straight to the DB
"""

from __future__ import annotations

import argparse
import asyncio
import time
from datetime import UTC, datetime
from typing import Any
from uuid import uuid4

import psycopg
from psycopg.rows import dict_row

from svcforge_core.domain.catalog import load_catalog
from svcforge_core.settings import load_settings

_SERVICE_TYPE = "elasticsearch"


async def _seed_direct(dsn: str, count: int, chart_version: str) -> float:
    """Insert `count` instances + provision tasks. Returns seconds taken.

    --direct exists to separate two questions that a single POST run conflates: "how fast can
    the API accept work" and "how fast can workers drain it". Measure them apart or you will
    tune the wrong one.

    `chart_version` comes from the catalog rather than a literal. Hardcoding it meant the
    seeded rows carried a version the catalog could not resolve, so every task failed fast
    and the drain measurement — the whole point of the script — timed the failure path.
    """
    started = time.monotonic()
    async with await psycopg.AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
        async with conn.transaction(), conn.cursor() as cur:
            for _ in range(count):
                iid = uuid4()
                await cur.execute(
                    """insert into instances (id, team, service_type, size, state, namespace,
                                              release_name, chart_version)
                       values (%s, 'loadtest', %s, 'small', 'requested',
                               'tenant-loadtest', %s, %s)""",
                    (iid, _SERVICE_TYPE, f"loadtest-{_SERVICE_TYPE}-{str(iid)[:8]}", chart_version),
                )
                await cur.execute(
                    "insert into tasks (instance_id, kind) values (%s, 'provision')",
                    (iid,),
                )
    return time.monotonic() - started


async def _depth(conn: psycopg.AsyncConnection[dict[str, Any]]) -> dict[str, int]:
    """Task counts by state, on a connection the caller owns."""
    cur = await conn.execute("select state, count(*) as n from tasks group by 1")
    return {str(r["state"]): int(r["n"]) for r in await cur.fetchall()}


async def _watch(dsn: str, timeout_s: float) -> None:
    """Print queue depth once a second until it drains. The slope is the number you want.

    One connection for the whole loop, held open. Reconnecting every second added a
    connection to the pooler on every tick of a script whose entire purpose is finding the
    connection ceiling — the measurement was perturbing the thing being measured.
    """
    started = time.monotonic()
    peak = 0
    print(f"{'t(s)':>6}  {'queued':>7} {'running':>8} {'done':>6} {'failed':>7}  slope/s")
    prev_done, prev_t = 0, started

    # autocommit: a held connection without it sits idle-in-transaction between polls, which
    # pins a snapshot on the pooler and is exactly the pathology this script hunts for.
    async with await psycopg.AsyncConnection.connect(dsn, row_factory=dict_row, autocommit=True) as conn:
        while time.monotonic() - started < timeout_s:
            d = await _depth(conn)
            queued, running = d.get("queued", 0), d.get("running", 0)
            done, failed = d.get("done", 0), d.get("failed", 0)
            peak = max(peak, queued + running)

            now = time.monotonic()
            slope = (done - prev_done) / max(now - prev_t, 1e-9)
            prev_done, prev_t = done, now

            print(f"{now - started:6.1f}  {queued:7d} {running:8d} {done:6d} {failed:7d}  {slope:7.1f}")

            if queued == 0 and running == 0:
                elapsed = now - started
                print(
                    f"\ndrained in {elapsed:.1f}s   peak depth {peak}   "
                    f"throughput {done / elapsed:.1f} task/s"
                )
                if failed:
                    print(f"WARNING: {failed} tasks failed — the number above is not a clean drain")
                return
            await asyncio.sleep(1.0)

    print(f"\nstill draining after {timeout_s}s — that IS the result. Record it.")


async def _amain() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--count", type=int, default=200)
    ap.add_argument("--direct", action="store_true", help="enqueue via SQL instead of the API")
    ap.add_argument("--watch", action="store_true", help="poll queue depth until drained")
    ap.add_argument("--timeout", type=float, default=600.0)
    ap.add_argument("--cleanup", action="store_true", help="delete loadtest rows and exit")
    args = ap.parse_args()

    settings = load_settings()
    dsn = settings.pg_dsn.unicode_string()

    if args.cleanup:
        async with await psycopg.AsyncConnection.connect(dsn, autocommit=True) as conn:
            await conn.execute("delete from instances where team = 'loadtest'")  # tasks cascade
        print("loadtest rows deleted")
        return

    if not args.direct:
        raise SystemExit(
            "POST mode needs a token; use --direct for the drain measurement, or drive the API "
            "with k6 (one dependency, not two — do not add locust for this)."
        )

    # The version the workers will actually resolve. Read it rather than restate it: a seeded
    # row whose chart_version disagrees with the catalog drains through the failure path.
    catalog = load_catalog(settings.catalog_path)
    entry = catalog.get(_SERVICE_TYPE)
    if entry is None:
        raise SystemExit(f"{settings.catalog_path} has no '{_SERVICE_TYPE}' entry to load-test with")

    print(f"seeding {args.count} instances at {datetime.now(UTC).isoformat()} ...")
    print(f"  service_type={_SERVICE_TYPE} chart_version={entry.chart_version} (from catalog)")
    took = await _seed_direct(dsn, args.count, entry.chart_version)
    print(f"enqueued {args.count} in {took:.2f}s ({args.count / took:.0f}/s)\n")

    if args.watch:
        await _watch(dsn, args.timeout)
        print("\nremember: `python -m scripts.load --cleanup` when you are done.")


if __name__ == "__main__":
    asyncio.run(_amain())