335 lines
12 KiB
Python
335 lines
12 KiB
Python
"""通过真实 HTTP、MySQL 和 Redis 测量注册并发性能并清理数据"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import math
|
|
import os
|
|
import secrets
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import psutil
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
|
|
def summarize(samples: list[float]) -> dict:
|
|
"""汇总毫秒延迟并保留尾部耗时"""
|
|
values = sorted(samples)
|
|
if not values:
|
|
return {}
|
|
return {
|
|
"count": len(values),
|
|
"mean_ms": round(sum(values) / len(values), 2),
|
|
"p50_ms": round(values[math.ceil(len(values) * 0.50) - 1], 2),
|
|
"p95_ms": round(values[math.ceil(len(values) * 0.95) - 1], 2),
|
|
"p99_ms": round(values[math.ceil(len(values) * 0.99) - 1], 2),
|
|
"max_ms": round(values[-1], 2),
|
|
}
|
|
|
|
|
|
async def cleanup(prefix: str) -> dict:
|
|
"""删除本轮用户及其 Redis 会话并核验残留"""
|
|
from sqlalchemy import delete, func, select
|
|
|
|
from app.core.config import get_settings
|
|
from app.core.database import async_session_factory, close_database
|
|
from app.core.redis import close_redis, redis_client
|
|
from app.model.user import User
|
|
|
|
try:
|
|
async with async_session_factory() as session:
|
|
condition = User.username.startswith(prefix, autoescape=True)
|
|
ids = set(map(str, (await session.scalars(select(User.id).where(condition))).all()))
|
|
keys_deleted = 0
|
|
pattern = f"{get_settings().redis_key_prefix}:session:*"
|
|
async for key in redis_client.scan_iter(match=pattern, count=200):
|
|
if await redis_client.get(key) in ids:
|
|
keys_deleted += await redis_client.delete(key)
|
|
result = await session.execute(delete(User).where(condition))
|
|
await session.commit()
|
|
remaining = await session.scalar(
|
|
select(func.count()).select_from(User).where(condition)
|
|
)
|
|
remaining_sessions = 0
|
|
async for key in redis_client.scan_iter(match=pattern, count=200):
|
|
if await redis_client.get(key) in ids:
|
|
remaining_sessions += 1
|
|
return {
|
|
"users_deleted": result.rowcount,
|
|
"sessions_deleted": keys_deleted,
|
|
"users_remaining": remaining,
|
|
"sessions_remaining": remaining_sessions,
|
|
}
|
|
finally:
|
|
await close_redis()
|
|
await close_database()
|
|
|
|
|
|
async def baseline() -> dict:
|
|
"""分别测量热连接往返和本机密码计算作为诊断基线"""
|
|
from sqlalchemy import text
|
|
|
|
from app.core.database import engine
|
|
from app.core.redis import redis_client
|
|
from app.core.security import hash_password
|
|
|
|
results = {}
|
|
async with engine.connect() as connection:
|
|
await connection.execute(text("SELECT 1"))
|
|
samples = []
|
|
for _ in range(10):
|
|
start = time.perf_counter()
|
|
await connection.execute(text("SELECT 1"))
|
|
samples.append((time.perf_counter() - start) * 1000)
|
|
results["mysql_select_1"] = summarize(samples)
|
|
await redis_client.ping()
|
|
samples = []
|
|
for _ in range(10):
|
|
start = time.perf_counter()
|
|
await redis_client.ping()
|
|
samples.append((time.perf_counter() - start) * 1000)
|
|
results["redis_ping"] = summarize(samples)
|
|
samples = []
|
|
for _ in range(10):
|
|
start = time.perf_counter()
|
|
await hash_password("BenchmarkOnly123")
|
|
samples.append((time.perf_counter() - start) * 1000)
|
|
results["local_password_hash"] = summarize(samples)
|
|
return results
|
|
|
|
|
|
async def stage(
|
|
client: httpx.AsyncClient,
|
|
process: psutil.Process,
|
|
prefix: str,
|
|
password: str,
|
|
token: str,
|
|
concurrency: int,
|
|
seconds: float,
|
|
) -> dict:
|
|
"""并行注册并采样已登录请求延迟与服务资源"""
|
|
latencies: list[float] = []
|
|
probe_latencies: list[float] = []
|
|
errors: Counter = Counter()
|
|
probe_errors: Counter = Counter()
|
|
cpu: list[float] = []
|
|
memory: list[float] = []
|
|
sequence = 0
|
|
completed = 0
|
|
stop = asyncio.Event()
|
|
start = time.perf_counter()
|
|
deadline = start + seconds
|
|
processes = [process, *process.children(recursive=True)]
|
|
for item in processes:
|
|
item.cpu_percent()
|
|
|
|
async def worker() -> None:
|
|
"""保持一个在途注册请求直到本阶段停止发起请求"""
|
|
nonlocal sequence, completed
|
|
while time.perf_counter() < deadline:
|
|
sequence += 1
|
|
name = f"{prefix}{concurrency}_{sequence}"
|
|
tick = time.perf_counter()
|
|
try:
|
|
response = await client.post(
|
|
"/api/auth/register",
|
|
json={
|
|
"username": name,
|
|
"email": f"{name}@example.com",
|
|
"password": password,
|
|
},
|
|
)
|
|
if response.status_code == 200 and response.json().get("code") == 0:
|
|
completed += 1
|
|
else:
|
|
errors[str(response.status_code)] += 1
|
|
except httpx.HTTPError as exc:
|
|
errors[type(exc).__name__] += 1
|
|
latencies.append((time.perf_counter() - tick) * 1000)
|
|
|
|
async def probe() -> None:
|
|
"""连续探测涉及数据库和 Redis 的普通认证请求"""
|
|
while not stop.is_set():
|
|
tick = time.perf_counter()
|
|
try:
|
|
response = await client.get(
|
|
"/api/users/me", headers={"Authorization": f"Bearer {token}"}
|
|
)
|
|
if response.status_code != 200 or response.json().get("code") != 0:
|
|
probe_errors[str(response.status_code)] += 1
|
|
except httpx.HTTPError as exc:
|
|
probe_errors[type(exc).__name__] += 1
|
|
probe_latencies.append((time.perf_counter() - tick) * 1000)
|
|
cpu.append(sum(item.cpu_percent() for item in processes if item.is_running()))
|
|
memory.append(
|
|
sum(item.memory_info().rss for item in processes if item.is_running()) / 1024**2
|
|
)
|
|
await asyncio.sleep(0.1)
|
|
|
|
probe_task = asyncio.create_task(probe())
|
|
try:
|
|
await asyncio.gather(*(worker() for _ in range(concurrency)))
|
|
finally:
|
|
elapsed = time.perf_counter() - start
|
|
stop.set()
|
|
await probe_task
|
|
return {
|
|
"concurrency": concurrency,
|
|
"elapsed_s": round(elapsed, 2),
|
|
"success": completed,
|
|
"errors": dict(errors),
|
|
"success_rps": round(completed / elapsed, 2),
|
|
"registration": summarize(latencies),
|
|
"authenticated_probe": summarize(probe_latencies),
|
|
"probe_errors": dict(probe_errors),
|
|
"cpu_mean_percent_one_core": round(sum(cpu) / len(cpu), 1),
|
|
"cpu_peak_percent_one_core": round(max(cpu), 1),
|
|
"rss_peak_mb": round(max(memory), 1),
|
|
}
|
|
|
|
|
|
async def run(args: argparse.Namespace) -> None:
|
|
"""启动独立单进程服务完成压测并确保结束后清理"""
|
|
prefix = f"bench{secrets.token_hex(6)}"
|
|
password = secrets.token_urlsafe(20) + "a9"
|
|
with socket.socket() as bound:
|
|
bound.bind(("127.0.0.1", 0))
|
|
port = bound.getsockname()[1]
|
|
report = {
|
|
"run_prefix": prefix,
|
|
"workers": 1,
|
|
"logical_cpus": os.cpu_count(),
|
|
"physical_cpus": psutil.cpu_count(logical=False),
|
|
"stages": [],
|
|
}
|
|
process = await asyncio.to_thread(
|
|
subprocess.Popen,
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"uvicorn",
|
|
"app.main:app",
|
|
"--host",
|
|
"127.0.0.1",
|
|
"--port",
|
|
str(port),
|
|
"--no-access-log",
|
|
"--log-level",
|
|
"error",
|
|
],
|
|
cwd=ROOT,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
|
)
|
|
try:
|
|
async with httpx.AsyncClient(
|
|
base_url=f"http://127.0.0.1:{port}",
|
|
timeout=60,
|
|
limits=httpx.Limits(max_connections=256),
|
|
trust_env=False,
|
|
) as client:
|
|
for _ in range(150):
|
|
if process.poll() is not None:
|
|
raise RuntimeError("压测服务启动失败，请检查依赖和配置")
|
|
try:
|
|
if (await client.get("/openapi.json")).status_code == 200:
|
|
break
|
|
except httpx.HTTPError:
|
|
pass
|
|
await asyncio.sleep(0.2)
|
|
else:
|
|
raise RuntimeError("压测服务启动超时")
|
|
report["baseline"] = await baseline()
|
|
print(json.dumps(report["baseline"], ensure_ascii=False), flush=True)
|
|
name = prefix + "probe"
|
|
response = await client.post(
|
|
"/api/auth/register",
|
|
json={
|
|
"username": name,
|
|
"email": f"{name}@example.com",
|
|
"password": password,
|
|
},
|
|
)
|
|
if response.status_code != 200:
|
|
raise RuntimeError(f"探针用户注册失败 HTTP {response.status_code}")
|
|
response = await client.post(
|
|
"/api/auth/login", json={"account": name, "password": password}
|
|
)
|
|
if response.status_code != 200:
|
|
raise RuntimeError(f"探针登录失败 HTTP {response.status_code}")
|
|
token = response.json()["data"]["token"]
|
|
samples = []
|
|
for _ in range(10):
|
|
tick = time.perf_counter()
|
|
response = await client.get(
|
|
"/api/users/me", headers={"Authorization": f"Bearer {token}"}
|
|
)
|
|
if response.status_code != 200:
|
|
raise RuntimeError("认证基线请求失败")
|
|
samples.append((time.perf_counter() - tick) * 1000)
|
|
report["baseline"]["authenticated_request"] = summarize(samples)
|
|
for concurrency in args.concurrency:
|
|
result = await stage(
|
|
client,
|
|
psutil.Process(process.pid),
|
|
prefix,
|
|
password,
|
|
token,
|
|
concurrency,
|
|
args.seconds,
|
|
)
|
|
report["stages"].append(result)
|
|
print(json.dumps(result, ensure_ascii=False), flush=True)
|
|
except Exception as exc:
|
|
report["error_type"] = type(exc).__name__
|
|
raise
|
|
finally:
|
|
children = []
|
|
if process.poll() is None:
|
|
children = psutil.Process(process.pid).children(recursive=True)
|
|
for child in children:
|
|
child.terminate()
|
|
process.terminate()
|
|
await asyncio.to_thread(process.wait)
|
|
await asyncio.to_thread(psutil.wait_procs, children, timeout=10)
|
|
try:
|
|
report["cleanup"] = await cleanup(prefix)
|
|
except Exception as exc:
|
|
report["cleanup"] = {"error_type": type(exc).__name__, "verified": False}
|
|
raise
|
|
finally:
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(
|
|
json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
|
|
)
|
|
print(json.dumps(report["cleanup"], ensure_ascii=False), flush=True)
|
|
|
|
|
|
def main() -> None:
|
|
"""解析压测档位和时长并运行真实依赖测试"""
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--concurrency", nargs="+", type=int, default=[1, 4, 8, 16, 32, 64])
|
|
parser.add_argument("--seconds", type=float, default=12)
|
|
parser.add_argument("--output", type=Path, default=ROOT / "artifacts/load-test.json")
|
|
args = parser.parse_args()
|
|
if args.seconds <= 0 or any(value < 1 or value > 200 for value in args.concurrency):
|
|
parser.error("时长必须为正数，并发必须在 1 到 200 之间")
|
|
try:
|
|
asyncio.run(run(args))
|
|
except Exception as exc:
|
|
print(f"压测失败：{type(exc).__name__}，请检查报告和服务配置", file=sys.stderr)
|
|
raise SystemExit(1) from None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|