Add pending request backpressure for subscriber. (googleapis#4892)

Jon Wayne Parrott · web-flow · commit bae0f754f65a · 2018-02-16T11:29:33.000-08:00
This change includes the consumer's pending request backlog in the policy's load calculations. This allows the policy to pause (and resume) the response stream if there are a large number of outstanding requests to be send on the stream. This additionally adds `max_requests` to `FlowControl` to allow tweaking this number. Resolves: googleapis#4792 Related: googleapis#4841
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_consumer.py b/pubsub/google/cloud/pubsub_v1/subscriber/_consumer.py
@@ -187,6 +187,10 @@ def __init__(self):
     def active(self):
         """bool: Indicates if the consumer is active.
 
+        *Active* means that the stream is open and that it is possible to
+        send and receive messages. This is distinct from *pausing* which just
+        pauses *response* consumption.
+
         This is intended to be an implementation independent way of indicating
         that the consumer is stopped. (E.g. so a policy that owns a consumer
         doesn't need to know what a ``threading.Event`` is.)
@@ -202,6 +206,14 @@ def send_request(self, request):
         with self._put_lock:
             self._request_queue.put(request)
 
+    @property
+    def pending_requests(self):
+        """int: An approximate count of the outstanding requests.
+
+        This can be used to determine if the consumer should be paused if there
+        are too many outstanding requests."""
+        return self._request_queue.qsize()
+
     def _request_generator_thread(self, policy):
         """Generate requests for the stream.
 
@@ -231,8 +243,9 @@ def _request_generator_thread(self, policy):
                 _LOGGER.debug('Request generator signaled to stop.')
                 break
 
-            _LOGGER.debug('Sending request:\n%r', request)
+            _LOGGER.debug('Sending request on stream')
             yield request
+            policy.on_request(request)
 
     def _stop_request_generator(self, request_generator, response_generator):
         """Ensure a request generator is closed.
@@ -325,7 +338,7 @@ def _blocking_consume(self, policy):
             # checks to make sure we're not exiting before opening a new
             # stream.
             if self._stopped.is_set():
-                _LOGGER.debug('Event signalled consumer exit.')
+                _LOGGER.debug('Event signaled consumer exit.')
                 break
 
             request_generator = self._request_generator_thread(policy)
@@ -334,14 +347,14 @@ def _blocking_consume(self, policy):
                 response_generator, self._can_consume)
             try:
                 for response in responses:
-                    _LOGGER.debug('Received response:\n%r', response)
+                    _LOGGER.debug('Received response on stream')
                     policy.on_response(response)
 
                 # If the loop above exits without an exception, then the
                 # request stream terminated cleanly, which should only happen
                 # when it was signaled to do so by stop_consuming. In this
                 # case, break out of the while loop and exit this thread.
-                _LOGGER.debug('Clean RPC loop exit signalled consumer exit.')
+                _LOGGER.debug('Clean RPC loop exit signaled consumer exit.')
                 break
             except Exception as exc:
                 recover = policy.on_exception(exc)
@@ -364,7 +377,8 @@ def pause(self):
 
         This will clear the ``_can_consume`` event which is checked
         every time :meth:`_blocking_consume` consumes a response from the
-        bidirectional streaming pull.
+        bidirectional streaming pull. *requests* can still be sent along
+        the stream.
 
         Complement to :meth:`resume`.
         """
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/policy/base.py b/pubsub/google/cloud/pubsub_v1/subscriber/policy/base.py
@@ -156,8 +156,21 @@ def _load(self):
         return max([
             len(self.managed_ack_ids) / self.flow_control.max_messages,
             self._bytes / self.flow_control.max_bytes,
+            self._consumer.pending_requests / self.flow_control.max_requests
         ])
 
+    def _maybe_resume_consumer(self):
+        """Check the current load and resume the consumer if needed."""
+        # If we have been paused by flow control, check and see if we are
+        # back within our limits.
+        #
+        # In order to not thrash too much, require us to have passed below
+        # the resume threshold (80% by default) of each flow control setting
+        # before restarting.
+        if (self._consumer.paused and
+                self._load < self.flow_control.resume_threshold):
+            self._consumer.resume()
+
     def ack(self, ack_id, time_to_ack=None, byte_size=None):
         """Acknowledge the message corresponding to the given ack_id.
 
@@ -216,15 +229,7 @@ def drop(self, ack_id, byte_size):
                     'Bytes was unexpectedly negative: %d', self._bytes)
                 self._bytes = 0
 
-        # If we have been paused by flow control, check and see if we are
-        # back within our limits.
-        #
-        # In order to not thrash too much, require us to have passed below
-        # the resume threshold (80% by default) of each flow control setting
-        # before restarting.
-        if (self._consumer.paused and
-                self._load < self.flow_control.resume_threshold):
-            self._consumer.resume()
+        self._maybe_resume_consumer()
 
     def get_initial_request(self, ack_queue=False):
         """Return the initial request.
@@ -397,6 +402,18 @@ def on_exception(self, exception):
         """
         raise NotImplementedError
 
+    def on_request(self, request):
+        """Called whenever a request has been sent to gRPC.
+
+        This allows the policy to measure the rate of requests sent along the
+        stream and apply backpressure by pausing or resuming the consumer
+        if needed.
+
+        Args:
+            request (Any): The protobuf request that was sent to gRPC.
+        """
+        self._maybe_resume_consumer()
+
     @abc.abstractmethod
     def on_response(self, response):
         """Process a response from gRPC.
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/policy/thread.py b/pubsub/google/cloud/pubsub_v1/subscriber/policy/thread.py
@@ -19,7 +19,6 @@
 import sys
 import threading
 
-import grpc
 from six.moves import queue as queue_mod
 
 from google.cloud.pubsub_v1 import types
@@ -33,19 +32,6 @@
 _CALLBACK_WORKER_NAME = 'Thread-Consumer-CallbackRequestsWorker'
 
 
-def _callback_completed(future):
-    """Simple callback that just logs a future's result.
-
-    Used on completion of processing a message received by a
-    subscriber.
-
-    Args:
-        future (concurrent.futures.Future): A future returned
-            from :meth:`~concurrent.futures.Executor.submit`.
-    """
-    _LOGGER.debug('Result: %s', future.result())
-
-
 def _do_nothing_callback(message):
     """Default callback for messages received by subscriber.
 
@@ -332,8 +318,7 @@ def on_response(self, response):
         """
         for msg in response.received_messages:
             _LOGGER.debug(
-                'Using %s to process new message received:\n%r',
-                self._callback, msg)
+                'Using %s to process message with ack_id %s.',
+                self._callback, msg.ack_id)
             message = Message(msg.message, msg.ack_id, self._request_queue)
-            future = self._executor.submit(self._callback, message)
-            future.add_done_callback(_callback_completed)
+            self._executor.submit(self._callback, message)
diff --git a/pubsub/google/cloud/pubsub_v1/types.py b/pubsub/google/cloud/pubsub_v1/types.py
@@ -53,12 +53,13 @@
 # The defaults should be fine for most use cases.
 FlowControl = collections.namedtuple(
     'FlowControl',
-    ['max_bytes', 'max_messages', 'resume_threshold'],
+    ['max_bytes', 'max_messages', 'resume_threshold', 'max_requests'],
 )
 FlowControl.__new__.__defaults__ = (
     psutil.virtual_memory().total * 0.2,  # max_bytes: 20% of total RAM
     float('inf'),                         # max_messages: no limit
     0.8,                                  # resume_threshold: 80%
+    100,                                  # max_requests: 100
 )
 
 
diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_consumer.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_consumer.py
@@ -81,7 +81,7 @@ def test_blocking_consume_when_exiting(_LOGGER):
     # Make sure method cleanly exits.
     assert consumer._blocking_consume(None) is None
 
-    _LOGGER.debug.assert_called_once_with('Event signalled consumer exit.')
+    _LOGGER.debug.assert_called_once_with('Event signaled consumer exit.')
 
 
 class OnException(object):
diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_base.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_base.py
@@ -169,7 +169,30 @@ def test_drop_below_threshold():
     assert consumer.paused is False
 
 
-def test_load():
+def test_on_request_below_threshold():
+    """Establish that we resume a paused subscription when the pending
+    requests count is below threshold."""
+    flow_control = types.FlowControl(max_requests=100)
+    policy = create_policy(flow_control=flow_control)
+    consumer = policy._consumer
+
+    assert consumer.paused is True
+
+    pending_requests_patch = mock.patch.object(
+        consumer.__class__, 'pending_requests', new_callable=mock.PropertyMock)
+    with pending_requests_patch as pending_requests:
+        # should still be paused, not under the threshold.
+        pending_requests.return_value = 90
+        policy.on_request(None)
+        assert consumer.paused is True
+
+        # should unpause, we're under the resume threshold
+        pending_requests.return_value = 50
+        policy.on_request(None)
+        assert consumer.paused is False
+
+
+def test_load_w_lease():
     flow_control = types.FlowControl(max_messages=10, max_bytes=1000)
     policy = create_policy(flow_control=flow_control)
     consumer = policy._consumer
@@ -191,6 +214,26 @@ def test_load():
         pause.assert_called_once_with()
 
 
+def test_load_w_requests():
+    flow_control = types.FlowControl(max_bytes=100, max_requests=100)
+    policy = create_policy(flow_control=flow_control)
+    consumer = policy._consumer
+
+    pending_requests_patch = mock.patch.object(
+        consumer.__class__, 'pending_requests', new_callable=mock.PropertyMock)
+    with pending_requests_patch as pending_requests:
+        pending_requests.return_value = 0
+        assert policy._load == 0
+
+        pending_requests.return_value = 100
+        print(consumer.pending_requests)
+        assert policy._load == 1
+
+        # If bytes count is higher, it should return that.
+        policy._bytes = 110
+        assert policy._load == 1.1
+
+
 def test_modify_ack_deadline():
     policy = create_policy()
     with mock.patch.object(policy._consumer, 'send_request') as send_request:
diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_thread.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_thread.py
@@ -25,7 +25,6 @@
 
 from google.cloud.pubsub_v1 import subscriber
 from google.cloud.pubsub_v1 import types
-from google.cloud.pubsub_v1.subscriber import _helper_threads
 from google.cloud.pubsub_v1.subscriber import message
 from google.cloud.pubsub_v1.subscriber.futures import Future
 from google.cloud.pubsub_v1.subscriber.policy import thread
@@ -229,16 +228,3 @@ def test_on_response():
     for call in submit_calls:
         assert call[1][0] == callback
         assert isinstance(call[1][1], message.Message)
-
-    add_done_callback_calls = [
-        m for m in future.method_calls if m[0] == 'add_done_callback']
-    assert len(add_done_callback_calls) == 2
-    for call in add_done_callback_calls:
-        assert call[1][0] == thread._callback_completed
-
-
-def test__callback_completed():
-    future = mock.Mock()
-    thread._callback_completed(future)
-    result_calls = [m for m in future.method_calls if m[0] == 'result']
-    assert len(result_calls) == 1

Original file line number	Diff line number	Diff line change
`@@ -53,12 +53,13 @@`
`53`	`53`	`# The defaults should be fine for most use cases.`
`54`	`54`	`FlowControl = collections.namedtuple(`
`55`	`55`	`'FlowControl',`
`56`		`- ['max_bytes', 'max_messages', 'resume_threshold'],`
	`56`	`+ ['max_bytes', 'max_messages', 'resume_threshold', 'max_requests'],`
`57`	`57`	`)`
`58`	`58`	`FlowControl.__new__.__defaults__ = (`
`59`	`59`	`psutil.virtual_memory().total * 0.2, # max_bytes: 20% of total RAM`
`60`	`60`	`float('inf'), # max_messages: no limit`
`61`	`61`	`0.8, # resume_threshold: 80%`
	`62`	`+ 100, # max_requests: 100`
`62`	`63`	`)`
`63`	`64`
`64`	`65`