benchtcpasynciostreams.py

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

import asyncio
#import uvloop
#asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())


async def read_frame(stream):
    nbytes = await stream.readexactly(8)
    nbytes = struct.unpack('L', nbytes)[0]
    data = await stream.readexactly(nbytes)
    return data


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


async def run_server(reader, writer):
    msg = await read_frame(reader)
    print('server', len(msg))
    await write_frame(writer, msg)


async def run_client(reader, writer, msg):
    await write_frame(writer, msg)
    msg = await read_frame(reader)
    print('client', len(msg))
    writer.close()


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

    loop = asyncio.get_event_loop()
    #evt_done = asyncio.Event()
    server = await asyncio.start_server(run_server, '127.0.0.1', 8000)

    start = clock()

    for i in range(niters):
        #evt_done.clear()
        reader, writer = await asyncio.open_connection('127.0.0.1', 8000)
        await run_client(reader, writer, msg=data)

    end = clock()

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


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(f())
    loop.close()