feat: add flow control for message publishing

google/cloud/pubsub_v1/publisher/client.py

@@ -34,6 +34,7 @@
 from google.cloud.pubsub_v1.publisher._batch import thread
 from google.cloud.pubsub_v1.publisher._sequencer import ordered_sequencer
 from google.cloud.pubsub_v1.publisher._sequencer import unordered_sequencer
+from google.cloud.pubsub_v1.publisher.flow_controller import FlowController
 
 __version__ = pkg_resources.get_distribution("google-cloud-pubsub").version
 
@@ -93,7 +94,11 @@ class Client(object):
 
             # Optional
             publisher_options = pubsub_v1.types.PublisherOptions(
-                enable_message_ordering=False
+                enable_message_ordering=False,
+                flow_control=pubsub_v1.types.PublishFlowControl(
+                    message_limit=2000,
+                    limit_exceeded_behavior=pubsub_v1.types.LimitExceededBehavior.BLOCK,
+                ),
             ),
 
             # Optional
@@ -198,6 +203,9 @@ def __init__(self, batch_settings=(), publisher_options=(), **kwargs):
         # Thread created to commit all sequencers after a timeout.
         self._commit_thread = None
 
+        # The object controlling the message publishing flow
+        self._flow_controller = FlowController(self.publisher_options.flow_control)
+
     @classmethod
     def from_service_account_file(cls, filename, batch_settings=(), **kwargs):
         """Creates an instance of this client using the provided credentials
@@ -333,6 +341,11 @@ def publish(self, topic, data, ordering_key="", **attrs):
 
             pubsub_v1.publisher.exceptions.MessageTooLargeError: If publishing
                 the ``message`` would exceed the max size limit on the backend.
+
+            :exception:`~pubsub_v1.publisher.exceptions.FlowControlLimitError`:
+                If publishing a new message would exceed the publish flow control
+                limits and the desired action on overflow is
+                :attr:`~google.cloud.pubsub_v1.types.LimitExceededBehavior.ERROR`.
         """
         # Sanity check: Is the data being sent as a bytestring?
         # If it is literally anything else, complain loudly about it.
@@ -364,6 +377,13 @@ def publish(self, topic, data, ordering_key="", **attrs):
             data=data, ordering_key=ordering_key, attributes=attrs
         )
 
+        # Messages should go through flow control to prevent excessive
+        # queuing on the client side (depending on the settings).
+        self._flow_controller.add(message)
+
+        def on_publish_done(future):
+            self._flow_controller.release(message)
+
         with self._batch_lock:
             if self._is_stopped:
                 raise RuntimeError("Cannot publish on a stopped publisher.")
@@ -372,6 +392,7 @@ def publish(self, topic, data, ordering_key="", **attrs):
 
             # Delegate the publishing to the sequencer.
             future = sequencer.publish(message)
+            future.add_done_callback(on_publish_done)
 
             # Create a timer thread if necessary to enforce the batching
             # timeout.

google/cloud/pubsub_v1/publisher/exceptions.py

@@ -38,7 +38,12 @@ def __init__(self, ordering_key):
         super(PublishToPausedOrderingKeyException, self).__init__()
 
 
+class FlowControlLimitError(Exception):
+    """An action resulted in exceeding the flow control limits."""
+
+
 __all__ = (
+    "FlowControlLimitError",
     "MessageTooLargeError",
     "PublishError",
     "TimeoutError",

google/cloud/pubsub_v1/publisher/flow_controller.py

@@ -0,0 +1,145 @@
+# Copyright 2020, Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import threading
+import warnings
+
+from google.cloud.pubsub_v1 import types
+from google.cloud.pubsub_v1.publisher import exceptions
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class FlowController(object):
+    """A class used to control the flow of messages passing through it.
+
+    Args:
+        settings (~google.cloud.pubsub_v1.types.PublishFlowControl):
+            Desired flow control configuration.
+    """
+
+    def __init__(self, settings):
+        self._settings = settings
+
+        self._message_count = 0
+        self._total_bytes = 0
+
+        # The lock is used to protect the internal state (message and byte count).
+        self._operational_lock = threading.Lock()
+
+        # The condition for blocking the flow if capacity is exceeded.
+        self._has_capacity = threading.Condition(lock=self._operational_lock)
+
+    def add(self, message):
+        """Add a message to flow control.
+
+        Adding a message updates the internal load statistics, and an action is
+        taken if these limits are exceeded (depending on the flow control settings).
+
+        Args:
+            message (:class:`~google.cloud.pubsub_v1.types.PubsubMessage`):
+                The message entering the flow control.
+
+        Raises:
+            :exception:`~pubsub_v1.publisher.exceptions.FlowControlLimitError`:
+                If adding a message exceeds flow control limits and the desired
+                action is :attr:`~google.cloud.pubsub_v1.types.LimitExceededBehavior.ERROR`.
+        """
+        if self._settings.limit_exceeded_behavior == types.LimitExceededBehavior.IGNORE:
+            return
+
+        with self._operational_lock:
+            self._message_count += 1
+            self._total_bytes += message.ByteSize()
+
+            if not self._is_overflow():
+                return
+
+            # We have an overflow, react.
+            if (
+                self._settings.limit_exceeded_behavior
+                == types.LimitExceededBehavior.ERROR
+            ):
+                msg = (
+                    "Flow control limits exceeded "
+                    "(messages: {} / {}, bytes: {} / {})."
+                ).format(
+                    self._message_count,
+                    self._settings.message_limit,
+                    self._total_bytes,
+                    self._settings.byte_limit,
+                )
+                error = exceptions.FlowControlLimitError(msg)
+
+                # Raising an error means rejecting a message, thus we need to deduct
+                # the latter's contribution to the total load.
+                self._message_count -= 1
+                self._total_bytes -= message.ByteSize()
+                raise error
+
+            assert (
+                self._settings.limit_exceeded_behavior
+                == types.LimitExceededBehavior.BLOCK
+            )
+
+            while self._is_overflow():
+                _LOGGER.debug(
+                    "Blocking until there is enough free capacity in the flow."
+                )
+                self._has_capacity.wait()
+                _LOGGER.debug("Woke up from waiting on free capacity in the flow.")
+
+    def release(self, message):
+        """Release a mesage from flow control.
+
+        Args:
+            message (:class:`~google.cloud.pubsub_v1.types.PubsubMessage`):
+                The message entering the flow control.
+        """
+        if self._settings.limit_exceeded_behavior == types.LimitExceededBehavior.IGNORE:
+            return
+
+        with self._operational_lock:
+            was_overflow = self._is_overflow()
+
+            self._message_count -= 1
+            self._total_bytes -= message.ByteSize()
+
+            if self._message_count < 0 or self._total_bytes < 0:
+                warnings.warn(
+                    "Releasing a message that was never added or already released.",
+                    category=RuntimeWarning,
+                    stacklevel=2,
+                )
+                self._message_count = max(0, self._message_count)
+                self._total_bytes = max(0, self._total_bytes)
+
+            if was_overflow and not self._is_overflow():
+                _LOGGER.debug("Notifying threads waiting to add messages to flow.")
+                self._has_capacity.notify_all()
+
+    def _is_overflow(self):
+        """Determine if the current message load exceeds flow control limits.
+
+        The method assumes that the caller has obtained ``_operational_lock``.
+
+        Returns:
+            bool
+        """
+        return (
+            self._message_count > self._settings.message_limit
+            or self._total_bytes > self._settings.byte_limit
+        )

google/cloud/pubsub_v1/types.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+
 import collections
+import enum
 import sys
 
 from google.api import http_pb2
@@ -30,25 +32,6 @@
 from google.cloud.pubsub_v1.proto import pubsub_pb2
 
 
-# Define the default publisher options.
-#
-# This class is used when creating a publisher client to pass in options
-# to enable/disable features.
-PublisherOptions = collections.namedtuple(
-    "PublisherConfig", ["enable_message_ordering"]
-)
-PublisherOptions.__new__.__defaults__ = (False,)  # enable_message_ordering: False
-
-if sys.version_info >= (3, 5):
-    PublisherOptions.__doc__ = "The options for the publisher client."
-    PublisherOptions.enable_message_ordering.__doc__ = (
-        "Whether to order messages in a batch by a supplied ordering key."
-        "EXPERIMENTAL: Message ordering is an alpha feature that requires "
-        "special permissions to use. Please contact the Cloud Pub/Sub team for "
-        "more information."
-    )
-
-
 # Define the default values for batching.
 #
 # This class is used when creating a publisher or subscriber client, and
@@ -81,6 +64,63 @@
     )
 
 
+class LimitExceededBehavior(str, enum.Enum):
+    """The possible actions when exceeding the publish flow control limits."""
+
+    IGNORE = "ignore"
+    BLOCK = "block"
+    ERROR = "error"
+
+
+PublishFlowControl = collections.namedtuple(
+    "PublishFlowControl", ["message_limit", "byte_limit", "limit_exceeded_behavior"]
+)
+PublishFlowControl.__new__.__defaults__ = (
+    10 * BatchSettings.__new__.__defaults__[2],  # message limit
+    10 * BatchSettings.__new__.__defaults__[0],  # byte limit
+    LimitExceededBehavior.IGNORE,  # desired behavior
+)
+
+if sys.version_info >= (3, 5):
+    PublishFlowControl.__doc__ = (
+        "The client flow control settings for message publishing."
+    )
+    PublishFlowControl.message_limit.__doc__ = (
+        "The maximum number of messages awaiting to be published."
+    )
+    PublishFlowControl.byte_limit.__doc__ = (
+        "The maximum total size of messages awaiting to be published."
+    )
+    PublishFlowControl.limit_exceeded_behavior.__doc__ = (
+        "The action to take when publish flow control limits are exceeded."
+    )
+
+# Define the default publisher options.
+#
+# This class is used when creating a publisher client to pass in options
+# to enable/disable features.
+PublisherOptions = collections.namedtuple(
+    "PublisherConfig", ["enable_message_ordering", "flow_control"]
+)
+PublisherOptions.__new__.__defaults__ = (
+    False,  # enable_message_ordering: False
+    PublishFlowControl(),  # default flow control settings
+)
+
+if sys.version_info >= (3, 5):
+    PublisherOptions.__doc__ = "The options for the publisher client."
+    PublisherOptions.enable_message_ordering.__doc__ = (
+        "Whether to order messages in a batch by a supplied ordering key."
+        "EXPERIMENTAL: Message ordering is an alpha feature that requires "
+        "special permissions to use. Please contact the Cloud Pub/Sub team for "
+        "more information."
+    )
+    PublisherOptions.flow_control.__doc__ = (
+        "Flow control settings for message publishing by the client. By default "
+        "the publisher client does not do any throttling."
+    )
+
+
 # Define the type class and default values for flow control settings.
 #
 # This class is used when creating a publisher or subscriber client, and