Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Fetching contributors…

Cannot retrieve contributors at this time

293 lines (244 sloc) 12.547 kb
import struct
from gearman.constants import PRIORITY_NONE, PRIORITY_LOW, PRIORITY_HIGH
from gearman.errors import ProtocolError
from gearman import compat
# Protocol specific constants
NULL_CHAR = '\x00'
MAGIC_RES_STRING = '%sRES' % NULL_CHAR
MAGIC_REQ_STRING = '%sREQ' % NULL_CHAR
COMMAND_HEADER_SIZE = 12
# Gearman commands 1-9
GEARMAN_COMMAND_CAN_DO = 1
GEARMAN_COMMAND_CANT_DO = 2
GEARMAN_COMMAND_RESET_ABILITIES = 3
GEARMAN_COMMAND_PRE_SLEEP = 4
GEARMAN_COMMAND_NOOP = 6
GEARMAN_COMMAND_SUBMIT_JOB = 7
GEARMAN_COMMAND_JOB_CREATED = 8
GEARMAN_COMMAND_GRAB_JOB = 9
# Gearman commands 10-19
GEARMAN_COMMAND_NO_JOB = 10
GEARMAN_COMMAND_JOB_ASSIGN = 11
GEARMAN_COMMAND_WORK_STATUS = 12
GEARMAN_COMMAND_WORK_COMPLETE = 13
GEARMAN_COMMAND_WORK_FAIL = 14
GEARMAN_COMMAND_GET_STATUS = 15
GEARMAN_COMMAND_ECHO_REQ = 16
GEARMAN_COMMAND_ECHO_RES = 17
GEARMAN_COMMAND_SUBMIT_JOB_BG = 18
GEARMAN_COMMAND_ERROR = 19
# Gearman commands 20-29
GEARMAN_COMMAND_STATUS_RES = 20
GEARMAN_COMMAND_SUBMIT_JOB_HIGH = 21
GEARMAN_COMMAND_SET_CLIENT_ID = 22
GEARMAN_COMMAND_CAN_DO_TIMEOUT = 23
GEARMAN_COMMAND_ALL_YOURS = 24
GEARMAN_COMMAND_WORK_EXCEPTION = 25
GEARMAN_COMMAND_OPTION_REQ = 26
GEARMAN_COMMAND_OPTION_RES = 27
GEARMAN_COMMAND_WORK_DATA = 28
GEARMAN_COMMAND_WORK_WARNING = 29
# Gearman commands 30-39
GEARMAN_COMMAND_GRAB_JOB_UNIQ = 30
GEARMAN_COMMAND_JOB_ASSIGN_UNIQ = 31
GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG = 32
GEARMAN_COMMAND_SUBMIT_JOB_LOW = 33
GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG = 34
GEARMAN_COMMAND_SUBMIT_JOB_EPOCH = 36
# Fake command code
GEARMAN_COMMAND_TEXT_COMMAND = 9999
GEARMAN_PARAMS_FOR_COMMAND = {
# Gearman commands 1-9
GEARMAN_COMMAND_CAN_DO: ['task'],
GEARMAN_COMMAND_CANT_DO: ['task'],
GEARMAN_COMMAND_RESET_ABILITIES: [],
GEARMAN_COMMAND_PRE_SLEEP: [],
GEARMAN_COMMAND_NOOP: [],
GEARMAN_COMMAND_SUBMIT_JOB: ['task', 'unique', 'data'],
GEARMAN_COMMAND_JOB_CREATED: ['job_handle'],
GEARMAN_COMMAND_GRAB_JOB: [],
# Gearman commands 10-19
GEARMAN_COMMAND_NO_JOB: [],
GEARMAN_COMMAND_JOB_ASSIGN: ['job_handle', 'task', 'data'],
GEARMAN_COMMAND_WORK_STATUS: ['job_handle', 'numerator', 'denominator'],
GEARMAN_COMMAND_WORK_COMPLETE: ['job_handle', 'data'],
GEARMAN_COMMAND_WORK_FAIL: ['job_handle'],
GEARMAN_COMMAND_GET_STATUS: ['job_handle'],
GEARMAN_COMMAND_ECHO_REQ: ['data'],
GEARMAN_COMMAND_ECHO_RES: ['data'],
GEARMAN_COMMAND_SUBMIT_JOB_BG: ['task', 'unique', 'data'],
GEARMAN_COMMAND_ERROR: ['error_code', 'error_text'],
# Gearman commands 20-29
GEARMAN_COMMAND_STATUS_RES: ['job_handle', 'known', 'running', 'numerator', 'denominator'],
GEARMAN_COMMAND_SUBMIT_JOB_HIGH: ['task', 'unique', 'data'],
GEARMAN_COMMAND_SET_CLIENT_ID: ['client_id'],
GEARMAN_COMMAND_CAN_DO_TIMEOUT: ['task', 'timeout'],
GEARMAN_COMMAND_ALL_YOURS: [],
GEARMAN_COMMAND_WORK_EXCEPTION: ['job_handle', 'data'],
GEARMAN_COMMAND_OPTION_REQ: ['option_name'],
GEARMAN_COMMAND_OPTION_RES: ['option_name'],
GEARMAN_COMMAND_WORK_DATA: ['job_handle', 'data'],
GEARMAN_COMMAND_WORK_WARNING: ['job_handle', 'data'],
# Gearman commands 30-39
GEARMAN_COMMAND_GRAB_JOB_UNIQ: [],
GEARMAN_COMMAND_JOB_ASSIGN_UNIQ: ['job_handle', 'task', 'unique', 'data'],
GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG: ['task', 'unique', 'data'],
GEARMAN_COMMAND_SUBMIT_JOB_LOW: ['task', 'unique', 'data'],
GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG: ['task', 'unique', 'data'],
GEARMAN_COMMAND_SUBMIT_JOB_EPOCH: ['task', 'unique', 'when_to_run', 'data'],
# Fake gearman command
GEARMAN_COMMAND_TEXT_COMMAND: ['raw_text']
}
GEARMAN_COMMAND_TO_NAME = {
GEARMAN_COMMAND_CAN_DO: 'GEARMAN_COMMAND_CAN_DO',
GEARMAN_COMMAND_CANT_DO: 'GEARMAN_COMMAND_CANT_DO',
GEARMAN_COMMAND_RESET_ABILITIES: 'GEARMAN_COMMAND_RESET_ABILITIES',
GEARMAN_COMMAND_PRE_SLEEP: 'GEARMAN_COMMAND_PRE_SLEEP',
GEARMAN_COMMAND_NOOP: 'GEARMAN_COMMAND_NOOP',
GEARMAN_COMMAND_SUBMIT_JOB: 'GEARMAN_COMMAND_SUBMIT_JOB',
GEARMAN_COMMAND_JOB_CREATED: 'GEARMAN_COMMAND_JOB_CREATED',
GEARMAN_COMMAND_GRAB_JOB: 'GEARMAN_COMMAND_GRAB_JOB',
# Gearman commands 10-19
GEARMAN_COMMAND_NO_JOB: 'GEARMAN_COMMAND_NO_JOB',
GEARMAN_COMMAND_JOB_ASSIGN: 'GEARMAN_COMMAND_JOB_ASSIGN',
GEARMAN_COMMAND_WORK_STATUS: 'GEARMAN_COMMAND_WORK_STATUS',
GEARMAN_COMMAND_WORK_COMPLETE: 'GEARMAN_COMMAND_WORK_COMPLETE',
GEARMAN_COMMAND_WORK_FAIL: 'GEARMAN_COMMAND_WORK_FAIL',
GEARMAN_COMMAND_GET_STATUS: 'GEARMAN_COMMAND_GET_STATUS',
GEARMAN_COMMAND_ECHO_REQ: 'GEARMAN_COMMAND_ECHO_REQ',
GEARMAN_COMMAND_ECHO_RES: 'GEARMAN_COMMAND_ECHO_RES',
GEARMAN_COMMAND_SUBMIT_JOB_BG: 'GEARMAN_COMMAND_SUBMIT_JOB_BG',
GEARMAN_COMMAND_ERROR: 'GEARMAN_COMMAND_ERROR',
# Gearman commands 20-29
GEARMAN_COMMAND_STATUS_RES: 'GEARMAN_COMMAND_STATUS_RES',
GEARMAN_COMMAND_SUBMIT_JOB_HIGH: 'GEARMAN_COMMAND_SUBMIT_JOB_HIGH',
GEARMAN_COMMAND_SET_CLIENT_ID: 'GEARMAN_COMMAND_SET_CLIENT_ID',
GEARMAN_COMMAND_CAN_DO_TIMEOUT: 'GEARMAN_COMMAND_CAN_DO_TIMEOUT',
GEARMAN_COMMAND_ALL_YOURS: 'GEARMAN_COMMAND_ALL_YOURS',
GEARMAN_COMMAND_WORK_EXCEPTION: 'GEARMAN_COMMAND_WORK_EXCEPTION',
GEARMAN_COMMAND_OPTION_REQ: 'GEARMAN_COMMAND_OPTION_REQ',
GEARMAN_COMMAND_OPTION_RES: 'GEARMAN_COMMAND_OPTION_RES',
GEARMAN_COMMAND_WORK_DATA: 'GEARMAN_COMMAND_WORK_DATA',
GEARMAN_COMMAND_WORK_WARNING: 'GEARMAN_COMMAND_WORK_WARNING',
# Gearman commands 30-39
GEARMAN_COMMAND_GRAB_JOB_UNIQ: 'GEARMAN_COMMAND_GRAB_JOB_UNIQ',
GEARMAN_COMMAND_JOB_ASSIGN_UNIQ: 'GEARMAN_COMMAND_JOB_ASSIGN_UNIQ',
GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG: 'GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG',
GEARMAN_COMMAND_SUBMIT_JOB_LOW: 'GEARMAN_COMMAND_SUBMIT_JOB_LOW',
GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG: 'GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG',
GEARMAN_COMMAND_SUBMIT_JOB_EPOCH: 'GEARMAN_COMMAND_SUBMIT_JOB_EPOCH',
GEARMAN_COMMAND_TEXT_COMMAND: 'GEARMAN_COMMAND_TEXT_COMMAND'
}
GEARMAN_SERVER_COMMAND_STATUS = 'status'
GEARMAN_SERVER_COMMAND_VERSION = 'version'
GEARMAN_SERVER_COMMAND_WORKERS = 'workers'
GEARMAN_SERVER_COMMAND_MAXQUEUE = 'maxqueue'
GEARMAN_SERVER_COMMAND_SHUTDOWN = 'shutdown'
def get_command_name(cmd_type):
return GEARMAN_COMMAND_TO_NAME.get(cmd_type, cmd_type)
def submit_cmd_for_background_priority_run_later(background, priority, run_later):
cmd_type_lookup = {
(True, PRIORITY_NONE, False): GEARMAN_COMMAND_SUBMIT_JOB_BG,
(True, PRIORITY_LOW, False): GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG,
(True, PRIORITY_HIGH, False): GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG,
(False, PRIORITY_NONE, False): GEARMAN_COMMAND_SUBMIT_JOB,
(False, PRIORITY_LOW, False): GEARMAN_COMMAND_SUBMIT_JOB_LOW,
(False, PRIORITY_HIGH, False): GEARMAN_COMMAND_SUBMIT_JOB_HIGH,
(True, PRIORITY_NONE, True): GEARMAN_COMMAND_SUBMIT_JOB_EPOCH
}
lookup_tuple = (background, priority, run_later)
cmd_type = cmd_type_lookup[lookup_tuple]
return cmd_type
def parse_binary_command(in_buffer, is_response=True):
"""Parse data and return (command type, command arguments dict, command size)
or (None, None, data) if there's not enough data for a complete command.
"""
in_buffer_size = len(in_buffer)
magic = None
cmd_type = None
cmd_args = None
cmd_len = 0
expected_packet_size = None
# If we don't have enough data to parse, error early
if in_buffer_size < COMMAND_HEADER_SIZE:
return cmd_type, cmd_args, cmd_len
# By default, we'll assume we're dealing with a gearman command
magic, cmd_type, cmd_len = struct.unpack('!4sII', in_buffer[:COMMAND_HEADER_SIZE])
received_bad_response = is_response and bool(magic != MAGIC_RES_STRING)
received_bad_request = not is_response and bool(magic != MAGIC_REQ_STRING)
if received_bad_response or received_bad_request:
raise ProtocolError('Malformed Magic')
expected_cmd_params = GEARMAN_PARAMS_FOR_COMMAND.get(cmd_type, None)
# GEARMAN_COMMAND_TEXT_COMMAND is a faked command that we use to support server text-based commands
if expected_cmd_params is None or cmd_type == GEARMAN_COMMAND_TEXT_COMMAND:
raise ProtocolError('Received unknown binary command: %s' % cmd_type)
# If everything indicates this is a valid command, we should check to see if we have enough stuff to read in our buffer
expected_packet_size = COMMAND_HEADER_SIZE + cmd_len
if in_buffer_size < expected_packet_size:
return None, None, 0
binary_payload = in_buffer[COMMAND_HEADER_SIZE:expected_packet_size]
split_arguments = []
if len(expected_cmd_params) > 0:
binary_payload = binary_payload.tostring()
split_arguments = binary_payload.split(NULL_CHAR, len(expected_cmd_params) - 1)
elif binary_payload:
raise ProtocolError('Expected no binary payload: %s' % get_command_name(cmd_type))
# This is a sanity check on the binary_payload.split() phase
# We should never be able to get here with any VALID gearman data
if len(split_arguments) != len(expected_cmd_params):
raise ProtocolError('Received %d argument(s), expecting %d argument(s): %s' % (len(split_arguments), len(expected_cmd_params), get_command_name(cmd_type)))
# Iterate through the split arguments and assign them labels based on their order
cmd_args = dict((param_label, param_value) for param_label, param_value in zip(expected_cmd_params, split_arguments))
return cmd_type, cmd_args, expected_packet_size
def pack_binary_command(cmd_type, cmd_args, is_response=False):
"""Packs the given command using the parameter ordering specified in GEARMAN_PARAMS_FOR_COMMAND.
*NOTE* Expects that all arguments in cmd_args are already str's.
"""
expected_cmd_params = GEARMAN_PARAMS_FOR_COMMAND.get(cmd_type, None)
if expected_cmd_params is None or cmd_type == GEARMAN_COMMAND_TEXT_COMMAND:
raise ProtocolError('Received unknown binary command: %s' % get_command_name(cmd_type))
expected_parameter_set = set(expected_cmd_params)
received_parameter_set = set(cmd_args.keys())
if expected_parameter_set != received_parameter_set:
raise ProtocolError('Received arguments did not match expected arguments: %r != %r' % (expected_parameter_set, received_parameter_set))
# Select the right expected magic
if is_response:
magic = MAGIC_RES_STRING
else:
magic = MAGIC_REQ_STRING
# !NOTE! str should be replaced with bytes in Python 3.x
# We will iterate in ORDER and str all our command arguments
if compat.any(type(param_value) != str for param_value in cmd_args.itervalues()):
raise ProtocolError('Received non-binary arguments: %r' % cmd_args)
data_items = [cmd_args[param] for param in expected_cmd_params]
# Now check that all but the last argument are free of \0 as per the protocol spec.
if compat.any('\0' in argument for argument in data_items[:-1]):
raise ProtocolError('Received arguments with NULL byte in non-final argument')
binary_payload = NULL_CHAR.join(data_items)
# Pack the header in the !4sII format then append the binary payload
payload_size = len(binary_payload)
packing_format = '!4sII%ds' % payload_size
return struct.pack(packing_format, magic, cmd_type, payload_size, binary_payload)
def parse_text_command(in_buffer):
"""Parse a text command and return a single line at a time"""
cmd_type = None
cmd_args = None
cmd_len = 0
if '\n' not in in_buffer:
return cmd_type, cmd_args, cmd_len
text_command, in_buffer = in_buffer.tostring().split('\n', 1)
if NULL_CHAR in text_command:
raise ProtocolError('Received unexpected character: %s' % text_command)
# Fake gearman command "TEXT_COMMAND" used to process server admin client responses
cmd_type = GEARMAN_COMMAND_TEXT_COMMAND
cmd_args = dict(raw_text=text_command)
cmd_len = len(text_command) + 1
return cmd_type, cmd_args, cmd_len
def pack_text_command(cmd_type, cmd_args):
"""Parse a text command and return a single line at a time"""
if cmd_type != GEARMAN_COMMAND_TEXT_COMMAND:
raise ProtocolError('Unknown cmd_type: Received %s, expecting %s' % (get_command_name(cmd_type), get_command_name(GEARMAN_COMMAND_TEXT_COMMAND)))
cmd_line = cmd_args.get('raw_text')
if cmd_line is None:
raise ProtocolError('Did not receive arguments any valid arguments: %s' % cmd_args)
return str(cmd_line)
Jump to Line
Something went wrong with that request. Please try again.