class ChildProcess(Process):
    def __init__(self, *args, **kwargs):
        super(ChildProcess, self).__init__(*args, **kwargs)
        self.daemon = True

    def run(self):
        # Make sure child processes don't catch ^C or eat input.
        os.close(0)
        signal(SIGINT, SIG_IGN)

        self.run_child_process() def debug(tmpl, *args):
    if _DEBUG:
        print '[DEBUG] ' + (tmpl % args) def traced_func(f):
    name = f.__name__

    def wrapped(*a, **kw):
        return trace_call(name, f, *a, **kw)

    return wrapped def traced_method(m):
    def wrapped(self, *a, **kw):
        name = self.__class__.__name__ + '.' + m.__name__
        return trace_call(name, m, self, *a, **kw)
    return wrapped class Tally(IntegerStatistic):
    def __init__(self):
        super(IntegerStatistic,self).__init__()

        self._grand_total = 0
        self.reset()

    def add(self):
        self._total += 1

    def calculate(self):
        self._grand_total += self._total

        stats = OrderedDict()

        stats["Current"] = self._total
        stats["Total"] = self._grand_total
        stats["Rate"] = self._total / (time.time() - self._start_time)

        return stats

    def reset(self):
        self._total = 0
        self._start_time = time.time() class Level(IntegerStatistic):
    def __init__(self):
        super(Level,self).__init__()

        self._level = 0
        self._levels = [0]

    def add(self, direction):
        if direction == "+":
            self._level += 1
        elif direction == "-":
            self._level -= 1

        self._levels.append(self._level)

    def reset(self):
        self._levels = []

    def calculate(self):
        stats = OrderedDict()

        stats['Current'] = self._level

        if self._levels:
            add_series_stats(stats, self._levels)

        return stats class Series(FloatStatistic):
    def __init__(self):
        super(Series, self).__init__()

        self._values = []

    def add(self, value):
        self._values.append(value)

    def reset(self):
        self._values = []

    def calculate(self):
        stats = OrderedDict()

        if self._values:
            stats['Current'] = self._values[-1]
            add_series_stats(stats, self._values)

        return stats def format_table(rows):
    if not rows:
        return

    num_cols = max(len(row) for row in rows)

    # make all rows as long as the longest
    rows = [row + [(ALIGN_LEFT, "")] * (num_cols - len(row)) for row in rows]

    # calculate max column widths
    widths = [max(len(r[i][1]) for r in rows) for i in xrange(num_cols)]

    output = ""
    for row in rows:
        for i, (align, text) in enumerate(row):
            if i > 0:
                output += " "

            output += "{0:{1}{2}}".format(text, align, widths[i])
        output += "\n"

    return output ALIGN_LEFT = "<" ALIGN_CENTER = "^" ALIGN_RIGHT = ">"
import optparse
import os
import re
import random
import socket
import sys
import tempfile
import cPickle
import MySQLdb
import time
import warnings
from datetime import datetime
from threading import Thread
from multiprocessing import Value, Queue
from multiprocessing.queues import Empty
from collections import defaultdict
from itertools import chain
from apiary.tools.childprocess import ChildProcess
from apiary.tools.debug import debug, traced_func, traced_method
from apiary.tools.stats import Tally, Level, Series
from apiary.tools.table import format_table, ALIGN_LEFT, ALIGN_CENTER, ALIGN_RIGHT
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# $/LicenseInfo$
#


'''
This module contains the main Apiary code.  The BeeKeeper spawns and manages
all subprocesses.  The QueenBee reads jobs from disk and enqueues them for
workers.  WorkerBee is a class to be extended by protocol-specific classes.  It
contains the basic workings of a worker client to make it simple to implement a
new protocol.  The StatsGatherer receives, aggregates, and prints status
messages from the WorkerBees.
'''



verbose = False


class BeeKeeper(object):
    """Manages the hive, including QueenBee, WorkerBees, and StatsGatherer."""

    def __init__(self, options, arguments):
        self.options = options
        self.arguments = arguments

        try:
            self.protocol = options.protocols[options.protocol]
        except KeyError:
            sys.exit('invalid protocol: %s (valid protocols: %s)' %
                     (options.protocol, " ".join(options.protocols)))

    def start(self):
        """Run the load test."""

        start_time = time.time()

        job_queue = Queue()
        stats_queue = Queue()

        workers = []

        delay = self.options.stagger_workers / 1000.0
        for i in xrange(self.options.workers):
            worker = WorkerBeeProcess(self.options, self.protocol, job_queue, stats_queue)
            worker.start()
            workers.append(worker)
            time.sleep(delay)

        # TODO: consider waiting until workers are ready
        #if self.options.startup_wait:
        #    print "workers started; waiting %d seconds..." % self.options.startup_wait
        #    time.sleep(self.options.startup_wait)

        stats_gatherer = StatsGatherer(self.options, stats_queue)
        stats_gatherer.start()

        queen = QueenBee(self.options, self.arguments, job_queue, stats_queue)
        queen.start()

        # Now wait while the queen does its thing.
        try:
            queen.join()
        except KeyboardInterrupt:
            print "Interrupted, shutting down..."
            queen.terminate()

        print "Waiting for workers to complete jobs and terminate (may take up to %d seconds)..." % self.options.max_ahead

        try:
            stop = Message(Message.STOP)
            for worker in xrange(self.options.workers * self.options.threads):
                job_queue.put(stop)

            # Now wait for the workers to get the message.  This may take a few
            # minutes as the QueenBee likes to stay ahead by a bit.

            for worker in workers:
                worker.join()

            # Tell the Stats Gatherer that it's done.
            stats_queue.put(Message(Message.STOP))

            # Wait for it to finish.
            stats_gatherer.join()

            print "Completed %d jobs in %0.2f seconds." % (queen.jobs_sent.value, time.time() - start_time)
        except KeyboardInterrupt:
            print "Interrupted before shutdown process completed."

            job_queue.cancel_join_thread()
            stats_queue.cancel_join_thread()


class StatsGatherer(ChildProcess):
    """Gather and present stats.

    The StatsGatherer process reads messages off of the stats queue sent by the
    workers and aggregates them.  It prints a report periodically.

    See tools.stats for a description of the kinds of statistics that are
    available.
    """

    def __init__(self, options, stats_queue):
        super(StatsGatherer, self).__init__()

        self._options = options
        self._verbose = options.verbose
        self._tallies = defaultdict(Tally)
        self._levels = defaultdict(Level)
        self._series = defaultdict(Series)
        self._last_report = time.time()
        self._worker_count = 0
        self._queue = stats_queue


    @traced_method