test_metrics.py

"""

Test Metrics: Executes identical workloads across a multiple clusters with the same configurations.

Usage:
    test_metrics.py [--num-runs=<num>] [--start-only=<instance, test case>]

Options:
    -h --help               Show this message.
    --version               Show version.
    --verbose               Show additional details about the operation being performed.
    --num-runs=<num>        The number of times to sample each test case [default: 3].
    --start-only=<instance,test case> 
        Start the cluster for a particular test case and run no tests  

"""
import analyze
import json
import os
import Queue
import random
import sys
import requests
import time
import traceback
import urllib2
import utils
import uuid

from datetime import datetime
from docopt import docopt
from fabric.api import env, execute, local, settings, sudo, task
from fabric.context_managers import lcd
from fabric.contrib.files import upload_template
from fabric.operations import get, put
from getpass import getpass
from riak import BucketType, RiakCluster, derive_bucket_uri
from threading import Thread


def get_argument(arguments, key, default_value):
    present = False
    value = arguments[key]
    if value:
        present = True
        value = utils.normalize_value_types(default_value, value)
    else:
        value = default_value

    if isinstance(value, tuple):
        return (present, ) + value
    return present, value


HARNESS_HOME_DIR = os.path.dirname(os.path.abspath(__file__))

SESSION_ID = datetime.now().strftime("%d%m%y_%H%M%S")
TEMPLATE_DIR = os.path.join(HARNESS_HOME_DIR, "templates")
RESULTS_BASE_DIR = os.path.join(HARNESS_HOME_DIR, "results")
DEFAULT_NUM_RUNS = 3

BENCH_HOME_DIR = "/usr/local/libexec/basho_bench"
BENCH_CONF_DIR = os.path.join(HARNESS_HOME_DIR, "bench_conf")

RESULT_FILENAME_TEMPLATE = "{cluster}-{scenario}-{count}.json"
TEST_INDEX_NAME = "test_index"


class ObjectConsumer:
    def __init__(self, id, cluster, node, bucket_type, bucket, object_queue):
        self._id = id
        self._object_queue = object_queue
        self._base_uri = "http://{host}:{port}/{bucket}" \
            .format(host=cluster.host, port=cluster.get_http_port(node),
                    bucket=derive_bucket_uri(bucket_type, bucket))
        self._object_count = 0

    def consume(self):
        print "Starting consumer {id} using base uri {base_uri}.".format(id=self._id,
                                                                         base_uri=self._base_uri)
        headers = {'Content-Type': 'application/json'}
        while True:
            (key, data) = self._object_queue.get()
            url = "{base_uri}/keys/{key}".format(base_uri=self._base_uri, key=key)
            try:
                r = requests.put(url=url, headers=headers, data=json.dumps(data))
                r.raise_for_status()

                self._update_progress()
            finally:
                self._object_queue.task_done()
                if self._object_queue.empty():
                    break

    def _update_progress(self):
        self._object_count += 1
        if self._object_count % 300 == 0:
            print "Consumer {id} has inserted {object_count} objects using base uri {base_uri}" \
                .format(id=self._id, object_count=self._object_count, base_uri=self._base_uri)


def generate_search_corpus(cluster, bucket_type, bucket, num_objects=100000,
                           num_consumers=50):
    object_queue = Queue.Queue(1000)

    nodes = xrange(1, 9)
    for i in xrange(1, num_consumers + 1):
        consumer = ObjectConsumer(i, cluster, random.choice(nodes), bucket_type, bucket, object_queue)
        thread = Thread(target=consumer.consume)
        thread.daemon = True
        thread.start()

    for counter in xrange(1, num_objects + 1):
        object_queue.put((uuid.uuid4(), {"counter": counter}))

    object_queue.join()


class StatsTestCase:
    def __init__(self, name, clusters, bench_conf_file, override_options={},
                 bench_home=BENCH_HOME_DIR, bucket_type=None, start_only=False,
                 data_generator=utils.noop):

        self._name = name
        self._bucket = "test"
        self._count = 0
        self._clusters = clusters
        self._bench_home = bench_home
        self._bench_conf_file = bench_conf_file
        self._override_options = override_options
        self._bucket_type = bucket_type
        self._start_only = start_only
        self._data_generator = data_generator

    @property
    def name(self):
        return self._name

    def run(self, results_dir, cluster_filter=None):

        self._count += 1

        clusters = self._clusters
        if cluster_filter:
            clusters = [cluster for cluster in self._clusters if cluster.name in cluster_filter]

        stats = {}
        for cluster in clusters:
            run_name = "-".join([cluster.name, self._name, str(self._count)])
            print "Executing run {count} of test case {test_case} against {cluster}" \
                .format(count=self._count, test_case=self._name, cluster=cluster.name)

            # Reset the cluster state for test execution
            cluster.reset(self._override_options)

            if self._bucket_type and "search_index" in self._bucket_type.options:
                cluster.create_index(self._bucket_type.options["search_index"])
                # time.sleep(30)

            if self._bucket_type:
                cluster.create_bucket_type(self._bucket_type)

            if not self._start_only:
                self._data_generator(cluster, self._bucket_type, self._bucket)

                # Execute Basho Bench ...
                with lcd(results_dir):
                    local(utils.build_command(os.path.join(self._bench_home, "basho_bench"),
                                              ["-n", run_name, self._bench_conf_file]))

                cluster.collect_logs(os.path.join(results_dir, "logs", run_name))

                stats[cluster.name] = cluster.capture_stats()

            print "Completed run {count} of test case {test_case} against {cluster}" \
                .format(count=self._count, test_case=self._name, cluster=cluster.name)

        return utils.TestResult(self._name, stats["folsom"], stats["exometer"])

    def __str__(self):
        return "Stats Test Case (name: {name}, count: {count}, riak options: {options})" \
            .format(name=self._name, count=self._count, options=self._riak_options)

    def __repr__(self):
        return self.__str__()


def main(arguments):
    try:

        username = raw_input("Username: ")
        env.password = getpass("sudo password: ")

        env.key_filename = "~/.ssh/id_rsa"

        _, num_runs = get_argument(arguments, "--num-runs", default_value=DEFAULT_NUM_RUNS)
        start_only, instance, test_case = get_argument(arguments, "--start-only",
                                                       default_value=(None, set()))
        print "start_only =", start_only, "instance =", instance, "test_case =", test_case

        instances = None
        if instance:
            instances = [instance]

        # TODO can this if block be removed?
        if not os.path.exists(RESULTS_BASE_DIR):
            os.makedirs(RESULTS_BASE_DIR)

        results_dir = os.path.join(RESULTS_BASE_DIR, SESSION_ID)
        os.makedirs(results_dir)
        print "Writing results to {results_dir}".format(results_dir=results_dir)

        # TODO Allow specifiction of the host(s) through a command line parameter ...
        clusters = [RiakCluster("folsom", "r2s29", "folsom", username, TEMPLATE_DIR, num_nodes=8),
                    RiakCluster("exometer", "r2s29", "exometer", username, TEMPLATE_DIR, num_nodes=8)]

        test_cases = [StatsTestCase("bitcask", clusters, BENCH_CONF_DIR + "/punisher.conf"),
                      StatsTestCase("leveldb", clusters, BENCH_CONF_DIR + "/punisher.conf", \
                                    {"backend": "leveldb"}),
                      StatsTestCase("memory", clusters, BENCH_CONF_DIR + "/punisher.conf", \
                                    {"backend": "memory"}),
                      # StatsTestCase("strong-consistency", clusters, BENCH_CONF_DIR + "/sc-punisher.conf", \
                      # override_options={ "consistency" : "on" }, \
                      # bucket_type=BucketType("sc", { "consistent" : True, "n_val" : 5 }),
                      # start_only=start_only)
                      StatsTestCase("yz_search", clusters, BENCH_CONF_DIR + "/yz-punisher.conf", \
                                    override_options={"search": "on"}, \
                                    bucket_type=BucketType("test_search", {"search_index": TEST_INDEX_NAME}),
                                    start_only=start_only, data_generator=generate_search_corpus)
                      ]

        print "Running each test case {num_runs} times.".format(num_runs=num_runs)
        for test_case in test_cases:
            for i in range(1, num_runs + 1):
                test_result = test_case.run(results_dir, instances)

                utils.write_file(os.path.join(results_dir, RESULT_FILENAME_TEMPLATE.format(cluster="folsom", \
                                                                                           scenario=test_result.name,
                                                                                           count=i)),
                                 test_result.folsom_stats, utils.json_formatter)
                utils.write_file(os.path.join(results_dir, RESULT_FILENAME_TEMPLATE.format(cluster="exometer", \
                                                                                           scenario=test_result.name,
                                                                                           count=i)),
                                 test_result.exometer_stats, utils.json_formatter)
                analyze.analyze_test_results([test_result], results_dir)

        print "Completed execution of {num_tests} tests. Results available in {results_dir}." \
            .format(num_tests=len(test_cases) * num_runs, results_dir=results_dir)

        return 0

    except Exception as e:

        print "Test failed due " + str(e)
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    arguments = docopt(__doc__, version="Test Metrics 0.1.0")
    sys.exit(main(arguments))