shopdb-flask/scripts/loadtest/poll_fleet.py

"""Async load generator for the endpoint command long-poll channel.

WHY this shape: the real agent (agent/ShopDbEndpointAgent.Core/CommandWorker.cs)
re-polls immediately on a 204, so every agent holds a connection essentially
continuously and fleet size IS the concurrency. This mirrors that loop rather
than issuing independent requests, so the measured concurrency matches a real
fleet of the same size.

Runs INSIDE a container on the compose network. Generating from the Windows host
would put host ephemeral-port and TLS-handshake limits in the measurement path,
which would look like server saturation but would not be.

While the fleet polls, a separate probe measures /api/dashboard/health - the
user-visible UI latency - because that, not agent latency, is the symptom that
matters for sizing.

Stdlib asyncio only, so it needs no dependency install in the container.

Usage:
    python poll_fleet.py --agents 300 --seconds 90 --base https://api:5001
"""
import argparse
import asyncio
import json
import os
import ssl
import statistics
import time

STOP = False


def percentile(values, fraction):
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(int(len(ordered) * fraction), len(ordered) - 1)
    return ordered[index]


async def request(host, port, context, path, api_key, timeout):
    """One HTTP/1.1 GET over a fresh connection. Returns (status, seconds)."""
    started = time.monotonic()
    writer = None
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(host, port, ssl=context), timeout=timeout)
        request_line = (
            'GET %s HTTP/1.1\r\nHost: %s\r\nX-API-Key: %s\r\n'
            'Accept: application/json\r\nConnection: close\r\n\r\n'
            % (path, host, api_key))
        writer.write(request_line.encode())
        await writer.drain()
        header = await asyncio.wait_for(reader.readuntil(b'\r\n'), timeout=timeout)
        status = int(header.split()[1])
        await asyncio.wait_for(reader.read(), timeout=timeout)
        return status, time.monotonic() - started
    except Exception as error:
        return type(error).__name__, time.monotonic() - started
    finally:
        if writer is not None:
            try:
                writer.close()
            except Exception:
                pass


async def agent_loop(index, host, port, context, api_key, results, timeout):
    """Imitate one agent: poll, and on completion immediately poll again."""
    hostname = 'LOADTEST%04d' % (index + 1)
    path = '/api/endpoints/commands/next?hostname=' + hostname
    while not STOP:
        status, elapsed = await request(
            host, port, context, path, api_key, timeout)
        results.append((status, elapsed))
        await asyncio.sleep(0)


async def main():
    global STOP
    parser = argparse.ArgumentParser()
    parser.add_argument('--agents', type=int, required=True)
    parser.add_argument('--seconds', type=int, default=90)
    parser.add_argument('--host', default='api')
    parser.add_argument('--port', type=int, default=5001)
    parser.add_argument('--timeout', type=float, default=120.0)
    parser.add_argument('--ramp', type=float, default=0.005)
    arguments = parser.parse_args()

    # A managed collector token, not the env key: the endpoints plugin has
    # plugin.endpoints.envkeyfallback=false (the agent kill switch), so the env
    # key no longer authorizes. See scripts/loadtest/loadtest_token.py.
    api_key = os.environ.get('LOADTEST_TOKEN')
    if not api_key:
        raise SystemExit('LOADTEST_TOKEN not present in environment')

    context = ssl.create_default_context()
    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE

    results = []
    tasks = []
    started = time.monotonic()

    for index in range(arguments.agents):
        tasks.append(asyncio.create_task(agent_loop(
            index, arguments.host, arguments.port, context, api_key,
            results, arguments.timeout)))
        if arguments.ramp:
            await asyncio.sleep(arguments.ramp)

    ramp_seconds = time.monotonic() - started

    await asyncio.sleep(arguments.seconds)
    STOP = True
    for task in tasks:
        task.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)

    poll_ok = [d for s, d in results if s == 204]
    poll_bad = [s for s, _ in results if s != 204]

    failures = {}
    for status in poll_bad:
        failures[str(status)] = failures.get(str(status), 0) + 1

    print(json.dumps({
        'agents': arguments.agents,
        'ramp_seconds': round(ramp_seconds, 1),
        'duration_seconds': arguments.seconds,
        'poll_completed': len(results),
        'poll_ok_204': len(poll_ok),
        'poll_failures': failures,
        'poll_hold_median': round(statistics.median(poll_ok), 2) if poll_ok else None,
        'poll_hold_p95': round(percentile(poll_ok, 0.95), 2) if poll_ok else None,
        'poll_hold_max': round(max(poll_ok), 2) if poll_ok else None,
    }, indent=2))


if __name__ == '__main__':
    asyncio.run(main())