Synnefo

AMQP Client's instatiated by this module silenty detect connection failures and

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)

class AMQPClient(object):

    This class will create an object of AMQPPukaClient or AMQPHaigha client

    depending on AMQP_BACKEND setting

    def __new__(cls, *args, **kwargs):

        return Client(*args, **kwargs)

# 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

# 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.

""" Module implementing connection and communication with an AMQP broker.

AMQP Client's implemented 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 logging

from puka import Client

from puka import spec_exceptions

from socket import error as socket_error

from time import sleep

from random import shuffle

from functools import wraps

from ordereddict import OrderedDict

from synnefo import settings

logging.basicConfig(level=logging.DEBUG,

                    format="[%(levelname)s %(asctime)s] %(message)s")

logger = logging.getLogger()

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, spec_exceptions.ConnectionForced) as e:

            logger.error('Connection Closed while in %s: %s', func.__name__, e)

            self.consume_promises = []

            self.connect()

    return wrapper

class AMQPPukaClient(object):

    """

    AMQP generic client implementing most of the basic AMQP operations.

    """

    def __init__(self, hosts=settings.AMQP_HOSTS, max_retries=30,

                 confirms=True, confirm_buffer=100):

        """Format hosts as "amqp://username:pass@host:port" """

        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()

        self.unsend = OrderedDict()

        self.consume_promises = []

    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)

        self.client = Client(host, pubacks=self.confirms)

        host = host.split('@')[-1]

        logger.debug('Connecting to node %s' % host)

        try:

            promise = self.client.connect()

            self.client.wait(promise)

        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')

        if self.unacked:

            self._resend_unacked_messages()

        if self.unsend:

            self._resend_unsend_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)

        promise = self.client.exchange_declare(exchange=exchange,

                                               type=type,

                                               durable=True,

                                               auto_delete=False)

        self.client.wait(promise)

    @reconnect_decorator

    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 = {}

        promise = self.client.queue_declare(queue=queue, durable=True,

                                            exclusive=exclusive,

                                            auto_delete=False,

                                            arguments=arguments)

        self.client.wait(promise)

    def queue_bind(self, queue, exchange, routing_key):

        logger.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)

    @reconnect_decorator

    def basic_publish(self, exchange, routing_key, body):

        """Publish a message with a specific routing key """

        self._publish(exchange, routing_key, body)

        self.flush_buffer()

        if self.confirms and len(self.unacked) >= self.confirm_buffer:

            self.get_confirms()

    @reconnect_decorator

    def basic_publish_multi(self, exchange, routing_key, bodies):

        for body in bodies:

            self.unsend[body] = (exchange, routing_key)

        for body in bodies:

            self._publish(exchange, routing_key, body)

            self.unsend.pop(body)

        self.flush_buffer()

        if self.confirms:

            self.get_confirms()

    def _publish(self, exchange, routing_key, body):

        # Persisent messages by default!

        headers = {}

        headers['delivery_mode'] = 2

        promise = self.client.basic_publish(exchange=exchange,

                                            routing_key=routing_key,

                                            body=body, headers=headers)

        if self.confirms:

            self.unacked[promise] = (exchange, routing_key, body)

        return promise

    @reconnect_decorator

    def flush_buffer(self):

        while self.client.needs_write():

            self.client.on_write()

    @reconnect_decorator

    def get_confirms(self):

        for promise in self.unacked.keys():

            self.client.wait(promise)

            self.unacked.pop(promise)

    @reconnect_decorator

    def _resend_unacked_messages(self):

        """Resend unacked messages in case of a connection failure."""

        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 _resend_unsend_messages(self):

        """Resend unsend messages in case of a connection failure."""

        for body in self.unsend.keys():

            (exchange, routing_key) = self.unsend[body]

            self.basic_publish(exchange, routing_key, body)

            self.unsend.pop(body)

    @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

        """

        # Store the queues and the callback

        self.consumers[queue] = callback

        def handle_delivery(promise, msg):

            """Hide promises and messages without body"""

            if 'body' in msg:

                callback(self, msg)

            else:

                logger.debug("Message without body %s" % msg)

                raise socket_error

        consume_promise = \

                self.client.basic_consume(queue=queue,

                                          prefetch_count=prefetch_count,

                                          callback=handle_delivery)

        self.consume_promises.append(consume_promise)

        return consume_promise

    @reconnect_decorator

    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.

        """

        if promise is not None:

            self.client.wait(promise, timeout)

        else:

            self.client.wait(self.consume_promises)

    @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

    @reconnect_decorator

    def basic_ack(self, message):

        self.client.basic_ack(message)

    @reconnect_decorator

    def basic_nack(self, message):

        #TODO:

        pass

    def close(self):

        """Check that messages have been send and close the connection."""

        try:

            if self.confirms:

                self.get_confirms()

            close_promise = self.client.close()

            self.client.wait(close_promise)

        except (socket_error, spec_exceptions.ConnectionForced) as e:

            logger.error('Connection closed while closing connection:%s',

                          e)

    def queue_delete(self, queue, if_unused=True, if_empty=True):

        """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 spec_exceptions.NotFound:

            logger.info("Queue %s does not exist", queue)

            return False

    def exchange_delete(self, exchange, if_unused=True):

        """Delete an exchange."""

        try:

            promise = self.client.exchange_delete(exchange=exchange,

                                                  if_unused=if_unused)

            self.client.wait(promise)

            return True

        except spec_exceptions.NotFound:

            logger.info("Exchange %s does not exist", exchange)

            return False

    @reconnect_decorator

    def basic_cancel(self, promise=None):

        """Cancel consuming from a queue. """

        if promise is not None:

            self.client.basic_cancel(promise)

        else:

            for promise in self.consume_promises:

                self.client.basic_cancel(promise)

class AMQPConnectionError():

    pass

# AMQP Backend Client. One of: 'puka', 'haigha'

AMQP_BACKEND='puka'

            self.client.exchange_declare(exchange=exchange,

                                         type="topic")

            self.client.queue_declare(queue=queue, mirrored=True)

        self.client = AMQPClient(confirm_buffer=100)

        self.client.exchange_declare(settings.EXCHANGE_GANETI, type='topic')

            self.client.basic_publish(settings.EXCHANGE_GANETI,

                                      routekey,

                                      msg)