zhifu-platform/tests/capacity_acceptance.py

import asyncio
import json
import platform
from collections import Counter
from pathlib import Path
from statistics import median, quantiles
from time import perf_counter
from uuid import uuid4

import httpx
from retrieval_evaluation import publish_dataset
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.identity.access_models import Principal, StoreGrant
from zhifu.modules.identity.auth import password_hasher
from zhifu.modules.knowledge.vector import connection


def percentiles(values):
    if not values:
        return None
    boundaries = quantiles(values, n=100, method="inclusive") if len(values) > 1 else values * 99
    return {
        "p50_ms": median(values) * 1000,
        "p95_ms": boundaries[94] * 1000,
        "p99_ms": boundaries[98] * 1000,
    }


async def accounts(context):
    password = "capacity-" + str(uuid4())
    encoded = await asyncio.to_thread(password_hasher.hash, password)
    identifiers = [uuid4() for _ in range(16)]
    async with context["app"].state.sessions.begin() as session:
        for identifier in identifiers:
            session.add(
                Principal(
                    id=identifier.bytes,
                    tenant_id=context["tenant"].bytes,
                    username="capacity-" + identifier.hex,
                    kind="CUSTOMER",
                    password_hash=encoded,
                )
            )
            await session.flush()
            session.add(
                StoreGrant(
                    tenant_id=context["tenant"].bytes,
                    principal_id=identifier.bytes,
                    store_id=context["store"].bytes,
                    permission="orders:read",
                )
            )
    headers = []
    for identifier in identifiers:
        response = await context["client"].post(
            "/openapi/v1/auth/sessions",
            json={
                "tenant_id": str(context["tenant"]),
                "username": "capacity-" + identifier.hex,
                "password": password,
            },
        )
        assert response.status_code == 200, response.text
        headers.append({"Authorization": "Bearer " + response.json()["access_token"]})
        response = await context["call"](
            "PUT",
            f"/internal/v1/orders/{uuid4()}",
            "connector",
            {**context["source"], "subject_id": str(identifier)},
        )
        assert response.status_code == 200, response.text
    return headers


async def measure(
    client, method, path, body, concurrency, identities, seconds=5, maximum_requests=2000
):
    if not 1 <= concurrency <= 64 or not 1 <= maximum_requests <= 100000:
        raise ValueError("并发或请求上限无效")
    samples = []
    issued = 0
    started = perf_counter()
    deadline = started + seconds

    async def worker(index):
        nonlocal issued
        while perf_counter() < deadline and issued < maximum_requests:
            # 当前事件循环内先占用请求名额，所有并发工作共享同一总数上限。
            issued += 1
            begin = perf_counter()
            result = (
                await asyncio.gather(
                    client.request(
                        method, path, json=body, headers=identities[index % len(identities)]
                    ),
                    return_exceptions=True,
                )
            )[0]
            elapsed = perf_counter() - begin
            if isinstance(result, httpx.HTTPError):
                samples.append(
                    {
                        "status": "transport_error",
                        "seconds": elapsed,
                        "error": type(result).__name__,
                    }
                )
            elif isinstance(result, BaseException):
                raise result
            else:
                samples.append({"status": str(result.status_code), "seconds": elapsed})

    async with asyncio.timeout(seconds + 35):
        await asyncio.gather(*(worker(index) for index in range(concurrency)))
    duration = perf_counter() - started
    successes = [row["seconds"] for row in samples if row["status"] == "200"]
    return {
        "concurrency": concurrency,
        "account_count": len(identities),
        "duration_seconds": duration,
        "requests": len(samples),
        "maximum_requests": maximum_requests,
        "maximum_seconds": seconds + 35,
        "statuses": dict(Counter(row["status"] for row in samples)),
        "all_rps": len(samples) / duration,
        "success_rps": len(successes) / duration,
        "success_latency": percentiles(successes),
        "all_latency": percentiles([row["seconds"] for row in samples]),
    }


async def test_capacity(scenario):  # noqa: F811
    try:
        await publish_dataset(scenario)
        identities = await accounts(scenario)
        async with httpx.AsyncClient(
            base_url="http://127.0.0.1:18000",
            headers=scenario["headers"]["customer"],
            timeout=30,
            trust_env=False,
            limits=httpx.Limits(max_connections=32),
        ) as client:
            report = {
                "environment": platform.platform(),
                "python": platform.python_version(),
                "api_workers": 1,
                "account_count": 17,
                "database_pool": "5+5",
                "inference_device": "cpu",
                "documents": 12,
                "stages": [],
            }
            for name, method, path, body, levels, users in (
                (
                    "orders",
                    "GET",
                    f"/openapi/v1/orders?store_id={scenario['store']}",
                    None,
                    (1, 4, 16, 32),
                    [scenario["headers"]["customer"]],
                ),
                (
                    "orders_multi_account",
                    "GET",
                    f"/openapi/v1/orders?store_id={scenario['store']}",
                    None,
                    (4, 16, 32),
                    identities,
                ),
                (
                    "knowledge",
                    "POST",
                    f"/openapi/v1/stores/{scenario['store']}/knowledge/search",
                    {"query": "键盘脏了能放水里洗吗"},
                    (1, 4, 16),
                    [scenario["headers"]["customer"]],
                ),
            ):
                response = await client.request(method, path, json=body)
                assert response.status_code == 200, response.text
                for concurrency in levels:
                    result = await measure(client, method, path, body, concurrency, users)
                    report["stages"].append({"workload": name, **result})
        await asyncio.to_thread(
            Path("docs/verification/capacity-results.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        for stage in report["stages"]:
            assert set(stage["statuses"]) <= {"200", "429"}, stage
            assert stage["statuses"].get("200", 0) > 0, stage
    finally:
        with connection() as (client, collection):
            client.delete(collection, filter=f'tenant == "{scenario["tenant"].hex}"')