shopdb-flask/scripts/loadtest/mixed_workload.py

"""Mixed-workload generator for sizing simulations.

WHY this shape: sizing cannot be judged from long-poll agents alone. The four
traffic classes have different concurrency costs, and only their SUM decides
whether a worker model holds:

  fleet       long-poll agents, one continuously-held connection each
  ui          human users clicking, latency-sensitive, the actual symptom
  integration Maximo / ServiceNow style collector clients on a fixed interval
  kiosk       unauthenticated display endpoints, public routes, low rate

The fleet class is what saturates a thread pool; the ui class is what a person
notices. Measuring them together is the point - a model that keeps agents happy
while starving the UI has failed even though every agent request succeeded.

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.

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

Usage:
    python mixed_workload.py --fleet 300 --users 5 --seconds 120
"""
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, method, path, timeout,
                  api_key=None, bearer=None, body=None):
    """One HTTP/1.1 request 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)
        headers = [
            '%s %s HTTP/1.1' % (method, path),
            'Host: %s' % host,
            'Accept: application/json',
            'Connection: close',
        ]
        if api_key:
            headers.append('X-API-Key: %s' % api_key)
        if bearer:
            headers.append('Authorization: Bearer %s' % bearer)
        payload = b''
        if body is not None:
            payload = json.dumps(body).encode()
            headers.append('Content-Type: application/json')
            headers.append('Content-Length: %d' % len(payload))
        blob = ('\r\n'.join(headers) + '\r\n\r\n').encode() + payload
        writer.write(blob)
        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


def record(results, bucket, status, elapsed):
    entry = results.setdefault(bucket, {'times': [], 'status': {}})
    entry['times'].append(elapsed * 1000.0)
    key = str(status)
    entry['status'][key] = entry['status'].get(key, 0) + 1


async def fleet_agent(index, host, port, context, api_key, results):
    """Mirrors CommandWorker.cs: re-poll immediately after each response."""
    hostname = 'LOADTEST%04d' % (index + 1)
    path = '/api/endpoints/commands/next?hostname=%s' % hostname
    while not STOP:
        status, elapsed = await request(
            host, port, context, 'GET', path, timeout=120.0, api_key=api_key)
        record(results, 'fleet', status, elapsed)
        await asyncio.sleep(0.05)


async def ui_user(host, port, context, bearer, results, think_seconds):
    """A human clicking around. Latency here is the user-visible symptom."""
    paths = ['/api/dashboard/health', '/api/endpoints?page=1&per_page=25',
             '/api/assets?page=1&per_page=25']
    cursor = 0
    while not STOP:
        path = paths[cursor % len(paths)]
        cursor += 1
        # Unauthenticated probe still measures queueing; auth only changes the
        # handler cost, not whether a worker slot was available to accept it.
        status, elapsed = await request(
            host, port, context, 'GET', path, timeout=120.0, bearer=bearer)
        record(results, 'ui', status, elapsed)
        await asyncio.sleep(think_seconds)


async def integration_client(name, host, port, context, api_key, results,
                             interval):
    """Maximo / ServiceNow style poller on a fixed interval."""
    while not STOP:
        status, elapsed = await request(
            host, port, context, 'GET', '/api/collector/status',
            timeout=120.0, api_key=api_key)
        record(results, 'integration', status, elapsed)
        await asyncio.sleep(interval)


async def kiosk_endpoint(index, host, port, context, results, interval):
    """Unauthenticated display screen. Public route, low rate."""
    path = '/api/dashboarddefaults/display-role'
    while not STOP:
        status, elapsed = await request(
            host, port, context, 'GET', path, timeout=120.0)
        record(results, 'kiosk', status, elapsed)
        await asyncio.sleep(interval)


async def main():
    global STOP
    parser = argparse.ArgumentParser()
    parser.add_argument('--fleet', type=int, default=300)
    parser.add_argument('--users', type=int, default=5)
    parser.add_argument('--integrations', type=int, default=2)
    parser.add_argument('--kiosks', type=int, default=8)
    parser.add_argument('--seconds', type=int, default=120)
    parser.add_argument('--host', default='api')
    parser.add_argument('--port', type=int, default=5001)
    parser.add_argument('--think', type=float, default=3.0)
    parser.add_argument('--integration-interval', type=float, default=30.0)
    parser.add_argument('--kiosk-interval', type=float, default=15.0)
    parser.add_argument('--rampseconds', type=float, default=20.0)
    arguments = parser.parse_args()

    api_key = os.environ.get('LOADTEST_TOKEN', '')
    if not api_key:
        raise SystemExit('LOADTEST_TOKEN not set')
    bearer = os.environ.get('LOADTEST_JWT') or None

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

    results = {}
    tasks = []

    # Ramp the fleet in rather than opening every connection at once. A
    # thundering herd measures handshake burst capacity, not steady state.
    stagger = arguments.rampseconds / max(arguments.fleet, 1)

    async def staggered(index):
        await asyncio.sleep(index * stagger)
        await fleet_agent(index, arguments.host, arguments.port, context,
                          api_key, results)

    for index in range(arguments.fleet):
        tasks.append(asyncio.ensure_future(staggered(index)))
    for _ in range(arguments.users):
        tasks.append(asyncio.ensure_future(ui_user(
            arguments.host, arguments.port, context, bearer, results,
            arguments.think)))
    for index in range(arguments.integrations):
        tasks.append(asyncio.ensure_future(integration_client(
            'integration%d' % index, arguments.host, arguments.port, context,
            api_key, results, arguments.integration_interval)))
    for index in range(arguments.kiosks):
        tasks.append(asyncio.ensure_future(kiosk_endpoint(
            index, arguments.host, arguments.port, context, results,
            arguments.kiosk_interval)))

    await asyncio.sleep(arguments.seconds)
    STOP = True
    await asyncio.sleep(2.0)
    for task in tasks:
        task.cancel()

    summary = {
        'fleet': arguments.fleet,
        'users': arguments.users,
        'integrations': arguments.integrations,
        'kiosks': arguments.kiosks,
        'seconds': arguments.seconds,
        'classes': {},
    }
    for bucket, entry in sorted(results.items()):
        times = entry['times']
        summary['classes'][bucket] = {
            'samples': len(times),
            'status': entry['status'],
            'median_ms': round(statistics.median(times), 1) if times else 0.0,
            'p95_ms': round(percentile(times, 0.95), 1),
            'max_ms': round(max(times), 1) if times else 0.0,
        }
    print(json.dumps(summary, indent=2))


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