Curio Benchmark of Large Data Transfer

curbench.py

# A Curio adaption of benchmark code at:
# 
# https://gist.github.com/pitrou/202221ca9c9c74c0b48373ac89e15fd7

import struct
try:
    from time import perf_counter as clock
except ImportError:
    from time import time as clock

import sys
sys.path.insert(0, '..')

import curio

async def read_frame(stream):
    nbytes = await stream.read_exactly(8)
    nbytes = struct.unpack('L', nbytes)[0]
    data = bytearray(nbytes)
    nread = await stream.readinto(data)
    assert nread == nbytes
    return data


async def write_frame(stream, msg):
    await stream.write(struct.pack('L', len(msg)))
    await stream.write(msg)


async def handler(sock, addr):
    async with sock:
        s = sock.as_stream()
        msg = await read_frame(s)
        print('server', len(msg))
        await write_frame(s, msg)

async def run_client(sock, msg):
    async with sock:
        s = sock.as_stream()
        await write_frame(s, msg)
        msg2 = await read_frame(s)
        print('client', len(msg))

async def f():
    data = b"x" * (100 * 1000**2)  # 100 MB
    niters = 5

    serv_task = await curio.spawn(curio.tcp_server('127.0.0.1', 8000, handler))

    start = clock()

    for i in range(niters):
        sock = await curio.open_connection('127.0.0.1', 8000)
        await run_client(sock, msg=data)

    end = clock()
    dt = end - start
    rate = len(data) * niters / dt
    print("duration: %s => rate: %d MB/s"
          % (dt, rate / 1e6))
    await serv_task.cancel()

if __name__ == "__main__":
    curio.run(f())