benchtcpgen.py

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

from tornado.tcpserver import TCPServer
from tornado.tcpclient import TCPClient
from tornado.ioloop import IOLoop
from tornado import gen

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


@gen.coroutine
def read(stream):
    nbytes = yield stream.read_bytes(8)
    nbytes = struct.unpack('L', nbytes)[0]
    data = yield stream.read_bytes(nbytes)
    raise gen.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 = b"x" * (100 * 1000**2)  # 100 MB
    niters = 5

    server = MyServer(max_buffer_size=int(1e9))
    server.listen(8000)

    client = TCPClient()

    start = clock()

    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('client', len(msg))
    print(len(msg))

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


if __name__ == '__main__':
    IOLoop().run_sync(f)