benchtcpplain.py

from __future__ import print_function

import multiprocessing
import numpy as np
import os
import socket
try:
    import socketserver
except ImportError:
    import SocketServer as socketserver
import struct
import threading
try:
    from time import perf_counter as clock
except ImportError:
    from time import time as clock


def read(sock):
    header = bytearray()
    while len(header) < 8:
        header += sock.recv(8 - len(header))
    msg_size = struct.unpack('L', header)[0]

    msg = bytearray(msg_size)
    pos = 0
    buf = memoryview(msg)
    while pos < msg_size:
        nbytes = sock.recv_into(buf[pos:])
        pos += nbytes

    return msg


def write(sock, msg):
    buffers = [struct.pack('L', len(msg)), memoryview(msg)]
    if 0:
        sock.sendmsg(buffers)
    else:
        for buf in buffers:
            sock.sendall(buf)


def serve_client(sock):
    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    data = read(sock)
    print('server', len(data))
    write(sock, data)


def server_loop(serv_sock, niters):
    serv_sock.listen(5)
    for i in range(niters):
        sock, addr = serv_sock.accept()
        serve_client(sock)


class RequestHandler(socketserver.BaseRequestHandler):

    def handle(self):
        serve_client(self.request)


class Server(socketserver.TCPServer):
    pass


USE_PROCESS = 0


def run_bench(nbytes, niters):
    data = b"x" * nbytes

    server = Server(('127.0.0.1', 0), RequestHandler)
    addr = server.server_address
    if USE_PROCESS:
        process = multiprocessing.Process(target=server.serve_forever)
    else:
        process = threading.Thread(target=server.serve_forever)
    process.start()

    start = clock()

    for i in range(niters):
        client = socket.create_connection(addr)
        client.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        write(client, data)
        msg = read(client)
        client.close()
        assert len(msg) == nbytes
        assert msg[-1:] == b"x"
        #print('.', end='', flush=True)
        print('client', len(msg))
        del msg
    print()

    end = clock()

    if USE_PROCESS:
        process.terminate()
    else:
        server.shutdown()
        process.join()

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


if __name__ == '__main__':
    run_bench(100 * 1000**2, 5)  # 100 MB