def __init__(self,
host='localhost',
port=11300,
socket_timeout=None,
socket_connect_timeout=None,
socket_keepalive=None,
retry_on_timeout=False,
socket_keepalive_options=None,
max_connections=None,
connection=None):
if not connection:
self.socket_timeout = socket_timeout
kwargs = {
'host': host,
'port': port,
'socket_connect_timeout': socket_connect_timeout,
'socket_keepalive': socket_keepalive,
'socket_keepalive_options': socket_keepalive_options,
'socket_timeout': socket_timeout,
'retry_on_timeout': retry_on_timeout,
'max_connections': max_connections
}
connection = Connection(**kwargs)
self.conn_queue = LifoQueue()
self.conn_queue.put_nowait(connection)
self._connection = connection
self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy()
self.expected_ok = self.__class__.EXPECTED_OK.copy()
self.expected_err = self.__class__.EXPECTED_ERR.copy()
class BarkerConsumer(clusto.script_helper.Script):
def __init__(self):
clusto.script_helper.Script.__init__(self)
self.queue = None
def callback(self, body, message):
if self.queue.qsize() > 500:
log.warning('Dropping message, queue size is over 500')
return
self.queue.put(body)
def run(self, args):
self.queue = Queue()
for hostname in QUEUE_HOSTS:
eventlet.spawn_n(self.consumer, hostname)
while True:
body = self.queue.get()
log.debug('Queue size %s' % self.queue.qsize())
barker_callback(body)
def consumer(self, hostname):
exchange = kombu.Exchange(QUEUE_EXCHANGE, type='fanout',
delivery_mode='transient')
queue = kombu.Queue(QUEUE_NAME, exchange)
try:
connection = kombu.BrokerConnection(
hostname=hostname,
userid=QUEUE_USER,
password=QUEUE_PASSWORD,
virtual_host=QUEUE_VHOST)
channel = connection.channel()
consumer = kombu.Consumer(channel, queue,
callbacks=[self.callback],
no_ack=True)
consumer.consume()
log.info('%s consumer running' % hostname)
while True:
try:
connection.drain_events()
except Exception, e:
log.error(str(e))
except Exception, e:
log.error(format_exc())
raise e
finally:
def run(self, args):
self.queue = Queue()
for hostname in QUEUE_HOSTS:
eventlet.spawn_n(self.consumer, hostname)
while True:
body = self.queue.get()
log.debug('Queue size %s' % self.queue.qsize())
barker_callback(body)
class Beanstalk(object):
RESPONSE_CALLBACKS = dict_merge({
'reserve': parse_body,
'stats-tube': parse_yaml
})
EXPECTED_OK = dict_merge({
'reserve': ['RESERVED'],
'delete': ['DELETED'],
'release': ['RELEASED'],
'bury': ['BURIED'],
'put': ['INSERTED'],
'use': ['USING'],
'watch': ['WATCHING'],
'stats-tube': ['OK'],
'kick': ['KICKED'],
'kick-job': ['KICKED']
})
EXPECTED_ERR = dict_merge({
'reserve': ['DEADLINE_SOON', 'TIMED_OUT'],
'delete': ['NOT_FOUND'],
'release': ['BURIED', 'NOT_FOUND', 'OUT_OF_MEMORY'],
'bury': ['NOT_FOUND', 'OUT_OF_MEMORY'],
'stats-tube': ['NOT_FOUND'],
'use': [],
'watch': [],
'put': ['JOB_TOO_BIG', 'BURIED', 'DRAINING', 'OUT_OF_MEMORY'],
'kick': ['OUT_OF_MEMORY'],
'kick-job': ['NOT_FOUND', 'OUT_OF_MEMORY']
})
@classmethod
def from_url(cls, url, **kwargs):
connection = Connection.from_url(url, **kwargs)
return cls(connection=connection)
def __init__(self,
host='localhost',
port=11300,
socket_timeout=None,
socket_connect_timeout=None,
socket_keepalive=None,
retry_on_timeout=False,
socket_keepalive_options=None,
max_connections=None,
connection=None):
if not connection:
self.socket_timeout = socket_timeout
kwargs = {
'host': host,
'port': port,
'socket_connect_timeout': socket_connect_timeout,
'socket_keepalive': socket_keepalive,
'socket_keepalive_options': socket_keepalive_options,
'socket_timeout': socket_timeout,
'retry_on_timeout': retry_on_timeout,
'max_connections': max_connections
}
connection = Connection(**kwargs)
self.conn_queue = LifoQueue()
self.conn_queue.put_nowait(connection)
self._connection = connection
self.response_callbacks = self.__class__.RESPONSE_CALLBACKS.copy()
self.expected_ok = self.__class__.EXPECTED_OK.copy()
self.expected_err = self.__class__.EXPECTED_ERR.copy()
def _get_connection(self):
try:
connection = self.conn_queue.get(block=True, timeout=None)
except Empty:
raise ConnectionError("No connection available")
return connection
def _release_connection(self, connection):
self.conn_queue.put_nowait(connection)
def execute_command(self, *args, **kwargs):
connection = self._get_connection()
command_name = args[0]
try:
connection.send_command(*args, **kwargs)
return self.parse_response(connection, command_name, **kwargs)
except (ConnectionError, TimeoutError):
connection.disconnect()
raise
finally:
self._release_connection(connection)
def parse_response(self, connection, command_name, **kwargs):
response = connection.read_response()
status, results = response
if status in self.expected_ok[command_name]:
if command_name in self.response_callbacks:
return self.response_callbacks[command_name](connection,
response,
**kwargs)
return response
elif status in self.expected_err[command_name]:
raise ResponseError(command_name, status, results)
else:
raise InvalidResponse(command_name, status, results)
return response
def put(self, body, priority=DEFAULT_PRIORITY, delay=0, ttr=DEFAULT_TTR):
assert isinstance(body, str), 'body must be str'
job_id = self.execute_command('put', priority, delay, ttr, body=body)
return job_id
def use(self, tube):
self._connection.use(tube)
def watch(self, tube):
self._connection.watch(tube)
def reserve(self, timeout=None):
if timeout is not None:
return self.execute_command('reserve-with-timeout', timeout)
else:
return self.execute_command('reserve')
def bury(self, job_id, priority=DEFAULT_PRIORITY):
self.execute_command('bury', job_id, priority)
def release(self, job_id, priority=DEFAULT_PRIORITY, delay=0):
self.execute_command('release', job_id, priority, delay)
def delete(self, job_id):
self.execute_command('delete', job_id)
def kick_job(self, job_id):
"""
Variant of` kick` that operates with a single job.
:param job_id: the job id to kick
:type job_id: `str`
"""
self.execute_command('kick-job', job_id)
def kick(self, bound=1000):
"""
Move jobs into the ready queue.
If there are any buried jobs, it will only kick buried jobs.
Otherwise it will kick delayed jobs.
:param bound: upper bound on the number of jobs to kick
:type bound: `int`
"""
kicked = int(self.execute_command('kick', str(bound))[1][0])
return kicked
def stats_tube(self, tube):
return self.execute_command('stats-tube', tube)
def close(self):
if self._connection:
self._connection.disconnect()
