iostream-benchmark.py

import struct
import time

import numpy as np
from tornado.tcpserver import TCPServer
from tornado.tcpclient import TCPClient
from tornado.ioloop import IOLoop
from tornado import gen


@gen.coroutine
def read(stream):
    nbytes = yield stream.read_bytes(8)
    nbytes = struct.unpack('L', nbytes)[0]
    data = yield stream.read_bytes(nbytes)
    return data


@gen.coroutine
def write(stream, msg):
    yield stream.write(struct.pack('L', len(msg)))
    yield stream.write(msg)


class MyServer(TCPServer):
    @gen.coroutine
    def handle_stream(self, stream, address):
        data = yield read(stream)
        print('server', len(data))
        yield write(stream, data)


@gen.coroutine
def f():
    data = bytes(np.random.randint(0, 255, dtype='u1', size=100000000).data)  # 100M
    server = MyServer()
    server.listen(8000)

    client = TCPClient()
    for i in range(5):
        stream = yield client.connect('127.0.0.1', 8000, max_buffer_size=int(1e9))
        yield write(stream, data)
        msg = yield read(stream)
    print(len(msg))


if __name__ == '__main__':
    start = time.time()
    IOLoop().run_sync(f)
    end = time.time()
    print(end - start)


# To profile
# python -m cProfile -o prof.out server.py
# snakeviz prof.out