├── LICENSE ├── README.md ├── aws-iot-publish.py ├── locustfile.py └── mqtt_locust.py /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Concurrency Labs 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | For more details, read this article: 3 | https://www.concurrencylabs.com/blog/hatch-a-swarm-of-things-using-locust-and-ec2/ 4 | 5 | 6 | ## Load test for AWS IoT applications using Locust.io and Paho MQTT client 7 | 8 | This repo is based on code found originally here 9 | 10 | 11 | ### Architecture 12 | 13 | ![IoT Locust Architecture](https://www.concurrencylabs.com/img/posts/4-locust-mqtt/cloudcraftLocustAWSIoT.png) 14 | 15 | This setup creates virtual IoT devices using Locust, which publishes messages to the AWS IoT message broker. 16 | Once messages arrive to AWS IoT, you can configure rules that integrate with other AWS services, such as Dynamo DB, Kinesis, Lambda, S3, SNS, SQS or CloudWatch. 17 | 18 | 19 | 20 | ### Dependencies 21 | 22 | * Highly recommended to use virtualenv when installing and running this test. 23 | * Python 2.7.9 or later. The code in this repo requires support for TLSv1.2, which is available starting in Python 2.7.9 24 | * The code in this repo does NOT work with Python 3 due to Locust not supporting Python 3. 25 | * GCC (GNU Compiler Collection), required by Locust: ```sudo yum install gcc``` 26 | * GCC C++, required by Locust: ```sudo yum install gcc-c++``` 27 | * Locust: ```pip install locustio``` 28 | * OpenSSL v1.0.1 or greater (also due to TLSv1.2) 29 | * Paho MQTT client: ```pip install paho-mqtt``` 30 | 31 | 32 | ### mqtt_locust.py 33 | 34 | This file contains logic to create an MQTT client, initiate a connection, publish messages 35 | to a topic and subscribe to a topic. It also reports failure or success results to Locust. 36 | 37 | At the time, it only supports QoS=0, which is what is needed for connecting to the AWS IoT 38 | service. 39 | 40 | Another important feature is the support for TLSv1.2, which is required by AWS IoT in order 41 | to establish a secure connection. You must create certificate files for AWS IoT, attach an 42 | IAM policy to them in AWS and your locust tests must indicate the location of the certificate 43 | files in your system. 44 | 45 | ### aws-iot-publish.py 46 | 47 | This file creates a basic locust test that publishes messages to an AWS IoT topic. The payload 48 | is a random value between 0 and 10, which represents a temperature reading. 49 | 50 | ### Certificates 51 | Follow the instructions in the AWS IoT documentation in order 52 | to create certificate files. Follow the instructions here to create a CA certificate. 53 | Set environment variables CA_CERT, IOT_CERT and IOT_PRIVATE_KEY with the full location of your .pem files 54 | 55 | ### Topics 56 | Set the environment variable ```MQTT_TOPIC``` with the topic you want to publish messages to. 57 | 58 | 59 | ### Running the test 60 | 61 | First you need to get the endpoint AWS IoT has assigned to your account, using the AWS CLI: 62 | ```aws iot describe-endpoint``` 63 | 64 | then run the locust command from your environment: 65 | 66 | ```locust --host= -f aws-iot-publish.py``` 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | -------------------------------------------------------------------------------- /aws-iot-publish.py: -------------------------------------------------------------------------------- 1 | import json 2 | import random 3 | import resource 4 | import ssl 5 | import time 6 | import os 7 | 8 | from locust import TaskSet, task 9 | 10 | from mqtt_locust import MQTTLocust 11 | 12 | #this value is the number of seconds to be used before retrying operations (valid for QoS >1) 13 | RETRY = 5 14 | 15 | #ms 16 | PUBLISH_TIMEOUT = 10000 17 | SUBSCRIBE_TIMEOUT = 10000 18 | 19 | 20 | class ThingBehavior(TaskSet): 21 | @task 22 | def pubqos0(self): 23 | topic = os.getenv('MQTT_TOPIC','') 24 | if topic == '': 25 | raise ValueError("Please set environment variable MQTT_TOPIC") 26 | self.client.publish(topic, payload=self.payload(), qos=0, name='publish:qos0:'+topic, timeout=PUBLISH_TIMEOUT) 27 | 28 | def on_start(self): 29 | #allow for the connection to be established before doing anything (publishing or subscribing) to the MQTT topic 30 | time.sleep(5) 31 | 32 | 33 | def payload(self): 34 | payload = { 35 | 'temperature': random.randrange(0,10,1) #set temperature between 0 and 10 36 | } 37 | return json.dumps(payload) 38 | 39 | 40 | """ 41 | Locust hatches several instances of this class, according to the number of simulated users 42 | that we define in the GUI. Each instance of MyThing represents a device that will connect to AWS IoT. 43 | """ 44 | class MyThing(MQTTLocust): 45 | ca_cert = os.getenv('CA_CERT','') 46 | iot_cert = os.getenv('IOT_CERT','') 47 | iot_private_key = os.getenv('IOT_PRIVATE_KEY','') 48 | if ca_cert == '' or iot_cert == '' or iot_private_key == '': 49 | raise ValueError("Make sure the following environment variables are set: CA_CERT, IOT_CERT, IOT_PRIVATE_KEY") 50 | task_set = ThingBehavior 51 | min_wait = 1000 52 | max_wait = 1500 53 | -------------------------------------------------------------------------------- /locustfile.py: -------------------------------------------------------------------------------- 1 | from locust import HttpLocust, TaskSet 2 | 3 | def index(l): 4 | l.client.get("/") 5 | 6 | class UserBehavior(TaskSet): 7 | tasks = {index:1} 8 | 9 | def on_start(self): 10 | index(self) 11 | 12 | class WebsiteUser(HttpLocust): 13 | task_set = UserBehavior 14 | min_wait=1000 15 | max_wait=2000 16 | 17 | -------------------------------------------------------------------------------- /mqtt_locust.py: -------------------------------------------------------------------------------- 1 | import random 2 | import time 3 | import sys 4 | import ssl 5 | 6 | import paho.mqtt.client as mqtt 7 | from locust import Locust, task, TaskSet, events 8 | 9 | REQUEST_TYPE = 'MQTT' 10 | MESSAGE_TYPE_PUB = 'PUB' 11 | MESSAGE_TYPE_SUB = 'SUB' 12 | 13 | def time_delta(t1, t2): 14 | return int((t2 - t1) * 1000) 15 | 16 | 17 | def fire_locust_failure(**kwargs): 18 | events.request_failure.fire(**kwargs) 19 | 20 | 21 | def fire_locust_success(**kwargs): 22 | events.request_success.fire(**kwargs) 23 | 24 | 25 | class LocustError(Exception): 26 | pass 27 | 28 | 29 | class TimeoutError(ValueError): 30 | pass 31 | 32 | 33 | class ConnectError(Exception): 34 | pass 35 | 36 | class DisconnectError(Exception): 37 | pass 38 | 39 | 40 | class Message(object): 41 | 42 | def __init__(self, type, qos, topic, payload, start_time, timeout, name): 43 | self.type = type, 44 | self.qos = qos, 45 | self.topic = topic 46 | self.payload = payload 47 | self.start_time = start_time 48 | self.timeout = timeout 49 | self.name = name 50 | 51 | def timed_out(self, total_time): 52 | return self.timeout is not None and total_time > self.timeout 53 | 54 | 55 | class MQTTClient(mqtt.Client): 56 | 57 | def __init__(self, *args, **kwargs): 58 | super(MQTTClient, self).__init__(*args, **kwargs) 59 | self.on_publish = self.locust_on_publish 60 | self.on_subscribe = self.locust_on_subscribe 61 | self.on_disconnect = self.locust_on_disconnect 62 | self.on_connect = self.locust_on_connect 63 | self.pubmmap = {} 64 | self.submmap = {} 65 | self.defaultQoS = 0 66 | 67 | def tls_set(self, ca_certs, certfile=None, keyfile=None, cert_reqs=ssl.CERT_REQUIRED, 68 | tls_version=ssl.PROTOCOL_TLSv1_2, ciphers=None): 69 | start_time = time.time() 70 | try: 71 | super(MQTTClient,self).tls_set(ca_certs, 72 | certfile, 73 | keyfile, 74 | cert_reqs, 75 | tls_version, 76 | ciphers) 77 | except Exception as e: 78 | fire_locust_failure( 79 | request_type=REQUEST_TYPE, 80 | name='tls_set', 81 | response_time=time_delta(start_time, time.time()), 82 | exception=e) 83 | 84 | #retry is not used at the time since this implementation only supports QoS 0 85 | def publish(self, topic, payload=None, qos=0, retry=5, name='publish', **kwargs): 86 | timeout = kwargs.pop('timeout', 10000) 87 | start_time = time.time() 88 | try: 89 | err, mid = super(MQTTClient, self).publish( 90 | topic, 91 | payload=payload, 92 | qos=qos, 93 | **kwargs 94 | ) 95 | if err: 96 | fire_locust_failure( 97 | request_type=REQUEST_TYPE, 98 | name=name, 99 | response_time=time_delta(start_time, time.time()), 100 | exception=ValueError(err) 101 | ) 102 | 103 | print ("publish: err,mid:["+str(err)+","+str(mid)+"]") 104 | self.pubmmap[mid] = Message( 105 | MESSAGE_TYPE_PUB, qos, topic, payload, start_time, timeout, name 106 | ) 107 | #print ("publish: Saved message - mqtt client obj id:["+str(id(self))+"] - pubmmap obj id:["+str(id(self.pubmmap))+"] in dict - mid:["+str(mid)+"] - message object id:["+str(id(self.pubmmap[mid]))+"]") 108 | except Exception as e: 109 | fire_locust_failure( 110 | request_type=REQUEST_TYPE, 111 | name=name, 112 | response_time=time_delta(start_time, time.time()), 113 | exception=e, 114 | ) 115 | print (str(e)) 116 | 117 | #retry is not used at the time since this implementation only supports QoS 0 118 | def subscribe(self, topic, qos=0, retry=5, name='subscribe', timeout=15000): 119 | #print ("subscribing to topic:["+topic+"]") 120 | start_time = time.time() 121 | try: 122 | err, mid = super(MQTTClient, self).subscribe( 123 | topic, 124 | qos=qos 125 | ) 126 | self.submmap[mid] = Message( 127 | MESSAGE_TYPE_SUB, qos, topic, "", start_time, timeout, name 128 | ) 129 | if err: 130 | raise ValueError(err) 131 | print ("Subscribed to topic with err:["+str(err)+"]messageId:["+str(mid)+"]") 132 | except Exception as e: 133 | total_time = time_delta(start_time, time.time()) 134 | fire_locust_failure( 135 | request_type=REQUEST_TYPE, 136 | name=name, 137 | response_time=total_time, 138 | exception=e, 139 | ) 140 | print ("Exception when subscribing to topic:["+str(e)+"]") 141 | 142 | 143 | def locust_on_connect(self, client, flags_dict, userdata, rc): 144 | #print("Connection returned result: "+mqtt.connack_string(rc)) 145 | fire_locust_success( 146 | request_type=REQUEST_TYPE, 147 | name='connect', 148 | response_time=0, 149 | response_length=0 150 | ) 151 | 152 | 153 | """ 154 | Paho documentation regarding on_publish event: 155 | 'For messages with QoS levels 1 and 2, this means that the appropriate handshakes have 156 | completed. For QoS 0, this simply means that the message has left the client.' 157 | 158 | This means that the value we record in fire_locust_success for QoS 0 will always 159 | be very low and not a meaningful metric. The interesting part comes when we analyze 160 | metrics emitted by the system on the other side of the MQTT broker (the systems processing 161 | incoming data from things). 162 | """ 163 | 164 | def locust_on_publish(self, client, userdata, mid): 165 | end_time = time.time() 166 | 167 | if self.defaultQoS == 0: 168 | #if QoS=0, we reach the callback before the publish() has enough time to update the pubmmap dictionary 169 | time.sleep(float(0.5)) 170 | 171 | message = self.pubmmap.pop(mid, None) 172 | #print ("on_publish - mqtt client obj id:["+str(id(self))+"] - pubmmap obj id:["+str(id(self.pubmmap))+"] - mid:["+str(mid)+"] - message obj id:["+str(id(message))+"]") 173 | if message is None: 174 | fire_locust_failure( 175 | request_type=REQUEST_TYPE, 176 | name="message_found", 177 | response_time=0, 178 | exception=ValueError("Published message could not be found"), 179 | ) 180 | return 181 | 182 | total_time = time_delta(message.start_time, end_time) 183 | if message.timed_out(total_time): 184 | fire_locust_failure( 185 | request_type=REQUEST_TYPE, 186 | name=message.name, 187 | response_time=total_time, 188 | exception=TimeoutError("publish timed out"), 189 | ) 190 | #print("report publish failure - response_time:["+str(total_time)+"]") 191 | else: 192 | fire_locust_success( 193 | request_type=REQUEST_TYPE, 194 | name=message.name, 195 | response_time=total_time, 196 | response_length=len(message.payload), 197 | ) 198 | #print("report publish success - response_time:["+str(total_time)+"]") 199 | 200 | 201 | def locust_on_subscribe(self, client, userdata, mid, granted_qos): 202 | end_time = time.time() 203 | message = self.submmap.pop(mid, None) 204 | if message is None: 205 | print ("did not find message for on_subscribe") 206 | return 207 | total_time = time_delta(message.start_time, end_time) 208 | if message.timed_out(total_time): 209 | fire_locust_failure( 210 | request_type=REQUEST_TYPE, 211 | name=message.name, 212 | response_time=total_time, 213 | exception=TimeoutError("subscribe timed out"), 214 | ) 215 | print("report subscribe failure - response_time:["+str(total_time)+"]") 216 | else: 217 | fire_locust_success( 218 | request_type=REQUEST_TYPE, 219 | name=message.name, 220 | response_time=total_time, 221 | response_length=0, 222 | ) 223 | print("report subscribe success - response_time:["+str(total_time)+"]") 224 | 225 | 226 | def locust_on_disconnect(self, client, userdata, rc): 227 | fire_locust_failure( 228 | request_type=REQUEST_TYPE, 229 | name='disconnect', 230 | response_time=0, 231 | exception=DisconnectError("disconnected"), 232 | ) 233 | self.reconnect() 234 | 235 | 236 | class MQTTLocust(Locust): 237 | 238 | def __init__(self, *args, **kwargs): 239 | super(Locust, self).__init__(*args, **kwargs) 240 | if self.host is None: 241 | raise LocustError("You must specify a host") 242 | 243 | #TODO: Current implementation sets an empty client_id when the connection is initialized, 244 | # which Paho handles by creating a random client_id. 245 | # Ideally we want to control the client_id that is set in Paho. Each client_id 246 | # should match a thing_id in the AWS IoT Thing Registry 247 | #self.client = MQTTClient(self.client_id) 248 | self.client = MQTTClient() 249 | try: 250 | [host, port] = self.host.split(":") 251 | except: 252 | host, port = self.host, 8883 253 | 254 | try: 255 | self.client.tls_set(self.ca_cert, self.iot_cert, self.iot_private_key, tls_version=ssl.PROTOCOL_TLSv1_2) 256 | #It is important to do an asynchronous connect, given that we will have 257 | #multiple connections happening in a single server during a Locust test 258 | self.client.connect_async(host, port) 259 | self.client.loop_start() 260 | except Exception as e: 261 | fire_locust_failure( 262 | request_type=REQUEST_TYPE, 263 | name='connect', 264 | response_time=time_delta(start_time, time.time()), 265 | exception=ConnectError("Could not connect to host:["+host+"]") 266 | ) 267 | 268 | --------------------------------------------------------------------------------