shopdb-flask/scripts/loadtest/burst_arrival.py

"""Patch-day burst simulator: synchronized cold start of an endpoint fleet.

WHY this is separate from mixed_workload.py: that harness RAMPS agents in over
~20s and measures steady state, which deliberately avoids a thundering herd.
Patch day is the opposite case and stresses different machinery:

  - a fresh TLS handshake per agent, all at once (handshakes are CPU-heavy)
  - _resolve_asset_for_hostname runs on every arrival, so the DB pool is hit by
    a synchronized burst rather than the smooth churn of steady state
  - no warm connections in the pool, no warm caches

Modelled on the real agent (verified in agent/ShopDbEndpointAgent.Core/):
  - CommandWorker.ExecuteAsync polls IMMEDIATELY, no startup delay, no jitter
  - TelemetryWorker also reports immediately; its jitter applies only BETWEEN
    runs, so first contact is synchronized too. Patch day is therefore TWO
    simultaneous bursts per endpoint, not one.
  - transient failures retry at a FIXED 10s (TransientRetryDelay), not
    exponential-with-jitter, so a herd that fails together re-arrives together

--arrivalseconds is the spread of the reboot window. 0 means every agent hits at
once (single maintenance window, worst case); 60 approximates a staged ring.

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

Usage:
    python burst_arrival.py --fleet 700 --arrivalseconds 0 --seconds 120
"""
import argparse
import asyncio
import json
import os
import random
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, body=None):
    """One HTTP/1.1 request over a FRESH connection, so each call pays a full
    TLS handshake - which is the point on patch day."""
    started = time.monotonic()
    writer = None
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(host, port, ssl=context), timeout=timeout)
        handshake = time.monotonic() - started
        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)
        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))
        writer.write(('\r\n'.join(headers) + '\r\n\r\n').encode() + payload)
        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, handshake
    except Exception as error:
        return type(error).__name__, time.monotonic() - started, None
    finally:
        if writer is not None:
            try:
                writer.close()
            except Exception:
                pass


def record(results, bucket, status, elapsed, handshake=None, first=False):
    entry = results.setdefault(bucket, {
        'times': [], 'status': {}, 'handshakes': [], 'firsttimes': []})
    entry['times'].append(elapsed * 1000.0)
    if handshake is not None:
        entry['handshakes'].append(handshake * 1000.0)
    if first:
        entry['firsttimes'].append(elapsed * 1000.0)
    key = str(status)
    entry['status'][key] = entry['status'].get(key, 0) + 1


async def agent(index, host, port, context, api_key, results, arrivalseconds):
    """One endpoint coming back from a patch reboot.

    Fires its telemetry report and its first command poll at the same instant,
    mirroring the two workers both starting with no delay.
    """
    if arrivalseconds > 0:
        await asyncio.sleep(random.uniform(0, arrivalseconds))

    hostname = 'LOADTEST%04d' % (index + 1)

    async def telemetry():
        status, elapsed, handshake = await request(
            host, port, context, 'POST', '/api/collector/heartbeat',
            timeout=120.0, api_key=api_key, body={'hostname': hostname})
        record(results, 'telemetry_firstcontact', status, elapsed, handshake,
               first=True)

    telemetry_task = asyncio.ensure_future(telemetry())

    path = '/api/endpoints/commands/next?hostname=%s' % hostname
    first = True
    while not STOP:
        status, elapsed, handshake = await request(
            host, port, context, 'GET', path, timeout=120.0, api_key=api_key)
        record(results, 'fleet', status, elapsed, handshake, first=first)
        if first:
            first = False
        if not isinstance(status, int):
            # Mirrors CommandWorker TransientRetryDelay: fixed 10s, no jitter,
            # so a herd that fails together re-arrives together.
            await asyncio.sleep(10.0)
        else:
            await asyncio.sleep(0.05)

    await telemetry_task


async def ui_user(host, port, context, results, think_seconds):
    while not STOP:
        status, elapsed, handshake = await request(
            host, port, context, 'GET', '/api/dashboard/health', timeout=120.0)
        record(results, 'ui', status, elapsed, handshake)
        await asyncio.sleep(think_seconds)


async def main():
    global STOP
    parser = argparse.ArgumentParser()
    parser.add_argument('--fleet', type=int, default=700)
    parser.add_argument('--users', type=int, default=5)
    parser.add_argument('--arrivalseconds', type=float, default=0.0)
    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)
    arguments = parser.parse_args()

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

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

    results = {}
    tasks = []
    for _ in range(arguments.users):
        tasks.append(asyncio.ensure_future(ui_user(
            arguments.host, arguments.port, context, results,
            arguments.think)))
    # Let the UI baseline establish before the storm hits.
    await asyncio.sleep(5.0)

    burst_started = time.monotonic()
    for index in range(arguments.fleet):
        tasks.append(asyncio.ensure_future(agent(
            index, arguments.host, arguments.port, context, api_key, results,
            arguments.arrivalseconds)))

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

    summary = {
        'fleet': arguments.fleet,
        'arrivalseconds': arguments.arrivalseconds,
        'users': arguments.users,
        'seconds': arguments.seconds,
        'classes': {},
    }
    for bucket, entry in sorted(results.items()):
        times = entry['times']
        block = {
            '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,
        }
        if entry['handshakes']:
            block['handshake_median_ms'] = round(
                statistics.median(entry['handshakes']), 1)
            block['handshake_p95_ms'] = round(
                percentile(entry['handshakes'], 0.95), 1)
            block['handshake_max_ms'] = round(max(entry['handshakes']), 1)
        if entry['firsttimes']:
            block['firstcontact_median_ms'] = round(
                statistics.median(entry['firsttimes']), 1)
            block['firstcontact_p95_ms'] = round(
                percentile(entry['firsttimes'], 0.95), 1)
            block['firstcontact_max_ms'] = round(max(entry['firsttimes']), 1)
        summary['classes'][bucket] = block
    print(json.dumps(summary, indent=2))


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