wrk.py

#!/usr/bin/env python
# coding: utf-8
# module for http benchmarking using wrk
# example command like
# ./wrk.py --connections '[1,2,3,4]' --filename ./benchmark_qa/query.json  --url http://x.y.z.k/hello.php --output 1.csv --duration 600

import subprocess
import tempfile
import urlparse
from contextlib import contextmanager
import os
from itertools import product
import inspect
import time

def format_script(url, filename, shuffle):
    """
    generate wrk lua script
    """
    params = {
        u'shuffle': int(shuffle),
        u'url': repr(url),
        u'filename': repr(filename),
    }

    template = ur'''
    counter = 0
    math.randomseed(os.time())

    function file_exists(file)
      local f = io.open(file, "rb")
      if f then f:close() end
      return f ~= nil
    end

    if {shuffle} then
        shuffle = function(paths)
          local j, k
          local n = #paths
          for i = 1, n do
            j, k = math.random(n), math.random(n)
            paths[j], paths[k] = paths[k], paths[j]
          end
          return paths
        end
    else
        shuffle = function(paths)
            return paths
        end
    end

    function read_lines(file)
      if not file_exists(file) then return {{}} end
      lines = {{}}
      for line in io.lines(file) do
        if not (line == '') then
          lines[#lines + 1] = line
        end
      end
      return shuffle(lines)
    end

    bodies = read_lines({filename})

    url = {url}
    counter = 0
    request = function()
        body = bodies[counter]
        counter = counter + 1
        if counter > #bodies then
            counter = 0
        end
        return wrk.format('POST', url, nil, body)
    end

    done = function(summary, latency, requests)
       io.write("------------------------------\n")
       for _, p in pairs({{50, 75, 90, 99, 99.999 }}) do
          n = latency:percentile(p)
          io.write(string.format("%g%%\t%d\n", p, n))
       end
       io.write(string.format("min\t%d\n", latency.min))
       io.write(string.format("max\t%d\n", latency.max))
       io.write(string.format("mean\t%d\n", latency.mean))
       io.write(string.format("requests\t%d\n", summary.requests))
       io.write(string.format("duration\t%d\n", summary.duration))
    end
'''
    return template.format(**params)


# In[3]:


@contextmanager
def generate_script(url, filename, shuffle=True, delete=True):
    content = format_script(url, filename, shuffle)
    f = tempfile.NamedTemporaryFile(delete=delete)
    f.write(content)
    f.flush()
    if not delete:
        print 'script name: ', f.name
    yield f.name
    f.close()

def run(filename, url, connections, threads, shuffle,
         timeout, duration, bin_path, delete, show=None):
    params = locals().copy()
    threads = min(12 if threads is None else treads, connections)
    rawurl = url
    url = urlparse.urlparse(url)
    host = url.scheme + '://' + url.netloc
    with generate_script(url.path, filename, shuffle=shuffle, delete=delete) as script_filename:
        cmd = [
            bin_path,
            '--latency',
            '-c', str(connections),
            '-t', str(threads),
            '--timeout', str(timeout) + 's',
            '--duration', str(duration) + 's',
            '-s', script_filename,
            host,
        ]
        content = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
        lines = content.split('\n')
        before = []
        after = []
        dest = before
        for line in lines:
            if line.startswith('----'):
                dest = after
                continue
            else:
                dest.append(line)
        if show or (show is None and os.stdin.isatty()):
            print '\n'.join(before)
        payload = dict([line.split('\t') for line in after if line])
        payload.update(params)
        return payload

def to_list(x):
    if type(x) == list:
        return x
    else:
        return [x]



def wrk(filename, url, connections=1, threads=None,
        shuffle=True, timeout=1, duration=30, bin_path='wrk',
        output=None, show=None, delete=True
):
    params = dict([(k, to_list(v)) for (k,v) in locals().items()])
    args = inspect.getargspec(run).args
    params = dict([(k,v) for (k,v) in params.items() if k in args])

    assert os.path.isfile(filename), 'filename {} not exists'.format(filename)

    keys = params.keys()
    values = params.values()

    timestamp = time.time()
    results = []
    for value in product(*values):
        params = dict(zip(keys, value))
        print params
        item = run(**params)
        item['timestamp'] = timestamp
        results.append(item)

    if output:
        import pandas as pd
        df = pd.DataFrame(results)
        df.to_csv(output)
    if not show:
        return results

if __name__ == "__main__":
    import fire
    fire.Fire(wrk)