from concurrent.futures import ProcessPoolExecutor
import os
import time

import zmq

NMSGS = 1000000
MSG_SIZE = 200
URL = 'ipc://pyzmq-test.sock'
# URL = 'tcp://127.0.0.1:5555'

def puller(n):
    ctx = zmq.Context()
    pull = ctx.socket(zmq.PULL)
    pull.bind(URL)
    tic = time.time()
    for i in range(n):
        pull.recv()
    toc = time.time()
    pull.close()
    ctx.term()
    return toc

def pusher(n):
    ctx = zmq.Context()
    push = ctx.socket(zmq.PUSH)
    msg = os.urandom(MSG_SIZE)
    push.connect(URL)
    tic = time.time()
    for i in range(n):
        push.send(msg)
    push.close()
    ctx.term()
    return tic

pool = ProcessPoolExecutor(1)
t = 0
n = 64

# keep running until a sample takes at least a few seconds
print("%8s / %8s = %8s msgs/sec" % ('n msgs', 't (sec)', ''))
while t < 4:
    f = pool.submit(puller, n)
    time.sleep(1)
    start = pusher(n)
    stop = f.result()
    t = stop - start
    msgs_sec = n / t
    print("%8i / %8.3g = %8i msgs/sec = %5i MB/sec" % (n, t, msgs_sec, 1e-6 * msgs_sec * MSG_SIZE))
    n *= 2