zhifu-platform/tests/scale_support.py

import asyncio
import json
import os
import subprocess
from pathlib import Path
from uuid import uuid4

import httpx
from sqlalchemy.engine import make_url

from zhifu.core.config import Settings
from zhifu.modules.knowledge.config import KnowledgeSettings
from zhifu.operations.scale_capacity import ScaleRequest, inspect_capacity

ROOT = Path(__file__).resolve().parents[1]


class ScaleEnvironment:
    """仅管理本实例随机命名、带用途标识的 API 容器、网络和网关。"""

    def __init__(self):
        self.prefix = "zhifu-scale-" + uuid4().hex[:10]
        self.network = self.prefix
        self.created = []
        self.commands = 0
        self.monitor_requests = 0
        self.base_url = None
        self.monitor_url = None
        self.addresses = {}
        self.network_created = False
        self.budget = None

    async def docker(self, *arguments, environment=None):
        self.commands += 1
        if self.commands > 100:
            raise ValueError("容量验证 Docker 命令超过 100 次上限")
        result = await asyncio.to_thread(
            subprocess.run, ["docker", *arguments], check=True, capture_output=True,
            text=True, timeout=30, env=environment,
        )
        return result.stdout.strip()

    async def start_api(self, suffix, *, address=None):
        settings = Settings()
        knowledge = KnowledgeSettings()
        environment = dict(os.environ)
        for name in ("database_url", "cache_url", "state_url", "broker_url"):
            url = make_url(getattr(settings, name).get_secret_value())
            if url.host in {"localhost", "127.0.0.1"}:
                host, port = {
                    "database_url": ("mysql", 3306), "cache_url": ("redis-cache", 6379),
                    "state_url": ("redis-state", 6379), "broker_url": ("rabbitmq", 5672),
                }[name]
                url = url.set(host=host, port=port)
            environment["ZHIFU_" + name.upper()] = url.render_as_string(hide_password=False)
        environment.update({
            "ZHIFU_ENVIRONMENT": "test", "ZHIFU_POOL_SIZE": "8", "ZHIFU_MAX_OVERFLOW": "2",
            "ZHIFU_SSE_GLOBAL_LIMIT": "3000", "ZHIFU_SSE_ACTOR_LIMIT": "3",
            "KNOWLEDGE_MILVUS_URL": str(httpx.URL(knowledge.milvus_url).copy_with(
                host="milvus"
            )),
            "KNOWLEDGE_INFERENCE_URL": str(httpx.URL(knowledge.inference_url).copy_with(
                host="host.docker.internal"
            )),
            "KNOWLEDGE_COLLECTION": knowledge.collection, "PYTHONPATH": "/app/src",
        })
        exported = [
            "ZHIFU_DATABASE_URL", "ZHIFU_CACHE_URL", "ZHIFU_STATE_URL", "ZHIFU_BROKER_URL",
            "ZHIFU_ENVIRONMENT", "ZHIFU_POOL_SIZE", "ZHIFU_MAX_OVERFLOW",
            "ZHIFU_SSE_GLOBAL_LIMIT", "ZHIFU_SSE_ACTOR_LIMIT", "KNOWLEDGE_MILVUS_URL",
            "KNOWLEDGE_INFERENCE_URL", "KNOWLEDGE_COLLECTION", "PYTHONPATH",
        ]
        name = self.prefix + "-" + suffix
        arguments = [
            "create", "--name", name, "--network", self.network, "--network-alias", "api",
            "--label", "zhifu.purpose=scale-acceptance", "--memory", "1g", "--cpus", "2",
            "--read-only", "--tmpfs", "/app/.work:uid=10001,gid=10001,mode=0700",
            "--ulimit", "nofile=16384:16384", "--mount",
            f"type=bind,src={ROOT / 'src'},dst=/app/src,readonly",
        ]
        for key in exported:
            arguments.extend(["-e", key])
        if address is not None:
            arguments.extend(["--ip", str(address)])
        arguments.append(os.environ.get("SCALE_API_IMAGE", "zhifu-platform:local"))
        await self.docker(*arguments, environment=environment)
        self.created.append(name)
        await self.docker(
            "network", "connect",
            os.environ.get("SCALE_DEPENDENCY_NETWORK", "zhifu-platform_default"), name,
        )
        await self.docker("start", name)
        metadata = json.loads(await self.docker("inspect", name))[0]
        self.addresses[name] = metadata["NetworkSettings"]["Networks"][self.network]["IPAddress"]
        return name

    async def start(self, replicas=2):
        if not 2 <= replicas <= 16:
            raise ValueError("验证副本数必须在 2–16 之间")
        self.budget = await inspect_capacity(ScaleRequest(replicas=replicas))
        if self.budget["status"] != "READY":
            raise ValueError("当前 MySQL 连接额度不足以同时运行指定 API 副本")
        await self.docker(
            "network", "create", "--label", "zhifu.purpose=scale-acceptance", self.network
        )
        self.network_created = True
        first = await self.start_api("api-a")
        second = await self.start_api("api-b")
        for index in range(2, replicas):
            await self.start_api("api-" + str(index))
        name = self.prefix + "-gateway"
        await self.docker(
            "run", "-d", "--name", name, "--network", self.network,
            "--label", "zhifu.purpose=scale-acceptance", "--memory", "256m", "--cpus", "1",
            "--ulimit", "nofile=32768:32768", "-p", "127.0.0.1::8080", "-p", "127.0.0.1::8404",
            "--mount", f"type=bind,src={ROOT / 'ops/haproxy/runtime.cfg'},"
            "dst=/usr/local/etc/haproxy/haproxy.cfg,readonly", "haproxy:3.2-alpine",
        )
        self.created.append(name)
        metadata = json.loads(await self.docker("inspect", name))[0]["NetworkSettings"]["Ports"]
        self.base_url = "http://127.0.0.1:" + metadata["8080/tcp"][0]["HostPort"]
        self.monitor_url = "http://127.0.0.1:" + metadata["8404/tcp"][0]["HostPort"]
        await self.wait_servers(replicas)
        return first, second

    async def wait_servers(self, expected):
        # 只按真实网关统计的健康状态继续，最多 45 次查询、90 秒等待。
        async with asyncio.timeout(90), httpx.AsyncClient(timeout=3, trust_env=False) as client:
            for _ in range(45):
                self.monitor_requests += 1
                if self.monitor_requests > 180:
                    raise ValueError("监控请求超过 180 次上限")
                response = (await asyncio.gather(
                    client.get(self.monitor_url + "/stats;json"), return_exceptions=True
                ))[0]
                if isinstance(response, httpx.TransportError):
                    await asyncio.sleep(2)
                    continue
                if isinstance(response, BaseException):
                    raise response
                response.raise_for_status()
                # HAProxy JSON 统计采用字段数组，使用标准 JSON 解码。
                rows = [{field["field"]["name"]: field["value"]["value"] for field in row}
                        for row in response.json()]
                servers = [row for row in rows if row.get("pxname") == "api_replicas"
                           and str(row.get("svname", "")).startswith("api")
                           and row.get("status") == "UP"]
                if len(servers) == expected:
                    return {row["svname"] for row in servers}
                await asyncio.sleep(2)
        raise TimeoutError("动态 API 副本未在 90 秒内达到预期健康数量")

    async def close(self):
        # 清理不使用测试请求额度，异常路径仍保证本次资源能够释放。
        for name in reversed(self.created):
            logs = await asyncio.to_thread(
                subprocess.run, ["docker", "logs", "--tail", "150", name],
                check=True, capture_output=True, timeout=15,
            )
            await asyncio.to_thread(
                (ROOT / ".work" / (name + ".log")).write_bytes, logs.stdout + logs.stderr
            )
            await asyncio.to_thread(
                subprocess.run, ["docker", "rm", "-fv", name], check=True,
                capture_output=True, timeout=30,
            )
        if self.network_created:
            await asyncio.to_thread(
                subprocess.run, ["docker", "network", "rm", self.network], check=True,
                capture_output=True, timeout=30,
            )