stress_test.py

import asyncio
import json
import os
import statistics
from multiprocessing.pool import Pool

import aiohttp

URL = "http://geth.mainnet.ethnodes.brainbot.com:8545"
HEADERS = {
    "Content-Type": "application/json",
    "User-Agent": "Web3.py/5.7.0/<class 'web3.providers.rpc.HTTPProvider'>",
}
LATEST_BLOCK_REQ = b'{"jsonrpc": "2.0", "method": "eth_blockNumber", "params": [], "id": 0}'
TIMEOUT = aiohttp.ClientTimeout(sock_read=10)


async def time_request(loop, session, data):
    start = loop.time()
    async with session.post(URL, data=data, timeout=TIMEOUT, headers=HEADERS) as response:
        await response.text()
        return loop.time() - start


async def timing_worker(number_of_requests: int, from_block_hex: str, to_block_hex: str):
    loop = asyncio.get_event_loop()

    async with aiohttp.ClientSession() as session:
        request = {
            "jsonrpc": "2.0",
            "method": "eth_getLogs",
            "params": [{"fromBlock": from_block_hex, "toBlock": to_block_hex}],
            "id": 1,
        }
        data = json.dumps(request)
        done, _ = await asyncio.wait(
            [time_request(loop, session, data) for _ in range(number_of_requests)]
        )
        return {future.result() for future in done}


def worker(number_of_requests: int, from_block_hex: str, to_block_hex: str):
    loop = asyncio.get_event_loop()
    timings = loop.run_until_complete(
        timing_worker(number_of_requests, from_block_hex, to_block_hex)
    )
    return timings


async def refresh_block_number():
    async with aiohttp.ClientSession() as session:
        request = await session.post(URL, data=LATEST_BLOCK_REQ, timeout=TIMEOUT, headers=HEADERS)

        response = await request.json()

        from_block_hex = response["result"]
        to_block_hex = hex(int(from_block_hex[2:], 16) - 1)

        return (from_block_hex, to_block_hex)


async def orchestrator(pool, cpu_count, number_of_requests):
    latest_block_result = await refresh_block_number()
    from_block_hex, to_block_hex = latest_block_result

    all_timings = list()
    worker_paramenters = [
        (number_of_requests, from_block_hex, to_block_hex) for _ in range(cpu_count)
    ]

    for worker_timings in pool.starmap(worker, worker_paramenters):
        all_timings.extend(worker_timings)

    quantiles = list(statistics.quantiles(all_timings, n=100))
    return (number_of_requests * cpu_count, quantiles[49], quantiles[89], quantiles[98])


def main():
    loop = asyncio.get_event_loop()
    cpu_count = os.cpu_count()

    with Pool(cpu_count) as pool:
        for number_of_requests in range(2, 1_000_000, 10):
            result = loop.run_until_complete(orchestrator(pool, cpu_count, number_of_requests))
            data = f"{result[0]};{result[1]};{result[2]};{result[3]}"
            print(data, flush=True)


if __name__ == "__main__":
    main()