object-bench.py

#!/usr/bin/env python
"""
Direct object-server simple bench test
"""
from __future__ import print_function

import argparse
from multiprocessing import Process
import sys
import time

try:
    import http.client as httplib
except ImportError:
    import httplib


def do_get(p, args):
    conn = httplib.HTTPConnection('127.0.0.1', args.port)
    conn.connect()
    for n in range(args.operations):
        conn.request('GET', '/sdb1/{0}/a/c/o-{1}'.format(p, n))
        resp = conn.getresponse()
        etag = resp.getheader('etag')
        first_chunk = chunk = resp.read(args.chunk_size)
        while chunk:
            chunk = resp.read(args.chunk_size)
        if resp.status != 200:
            print(resp.status)
            print(first_chunk)
        else:
            print(etag)


def do_put(p, args):
    conn = httplib.HTTPConnection('127.0.0.1', args.port)
    conn.connect()
    for n in range(args.operations):
        url = '/sdb1/{0}/a/c/o-{1}'.format(p, n)
        conn.putrequest('PUT', url)
        conn.putheader('Content-Length', args.num_chunks * args.chunk_size)
        conn.putheader('X-Timestamp', str(time.time()))
        conn.putheader('Content-Type', 'application/bytes')
        conn.endheaders()
        chunk = '\x00' * args.chunk_size
        for i in range(args.num_chunks):
            conn.send(chunk)
        resp = conn.getresponse()
        etag = resp.getheader('etag')
        print(etag)
        first_chunk = chunk = resp.read(args.chunk_size)
        while chunk:
            chunk = resp.read(args.chunk_size)
        if resp.status != 201:
            print(resp.status)
            print(chunk)


method_map = {
    'GET': do_get,
    'PUT': do_put,
}

parser = argparse.ArgumentParser()
parser.add_argument('--method', choices=method_map.keys(),
                    default='PUT', help='type of benchmark')
parser.add_argument('--port', type=int, default=6010,
                    help='port of serverish like thing')
parser.add_argument('--workers', type=int, default=10,
                    help='number of processes')
parser.add_argument('--num-chunks', type=int, default=1000,
                    help='number of the network chunks to send')
parser.add_argument('--chunk-size', type=int, default=64 * 2 ** 10,
                    help='size of the network chunks to send')
parser.add_argument('--operations', type=int, default=1,
                    help="number of operations to perform")



def main():
    args = parser.parse_args()
    target = method_map[args.method]
    print('Begin!')
    workers = []
    start = time.time()
    for i in range(args.workers):
        p = Process(target=target, args=(i, args))
        p.start()
        workers.append(p)
    for p in workers:
        p.join()
    print('Finished: {0}'.format(time.time() - start))


if __name__ == '__main__':
    sys.exit(main())