Reconnect to rabbitmq in our queue client if the connection is lost

(imported from commit 82ed67f0c0f93e9f79ff65ef55b397815d20aaeb)

zephyr/lib/queue.py

@@ -3,6 +3,8 @@ import pika
 import logging
 import simplejson
 import random
+import time
+from collections import defaultdict
 
 # This simple queuing library doesn't expose much of the power of
 # rabbitmq/pika's queuing system; its purpose is to just provide an
@@ -13,6 +15,7 @@ class SimpleQueueClient(object):
         self.log = logging.getLogger('humbug.queue')
         self.queues = set()
         self.channel = None
+        self.consumers = defaultdict(set)
         self._connect()
 
     def _connect(self):
@@ -20,6 +23,12 @@ class SimpleQueueClient(object):
         self.channel    = self.connection.channel()
         self.log.info('SimpleQueueClient connected')
 
+    def _reconnect(self):
+        self.connection = None
+        self.channel = None
+        self.queues = set()
+        self._connect()
+
     def _get_parameters(self):
         return pika.ConnectionParameters('localhost',
             credentials = pika.PlainCredentials(
@@ -28,12 +37,23 @@ class SimpleQueueClient(object):
     def _generate_ctag(self, queue_name):
         return "%s_%s" % (queue_name, str(random.getrandbits(16)))
 
+    def _reconnect_callbacks(self):
+        for queue, consumers in self.consumers.items():
+            for consumer in consumers:
+                self.ensure_queue(queue, lambda: self.channel.basic_consume(
+                                                        consumer,
+                                                        queue=queue,
+                                                        consumer_tag=self._generate_ctag(queue)))
+
     def ready(self):
         return self.channel is not None
 
     def ensure_queue(self, queue_name, callback):
         '''Ensure that a given queue has been declared, and then call
            the callback with no arguments.'''
+        if not self.connection.is_open:
+            self._connect()
+
         if queue_name not in self.queues:
             self.channel.queue_declare(queue=queue_name, durable=True)
             self.queues.add(queue_name)
@@ -48,13 +68,20 @@ class SimpleQueueClient(object):
                 body=body))
 
     def json_publish(self, queue_name, body):
-        return self.publish(queue_name, simplejson.dumps(body))
+        try:
+            return self.publish(queue_name, simplejson.dumps(body))
+        except (AttributeError, pika.exceptions.AMQPConnectionError):
+            self.log.warning("Failed to send to rabbitmq, trying to reconnect and send again")
+            self._reconnect()
+
+            return self.publish(queue_name, simplejson.dumps(body))
 
     def register_consumer(self, queue_name, consumer):
         def wrapped_consumer(ch, method, properties, body):
             consumer(ch, method, properties, body)
             ch.basic_ack(delivery_tag=method.delivery_tag)
 
+        self.consumers[queue_name].add(wrapped_consumer)
         self.ensure_queue(queue_name,
             lambda: self.channel.basic_consume(wrapped_consumer, queue=queue_name,
                 consumer_tag=self._generate_ctag(queue_name)))
@@ -74,20 +101,46 @@ class TornadoQueueClient(SimpleQueueClient):
     # Based on:
     # https://pika.readthedocs.org/en/0.9.8/examples/asynchronous_consumer_example.html
 
-    def _connect(self):
+    def _connect(self, on_open_cb = None):
+        self.log.info("Beginning TornadoQueueClient connection")
+        self._on_open_cb = on_open_cb
         self.connection = pika.adapters.TornadoConnection(
             self._get_parameters(),
             on_open_callback = self._on_open,
             stop_ioloop_on_close = False)
 
+    def _reconnect(self, on_open_cb = None):
+        self.connection = None
+        self.channel = None
+        self.queues = set()
+        self._connect(on_open_cb)
+
     def _on_open(self, connection):
+        self.connection.add_on_close_callback(self._on_connection_closed)
         self.connection.channel(
             on_open_callback = self._on_channel_open)
 
     def _on_channel_open(self, channel):
         self.channel = channel
+        if self._on_open_cb:
+            self._on_open_cb()
         self.log.info('TornadoQueueClient connected')
 
+    def _on_connection_closed(self, method_frame):
+        self.log.warning("TornadoQueueClient lost connection to RabbitMQ, reconnecting...")
+        from tornado import ioloop
+
+        # Try to reconnect in two seconds
+        retry_seconds = 2
+        def on_timeout():
+            try:
+                self._reconnect(self._reconnect_callbacks)
+            except pika.exceptions.AMQPConnectionError:
+                self.log.critical("Failed to reconnect to RabbitMQ, retrying...")
+                ioloop.IOLoop.instance().add_timeout(time.time() + retry_seconds, on_timeout)
+
+        ioloop.IOLoop.instance().add_timeout(time.time() + retry_seconds, on_timeout)
+
     def ensure_queue(self, queue_name, callback):
         def finish(frame):
             self.queues.add(queue_name)