worker_handler.py

import logging

from gearman.command_handler import GearmanCommandHandler
from gearman.errors import InvalidWorkerState
from gearman.protocol import GEARMAN_COMMAND_PRE_SLEEP, GEARMAN_COMMAND_RESET_ABILITIES, GEARMAN_COMMAND_CAN_DO, GEARMAN_COMMAND_SET_CLIENT_ID, GEARMAN_COMMAND_GRAB_JOB_UNIQ, \
    GEARMAN_COMMAND_WORK_STATUS, GEARMAN_COMMAND_WORK_COMPLETE, GEARMAN_COMMAND_WORK_FAIL, GEARMAN_COMMAND_WORK_EXCEPTION, GEARMAN_COMMAND_WORK_WARNING, GEARMAN_COMMAND_WORK_DATA

gearman_logger = logging.getLogger(__name__)

class GearmanWorkerCommandHandler(GearmanCommandHandler):
    """GearmanWorker state machine on a per connection basis

    A worker can be in the following distinct states:
        SLEEP         -> Doing nothing, can be awoken
        AWAKE         -> Transitional state (for NOOP)
        AWAITING_JOB  -> Holding worker level job lock and awaiting a server response
        EXECUTING_JOB -> Transitional state (for ASSIGN_JOB)
    """
    def __init__(self, connection_manager=None):
        super(GearmanWorkerCommandHandler, self).__init__(connection_manager=connection_manager)

        self._handler_abilities = []
        self._client_id = None

    def initial_state(self, abilities=None, client_id=None):
        self.set_client_id(client_id)
        self.set_abilities(abilities)

        self._sleep()

    ##################################################################
    ##### Public interface methods to be called by GearmanWorker #####
    ##################################################################
    def set_abilities(self, connection_abilities_list):
        assert type(connection_abilities_list) in (list, tuple)
        self._handler_abilities = connection_abilities_list

        self.send_command(GEARMAN_COMMAND_RESET_ABILITIES)
        for task in self._handler_abilities:
            self.send_command(GEARMAN_COMMAND_CAN_DO, task=task)

    def set_client_id(self, client_id):
        self._client_id = client_id

        if self._client_id is not None:
            self.send_command(GEARMAN_COMMAND_SET_CLIENT_ID, client_id=self._client_id)

    ###############################################################
    #### Convenience methods for typical gearman jobs to call #####
    ###############################################################
    def send_job_status(self, current_job, numerator, denominator):
        assert type(numerator) in (int, float), 'Numerator must be a numeric value'
        assert type(denominator) in (int, float), 'Denominator must be a numeric value'
        self.send_command(GEARMAN_COMMAND_WORK_STATUS, job_handle=current_job.handle, numerator=str(numerator), denominator=str(denominator))

    def send_job_complete(self, current_job, data):
        """Removes a job from the queue if its backgrounded"""
        self.send_command(GEARMAN_COMMAND_WORK_COMPLETE, job_handle=current_job.handle, data=self.encode_data(data))

    def send_job_failure(self, current_job):
        """Removes a job from the queue if its backgrounded"""
        self.send_command(GEARMAN_COMMAND_WORK_FAIL, job_handle=current_job.handle)

    def send_job_exception(self, current_job, data):
        # Using GEARMAND_COMMAND_WORK_EXCEPTION is not recommended at time of this writing [2010-02-24]
        # http://groups.google.com/group/gearman/browse_thread/thread/5c91acc31bd10688/529e586405ed37fe
        #
        self.send_command(GEARMAN_COMMAND_WORK_EXCEPTION, job_handle=current_job.handle, data=self.encode_data(data))

    def send_job_data(self, current_job, data):
        self.send_command(GEARMAN_COMMAND_WORK_DATA, job_handle=current_job.handle, data=self.encode_data(data))

    def send_job_warning(self, current_job, data):
        self.send_command(GEARMAN_COMMAND_WORK_WARNING, job_handle=current_job.handle, data=self.encode_data(data))

    ###########################################################
    ### Callbacks when we receive a command from the server ###
    ###########################################################
    def _grab_job(self):
        self.send_command(GEARMAN_COMMAND_GRAB_JOB_UNIQ)

    def _sleep(self):
        self.send_command(GEARMAN_COMMAND_PRE_SLEEP)

    def _check_job_lock(self):
        return self.connection_manager.check_job_lock(self)

    def _acquire_job_lock(self):
        return self.connection_manager.set_job_lock(self, lock=True)

    def _release_job_lock(self):
        if not self.connection_manager.set_job_lock(self, lock=False):
            raise InvalidWorkerState("Unable to release job lock for %r" % self)

        return True

    def recv_noop(self):
        """Transition from being SLEEP --> AWAITING_JOB / SLEEP

          AWAITING_JOB -> AWAITING_JOB :: Noop transition, we're already awaiting a job
        SLEEP -> AWAKE -> AWAITING_JOB :: Transition if we can acquire the worker job lock
        SLEEP -> AWAKE -> SLEEP        :: Transition if we can NOT acquire a worker job lock
        """
        if self._check_job_lock():
            pass
        elif self._acquire_job_lock():
            self._grab_job()
        else:
            self._sleep()

        return True

    def recv_no_job(self):
        """Transition from being AWAITING_JOB --> SLEEP

        AWAITING_JOB -> SLEEP :: Always transition to sleep if we have nothing to do
        """
        self._release_job_lock()
        self._sleep()

        return True

    def recv_job_assign_uniq(self, job_handle, task, unique, data):
        """Transition from being AWAITING_JOB --> EXECUTE_JOB --> SLEEP

        AWAITING_JOB -> EXECUTE_JOB -> SLEEP :: Always transition once we're given a job
        """
        assert task in self._handler_abilities, '%s not found in %r' % (task, self._handler_abilities)

        # After this point, we know this connection handler is holding onto the job lock so we don't need to acquire it again
        if not self.connection_manager.check_job_lock(self):
            raise InvalidWorkerState("Received a job when we weren't expecting one")

        gearman_job = self.connection_manager.create_job(self, job_handle, task, unique, self.decode_data(data))

        # Create a new job
        self.connection_manager.on_job_execute(gearman_job)

        # Release the job lock once we're doing and go back to sleep
        self._release_job_lock()
        self._sleep()

        return True

    def recv_job_assign(self, job_handle, task, data):
        """JOB_ASSIGN and JOB_ASSIGN_UNIQ are essentially the same"""
        return self.recv_job_assign_uniq(job_handle=job_handle, task=task, unique=None, data=data)