如何在Locust中按端点统计被限流请求并扩展自定义统计

我正在基于Locust创建自定义TUI仪表盘，扩展终端可查看的数据。目前通过以下逻辑捕获数据：

@events.init.add_listener
def on_locust_init(environment, **_kwargs):
    if isinstance(environment.runner, (MasterRunner, LocalRunner)):
        stream_handler = prepare_handler()
        gevent.spawn(Dashboard(stream_handler).run, environment)

之后通过environment.runner.stats.entries获取所有RequestStats数据并生成自定义表格，但没办法把被限流的请求统计进去。

我在HttpUser的task里做了如下处理，把429响应标记为失败，并给统计条目加了自定义属性：

@task
def endpoint(self, method, url, *args, **kwargs):
    with self.client.get("/some/url/", catch_response=True, *args, **kwargs) as response:
        if response.status_code == 429:
            response.failure("Throttled!")
            if entry := self.environment.runner.stats.entries.get((url, method)):
                if not getattr(entry, "throttled"):
                    entry.throttled = 1
                else:
                    entry.throttled += 1

但在on_locust_init监听的下游逻辑里，environment.runner.stats.entries里的条目并没有这个新增的throttled属性。想知道怎么在请求处理时添加自定义统计，方便后续生成自定义表格？

方法1：继承RequestStats类扩展字段 #

Locust的RequestStats支持继承扩展，你可以自定义带throttled字段的子类，替换默认统计类：

1. 定义自定义RequestStats：

from locust.stats import RequestStats

class CustomRequestStats(RequestStats):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.throttled = 0  # 初始化限流统计字段

2. 在初始化时替换默认统计类：
   在on_locust_init中修改runner的统计条目类：

@events.init.add_listener
def on_locust_init(environment, **_kwargs):
    if isinstance(environment.runner, (MasterRunner, LocalRunner)):
        # 替换默认的RequestStats为自定义类
        environment.runner.stats.EntryClass = CustomRequestStats
        stream_handler = prepare_handler()
        gevent.spawn(Dashboard(stream_handler).run, environment)

3. 简化task里的统计逻辑：
   现在可以直接访问entry的throttled字段，无需判断字段是否存在：

@task
def endpoint(self, method, url, *args, **kwargs):
    with self.client.get("/some/url/", catch_response=True, *args, **kwargs) as response:
        if response.status_code == 429:
            response.failure("Throttled!")
            if entry := self.environment.runner.stats.entries.get((url, method)):
                entry.throttled += 1

方法2：用事件钩子统一统计限流请求 #

使用Locust的request事件钩子，统一处理所有请求的限流统计，避免在每个task里重复写逻辑：

from locust import events

@events.request.add_listener
def on_request(request_type, name, response_time, response_length, exception, context, **kwargs):
    # 判断是否为429限流响应
    if hasattr(context, "response") and context.response.status_code == 429:
        # 获取对应统计条目
        entry = context.environment.runner.stats.entries.get((name, request_type))
        if entry:
            # 确保throttled字段存在并累加
            if not getattr(entry, "throttled", None):
                setattr(entry, "throttled", 0)
            entry.throttled += 1

注意事项 #

• 分布式运行（Master-Worker模式）时，Worker节点的自定义统计需要同步到Master。可以通过events.report_to_master和events.worker_report事件传递统计数据，确保Master端能汇总所有Worker的限流数据。
• 操作统计条目时要注意协程安全，尽量保持逻辑简洁，避免阻塞。

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