如何在Locust主从配置下实现负载测试前的用户批量注册预处理

你原来的逻辑在单机模式跑通的核心，是所有代码、内存变量都在同一个进程里，test_start事件只触发一次，生成的token_list可以被所有用户实例直接访问。换成master/worker分布式模式后，有两个本质断裂点：

• 事件触发逻辑变化：test_start事件会在master进程、每个worker进程上各自独立触发一次，原来的批量注册逻辑会被重复执行N+1次（1个master+N个worker），产生大量冗余的无效注册请求。
• 进程内存隔离：每个worker有独立的内存空间，你在某个进程里生成的token列表，其他进程完全访问不到，会出现用户启动时从空的token列表里pop数据、直接报错的问题。

整体逻辑要做三个调整，完全适配你当前使用的locust==1.4.3版本的API：

• 用内置的节点判断逻辑，把批量注册的预处理代码只在master/单机节点执行一次，从根源上避免重复注册。
• 利用Locust自带的ZMQ跨节点通信通道，把master生成的完整token列表，推送到所有连接的worker节点上，存在每个worker的本地内存里。
• worker侧监听自定义消息，收到token列表之后再允许用户实例启动，保持和原来单机模式一致的取token逻辑，不需要改动原有测试任务的代码。

import logging
import uuid
import requests
from locust import between, task, events
from locust.contrib.fasthttp import FastHttpUser
# 导入runner类用于判断当前运行节点类型
from locust.runners import WorkerRunner, MasterRunner

# 全局token列表，每个进程独立维护本地副本
token_list = []

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    global token_list
    # 跳过worker节点的预处理逻辑，只在master/单机节点执行
    if isinstance(environment.runner, WorkerRunner):
        logging.info("当前为worker节点，跳过本地批量注册逻辑")
        return

    logging.info("开始执行批量账号注册预处理")
    account_list = []
    base_user_id = str(uuid.uuid4())
    # 从环境对象读取全集群的目标用户总数
    total_user_count = environment.runner.target_user_count
    idx = 0
    max_batch = 499
    account_path = f'{environment.host}/v1/accounts'

    # 分批次执行批量账号注册（和原有逻辑完全保持一致）
    for _ in range(((total_user_count - 1) // max_batch) + 1):
        batch_num = min(max_batch, total_user_count)
        external_user_ids = [f'{base_user_id}_{n}' for n in range(idx, idx + batch_num)]
        payload = create_initial_account_api_request(external_user_ids)
        access_token = create_initial_token()
        resp = requests.post(account_path, json=payload, headers=access_token)
        logging.info(f"批量注册响应码：{resp.status_code}, 响应内容：{resp.text}")
        res_json = resp.json()
        account_list += res_json['response']
        idx += batch_num
        total_user_count -= batch_num

    # 为每个注册好的账号获取业务token（和原有逻辑完全保持一致）
    auth_header = make_auth_header()
    token_path = f'{environment.host}/v1/oauth/token'
    for account in account_list:
        external_user_id = account['body']['external_user_id']
        payload = create_initial_user_token_api_request(external_user_id)
        resp = requests.post(token_path, data=payload, headers=auth_header)
        logging.info(f"换绑token响应码：{resp.status_code}, 响应内容：{resp.text}")
        token_list.append(resp.json())

    # 如果当前是master节点，把生成好的token列表推送给所有连接的worker
    if isinstance(environment.runner, MasterRunner):
        logging.info(f"预处理完成，共生成{len(token_list)}个token，开始推送到所有worker节点")
        environment.runner.send_message(
            msg_type="init_task_tokens",
            data=token_list
        )

@events.message.add_listener
def on_message_received(environment, msg, **kwargs):
    # worker节点监听来自master的自定义消息
    global token_list
    if isinstance(environment.runner, WorkerRunner):
        if msg["type"] == "init_task_tokens":
            token_list = msg["data"]
            logging.info(f"从master收到{len(token_list)}个初始化token，本地预处理完成")

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    global token_list
    token_list = []

class MyUser(FastHttpUser):
    wait_time = between(3, 6)

    def on_start(self):
        global token_list
        if not token_list:
            logging.error("本地token列表为空，无法初始化用户实例")
            return
        token_data = token_list.pop()
        logging.info(f"用户实例初始化，账号信息：{str(token_data)}")
        self.account_id = token_data['account_id']
        self.refresh_token = token_data['refresh_token']
        self.access_token = token_data['access_token']

    @task
    def normal_task(self):
        self.refresh_task()
        self.first_task()
        self.second_task()

• 启动测试前要确保所有worker节点和master的连接正常，如果有未连接的worker，master推送token的时候会漏发，导致对应worker启动时拿不到token报错。
• 你原来的shell脚本里LOCUS_OPTS变量名有拼写错误，建议改成LOCUST_OPTS，避免后续维护时出现参数传递错误。
• 如果预处理生成的token数量在十万级以上，不建议全量通过ZMQ通道传输，可以把生成的token存在共享缓存里，每个worker启动时自行拉取，降低master节点的通信压力。
• 测试停止时不需要额外调整清理逻辑，和原有单机模式的行为保持一致即可。

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