Merge pull request #20 from nisanharamati/discharging_queue

Auto-flushing queue by @nisanharamati

Author: borntyping

README.rst

@@ -78,9 +78,28 @@ Requirements
 * `protobuf`_ (when using Python 2)
 * `protobuf-py3`_ (when using Python 3)
 
+Testing (Linux/OSX)
+-------
+
+Testing is done with `tox`_::
+
+    tox
+
+.. _tox: https://tox.readthedocs.org/en/latest/
+
 Changelog
 ---------
 
+Version 6.1.1
+^^^^^^^^^^^^^
+
+* Fixed socket error handling in ``riemann_client.client.AutoFlushingQueuedClient``.
+
+Version 6.1.0
+^^^^^^^^^^^^^
+
+* ``riemann_client.client.AutoFlushingQueuedClient`` added.
+
 Version 6.0.0
 ^^^^^^^^^^^^^
 

riemann_client/__init__.py

@@ -1,4 +1,4 @@
 """A Python Riemann client and command line tool"""
 
-__version__ = '6.0.2'
+__version__ = '6.1.1'
 __author__ = 'Sam Clements <sam.clements@datasift.com>'

riemann_client/client.py

@@ -6,12 +6,32 @@
 
 from __future__ import absolute_import
 
+import logging
+try:
+    from logging import NullHandler
+except ImportError:
+    # Create a NullHandler class in logging for python 2.6
+    class NullHandler(logging.Handler):
+        def emit(self, record):
+            pass
+
 import socket
+try:
+    from threading import RLock
+    from threading import Timer
+except ImportError:
+    RLock = None
+    Timer = None
+import time
 
 import riemann_client.riemann_pb2
 import riemann_client.transport
 
 
+logger = logging.getLogger(__name__)
+logger.addHandler(NullHandler())
+
+
 class Client(object):
     """A client for sending events and querying a Riemann server.
 
@@ -200,3 +220,147 @@ def send_events(self, events):
     def clear_queue(self):
         """Resets the message/queue to a blank :py:class:`.Msg` object"""
         self.queue = riemann_client.riemann_pb2.Msg()
+
+
+if RLock and Timer:  # noqa
+    class AutoFlushingQueuedClient(QueuedClient):
+        """A Riemann client using a queue and a timer that will automatically
+        flush its contents if either:
+            - the queue size exceeds :param max_batch_size: or
+            - more than :param max_delay: has elapsed since the last flush and
+             the queue is non-empty.
+
+        if :param stay_connected: is False, then the transport will be
+        disconnected after each flush and reconnected at the beginning of
+        the next flush.
+        if :param clear_on_fail: is True, then the client will discard its
+        buffer after the second retry in the event of a socket error.
+
+        A message object is used as a queue, and the following methods are
+        given:
+            - :py:meth:`.send_event` - add a new event to the queue
+            - :py:meth:`.send_events` add a tuple of new events to the queue
+            - :py:meth:`.event` - add a new event to the queue from
+               keyword arguments
+            - :py:meth:`.events` - add new events to the queue from
+               dictionaries
+            - :py:meth:`.flush` - manually force flush the queue to the
+               transport
+        """
+
+        def __init__(self, transport, max_delay=0.5, max_batch_size=100,
+                     stay_connected=False, clear_on_fail=False):
+            super(AutoFlushingQueuedClient, self).__init__(transport)
+            self.stay_connected = stay_connected
+            self.clear_on_fail = clear_on_fail
+            self.max_delay = max_delay
+            self.max_batch_size = max_batch_size
+            self.lock = RLock()
+            self.event_counter = 0
+            self.last_flush = time.time()
+            self.timer = None
+
+            # start the timer
+            self.start_timer()
+
+        def connect(self):
+            """Connect the transport if it is not already connected."""
+            if not self.is_connected():
+                self.transport.connect()
+
+        def is_connected(self):
+            """Check whether the transport is connected."""
+            try:
+                # this will throw an exception whenever socket isn't connected
+                self.transport.socket.type
+                return True
+            except (AttributeError, RuntimeError, socket.error):
+                return False
+
+        def event(self, **data):
+            """Enqueues an event, using keyword arguments to create an Event
+
+            >>> client.event(service='riemann-client', state='awesome')
+
+            :param \*\*data: keyword arguments used for :py:func:`create_event`
+            """
+            self.send_events((self.create_event(data),))
+
+        def events(self, *events):
+            """Enqueues multiple events in a single message
+
+            >>> client.events({'service': 'riemann-client',
+            >>>                'state': 'awesome'})
+
+             :param \*events: event dictionaries for :py:func:`create_event`
+             :returns: The response message from Riemann
+            """
+            self.send_events(self.create_event(evd) for evd in events)
+
+        def send_events(self, events):
+            """Enqueues multiple events
+
+            :param events: A list or iterable of ``Event`` objects
+            :returns: The response message from Riemann
+            """
+            with self.lock:
+                for event in events:
+                    self.queue.events.add().MergeFrom(event)
+                    self.event_counter += 1
+                    self.check_for_flush()
+
+        def flush(self):
+            """Sends the events in the queue to Riemann in a single protobuf
+            message
+
+            :returns: The response message from Riemann
+            """
+            response = None
+            with self.lock:
+                if not self.is_connected():
+                    self.connect()
+                try:
+                    response = super(AutoFlushingQueuedClient, self).flush()
+                except socket.error:
+                    # log and retry
+                    logger.warn("Socket error on flushing. "
+                                "Attempting reconnect and retry...")
+                    try:
+                        self.transport.disconnect()
+                        self.connect()
+                        response = (
+                            super(AutoFlushingQueuedClient, self).flush())
+                    except:
+                        logger.warn("Socket error on flushing "
+                                    "second attempt. Batch discarded.")
+                        self.transport.disconnect()
+                        if self.clear_on_fail:
+                            self.clear_queue()
+                self.event_counter = 0
+                if not self.stay_connected:
+                    self.transport.disconnect()
+                self.last_flush = time.time()
+            self.start_timer()
+            return response
+
+        def check_for_flush(self):
+            """Checks the conditions for flushing the queue"""
+            if (self.event_counter >= self.max_batch_size or
+                    (time.time() - self.last_flush) >= self.max_delay):
+                self.flush()
+
+        def start_timer(self):
+            """Cycle the timer responsible for periodically flushing the queue
+            """
+            if self.timer:
+                self.timer.cancel()
+            self.timer = Timer(self.max_delay, self.check_for_flush)
+            self.timer.daemon = True
+            self.timer.start()
+
+        def stop_timer(self):
+            """Stops the current timer
+
+            a :py:meth:`.flush` event will reactviate the timer
+            """
+            self.timer.cancel()

riemann_client/transport.py

@@ -169,27 +169,34 @@ def connect(self):
 
 
 class BlankTransport(Transport):
-    """A transport that collects messages in a list, and has no connection
+    """A transport that collects events in a list, and has no connection
 
     Used by ``--transport none``, which is useful for testing commands without
     contacting a Riemann server. This is also used by the automated tests in
     ``riemann_client/tests/test_riemann_command.py``.
     """
 
+    def __init__(self, *args, **kwargs):
+        self.events = []
+
     def connect(self):
         """Creates a list to hold messages"""
-        self.messages = []
+        pass
 
     def send(self, message):
         """Adds a message to the list, returning a fake 'ok' response
 
         :returns: A response message with ``ok = True``
         """
-        self.messages.append(message)
+        for event in message.events:
+            self.events.append(event)
         reply = riemann_client.riemann_pb2.Msg()
         reply.ok = True
         return reply
 
     def disconnect(self):
         """Clears the list of messages"""
-        del self.messages
+        pass
+
+    def __len__(self):
+        return len(self.events)

setup.py

@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name='riemann-client',
-    version='6.0.2',
+    version='6.1.1',
 
     author="Sam Clements",
     author_email="sam.clements@datasift.com",
@@ -23,7 +23,6 @@
 
     packages=[
         'riemann_client',
-        'riemann_client.tests'
     ],
 
     install_requires=[