Locust同步操作内WebSocket异步调用失效问题的解决方法

核心原因分析 #

你遇到的问题本质是gevent协程切换未正常触发，导致监听WebSocket的协程占据CPU资源，Locust主线程的任务调度被阻塞。虽然使用了gevent.spawn()创建协程，但如果监听逻辑未主动让出CPU、或猴子补丁未生效，都会导致主线程无法切换到下一个任务。

具体解决步骤 #

1. 确保猴子补丁优先加载 #

Locust依赖gevent的协程机制，必须在代码最开头打猴子补丁，替换标准库的阻塞IO操作，否则WebSocket的recv()会变成真正的阻塞调用，卡住整个协程上下文。

# 放在所有导入的最前面
from gevent import monkey
monkey.patch_all()

from locust import TaskSet, task, CustomWebsocketLocust
# 其他业务导入...

2. 改造监听协程，主动触发协程切换 #

在监听WebSocket消息的死循环中，每次处理完消息后调用gevent.sleep(0)强制让出CPU，让Locust的任务调度协程有机会执行。

class WebSocketTaskSet(TaskSet):
    def on_start(self):
        # 初始化WebSocket连接
        self.ws = self.locust.connect_websocket()
        # 启动监听协程并保存引用
        self.listen_coroutine = gevent.spawn(self._listen_messages)

    def _listen_messages(self):
        while True:
            try:
                # 接收WebSocket消息
                msg = self.ws.recv()
                if not msg:
                    break
                # 消息处理逻辑（如解析、日志记录）
                print(f"Received: {msg}")
            except Exception as e:
                print(f"Listen error: {str(e)}")
                break
            # 关键：主动让出CPU，触发协程切换
            gevent.sleep(0)

    @task
    def init_task(self):
        # 第一个任务：发送初始化指令
        self.ws.send('{"action": "init"}')
        # 任务逻辑执行完毕立即返回，不阻塞

    @task
    def business_task(self):
        # 第二个任务：发送业务请求
        self.ws.send('{"action": "business_query"}')

3. 避免任务函数内的阻塞操作 #

所有@task装饰的函数必须是非阻塞的，不要在任务内写无限循环或长时间阻塞的逻辑。持续运行的监听、轮询等逻辑，全部放到gevent.spawn()创建的独立协程中，让任务函数快速执行完毕，Locust才能正常调度下一个任务。

4. 统一初始化时机 #

建议在on_start()方法中完成WebSocket连接创建和监听协程启动，而非在第一个任务内。这样监听逻辑会在用户实例初始化时就启动，后续任务可以不受干扰地被调度执行。

5. 验证协程运行状态 #

可以在任务中添加协程状态检查，确认监听协程是否正常存活：

@task
def check_coroutine_status(self):
    print(f"Listen coroutine alive: {self.listen_coroutine.is_alive()}")

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