#WorkBuddy# WebSocket 压测报告全是假的？80 连接虚报成功率 89.69 个百分点，五道关口校出真值

# wsbench_server.py — 自建网关，独立进程运行
import asyncio, json, random, logging
from websockets.asyncio.server import serve

logging.getLogger("websockets").setLevel(logging.CRITICAL)
random.seed(20261114)
S = {"conns": 0, "echoed": 0, "limited": 0}   # 网关侧计数，供对账

async def process_one(ws, msg):                # 并发处理：每条消息独立任务
    await asyncio.sleep(0.002 + random.expovariate(1 / 0.003))
    S["echoed"] += 1
    await ws.send(json.dumps({"id": msg.get("id"), "ok": True}))

async def gateway(ws):
    S["conns"] += 1
    serial = ws.request.path.startswith("/serial")   # 串行模式供 D 阶段用
    tokens, last = 100.0, None                        # 每连接令牌桶
    async for raw in ws:
        now = __import__("time").perf_counter()
        if last is not None:
            tokens = min(100.0, tokens + (now - last) * 400.0)  # 补充 400/s
        last = now
        msg = json.loads(raw)
        if "cmd" in msg:                              # 控制通道：回报统计
            await ws.send(json.dumps({"server_stats": dict(S)})); continue
        if serial:
            await asyncio.sleep(0.005); S["echoed"] += 1
            await ws.send(json.dumps({"id": msg["id"], "ok": True}))
        elif tokens >= 1.0:
            tokens -= 1.0
            asyncio.create_task(process_one(ws, msg))
        else:
            S["limited"] += 1
            asyncio.create_task(ws.send(json.dumps(
                {"id": msg["id"], "ok": False, "code": "RATE_LIMITED"})))

async def main():
    async with serve(gateway, "127.0.0.1", 18914):
        await asyncio.sleep(240)
asyncio.run(main())

# wsbench_run.py（节选）— 治理式压测连接：开环速率 + in-flight 上限 + 按 id 配对
async def eng_conn(offer_rate, dur_s, results):
    ok_ids, err_ids, rtts, in_flight = set(), set(), [], {}
    next_send, seq = time.perf_counter(), 0
    deadline = next_send + dur_s
    async with connect("ws://127.0.0.1:18914/gateway") as ws:
        async def reader():                        # 关口①：按 id 配对，ok/error 分账
            while True:
                rep = json.loads(await ws.recv())
                sid = rep["id"]
                if sid in in_flight:
                    t = in_flight.pop(sid)
                    (ok_ids.add(sid), rtts.append(time.perf_counter() - t)) if rep.get("ok") else err_ids.add(sid)
        rtask = asyncio.create_task(reader())
        while time.perf_counter() < deadline:
            now = time.perf_counter()
            if now < next_send:                    # 关口②：开环——严格按名义速率发
                await asyncio.sleep(min(next_send - now, 0.005)); continue
            if len(in_flight) >= 32:               # 关口③：未决请求上限，背压感知
                await asyncio.sleep(0.005); continue
            in_flight[seq] = time.perf_counter()
            await ws.send(json.dumps({"id": seq, "payload": "x" * 256}))
            seq += 1; next_send += 1.0 / offer_rate
        await asyncio.sleep(1.5); rtask.cancel()   # 排空响应再收尾
    results.append({"ok": len(ok_ids), "err": len(err_ids), "rtts": rtts})

# 建连节奏（关口④）与对账（关口⑤）的主流程骨架
async def main():
    tasks = []
    for start in range(0, 80, 8):              # 阶梯爬坡：每 0.5s 起 8 条连接
        for _ in range(8):
            tasks.append(asyncio.create_task(eng_conn(100, 10, results)))
        await asyncio.sleep(0.5)
    await asyncio.gather(*tasks)
    async with connect("ws://127.0.0.1:18914/gateway") as ws:   # 控制通道对账
        await ws.send(json.dumps({"cmd": "stats"}))
        print(json.loads(await ws.recv())["server_stats"])      # 必须与客户端逐项对上