LedgerLine/tests/harness/load.py

"""Concurrent request driver + LatencyReport, per phase-0a-test-suite-spec.md §7.

A closed-loop driver (fixed concurrency) is used deliberately for its
simplicity here. This under-measures latency under saturation (coordinated
omission) -- Phase 6 replaces it with k6's arrival-rate executors. Don't try
to fix that in this suite; note it, live with it.
"""
import asyncio
import math
import statistics
import time
from dataclasses import dataclass, field
from typing import Callable

import httpx


@dataclass
class LatencyReport:
    count: int
    ok: int
    http_errors: int  # non-2xx HTTP responses
    exceptions: int  # requests that raised (timeout, connection reset, ...)
    errors: int  # http_errors + exceptions
    error_rate: float
    status_counts: dict
    p50: float
    p95: float
    p99: float
    max: float
    throughput_rps: float
    wall_seconds: float
    latencies_ms: list = field(default_factory=list, repr=False)

    def __str__(self):
        return (
            f"n={self.count} ok={self.ok} http_errors={self.http_errors} "
            f"exceptions={self.exceptions} error_rate={self.error_rate:.1%} "
            f"p50={self.p50:.1f}ms p95={self.p95:.1f}ms p99={self.p99:.1f}ms "
            f"max={self.max:.1f}ms throughput={self.throughput_rps:.1f}rps "
            f"wall={self.wall_seconds:.2f}s status_counts={self.status_counts}"
        )


def _nearest_rank(sorted_vals, p):
    """Nearest-rank percentile: no interpolation, no averaging."""
    n = len(sorted_vals)
    if n == 0:
        return 0.0
    rank = max(1, min(n, math.ceil(p * n)))
    return sorted_vals[rank - 1]


async def drive(
    method: str,
    url_factory: Callable[[], tuple],
    total: int,
    concurrency: int,
    warmup: int = 20,
    timeout: float = 15.0,
) -> LatencyReport:
    """Drive `total` requests at fixed `concurrency`, discarding `warmup` first.

    `url_factory()` returns `(url, kwargs)`, called once per request so a test
    can vary the target account/params per call.
    """
    async with httpx.AsyncClient(timeout=timeout) as client:

        async def one():
            url, kwargs = url_factory()
            t0 = time.perf_counter()
            try:
                resp = await client.request(method, url, **kwargs)
                return (time.perf_counter() - t0) * 1000, resp.status_code
            except Exception:
                return (time.perf_counter() - t0) * 1000, None

        sem = asyncio.Semaphore(concurrency)

        async def bound(collector):
            async with sem:
                result = await one()
                if collector is not None:
                    collector.append(result)

        if warmup:
            await asyncio.gather(*[bound(None) for _ in range(warmup)])

        results: list = []
        wall_start = time.perf_counter()
        await asyncio.gather(*[bound(results) for _ in range(total)])
        wall_seconds = time.perf_counter() - wall_start

    latencies = sorted(r[0] for r in results)
    status_counts: dict = {}
    ok = http_errors = exceptions = 0
    for _, status in results:
        if status is None:
            exceptions += 1
            status_counts["exception"] = status_counts.get("exception", 0) + 1
        elif 200 <= status < 300:
            ok += 1
            status_counts[status] = status_counts.get(status, 0) + 1
        else:
            http_errors += 1
            status_counts[status] = status_counts.get(status, 0) + 1

    count = len(results)
    errors = http_errors + exceptions
    return LatencyReport(
        count=count,
        ok=ok,
        http_errors=http_errors,
        exceptions=exceptions,
        errors=errors,
        error_rate=(errors / count if count else 0.0),
        status_counts=status_counts,
        p50=_nearest_rank(latencies, 0.50),
        p95=_nearest_rank(latencies, 0.95),
        p99=_nearest_rank(latencies, 0.99),
        max=(latencies[-1] if latencies else 0.0),
        throughput_rps=(count / wall_seconds if wall_seconds > 0 else 0.0),
        wall_seconds=wall_seconds,
        latencies_ms=latencies,
    )


async def measure_median(n: int, fn: Callable, key=lambda x: x):
    """Run async `fn()` n times and return the run whose `key(result)` is the
    median -- e.g. key=lambda r: r.p95 picks the median-p95 LatencyReport
    whole, rather than averaging percentiles across runs (never do that)."""
    results = [await fn() for _ in range(n)]
    return sorted(results, key=key)[len(results) // 2]