#!/usr/bin/python
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# expects a Cassandra server to be running and listening on port 9160.
# (read tests expect insert tests to have run first too.)

from __future__ import with_statement

have_multiproc = False
try:
    from multiprocessing import Array as array, Process as Thread
    from uuid import uuid1 as get_ident
    array('i', 1) # catch "This platform lacks a functioning sem_open implementation"
    Thread.isAlive = Thread.is_alive
    have_multiproc = True
except ImportError:
    from threading import Thread
    from thread import get_ident
    from array import array
from hashlib import md5
import time, random, sys, os
from random import randint, gauss
from optparse import OptionParser

from thrift.transport import TTransport
from thrift.transport import TSocket
from thrift.transport import THttpClient
from thrift.protocol import TBinaryProtocol

try:
    from cassandra import Cassandra
    from cassandra.ttypes import *
except ImportError:
    # add cassandra directory to sys.path
    L = os.path.abspath(__file__).split(os.path.sep)[:-3]
    root = os.path.sep.join(L)
    _ipath = os.path.join(root, 'interface', 'thrift', 'gen-py')
    sys.path.append(os.path.join(_ipath, 'cassandra'))
    import Cassandra
    from ttypes import *
except ImportError:
    print "Cassandra thrift bindings not found, please run 'ant gen-thrift-py'"
    sys.exit(2)

try:
    from thrift.protocol import fastbinary
except ImportError:
    print "WARNING: thrift binary extension not found, benchmark will not be accurate!"

parser = OptionParser()
parser.add_option('-n', '--num-keys', type="int", dest="numkeys",
                  help="Number of keys", default=1000**2)
parser.add_option('-N', '--skip-keys', type="float", dest="skipkeys",
                  help="Fraction of keys to skip initially", default=0)
parser.add_option('-t', '--threads', type="int", dest="threads",
                  help="Number of threads/procs to use", default=50)
parser.add_option('-c', '--columns', type="int", dest="columns",
                  help="Number of columns per key", default=5)
parser.add_option('-S', '--column-size', type="int", dest="column_size",
                  help="Size of column values in bytes", default=34)
parser.add_option('-C', '--cardinality', type="int", dest="cardinality",
                  help="Number of unique values stored in columns", default=50)
parser.add_option('-d', '--nodes', type="string", dest="nodes",
                  help="Host nodes (comma separated)", default="localhost")
parser.add_option('-D', '--nodefile', type="string", dest="nodefile",
                  help="File containing list of nodes (one per line)", default=None)
parser.add_option('-s', '--stdev', type="float", dest="stdev", default=0.1,
                  help="standard deviation factor")
parser.add_option('-r', '--random', action="store_true", dest="random",
                  help="use random key generator (stdev will have no effect)")
parser.add_option('-f', '--file', type="string", dest="file", 
                  help="write output to file")
parser.add_option('-p', '--port', type="int", default=9160, dest="port",
                  help="thrift port")
parser.add_option('-m', '--unframed', action="store_true", dest="unframed",
                  help="use unframed transport")
parser.add_option('-o', '--operation', type="choice", dest="operation",
                  default="insert", choices=('insert', 'read', 'rangeslice',
                  'indexedrangeslice', 'multiget'),
                  help="operation to perform")
parser.add_option('-u', '--supercolumns', type="int", dest="supers", default=1,
                  help="number of super columns per key")
parser.add_option('-y', '--family-type', type="choice", dest="cftype",
                  choices=('regular','super'), default='regular',
                  help="column family type")
parser.add_option('-k', '--keep-going', action="store_true", dest="ignore",
                  help="ignore errors inserting or reading")
parser.add_option('-i', '--progress-interval', type="int", default=10,
                  dest="interval", help="progress report interval (seconds)")
parser.add_option('-g', '--keys-per-call', type="int", default=1000,
                  dest="rangecount",
                  help="amount of keys to get_range_slices or multiget per call")
parser.add_option('-l', '--replication-factor', type="int", default=1,
                  dest="replication",
                  help="replication factor to use when creating needed column families")
parser.add_option('-e', '--consistency-level', type="str", default='ONE',
                  dest="consistency", help="consistency level to use")
parser.add_option('-x', '--create-index', type="choice",
                  choices=('keys','keys_bitmap', 'none'), default='none',
                  dest="index", help="type of index to create on needed column families")

(options, args) = parser.parse_args()
 
total_keys = options.numkeys
n_threads = options.threads
keys_per_thread = total_keys / n_threads
columns_per_key = options.columns
supers_per_key = options.supers
# this allows client to round robin requests directly for
# simple request load-balancing
nodes = options.nodes.split(',')
if options.nodefile != None:
    with open(options.nodefile) as f:
        nodes = [n.strip() for n in f.readlines() if len(n.strip()) > 0]

#format string for keys
fmt = '%0' + str(len(str(total_keys))) + 'd'

# a generator that generates all keys according to a bell curve centered
# around the middle of the keys generated (0..total_keys).  Remember that
# about 68% of keys will be within stdev away from the mean and 
# about 95% within 2*stdev.
stdev = total_keys * options.stdev
mean = total_keys / 2

consistency = getattr(ConsistencyLevel, options.consistency, None)
if consistency is None:
    print "%s is not a valid consistency level" % options.consistency
    sys.exit(3)

# generates a list of unique, deterministic values
def generate_values():
    values = []
    for i in xrange(0, options.cardinality):
        h = md5(str(i)).hexdigest()
        values.append(h * int(options.column_size/len(h)) + h[:options.column_size % len(h)])
    return values

def key_generator_gauss():
    while True:
        guess = gauss(mean, stdev)
        if 0 <= guess < total_keys:
            return fmt % int(guess)
    
# a generator that will generate all keys w/ equal probability.  this is the
# worst case for caching.
def key_generator_random():
    return fmt % randint(0, total_keys - 1)

key_generator = key_generator_gauss
if options.random:
    key_generator = key_generator_random


def get_client(host='127.0.0.1', port=9160):
    socket = TSocket.TSocket(host, port)
    if options.unframed:
        transport = TTransport.TBufferedTransport(socket)
    else:
        transport = TTransport.TFramedTransport(socket)
    protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
    client = Cassandra.Client(protocol)
    client.transport = transport
    return client

def make_keyspaces():
    colms = []
    if options.index == 'keys':
        colms = [ColumnDef(name='C1', validation_class='UTF8Type', index_type=IndexType.KEYS)]
    elif options.index == 'keys_bitmap':
        colms = [ColumnDef(name='C1', validation_class='UTF8Type', index_type=IndexType.KEYS_BITMAP)]
    cfams = [CfDef(keyspace='Keyspace1', name='Standard1', column_metadata=colms),
             CfDef(keyspace='Keyspace1', name='Super1', column_type='Super')]
    keyspace = KsDef(name='Keyspace1',
                     strategy_class='org.apache.cassandra.locator.SimpleStrategy',
                     strategy_options={'replication_factor': str(options.replication)}, 
                     cf_defs=cfams)
    client = get_client(nodes[0], options.port)
    client.transport.open()
    try:
        client.system_add_keyspace(keyspace)
        print "Created keyspaces.  Sleeping %ss for propagation." % len(nodes)
        time.sleep(len(nodes))
    except InvalidRequestException, e:
        print e.why
    client.transport.close()

class Operation(Thread):
    def __init__(self, i, opcounts, keycounts, latencies):
        Thread.__init__(self)
        # generator of the keys to be used
        self.range = xrange(int(keys_per_thread * (i + options.skipkeys)), 
                            keys_per_thread * (i + 1))
        # we can't use a local counter, since that won't be visible to the parent
        # under multiprocessing.  instead, the parent passes a "opcounts" array
        # and an index that is our assigned counter.
        self.idx = i
        self.opcounts = opcounts
        # similarly, a shared array for latency and key totals
        self.latencies = latencies
        self.keycounts = keycounts
        # random host for pseudo-load-balancing
        [hostname] = random.sample(nodes, 1)
        # open client
        self.cclient = get_client(hostname, options.port)
        self.cclient.transport.open()
        self.cclient.set_keyspace('Keyspace1')

class Inserter(Operation):
    def run(self):
        values = generate_values()
        columns = [Column('C' + str(j), 'unset', time.time() * 1000000) for j in xrange(columns_per_key)]
        if 'super' == options.cftype:
            supers = [SuperColumn('S' + str(j), columns) for j in xrange(supers_per_key)]
        for i in self.range:
            key = fmt % i
            if 'super' == options.cftype:
                cfmap= {key: {'Super1' : [Mutation(ColumnOrSuperColumn(super_column=s)) for s in supers]}}
            else:
                cfmap = {key: {'Standard1': [Mutation(ColumnOrSuperColumn(column=c)) for c in columns]}}
            # set the correct column values for this row
            value = values[i % len(values)]
            for column in columns:
                column.value = value
            start = time.time()
            try:
                self.cclient.batch_mutate(cfmap, consistency)
            except KeyboardInterrupt:
                raise
            except Exception, e:
                if options.ignore:
                    print e
                else:
                    raise
            self.latencies[self.idx] += time.time() - start
            self.opcounts[self.idx] += 1
            self.keycounts[self.idx] += 1


class Reader(Operation):
    def run(self):
        p = SlicePredicate(slice_range=SliceRange('', '', False, columns_per_key))
        if 'super' == options.cftype:
            for i in xrange(keys_per_thread):
                key = key_generator()
                for j in xrange(supers_per_key):
                    parent = ColumnParent('Super1', 'S' + str(j))
                    start = time.time()
                    try:
                        r = self.cclient.get_slice(key, parent, p, consistency)
                        if not r: raise RuntimeError("Key %s not found" % key)
                    except KeyboardInterrupt:
                        raise
                    except Exception, e:
                        if options.ignore:
                            print e
                        else:
                            raise
                    self.latencies[self.idx] += time.time() - start
                    self.opcounts[self.idx] += 1
                    self.keycounts[self.idx] += 1
        else:
            parent = ColumnParent('Standard1')
            for i in xrange(keys_per_thread):
                key = key_generator()
                start = time.time()
                try:
                    r = self.cclient.get_slice(key, parent, p, consistency)
                    if not r: raise RuntimeError("Key %s not found" % key)
                except KeyboardInterrupt:
                    raise
                except Exception, e:
                    if options.ignore:
                        print e
                    else:
                        raise
                self.latencies[self.idx] += time.time() - start
                self.opcounts[self.idx] += 1
                self.keycounts[self.idx] += 1

class RangeSlicer(Operation):
    def run(self):
        begin = self.range[0]
        end = self.range[-1]
        current = begin
        last = current + options.rangecount
        p = SlicePredicate(slice_range=SliceRange('', '', False, columns_per_key))
        if 'super' == options.cftype:
            while current < end:
                keyrange = KeyRange(fmt % current, fmt % last, count = options.rangecount)
                res = []
                for j in xrange(supers_per_key):
                    parent = ColumnParent('Super1', 'S' + str(j)) 
                    begin = time.time()
                    try:
                        res = self.cclient.get_range_slices(parent, p, keyrange, consistency)
                        if not res: raise RuntimeError("Key %s not found" % key)
                    except KeyboardInterrupt:
                        raise
                    except Exception, e:
                        if options.ignore:
                            print e
                        else:
                            raise
                    self.latencies[self.idx] += time.time() - begin
                    self.opcounts[self.idx] += 1
                current += len(r) + 1
                last = current + len(r) + 1
                self.keycounts[self.idx] += len(r)
        else:
            parent = ColumnParent('Standard1')
            while current < end:
                start = fmt % current 
                finish = fmt % last
                keyrange = KeyRange(start, finish, count = options.rangecount)
                begin = time.time()
                try:
                    r = self.cclient.get_range_slices(parent, p, keyrange, consistency)
                    if not r: raise RuntimeError("Range not found:", start, finish)
                except KeyboardInterrupt:
                    raise
                except Exception, e:
                    if options.ignore:
                        print e
                    else:
                        print start, finish
                        raise
                current += len(r) + 1
                last = current + len(r)  + 1
                self.latencies[self.idx] += time.time() - begin
                self.opcounts[self.idx] += 1
                self.keycounts[self.idx] += len(r)

# Each thread queries for a portion of the unique values
# TODO: all threads start at the same key: implement wrapping, and start
# from the thread's appointed range
class IndexedRangeSlicer(Operation):
    def run(self):
        p = SlicePredicate(slice_range=SliceRange('', '', False, columns_per_key))
        values = generate_values()
        parent = ColumnParent('Standard1')
        # the number of rows with a particular value and the number of values we should query for
        expected_per_value = total_keys // len(values)
        valuebegin = self.range[0] // expected_per_value
        valuecount = len(self.range) // expected_per_value
        for valueidx in xrange(valuebegin, valuebegin + valuecount):
            received = 0
            start = fmt % 0
            value = values[valueidx % len(values)]
            expressions = [IndexExpression(column_name='C1', op=IndexOperator.EQ, value=value)]
            while received < expected_per_value:
                clause = IndexClause(start_key=start, count=options.rangecount, expressions=expressions)
                begin = time.time()
                try:
                    r = self.cclient.get_indexed_slices(parent, clause, p, consistency)
                    if not r: raise RuntimeError("No indexed values from offset received:", start)
                except KeyboardInterrupt:
                    raise
                except Exception, e:
                    if options.ignore:
                        print e
                        continue
                    else:
                        raise
                received += len(r)
                # convert max key found back to an integer, and increment it
                start = fmt % (1 + max([int(keyslice.key) for keyslice in r]))
                self.latencies[self.idx] += time.time() - begin
                self.opcounts[self.idx] += 1
                self.keycounts[self.idx] += len(r)


class MultiGetter(Operation):
    def run(self):
        p = SlicePredicate(slice_range=SliceRange('', '', False, columns_per_key))
        offset = self.idx * keys_per_thread
        count = (((self.idx+1) * keys_per_thread) - offset) / options.rangecount
        if 'super' == options.cftype:
            for x in xrange(count):
                keys = [key_generator() for i in xrange(offset, offset + options.rangecount)]
                for j in xrange(supers_per_key):
                    parent = ColumnParent('Super1', 'S' + str(j))
                    start = time.time()
                    try:
                        r = self.cclient.multiget_slice(keys, parent, p, consistency)
                        if not r: raise RuntimeError("Keys %s not found" % keys)
                    except KeyboardInterrupt:
                        raise
                    except Exception, e:
                        if options.ignore:
                            print e
                        else:
                            raise
                    self.latencies[self.idx] += time.time() - start
                    self.opcounts[self.idx] += 1
                    self.keycounts[self.idx] += len(keys)
                    offset += options.rangecount
        else:
            parent = ColumnParent('Standard1')
            for x in xrange(count):
                keys = [key_generator() for i in xrange(offset, offset + options.rangecount)]
                start = time.time()
                try:
                    r = self.cclient.multiget_slice(keys, parent, p, consistency)
                    if not r: raise RuntimeError("Keys %s not found" % keys)
                except KeyboardInterrupt:
                    raise
                except Exception, e:
                    if options.ignore:
                        print e
                    else:
                        raise
                self.latencies[self.idx] += time.time() - start
                self.opcounts[self.idx] += 1
                self.keycounts[self.idx] += len(keys)
                offset += options.rangecount


class OperationFactory:
    @staticmethod
    def create(type, i, opcounts, keycounts, latencies):
        if type == 'read':
            return Reader(i, opcounts, keycounts, latencies)
        elif type == 'insert':
            return Inserter(i, opcounts, keycounts, latencies)
        elif type == 'rangeslice':
            return RangeSlicer(i, opcounts, keycounts, latencies)
        elif type == 'indexedrangeslice':
            return IndexedRangeSlicer(i, opcounts, keycounts, latencies)
        elif type == 'multiget':
            return MultiGetter(i, opcounts, keycounts, latencies)
        else:
            raise RuntimeError, 'Unsupported op!'


class Stress(object):
    opcounts = array('i', [0] * n_threads)
    latencies = array('d', [0] * n_threads)
    keycounts = array('i', [0] * n_threads)

    def create_threads(self,type):
        threads = []
        for i in xrange(n_threads):
            th = OperationFactory.create(type, i, self.opcounts, self.keycounts, self.latencies)
            threads.append(th)
            th.start()
        return threads

    def run_test(self,filename,threads):
        start_t = time.time()
        if filename:
            outf = open(filename,'w')
        else:
            outf = sys.stdout
        outf.write('total,interval_op_rate,interval_key_rate,avg_latency,elapsed_time\n')
        epoch = total = old_total = latency = keycount = old_keycount = old_latency = 0
        epoch_intervals = (options.interval * 10) # 1 epoch = 1 tenth of a second
        terminate = False
        while not terminate:
            time.sleep(0.1)
            if not [th for th in threads if th.isAlive()]:
                terminate = True
            epoch = epoch + 1
            if terminate or epoch > epoch_intervals:
                epoch = 0
                old_total, old_latency, old_keycount = total, latency, keycount
                total = sum(self.opcounts[th.idx] for th in threads)
                latency = sum(self.latencies[th.idx] for th in threads)
                keycount = sum(self.keycounts[th.idx] for th in threads)
                opdelta = total - old_total
                keydelta = keycount - old_keycount
                delta_latency = latency - old_latency
                if opdelta > 0:
                    delta_formatted = (delta_latency / opdelta)
                else:
                    delta_formatted = 'NaN'
                elapsed_t = int(time.time() - start_t)
                outf.write('%d,%d,%d,%s,%d\n' 
                           % (total, opdelta / options.interval, keydelta / options.interval, delta_formatted, elapsed_t))

    def insert(self):
        threads = self.create_threads('insert')
        self.run_test(options.file,threads);

    def read(self):
        threads = self.create_threads('read')
        self.run_test(options.file,threads);
        
    def rangeslice(self):
        threads = self.create_threads('rangeslice')
        self.run_test(options.file,threads);

    def indexedrangeslice(self):
        threads = self.create_threads('indexedrangeslice')
        self.run_test(options.file,threads);

    def multiget(self):
        threads = self.create_threads('multiget')
        self.run_test(options.file,threads);

stresser = Stress()
benchmark = getattr(stresser, options.operation, None)
if not have_multiproc:
    print """WARNING: multiprocessing not present, threading will be used.
        Benchmark may not be accurate!"""
if options.operation == 'insert':
    make_keyspaces()
benchmark()