nfr_bench.py

#!/usr/bin/env python
from __future__ import with_statement

# Requires httplib2:
#     easy_install httplib2
# or:
#     pip install httplib2

import re
import sys
import time

from httplib2 import Http
from threading import Thread

METHOD = 'GET'
TIMEOUT = 30
MAX_REDIRECTIONS = 3

HEADERS = {
    'User-Agent': 'Mozilla/5.0 (compatible; NFR Benchmark)',
    'Accept-Charset': "utf8,ISO-8859-1;q=.5",
}

def response_time_for_url(url):
    """Returns the response time of the given url in ms"""
    http = Http(timeout=TIMEOUT)
    before = time.time()
    r, c = http.request(url,
                        method=METHOD,
                        headers=HEADERS,
                        redirections=MAX_REDIRECTIONS)
    if int(r['status']) == 200:
        return int((time.time() - before) * 1000)
    else:
        return -1

def response_times_for_urls(urls, threads, process_until):
    result = [] # Python lists are thread safe.
    finished = []

    def worker(all_urls):
        while True:
            for url in all_urls:
                if time.time() > process_until:
                    finished.append(True)
                    return
                result.append((url, response_time_for_url(url),))

    for i in range(threads):
        w = Thread(target=worker, args=(list(urls),))
        w.setDaemon(True)
        w.start()

    while True:
        if len(finished) == threads:
            return result
        time.sleep(.1)

def print_stats(urls, runtime):
    print "\t".join(['Avg', '75%', 'URL',])
    print "-" * 78

    all = []

    def average75(numbers):
        quarter = int(round(len(numbers)/4.0))
        
        avg = sum(numbers) / len(numbers)
        return sum(numbers[-quarter:]) / quarter 

    for url, times in urls.items():
        times.sort()
        all += times

        avg = sum(times) / len(times)
        avg75 = average75(times)
        print "%6d %6d %s" % (avg, avg75, url,)

    avg = sum(all) / len(all)
    avg75 = average75(all)

    reqps = len(all) / float(runtime)

    print "=" * 78
    print "%6d %6d Total %d req over %ds (%.2f req/s)" % (avg, avg75, len(all), runtime, reqps,)


if __name__ == "__main__":
    if len(sys.argv) != 4:
        print "Usage: nfr_bench.py <input file> <thread num> <runtime (sec)>"
        sys.exit(1)

    runtime = int(sys.argv.pop())
    threads = int(sys.argv.pop())
    input = sys.argv.pop()

    urls = []
    with open(input) as file:
        urls = [re.sub(r"\n$", "", line) for line in file.readlines()]

    res = response_times_for_urls(urls, threads, time.time() + runtime)

    url_dict = {}
    for url, ms in res:
        if url in url_dict:
            url_dict[url].append(ms)
        else:
            url_dict[url] = [ms]

    print_stats(url_dict, runtime)