/snf-common/synnefo/lib/amqp.py - Diff - Synnefo - Greek Research and Technology Network's projects

Revision db400d82 snf-common/synnefo/lib/amqp.py

     """ Module implementing connection and communication with an AMQP broker.
     AMQP Client's implemented by this module silenty detect connection failures and
     AMQP Client's instatiated by this module silenty detect connection failures and
     try to reconnect to any available broker. Also publishing takes advantage of
     publisher-confirms in order to guarantee that messages are properly delivered
     to the broker.
     """
     import puka
     import logging
     import socket
     from time import sleep
     from random import shuffle
     from functools import wraps
     from ordereddict import OrderedDict
     from synnefo import settings
     AMQP_HOSTS = settings.AMQP_HOSTS
     MAX_RETRIES = 20
     if settings.AMQP_BACKEND == 'puka':
         from amqp_puka import AMQPPukaClient as Client
     elif settings.AMQP_BACKEND == 'haigha':
         from amqp_haigha import AMQPHaighaClient as Client
     else:
         raise Exception('Unknown Backend %s' % settings.AMQP_BACKEND)
     log = logging.getLogger()
     def reconnect_decorator(func):
         """
         Decorator for persistent connection with one or more AMQP brokers.
         """
         @wraps(func)
         def wrapper(self, *args, **kwargs):
             try:
                 if self.client.sd is None:
                     self.connect()
                 return func(self, *args, **kwargs)
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 log.debug("Connection closed. Reconnecting")
                 self.connect()
                 return wrapper(self, *args, **kwargs)
             except Exception:
                 if self.client.sd is None:
                     log.debug("Connection closed. Reconnecting")
                     self.connect()
                     return wrapper(self, *args, **kwargs)
                 else:
                     raise
         return wrapper
     class AMQPClient():
     class AMQPClient(object):
         """
         AMQP generic client implementing most of the basic AMQP operations.
         This client confirms delivery of each published message before publishing
         the next one, which results in low performance. Better performance can be
         achieved by using AMQPAsyncClient.
         This class will create an object of AMQPPukaClient or AMQPHaigha client
         depending on AMQP_BACKEND setting
         """
         def __init__(self, hosts=AMQP_HOSTS, max_retries=MAX_RETRIES):
             """Format hosts as "amqp://username:pass@host:port" """
             # Shuffle the elements of the host list for better load balancing
             self.hosts = hosts
             shuffle(self.hosts)
             self.max_retries = max_retries
             self.promises = []
             self.consume_promises = []
             self.consume_info = {}
         def connect(self, retries=0):
             # Pick up a host
             url = self.hosts.pop()
             self.hosts.insert(0, url)
             if retries > self.max_retries:
                 raise AMQPError("Cannot connect to any node after %s attemps" % retries)
             if retries > 2 * len(self.hosts):
                 sleep(1)
             self.client = puka.Client(url, pubacks=True)
             host = url.split('@')[-1]
             log.debug('Connecting to node %s' % host)
             try:
                 promise = self.client.connect()
                 self.client.wait(promise)
             except socket.error as e:
                 log.debug('Cannot connect to node %s.' % host)
                 return self.connect(retries+1)
         @reconnect_decorator
         def exchange_declare(self, exchange_name, exchange_type='direct',
                              durable=True, auto_delete=False):
             """Declare an exchange
             @type exchange_name: string
             @param exchange_name: name of the exchange
             @type exchange_type: string
             @param exhange_type: one of 'direct', 'topic', 'fanout'
             """
             log.debug('Declaring exchange %s of %s type.'
                       %(exchange_name, exchange_type))
             promise = self.client.exchange_declare(exchange=exchange_name,
                                                    type=exchange_type,
                                                    durable=durable,
                                                    auto_delete=auto_delete)
             self.client.wait(promise)
             log.debug('Exchange %s declared succesfully ' % exchange_name)
         @reconnect_decorator
         def queue_declare(self, queue, durable=True, exclusive=False,
                           auto_delete=False, mirrored=True, mirrored_nodes='all'):
             """Declare a queue
             @type queue: string
             @param queue: name of the queue
             @param mirrored: whether the queue will be mirrored to other brokers
             @param mirrored_nodes: the policy for the mirrored queue.
                 Available policies:
                     - 'all': The queue is mirrored to all nodes and the
                       master node is the one to which the client is
                       connected
                     - a list of nodes. The queue will be mirrored only to
                       the specified nodes, and the master will be the
                       first node in the list. Node names must be provided
                       and not host IP. example: [node1@rabbit,node2@rabbit]
             """
             log.debug('Declaring queue %s' % queue)
             if mirrored:
                 if mirrored_nodes == 'all':
                     arguments={'x-ha-policy':'all'}
                 elif isinstance(mirrored_nodes, list):
                     arguments={'x-ha-policy':'nodes', 'x-ha-policy-params':mirrored_nodes}
                 else:
                     raise AttributeError
             else:
                 arguments = {}
             promise = self.client.queue_declare(queue=queue, durable=durable,
                                                 exclusive=exclusive,
                                                 auto_delete=auto_delete,
                                                 arguments=arguments)
             self.client.wait(promise)
             log.debug('Queue %s declared successfully.' % queue)
         @reconnect_decorator
         def queue_bind(self, queue, exchange, routing_key):
             log.debug('Binding queue %s to exchange %s with key %s'
                      % (queue, exchange, routing_key))
             promise = self.client.queue_bind(exchange=exchange, queue=queue,
                                              routing_key=routing_key)
             self.client.wait(promise)
             log.debug('Binding completed successfully')
         def basic_publish_multi(self, exhange, routing_key, msgs, headers={}):
             """Send many messages to the  same exchange and with the same key """
             for msg in msgs:
                 self.basic_publish(exchange, routing_key, headers, msg)
         @reconnect_decorator
         def basic_publish(self, exchange, routing_key, body, headers={}):
             """Publish a message with a specific routing key """
             # Persisent messages by default!
             if not 'delivery_mode' in headers:
                 headers['delivery_mode'] = 2
             promise = self.client.basic_publish(exchange=exchange,
                                                 routing_key=routing_key,
                                                 body=body, headers=headers)
             self.client.wait(promise)
         @reconnect_decorator
         def basic_consume(self, queue, callback, prefetch_count=0):
             """Consume from a queue.
             @type queue: string or list of strings
             @param queue: the name or list of names from the queues to consume
             @type callback: function
             @param callback: the callback function to run when a message arrives
             """
             if isinstance(queue, str):
                 queue = [queue]
             elif isinstance(queue, list):
                 pass
             else:
                 raise AttributeError
             # Store the queues and the callback
             for q in queue:
                 self.consume_info[q] = callback
             def handle_delivery(promise, result):
                 """Hide promises and messages without body"""
                 if 'body' in result:
                     callback(self, result)
                 else:
                     log.debug("Message without body %s" % result)
                     raise socket.error
             consume_promise = \
                     self.client.basic_consume_multi(queues=queue,
                                                     prefetch_count=prefetch_count,
                                                     callback=handle_delivery)
             self.consume_promises.append(consume_promise)
             return consume_promise
         def basic_wait(self, promise=None, timeout=0):
             """Wait for messages from the queues declared by basic_consume.
             This function will block until a message arrives from the queues that
             have been declared with basic_consume. If the optional arguments
             'promise' is given, only messages for this promise will be delivered.
             """
             try:
                 if promise is not None:
                    self.client.wait(promise, timeout)
                 else:
                    self.client.wait(self.consume_promises)
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 log.debug('Connection closed while receiving messages.')
                 self.consume_promises = []
                 try:
                     self.client.close()
                 except:
                     pass
                 self.connect()
                 for queues, callback in self.consume_info.items():
                     self.basic_consume(queues, callback)
                 self.basic_wait(timeout)
             except Exception as e:
                 if self.client.sd is None:
                     log.debug('Connection closed while receiving messages.')
                     self.consume_promises = []
                     self.connect()
                     for queues, callback in self.consume_info.items():
                         self.basic_consume(queues, callback)
                     self.basic_wait(timeout)
                 else:
                     log.error("Exception while waiting for messages ",e)
                     raise
         def basic_cancel(self, promise=None):
             """Cancel consuming from a queue. """
             try:
                 if promise is not None:
                     self.client.basic_cancel(promise)
                 else:
                     for promise in self.consume_promises:
                         self.client.basic_cancel(promise)
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 pass
             except Exception as e:
                 if self.client.sd is None:
                     pass;
                 else:
                     log.error("Exception while canceling client ",e)
                     raise
         @reconnect_decorator
         def basic_get(self, queue):
             """Get a single message from a queue.
             This is a non-blocking method for getting messages from a queue.
             It will return None if the queue is empty.
             """
             get_promise = self.client.basic_get(queue=queue)
             result = self.client.wait(get_promise)
             if 'empty' in result:
                 # The queue is empty
                 return None
             else:
                 return result
         def basic_ack(self, message):
             """Acknowledge a message. """
             try:
                 self.client.basic_ack(message)
                 return True
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 return False
             except Exception as e:
                 if self.client.sd is None:
                     return False
                 else:
                     log.error("Exception while acknowleding message ",e)
                     raise
         def close(self):
             """Close the connection with the AMQP broker. """
             try:
                 close_promise = self.client.close()
                 self.client.wait(close_promise)
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 pass
         def queue_delete(self, queue, if_unused=False, if_empty=False):
             """Delete a queue.
             Returns False if the queue does not exist
             """
             try:
                 promise = self.client.queue_delete(queue=queue, if_unused=if_unused,
                                                    if_empty=if_empty)
                 self.client.wait(promise)
                 return True
             except puka.spec_exceptions.NotFound:
                 log.debug("Queue %s does not exist", queue)
                 return False
         def exchange_delete(self, exchange, if_unused=False):
             """Delete an exchange."""
             try:
                 promise = self.client.exchange_delete(exchange=exchange,
                                                       if_unused=if_unused)
                 self.client.wait(promise)
                 return True
             except puka.spec_exceptions.NotFound:
                 log.debug("Exchange %s does not exist", exchange)
                 return False
     class AMQPAsyncClient(AMQPClient):
         """AMQP client implementing asynchronous sending of messages.
         This client is more efficient that AMQPClient in sending large number
         of messages. Messages are confirmed to be sent to the broker in batches
         of a size specified by the 'max_buffer' argument.
         Messages are kept to an internal buffer until the max_buffer messages are
         sent or until the connection closes. Explicit delivering of messages can be
         achieved by calling 'wait_for_promises' method.
         Always remember to close the connection, or messages may be lost.
         """
         def __init__(self, hosts=AMQP_HOSTS, max_buffer=5000,
                      max_retries=MAX_RETRIES):
             AMQPClient.__init__(self, hosts, max_retries)
             self.published_msgs = OrderedDict()
             self._promise_counter = 0
             self.max_buffer=max_buffer
         def basic_publish_multi(self, exhange, routing_key, msgs, headers={}):
             while msgs:
                 msg = msgs.pop[0]
                 self.basic_publish(exchange, routing_key, msg, headers)
         def basic_publish(self, exchange, routing_key, body, headers={}):
             """Publish a message.
             The message will not be actually published to the broker until
             'max_buffer' messages are published or wait_for_promises is called.
             """
             try:
                 if not 'delivery_mode' in headers:
                     headers['delivery_mode'] = 2
                 promise = self.client.basic_publish(exchange=exchange,
                                                     routing_key=routing_key,
                                                     body=body,
                                                     headers=headers)
                 self._promise_counter += 1
                 self.published_msgs[promise] = {'exchange':exchange,
                                                 'routing_key':routing_key,
                                                 'body':body,
                                                 'headers':headers}
                 if self._promise_counter > self.max_buffer:
                     self.wait_for_promises()
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 log.debug('Connection closed while sending message %s.\
                           Reconnecting and retrying' % body)
                 self.connect()
                 self.basic_publish(exchange, routing_key, body, headers)
                 return self._retry_publish_msgs()
             except Exception as e:
                 if self.client.sd is None:
                     log.debug('Connection closed while sending message %s.\
                                  Reconnecting and retrying' % body)
                     self.connect()
                     self.basic_publish(exchange, routing_key, body, headers)
                     return self._retry_publish_msgs()
                 else:
                     log.error("Exception while publishing message ",e)
                     raise
         def wait_for_promises(self):
             """Wait for confirm that all messages are sent."""
             try:
                 promises = self.published_msgs.keys()
                 for promise in promises:
                     self.client.wait(promise)
                     self.published_msgs.pop(promise)
                 self._promise_counter = 0
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 log.debug('Connection closed while waiting from promises')
                 self.connect()
                 self._retry_publish_msgs()
             except Exception as e:
                 if self.client.sd is None:
                     log.debug('Connection closed while waiting from promises')
                     self.connect()
                     self._retry_publish_msgs()
                 else:
                     log.error("Exception while waiting for promises ",e)
                     raise
         def _retry_publish_msgs(self):
             """Resend messages in case of a connection failure."""
             values = self.published_msgs.values()
             self.published_msgs = OrderedDict()
             for message in values:
                 exchange = message['exchange']
                 key = message['routing_key']
                 body = message['body']
                 headers = message['headers']
                 log.debug('Resending message %s' % body)
                 self.basic_publish(exchange, key, body, headers)
         def close(self):
             """Check that messages have been send and close the connection."""
             try:
                 self.wait_for_promises()
                 close_promise = self.client.close()
                 self.client.wait(close_promise)
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 pass
             except Exception as e:
                 if self.client.sd is None:
                     pass
                 else:
                     log.error("Exception while closing the connection ",e)
                     raise
         def flush_buffer(self):
             try:
                 self.client.on_write()
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 log.debug('Connection closed while clearing buffer')
                 self.connect()
                 return self._retry_publish_msgs()
             except Exception as e:
                 if self.client.sd is None:
                     log.debug('Connection closed while clearing buffer')
                     self.connect()
                     return self._retry_publish_msgs()
                 else:
                     log.error("Exception while clearing buffer ",e)
                     raise
     class AMQPConsumer(AMQPClient):
         """AMQP client implementing a consumer without callbacks.
         """
         def __init__(self, hosts=AMQP_HOSTS, max_buffer=5000,
                      max_retries=MAX_RETRIES):
             AMQPClient.__init__(self, hosts, max_retries)
             self.consume_queues = []
             self.consume_promises = []
         @reconnect_decorator
         def basic_consume(self, queue, prefetch_count=0):
             """Consume from a queue.
             @type queue: string or list of strings
             @param queue: the name or list of names from the queues to consume
             @type callback: function
             @param callback: the callback function to run when a message arrives
             """
             if isinstance(queue, str):
                 queue = [queue]
             elif isinstance(queue, list):
                 pass
             else:
                 raise AttributeError
             # Store the queues and the callback
             for q in queue:
                 self.consume_queues.append(q)
             consume_promise = \
                     self.client.basic_consume_multi(queues=queue,
                                                     prefetch_count=prefetch_count)
             self.consume_promises.append(consume_promise)
             return consume_promise
         def basic_wait(self, promise=None, timeout=0):
             """Wait for messages from the queues declared by basic_consume.
             This function will block until a message arrives from the queues that
             have been declared with basic_consume. If the optional arguments
             'promise' is given, only messages for this promise will be delivered.
             """
             try:
                 if promise is not None:
                    return self.client.wait(promise, timeout)
                 else:
                    return self.client.wait(self.consume_promises)
             except (socket.error, puka.spec_exceptions.ConnectionForced):
                 log.debug('Connection closed while receiving messages.')
                 self.consume_promises = []
                 try:
                     self.client.close()
                 except:
                     pass
                 self.connect()
                 for queues in self.consume_queues:
                     self.basic_consume(queues)
                 self.basic_wait(timeout)
             except Exception as e:
                 if self.client.sd is None:
                     log.debug('Connection closed while receiving messages.')
                     self.consume_promises = []
                     self.connect()
                     for queues in self.consume_queues:
                         self.basic_consume(queues)
                     self.basic_wait(timeout)
                 else:
                     log.error("Exception while waiting for messages ",e)
                     raise
     class AMQPError(Exception):
         def __init__(self, msg):
             self.msg = msg
         def __str__(self):
             return repr(self.msg)
         def __new__(cls, *args, **kwargs):
             return Client(*args, **kwargs)

Also available in: Unified diff