svcforge/scripts/load.py

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

The point is not a benchmark. It 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, not mysterious:

    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 uuid import uuid4

import psycopg
from psycopg.rows import dict_row

from svcforge_core.settings import load_settings


async def _seed_direct(dsn: str, count: int) -> 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.
    """
    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', 'elasticsearch', 'small', 'requested',
                               'tenant-loadtest', %s, '21.3.19')""",
                    (iid, f"loadtest-elasticsearch-{str(iid)[:8]}"),
                )
                await cur.execute(
                    "insert into tasks (instance_id, kind) values (%s, 'provision')",
                    (iid,),
                )
    return time.monotonic() - started


async def _depth(dsn: str) -> dict[str, int]:
    async with await psycopg.AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
        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."""
    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

    while time.monotonic() - started < timeout_s:
        d = await _depth(dsn)
        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}   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)."
        )

    print(f"seeding {args.count} instances at {datetime.now(UTC).isoformat()} ...")
    took = await _seed_direct(dsn, args.count)
    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())