# Copyright 2012 GRNET S.A. All rights reserved. # # Redistribution and use in source and binary forms, with or # without modification, are permitted provided that the following # conditions are met: # # 1. Redistributions of source code must retain the above # copyright notice, this list of conditions and the following # disclaimer. # # 2. Redistributions in binary form must reproduce the above # copyright notice, this list of conditions and the following # disclaimer in the documentation and/or other materials # provided with the distribution. # # THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS # OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED # WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR # PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR # CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF # USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED # AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT # LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN # ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE # POSSIBILITY OF SUCH DAMAGE. # # The views and conclusions contained in the software and # documentation are those of the authors and should not be # interpreted as representing official policies, either expressed # or implied, of GRNET S.A. from haigha.connections import RabbitConnection from haigha.message import Message from haigha import exceptions from random import shuffle from time import sleep import logging import socket from synnefo import settings from ordereddict import OrderedDict import gevent from gevent import monkey from functools import wraps logging.basicConfig(level=logging.INFO, format="[%(levelname)s %(asctime)s] %(message)s" ) logger = logging.getLogger('haigha') sock_opts = { (socket.IPPROTO_TCP, socket.TCP_NODELAY): 1, } def reconnect_decorator(func): """ Decorator for persistent connection with one or more AMQP brokers. """ @wraps(func) def wrapper(self, *args, **kwargs): try: func(self, *args, **kwargs) except (socket.error, exceptions.ConnectionError) as e: logger.error('Connection Closed while in %s: %s', func.__name__, e) self.connect() return wrapper class AMQPHaighaClient(): def __init__(self, hosts=settings.AMQP_HOSTS, max_retries=30, confirms=True, confirm_buffer=200): self.hosts = hosts shuffle(self.hosts) self.max_retries = max_retries self.confirms = confirms self.confirm_buffer = confirm_buffer self.connection = None self.channel = None self.consumers = {} self.unacked = OrderedDict() def connect(self, retries=0): if retries > self.max_retries: logger.error("Aborting after %s retries", retries - 1) raise AMQPConnectionError('Aborting after %d connection failures.'\ % (retries - 1)) return # Pick up a host host = self.hosts.pop() self.hosts.insert(0, host) #Patch gevent monkey.patch_all() try: self.connection = \ RabbitConnection(logger=logger, debug=True, user='rabbit', password='r@bb1t', vhost='/', host=host, heartbeat=None, sock_opts=sock_opts, transport='gevent') except socket.error as e: logger.error('Cannot connect to host %s: %s', host, e) if retries > 2 * len(self.hosts): sleep(1) return self.connect(retries + 1) logger.info('Successfully connected to host: %s', host) logger.info('Creating channel') self.channel = self.connection.channel() if self.confirms: self._confirm_select() if self.unacked: self._resend_unacked_messages() if self.consumers: for queue, callback in self.consumers.items(): self.basic_consume(queue, callback) def exchange_declare(self, exchange, type='direct'): """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' """ logger.info('Declaring %s exchange: %s', type, exchange) self.channel.exchange.declare(exchange, type, auto_delete=False, durable=True) def queue_declare(self, queue, exclusive=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] """ logger.info('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 = {} self.channel.queue.declare(queue, durable=True, exclusive=exclusive, auto_delete=False, arguments=arguments) def queue_bind(self, queue, exchange, routing_key): logger.info('Binding queue %s to exchange %s with key %s', queue, exchange, routing_key) self.channel.queue.bind(queue=queue, exchange=exchange, routing_key=routing_key) def _confirm_select(self): logger.info('Setting channel to confirm mode') self.channel.confirm.select() self.channel.basic.set_ack_listener(self._ack_received) self.channel.basic.set_nack_listener(self._nack_received) @reconnect_decorator def basic_publish(self, exchange, routing_key, body): msg = Message(body, delivery_mode=2) mid = self.channel.basic.publish(msg, exchange, routing_key) if self.confirms: self.unacked[mid] = (exchange, routing_key, body) if len(self.unacked) > self.confirm_buffer: self.get_confirms() logger.debug('Published message %s with id %s', body, mid) @reconnect_decorator def get_confirms(self): self.connection.read_frames() @reconnect_decorator def _resend_unacked_messages(self): msgs = self.unacked.values() self.unacked.clear() for exchange, routing_key, body in msgs: logger.debug('Resending message %s', body) self.basic_publish(exchange, routing_key, body) @reconnect_decorator def _ack_received(self, mid): print mid logger.debug('Received ACK for message with id %s', mid) self.unacked.pop(mid) @reconnect_decorator def _nack_received(self, mid): logger.error('Received NACK for message with id %s. Retrying.', mid) (exchange, routing_key, body) = self.unacked[mid] self.basic_publish(exchange, routing_key, body) def basic_consume(self, queue, callback): """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 """ self.consumers[queue] = callback self.channel.basic.consume(queue, consumer=callback, no_ack=False) @reconnect_decorator def basic_wait(self): """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. """ self.connection.read_frames() gevent.sleep(0) @reconnect_decorator def basic_get(self, queue): self.channel.basic.get(queue, no_ack=False) @reconnect_decorator def basic_ack(self, message): delivery_tag = message.delivery_info['delivery_tag'] self.channel.basic.ack(delivery_tag) @reconnect_decorator def basic_nack(self, message): delivery_tag = message.delivery_info['delivery_tag'] self.channel.basic.ack(delivery_tag) def close(self): try: if self.confirms: while self.unacked: print self.unacked self.get_confirms() self.channel.close() close_info = self.channel.close_info logger.info('Successfully closed channel. Info: %s', close_info) self.connection.close() except socket.error as e: logger.error('Connection closed while closing connection:%s', e) def queue_delete(self, queue, if_unused=True, if_empty=True): self.channel.queue.delete(queue, if_unused, if_empty) def exchange_delete(self, exchange, if_unused=True): self.channel.exchange.delete(exchange, if_unused) def basic_class(self): pass class AMQPConnectionError(): pass