feat: add flow control for message publishing (#96)

* feat: add publish flow control settings

* Add flow control logic to publisher client

* Add flow control support for multiple add() threads

* Raise publish flow control errors through futures

* Include load info in debug log messages

* Remove incorrect comment in a test

* Remove comment about an error not directly raised

* Remove redundant check for reservation exsistence

* Change exception for publishing too large a message

* Add internal sanity check for byte reservations

* Reword the docstring on flow control limits error

google/cloud/pubsub_v1/publisher/client.py

@@ -31,9 +31,12 @@
 from google.cloud.pubsub_v1 import types
 from google.cloud.pubsub_v1.gapic import publisher_client
 from google.cloud.pubsub_v1.gapic.transports import publisher_grpc_transport
+from google.cloud.pubsub_v1.publisher import exceptions
+from google.cloud.pubsub_v1.publisher import futures
 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 +96,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 +205,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
@@ -364,6 +374,18 @@ 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).
+        try:
+            self._flow_controller.add(message)
+        except exceptions.FlowControlLimitError as exc:
+            future = futures.Future()
+            future.set_exception(exc)
+            return future
+
+        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 +394,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,297 @@
+# 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.
+
+from collections import deque
+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 _QuantityReservation(object):
+    """A (partial) reservation of a quantifiable resource."""
+
+    def __init__(self, reserved, needed):
+        self.reserved = reserved
+        self.needed = needed
+
+
+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
+
+        # Load statistics. They represent the number of messages added, but not
+        # yet released (and their total size).
+        self._message_count = 0
+        self._total_bytes = 0
+
+        # A FIFO queue of threads blocked on adding a message, from first to last.
+        # Only relevant if the configured limit exceeded behavior is BLOCK.
+        self._waiting = deque()
+
+        # Reservations of available flow control bytes by the waiting threads.
+        # Each value is a _QuantityReservation instance.
+        self._byte_reservations = dict()
+        self._reserved_bytes = 0
+
+        # The lock is used to protect all internal state (message and byte count,
+        # waiting threads to add, etc.).
+        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`:
+                Raised when the desired action is
+                :attr:`~google.cloud.pubsub_v1.types.LimitExceededBehavior.ERROR` and
+                the message would exceed flow control limits, or when the desired action
+                is :attr:`~google.cloud.pubsub_v1.types.LimitExceededBehavior.BLOCK` and
+                the message would block forever against the flow control limits.
+        """
+        if self._settings.limit_exceeded_behavior == types.LimitExceededBehavior.IGNORE:
+            return
+
+        with self._operational_lock:
+            if not self._would_overflow(message):
+                self._message_count += 1
+                self._total_bytes += message.ByteSize()
+                return
+
+            # Adding a message would overflow, react.
+            if (
+                self._settings.limit_exceeded_behavior
+                == types.LimitExceededBehavior.ERROR
+            ):
+                # Raising an error means rejecting a message, thus we do not
+                # add anything to the existing load, but we do report the would-be
+                # load if we accepted the message.
+                load_info = self._load_info(
+                    message_count=self._message_count + 1,
+                    total_bytes=self._total_bytes + message.ByteSize(),
+                )
+                error_msg = "Flow control limits would be exceeded - {}.".format(
+                    load_info
+                )
+                raise exceptions.FlowControlLimitError(error_msg)
+
+            assert (
+                self._settings.limit_exceeded_behavior
+                == types.LimitExceededBehavior.BLOCK
+            )
+
+            # Sanity check - if a message exceeds total flow control limits all
+            # by itself, it would block forever, thus raise error.
+            if (
+                message.ByteSize() > self._settings.byte_limit
+                or self._settings.message_limit < 1
+            ):
+                load_info = self._load_info(
+                    message_count=1, total_bytes=message.ByteSize()
+                )
+                error_msg = (
+                    "Total flow control limits too low for the message, "
+                    "would block forever - {}.".format(load_info)
+                )
+                raise exceptions.FlowControlLimitError(error_msg)
+
+            current_thread = threading.current_thread()
+
+            while self._would_overflow(message):
+                if current_thread not in self._byte_reservations:
+                    self._waiting.append(current_thread)
+                    self._byte_reservations[current_thread] = _QuantityReservation(
+                        reserved=0, needed=message.ByteSize()
+                    )
+
+                _LOGGER.debug(
+                    "Blocking until there is enough free capacity in the flow - "
+                    "{}.".format(self._load_info())
+                )
+
+                self._has_capacity.wait()
+
+                _LOGGER.debug(
+                    "Woke up from waiting on free capacity in the flow - "
+                    "{}.".format(self._load_info())
+                )
+
+            # Message accepted, increase the load and remove thread stats.
+            self._message_count += 1
+            self._total_bytes += message.ByteSize()
+            self._reserved_bytes -= self._byte_reservations[current_thread].reserved
+            del self._byte_reservations[current_thread]
+            self._waiting.remove(current_thread)
+
+    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:
+            # Releasing a message decreases the load.
+            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)
+
+            self._distribute_available_bytes()
+
+            # If at least one thread waiting to add() can be unblocked, wake them up.
+            if self._ready_to_unblock():
+                _LOGGER.debug("Notifying threads waiting to add messages to flow.")
+                self._has_capacity.notify_all()
+
+    def _distribute_available_bytes(self):
+        """Distribute availalbe free capacity among the waiting threads in FIFO order.
+
+        The method assumes that the caller has obtained ``_operational_lock``.
+        """
+        available = self._settings.byte_limit - self._total_bytes - self._reserved_bytes
+
+        for thread in self._waiting:
+            if available <= 0:
+                break
+
+            reservation = self._byte_reservations[thread]
+            still_needed = reservation.needed - reservation.reserved
+
+            # Sanity check for any internal inconsistencies.
+            if still_needed < 0:
+                msg = "Too many bytes reserved: {} / {}".format(
+                    reservation.reserved, reservation.needed
+                )
+                warnings.warn(msg, category=RuntimeWarning)
+                still_needed = 0
+
+            can_give = min(still_needed, available)
+            reservation.reserved += can_give
+            self._reserved_bytes += can_give
+            available -= can_give
+
+    def _ready_to_unblock(self):
+        """Determine if any of the threads waiting to add a message can proceed.
+
+        The method assumes that the caller has obtained ``_operational_lock``.
+
+        Returns:
+            bool
+        """
+        if self._waiting:
+            # It's enough to only check the head of the queue, because FIFO
+            # distribution of any free capacity.
+            reservation = self._byte_reservations[self._waiting[0]]
+            return (
+                reservation.reserved >= reservation.needed
+                and self._message_count < self._settings.message_limit
+            )
+
+        return False
+
+    def _would_overflow(self, message):
+        """Determine if accepting a message would exceed flow control limits.
+
+        The method assumes that the caller has obtained ``_operational_lock``.
+
+        Args:
+            message (:class:`~google.cloud.pubsub_v1.types.PubsubMessage`):
+                The message entering the flow control.
+
+        Returns:
+            bool
+        """
+        reservation = self._byte_reservations.get(threading.current_thread())
+
+        if reservation:
+            enough_reserved = reservation.reserved >= reservation.needed
+        else:
+            enough_reserved = False
+
+        bytes_taken = self._total_bytes + self._reserved_bytes + message.ByteSize()
+        size_overflow = bytes_taken > self._settings.byte_limit and not enough_reserved
+        msg_count_overflow = self._message_count + 1 > self._settings.message_limit
+
+        return size_overflow or msg_count_overflow
+
+    def _load_info(self, message_count=None, total_bytes=None, reserved_bytes=None):
+        """Return the current flow control load information.
+
+        The caller can optionally adjust some of the values to fit its reporting
+        needs.
+
+        The method assumes that the caller has obtained ``_operational_lock``.
+
+        Args:
+            message_count (Optional[int]):
+                The value to override the current message count with.
+            total_bytes (Optional[int]):
+                The value to override the current total bytes with.
+            reserved_bytes (Optional[int]):
+                The value to override the current number of reserved bytes with.
+
+        Returns:
+            str
+        """
+        msg = "messages: {} / {}, bytes: {} / {} (reserved: {})"
+
+        if message_count is None:
+            message_count = self._message_count
+
+        if total_bytes is None:
+            total_bytes = self._total_bytes
+
+        if reserved_bytes is None:
+            reserved_bytes = self._reserved_bytes
+
+        return msg.format(
+            message_count,
+            self._settings.message_limit,
+            total_bytes,
+            self._settings.byte_limit,
+            reserved_bytes,
+        )