locust压测框架实战案例2 locust压测websocket

一、接上回的内容：https://m.toutiao.com/is/RQbHxWg/

本次要实现的是新增自定义的图表到locust的report.html中；

二、准备工作

1. 先下载代码：https://github.com/locustio/locust 检出release分支，这次用的是2.4.3；

2. 之前没怎么看过生成报告相关的代码，这次看完之后基本了解得差不多，如果要自行阅读，推荐一个顺序：stats.py->web.py->static/locust.js->templates/stats_data.html->templates/report.html，locust的web服务用的是flask，报表生成用的echarts，中间还用到了jinja2；

3. 整理一下流程

a) env初始化了stats对象，create_runner的时候注册request回调将数据提交到stats对象，压测过程中的request.fire会触发回调将数据提交(比如response.success())，网页端每2秒会从webserver请求一次stats的数据，然后调用update_stats_charts函数生成Charts页面的图表；

b) Download Data下的Download Report按钮在上述流程中用到了jinja2来生成静态html，模板位于templates的report.html和stats_data.html，注意js使用的几个变量是在stats_data.html页面用jinja2语法申明和初始化的，自定义的变量也要追加到这个模板文件里；

4. 业务需求，1) 要在locust web网页的Charts下面新增目标业务耗时打点统计的堆叠图，数据源是每次业务请求响应内容里的一个字典，包含了5个k-v，分别记录了一次业务执行过程中5个阶段的耗时；2) locust Download Report按钮生成的report.html中新增耗时统计的堆叠图；

三、大致实现

a) 先说几个遇到的需要注意的点

i. 下载代码后还要安装依赖进行测试和编译pip install wheel setuptools tox，编译命令看Makefile；编译完成之后把build/lib/locust文件夹复制到项目工程的根目录进行调试；

ii. response触发fire回调是在self.__exit__()内部做的，success()只是标记一下状态，所以不要再with ... as response:的函数体内部使用self.environment.events.request.fire()，这样会导致重复统计；

iii. 需要多了解下echarts的内容，改图表才能更效率；jinja2的语法有些行如果有换行会报错；vscode对js文件的自动格式化，也会导致运行报错；

b) 因为涉及到部分公司业务，所以只列出需要改动的地方;

1. runners.py

class Runner:
    def __init__(self, environment):
        def on_request_success(request_type, name, response_time, response_length, **_kwargs):
            self.stats.log_request(request_type, name, response_time, response_length, **_kwargs)

2. 压测脚本的task

with self.client.post(self.url, json=data, headers=headers, catch_response=True) as response:
    if response.status_code == 200:
        response.request_meta.update({"metrics": response.json()["data"]["metrics"]})
        response.success()

3. stats.py

class RequestStats:
    def log_request(self, method, name, response_time, content_length, **kwargs):
        metrics = kwargs.get("metrics")
        self.total.log(response_time, content_length, metrics)
        self.get(name, method).log(response_time, content_length, metrics)

class StatsEntry:
    metrics = None
    metrics_cache = None

    def log(self, response_time, content_length, metrics):
        ...
        if metrics:
            self.cache_metrics(metrics)

    def extend(self, other):
        …
        self.cache_metrics(other.metrics)

    def serialize(self):
        return {
            …
            "metrics": self.metrics
        }

    def unserialize(cls, data):
        ...
        for key in […, "metrics"]:
            …
        ...

    def cache_metrics(self, ms):
        if not ms:
            return
        if not self.metrics:
            self.metrics = {}
        for x in ms:
            self.metrics[x] += ms[x]

    def get_metrics(self):
        if not self.metrics:
            return {}
        count = self.metrics.get("count")
        if count == 0:
            if self.metrics_cache:
                return self.metrics_cache
            return self.metrics
        tmp = {}
        for x in self.metrics:
            tmp[x] = int(self.metrics[x] / self.metrics["count"])
        self.metrics_cache = tmp
        self.metrics = {}
        return tmp

def stats_history(runner):
    while True:
        ...
        if runner.state != "stopped":
            metrics = stats.total.get_metrics()
            r = {
                ...
                "cost": metrics.get("cost") or 0
            }
        ...

4.web.py

class WebUI:
    def __init__(...):
        def request_stats():
			      ...
			      metrics = environment.runner.stats.total.get_metrics()
            report.update(metrics)
            return jsonify(report)

5.static/locust.js

function update_stats_charts() {
    if (stats_history["time"].length > 0) {
        renderCostChart.chart.setOption({
            xAxis: { data: stats_history["time"] },
            series: [
                { data: stats_history["cost"], markLine: createMarkLine(), stack: 'Total', emphasis: { focus: 'series' }, areaStyle: {} },
            ]
        });
    }
}
...
var CostChart = new LocustLineChart($(".charts-container"), "Cost Times (ms)", ["cost"], "ms");
...
charts.push(rpsChart, responseTimeChart, CostChart, usersChart);
...

function updateStats() {
    $.get('./stats/requests', function (report) {
        try {
            ...
            stats_history["cost"].push({ "value": report.cost, "users": report.user_count });
			      update_stats_charts();
		    ...

6.templates/stats-data.html

{% set cost_data = [] %}
{% for r in history}...{% do cost_data.append({"value": r.cost, "users": r.user_count}) %}...var stats_history = {..."cost": {{ cost_data | tojson }}...};

7.templates/report.html

<script>
        {% include 'stats_data.html' %}
        var CostChart = new LocustLineChart($(".charts-container"), "Cost Times (ms)", ["cost"], "ms");
        ...
        if(stats_history["time"].length > 0){
            CostChart.chart.setOption({
                xAxis: {data: stats_history["time"]},
                series: [
                    {data: stats_history["cost"], stack: 'Total', emphasis: { focus: 'series' }, areaStyle: {}},
                ]
            });
        ...

来看看最终效果：