504.py

import asyncio
from collections import defaultdict
from pprint import pprint
import socket
import sys

from aiohttp import ClientSession, TCPConnector

TIMEOUT = 30
MAX_CONNECTIONS = 100


async def get(url, session: ClientSession):
    try:
        print(f'Requesting {url}')
        async with session.get(url, timeout=TIMEOUT) as response:
            return response.status
    except asyncio.TimeoutError:
        return "TIMEOUT"


async def count_status_codes(url, count=1000):
    connector = TCPConnector(family=socket.AF_INET,
                             verify_ssl=False, limit=MAX_CONNECTIONS)
    results = defaultdict(int)
    async with ClientSession(connector=connector) as session:
        coros = [asyncio.ensure_future(
            get(url, session))
            for _ in range(count)]
        codes = await asyncio.gather(*coros)
        for code in codes:
            results[code] += 1
    pprint(dict(results))


if __name__ == '__main__':
    url = sys.argv[1]
    count = int(sys.argv[2])
    loop = asyncio.get_event_loop()
    loop.run_until_complete(asyncio.ensure_future(count_status_codes(url, count)))