mqtt_core.py -- mqtt核心类
# -*- coding: utf-8 -*-
# @Time : 2020/7/8 3:56 下午
# @Author : xiaobin
# @File : mqtt_core.py
# @Software: PyCharm
import paho.mqtt.client as mqtt
from locust import (TaskSet,task,events,Locust)
from gevent._semaphore import Semaphore
import json
import time
from util.data_params import *
all_locusts_spawned = Semaphore()
all_locusts_spawned.acquire()
def on_hatch_complete(**kwargs):
all_locusts_spawned.release()
events.hatch_complete += on_hatch_complete
class MqttClient(object):
"""
订阅topic
"""
def set_mqtt_params(self, data):
self.size = 1
self.start_time = time.time()
self.data = data
self._qos = 1
self._port = 1883
self._timeout = 1000
self._clean_session = False
self._topics = []
if self.data:
data_json = json.loads(self.data)
self._host = data_json.get('data').get('req').get('endPoint')
self._user = data_json.get('data').get('username')
self._password = data_json.get('data').get('password')
self._client_id = data_json.get('data').get('req').get('clientId')
topics = data_json.get('data').get('req').get('topics')
if type(topics) == dict:
for k, v in topics.items():
for i in v:
topic = f'{k}/{i}'
if topic:
self._topics.append(topic)
if self._client_id:
self.file_url = f'../files/{self._client_id}'
def _on_connect(self, client, userdata, flags, rc):
print(f'client_id: {self._client_id} Connected with result code:{rc}, time:{time.time()}')
total_time = int((time.time() - self.start_time) * 1000)
if self.size > 0:
events.request_success.fire(
request_type='mqtt',
name='连接数',
response_time=total_time,
response_length=0)
save_file_size(self.file_url, f'Connected with result code:{rc}, time:{time.time()}')
self.size = self.size - 1
def _on_message(self, client, userdata, msg):
# print(f'msg.topic: {msg.topic}, msg.payload: {msg.payload}')
data = msg.payload
response_length = (len(data))
data = json.loads(data)
start_time = int(data.get('start_time'))
end_time = int(time.time() * 1000)
total_time = end_time - start_time
msg_id = data.get('msg_id')
events.request_success.fire(
request_type='mqtt',
name='收到消息',
response_time=total_time,
response_length=response_length)
save_file_size(self.file_url, f'{self._client_id}, {msg_id}')
def client_loop(self):
client = mqtt.Client(client_id=self._client_id, clean_session=self._clean_session)
client.username_pw_set(username=self._user, password=self._password)
client.on_connect = self._on_connect
client.on_message = self._on_message
client.connect(host=self._host, port=self._port, keepalive=self._timeout)
for topic in self._topics:
client.subscribe(topic, qos=self._qos)
client.loop_forever()
locust_mqtt.py -- 继承locust父类
# -*- coding: utf-8 -*-
# @Time : 2020/7/8 4:02 下午
# @Author : xiaobin
# @File : locust_mqtt_test.py
# @Software: PyCharm
from locust import (TaskSet,task,events,Locust)
import core.mqtt_core
class locust_mqtt(Locust):
def __init__(self):
super(locust_mqtt, self).__init__()
self.client = core.mqtt_core.MqttClient()
test_mqtt.py --locust 执行文件
# -*- coding: utf-8 -*-
# @Time : 2020/7/8 4:04 下午
# @Author : xiaobin
# @File : test_mqtt.py
# @Software: PyCharm
import sys
import os
curPath = os.path.abspath(os.path.dirname(__file__))
rootPath = os.path.split(curPath)[0]
sys.path.append(rootPath)
from locust import *
import core.locust_mqtt
from util.parasm_size_lock import *
class test_mqtt(TaskSet):
@task(weight=1)
def test_mqtt_connect(self):
if self.locust.token_list:
data = self.locust.token_list[0]
self.client.set_mqtt_params(data)
self.locust.token_list.pop(0)
self.client.client_loop()
class myrun(core.locust_mqtt.locust_mqtt):
task_set = test_mqtt
host = 'http://127.0.0.1:8089'
file_name = '../files/mqtt_token_005'
token_list = []
with open(file_name, 'rb') as file_to_read:
print('start...')
while True:
lines = file_to_read.readline()
if lines:
data = lines
if isinstance(data, bytes):
data = str(data, encoding='utf-8')
data = data.replace('\n', '')
token_list.append(data)
else:
print('end')
break
lock = SizeLock(slave_size=16, data_list_size=len(token_list))
start_size = lock.get_start_size()
end_size = lock.get_end_size()
token_list = token_list[start_size:end_size]
print(f'启动获取mqtt动态唯一参数,start_size:{start_size}, end_size:{end_size} 可用token size:{len(token_list)}')
min_wait = 0
max_wait = 0
file_lock -- 分布式唯一参数控制，通过本地文件，每次服务m/s启动需要重置file中size为0或则del file
# -*- coding: utf-8 -*-
# @Time : 2020/7/9 1:49 下午
# @Author : xiaobin
# @File : parasm_size_lock.py
# @Software: PyCharm
import io
class SizeLock(object):
"""
用本地文件落地资源文件控制获取唯一参数
:param slave_size 启动节点数
:param data_list_size 初始化参数不唯一数
0 : 加载全部,从1开始
根据 启动节点数和初始化参数 和文件 size 计算得到该slave节点启动加载到的参数范围
"""
def __init__(self, slave_size, data_list_size):
self.slave_size = slave_size
self.data_list_size = data_list_size
self.file_url = '../files/size.lock'
self._set_scale_value()
def _set_scale_value(self):
self.scale_value = self.data_list_size//self.slave_size
self.scale_value = int(self.scale_value)
def save_file_size(self, data):
try:
fl = io.open(self.file_url, 'w')
fl.write(data)
fl.close()
except Exception as e:
print(e)
fl.close()
def get_file_size(self):
size = 0
try:
with open(self.file_url, 'rb') as file_to_read:
while True:
lines = file_to_read.readline()
if lines:
data = lines
if isinstance(data, bytes):
data = str(data, encoding='utf-8')
size = int(data)
else:
break
except Exception as e:
print(e)
return size
def get_start_size(self):
"""
开始兼容为0的情况
"""
file_size = self.get_file_size()
if file_size == 0:
return 0
elif file_size * self.scale_value < self.data_list_size:
return file_size * self.scale_value - self.scale_value
else:
print(f'获取index_size异常了')
return self.data_list_size
def get_end_size(self):
"""
兼容最后一次多赋值情况，比如 10 /3 0，1，2 4，5，6 7，8，9，10
:return:
"""
file_size = self.get_file_size()
end_size = None
if file_size == 0:
end_size = self.scale_value
else:
if (file_size+1) * self.scale_value <= self.data_list_size:
end_size = file_size * self.scale_value
else:
end_size = self.data_list_size
if end_size is not None:
self.save_file_size(f'{file_size+1}')
return end_size
if __name__ == '__main__':
list_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
s = SizeLock(3, len(list_data))
start_size = s.get_start_size()
end_size = s.get_end_size()
print(start_size, end_size)
print(list_data[start_size:end_size])
