如何让Locust在多进程下执行指定请求数后干净终止？

需求：在多用户、多进程的Locust测试场景中，执行完指定数量的请求后终止测试，且不希望Locust重启已完成任务的用户。

当前实现思路 #

• 通过init事件初始化的Distributor，让每个User实例知晓自身需执行的请求总数并计数：

class MyUser:
    def __init__(self):
        self.id = next(distributors["user_id"])
        self.limit = ... # 根据用户ID计算请求上限
    @task
    def request(self):
        # 执行请求逻辑
        self.requests += 1

• 每次请求后检查完成数，达到上限则停止当前用户；若所有用户都已停止则终止Locust，逻辑参考iteration_limit_wrapper：

# request()方法的收尾逻辑
if self.requests == self.limit:
    runner = self.environment.runner
    if runner.user_count == 1:
        logging.info("最后一个用户已停止，正在终止测试")
        if isinstance(runner, WorkerRunner):
            runner._send_stats()  # 发送最终统计报告
        # 在单独的greenlet中触发终止，避免异步处理阻塞
        gevent.spawn_later(0.1, runner.quit)
    raise StopUser()

遇到的问题 #

上述实现虽能终止测试，但过程不规范：Locust会因为发现部分指定数量的用户已退出，尝试重启他们，导致Distributor被调用次数超出预期（主进程抛出StopIteration）。

以--users=5和--processes=2运行时，包含进程ID的日志显示：

[2024-05-03 12:41:10,067] localhost/70327/INFO/root: 最后一个用户已停止，正在终止测试运行器
[2024-05-03 12:41:10,237] localhost/70327/INFO/locust.main: 正在关闭（退出码0）
[2024-05-03 12:41:10,237] localhost/70322/INFO/locust.runners: 向1个就绪worker发送生成5个用户的任务，生成速率为5.00
...
[2024-05-03 12:41:10,267] localhost/70328/DEBUG/locust.runners: 正在额外生成{"MyUser": 2}（当前已运行{"MyUser": 3}）...
...
[2024-05-03 12:41:10,519] localhost/70328/DEBUG/locust.runners: 正在停止所有用户（调用自.venv/lib/python3.11/site-packages/locust/runners.py:411）

提问 #

是否有机制能让Locust仅生成指定数量的用户，在用户完成任务后不再重启？或是有比raise StopUser()更明确的方式告知“该用户已完成任务，请勿重启”？

方法1：提前通知主进程停止补充用户 #

直接在用户完成任务时，通知主进程停止生成新用户，再执行终止逻辑。主进程调用runner.stop()后，就不会再向worker发送生成新用户的指令：

if self.requests == self.limit:
    runner = self.environment.runner
    # 先停止主进程的用户生成逻辑
    if isinstance(runner, MasterRunner):
        runner.stop()
    elif isinstance(runner, WorkerRunner):
        # 给主进程发送停止信号
        runner.send_message("stop")
    
    if runner.user_count == 1:
        logging.info("最后一个用户已停止，正在终止测试")
        if isinstance(runner, WorkerRunner):
            runner._send_stats()
        gevent.spawn_later(0.1, runner.quit)
    raise StopUser()

方法2：自定义用户生成逻辑，禁用自动重启 #

通过全局标志控制用户生命周期，当测试即将完成时，阻止新用户的初始化：

# 在init事件中初始化全局标志
def on_init(environment, **kwargs):
    environment.test_completed = False

events.init.add_listener(on_init)

class MyUser:
    def __init__(self):
        self.id = next(distributors["user_id"])
        self.limit = ...
        # 如果测试已完成，直接停止用户
        if self.environment.test_completed:
            raise StopUser()
    
    @task
    def request(self):
        # 执行请求逻辑
        self.requests += 1
        if self.requests == self.limit:
            runner = self.environment.runner
            if runner.user_count == 1:
                self.environment.test_completed = True
                logging.info("最后一个用户已停止，正在终止测试")
                if isinstance(runner, WorkerRunner):
                    runner._send_stats()
                gevent.spawn_later(0.1, runner.quit)
            raise StopUser()

方法3：结合内置参数快速收尾 #

使用--stop-timeout参数，确保测试终止时不会有新用户生成，同时在代码中完成任务后立即触发终止，减少不必要的等待：

locust -f your_test_file.py --users=5 --processes=2 --stop-timeout=1

配合代码中的runner.quit()调用，可快速结束测试，避免用户重启问题。

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