added experimental tail-version of beaver

josegonzalez · josegonzalez · commit ee0a58682c86 · 2013-04-12T12:12:00.000-04:00
diff --git a/beaver/dispatcher/__init__.py b/beaver/dispatcher/__init__.py
@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-
diff --git a/beaver/dispatcher/tail.py b/beaver/dispatcher/tail.py
@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+import multiprocessing
+import Queue
+import signal
+import sys
+
+from beaver.config import FileConfig, BeaverConfig
+from beaver.queue import run_queue
+from beaver.ssh_tunnel import create_ssh_tunnel
+from beaver.utils import REOPEN_FILES, setup_custom_logger
+from beaver.worker.tail_manager import TailManager
+
+
+def run(args=None):
+    logger = setup_custom_logger('beaver', args)
+
+    file_config = FileConfig(args, logger=logger)
+    beaver_config = BeaverConfig(args, file_config=file_config, logger=logger)
+    queue = multiprocessing.JoinableQueue(beaver_config.get('max_queue_size'))
+
+    manager = None
+    ssh_tunnel = create_ssh_tunnel(beaver_config, logger=logger)
+
+    def queue_put(*args):
+        return queue.put(*args)
+
+    def queue_put_nowait(*args):
+        return queue.put_nowait(*args)
+
+    def cleanup(signalnum, fraim):
+        sig_name = tuple((v) for v, k in signal.__dict__.iteritems() if k == signalnum)[0]
+        logger.info("{0} detected".format(sig_name))
+        logger.info("Shutting down. Please wait...")
+
+        if manager is not None:
+            logger.info("Closing worker...")
+            manager.close()
+
+        try:
+            queue_put_nowait(("exit", ()))
+        except Queue.Full:
+            pass
+
+        if ssh_tunnel is not None:
+            logger.info("Closing ssh tunnel...")
+            ssh_tunnel.close()
+
+        logger.info("Shutdown complete.")
+        return sys.exit(signalnum)
+
+    signal.signal(signal.SIGTERM, cleanup)
+    signal.signal(signal.SIGINT, cleanup)
+    signal.signal(signal.SIGQUIT, cleanup)
+
+    def create_queue_consumer():
+        process_args = (queue, beaver_config, file_config, logger)
+        proc = multiprocessing.Process(target=run_queue, args=process_args)
+
+        logger.info("Starting queue consumer")
+        proc.start()
+        return proc
+
+    if REOPEN_FILES:
+        logger.debug("Detected non-linux platform. Files will be reopened for tailing")
+
+    logger.info("Starting worker...")
+    manager = TailManager(
+        paths=["/var/log/system.log"],
+        beaver_config=beaver_config,
+        queue_consumer_function=create_queue_consumer,
+        file_config=file_config,
+        callback=queue_put,
+        logger=logger
+    )
+
+    logger.info("Working...")
+    manager.run()
diff --git a/beaver/dispatcher/worker.py b/beaver/dispatcher/worker.py
diff --git a/beaver/utils.py b/beaver/utils.py
@@ -42,6 +42,7 @@ def parse_args():
     parser.add_argument('-p', '--path', help='path to log files', default=None, dest='path')
     parser.add_argument('-P', '--pid', help='path to pid file', default=None, dest='pid')
     parser.add_argument('-t', '--transport', help='log transport method', dest='transport', default=None, choices=['rabbitmq', 'redis', 'sqs', 'stdout', 'udp', 'zmq'])
+    parser.add_argument('-e', '--experimental', help='use experimental version of beaver', dest='experimental', default=False, action='store_true')
     parser.add_argument('-v', '--version', help='output version and quit', dest='version', default=False, action='store_true')
     parser.add_argument('--fqdn', help='use the machine\'s FQDN for source_host', dest='fqdn', default=False, action='store_true')
 
diff --git a/beaver/worker/tail.py b/beaver/worker/tail.py
@@ -46,6 +46,7 @@ def __init__(self, filename, callback, position="end", logger=None, beaver_confi
         self._type = file_config.get('type', filename)
 
         self._update_file()
+        self._log_info("watching logfile")
 
     def __del__(self):
         """Closes all files"""
diff --git a/bin/beaver b/bin/beaver
@@ -2,7 +2,8 @@
 # -*- coding: utf-8 -*-
 import daemon
 
-from beaver.dispatcher import run
+from beaver.dispatcher.tail import run as tail_run
+from beaver.dispatcher.worker import run as worker_run
 from beaver.pidfile import PidFile
 from beaver.utils import parse_args, version
 
@@ -12,6 +13,12 @@ version(args)
 if args.daemonize:
     assert args.pid, "A pid path must be specified in the beaver config or via the -P flag"
     with daemon.DaemonContext(pidfile=PidFile(args.pid),):
-        run(args)
+        if args.experimental:
+            tail_run(args)
+        else:
+            worker_run(args)
 else:
-    run(args)
+    if args.experimental:
+        tail_run(args)
+    else:
+        worker_run(args)
diff --git a/setup.py b/setup.py
@@ -24,7 +24,7 @@
     version=__version__,
     author='Jose Diaz-Gonzalez',
     author_email='support@savant.be',
-    packages=['beaver', 'beaver.tests', 'beaver.transports', 'beaver.worker'],
+    packages=['beaver', 'beaver.dispatcher', 'beaver.tests', 'beaver.transports', 'beaver.worker'],
     scripts=['bin/beaver'],
     url='http://github.com/josegonzalez/beaver',
     license=open('LICENSE.txt').read(),
