zhifu-platform/tests/dynamic_scale_acceptance.py

import asyncio
import hashlib
import json
import secrets
from collections import Counter
from datetime import timedelta
from ipaddress import ip_network
from pathlib import Path
from time import perf_counter
from uuid import uuid4

import httpx
from capacity_acceptance import accounts, measure, percentiles
from scale_support import ScaleEnvironment
from sqlalchemy import func, select
from sse_capacity_acceptance import wait_released
from test_after_sales import scenario  # noqa: F401

from zhifu.modules.after_sales.models import Ticket
from zhifu.modules.agent_runtime.models import ModelAttempt
from zhifu.modules.identity.access_models import AccessSession, Principal, StoreGrant
from zhifu.modules.identity.auth import now


async def test_dynamic_discovery_replacement_and_idempotency(scenario):  # noqa: F811
    environment = ScaleEnvironment()
    report = {"generation_attempts": 0, "maximum_business_requests": 240,
              "maximum_monitor_requests": 180, "maximum_seconds": 360, "phases": [],
              "status": "INCOMPLETE"}
    started = perf_counter()
    try:
        async with asyncio.timeout(360):
            first, _ = await environment.start()
            report["database_budget"] = environment.budget
            async with httpx.AsyncClient(
                base_url=environment.base_url, headers=scenario["headers"]["customer"],
                timeout=15, trust_env=False,
            ) as client:
                key = uuid4().hex
                for phase, expected in (("both", 2), ("failed", 1), ("replaced", 2)):
                    if phase == "failed":
                        await environment.docker("kill", first)
                        recovery_started = perf_counter()
                    elif phase == "replaced":
                        network = json.loads(await environment.docker(
                            "network", "inspect", environment.network
                        ))[0]
                        address = ip_network(network["IPAM"]["Config"][0]["Subnet"])[10]
                        replacement = await environment.start_api("api-c", address=address)
                        assert environment.addresses[replacement] != environment.addresses[first]
                        recovery_started = perf_counter()
                    servers = await environment.wait_servers(expected)
                    reads = await asyncio.gather(*(
                        client.get(f"/openapi/v1/orders?store_id={scenario['store']}")
                        for _ in range(64)
                    ))
                    writes = await asyncio.gather(*(
                        client.post("/openapi/v1/after-sales/proposals",
                                    headers={"Idempotency-Key": key}, json=scenario["proposal"])
                        for _ in range(16)
                    ))
                    assert all(row.status_code == 200 for row in reads)
                    assert all(row.status_code == 201 for row in writes), [r.text for r in writes]
                    observed = {row.headers["X-Zhifu-Replica"] for row in reads + writes}
                    assert observed == servers
                    assert len({row.json()["ticket_id"] for row in writes}) == 1
                    report["phases"].append({
                        "phase": phase, "replicas": sorted(servers),
                        "statuses": dict(Counter(row.status_code for row in reads + writes)),
                        "read_replicas": dict(Counter(
                            row.headers["X-Zhifu-Replica"] for row in reads
                        )),
                        "recovery_seconds": perf_counter() - recovery_started
                        if phase != "both" else None,
                    })
                async with scenario["app"].state.sessions() as session:
                    tickets = await session.scalar(select(func.count()).select_from(Ticket).where(
                        Ticket.tenant_id == scenario["tenant"].bytes
                    ))
                assert tickets == 1
                report.update(status="PASSED", tickets=tickets, addresses=environment.addresses)
    finally:
        await environment.close()
        report.update(seconds=perf_counter() - started, temporary_resources_remaining=0,
                      docker_commands=environment.commands,
                      monitor_requests=environment.monitor_requests)
        await asyncio.to_thread(Path(".work/dynamic-discovery.json").write_text,
            json.dumps(report, indent=2) + "\n", encoding="utf-8"
        )


async def sse_accounts(context, count):
    sessions = context["app"].state.sessions
    identities = []
    actors = [uuid4() for _ in range(count)]
    async with sessions.begin() as session:
        customer = await session.get(Principal, context["accounts"]["customer"].bytes)
        session.add_all([Principal(
            id=actor.bytes, tenant_id=context["tenant"].bytes, username="scale-" + actor.hex,
            password_hash=customer.password_hash, kind="CUSTOMER",
        ) for actor in actors])
        await session.flush()
        for actor in actors:
            token = secrets.token_urlsafe(32)
            session.add(StoreGrant(
                tenant_id=context["tenant"].bytes, principal_id=actor.bytes,
                store_id=context["store"].bytes, permission="orders:read",
            ))
            session.add(AccessSession(
                token_hash=hashlib.sha256(token.encode()).digest(),
                tenant_id=context["tenant"].bytes, principal_id=actor.bytes,
                expires_at=now() + timedelta(minutes=15),
            ))
            identities.append({"Authorization": "Bearer " + token})
    return identities


async def test_three_thousand_sse_connections(scenario):  # noqa: F811
    environment = ScaleEnvironment()
    state = scenario["app"].state.state
    assert await state.zcard("sse:global") == 0, "验收开始时需要无其他 SSE 连接"
    report = {"connections_target": 3000, "account_count": 1001,
              "maximum_sse_requests": 4005, "admission_concurrency": 256,
              "stream_clients": 8, "connections_per_client": 400,
              "maximum_seconds": 360, "api_replicas": 8,
              "generation_attempts": 0, "status": "INCOMPLETE"}
    started = perf_counter()
    tasks = []
    finish = asyncio.Event()
    latencies = []
    clients = []
    try:
        async with asyncio.timeout(360):
            await environment.start(replicas=8)
            report["database_budget"] = environment.budget
            identities = await sse_accounts(scenario, 1001)
            ready = asyncio.Queue()
            admission = asyncio.Semaphore(256)
            preparation = asyncio.Semaphore(64)
            clients = [httpx.AsyncClient(
                base_url=environment.base_url, timeout=35, trust_env=False,
                limits=httpx.Limits(max_connections=400, max_keepalive_connections=400),
            ) for _ in range(8)]
            async with httpx.AsyncClient(
                base_url=environment.base_url, timeout=35, trust_env=False,
                limits=httpx.Limits(max_connections=256, max_keepalive_connections=256),
            ) as client:
                async def create(headers):
                    async with preparation:
                        response = await client.post(
                            "/openapi/v1/conversations", json={"store_id": str(scenario["store"])},
                            headers={**headers, "Idempotency-Key": uuid4().hex},
                        )
                        assert response.status_code == 201, response.text
                        return "/openapi/v1/conversations/" + response.json()["conversation_id"]

                paths = await asyncio.gather(*(create(headers) for headers in identities))

                async def subscribe(index):
                    begin = perf_counter()
                    async with admission:
                        selected = clients[index % len(clients)]
                        request = selected.build_request(
                            "GET", paths[index // 3] + "/events", headers=identities[index // 3]
                        )
                        response = await selected.send(request, stream=True)
                    try:
                        assert response.status_code == 200, (await response.aread()).decode()
                        observed = False
                        async for _chunk in response.aiter_raw():
                            if not observed:
                                observed = True
                                latencies.append(perf_counter() - begin)
                                ready.put_nowait(index)
                            if finish.is_set():
                                break
                        assert finish.is_set(), "保持窗口内 SSE 连接提前结束"
                    finally:
                        await response.aclose()

                async with asyncio.TaskGroup() as group:
                    for index in range(3000):
                        tasks.append(group.create_task(subscribe(index)))
                    for _ in range(3000):
                        await ready.get()
                    report["admission_latency"] = percentiles(latencies)
                    assert await state.zcard("sse:global") == 3000
                    response = await client.get(paths[-1] + "/events", headers=identities[-1])
                    assert response.status_code == 429 and response.json()["code"] == "SSE_LIMIT"
                    response = await client.get(paths[0] + "/events", headers=identities[0])
                    assert response.status_code == 429
                    await asyncio.sleep(10)
                    assert await state.zcard("sse:global") == 3000
                    assert (await client.get("/health/ready")).status_code == 200
                    disconnected_at = perf_counter()
                    finish.set()
                    for task in tasks:
                        task.cancel()
                await wait_released(state, 65)
                report["disconnect_release_seconds"] = perf_counter() - disconnected_at
                report["recovery_observation_window_seconds"] = 65
                async with client.stream("GET", paths[0] + "/events", headers=identities[0]) as row:
                    assert row.status_code == 200
            await wait_released(state, 65)
            report.update(sse_status="PASSED", connections=3000, held_seconds=10,
                          http_429=2, disconnect_recovery=True)
            async with scenario["app"].state.sessions() as session:
                attempts = await session.scalar(
                    select(func.count()).select_from(ModelAttempt).where(
                        ModelAttempt.tenant_id == scenario["tenant"].bytes
                    )
                )
            assert attempts == 0
            report["status"] = "PASSED"
    finally:
        finish.set()
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        await asyncio.gather(*(client.aclose() for client in clients))
        try:
            await wait_released(state, 65)
        finally:
            await environment.close()
            report.update(seconds=perf_counter() - started, observed_connections=len(latencies),
                          temporary_resources_remaining=0)
            await asyncio.to_thread(Path(".work/dynamic-capacity.json").write_text,
                json.dumps(report, indent=2) + "\n", encoding="utf-8"
            )


async def test_bounded_orders_with_dynamic_replicas(scenario):  # noqa: F811
    environment = ScaleEnvironment()
    report = {"maximum_load_requests": 6000, "maximum_seconds": 180, "api_replicas": 2,
              "generation_attempts": 0, "status": "INCOMPLETE"}
    try:
        async with asyncio.timeout(180):
            await environment.start()
            report["database_budget"] = environment.budget
            users = await accounts(scenario)
            async with httpx.AsyncClient(
                base_url=environment.base_url, timeout=15, trust_env=False,
                limits=httpx.Limits(max_connections=32),
            ) as client:
                report["orders"] = await measure(
                    client, "GET", f"/openapi/v1/orders?store_id={scenario['store']}",
                    None, 32, users, seconds=30, maximum_requests=6000,
                )
            assert set(report["orders"]["statuses"]) == {"200"}
            report["status"] = "PASSED"
    finally:
        await environment.close()
        report["temporary_resources_remaining"] = 0
        await asyncio.to_thread(
            Path(".work/dynamic-orders.json").write_text,
            json.dumps(report, indent=2) + "\n", encoding="utf-8",
        )