#------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. #-------------------------------------------------------------------------- import time import threading from datetime import datetime, timedelta import concurrent import sys import uuid import logging try: import psutil except ImportError: pass # If psutil isn't installed, simply does not capture process stats. from azure.servicebus import ServiceBusClient, ServiceBusMessage, ServiceBusMessageBatch from azure.servicebus.exceptions import MessageAlreadySettled from utilities import _build_logger _logger = _build_logger("stress-test", logging.WARN) class ReceiveType: push="push" pull="pull" none=None class StressTestResults(object): def __init__(self): self.total_sent=0 self.total_received=0 self.time_elapsed=None self.state_by_sender={} self.state_by_receiver={} def __repr__(self): return str(vars(self)) class StressTestRunnerState(object): '''Per-runner state, e.g. if you spawn 3 senders each will have this as their state object, which will be coalesced at completion into StressTestResults''' def __init__(self): self.total_sent=0 self.total_received=0 self.cpu_percent = None self.memory_bytes = None self.timestamp = None self.exceptions = [] def __repr__(self): return str(vars(self)) def populate_process_stats(self): self.timestamp = datetime.utcnow() try: self.cpu_percent = psutil.cpu_percent() self.memory_bytes = psutil.virtual_memory().total except NameError: return # psutil was not installed, fall back to simply not capturing these stats. class StressTestRunner: '''Framework for running a service bus stress test. Duration can be overriden via the --stress_test_duration flag from the command line''' def __init__(self, senders, receivers, duration = timedelta(minutes=15), receive_type = ReceiveType.push, send_batch_size = None, message_size = 10, max_wait_time = 10, send_delay = .01, receive_delay = 0, should_complete_messages = True, max_message_count = 1, send_session_id = None, fail_on_exception = True): self.senders = senders self.receivers = receivers self.duration=duration self.receive_type = receive_type self.message_size = message_size self.send_batch_size = send_batch_size self.max_wait_time = max_wait_time self.send_delay = send_delay self.receive_delay = receive_delay self.should_complete_messages = should_complete_messages self.max_message_count = max_message_count self.fail_on_exception = fail_on_exception self.send_session_id = send_session_id # Because of pickle we need to create a state object and not just pass around ourselves. # If we ever require multiple runs of this one after another, just make Run() reset this. self._state = StressTestRunnerState() self._duration_override = None for arg in sys.argv: if arg.startswith('--stress_test_duration_seconds='): self._duration_override = timedelta(seconds=int(arg.split('=')[1])) self._should_stop = False # Plugin functions the caller can override to further tailor the test. def on_send(self, state, sent_message, sender): '''Called on every successful send, per message''' pass def on_receive(self, state, received_message, receiver): '''Called on every successful receive, per message''' pass def on_receive_batch(self, state, batch, receiver): '''Called on every successful receive, at the batch or iterator level rather than per-message''' pass def post_receive(self, state, receiver): '''Called after completion of every successful receive''' pass def on_complete(self, send_results=[], receive_results=[]): '''Called on stress test run completion''' pass def pre_process_message(self, message): '''Allows user to transform the message before batching or sending it.''' pass def pre_process_message_batch(self, message): '''Allows user to transform the batch before sending it.''' pass def pre_process_message_body(self, payload): '''Allows user to transform message payload before sending it.''' return payload def _schedule_interval_logger(self, end_time, description="", interval_seconds=30): def _do_interval_logging(): if end_time > datetime.utcnow() and not self._should_stop: self._state.populate_process_stats() _logger.critical("{} RECURRENT STATUS: {}".format(description, self._state)) self._schedule_interval_logger(end_time, description, interval_seconds) t = threading.Timer(interval_seconds, _do_interval_logging) t.start() def _construct_message(self): if self.send_batch_size != None: batch = ServiceBusMessageBatch() for _ in range(self.send_batch_size): message = ServiceBusMessage(self.pre_process_message_body("a" * self.message_size)) self.pre_process_message(message) batch.add_message(message) self.PreProcessMessageBatch(batch) return batch else: message = ServiceBusMessage(self.pre_process_message_body("a" * self.message_size)) self.pre_process_message(message) return message def _send(self, sender, end_time): self._schedule_interval_logger(end_time, "Sender " + str(self)) try: _logger.info("STARTING SENDER") with sender: while end_time > datetime.utcnow() and not self._should_stop: _logger.info("SENDING") try: message = self._construct_message() if self.send_session_id != None: message.session_id = self.send_session_id sender.send_messages(message) self.on_send(self._state, message, sender) except Exception as e: _logger.exception("Exception during send: {}".format(e)) self._state.exceptions.append(e) if self.fail_on_exception: raise self._state.total_sent += 1 time.sleep(self.send_delay) self._state.timestamp = datetime.utcnow() return self._state except Exception as e: _logger.exception("Exception in sender: {}".format(e)) self._should_stop = True raise def _receive(self, receiver, end_time): self._schedule_interval_logger(end_time, "Receiver " + str(self)) try: with receiver: while end_time > datetime.utcnow() and not self._should_stop: _logger.info("RECEIVE LOOP") try: if self.receive_type == ReceiveType.pull: batch = receiver.receive_messages(max_message_count=self.max_message_count, max_wait_time=self.max_wait_time) elif self.receive_type == ReceiveType.push: batch = receiver._get_streaming_message_iter(max_wait_time=self.max_wait_time) else: batch = [] for message in batch: self.on_receive(self._state, message, receiver) try: if self.should_complete_messages: receiver.complete_message(message) except MessageAlreadySettled: # It may have been settled in the plugin callback. pass self._state.total_received += 1 #TODO: Get EnqueuedTimeUtc out of broker properties and calculate latency. Should properties/app properties be mostly None? if end_time <= datetime.utcnow(): break time.sleep(self.receive_delay) self.post_receive(self._state, receiver) except Exception as e: _logger.exception("Exception during receive: {}".format(e)) self._state.exceptions.append(e) if self.fail_on_exception: raise self._state.timestamp = datetime.utcnow() return self._state except Exception as e: _logger.exception("Exception in receiver {}".format(e)) self._should_stop = True raise def run(self): start_time = datetime.utcnow() end_time = start_time + (self._duration_override or self.duration) sent_messages = 0 received_messages = 0 with concurrent.futures.ThreadPoolExecutor(max_workers=4) as proc_pool: _logger.info("STARTING PROC POOL") senders = [proc_pool.submit(self._send, sender, end_time) for sender in self.senders] receivers = [proc_pool.submit(self._receive, receiver, end_time) for receiver in self.receivers] result = StressTestResults() for each in concurrent.futures.as_completed(senders + receivers): _logger.info("SOMETHING FINISHED") if each in senders: result.state_by_sender[each] = each.result() if each in receivers: result.state_by_receiver[each] = each.result() # TODO: do as_completed in one batch to provide a way to short-circuit on failure. result.state_by_sender = {s:f.result() for s,f in zip(self.senders, concurrent.futures.as_completed(senders))} result.state_by_receiver = {r:f.result() for r,f in zip(self.receivers, concurrent.futures.as_completed(receivers))} _logger.info("got receiever results") result.total_sent = sum([r.total_sent for r in result.state_by_sender.values()]) result.total_received = sum([r.total_received for r in result.state_by_receiver.values()]) result.time_elapsed = end_time - start_time _logger.critical("Stress test completed. Results:\n{}".format(result)) return result