test(#94): 性能压测套件 - Locust REST API梯度压测 + WebSocket长连接并发压测(100-5000) + 压测方案文档

tests/performance/websocket_stress.py

@@ -0,0 +1,83 @@
+"""WebSocket 长连接并发压测（Issue #94）
+
+支持 100 - 5000 并发连接梯度，统计建连成功率、消息往返延迟（P50/P95/P99）。
+
+运行：
+    python tests/performance/websocket_stress.py --url ws://127.0.0.1:8000/ws \
+        --concurrency 500 --duration 60
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import statistics
+import time
+
+try:
+    import websockets
+except ImportError:  # pragma: no cover
+    websockets = None
+
+
+class Stats:
+    def __init__(self) -> None:
+        self.connected = 0
+        self.failed = 0
+        self.latencies: list[float] = []
+
+    def summary(self) -> str:
+        if not self.latencies:
+            lat = "无消息样本"
+        else:
+            ls = sorted(self.latencies)
+            p = lambda q: ls[min(len(ls) - 1, int(q * len(ls)))]
+            lat = (f"P50={p(0.5)*1000:.1f}ms P95={p(0.95)*1000:.1f}ms "
+                   f"P99={p(0.99)*1000:.1f}ms avg={statistics.mean(ls)*1000:.1f}ms")
+        return (f"建连成功={self.connected} 失败={self.failed} | 往返延迟 {lat}")
+
+
+async def worker(url: str, duration: float, stats: Stats) -> None:
+    try:
+        async with websockets.connect(url, open_timeout=10) as ws:
+            stats.connected += 1
+            end = time.time() + duration
+            while time.time() < end:
+                t0 = time.perf_counter()
+                await ws.send('{"type":"ping"}')
+                try:
+                    await asyncio.wait_for(ws.recv(), timeout=5)
+                    stats.latencies.append(time.perf_counter() - t0)
+                except asyncio.TimeoutError:
+                    pass
+                await asyncio.sleep(1)
+    except Exception:
+        stats.failed += 1
+
+
+async def main_async(url: str, concurrency: int, duration: float, ramp: float) -> None:
+    stats = Stats()
+    tasks = []
+    for i in range(concurrency):
+        tasks.append(asyncio.create_task(worker(url, duration, stats)))
+        await asyncio.sleep(ramp)  # 平滑加压，避免瞬时 SYN 风暴
+    await asyncio.gather(*tasks)
+    print(stats.summary())
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--url", default="ws://127.0.0.1:8000/ws")
+    ap.add_argument("--concurrency", type=int, default=100, help="并发连接数(100-5000)")
+    ap.add_argument("--duration", type=float, default=60, help="每连接压测时长(秒)")
+    ap.add_argument("--ramp", type=float, default=0.01, help="建连间隔(秒)")
+    args = ap.parse_args()
+    if websockets is None:
+        raise SystemExit("请先 pip install websockets")
+    if not 1 <= args.concurrency <= 5000:
+        raise SystemExit("concurrency 需在 1-5000 之间")
+    asyncio.run(main_async(args.url, args.concurrency, args.duration, args.ramp))
+
+
+if __name__ == "__main__":
+    main()