Merge pull request #52 from rafaelmagu/master

Jose Diaz-Gonzalez · Jose Diaz-Gonzalez · commit 01f080d07c7c · 2012-11-28T11:13:15.000-08:00
Added resiliency to RabbitMQ transport
diff --git a/beaver/rabbitmq_transport.py b/beaver/rabbitmq_transport.py
@@ -3,6 +3,7 @@
 import pika
 
 import beaver.transport
+from beaver.transport import TransportException
 
 
 class RabbitmqTransport(beaver.transport.Transport):
@@ -52,18 +53,30 @@ def __init__(self, configfile):
     def callback(self, filename, lines):
         timestamp = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ")
         for line in lines:
-            self.channel.basic_publish(
-                exchange=self.rabbitmq_exchange,
-                routing_key=self.rabbitmq_key,
-                body=self.format(filename, timestamp, line),
-                properties=pika.BasicProperties(
-                    content_type="text/json",
-                    delivery_mode=1
-                )
-            )
+            try:
+                import warnings
+                with warnings.catch_warnings():
+                    warnings.simplefilter("error")
+                    self.channel.basic_publish(
+                        exchange=self.rabbitmq_exchange,
+                        routing_key=self.rabbitmq_key,
+                        body=self.format(filename, timestamp, line),
+                        properties=pika.BasicProperties(
+                            content_type="text/json",
+                            delivery_mode=1
+                        )
+                    )
+            except UserWarning:
+                raise TransportException("Connection appears to have been lost")
+            except Exception, e:
+                try:
+                    raise TransportException(e.strerror)
+                except AttributeError:
+                    raise TransportException("Unspecified exception encountered")  # TRAP ALL THE THINGS!
 
     def interrupt(self):
-        self.connection.close()
+        if self.connection:
+            self.connection.close()
 
     def unhandled(self):
         return True
diff --git a/beaver/transport.py b/beaver/transport.py
@@ -61,3 +61,7 @@ def format(self, filename, timestamp, line):
             })
 
         return "[{0}] [{1}] {2}".format(self.current_host, timestamp, line)
+
+
+class TransportException(Exception):
+    pass
diff --git a/beaver/worker.py b/beaver/worker.py
@@ -5,6 +5,7 @@
 import stat
 import sys
 import time
+from transport import TransportException
 
 
 class Worker(object):
@@ -225,6 +226,8 @@ def run_worker(configfile, args):
         l = Worker(configfile, args, transport.callback)
         logger.info("Working...")
         l.loop()
+    except TransportException, e:
+        raise TransportException(e.message)
     except KeyboardInterrupt:
         logger.info("Shutting down. Please wait.")
         transport.interrupt()
diff --git a/bin/beaver b/bin/beaver
@@ -2,11 +2,15 @@
 
 import argparse
 import os
+import sys
 
 import beaver.config
 import beaver.utils
 import beaver.worker
 
+from beaver.transport import TransportException
+from time import sleep
+
 epilog_example = """
 Beaver provides an lightweight method for shipping local log
 files to Logstash. It does this using either redis, stdin,
@@ -62,7 +66,6 @@ parser.set_defaults(
     debug=False
 )
 
-
 args = parser.parse_args()
 args.globs = files
 
@@ -76,4 +79,26 @@ except AttributeError:
     args.files = configfile.getfilepaths()
     args.globs = configfile.getglobs()
 
-beaver.worker.run_worker(configfile, args)
+failure_count = 0
+respawn_delay = 3
+max_failure = 7
+
+while 1:
+    try:
+        worker = beaver.worker.run_worker(configfile, args)
+    except TransportException, e:
+        failure_count = failure_count + 1
+        if failure_count > max_failure:
+            failure_count = max_failure
+        sleep_time = respawn_delay ** failure_count
+        logger.info("Caught transport exception, respawning in %d seconds" % sleep_time)
+        try:
+            sleep(sleep_time)
+        except KeyboardInterrupt:
+            logger.info("User cancelled respawn.")
+            sys.exit(0)
+    except KeyboardInterrupt:
+        logger.info("Shutting down. Please wait.")
+        worker.close()
+        logger.info("Shutdown complete.")
+        sys.exit(0)
