Benchmarking tool for OpenStack Swift

swift-benchmark.py

#!/usr/bin/env python

import argparse
import keystoneclient  # From python-keystoneclient
import logging
import math
import os
import random
import requests  # from requests
import string
import statsd

from collections import Counter

LOG = None


def randname():
    """ Generates a random uppercase name """
    return "".join(random.choice(string.ascii_uppercase) for i in range(12))


def get_container_url(args):
    """ Returns the path to a container, and a token that could be used
    to access that container. """
    keystone = keystoneclient.v2_0.client.Client(username=args.username,
                                                 password=args.password,
                                                 tenant_name=args.tenant,
                                                 auth_url=args.auth_url)

    endpoints = keystone.service_catalog.get_endpoints('object-store')
    base_url = endpoints['object-store'][0][args.url_type]
    if args.force_http:
        base_url = base_url.replace('https', 'http')

    return base_url + "/" + args.container, keystone.auth_token


def classify(value):
    """ Rounds the value up to a humanly locigal value, e.g. 100, 2000 or 0.5
    """
    log = math.log10(value)
    log_base = math.trunc(log)
    log_sub = 10**(log - log_base)
    if log_base <= 0:
        value = int(10**log_base)
    else:
        value = 10**log_base
    if log_sub <= 1.0001:
        return value
    if log_sub < 2:
        return value*2
    if log_sub < 5:
        return value*5
    return 10 * value


def benchmark(args, container_url, token):
    """ THe workhorse of this benchmark script """

    # Generate some random data to be uploaded
    data = os.urandom(args.filesize * 1024)

    if args.keep_alive:
        session = requests.Session()
    else:
        session = requests

    session.put(container_url,
                auth=(args.username, args.password)).raise_for_status()

    put_times = []
    get_times = []

    for i in range(0, args.iterations):
        file_url = container_url + "/" + randname()
        response_put = session.put(
            file_url,
            auth=(args.username, args.password),
            headers={'Content-Type': 'image/jpeg'},
            data=data)

        seconds_put = response_put.elapsed.seconds * 1000
        microseconds_put = response_put.elapsed.microseconds / 1000.0
        total_put = seconds_put + microseconds_put
        put_times.append(total_put)

        # Download
        response_get = session.get(file_url,
                                   auth=(args.username, args.password))
        response_get.raise_for_status()
        assert response_get.content == data

        seconds_get = response_get.elapsed.seconds * 1000
        microseconds_get = response_get.elapsed.microseconds / 1000.0
        total_get = seconds_get + microseconds_get
        get_times.append(total_get)

        args.csv.write('%.0f;%.0f\n' % (total_put, total_get))

    gauge = statsd.Gauge('openstack.swift.benchmark')

    m = Counter(classify(t) for t in get_times)
    for order, count in sorted(m.items()):
        pct = (100.0 * count) / args.iterations
        args.histogram.write(
            "GET %4.0f%% <%ims\n" % (pct, order))
        gauge.send("get_%ims" % order, pct)

    m = Counter(classify(t) for t in put_times)
    for order, count in sorted(m.items()):
        pct = (100.0 * count) / args.iterations
        args.histogram.write(
            "PUT %4.0f%% <%ims\n" % (pct, order))
        gauge.send("put_%ims" % order, pct)

    get_times.sort()
    put_times.sort()

    args.histogram.write("\n  % GET (ms) PUT (ms)\n")

    for pct in [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]:
        get_percentile = get_times[int(pct*(len(get_times)-1) / 100.0)]
        put_percentile = put_times[int(pct*(len(put_times)-1) / 100.0)]
        args.histogram.write("%3i%8i%8i\n" % (pct, get_percentile, 
                                              put_percentile))
        gauge.send("get_%03ipct" % pct, get_percentile)
        gauge.send("put_%03ipct" % pct, put_percentile)


def delete_all(args, container_url, token):
    """ Cleans up the container created elsewhere """
    LOG.info("Cleaning up %s container", container_url)
    session = requests.Session()
    session.auth = (args.username, args.password)

    while True:
        response = session.get(container_url,
                               headers={'Accept': 'application/json'})

        if response.status_code == 404:
            return

        for file in response.json():
            session.delete(container_url + "/" + file['name'])
        else:
            break

    session.delete(container_url)


def main():
    global LOG
    parser = argparse.ArgumentParser(
        description="Detect existing resources for disabled tenants")

    parser.add_argument('--username', default=os.environ.get('OS_USERNAME'),
                        required='OS_USERNAME' not in os.environ,
                        help='Openstack username')
    parser.add_argument('--password', default=os.environ.get('OS_PASSWORD'),
                        required='OS_PASSWORD' not in os.environ,
                        help='Openstack password')
    parser.add_argument('--tenant', default=os.environ.get('OS_TENANT_ID'),
                        required='OS_TENANT_ID' not in os.environ,
                        help='Openstack tenant')
    parser.add_argument('--auth-url', default=os.environ.get('OS_AUTH_URL'),
                        required='OS_AUTH_URL' not in os.environ,
                        help='Openstack authentication url')
    parser.add_argument('--iterations', default=100, type=int,
                        help="The number of iterations to be tested")
    parser.add_argument('--filesize', default=2200, type=int,
                        help="Size if the file (in Kb) to be tested.")
    parser.add_argument('--container', default=randname(), 
                        help="Name of the container to do the benchmark on")
    parser.add_argument('--force-http', action='store_true',
                        help="use http, even if keystone indicates http")
    parser.add_argument('--keep-alive', action='store_true',
                        help="Keep connection alive between requests.")
    parser.add_argument('--csv', default='-', type=argparse.FileType('w'),
                        help="File to output the CSV perfdata to")
    parser.add_argument('--histogram', type=argparse.FileType('w'),
                        help="File to output the histogram to", default='-')
    parser.add_argument('--url-type', default="publicURL",
                        help='What URL to get from the service catalog')
    parser.add_argument('--statsd-host', help='Statsd host')
    parser.add_argument('--statsd-port', default=8125, type=int,
                        help='Statsd host')

    parser.add_argument('--verbose', action='store_true')
    parser.add_argument('--debug', action='store_true')

    args = parser.parse_args()

    if args.debug:
        log_level = 'DEBUG'
    elif args.verbose:
        log_level = 'INFO'
    else:
        log_level = 'WARNING'

    logging.basicConfig(level=log_level)
    LOG = logging.getLogger()

    LOG.debug("Initializing Statsd connection settings")
    statsd.Connection.set_defaults(host=args.statsd_host or '127.0.0.1',
                                   port=args.statsd_port,
                                   disabled=not args.statsd_host)

    url, token = get_container_url(args)

    try:
        benchmark(args, url, token)
    finally:
        delete_all(args, url, token)

if __name__ == '__main__':
    main()