runner.py

"""Performance-run orchestrator (CLI).

Times the **native** and **server-side** paths across the query set (cold +
warm), aggregates to percentiles + a bootstrap median CI, and writes one
structured JSON document under ``benchmarks/results/``.

    uv run python -m benchmarks.perf.runner \
        --database-url "$MCPG_TEST_DATABASE_URL" \
        --scale-factor 1 --iterations 50 --output benchmarks/results/perf.json

It also runs the **overhead decomposition** (perf/decompose.py) on the
server-side path and records the load-bearing ``t_db == native`` assertion.
With ``--e2e`` it additionally measures the **end-to-end paths** through the
real MCP protocol (perf/e2e.py): in-memory + stdio subprocess, plus streamable
HTTP against an operator-started server via ``--e2e-http-url``. With
``--concurrency`` it runs the **throughput sweep** (perf/concurrency.py) at
1/4/16/64 clients. Provenance (git SHA, timestamp) is passed in so a result
always carries the exact conditions that produced it.

Operator tool — not unit-tested (needs a live PostgreSQL); the pure helpers it
calls (stats, queries, schema, decompose) are.
"""

from __future__ import annotations

import argparse
import asyncio
import gc
import json
import logging
import platform
import sys
from pathlib import Path
from typing import Any

from benchmarks.perf import stats
from benchmarks.perf.concurrency import CONCURRENCY_LEVELS, ConcurrencyResult, sweep_level
from benchmarks.perf.decompose import DecompositionRunner, SegmentSample, summarize_segments, t_db_within_native
from benchmarks.perf.e2e import E2EHttpRunner, E2EInMemoryRunner, E2ERunner, E2EStdioRunner
from benchmarks.perf.paths import NativeRunner, PathRunner, ServerSideRunner
from benchmarks.perf.queries import BenchQuery, all_queries
from benchmarks.perf.schema import Assertion, Decomposition, LatencyBlock, PerfRun, ResultRow
from mcpg import __version__
from mcpg.config import load_settings
from mcpg.database import Database
from mcpg.query import DEFAULT_TIMEOUT_SECONDS

logger = logging.getLogger(__name__)

_WARMUP = 5


async def _sample_path(runner: PathRunner, query: BenchQuery, iterations: int) -> tuple[list[int], int]:
    """Return (warm samples, cold sample) for one path x query.

    GC is disabled around each timed call so a collection pause never lands
    inside a measurement; collected in the gaps instead.
    """
    gc.collect()
    cold = await runner.run_once(query.sql, max_rows=query.max_rows)  # first call = cold bucket
    samples: list[int] = []
    for i in range(iterations + _WARMUP):
        gc.disable()
        try:
            elapsed = await runner.run_once(query.sql, max_rows=query.max_rows)
        finally:
            gc.enable()
        samples.append(elapsed)
        if i % 8 == 7:
            gc.collect()
    warm = stats.drop_warmup(samples, _WARMUP)
    return warm, cold


async def _sample_decomposition(runner: DecompositionRunner, query: BenchQuery, iterations: int) -> Decomposition:
    """Sample the server-path waterfall, returning per-segment medians (ns).

    Warm-up is discarded exactly like :func:`_sample_path`; GC is disabled
    around each timed call so no collection pause lands inside a segment.
    """
    gc.collect()
    samples: list[SegmentSample] = []
    for i in range(iterations + _WARMUP):
        gc.disable()
        try:
            sample = await runner.run_once(query.sql, max_rows=query.max_rows)
        finally:
            gc.enable()
        samples.append(sample)
        if i % 8 == 7:
            gc.collect()
    return summarize_segments(samples, warmup=_WARMUP)


async def _sample_native_db(native: NativeRunner, query: BenchQuery, iterations: int) -> float:
    """Median (ns) of native's pure DB segment — the ``t_db == native`` anchor."""
    gc.collect()
    samples: list[int] = []
    for i in range(iterations + _WARMUP):
        gc.disable()
        try:
            samples.append(await native.db_segment_once(query.sql))
        finally:
            gc.enable()
        if i % 8 == 7:
            gc.collect()
    warm = stats.drop_warmup(samples, _WARMUP)
    return stats.percentile(sorted(warm), 50)


def _conc_row(path: str, query: BenchQuery, cr: ConcurrencyResult) -> ResultRow:
    """Build a ResultRow for one concurrency data-point.

    Raw per-call samples are omitted (``samples_ns=[]``) — at 64 clients x N
    iterations they'd bloat the JSON with little added value over the
    percentiles + throughput already summarised here.
    """
    s = stats.summarize(cr.latencies_ns)
    return ResultRow(
        path=path,
        query_id=query.id,
        compute_class=query.compute_class,
        result_size=query.result_size,
        temperature="warm",
        concurrency=cr.concurrency,
        n=s.n,
        latency_ms=LatencyBlock(
            p50=s.p50, p95=s.p95, p99=s.p99, mean=s.mean, stdev=s.stdev, min=s.min, max=s.max, median_ci95=s.median_ci95
        ),
        throughput_rps=cr.throughput_rps,
        samples_ns=[],
    )


async def _run_concurrency(database_url: str, iterations: int, timeout: float) -> list[ResultRow]:
    """Sweep the **ultralight** queries across the concurrency levels.

    Throughput-under-load exists to expose the *pool + per-call* overhead, so it
    only makes sense on trivially-cheap queries (point lookups, ``SELECT 1``).
    Anything heavier measures the database instead of MCPg: a heavy TPC-H query
    at 64 clients just times the DB's own execution, and even a "light" 90 ms
    GROUP BY (or a 100k-row fetch) run 64-way saturates CPU / serialization and
    can blow past the per-query timeout. So the sweep restricts to
    ``compute_class == "ultralight"``.

    Owns its own resources: a dedicated pool sized to the sweep ceiling (so the
    server-side path isn't starved) and one persistent native connection per
    concurrent client (opened once, sliced per level).

    Connection budget: at the top level the native path holds ``max_c``
    connections **while** the server pool opens up to ``max_c`` -- so PostgreSQL
    needs ``max_connections`` >= about ``2 * max(CONCURRENCY_LEVELS)`` plus
    headroom (e.g. 150+ for the 64-client sweep). Below that, connection
    checkouts fail with "too many clients" and the numbers are degraded; size
    the server for the sweep before running it.
    """
    max_c = max(CONCURRENCY_LEVELS)
    conc_settings = load_settings(
        {
            "MCPG_DATABASE_URL": database_url,
            "MCPG_POOL_MIN_SIZE": "1",
            "MCPG_POOL_MAX_SIZE": str(max_c),
            "MCPG_STATEMENT_TIMEOUT_MS": str(int(timeout * 1000)),
        }
    )
    conc_db = Database(conc_settings)
    await conc_db.connect()
    native_conns = [await NativeRunner.connect(database_url) for _ in range(max_c)]
    rows: list[ResultRow] = []
    try:
        for query in all_queries():
            if query.compute_class != "ultralight":
                continue
            for level in CONCURRENCY_LEVELS:
                native_result = await sweep_level(
                    list(native_conns[:level]), query, iterations_per_worker=iterations, warmup_per_worker=_WARMUP
                )
                rows.append(_conc_row("native", query, native_result))
                server_runners: list[PathRunner] = [ServerSideRunner(conc_db) for _ in range(level)]
                server_result = await sweep_level(
                    server_runners, query, iterations_per_worker=iterations, warmup_per_worker=_WARMUP
                )
                rows.append(_conc_row("server_side", query, server_result))
    finally:
        for conn in native_conns:
            await conn.close()
        await conc_db.close()
    return rows


def _row(
    path: str,
    query: BenchQuery,
    temperature: str,
    samples_ns: list[int],
    decomposition: Decomposition | None = None,
) -> ResultRow:
    s = stats.summarize(samples_ns)
    return ResultRow(
        path=path,
        query_id=query.id,
        compute_class=query.compute_class,
        result_size=query.result_size,
        temperature=temperature,
        concurrency=1,
        n=s.n,
        latency_ms=LatencyBlock(
            p50=s.p50, p95=s.p95, p99=s.p99, mean=s.mean, stdev=s.stdev, min=s.min, max=s.max, median_ci95=s.median_ci95
        ),
        samples_ns=samples_ns,
        decomposition_ns=decomposition,
    )


def _checkpoint(
    args: argparse.Namespace,
    pg_meta: dict[str, Any],
    results: list[ResultRow],
    native_db_ns_by_query: dict[str, float],
    e2e_failures: list[dict[str, str]],
    *,
    complete: bool,
) -> PerfRun:
    """Write the run so far to ``args.output``, overwriting the prior checkpoint.

    Heavy TPC-H queries at large scale factors can each take minutes across
    every path x iteration x decomposition sample, and this environment has
    seen the run process killed externally (outside any Python exception) more
    than once over a multi-hour run. Checkpointing after every query means a
    kill loses at most one query's worth of measurement, not the whole run —
    the file is always a valid, loadable ``PerfRun`` JSON, just possibly
    ``metadata.complete: false`` if interrupted before the final query.
    """
    metadata: dict[str, Any] = {
        "timestamp": args.timestamp,
        "git_sha": args.git_sha,
        "mcpg_version": __version__,
        "postgres": pg_meta,
        "scale_factor": args.scale_factor,
        "host": {
            "python": platform.python_version(),
            "os": platform.platform(),
            "machine": platform.machine(),
        },
        "iterations": args.iterations,
        "server_side_timeout_seconds": args.timeout,
        "warmup_discarded": _WARMUP,
        "concurrency_levels": list(CONCURRENCY_LEVELS) if args.concurrency else [],
        "e2e_failures": e2e_failures,
        "complete": complete,
    }
    run = PerfRun(metadata=metadata, results=results, assertions=_assertions(results, native_db_ns_by_query))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(run.to_dict(), indent=2) + "\n")
    return run


async def _run(args: argparse.Namespace) -> PerfRun:
    # MCPG_STATEMENT_TIMEOUT_MS sets Postgres's own statement_timeout GUC — a
    # second, independent ceiling from the asyncio guard --timeout controls
    # (mcpg.query.run_select's own `timeout` kwarg). Both must be raised
    # together for heavy queries at large scale factors; deriving this one
    # from --timeout keeps the single flag authoritative rather than needing
    # two flags kept in sync by hand. Applies to server_side and e2e_inmemory
    # (both share this Settings/Database) — not native (raw psycopg, no
    # MCPg settings) or e2e_stdio (a real subprocess with its own env, which
    # inherits MCPG_STATEMENT_TIMEOUT_MS from the shell if set there).
    settings = load_settings(
        {
            "MCPG_DATABASE_URL": args.database_url,
            "MCPG_POOL_MIN_SIZE": "1",
            "MCPG_POOL_MAX_SIZE": "4",
            "MCPG_STATEMENT_TIMEOUT_MS": str(int(args.timeout * 1000)),
        }
    )
    database = Database(settings)
    await database.connect()
    native = await NativeRunner.connect(args.database_url)
    results: list[ResultRow] = []
    # Pre-initialised so a failure in the metadata query below can't mask the
    # original exception with an UnboundLocalError when `metadata` is built.
    pg_meta: dict[str, Any] = {}
    e2e_runners: list[E2ERunner] = []
    try:
        pg = await database.driver().execute_query(
            "SELECT current_setting('server_version') AS v, current_setting('server_version_num')::int AS num",
            force_readonly=True,
        )
        if pg:
            pg_meta = {"version_string": pg[0].cells["v"], "server_version_num": pg[0].cells["num"]}

        # Opt-in end-to-end paths (through the real MCP protocol). Started once
        # and reused across every query; torn down in the finally.
        e2e_paths: list[tuple[str, PathRunner]] = []
        # Register each runner for teardown *before* starting it, so a failure
        # partway through start() still gets closed by the finally (close() is a
        # no-op on an un-entered stack).
        if args.e2e:
            inmem = E2EInMemoryRunner(settings)
            e2e_runners.append(inmem)
            await inmem.start()
            e2e_paths.append(("e2e_inmemory", inmem))
            stdio = E2EStdioRunner(args.database_url)
            e2e_runners.append(stdio)
            await stdio.start()
            e2e_paths.append(("e2e_stdio", stdio))
        if args.e2e_http_url:
            http = E2EHttpRunner(args.e2e_http_url)
            e2e_runners.append(http)
            await http.start()
            e2e_paths.append(("e2e_http", http))

        paths: list[tuple[str, PathRunner]] = [
            ("native", native),
            ("server_side", ServerSideRunner(database, timeout=args.timeout)),
            *e2e_paths,
        ]
        decomposer = DecompositionRunner(database)
        native_db_ns_by_query: dict[str, float] = {}
        e2e_failures: list[dict[str, str]] = []
        for query in all_queries():
            for label, runner in paths:
                # e2e paths call the real, unmodified server tool, which enforces
                # the product's fixed 30s query timeout with no override (unlike
                # the server-side path here, which accepts --timeout). A heavy
                # query at a large scale factor can genuinely exceed that on
                # slower hardware — a real product limit, not a harness bug — so
                # it must not discard every other path/query already measured.
                # native/server_side failures are not caught: those would be a
                # real harness or product regression, not an expected ceiling.
                try:
                    warm, cold = await _sample_path(runner, query, args.iterations)
                except Exception as exc:
                    if not label.startswith("e2e"):
                        raise
                    logger.warning("e2e path %r timed out/failed on query %r: %s", label, query.id, exc)
                    e2e_failures.append({"path": label, "query_id": query.id, "error": str(exc)})
                    continue
                # Attach the overhead waterfall to the server-side warm row —
                # the one the report reads t_db from for the native comparison.
                decomposition = (
                    await _sample_decomposition(decomposer, query, args.iterations) if label == "server_side" else None
                )
                results.append(_row(label, query, "warm", warm, decomposition))
                results.append(_row(label, query, "cold", [cold]))
            # The native DB segment (execute + fetch only) anchors t_db == native.
            native_db_ns_by_query[query.id] = await _sample_native_db(native, query, args.iterations)
            _ = _checkpoint(args, pg_meta, results, native_db_ns_by_query, e2e_failures, complete=False)

        # Throughput-under-concurrency sweep (opt-in; owns its own pool sized to
        # the sweep ceiling so the server-side path isn't starved). It runs last
        # and is non-essential, so a failure here (e.g. connection exhaustion or
        # a per-query timeout under load) must not discard the expensive
        # core + e2e results already collected — log and carry on.
        if args.concurrency:
            try:
                results.extend(await _run_concurrency(args.database_url, args.iterations, args.timeout))
            except Exception as exc:
                logger.warning("Concurrency sweep failed; writing results without it: %s", exc)
            _ = _checkpoint(args, pg_meta, results, native_db_ns_by_query, e2e_failures, complete=False)
    finally:
        # Close e2e runners in REVERSE start order. Each one opens an internal
        # anyio task-group / cancel scope when started (in-memory, then stdio,
        # then http), nested in that order; anyio requires the inner scope exit
        # before the outer, so tearing down forward raises "Attempted to exit a
        # cancel scope that isn't the current task's current cancel scope" and
        # (from the finally) aborts the whole run before the JSON is written.
        for e2e in reversed(e2e_runners):
            await e2e.close()
        await native.close()
        await database.close()

    return _checkpoint(args, pg_meta, results, native_db_ns_by_query, e2e_failures, complete=True)


def _assertions(results: list[ResultRow], native_db_ns_by_query: dict[str, float]) -> list[Assertion]:
    """Overhead + the load-bearing ``t_db == native`` gate.

    Two assertions per query: an informational warm total-latency delta
    (``server_side_overhead_p50_ms``), and the machine-checkable claim that the
    server path's DB segment matches the native baseline
    (``t_db_matches_native``) — the result the whole performance objective turns
    on.
    """
    out: list[Assertion] = []
    # Only the single-client baseline rows feed the overhead and t_db
    # assertions. A row is the baseline when it is warm, concurrency == 1, AND
    # carries no throughput (throughput_rps is None) — the concurrency sweep also
    # emits a concurrency==1 row (its level-1 point), which shares (path,
    # query_id) but has throughput set and no decomposition; without the
    # throughput guard it overwrites the real baseline and the t_db gate silently
    # drops those queries.
    by_key = {
        (r.path, r.query_id): r
        for r in results
        if r.temperature == "warm" and r.concurrency == 1 and r.throughput_rps is None
    }
    query_ids = {r.query_id for r in results}
    for qid in sorted(query_ids):
        native = by_key.get(("native", qid))
        server = by_key.get(("server_side", qid))
        if native is None or server is None:
            continue
        overhead_ms = server.latency_ms.p50 - native.latency_ms.p50
        out.append(
            Assertion(
                name="server_side_overhead_p50_ms",
                query_id=qid,
                passed=True,  # informational
                detail={
                    "native_p50_ms": native.latency_ms.p50,
                    "server_side_p50_ms": server.latency_ms.p50,
                    "overhead_p50_ms": overhead_ms,
                },
            )
        )
        server_t_db = server.decomposition_ns.t_db if server.decomposition_ns else None
        native_t_db = native_db_ns_by_query.get(qid)
        if server_t_db is not None and native_t_db is not None:
            out.append(
                Assertion(
                    name="t_db_matches_native",
                    query_id=qid,
                    passed=t_db_within_native(server_t_db, native_t_db),
                    detail={
                        "native_t_db_ms": native_t_db / 1e6,
                        "server_t_db_ms": server_t_db / 1e6,
                        "delta_ms": (server_t_db - native_t_db) / 1e6,
                    },
                )
            )
    return out


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="MCPg performance benchmark (native vs server-side).")
    parser.add_argument("--database-url", required=True, help="PostgreSQL DSN (a TPC-H-loaded database).")
    parser.add_argument(
        "--iterations", type=int, default=50, help=f"Warm iterations per pathxquery (>= {_WARMUP + 1})."
    )
    parser.add_argument("--scale-factor", type=int, default=1, help="TPC-H scale factor the DB was loaded at.")
    parser.add_argument(
        "--timeout",
        type=float,
        default=DEFAULT_TIMEOUT_SECONDS,
        help="Per-query timeout (seconds) for the in-process server-side path. Defaults to run_select's own "
        "product default; raise it for heavy-tier queries at large scale factors whose native runtime "
        "approaches that ceiling on slower hardware. Only affects the server-side path — the e2e paths call "
        "the real, unmodified server tool, which always enforces the product default.",
    )
    parser.add_argument("--output", type=Path, required=True, help="Path to write the result JSON.")
    parser.add_argument(
        "--e2e",
        action="store_true",
        help="Also measure the end-to-end paths through the MCP protocol (in-memory + stdio subprocess).",
    )
    parser.add_argument(
        "--e2e-http-url",
        default=None,
        help="Add the streamable-HTTP e2e path against an operator-started mcpg server at this URL (e.g. "
        "http://127.0.0.1:8000/mcp). Implies the HTTP transport is already running.",
    )
    parser.add_argument(
        "--concurrency",
        action="store_true",
        help=f"Also run the throughput sweep at {'/'.join(map(str, CONCURRENCY_LEVELS))} concurrent clients "
        "(native + server-side). Multiplies DB load — each level runs --iterations per client.",
    )
    parser.add_argument("--git-sha", default="unknown", help="Provenance: the commit under test.")
    parser.add_argument("--timestamp", default="unknown", help="Provenance: ISO-8601 run timestamp.")
    args = parser.parse_args(argv)
    if args.iterations < _WARMUP + 1:
        # Below this the warm bucket is empty after dropping warmup, and
        # summarize() would report misleading all-zero stats that look valid.
        parser.error(f"--iterations must be >= {_WARMUP + 1} (warm measurements would be empty)")

    run = asyncio.run(_run(args))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(run.to_dict(), indent=2) + "\n")
    print(f"wrote {args.output} ({len(run.results)} result rows)")
    return 0


if __name__ == "__main__":  # pragma: no cover
    sys.exit(main())