Locust生成指定数量WebSocket用户后暂停生成新用户的方案咨询

1. 解决任务反复执行问题 #

你遇到的单用户下任务反复执行是Locust的默认运行逻辑：用户实例完成当前任务后，会等待你配置的constant(2)即2秒后，再次循环执行任务队列中的任务。同时你将WebSocket连接创建逻辑写在了@task方法内，导致每次执行任务都会新建连接，才会出现每秒重复生成连接的现象。

修复方法如下：

• 如果需要单次对话完成后就终止当前用户，在send_convo方法的末尾添加self.stop()，主动结束当前用户实例的生命周期
• 如果需要长时间维持WebSocket连接，将连接创建逻辑移到on_start方法中，每个用户实例仅在启动时创建一次连接，后续任务复用该连接，避免重复建连

2. 实现1000个用户一次性生成后暂停新建连接 #

Locust本身的Swarm机制在达到你设置的最大用户数后，就不会再生成新的用户实例，你之前遇到的持续生成新连接是任务反复执行、每次任务都新建连接导致的，调整代码逻辑后即可满足需求：

• 调整后设置最大用户数为1000，孵化率为1000，Locust会在1秒内生成全部1000个用户，每个用户仅建一次连接，达到上限后不会生成新用户，自然也不会新建连接
• 如果需要在1000个用户运行一段时间后，再批量生成新的1000个用户，可以通过Locust的Web UI手动调整最大用户数，或者调用Locust的API控制用户数变更

import json
import random
import time
from locust import TaskSet, task, HttpUser, constant, events
from websocket import create_connection

class UserBehavior(TaskSet):
    statements = [
        "Do you like coffee?",
        "What's your favorite book?",
        "Do you invest in crypto?",
        "Who will host the Superbowl next year?",
        "Have you listened to the new Adele?",
        "Coldplay released a new album",
        "I watched the premiere of Succession season 3 last night",
        "Who is your favorite team in the NBA?",
        "I want to buy the new Travis Scott x Jordan shoes",
        "I want a Lamborghini Urus",
        "Have you been to the Philippines?",
        "Did you sign up for a Netflix account?"
    ]

    def on_start(self):
        # 每个用户启动时仅创建一次WebSocket连接
        ws_url = "ws://xx.xx.xx.xx:8080/websocket"
        self.ws = create_connection(ws_url)
        self.end = False

    def on_quit(self):
        # 用户退出时关闭连接
        if not self.ws.closed:
            self.ws.close()

    @task
    def send_convo(self):
        if self.end:
            self.stop()
            return
        
        body = json.dumps({"text": "start blender"})
        self.ws.send(body)
        while True:
            response = self.ws.recv()
            if response != None:
                if "Sorry, this world closed" in response:
                    self.end = True
            break

        if not self.end:
            body = json.dumps({"text": "begin"})
            self.ws.send(body)
            while True:
                response = self.ws.recv()
                if response != None:
                    if "Sorry, this world closed" in response:
                        self.end = True
                        self.ws.close()
                    break

        if not self.end:
            body = json.dumps({"text": random.choice(self.statements)})
            start_at = time.time()
            self.ws.send(body)
            while True:
                response = self.ws.recv()
                if response != None:
                    if "Sorry, this world closed" not in response:
                        response_time = int((time.time() - start_at)*1000)
                        print(f"[BOT]Response: {response}")
                        response_length = len(response)
                        events.request_success.fire(
                            request_type='WebSocket Recv',
                            name='test/ws/echo',
                            response_time=response_time,
                            response_length=response_length,
                        )
                    else:
                        self.end = True
                        self.ws.close()
                    break

        if not self.end:
            body = json.dumps({"text": "[DONE]"})
            self.ws.send(body)
            while True:
                response = self.ws.recv()
                if response != None:
                    if "Sorry, this world closed" in response:
                        self.end = True
                        self.ws.close()
                    break

        if not self.end:
            time.sleep(1)
            body = json.dumps({"text": "EXIT"})
            self.ws.send(body)
            time.sleep(1)
            self.ws.close()
            self.end = True
        
        # 如果需要单次任务完成后就终止用户，取消下面注释即可
        # self.stop()

class WebsiteUser(HttpUser):
    tasks = [UserBehavior]
    # 如果需要维持长连接在线，不需要频繁执行任务可以把等待时间调长，比如constant(300)即每5分钟执行一次任务
    wait_time = constant(2)
    host = "ws://xx.xx.xx.xx:8080/websocket"

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