Add max_latency to BackgroundThreadTransport (#4762)

Author: tcwalther

google/cloud/logging/handlers/transports/background_thread.py

@@ -22,6 +22,7 @@
 import atexit
 import logging
 import threading
+import time
 
 from six.moves import range
 from six.moves import queue
@@ -30,12 +31,13 @@
 
 _DEFAULT_GRACE_PERIOD = 5.0  # Seconds
 _DEFAULT_MAX_BATCH_SIZE = 10
+_DEFAULT_MAX_LATENCY = 0  # Seconds
 _WORKER_THREAD_NAME = 'google.cloud.logging.Worker'
 _WORKER_TERMINATOR = object()
 _LOGGER = logging.getLogger(__name__)
 
 
-def _get_many(queue_, max_items=None):
+def _get_many(queue_, max_items=None, max_latency=0):
     """Get multiple items from a Queue.
 
     Gets at least one (blocking) and at most ``max_items`` items
@@ -48,14 +50,22 @@ def _get_many(queue_, max_items=None):
     :param max_items: The maximum number of items to get. If ``None``, then all
         available items in the queue are returned.
 
+    :type max_latency: float
+    :param max_latency: The maximum number of seconds to wait for more than one
+        item from a queue. This number includes the time required to retrieve
+        the first item.
+
     :rtype: Sequence
     :returns: A sequence of items retrieved from the queue.
     """
+    start = time.time()
     # Always return at least one item.
     items = [queue_.get()]
     while max_items is None or len(items) < max_items:
         try:
-            items.append(queue_.get_nowait())
+            elapsed = time.time() - start
+            timeout = max(0, max_latency - elapsed)
+            items.append(queue_.get(timeout=timeout))
         except queue.Empty:
             break
     return items
@@ -74,13 +84,22 @@ class _Worker(object):
     :type max_batch_size: int
     :param max_batch_size: The maximum number of items to send at a time
         in the background thread.
+
+    :type max_latency: float
+    :param max_latency: The amount of time to wait for new logs before
+        sending a new batch. It is strongly recommended to keep this smaller
+        than the grace_period. This means this is effectively the longest
+        amount of time the background thread will hold onto log entries
+        before sending them to the server.
     """
 
     def __init__(self, cloud_logger, grace_period=_DEFAULT_GRACE_PERIOD,
-                 max_batch_size=_DEFAULT_MAX_BATCH_SIZE):
+                 max_batch_size=_DEFAULT_MAX_BATCH_SIZE,
+                 max_latency=_DEFAULT_MAX_LATENCY):
         self._cloud_logger = cloud_logger
         self._grace_period = grace_period
         self._max_batch_size = max_batch_size
+        self._max_latency = max_latency
         self._queue = queue.Queue(0)
         self._operational_lock = threading.Lock()
         self._thread = None
@@ -112,7 +131,9 @@ def _thread_main(self):
         quit_ = False
         while True:
             batch = self._cloud_logger.batch()
-            items = _get_many(self._queue, max_items=self._max_batch_size)
+            items = _get_many(
+                self._queue, max_items=self._max_batch_size,
+                max_latency=self._max_latency)
 
             for item in items:
                 if item is _WORKER_TERMINATOR:
@@ -249,15 +270,24 @@ class BackgroundThreadTransport(Transport):
     :type batch_size: int
     :param batch_size: The maximum number of items to send at a time in the
         background thread.
+
+    :type max_latency: float
+    :param max_latency: The amount of time to wait for new logs before
+        sending a new batch. It is strongly recommended to keep this smaller
+        than the grace_period. This means this is effectively the longest
+        amount of time the background thread will hold onto log entries
+        before sending them to the server.
     """
 
     def __init__(self, client, name, grace_period=_DEFAULT_GRACE_PERIOD,
-                 batch_size=_DEFAULT_MAX_BATCH_SIZE):
+                 batch_size=_DEFAULT_MAX_BATCH_SIZE,
+                 max_latency=_DEFAULT_MAX_LATENCY):
         self.client = client
         logger = self.client.logger(name)
         self.worker = _Worker(logger,
                               grace_period=grace_period,
-                              max_batch_size=batch_size)
+                              max_batch_size=batch_size,
+                              max_latency=max_latency)
         self.worker.start()
 
     def send(self, record, message, resource=None, labels=None):

tests/unit/handlers/transports/test_background_thread.py

@@ -78,19 +78,24 @@ def test_flush(self):
 
     def test_worker(self):
         client = _Client(self.PROJECT)
-        name = 'python_logger'        
+        name = 'python_logger'
         batch_size = 30
         grace_period = 20.
+        max_latency = 0.1
         transport, worker = self._make_one(client,
                                            name,
                                            grace_period=grace_period,
-                                           batch_size=batch_size)
+                                           batch_size=batch_size,
+                                           max_latency=max_latency)
         worker_grace_period = worker.call_args[1]['grace_period']  # **kwargs.
         worker_batch_size = worker.call_args[1]['max_batch_size']
+        worker_max_latency = worker.call_args[1]['max_latency']
         self.assertEqual(worker_grace_period,
                          grace_period)
         self.assertEqual(worker_batch_size,
                          batch_size)
+        self.assertEqual(worker_max_latency,
+                         max_latency)
 
 
 class Test_Worker(unittest.TestCase):
@@ -115,13 +120,16 @@ def test_constructor(self):
         logger = _Logger(self.NAME)
         grace_period = 50
         max_batch_size = 50
+        max_latency = 0.1
 
         worker = self._make_one(
-            logger, grace_period=grace_period, max_batch_size=max_batch_size)
+            logger, grace_period=grace_period, max_batch_size=max_batch_size,
+            max_latency=max_latency)
 
         self.assertEqual(worker._cloud_logger, logger)
         self.assertEqual(worker._grace_period, grace_period)
         self.assertEqual(worker._max_batch_size, max_batch_size)
+        self.assertEqual(worker._max_latency, max_latency)
         self.assertFalse(worker.is_alive)
         self.assertIsNone(worker._thread)
 
@@ -264,6 +272,74 @@ def test__thread_main_batches(self):
         self.assertFalse(worker._cloud_logger._batch.commit_called)
         self.assertEqual(worker._queue.qsize(), 0)
 
+    @mock.patch('time.time', autospec=True, return_value=1)
+    def test__thread_main_max_latency(self, time):
+        # Note: this test is a bit brittle as it assumes the operation of
+        # _get_many invokes queue.get() followed by queue._get(). It fails
+        # the "change detector" test in that way. However, this is still a
+        # useful test to verify the queue timeout is appropriately calculated.
+        from six.moves import queue
+        from google.cloud.logging.handlers.transports import background_thread
+
+        # Use monotonically increasing time.
+        time.side_effect = range(1, 6)
+
+        worker = self._make_one(
+            _Logger(self.NAME), max_latency=2, max_batch_size=10)
+        worker._queue = mock.create_autospec(queue.Queue, instance=True)
+
+        worker._queue.get.side_effect = [
+            {'info': {'message': '1'}},  # Single record.
+            queue.Empty(),  # Emulate a queue.get() timeout.
+            {'info': {'message': '1'}},  # Second record.
+            background_thread._WORKER_TERMINATOR,  # Stop the thread.
+            queue.Empty(),  # Emulate a queue.get() timeout.
+        ]
+
+        worker._thread_main()
+
+        self.assertEqual(worker._cloud_logger._num_batches, 2)
+        self.assertTrue(worker._cloud_logger._batch.commit_called)
+        self.assertEqual(worker._cloud_logger._batch.commit_count, 1)
+
+        # Time should have been called five times.
+        #
+        #   For the first batch, it should have been called:
+        #       * Once to get the start time. (1)
+        #       * Once to get the elapsed time while grabbing the second item.
+        #         (2)
+        #
+        #   For the second batch, it should have been called:
+        #       * Once to get start time. (3)
+        #       * Once to get the elapsed time while grabbing the second item.
+        #         (3)
+        #       * Once to get the elapsed time while grabbing the final
+        #         item. (4)
+        #       * Once final time to get the elapsed time while receiving
+        #         the empty queue.
+        #
+        self.assertEqual(time.call_count, 5)
+
+        # Queue.get should've been called 5 times as well, but with different
+        # timeouts due to the monotonically increasing time.
+        #
+        #   For the first batch, it will be called once without a timeout
+        #   (for the first item) and then with timeout=1, as start will be
+        #   1 and now will be 2.
+        #
+        #   For the second batch, it will be called once without a timeout
+        #   (for the first item) and then with timeout=1, as start will be
+        #   3 and now will be 4, and finally with timeout=0 as start will be 3
+        #   and now will be 5.
+        #
+        worker._queue.get.assert_has_calls([
+            mock.call(),
+            mock.call(timeout=1),
+            mock.call(),
+            mock.call(timeout=1),
+            mock.call(timeout=0)
+        ])
+
     def test_flush(self):
         worker = self._make_one(_Logger(self.NAME))
         worker._queue = mock.Mock(spec=queue.Queue)
@@ -331,9 +407,11 @@ def __init__(self, name):
         self.name = name
         self._batch_cls = _Batch
         self._batch = None
+        self._num_batches = 0
 
     def batch(self):
         self._batch = self._batch_cls()
+        self._num_batches += 1
         return self._batch