基于Python多线程的性能测试脚本

背景：模拟发送大量消息到Kafka，测试服务对它的消费速度

import time
import threading
from kafka import KafkaProducer
from kafka.errors import kafka_errors
import traceback
import json
import random

# 测试环境是json字符串，需要根据实际情况配置
producer = KafkaProducer(
    bootstrap_servers=['192.168.1.2:9093'],
    key_serializer=lambda k: json.dumps(k).encode(),
    value_serializer=lambda v: json.dumps(v).encode())

"""1.Kafka消息定义"""
# 模拟设备运行状态
device_status = {
        "deviceStatusInfo": {
          "deviceName": "",
          "deviceId": "",
          "result": "1",
          "createAt": 1629702387,
          "pushTime": 1629702387,
        }
}

# 模拟设备上报数据
device_infos = {
        "deviceInfo": {
          "msg" : "启动中",
          "createAt": 1629702387,
          "pushTime": 1629702387,
        }
}


#######################################################################################################################

# 模拟思考时间
THINK_TIME = 0.1

# 单个user发送的数据条数
TOTAL_NUMBER = 2000

# 频率
quantity = 40

#######################################################################################################################


def fun(args, random_number):
    print("线程%s已启动" % args)
    num = random_number
    times = TOTAL_NUMBER / quantity
    for _ in range(int(times)):
        time.sleep(THINK_TIME)
        for _ in range(0, quantity):
            num += 1
            future = producer.send(
                'device_report',
                key=num,
                value=str(device_status),
                partition=0)
                
            future2 = producer.send(
                'device_event',
                key=num,
                value=str(device_infos),
                partition=0)  # 向第一个分区发送消息

            try:
                future.get(timeout=10)  # 监控是否发送成功
                future2.get(timeout=10)  # 监控是否发送成功
            except kafka_errors:  # 发送失败抛出kafka_errors
                traceback.format_exc()
    print("线程%s执行完成" % args)

if __name__ == '__main__':
    """模拟多用户"""
    USER = 5
    start = time.time()
    print(f"本次测试发送的数据量为{TOTAL_NUMBER*USER}")

    # kafka的消息key，需要定义一个值，从它开始累加，而且不能重复
    INIT_RANDOM_NUM = 10086

    thread_group = []
    for i in range(USER):
        t = threading.Thread(target=fun, args=(i,INIT_RANDOM_NUM))
        thread_group.append(t)

    for t in thread_group:
        t.start()

    for t in thread_group:
        t.join()

    print(f"总耗时{time.time() - start}秒")