from __future__ import print_function

import os
import socket
import time

COUNT = 1000000
LINELEN = 1000
ADDR = ('127.0.0.1', 12345)

def experiment(TCP, makefile):
    if TCP:
        ls = socket.socket()
        ls.bind(ADDR)
        ls.listen(1)
    else:
        (ss, cs) = socket.socketpair()

    child = os.fork()
    if child:
        if TCP:
            (ss, addr) = ls.accept()
            ls.close()
        else:
            cs.close()
        if makefile:
            fileobj = ss.makefile('rb', 4096*2)
        else:
            fileobj = os.fdopen(ss.fileno(), 'rb', 4096*2)
        start = time.clock()
        for line in fileobj:
            pass
        end = time.clock()
        fileobj.close()
        ss.close()

        delay = end-start
        print ('{0} mode, {1} method. {2:.1f} lines per second ({3:.1f} MB/s). Delay is {4:.2f} seconds'.format(
            'TCP ' if TCP else 'UNIX',
            'makefile' if makefile else '  fdopen',
            COUNT/delay,
            COUNT*LINELEN/(delay*1024*1024),
            delay))
        os.waitpid(child, 0)
    else:
        if TCP:
            ls.close()
            cs = socket.socket()
            cs.connect(ADDR)
        else:
            ss.close()

        line = b'a' * (LINELEN-1) + b'\n'
        line = line * 1000
        for i in range(COUNT//1000):
            cs.sendall(line)
        cs.close()
        os._exit(0)

def main():
    experiment(True, True)
    experiment(True, False)
    experiment(False, True)
    experiment(False, False)

if __name__ == '__main__':
    main()