Can has flow control.

pubsub/google/cloud/pubsub_v1/subscriber/consumer.py

@@ -200,10 +200,11 @@ def _request_generator_thread(self):
         """
         # First, yield the initial request. This occurs on every new
         # connection, fundamentally including a resumed connection.
+        initial_request = self._policy.get_initial_request(ack_queue=True)
         _LOGGER.debug('Sending initial request: {initial_request}'.format(
-            initial_request=self._policy.initial_request,
+            initial_request=initial_request,
         ))
-        yield self._policy.initial_request
+        yield initial_request
 
         # Now yield each of the items on the request queue, and block if there
         # are none. This can and must block to keep the stream open.

pubsub/google/cloud/pubsub_v1/subscriber/message.py

@@ -134,9 +134,9 @@ def ack(self):
         time_to_ack = math.ceil(time.time() - self._received_timestamp)
         self._request_queue.put(('ack', {
             'ack_id': self._ack_id,
+            'byte_size': self.size,
             'time_to_ack': time_to_ack,
         }))
-        self.drop()
 
     def drop(self):
         """Release the message from lease management.
@@ -196,5 +196,7 @@ def nack(self):
 
         This will cause the message to be re-delivered to the subscription.
         """
-        self.modify_ack_deadline(seconds=0)
-        self.drop()
+        self._request_queue.put(('nack', {
+            'ack_id': self._ack_id,
+            'byte_size': self.size,
+        }))

pubsub/google/cloud/pubsub_v1/subscriber/policy/base.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import absolute_import
+from __future__ import absolute_import, division
 
 import abc
 import logging
@@ -69,10 +69,15 @@ def __init__(self, client, subscription,
         self._consumer = consumer.Consumer(self)
         self._ack_deadline = 10
         self._last_histogram_size = 0
-        self._bytes = 0
         self.flow_control = flow_control
         self.histogram = histogram.Histogram(data=histogram_data)
 
+        # These are for internal flow control tracking.
+        # They should not need to be used by subclasses.
+        self._bytes = 0
+        self._ack_on_resume = set()
+        self._paused = False
+
     @property
     def ack_deadline(self):
         """Return the appropriate ack deadline.
@@ -92,18 +97,6 @@ def ack_deadline(self):
             self._ack_deadline = self.histogram.percentile(percent=99)
         return self._ack_deadline
 
-    @property
-    def initial_request(self):
-        """Return the initial request.
-
-        This defines the initial request that must always be sent to Pub/Sub
-        immediately upon opening the subscription.
-        """
-        return types.StreamingPullRequest(
-            stream_ack_deadline_seconds=self.histogram.percentile(99),
-            subscription=self.subscription,
-        )
-
     @property
     def managed_ack_ids(self):
         """Return the ack IDs currently being managed by the policy.
@@ -124,19 +117,54 @@ def subscription(self):
         """
         return self._subscription
 
-    def ack(self, ack_id, time_to_ack=None):
+    @property
+    def _load(self):
+        """Return the current load.
+
+        The load is represented as a float, where 1.0 represents having
+        hit one of the flow control limits, and values between 0.0 and 1.0
+        represent how close we are to them. (0.5 means we have exactly half
+        of what the flow control setting allows, for example.)
+
+        There are (currently) two flow control settings; this property
+        computes how close the subscriber is to each of them, and returns
+        whichever value is higher. (It does not matter that we have lots of
+        running room on setting A if setting B is over.)
+
+        Returns:
+            float: The load value.
+        """
+        return max([
+            len(self.managed_ack_ids) / self.flow_control.max_messages,
+            self._bytes / self.flow_control.max_bytes,
+        ])
+
+    def ack(self, ack_id, time_to_ack=None, byte_size=None):
         """Acknowledge the message corresponding to the given ack_id.
 
         Args:
             ack_id (str): The ack ID.
             time_to_ack (int): The time it took to ack the message, measured
                 from when it was received from the subscription. This is used
                 to improve the automatic ack timing.
+            byte_size (int): The size of the PubSub message, in bytes.
         """
+        # If we got timing information, add it to the histogram.
         if time_to_ack is not None:
             self.histogram.add(int(time_to_ack))
-        request = types.StreamingPullRequest(ack_ids=[ack_id])
-        self._consumer.send_request(request)
+
+        # Send the request to ack the message.
+        # However, if the consumer is inactive, then queue the ack_id here
+        # instead; it will be acked as part of the initial request when the
+        # consumer is started again.
+        if self._consumer.active:
+            request = types.StreamingPullRequest(ack_ids=[ack_id])
+            self._consumer.send_request(request)
+        else:
+            self._ack_on_resume.add(ack_id)
+
+        # Remove the message from lease management.
+        self.drop(ack_id=ack_id, byte_size=byte_size)
 
     def call_rpc(self, request_generator):
         """Invoke the Pub/Sub streaming pull RPC.
@@ -155,22 +183,89 @@ def drop(self, ack_id, byte_size):
             ack_id (str): The ack ID.
             byte_size (int): The size of the PubSub message, in bytes.
         """
+        # Remove the ack ID from lease management, and decrement the
+        # byte counter.
         if ack_id in self.managed_ack_ids:
             self.managed_ack_ids.remove(ack_id)
             self._bytes -= byte_size
             self._bytes = min([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._paused and self._load < self.flow_control.resume_threshold:
+            self._paused = False
+            self.open(self._callback)
+
+    def get_initial_request(self, ack_queue=False):
+        """Return the initial request.
+
+        This defines the initial request that must always be sent to Pub/Sub
+        immediately upon opening the subscription.
+
+        Args:
+            ack_queue (bool): Whether to include any acks that were sent
+                while the connection was paused.
+
+        Returns:
+            ~.pubsub_v1.types.StreamingPullRequest: A request suitable
+                for being the first request on the stream (and not suitable
+                for any other purpose).
+
+        .. note::
+            If ``ack_queue`` is set to True, this includes the ack_ids, but
+            also clears the internal set.
+
+            This means that calls to :meth:`get_initial_request` with
+            ``ack_queue`` set to True are not idempotent.
+        """
+        # Any ack IDs that are under lease management and not being acked
+        # need to have their deadline extended immediately.
+        ack_ids = set()
+        lease_ids = self.managed_ack_ids
+        if ack_queue:
+            ack_ids = self._ack_on_resume
+            lease_ids = lease_ids.difference(ack_ids)
+
+        # Put the request together.
+        request = types.StreamingPullRequest(
+            ack_ids=list(ack_ids),
+            modify_deadline_ack_ids=list(lease_ids),
+            modify_deadline_seconds=[self.ack_deadline] * len(lease_ids),
+            stream_ack_deadline_seconds=self.histogram.percentile(99),
+            subscription=self.subscription,
+        )
+
+        # Clear the ack_ids set.
+        # Note: If `ack_queue` is False, this just ends up being a no-op,
+        # since the set is just an empty set.
+        ack_ids.clear()
+
+        # Return the initial request.
+        return request
+
     def lease(self, ack_id, byte_size):
         """Add the given ack ID to lease management.
 
         Args:
             ack_id (str): The ack ID.
             byte_size (int): The size of the PubSub message, in bytes.
         """
+        # Add the ack ID to the set of managed ack IDs, and increment
+        # the size counter.
         if ack_id not in self.managed_ack_ids:
             self.managed_ack_ids.add(ack_id)
             self._bytes += byte_size
 
+        # Sanity check: Do we have too many things in our inventory?
+        # If we do, we need to stop the stream.
+        if self._load >= 1.0:
+            self._paused = True
+            self.close()
+
     def maintain_leases(self):
         """Maintain all of the leases being managed by the policy.
 
@@ -202,7 +297,7 @@ def maintain_leases(self):
         # it is more efficient to make a single request.
         ack_ids = list(self.managed_ack_ids)
         logger.debug('Renewing lease for %d ack IDs.' % len(ack_ids))
-        if len(ack_ids) > 0:
+        if len(ack_ids) > 0 and self._consumer.active:
             request = types.StreamingPullRequest(
                 modify_deadline_ack_ids=ack_ids,
                 modify_deadline_seconds=[p99] * len(ack_ids),
@@ -233,13 +328,33 @@ def modify_ack_deadline(self, ack_id, seconds):
         )
         self._consumer.send_request(request)
 
-    def nack(self, ack_id):
+    def nack(self, ack_id, byte_size=None):
         """Explicitly deny receipt of a message.
 
         Args:
             ack_id (str): The ack ID.
+            byte_size (int): The size of the PubSub message, in bytes.
+        """
+        self.modify_ack_deadline(ack_id=ack_id, seconds=0)
+        self.drop(ack_id=ack_id, byte_size=byte_size)
+
+    @abc.abstractmethod
+    def close(self):
+        """Close the existing connection."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def on_exception(self, exception):
+        """Called when a gRPC exception occurs.
+
+        If this method does nothing, then the stream is re-started. If this
+        raises an exception, it will stop the consumer thread.
+        This is executed on the response consumer helper thread.
+
+        Args:
+            exception (Exception): The exception raised by the RPC.
         """
-        return self.modify_ack_deadline(ack_id=ack_id, seconds=0)
+        raise NotImplementedError
 
     @abc.abstractmethod
     def on_response(self, response):
@@ -262,19 +377,6 @@ def on_response(self, response):
         """
         raise NotImplementedError
 
-    @abc.abstractmethod
-    def on_exception(self, exception):
-        """Called when a gRPC exception occurs.
-
-        If this method does nothing, then the stream is re-started. If this
-        raises an exception, it will stop the consumer thread.
-        This is executed on the response consumer helper thread.
-
-        Args:
-            exception (Exception): The exception raised by the RPC.
-        """
-        raise NotImplementedError
-
     @abc.abstractmethod
     def open(self, callback):
         """Open a streaming pull connection and begin receiving messages.

pubsub/google/cloud/pubsub_v1/types.py

@@ -44,11 +44,12 @@
 # these settings can be altered to tweak Pub/Sub behavior.
 # The defaults should be fine for most use cases.
 FlowControl = collections.namedtuple('FlowControl',
-    ['max_bytes', 'max_messages'],
+    ['max_bytes', 'max_messages', 'resume_threshold'],
 )
 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%
 )
 
 

pubsub/tests/unit/pubsub_v1/subscriber/test_message.py

@@ -57,9 +57,9 @@ def test_ack():
             msg.ack()
             put.assert_called_once_with(('ack', {
                 'ack_id': 'bogus_ack_id',
+                'byte_size': 25,
                 'time_to_ack': mock.ANY,
             }))
-            drop.assert_called_once_with()
 
 
 def test_drop():
@@ -93,9 +93,10 @@ def test_modify_ack_deadline():
 
 
 def test_nack():
-    msg = create_message(b'foo')
-    with mock.patch.object(message.Message, 'modify_ack_deadline') as mad:
-        with mock.patch.object(message.Message, 'drop') as drop:
-            msg.nack()
-            mad.assert_called_once_with(seconds=0)
-            drop.assert_called_once_with()
+    msg = create_message(b'foo', ack_id='bogus_id')
+    with mock.patch.object(msg._request_queue, 'put') as put:
+        msg.nack()
+        put.assert_called_once_with(('nack', {
+            'ack_id': 'bogus_id',
+            'byte_size': 25,
+        }))