Skip to content

Instantly share code, notes, and snippets.

@vangheem
Last active June 11, 2019 08:13
Show Gist options
  • Save vangheem/00666c9cf264f883db119cca6c59016e to your computer and use it in GitHub Desktop.
Save vangheem/00666c9cf264f883db119cca6c59016e to your computer and use it in GitHub Desktop.
# Usage
# python load.py http://localhost:8080 --concurrency=100 --total=10000
import asyncio
import time
from functools import partial
import aiohttp
import argparse
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('url')
parser.add_argument('--concurrency', type=int, default=100)
parser.add_argument('--total', type=int, default=10000)
class Runner:
def __init__(self, arguments):
self.arguments = arguments
self.stats = {
'start': time.time(),
'scheduled': 0,
'finished': 0,
'tasks': []
}
async def retrieve(self, session):
try:
async with session.get(self.arguments.url) as resp:
await resp.text()
assert resp.status == 200
except aiohttp.client_exceptions.ServerDisconnectedError:
print('Error, server disconnected')
def schedule(self, session):
task = asyncio.ensure_future(self.retrieve(session))
self.stats['scheduled'] += 1
self.stats['tasks'].append(task)
task.add_done_callback(partial(self.done, session))
def done(self, session, fut):
self.print_status()
self.stats['tasks'].remove(fut)
self.stats['finished'] += 1
if self.stats['scheduled'] < self.arguments.total:
self.schedule(session)
async def run(self):
async with aiohttp.ClientSession() as session:
for _ in range(self.arguments.concurrency):
self.schedule(session)
while len(self.stats['tasks']) > 0:
await asyncio.sleep(1)
def print_status(self):
if self.stats['finished'] > 0:
now = time.time()
secs = (now - self.stats['start'])
per_sec = self.stats['finished'] / secs
print(f'{self.stats["finished"]}({len(self.stats["tasks"])}) '
f'- {per_sec:.2f}/req/s', end="\r")
if __name__ == '__main__':
args = parser.parse_args()
loop = asyncio.get_event_loop()
runner = Runner(args)
loop.run_until_complete(runner.run())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment