Skip to content
Permalink
Branch: master
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
282 lines (235 sloc) 10.1 KB
import argparse
import asyncio
from collections import Counter
import datetime
import dateutil.parser
import json
import re
import time
import aiohttp
class LoadTester:
def __init__(self, loop, request_generator_factory, response_accumulator, host, port, speed_multiplier=1):
"""
:param loop: the event loop to run on
:param request_generator_factory: takes no args and returns a request_generator
:param response_accumulator: takes a dictionary of information about the response and collects useful information
:param host:
:param port:
:param speed_multiplier:
"""
self.loop = loop
self.request_generator_factory = request_generator_factory
self.response_accumulator = response_accumulator
self.host = host
self.port = port
self.speed_multiplier = speed_multiplier
self.log_initial_time = datetime.datetime.min # sentinel value meaning that this hasn't been evaluated
self.real_initial_time = None
self.log_delta_time = datetime.timedelta(0) # sentinel value meaning that this hasn't been evaluated
async def run(self, run_time_minutes):
run_time_secs = run_time_minutes*60
num_sent_requests = 0
start_time = time.time()
async with aiohttp.ClientSession(loop=self.loop) as session:
for request in self._requests():
delta_secs = time.time() - start_time
if delta_secs >= run_time_secs:
break
await asyncio.sleep(min(request['sleep_time'], run_time_secs - delta_secs))
num_sent_requests += 1
fut = asyncio.ensure_future(
self._fetch(session, request['url'], request['body'])
)
fut.add_done_callback(self._callback)
pending_tasks = [
task for task in asyncio.all_tasks()
if (
not task.done()
and not task == asyncio.current_task()
)
]
num_outstanding_requests = len(pending_tasks)
for task in pending_tasks:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
seconds_behind = max(-request['sleep_time'], 0)
return {
'run_time_minutes': delta_secs/60,
'num_sent_requests': num_sent_requests,
'average_requests_per_second': num_sent_requests/delta_secs,
'num_outstanding_requests': num_outstanding_requests,
'seconds_behind': seconds_behind, # the last request we send out should have been sent out this many seconds ago
'percentage_behind': seconds_behind/delta_secs, # ideally this should be zero meaning that the simulation is keeping up; max value is 1.0
}
def _requests(self):
while True:
request_generator = self.request_generator_factory()
if self.real_initial_time is None:
self.real_initial_time = datetime.datetime.now()
elif self.log_delta_time == datetime.timedelta(0):
# then the log has ended for the 1st time and we need to determine the length of the log using the last
# request datetime
self.log_delta_time = request.timestamp - self.log_initial_time
# every time we loop through the log we move log_initial_time backward the delta time of the log
self.log_initial_time = self.log_initial_time - self.log_delta_time
for request in request_generator:
yield {
'url': 'http://{host}:{port}/{url_path}'.format(host=self.host, port=self.port, url_path=request.url_path),
'body': request.json_body,
'sleep_time': self._get_sleep_time(request.timestamp),
}
def _get_sleep_time(self, log_time):
if self.log_initial_time == datetime.datetime.min:
self.log_initial_time = log_time
real_time = datetime.datetime.now()
# Get real time till next event considering
log_time_passed_secs = (log_time - self.log_initial_time).total_seconds()
scaled_time_passed = (real_time - self.real_initial_time).total_seconds() * self.speed_multiplier
scaled_time_till_next_event = log_time_passed_secs - scaled_time_passed
real_time_till_next_event = scaled_time_till_next_event / self.speed_multiplier
return real_time_till_next_event # a.k.a. sleep_time
@staticmethod
async def _fetch(session, url, body):
async with session.post(url, json=body) as response:
return Response(
status=response.status,
json=await response.json(),
)
def _callback(self, fut):
try:
self.response_accumulator.process_response_information(fut.result())
except asyncio.CancelledError:
pass
class Request:
"""A response object.
For now, this is basically a placeholder as the API isn't stable. This is consumed by
ElasticsearchResponseAccumulator.process_response_information
"""
def __init__(self, timestamp, method, url_path, json_body):
self.timestamp = timestamp
self.method = method
self.url_path = url_path
self.json_body = json_body
class Response:
"""A response object.
For now, this is basically a placeholder as the API isn't stable. This is consumed by
ElasticsearchResponseAccumulator.process_response_information
"""
def __init__(self, status, json):
self.status = status
self.json = json
class ElasticsearchRequestLogParser:
"""Parses request log file and for each request returns request datetime, method, url_path, and json_body of the request."""
def __init__(self, log_file):
self.log_file = log_file
def __iter__(self):
regex = re.compile(
r'^\[(?P<timestamp>.*?)\]\[.*?\]\[(?P<request_type>.*?)\]\s*\[(?P<index>.*?)\].*source\[(?P<source>.*)\],\s*extra_source\[(?P<extra_source>.*)\]')
for line in self.lines_of_log_file():
match = regex.match(line)
if not match:
raise Exception(line)
match_dict = match.groupdict()
request_type = match_dict['request_type'].split('.')[-1]
if request_type != 'query':
continue
yield Request(
timestamp=dateutil.parser.parse(match_dict['timestamp'].split(',')[0]),
method='POST',
url_path='/{index}/_search'.format(index=match_dict['index']),
json_body=self._get_body(match_dict),
)
def lines_of_log_file(self):
# this function is made for simpler testing
with open(self.log_file, 'r') as f:
for line in f:
yield line
@staticmethod
def _get_body(match_dict):
json_body = {}
if match_dict['source']:
json_body = json.loads(match_dict['source'])
if match_dict['extra_source']:
json_body.update(json.loads(match_dict['extra_source']))
return json_body
class ElasticsearchResponseAccumulator:
"""
Takes response dictionaries and accumulates important information.
The input to ElasticsearchResponseAccumulator.process_response_information comes from the output of
LoadTester._fetch, but the API here isn't stable yet
"""
def __init__(self):
self.total_took_time = 0
self.completion_status_counts = Counter()
def process_response_information(self, response):
self.completion_status_counts.update([response.status])
if response.status == 200:
self.total_took_time += response.json['took']
def get_summary(self):
return {
'completion_status_counts': dict(self.completion_status_counts),
'average_time_per_successful_request': self.total_took_time / self.completion_status_counts[200],
}
async def main(loop, log_file, host, port, speed_multiplier, run_time_minutes):
def request_generator_factory():
return ElasticsearchRequestLogParser(log_file=log_file)
response_accumulator = ElasticsearchResponseAccumulator()
load_tester = LoadTester(
loop=loop,
request_generator_factory=request_generator_factory,
response_accumulator=response_accumulator,
host=host,
port=port,
speed_multiplier=speed_multiplier,
)
run_information = await load_tester.run(run_time_minutes)
return {
'run_information': run_information,
'accumulator_information': response_accumulator.get_summary(),
}
if __name__ == '__main__':
# run with python load_tester_old.py --log_file='/Users/johnb/Desktop/elasticsearch_slowlog.log' --host='qa-core-es3' --port=9200 --speed_multiplier=100
parser = argparse.ArgumentParser(description='Load test Elasticsearch by replaying the slowlog file.')
parser.add_argument(
'--log_file',
help='the elasticsearch slowlog file',
)
parser.add_argument(
'--host',
help='the host to load test',
)
parser.add_argument(
'--port',
help='the port',
)
parser.add_argument(
'--speed_multiplier',
type=float,
help='the speed_multiplier - 1 is realtime, 0.5 is half time',
)
parser.add_argument(
'--run_time_minutes',
type=float,
help='amount of wallclock time to keep the load test running',
)
args = parser.parse_args()
loop = asyncio.get_event_loop()
log_file = '/Users/johnb/Desktop/elasticsearch_slowlog.log'
host = 'qa-core-es3'
port = 9200
speed_multiplier = 100
run_result = loop.run_until_complete(
main(
loop=loop,
log_file=args.log_file,
host=args.host,
port=args.port,
speed_multiplier=args.speed_multiplier,
run_time_minutes=args.run_time_minutes,
)
)
loop.close()
print(json.dumps(run_result, indent=4))
You can’t perform that action at this time.