GuestOps/deploy/capacity_probe.py

#!/usr/bin/env python3
"""Run a bounded, read-only capacity probe against an approved GuestOps sandbox."""

from __future__ import annotations

import argparse
import concurrent.futures
import datetime as dt
import http.cookiejar
import json
import os
from pathlib import Path
import re
import statistics
import time
import urllib.error
import urllib.parse
import urllib.request


PATHS = ("/health/ready", "/api/hotel", "/api/conversations/page")


def require(condition: bool, message: str) -> None:
    if not condition:
        raise ValueError(message)


def percentile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    position = max(0, min(len(ordered) - 1, int(len(ordered) * fraction + 0.999999) - 1))
    return ordered[position]


def summarize(results: list[tuple[bool, float]], concurrency: int) -> dict[str, object]:
    require(len(results) > 0, "At least one probe result is required.")
    latencies = [latency for _, latency in results]
    successes = sum(1 for success, _ in results if success)
    return {
        "concurrency": concurrency,
        "requests": len(results),
        "successes": successes,
        "failures": len(results) - successes,
        "errorRate": round((len(results) - successes) / len(results), 6),
        "latencyMs": {
            "median": round(statistics.median(latencies), 2),
            "p95": round(percentile(latencies, 0.95), 2),
            "maximum": round(max(latencies), 2),
        },
    }


def login(origin: str, email: str, password: str) -> str:
    jar = http.cookiejar.CookieJar()
    opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))
    with opener.open(origin + "/api/session", timeout=15) as response:
        csrf = json.load(response)["csrfToken"]
    body = json.dumps({"email": email, "password": password}).encode()
    request = urllib.request.Request(origin + "/api/auth/login", data=body, method="POST", headers={"Content-Type": "application/json", "X-CSRF-TOKEN": csrf})
    with opener.open(request, timeout=15) as response:
        require(response.status == 200, "Sandbox login failed.")
    cookies = "; ".join(f"{cookie.name}={cookie.value}" for cookie in jar)
    require("guestops.session=" in cookies, "Sandbox did not issue a GuestOps session cookie.")
    return cookies


def request_once(origin: str, cookie: str, number: int) -> tuple[bool, float]:
    path = PATHS[number % len(PATHS)]
    request = urllib.request.Request(origin + path, headers={"Cookie": cookie, "Accept": "application/json"})
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=20) as response:
            success = response.status == 200
            response.read(1024)  # Bound local processing; never retain response or guest content.
    except (OSError, urllib.error.HTTPError):
        success = False
    return success, (time.perf_counter() - started) * 1000


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--origin", required=True)
    parser.add_argument("--requests", type=int, default=100)
    parser.add_argument("--concurrency", type=int, default=5)
    parser.add_argument("--release-commit", required=True)
    parser.add_argument("--release-record-sha256", required=True)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--confirm-sandbox", action="store_true")
    args = parser.parse_args()

    parsed = urllib.parse.urlparse(args.origin)
    require(args.confirm_sandbox, "Use --confirm-sandbox after confirming the target and maintenance window.")
    require(parsed.scheme == "https" and parsed.hostname and parsed.path in ("", "/") and not parsed.query and not parsed.fragment and not parsed.username,
            "Origin must be an HTTPS origin without credentials, path, query or fragment.")
    require(parsed.hostname != "localhost" and not parsed.hostname.startswith("127."), "Use the deployed HTTPS sandbox, not a development server.")
    require(1 <= args.concurrency <= 20 and 1 <= args.requests <= 2000, "Probe bounds are 1–20 concurrent workers and 1–2000 requests.")
    require(re.fullmatch(r"[0-9a-f]{40}", args.release_commit) is not None, "Use a full lowercase release commit SHA.")
    require(re.fullmatch(r"[0-9a-f]{64}", args.release_record_sha256) is not None, "Use the release-record SHA-256.")
    require(not args.output.exists() and args.output.parent.is_dir(), "Output must be a new file in an existing restricted directory.")
    email = os.environ.get("CAPACITY_EMAIL", "")
    password = os.environ.get("CAPACITY_PASSWORD", "")
    require(email and password, "Set CAPACITY_EMAIL and CAPACITY_PASSWORD for a dedicated sandbox staff account.")

    cookie = login(args.origin.rstrip("/"), email, password)
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as pool:
        results = list(pool.map(lambda number: request_once(args.origin.rstrip("/"), cookie, number), range(args.requests)))
    report = {
        "schemaVersion": 1,
        "kind": "guestops-read-only-capacity",
        "recordedAt": dt.datetime.now(dt.timezone.utc).isoformat().replace("+00:00", "Z"),
        "originHost": parsed.hostname,
        "releaseCommit": args.release_commit,
        "releaseRecordSha256": args.release_record_sha256,
        "paths": list(PATHS),
        **summarize(results, args.concurrency),
    }
    descriptor = os.open(args.output, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(descriptor, "w", encoding="utf-8") as output:
        output.write(json.dumps(report, indent=2, sort_keys=True) + "\n")
    print(f"Capacity probe completed: {report['successes']}/{report['requests']} successful; p95 {report['latencyMs']['p95']} ms. Review against the approved target before release.")


if __name__ == "__main__":
    try:
        main()
    except (OSError, ValueError, json.JSONDecodeError) as error:
        print(f"Capacity probe stopped: {error}", file=__import__("sys").stderr)
        raise SystemExit(1)