如何基于第三方API返回结果控制Locust自动停止负载测试？

核心实现逻辑 #

• 用Locust内置的LoadTestShape自定义负载模型，实现并发用户数阶梯式逐步提升
• 绑定Locust初始化事件，测试启动时拉起独立的后台daemon线程，每隔5分钟调用内存监控API拉取服务端内存值
• 内存占用达到预设阈值时，调用environment.runner.quit()方法主动终止测试进程
• 所有全局状态标记采用线程安全的threading.Event实现，避免多线程读写冲突

完整可运行代码 #

import json
import time
import threading
import requests
from locust import SequentialTaskSet, task, HttpUser, between, LoadTestShape, events
from locust.env import Environment

# -------------------------- 配置项 可根据实际需求修改 --------------------------
# 内存占用阈值，单位：百分比
MEMORY_THRESHOLD = 80
# 内存监控间隔，单位：秒，5分钟即300
MONITOR_INTERVAL = 300
# 服务端内存查询API地址
MEM_MONITOR_API = "http://your-monitor-api.com/get_server_memory"
# 阶梯加压每步新增并发用户数
STEP_RAMP_UP = 10
# 每个并发阶梯维持时长，单位：秒，建议和监控间隔对齐
STEP_DURATION = 300
# 最大并发用户数上限，避免无限制加压
MAX_USERS = 100
# -----------------------------------------------------------------------------

# 线程安全的内存超标状态标记
mem_over_threshold = threading.Event()

def memory_monitor_task(environment: Environment):
    """后台内存监控任务"""
    while not mem_over_threshold.is_set():
        try:
            resp = requests.get(MEM_MONITOR_API, timeout=10)
            resp.raise_for_status()
            # 此处根据实际接口返回结构调整取值逻辑
            current_mem_usage = resp.json().get("memory_usage_percent")
            print(f"[监控日志] 当前服务端内存占用：{current_mem_usage}%")
            if current_mem_usage >= MEMORY_THRESHOLD:
                print(f"[监控告警] 内存占用达到阈值{MEMORY_THRESHOLD}%，自动停止测试")
                mem_over_threshold.set()
                if environment.runner:
                    environment.runner.quit()
        except Exception as e:
            print(f"[监控异常] 内存接口调用失败：{str(e)}")
        time.sleep(MONITOR_INTERVAL)

@events.init.add_listener
def on_locust_start(environment: Environment, **kwargs):
    """Locust初始化钩子，启动时自动拉起监控线程"""
    monitor_thread = threading.Thread(
        target=memory_monitor_task,
        args=(environment,),
        daemon=True
    )
    monitor_thread.start()

class StepLoadShape(LoadTestShape):
    """自定义阶梯加压负载模型"""
    def tick(self):
        # 内存超标时停止加压
        if mem_over_threshold.is_set():
            return None
        run_time = self.get_run_time()
        current_step = int(run_time // STEP_DURATION)
        target_users = min(current_step * STEP_RAMP_UP, MAX_USERS)
        # 返回值格式：(目标并发数, 用户 spawn 速率)
        return (target_users, STEP_RAMP_UP)

class UserBehavior(SequentialTaskSet):
    @task
    def api_call(self):
        with self.client.get('/test_api') as response:
            print(response)

class WebsiteUser(HttpUser):
    tasks = [UserBehavior]
    wait_time = between(2, 10)

使用说明 #

• 替换代码中配置项的MEM_MONITOR_API为实际的内存查询接口地址，根据接口返回结构调整current_mem_usage的取值逻辑
• 按需修改内存阈值、阶梯步长、最大并发等配置参数即可直接使用
• 启动命令和常规Locust脚本一致：locust -f 脚本文件名.py，自动开启阶梯加压和内存监控停止逻辑
• 若不需要阶梯加压，仅保留内存自动停止能力，删除StepLoadShape类即可，并发规则按默认的Locust参数配置

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