基於Locust實現MQTT協議服務的壓測腳本

最近在忙業務的間隙，穿插着做了些性能測試。

一、背景簡介

業務背景大概介紹一下，就是按照國標規定，車輛需要上傳一些指定的數據到ZF的指定平臺，同時車輛也會把數據傳到企業雲端服務上，於是乎就產生了一些性能需求。

目前我們只是先簡單的進行了一個性能場景的測試，就是評估目前服務是否能夠支持，預期的最大同時在線車輛上傳數據。經過評估，在線車輛數據按照預期的10倍來進行的，並且後面增加持續運行12h查看服務鏈路的穩定性。

本篇並不是一個嚴謹的性能測試過程結果分享，主要是分享下關於mqtt協議服務的壓測腳本的編寫。因爲之前我也沒接觸過MQTT協議的壓測，網上關於相關的壓測腳本的內容也比較雜亂，所以記錄一下，僅供參考。

捋一下鏈路就知道需要生成哪些數據（因爲服務還未上線使用，所以產生的壓測數據後面可以直接清理掉即可。）：

一些前置數據：比如數據庫、緩存裏涉及到的車輛數據，通信祕鑰數據等等，這些可以之前寫腳本一次性生成即可。
車輛上報的數據：車輛上報到雲端的數據，是經過一系列加密轉碼，期間還要設計到解密等，這個經過評估，可以簡化其中的某些環境，所以所有的車可以直接發送相同的數據即可。
車輛數據：最後就是生成對應的車輛數據，同時在線，按照評估的頻率發送數據。

其中第1、2的數據在之前針對性的分別生成即可，第3步的車輛發送數據就是壓測腳本要乾的事情了。

二、技術選型

這個倒是很快，搜索引擎大概搜了一下，內容很少，或者說對我有用的內容很少。有看到jmeter有相關插件的，但是這個方案基本上我都是否決的，一來我不擅長用，而來我覺得用起來肯定會比自己編碼要麻煩的多。

所以就繼續編碼好了，仍然首選python，想到了locust庫，後來看官方文檔的時候，看到locust也針對mqtt協議拓展了一些內容。但是我嘗試下來不太符合我這的需求，也可能當時我用的不對吧，所以就只能自己來從零開始編寫了。

搜索中又發現Python中用於mqtt協議的庫叫paho.mqtt，支持連接代理，消息的訂閱、收發等等，於是最後確定使用：locust+paho.mqtt的組合來實現本次的負載腳本。

三、代碼編寫

1. 腳本代碼

暫時沒做代碼分層，目前場景簡單，就直接都放一個模塊裏了，有點長，先貼上來，後面部分會對腳本的重點內容進行拆解。

腳本目前做了這些事情：

從db中查詢有效可用的所有測試車輛信息數據
根據命令行的輸入參數，指定啓動的車輛數，以及與broker代理建立連接的頻率
建立連接成功的車輛，就可以根據腳本里指定的頻次，來像broker發送數據
腳本統計連接數、請求數、響應時間等信息寫到報表中
調試遇到車輛會批量斷開連接的情況，增加了當車輛斷開連接時，把斷開時間、車輛信息寫到本地csv中，方便第二天來查看分析。

import csv
import datetime
import queue
import os
import sys
import time
import ssl

from paho.mqtt import client as mqtt_client

# 根據不同系統進行路徑適配
if os.name == "nt":
    path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
    sys.path.insert(0, path)
    from GB_test.utils.mysql_operating import DB
elif os.name == "posix":
    sys.path.append("/app/qa_test_app/")
    from GB_test.utils.mysql_operating import DB

from locust import User, TaskSet, events, task, between, run_single_user


BROKER_ADDRESS = "broker服務地址"
PORT = 1111
PASSWORD = "111111"
PUBLISH_TIMEOUT = 10000  # 超時時間
TEST_TOPIC = "test_topic"

TEST_VALUE = [16, 3, -26, 4, 0, 36,.......]  # 用來publish的測試數據，僅示意

BYTES_DATA = bytes(i % 256 for i in TEST_VALUE)  # 業務需要轉換成 byte 類型後再發送

# 創建隊列
client_queue = queue.Queue()

# 連接DB，讀取車輛數據
db = DB("db_vmd")
select_sql = "select xxxx"  
client_list = db.fetch_all(select_sql)
print("車輛數據查詢完畢，數據量：{}".format(len(client_list)))
for t in client_list:
    # 把可用的車輛信息存到隊列中去
    client_queue.put(t)


def fire_success(**kwargs):
    """請求成功時調用"""
    events.request.fire(**kwargs)


def calculate_resp_time(t1, t2):
    """計算響應時間"""
    return int((t2 - t1) * 1000)


class MQTTMessage:
    """已發送的消息實體類"""
    def __init__(self, _type, qos, topic, payload, start_time, timeout):
        self.type = _type,
        self.qos = qos,
        self.topic = topic
        self.payload = payload
        self.start_time = start_time
        self.timeout = timeout


# 統計總共發送成功的消息數量
total_published = 0
disconnect_record_list = []  # 定義存放連接斷開的記錄的列表容器


class PublishTask(TaskSet):

    @task
    def task_publish(self):
        self.client.loop_start()
        topic = TEST_TOPIC
        payload = BYTES_DATA
        # 記錄發送的開始時間
        start_time = time.time()
        mqtt_msg_info = self.client.publish(topic, payload, qos=1, retain=False)
        published_mid = mqtt_msg_info.mid
        # 將發送成功的消息內容，放入client實例的 published_message 字段
        self.client.published_message[published_mid] = MQTTMessage(REQUEST_TYPE,
                                                                   0,
                                                                   topic,
                                                                   payload,
                                                                   start_time,
                                                                   PUBLISH_TIMEOUT)
        # 發送成功回調
        self.client.on_publish = self.on_publish
        # 斷開連接回調
        self.client.on_disconnect = self.on_disconnect

    @staticmethod
    def on_disconnect(client, userdata, rc):
        """ broker連接斷開，放入列表容器"""
        disconnected_info = [str(client._client_id), rc, datetime.datetime.now()]
        disconnect_record_list.append(disconnected_info)
        print("rc狀態:{} - -".format(rc), "{}-broker連接已斷開".format(str(client._client_id)))

    @staticmethod
    def on_publish(client, userdata, mid):
        if mid:
            # 記錄消息發送成功的時間
            end_time = time.time()
            # 從已發送的消息容器中，取出消息
            message = client.published_message.pop(mid, None)
            # 計算開始發送到發送成功的耗時
            publish_resp_time = calculate_resp_time(message.start_time, end_time)
            fire_success(
                request_type="p_success",
                name="client_id: " + str(client._client_id),
                response_time=publish_resp_time,
                response_length=len(message.payload),
                exception=None,
                context=None
            )
            global total_published
            # 成功發送累加1
            total_published += 1


class MQTTLocustUser(User):
    tasks = [PublishTask]
    wait_time = between(2, 2)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 從隊列中獲取客戶端 username 和 client_id
        current_client = client_queue.get()

        self.client = mqtt_client.Client(current_client[1])
        self.client.username_pw_set(current_client[0], PASSWORD)
        # self.client.username_pw_set(current_client[0] + "1", PASSWORD)  # 模擬client連接報錯

        # 定義一個容器，存放已發送的消息
        self.client.published_message = {}

    def on_start(self):
        # 設置tls
        context = ssl.SSLContext(ssl.PROTOCOL_TLS)
        self.client.tls_set_context(context)

        self.client.connect(host=BROKER_ADDRESS, port=PORT, keepalive=60)
        self.client.on_connect = self.on_connect

    def on_stop(self):
        print("publish 成功, 當前已成功發送數量：{}".format(total_published))
        if len(disconnect_record_list) == 0:
            print("無斷開連接的client")
        else:
            # 把斷開記錄裏的信息寫入csv
            with open("disconnect_record.csv", "w", newline='', encoding='UTF8') as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow(['client_id', 'rc_status', 'disconnected_time'])
                for i in disconnect_record_list:
                    writer.writerow(i)
            print("斷開連接的client信息已寫入csv文件")

    @staticmethod
    def on_connect(client, userdata, flags, rc, props=None):
        if rc == 0:
            print("rc狀態:{} - -".format(rc), "{}-連接broker成功".format(str(client._client_id)))
            fire_success(
                request_type="c_success",
                name='count_connected',
                response_time=0,
                response_length=0,
                exception=None,
                context=None
            )
        else:
            print("rc狀態:{} - -".format(rc), "{}-連接broker失敗".format(str(client._client_id)))
            fire_success(
                request_type="c_fail",
                name="client_id: " + str(client._client_id),
                response_time=0,
                response_length=0,
                exception=None,
                context=None
            )


if __name__ == '__main__':
    run_single_user(MQTTLocustUser)