bench.py

# example zmq benchmark
# LICENSE: CC-0 (Public Domain)

from multiprocessing import Process
import os
import time

import zmq


def receiver(url):
    ctx = zmq.Context()
    s = ctx.socket(zmq.ROUTER)
    s.bind(url)
    # handshake to be sure we are ready
    msg = s.recv_multipart()
    assert msg[1] == b'BEGIN'
    s.send_multipart(msg)
    while True:
        msg = s.recv_multipart()
        if msg[1] == b'DONE':
            s.send_multipart(msg)
            break
    s.close()
    ctx.term()

def sender(url, n, size):
    ctx = zmq.Context()
    s = ctx.socket(zmq.DEALER)
    s.connect(url)
    msg = os.urandom(size)
    # make sure we are ready
    s.send(b'BEGIN')
    assert s.recv() == b'BEGIN'
    time.sleep(1)
    start = time.monotonic()
    for i in range(n):
        s.send(msg)
    sent = time.monotonic()
    s.send(b'DONE')
    assert s.recv() == b'DONE'
    s.close()
    ctx.term()
    done = time.monotonic()
    return (sent-start, done-start)

def main(n, size):
    url = 'tcp://127.0.0.1:5555'
    print("Sending %g msgs of size %g B" % (n, size))
    
    r = Process(target=receiver, args=(url,), daemon=True)
    r.daemon = True
    r.start()
    sent, done = sender(url, n, size)
    sent_sec = n // sent
    done_sec = n // done
    if done < 1:
        print("Test too short, try a bigger number by ~%.1fx" % (1/done))
    print("Sent %g msgs in %.3f seconds" % (n, done))
    print(" sent: %8i msgs/sec" % sent_sec)
    print("recvd: %8i msgs/sec" % done_sec)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('number', default=1000, nargs='?', type=int, help="Number of messages to send")
    parser.add_argument('size', default=1000, nargs='?', type=int, help="Size of messages (in bytes)")
    opts = parser.parse_args()
    main(opts.number, opts.size)