wx-win-agent/scripts/control-plane-scale-smoke.py

#!/usr/bin/env python3
"""Bounded control-plane readiness/read-task concurrency smoke test."""

from __future__ import annotations

import argparse
import concurrent.futures
import http.client
import json
import os
import ssl
import statistics
import sys
import time
import uuid
from urllib.parse import urlsplit


def request(
    base_url: str,
    path: str,
    method: str,
    body: object | None,
    headers: dict[str, str],
    timeout: float,
):
    url = base_url.rstrip("/") + path
    parts = urlsplit(url)
    if (
        parts.scheme.lower() not in {"http", "https"}
        or not parts.hostname
        or parts.username
        or parts.password
    ):
        raise ValueError("base URL must be an http/https URL without user information")
    target = parts.path or "/"
    if parts.query:
        target += "?" + parts.query
    data = None if body is None else json.dumps(body, ensure_ascii=False).encode()
    connection: http.client.HTTPConnection
    if parts.scheme.lower() == "https":
        connection = http.client.HTTPSConnection(
            parts.hostname,
            parts.port,
            timeout=timeout,
            context=ssl.create_default_context(),
        )
    else:
        connection = http.client.HTTPConnection(
            parts.hostname, parts.port, timeout=timeout
        )
    started = time.perf_counter()
    try:
        connection.request(
            method, target, body=data, headers={"Accept": "application/json", **headers}
        )
        response = connection.getresponse()
        return response.status, time.perf_counter() - started, response.read()
    except OSError as error:
        return 599, time.perf_counter() - started, str(error).encode()
    finally:
        connection.close()


def percentile(values: list[float], p: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, round((len(ordered) - 1) * p))
    return ordered[index]


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--base-url",
        default=os.getenv("WXAGENT_CONTROL_PLANE_URL", "http://127.0.0.1:8090"),
    )
    parser.add_argument("--mode", choices=("ready", "reads"), default="reads")
    parser.add_argument(
        "--web-user", default=os.getenv("WXAGENT_SCALE_WEB_USER", "admin")
    )
    parser.add_argument("--node-id", default=os.getenv("WXAGENT_SCALE_NODE_ID", ""))
    parser.add_argument(
        "--account-id", default=os.getenv("WXAGENT_SCALE_ACCOUNT_ID", "")
    )
    parser.add_argument("--requests", type=int, default=100)
    parser.add_argument("--workers", type=int, default=10)
    parser.add_argument("--timeout", type=float, default=15.0)
    args = parser.parse_args()
    if (
        args.requests < 1
        or args.requests > 5000
        or args.workers < 1
        or args.workers > 100
    ):
        parser.error("requests must be 1..5000 and workers must be 1..100")

    headers: dict[str, str] = {}
    if args.mode == "reads":
        password = os.getenv("WXAGENT_SCALE_WEB_PASSWORD")
        if not password or not args.node_id or not args.account_id:
            parser.error(
                "reads mode requires WXAGENT_SCALE_WEB_PASSWORD, --node-id and --account-id"
            )
        status, _, payload = request(
            args.base_url,
            "/v1/auth/login",
            "POST",
            {"username": args.web_user, "password": password},
            {"Content-Type": "application/json"},
            args.timeout,
        )
        if status != 200:
            print(
                json.dumps(
                    {
                        "ok": False,
                        "phase": "login",
                        "status": status,
                        "body": payload[:256].decode(errors="replace"),
                    }
                )
            )
            return 1
        try:
            login = json.loads(payload)
            access_token = login["access_token"]
        except (UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError) as error:
            print(
                json.dumps(
                    {"ok": False, "phase": "login-response", "error": str(error)}
                )
            )
            return 1
        headers = {
            "Authorization": "Bearer " + access_token,
            "Content-Type": "application/json",
        }

    def one(index: int) -> tuple[int, float]:
        if args.mode == "ready":
            status, elapsed, _ = request(
                args.base_url, "/readyz", "GET", None, {}, args.timeout
            )
        else:
            body = {
                "node_id": args.node_id,
                "account_id": args.account_id,
                "idempotency_key": f"scale-read-{uuid.uuid4().hex}-{index}",
                "limit": 1,
                "offset": 0,
            }
            status, elapsed, _ = request(
                args.base_url, "/v1/reads/sessions", "POST", body, headers, args.timeout
            )
        return status, elapsed

    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
        results = list(executor.map(one, range(args.requests)))
    elapsed = [item[1] for item in results]
    success = sum(200 <= status < 300 for status, _ in results)
    summary = {
        "ok": success == args.requests,
        "mode": args.mode,
        "requests": args.requests,
        "workers": args.workers,
        "success": success,
        "failed": args.requests - success,
        "p50_ms": round(statistics.median(elapsed) * 1000, 2) if elapsed else 0,
        "p95_ms": round(percentile(elapsed, 0.95) * 1000, 2),
        "statuses": {
            str(status): sum(item[0] == status for item in results)
            for status, _ in results
        },
    }
    print(json.dumps(summary, ensure_ascii=False, sort_keys=True))
    return 0 if summary["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())