84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
"""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()
|