如何用Locust编写WebSocket压测：保活、响应时间计算及问题排查

我来帮你解决这两个WebSocket压测的问题，先从你提到的响应时间异常的问题说起，再梳理长连接保活的正确实现方式：

问题1：响应时间计算不准确的原因与修复 #

你观察到平均响应时间过长是对的，当前代码里的计时逻辑只统计了WebSocket连接建立的耗时，而这本来就会比单次ping/pong交互久很多，更关键的是你需要统计的ping请求到pong响应的时间完全没有被记录。

问题出在这几点： #

1. 只有connect方法触发了request_success事件，循环里的ping/pong交互没有任何统计
2. 连接建立的耗时被错误当成了业务请求（ping）的响应时间
3. 每次执行searchtest任务都会重新创建连接，不是真正的长连接

修复后的计时逻辑 #

我们需要在发送ping到收到pong的整个过程单独计时，并触发Locust的统计事件：

import time
import json
import websocket
from locust import Locust, TaskSet, task, events

class WebSocketClient(object):
    def __init__(self, host):
        self.host = host
        self.ws = None
        self.connected = False

    def connect(self, url, name):
        if self.connected:
            return True
        start_time = time.time()
        try:
            self.ws = websocket.WebSocket()
            self.ws.connect(url)
        except websocket.WebSocketTimeoutException as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(
                request_type="websocket_connect",
                name=f"{name}_error",
                response_time=total_time,
                exception=e
            )
            self.connected = False
            return False
        else:
            total_time = int((time.time() - start_time) * 1000)
            events.request_success.fire(
                request_type="websocket_connect",
                name=f"{name}_success",
                response_time=total_time,
                response_length=0
            )
            self.connected = True
            return True

    def send_ping(self, name):
        if not self.connected:
            return False
        ping_data = json.dumps({"action": "ping"})
        start_time = time.time()
        try:
            self.ws.send(ping_data)
            response = self.ws.recv()
            response_json = json.loads(response)
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(
                request_type="websocket_ping",
                name=f"{name}_error",
                response_time=total_time,
                exception=e
            )
            self.connected = False
            return False
        else:
            total_time = int((time.time() - start_time) * 1000)
            if response_json.get("data") == "pong":
                events.request_success.fire(
                    request_type="websocket_ping",
                    name=f"{name}_success",
                    response_time=total_time,
                    response_length=len(response)
                )
                return True
            else:
                events.request_failure.fire(
                    request_type="websocket_ping",
                    name=f"{name}_invalid_response",
                    response_time=total_time,
                    exception=ValueError("Received non-pong response")
                )
                return False

    def close(self):
        if self.ws:
            self.ws.close()
        self.connected = False

class WebsocketLocust(Locust):
    def __init__(self, *args, **kwargs):
        super(WebsocketLocust, self).__init__(*args, **kwargs)
        self.client = WebSocketClient(self.host)

class SearchTest(TaskSet):
    def on_start(self):
        # 任务开始时只建立一次长连接
        self.url = 'wss://api.bbxapp.vip/v1/ifcontract/realTime'
        if not self.client.connect(self.url, name='chat_room'):
            # 连接失败则停止当前用户任务
            self.interrupt()

    @task
    def ping_pong_test(self):
        # 持续发送ping保活并统计响应时间
        self.client.send_ping(name='ping_pong')

    def on_stop(self):
        # 任务结束时关闭连接
        self.client.close()

关键修改点： #

• 新增send_ping方法，单独统计从发送ping到收到pong的耗时
• 区分websocket_connect和websocket_ping两种请求类型的统计，避免混淆
• 用connected状态标记维护长连接，防止重复创建连接

问题2：WebSocket长连接保活的正确实现 #

你的原代码每次执行searchtest任务都会重新调用connect，这其实是短连接而非长连接。正确的长连接实现应该：

1. 在TaskSet的on_start方法中只建立一次连接，每个Locust用户对应一个长连接
2. 在task中循环发送ping包（或者按固定间隔发送），维持连接活跃
3. 在on_stop方法中关闭连接，清理资源

上面的修改后代码已经实现了这个逻辑：

• on_start：每个用户启动时建立一次WebSocket连接，失败则中断任务
• ping_pong_test任务：持续发送ping包，触发保活和响应时间统计
• on_stop：用户停止时关闭连接

修改后的统计效果 #

现在你会在Locust的输出中看到两类统计：

• chat_room_success：连接建立的耗时（这一项本来就会偏长，属于初始化开销）
• ping_pong_success：每次ping/pong交互的响应时间，这才是你需要关注的业务请求耗时

内容的提问来源于stack exchange，提问作者Zoe Zou