def __init__(self, host="localhost", port=5672, virtualhost="/", credentials=None):
from adapter import TornadoPikaAdapter
if credentials is not None:
assert isinstance(credentials, (PlainCredentials, ExternalCredentials))
self.channel = TornadoPikaAdapter(
pika.ConnectionParameters(host=host, port=port, credentials=credentials, virtual_host=virtualhost)
)
client_uid = uuid()
self._res_queue = "crew.master.%s" % client_uid
self._pubsub_queue = "crew.subscribe.%s" % client_uid
self.callbacks_hash = {}
self._subscribe_cache = {}
tornado.ioloop.IOLoop.instance().add_callback(self.connect)
class Client(object):
SERIALIZERS = {
'json': 'application/json',
'pickle': 'application/python-pickle',
'text': 'text/plain',
}
def __init__(self, host='localhost', port=5672, virtualhost='/', credentials=None):
from adapter import TornadoPikaAdapter
if credentials is not None:
assert isinstance(credentials, (PlainCredentials, ExternalCredentials))
self.channel = TornadoPikaAdapter(pika.ConnectionParameters(
host=host,
port=port,
credentials=credentials,
virtual_host=virtualhost,
))
client_uid = uuid()
self.io_loop = tornado.ioloop.IOLoop.current()
self._res_queue = "crew.master.%s" % client_uid
self._pubsub_queue = "crew.subscribe.%s" % client_uid
self.callbacks_hash = {}
self._subscribe_cache = {}
tornado.ioloop.IOLoop.instance().add_callback(self.connect)
def parse_body(self, body, props):
content_type = getattr(props, 'content_type', 'text/plain')
if props.content_encoding == 'gzip':
body = zlib.decompress(body)
if 'application/json' in content_type:
return json.loads(body)
elif 'application/python-pickle' in content_type:
return pickle.loads(body)
def _on_result(self, channel, method, props, body):
log.debug('PikaCient: Result message received, tag #%i len %d', method.delivery_tag, len(body))
correlation_id = getattr(props, 'correlation_id', None)
if correlation_id not in self.callbacks_hash:
log.info('Got result for task "%d", but no has callback', correlation_id)
try:
cb = self.callbacks_hash.pop(correlation_id)
body = self.parse_body(body, props)
if isinstance(cb, Future):
if isinstance(body, Exception):
cb.set_exception(body)
else:
cb.set_result(body)
else:
out = cb(body, headers=props.headers)
return out
finally:
channel.basic_ack(delivery_tag=method.delivery_tag)
def _on_dlx_received(self, channel, method, props, body):
correlation_id = getattr(props, 'correlation_id', None)
if correlation_id in self.callbacks_hash:
cb = self.callbacks_hash.pop(correlation_id)
else:
log.error("Method callback %s is not found", correlation_id)
channel.basic_ack(delivery_tag=method.delivery_tag)
return
try:
dl = props.headers['x-death'][0]
body = ExpirationError(
"Dead letter received. Reason: {0}".format(dl.get('reason'))
)
body.reason = dl.get('reason')
body.time = dl.get('time')
body.expiration = int(dl.get('original-expiration')) / 1000
if isinstance(cb, Future):
self.io_loop.add_callback(cb.set_result, body)
elif callable(cb):
self.io_loop.add_callback(cb, body)
else:
log.error("Callback is not callable")
finally:
channel.basic_ack(delivery_tag=method.delivery_tag)
@tornado.gen.coroutine
def connect(self):
self.channel.exchange_declare("crew.PUBSUB", auto_delete=True, exchange_type="headers")
self.channel.exchange_declare("crew.DLX", auto_delete=True, exchange_type="headers")
self.channel.queue_declare(queue="crew.DLX", auto_delete=False)
self.channel.queue_declare(
queue=self._res_queue, exclusive=True,
auto_delete=True, arguments={"x-message-ttl": 60000}
)
self.channel.queue_declare(
queue=self._pubsub_queue, exclusive=True,
auto_delete=True, arguments={"x-message-ttl": 60000}
)
self.channel.queue_bind("crew.DLX", "crew.DLX", arguments={"x-original-sender": self._res_queue})
self.channel.consume(queue="crew.DLX", callback=self._on_dlx_received)
self.channel.consume(queue=self._pubsub_queue, callback=self._on_subscribed_message)
self.channel.consume(queue=self._res_queue, callback=self._on_result)
yield self.channel.connect()
def call(self, channel, data=None, callback=None, serializer='pickle', headers=None, persistent=True, priority=0,
expiration=86400, timestamp=None, gzip=None, gzip_level=6, set_cid=None, routing_key=None, exchange=''):
if not headers:
headers = {}
assert priority <= 255
assert isinstance(expiration, int) and expiration > 0
qname = "crew.tasks.%s" % channel
serializer, content_type = self.get_serializer(serializer)
if set_cid:
cid = str(set_cid)
if cid in self.callbacks_hash:
raise DuplicateTaskId('Task ID: {0} already exists'.format(cid))
else:
cid = "{0}.{1}".format(channel, uuid())
data = serializer(data)
if gzip is None and data is not None and len(data) > 1024 * 32:
gzip = True
data = zlib.compress(data, gzip_level) if gzip else data
headers.update({"x-original-sender": self._res_queue})
props = pika.BasicProperties(
content_encoding='gzip' if gzip else 'plain',
content_type=content_type,
reply_to=self._res_queue if not routing_key else routing_key,
correlation_id=cid,
headers=headers,
timestamp=int(time.time()),
delivery_mode=2 if persistent else None,
priority=priority,
expiration="%d" % (expiration * 1000),
)
if callback is None:
callback = Future()
self.callbacks_hash[props.correlation_id] = callback
self.channel.basic_publish(
exchange=exchange,
routing_key=qname,
properties=props,
body=data
)
if isinstance(callback, Future):
return callback
else:
return props.correlation_id
def _on_subscribed_message(self, channel, method, props, body):
key = props.headers['x-channel-name']
cb = self._subscribe_cache.get(key, None)
if not cb:
log.error("[PubSub] Method callback %s is not found", key)
channel.basic_ack(delivery_tag=method.delivery_tag)
return
body = self.parse_body(body, props)
channel.basic_ack(delivery_tag=method.delivery_tag)
if isinstance(cb, Future):
self.io_loop.add_callback(cb.set_result, body)
elif callable(cb):
self.io_loop.add_callback(cb, body)
else:
log.error("Callback is not callable")
def subscribe(self, channel, callback):
self.channel.queue_bind(self._pubsub_queue, exchange="crew.PUBSUB", arguments={"x-channel-name": channel})
self._subscribe_cache[channel] = callback
@tornado.gen.coroutine
def unsubscribe(self, qname, callback):
log.debug('Cancelling subscription for channel: "%s"', qname)
yield self.channel.cancel(qname)
def get_serializer(self, name):
assert name in self.SERIALIZERS
if name == 'pickle':
return (lambda x: pickle.dumps(x, protocol=2), self.SERIALIZERS[name])
elif name == 'json':
return (json.dumps, self.SERIALIZERS[name])
elif name == 'text':
return lambda x: str(x).encode('utf-8')
def publish(self, channel, message, serializer='pickle'):
assert serializer in self.SERIALIZERS
serializer, t = self.get_serializer(serializer)
self.channel.basic_publish(
exchange='crew.PUBSUB',
routing_key='',
body=serializer(message),
properties=pika.BasicProperties(
content_type=t, delivery_mode=1,
headers={"x-channel-name": channel}
)
)
def _on_custom_consume(self, callback, channel, method, props, body):
log.debug('PikaCient: Result message received, tag #%i len %d', method.delivery_tag, len(body))
try:
body = self.parse_body(body, props)
return callback(body, headers=props.headers)
finally:
channel.basic_ack(delivery_tag=method.delivery_tag)
def consume(self, queue, callback):
self.channel.consume(queue=queue, callback=lambda *a: self._on_custom_consume(callback, *a))
def parallel(self):
return MultitaskCall(self)
@tornado.gen.coroutine
def close(self):
yield self.channel.close()
yield self.connect.close()
