Python rabbit_helper.RabbitHelper类代码示例

def get_remote_phase_status(remoteIP, taskID, retry = 5):

    # assemble a request to remoteIP phase_status method
    rabbitHelper = RabbitHelper(mq_server = remoteIP)
    rcq = getResponseQueue(rabbitHelper)
    task_method = "app.systest_manager.get_phase_status"
    task_queue = "phase_status_"+cfg.CB_CLUSTER_TAG
    task_routing_queue = cfg.CB_CLUSTER_TAG+".phase.status"
    args = (taskID, rcq)

    # call phase_status task
    rawTaskPublisher(task_method,
                     args,
                     task_queue,
                     broker = remoteIP,
                     exchange="default",
                     routing_key = task_routing_queue)

    # retrieve status of phase
    rc = None
    while rc is None and retry > 0:
        rc = rabbitHelper.getMsg(rcq)
        time.sleep(2)
        retry = retry - 1

    rabbitHelper.delete(rcq)
    return rc == 'True'

def runRemotePhases(remote_phases, retry = 5):

    taskIds = []

    for remoteIP in remote_phases:

        # get handler to remote broker
        rabbitHelper = RabbitHelper(mq_server = remoteIP)
        rcq = getResponseQueue(rabbitHelper)
        args = (remote_phases[remoteIP], rcq)

        # call runPhase on remoteIP
        rawTaskPublisher("app.systest_manager.runPhase",
                         args,
                         "run_phase_"+cfg.CB_CLUSTER_TAG,
                         broker = remoteIP,
                         exchange = "default",
                         routing_key = cfg.CB_CLUSTER_TAG+".run.phase")

        # get taskID of phase running on remote broker
        taskId = None
        while taskId is None and retry > 0:
            time.sleep(2)
            taskId = rabbitHelper.getMsg(rcq)
            taskIds.append((remoteIP, taskId))
            retry = retry - 1

    return taskIds

def run_workload(args):
 
    workload = {}

    if args.name != None:
        # TODO: read in workload params from saved store
        # workload.update(cached_workload)
        pass

    if args.wait is not None:
        args.wait = conv_to_secs(args.wait) 

    if args.expires is not None:
        args.expires = conv_to_secs(args.expires)

    workload = { "bucket"      : args.bucket,
                 "ops_per_sec" : args.ops,
                 "create_perc" : args.create, 
                 "update_perc" : args.update, 
                 "get_perc"    : args.get, 
                 "del_perc"    : args.delete, 
                 "cc_queues"   : args.cc_queues,
                 "consume_queue" : args.consume_queue,
                 "postconditions" : args.postcondition,
                 "preconditions" : args.precondition,
                 "wait"  : args.wait,
                 "expires"  : args.expires,
                 "template"  : args.template}

    rabbitHelper = RabbitHelper(args.broker)
    rabbitHelper.putMsg("workload", json.dumps(workload))

def worker_init():
    # cleanup queues
    rabbitHelper = RabbitHelper()

    cached_queues = WorkloadCacher().queues +  TemplateCacher().cc_queues
    test_queues = ["workload","workload_template", "admin_tasks", "xdcr_tasks"] + cached_queues

    for queue in test_queues:
        try:
            if rabbitHelper.qsize(queue) > 0:
                print "Purge Queue: "+queue +" "+ str(rabbitHelper.qsize(queue))
                rabbitHelper.purge(queue)
        except Exception as ex:
            print ex

    cacheClean()

    # kill old background processes
    kill_procs=["sdkserver"]
    for proc in kill_procs:
        os.system("ps aux | grep %s | awk '{print $2}' | xargs kill" % proc)

    # start sdk servers
    os.system("ruby sdkserver.rb &")
    os.system("python sdkserver.py  &")

    # make sure logdir exists
    os.system("mkdir -p "+cfg.LOGDIR)

def run_workload(args):

    workload = {}

    if args.name != None:
        # TODO: read in workload params from saved store
        # workload.update(cached_workload)
        pass

    if args.wait is not None:
        args.wait = conv_to_secs(args.wait)

    workload = { "bucket"      : args.bucket,
                 "password"    : args.password,
                 "ops_per_sec" : args.ops,
                 "create_perc" : args.create,
                 "update_perc" : args.update,
                 "get_perc"    : args.get,
                 "del_perc"    : args.delete,
                 "exp_perc"    : args.expire,
                 "miss_perc"   : args.miss,
                 "ttl"         : args.ttl,
                 "cc_queues"   : args.cc_queues,
                 "consume_queue" : args.consume_queue,
                 "postconditions" : args.postcondition,
                 "preconditions" : args.precondition,
                 "wait"  : args.wait,
                 "template"  : args.template}
    cluster = args.cluster

    rabbitHelper = RabbitHelper(args.broker, cluster)
    workload['rcq'] = getResponseQueue(rabbitHelper)
    rabbitHelper.putMsg("workload_"+cluster, json.dumps(workload))
    receiveResponse(rabbitHelper, workload['rcq'])

def import_template(args):

    val = None

    if args.type == "json":
        json_val = {}
        for kv in args.kvpairs:
            pair = '{%s}' % kv 
            try:
                pair = json.loads(pair)
                json_val.update(pair)
            except ValueError as ex:
                print "ERROR: Unable to encode as valid json: %s " % kv
                print "make sure strings surrounded by double quotes"
                return
        val = json_val

    #TODO binary blobs

    template = { "name" : args.name,
                 "ttl" : args.ttl,
                 "flags" : args.flags,
                 "cc_queues" : args.cc_queues,
                 "size" : args.size,
                 "kv" : val}

    rabbitHelper = RabbitHelper(args.broker)
    rabbitHelper.putMsg("workload_template", json.dumps(template))

 def getKeyMapFromRemoteQueue(self, requeue = True):
     key_map = None
     mq = RabbitHelper()
     if mq.qsize(self.ccq) > 0:
         try:
             key_map = mq.getJsonMsg(self.ccq, requeue = requeue )
         except Exception:
             pass
     return key_map

    def requeueNonDeletedKeys(self):
        rabbitHelper = RabbitHelper()
        task_type = 'app.sdk_client_tasks.mdelete'

        # requeue pending delete keys so that they may be deleted in another workload
        while rabbitHelper.qsize(self.task_queue) > 0:
            task_set = rabbitHelper.getJsonMsg(self.task_queue)

            if len(task_set) > 0:
                keys = [task['args'] for task in task_set \
                            if task['task'] == task_type]
                if len(keys) > 0:
                    # put back on to consume_queue
                    msg = json.dumps(keys[0][0])
                    rabbitHelper.putMsg(self.consume_queue, msg)

        try:
            # delete task queue
            rabbitHelper.delete(self.task_queue)

            # delete consume queue if it was a miss_queue
            if self.miss_queue is not None and self.consume_queue is not None:
                rabbitHelper.delete(self.consume_queue)
        except:
            pass

def perform_query_tasks(args):
    queryMsg = {'queries_per_sec' : args.queries_per_sec,
                'ddoc' : args.ddoc,
                'view' : args.view,
                'bucket' : args.bucket,
                'password' : args.password}
    cluster = args.cluster

    rabbitHelper = RabbitHelper(args.broker)
    queryMsg['rcq'] = getResponseQueue(rabbitHelper)
    rabbitHelper.putMsg('query_'+cluster, json.dumps(queryMsg))
    receiveResponse(rabbitHelper, queryMsg['rcq'])

def perform_admin_tasks(args):

    actions = {'rebalance_in': args.rebalance_in,
               'rebalance_out': args.rebalance_out,
               'failover': args.failover,
               'soft_restart': args.soft_restart,
               'hard_restart': args.hard_restart,
               'only_failover': args.only_failover
              }

    #TODO: Validate the user inputs, before passing to rabbit
    print actions
    rabbitHelper = RabbitHelper(args.broker)
    rabbitHelper.putMsg("admin_tasks", json.dumps(actions))

def perform_xdcr_tasks(args):

    xdcrMsg = {'dest_cluster_ip': args.dest_cluster_ip,
               'dest_cluster_rest_username': args.dest_cluster_username,
               'dest_cluster_rest_pwd':  args.dest_cluster_pwd,
               'dest_cluster_name': args.dest_cluster_name,
               'replication_type': args.replication_type,
    }
    cluster = args.cluster

    #TODO: Validate the user inputs, before passing to rabbit
    print xdcrMsg
    rabbitHelper = RabbitHelper(args.broker, cluster)
    xdcrMsg['rcq'] = getResponseQueue(rabbitHelper)
    rabbitHelper.putMsg("xdcr_"+cluster, json.dumps(xdcrMsg))
    receiveResponse(rabbitHelper, xdcrMsg['rcq'])

def perform_admin_tasks(args):

    actions = {'rebalance_in': args.rebalance_in,
               'rebalance_out': args.rebalance_out,
               'failover': args.failover,
               'soft_restart': args.soft_restart,
               'hard_restart': args.hard_restart,
               'only_failover': args.only_failover
              }
    cluster = args.cluster

    #TODO: Validate the user inputs, before passing to rabbit
    rabbitHelper = RabbitHelper(args.broker, cluster)
    actions['rcq'] = getResponseQueue(rabbitHelper)
    rabbitHelper.putMsg("admin_"+cluster, json.dumps(actions))
    receiveResponse(rabbitHelper, actions['rcq'])

    def flushq(self, flush_hotkeys = False):

        mq = RabbitHelper()

        if self.ccq is not None:

            logging.info("[Thread %s] flushing %s items to %s" %
                         (self.name, self.memq.qsize(), self.ccq))

            # declare queue
            mq.declare(self.ccq)

            # empty the in memory queue
            while self.memq.empty() == False:
                try:
                    msg = self.memq.get_nowait()
                    msg = json.dumps(msg)
                    mq.putMsg(self.ccq, msg)
                except queue.Empty:
                    pass

        # hot keys
        if flush_hotkeys and (len(self.hotkey_batches) > 0):

            # try to put onto remote queue
            queue = self.consume_queue or self.ccq

            if queue is not None:
                key_map = {'start' : self.hotkey_batches[0][0],
                           'end' : self.hotkey_batches[-1][-1]}
                msg = json.dumps(key_map)
                mq.putMsg(queue, msg)
                self.hotkey_batches = []

    def getKeyMapFromRemoteQueue(self, requeue = True):

        key_map = None
        mq = RabbitHelper()

        # try to fetch from consume queue and
        # fall back to ccqueue
        queue = self.consume_queue

        if queue is None  or mq.qsize(queue) == 0:
            queue = self.ccq

        if mq.qsize(queue) > 0:
            try:
                key_map = mq.getJsonMsg(queue, requeue = requeue )
            except Exception:
                pass

        return key_map

def run_systemtest(args):

    cluster = args.cluster
    rabbitHelper = RabbitHelper(args.broker, cluster)

    test = {'suffix' : args.filesuffix}

    if args.fromfile is not None:

        # load json config
        json_data = open(args.fromfile)
        msg = json.load(json_data)

    elif args.name is not None:
        msg = { "localtestname" : args.name }

    test.update(msg)
    test['rcq'] = getResponseQueue(rabbitHelper)
    rabbitHelper.putMsg('systest_manager_'+cluster, json.dumps(test))
    receiveResponse(rabbitHelper, test['rcq'])