diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/leaser.py b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/leaser.py index bcb73352b537..40f17f30a7cb 100644 --- a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/leaser.py +++ b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/leaser.py @@ -64,30 +64,32 @@ def bytes(self): def add(self, items): """Add messages to be managed by the leaser.""" - for item in items: - # Add the ack ID to the set of managed ack IDs, and increment - # the size counter. - if item.ack_id not in self._leased_messages: - self._leased_messages[item.ack_id] = _LeasedMessage( - added_time=time.time(), size=item.byte_size - ) - self._bytes += item.byte_size - else: - _LOGGER.debug("Message %s is already lease managed", item.ack_id) + with self._operational_lock: + for item in items: + # Add the ack ID to the set of managed ack IDs, and increment + # the size counter. + if item.ack_id not in self._leased_messages: + self._leased_messages[item.ack_id] = _LeasedMessage( + added_time=time.time(), size=item.byte_size + ) + self._bytes += item.byte_size + else: + _LOGGER.debug("Message %s is already lease managed", item.ack_id) def remove(self, items): """Remove messages from lease management.""" - # Remove the ack ID from lease management, and decrement the - # byte counter. - for item in items: - if self._leased_messages.pop(item.ack_id, None) is not None: - self._bytes -= item.byte_size - else: - _LOGGER.debug("Item %s was not managed.", item.ack_id) - - if self._bytes < 0: - _LOGGER.debug("Bytes was unexpectedly negative: %d", self._bytes) - self._bytes = 0 + with self._operational_lock: + # Remove the ack ID from lease management, and decrement the + # byte counter. + for item in items: + if self._leased_messages.pop(item.ack_id, None) is not None: + self._bytes -= item.byte_size + else: + _LOGGER.debug("Item %s was not managed.", item.ack_id) + + if self._bytes < 0: + _LOGGER.debug("Bytes was unexpectedly negative: %d", self._bytes) + self._bytes = 0 def maintain_leases(self): """Maintain all of the leases being managed. diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py index 650e2f661915..445d23441a3d 100644 --- a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py +++ b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py @@ -21,6 +21,7 @@ import grpc import six +from six.moves import queue from google.api_core import bidi from google.api_core import exceptions @@ -52,24 +53,35 @@ def _maybe_wrap_exception(exception): return exception -def _wrap_callback_errors(callback, on_callback_error, message): - """Wraps a user callback so that if an exception occurs the message is +def _wrap_callback_errors(callback, on_callback_error, messages): + """Wrap a user callback so that if an exception occurs the messages are nacked. Args: - callback (Callable[None, Message]): The user callback. - message (~Message): The Pub/Sub message. + callback (Callable[Union[Message, List[Message]]]): The user callback. + on_callback_error (Callable[Exception]): The handler to invoke if + unhandled error occurs in ``callback``. + messages (Union[~Message, List[~Message]]): The Pub/Sub message(s). """ try: - callback(message) + callback(messages) except Exception as exc: # Note: the likelihood of this failing is extremely low. This just adds # a message to a queue, so if this doesn't work the world is in an # unrecoverable state and this thread should just bail. _LOGGER.exception( - "Top-level exception occurred in callback while processing a message" + "Top-level exception occurred in callback while processing message(s)" ) - message.nack() + + if isinstance(messages, google.cloud.pubsub_v1.subscriber.message.Message): + messages = [messages] + + for msg in messages: + try: + msg.nack() + except google.cloud.pubsub_v1.subscriber.message.AckStatusSentError: + pass # this message does not have to be NACK-ed + on_callback_error(exc) @@ -116,6 +128,12 @@ def __init__( else: self._scheduler = scheduler + # A FIFO queue for the messages that have been received from the server, + # but not yet added to the lease management (and not sent to user callback), + # because the FlowControl limits have been hit. + self._messages_on_hold = queue.Queue() + self._pause_resume_lock = threading.Lock() + # The threads created in ``.open()``. self._dispatcher = None self._leaser = None @@ -211,26 +229,82 @@ def add_close_callback(self, callback): def maybe_pause_consumer(self): """Check the current load and pause the consumer if needed.""" - if self.load >= 1.0: - if self._consumer is not None and not self._consumer.is_paused: - _LOGGER.debug("Message backlog over load at %.2f, pausing.", self.load) - self._consumer.pause() + with self._pause_resume_lock: + if self.load >= 1.0: + if self._consumer is not None and not self._consumer.is_paused: + _LOGGER.debug( + "Message backlog over load at %.2f, pausing.", self.load + ) + self._consumer.pause() 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 is None or not self._consumer.is_paused: + """Check the load and held messages and resume the consumer if needed. + + If there are messages held internally, release those messages before + resuming the consumer. That will avoid leaser overload. + """ + with self._pause_resume_lock: + # 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 is None or not self._consumer.is_paused: + return + + _LOGGER.debug("Current load: %.2f", self.load) + + # Before maybe resuming the background consumer, release any messages + # currently on hold, if the current load allows for it. + self._maybe_release_messages() + + if self.load < self.flow_control.resume_threshold: + _LOGGER.debug("Current load is %.2f, resuming consumer.", self.load) + self._consumer.resume() + else: + _LOGGER.debug("Did not resume, current load is %.2f.", self.load) + + def _maybe_release_messages(self): + """Release (some of) the held messages if the current load allows for it. + + The method tries to release as many messages as the current leaser load + would allow. Each released message is added to the lease management, + and the user callback is scheduled for them. + + If there are currently no messageges on hold, or if the leaser is + already overloaded, this method is effectively a no-op. + """ + released_messages = [] + + while True: + if self.load >= 1.0: + break # already overloaded + + try: + msg = self._messages_on_hold.get_nowait() + except queue.Empty: + break + + released_messages.append(msg) + + self.leaser.add( + [requests.LeaseRequest(ack_id=msg.ack_id, byte_size=msg.size)] + ) + _LOGGER.debug( + "Released held message to leaser, still on hold %s.", + self._messages_on_hold.qsize(), + ) + + if not released_messages: return - if self.load < self.flow_control.resume_threshold: - self._consumer.resume() - else: - _LOGGER.debug("Did not resume, current load is %s", self.load) + _LOGGER.debug( + "Scheduling %s for %s released message(s).", + "batch callback" if self._batch_callback else "callbacks", + len(released_messages), + ) + self._schedule_callbacks(released_messages) def _send_unary_request(self, request): """Send a request using a separate unary request instead of over the @@ -300,13 +374,24 @@ def heartbeat(self): if self._rpc is not None and self._rpc.is_active: self._rpc.send(types.StreamingPullRequest()) - def open(self, callback, on_callback_error): + def open(self, callback, batch, on_callback_error): """Begin consuming messages. Args: - callback (Callable[None, google.cloud.pubsub_v1.message.Message]): - A callback that will be called for each message received on the + callback ( + Callable[ + Union[ + google.cloud.pubsub_v1.message.Message, + List[google.cloud.pubsub_v1.message.Message], + ] + ] + ): + A callback that will be called for messages received on the stream. + batch (bool): If ``False``, invoke ``callback`` for each individual + message. If ``True``, invoke ``callback`` for each group of + messages received in a single server response (may be less than + a full group depending on the ``flow_control`` limits). on_callback_error (Callable[Exception]): A callable that will be called if an exception is raised in the provided `callback`. @@ -320,6 +405,7 @@ def open(self, callback, on_callback_error): self._callback = functools.partial( _wrap_callback_errors, callback, on_callback_error ) + self._batch_callback = batch # Create the RPC self._rpc = bidi.ResumableBidiRpc( @@ -429,11 +515,14 @@ def _on_response(self, response): redelivered multiple times. After the messages have all had their ack deadline updated, execute - the callback for each message using the executor. + the callback (individually or in a batch) using the executor for all + messages that have been added to the lease management, i.e. not put on + hold due to leaser overload. """ - _LOGGER.debug( - "Scheduling callbacks for %s messages.", len(response.received_messages) + "Processing %s received message(s), currenty on hold %s.", + len(response.received_messages), + self._messages_on_hold.qsize(), ) # Immediately modack the messages we received, as this tells the server @@ -443,12 +532,51 @@ def _on_response(self, response): for message in response.received_messages ] self._dispatcher.modify_ack_deadline(items) + + invoke_callbacks_for = [] + for received_message in response.received_messages: message = google.cloud.pubsub_v1.subscriber.message.Message( - received_message.message, received_message.ack_id, self._scheduler.queue + received_message.message, + received_message.ack_id, + self._scheduler.queue, + autolease=False, ) - # TODO: Immediately lease instead of using the callback queue. - self._scheduler.schedule(self._callback, message) + if self.load < 1.0: + req = requests.LeaseRequest( + ack_id=message.ack_id, byte_size=message.size + ) + self.leaser.add([req]) + invoke_callbacks_for.append(message) + self.maybe_pause_consumer() + else: + self._messages_on_hold.put(message) + + _LOGGER.debug( + "Scheduling %s for %s new messages, new total on hold %s.", + "batch callback" if self._batch_callback else "callbacks", + len(invoke_callbacks_for), + self._messages_on_hold.qsize(), + ) + self._schedule_callbacks(invoke_callbacks_for) + + def _schedule_callbacks(self, messages): + """Schedule callbacks for given messages for an async execution. + + If the streaming pull manager was opened with `batch` == True, a single + batch callback will be scheduled for all given messages. Conversely, if + `batch` == False, a callback will be scheduled once for each message + in `messages`. + + Args: + messages (Union[~Message, List[~Message]]): The Pub/Sub message(s) + to schedule the callbacks for. + """ + if self._batch_callback: + self._scheduler.schedule(self._callback, messages) + else: + for msg in messages: + self._scheduler.schedule(self._callback, msg) def _should_recover(self, exception): """Determine if an error on the RPC stream should be recovered. diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/client.py b/pubsub/google/cloud/pubsub_v1/subscriber/client.py index f2e8faa4fcf5..4da2142f5126 100644 --- a/pubsub/google/cloud/pubsub_v1/subscriber/client.py +++ b/pubsub/google/cloud/pubsub_v1/subscriber/client.py @@ -124,19 +124,25 @@ def api(self): """The underlying gapic API client.""" return self._api - def subscribe(self, subscription, callback, flow_control=(), scheduler=None): + def subscribe( + self, subscription, callback, batch=False, flow_control=(), scheduler=None + ): """Asynchronously start receiving messages on a given subscription. This method starts a background thread to begin pulling messages from a Pub/Sub subscription and scheduling them to be processed using the provided ``callback``. - The ``callback`` will be called with an individual - :class:`google.cloud.pubsub_v1.subscriber.message.Message`. It is the - responsibility of the callback to either call ``ack()`` or ``nack()`` - on the message when it finished processing. If an exception occurs in - the callback during processing, the exception is logged and the message - is ``nack()`` ed. + If ``batch`` is False (default), the ``callback`` will be called with an + individual :class:`google.cloud.pubsub_v1.subscriber.message.Message`. + If ``batch`` is True, the ``callback`` will be called with a list of + all messages received in a single response from the server, albeit + within limits imposed by the ``flow_control`` settings. + + It is the responsibility of the callback to either call ``ack()`` or + ``nack()`` on the message(s) when it finishes processing. If an exception + occurs in the callback during processing, the exception is logged and the + messages are ``nack()`` ed. The ``flow_control`` argument can be used to control the rate of at which messages are pulled. The settings are relatively conservative by @@ -186,10 +192,22 @@ def callback(message): subscription (str): The name of the subscription. The subscription should have already been created (for example, by using :meth:`create_subscription`). - callback (Callable[~google.cloud.pubsub_v1.subscriber.message.Message]): - The callback function. This function receives the message as - its only argument and will be called from a different thread/ - process depending on the scheduling strategy. + callback ( \ + Callable[ \ + Union[ \ + ~google.cloud.pubsub_v1.subscriber.message.Message, \ + List[~google.cloud.pubsub_v1.subscriber.message.Message], \ + ] \ + ] \ + ): + The callback function. This function receives the message (or + a list of messages if ``batch`` is ``True``) as its only argument + and will be called from a different thread/process depending on the + scheduling strategy. + batch (bool): If ``False`` (default), invoke ``callback`` for each + individual received message. If ``True``, invoke ``callback`` + for each _group_ of messages received in a single server response + (within limits imposed by the ``flow_control`` settings). flow_control (~google.cloud.pubsub_v1.types.FlowControl): The flow control settings. Use this to prevent situations where you are inundated with too many messages at once. @@ -209,6 +227,8 @@ def callback(message): future = futures.StreamingPullFuture(manager) - manager.open(callback=callback, on_callback_error=future.set_exception) + manager.open( + callback=callback, batch=batch, on_callback_error=future.set_exception + ) return future diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/message.py b/pubsub/google/cloud/pubsub_v1/subscriber/message.py index 56dde9a7f6b8..f712120418fe 100644 --- a/pubsub/google/cloud/pubsub_v1/subscriber/message.py +++ b/pubsub/google/cloud/pubsub_v1/subscriber/message.py @@ -17,6 +17,7 @@ import datetime import json import math +import threading import time from google.api_core import datetime_helpers @@ -49,6 +50,10 @@ def _indent(lines, prefix=" "): return "\n".join(indented) +class AckStatusSentError(Exception): + """An error raised on attempt to (N)ACK an already (N)ACK-ed Message.""" + + class Message(object): """A representation of a single Pub/Sub message. @@ -70,7 +75,7 @@ class Message(object): published. """ - def __init__(self, message, ack_id, request_queue): + def __init__(self, message, ack_id, request_queue, autolease=True): """Construct the Message. .. note:: @@ -85,6 +90,9 @@ def __init__(self, message, ack_id, request_queue): request_queue (queue.Queue): A queue provided by the policy that can accept requests; the policy is responsible for handling those requests. + autolease (bool): An optional flag determining whether a new Message + instance should automatically lease itself upon creation. + Defaults to :data:`True`. """ self._message = message self._ack_id = ack_id @@ -96,9 +104,14 @@ def __init__(self, message, ack_id, request_queue): # the default lease deadline. self._received_timestamp = time.time() + # whether or not an ACK/NACK request has been sent for the message + self._ack_status_sent = False + self._ack_status_lock = threading.Lock() + # The policy should lease this message, telling PubSub that it has # it until it is acked or otherwise dropped. - self.lease() + if autolease: + self.lease() def __repr__(self): # Get an abbreviated version of the data. @@ -179,13 +192,22 @@ def ack(self): Acks in Pub/Sub are best effort. You should always ensure that your processing code is idempotent, as you may receive any given message more than once. + + Raises: + ~.pubsub_v1.subscriber.message.AckStatusSentError: If the message + has already been ACK-ed or NACK-ed. """ - time_to_ack = math.ceil(time.time() - self._received_timestamp) - self._request_queue.put( - requests.AckRequest( - ack_id=self._ack_id, byte_size=self.size, time_to_ack=time_to_ack + with self._ack_status_lock: + if self._ack_status_sent: + raise AckStatusSentError("Message already ACK/NACK-ed") + + time_to_ack = math.ceil(time.time() - self._received_timestamp) + self._request_queue.put( + requests.AckRequest( + ack_id=self._ack_id, byte_size=self.size, time_to_ack=time_to_ack + ) ) - ) + self._ack_status_sent = True def drop(self): """Release the message from lease management. @@ -208,8 +230,8 @@ def lease(self): """Inform the policy to lease this message continually. .. note:: - This method is called by the constructor, and you should never - need to call it manually. + By default this method is called by the constructor, and you should + never need to call it manually. """ self._request_queue.put( requests.LeaseRequest(ack_id=self._ack_id, byte_size=self.size) @@ -238,7 +260,16 @@ def nack(self): """Decline to acknowldge the given message. This will cause the message to be re-delivered to the subscription. + + Raises: + ~.pubsub_v1.subscriber.message.AckStatusSentError: If the message + has already been ACK-ed or NACK-ed. """ - self._request_queue.put( - requests.NackRequest(ack_id=self._ack_id, byte_size=self.size) - ) + with self._ack_status_lock: + if self._ack_status_sent: + raise AckStatusSentError("Message already ACK/NACK-ed") + + self._request_queue.put( + requests.NackRequest(ack_id=self._ack_id, byte_size=self.size) + ) + self._ack_status_sent = True diff --git a/pubsub/tests/system.py b/pubsub/tests/system.py index e8921e039164..988b5e94d4cc 100644 --- a/pubsub/tests/system.py +++ b/pubsub/tests/system.py @@ -15,6 +15,7 @@ from __future__ import absolute_import import datetime +import itertools import threading import time @@ -24,7 +25,10 @@ import google.auth from google.cloud import pubsub_v1 - +from google.cloud.pubsub_v1 import exceptions +from google.cloud.pubsub_v1 import futures +from google.cloud.pubsub_v1 import types +from google.cloud.pubsub_v1.subscriber import message from test_utils.system import unique_resource_id @@ -206,6 +210,124 @@ class CallbackError(Exception): with pytest.raises(CallbackError): future.result(timeout=30) + def test_batch_callbacks( + self, publisher, topic_path, subscriber, subscription_path, cleanup + ): + # Make sure the topic and subscription get deleted. + cleanup.append((publisher.delete_topic, topic_path)) + cleanup.append((subscriber.delete_subscription, subscription_path)) + + # create a topic and subscribe to it + publisher.create_topic(topic_path) + subscriber.create_subscription(subscription_path, topic_path) + + # publish messages and wait until published + self._publish_messages(publisher, topic_path, batch_sizes=[3]) + + # start pulling messages and check that the callback is indeed invoked + # with batches of messages + callback_lock = threading.Lock() + callback_invoked = futures.Future() + + def callback(messages): + with callback_lock: + if not callback_invoked.done(): + callback_invoked.set_result(messages) + + subscription_future = subscriber.subscribe( + subscription_path, callback, batch=True + ) + + try: + result = callback_invoked.result(timeout=30) + except exceptions.TimeoutError: + pytest.fail("Subscription callback not invoked in time.") + else: + assert isinstance(result, list) + assert len(result) == 3 + assert all(isinstance(item, message.Message) for item in result) + finally: + subscription_future.cancel() # shutdown streaming pull + + def test_streaming_pull_max_messages( + self, publisher, topic_path, subscriber, subscription_path, cleanup + ): + # Make sure the topic and subscription get deleted. + cleanup.append((publisher.delete_topic, topic_path)) + cleanup.append((subscriber.delete_subscription, subscription_path)) + + # create a topic and subscribe to it + publisher.create_topic(topic_path) + subscriber.create_subscription(subscription_path, topic_path) + + batch_sizes = (7, 4, 8, 2, 10, 1, 3, 8, 6, 1) # total: 50 + self._publish_messages(publisher, topic_path, batch_sizes=batch_sizes) + + # now subscribe and do the main part, check for max pending messages + total_messages = sum(batch_sizes) + flow_control = types.FlowControl(max_messages=5) + callback = StreamingPullCallback( + processing_time=1, resolve_at_msg_count=total_messages + ) + + subscription_future = subscriber.subscribe( + subscription_path, callback, flow_control=flow_control + ) + + # Expected time to process all messages in ideal case: + # (total_messages / FlowControl.max_messages) * processing_time + # + # With total=50, max messages=5, and processing_time=1 this amounts to + # 10 seconds (+ overhead), thus a full minute should be more than enough + # for the processing to complete. If not, fail the test with a timeout. + try: + callback.done_future.result(timeout=60) + except exceptions.TimeoutError: + pytest.fail( + "Timeout: receiving/processing streamed messages took too long." + ) + + # The callback future gets resolved once total_messages have been processed, + # but we want to wait for just a little bit longer to possibly catch cases + # when the callback gets invoked *more* than total_messages times. + time.sleep(3) + + try: + # All messages should have been processed exactly once, and no more + # than max_messages simultaneously at any time. + assert callback.completed_calls == total_messages + assert sorted(callback.seen_message_ids) == list( + range(1, total_messages + 1) + ) + assert callback.max_pending_ack <= flow_control.max_messages + finally: + subscription_future.cancel() # trigger clean shutdown + + def _publish_messages(self, publisher, topic_path, batch_sizes): + """Publish ``count`` messages in batches and wait until completion.""" + publish_futures = [] + msg_counter = itertools.count(start=1) + + for batch_size in batch_sizes: + msg_batch = self._make_messages(count=batch_size) + for msg in msg_batch: + future = publisher.publish( + topic_path, msg, seq_num=str(next(msg_counter)) + ) + publish_futures.append(future) + time.sleep(0.1) + + # wait untill all messages have been successfully published + for future in publish_futures: + future.result(timeout=30) + + def _make_messages(self, count): + messages = [ + u"message {}/{}".format(i, count).encode("utf-8") + for i in range(1, count + 1) + ] + return messages + class AckCallback(object): def __init__(self): @@ -236,3 +358,33 @@ def __call__(self, message): # ``calls`` is incremented to do it. self.call_times.append(now) self.calls += 1 + + +class StreamingPullCallback(object): + def __init__(self, processing_time, resolve_at_msg_count): + self._lock = threading.Lock() + self._processing_time = processing_time + self._pending_ack = 0 + self.max_pending_ack = 0 + self.completed_calls = 0 + self.seen_message_ids = [] + + self._resolve_at_msg_count = resolve_at_msg_count + self.done_future = futures.Future() + + def __call__(self, message): + with self._lock: + self._pending_ack += 1 + self.max_pending_ack = max(self.max_pending_ack, self._pending_ack) + self.seen_message_ids.append(int(message.attributes["seq_num"])) + + time.sleep(self._processing_time) + + with self._lock: + self._pending_ack -= 1 + message.ack() + self.completed_calls += 1 + + if self.completed_calls >= self._resolve_at_msg_count: + if not self.done_future.done(): + self.done_future.set_result(None) diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py index 98a946ae75c6..e320dfe2ce1b 100644 --- a/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py +++ b/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py @@ -16,6 +16,7 @@ import time import mock +import pytest import pytz from six.moves import queue from google.protobuf import timestamp_pb2 @@ -33,7 +34,7 @@ PUBLISHED_SECONDS = datetime_helpers.to_milliseconds(PUBLISHED) // 1000 -def create_message(data, ack_id="ACKID", **attrs): +def create_message(data, ack_id="ACKID", autolease=True, **attrs): with mock.patch.object(message.Message, "lease") as lease: with mock.patch.object(time, "time") as time_: time_.return_value = RECEIVED_SECONDS @@ -48,8 +49,12 @@ def create_message(data, ack_id="ACKID", **attrs): ), ack_id, queue.Queue(), + autolease=autolease, ) - lease.assert_called_once_with() + if autolease: + lease.assert_called_once_with() + else: + lease.assert_not_called() return msg @@ -79,6 +84,11 @@ def test_publish_time(): assert msg.publish_time == PUBLISHED +def test_disable_autolease_on_creation(): + # the create_message() helper does the actual assertion + create_message(b"foo", autolease=False) + + def check_call_types(mock, *args, **kwargs): """Checks a mock's call types. @@ -154,6 +164,28 @@ def test_nack(): check_call_types(put, requests.NackRequest) +def test_repeated_ack_nack_not_allowed(): + msg = create_message(b"foo", ack_id="bogus_ack_id") + with mock.patch.object(msg._request_queue, "put") as put: + msg.ack() + put.reset_mock() + with pytest.raises(message.AckStatusSentError): + msg.ack() + with pytest.raises(message.AckStatusSentError): + msg.nack() + put.assert_not_called() + + msg = create_message(b"bar", ack_id="bogus_ack_id_2") + with mock.patch.object(msg._request_queue, "put") as put: + msg.nack() + put.reset_mock() + with pytest.raises(message.AckStatusSentError): + msg.ack() + with pytest.raises(message.AckStatusSentError): + msg.nack() + put.assert_not_called() + + def test_repr(): data = b"foo" msg = create_message(data, snow="cones", orange="juice") diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py index cbd02e28ac6c..dc54cc90799b 100644 --- a/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py +++ b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py @@ -13,9 +13,11 @@ # limitations under the License. import logging +import types as stdlib_types import mock import pytest +from six.moves import queue from google.api_core import bidi from google.api_core import exceptions @@ -60,7 +62,7 @@ def test__wrap_callback_errors_no_error(): on_callback_error.assert_not_called() -def test__wrap_callback_errors_error(): +def test__wrap_callback_errors_error_no_batch(): callback_error = ValueError("meep") msg = mock.create_autospec(message.Message, instance=True) @@ -73,6 +75,28 @@ def test__wrap_callback_errors_error(): on_callback_error.assert_called_once_with(callback_error) +def test__wrap_callback_errors_error_with_batch(): + callback_error = ValueError("meep") + + messages = [ + mock.create_autospec(message.Message, instance=True), + mock.create_autospec(message.Message, instance=True), + ] + # Some of the messages might have already been (N)ACK-ed in the callback, + # thus simulate that to verify the robustness of the MUT. + messages[0].nack.side_effect = message.AckStatusSentError + + callback = mock.Mock(side_effect=callback_error) + on_callback_error = mock.Mock() + + # If the call below does not raise AckStatusSentError, it means that the MUT + # handled the already (N)ACK-ed messages correctly. + streaming_pull_manager._wrap_callback_errors(callback, on_callback_error, messages) + + assert all(len(msg.nack.mock_calls) == 1 for msg in messages) + on_callback_error.assert_called_once_with(callback_error) + + def test_constructor_and_default_state(): manager = streaming_pull_manager.StreamingPullManager( mock.sentinel.client, mock.sentinel.subscription @@ -113,6 +137,23 @@ def make_manager(**kwargs): ) +def fake_leaser_add(leaser, init_msg_count=0, init_bytes=0): + """Add a simplified fake add() method to a leaser instance. + + The fake add() method actually increases the leaser's internal message count + by one for each message, and the total bytes by 10 for each message (hardcoded, + regardless of the actual message size). + """ + + def fake_add(self, items): + self.message_count += len(items) + self.bytes += len(items) * 10 + + leaser.message_count = init_msg_count + leaser.bytes = init_bytes + leaser.add = stdlib_types.MethodType(fake_add, leaser) + + def test_ack_deadline(): manager = make_manager() assert manager.ack_deadline == 10 @@ -208,6 +249,71 @@ def test_maybe_resume_consumer_wo_consumer_set(): manager.maybe_resume_consumer() # no raise +def test__maybe_release_messages_on_overload(): + manager = make_manager( + flow_control=types.FlowControl(max_messages=10, max_bytes=1000) + ) + manager._batch_callback = False + manager._schedule_callbacks = mock.Mock() + + # Ensure load is exactly 1.0 (to verify that >= condition is used) + _leaser = manager._leaser = mock.create_autospec(leaser.Leaser) + _leaser.message_count = 10 + _leaser.bytes = 1000 + + msg = mock.create_autospec(message.Message, instance=True, ack_id="ack", size=11) + manager._messages_on_hold.put(msg) + + manager._maybe_release_messages() + + assert manager._messages_on_hold.qsize() == 1 + manager._leaser.add.assert_not_called() + manager._schedule_callbacks.assert_not_called() + + +def test__maybe_release_messages_below_overload(): + manager = make_manager( + flow_control=types.FlowControl(max_messages=10, max_bytes=1000) + ) + manager._callback = mock.sentinel.callback + manager._batch_callback = False + manager._schedule_callbacks = mock.Mock() + + # init leaser message count to 8 to leave room for 2 more messages + _leaser = manager._leaser = mock.create_autospec(leaser.Leaser) + fake_leaser_add(_leaser, init_msg_count=8, init_bytes=200) + _leaser.add = mock.Mock(wraps=_leaser.add) # to spy on calls + + messages = [ + mock.create_autospec(message.Message, instance=True, ack_id="ack_foo", size=11), + mock.create_autospec(message.Message, instance=True, ack_id="ack_bar", size=22), + mock.create_autospec(message.Message, instance=True, ack_id="ack_baz", size=33), + ] + for msg in messages: + manager._messages_on_hold.put(msg) + + # the actual call of MUT + manager._maybe_release_messages() + + assert manager._messages_on_hold.qsize() == 1 + msg = manager._messages_on_hold.get_nowait() + assert msg.ack_id == "ack_baz" + + assert len(_leaser.add.mock_calls) == 2 + expected_calls = [ + mock.call([requests.LeaseRequest(ack_id="ack_foo", byte_size=11)]), + mock.call([requests.LeaseRequest(ack_id="ack_bar", byte_size=22)]), + ] + _leaser.add.assert_has_calls(expected_calls) + + calls = manager._schedule_callbacks.mock_calls + assert len(calls) == 1 + call_arg = calls[0][1][0] # the first positional argument of the call + assert len(call_arg) == 2 + assert all(isinstance(item, message.Message) for item in call_arg) + assert sorted(item.ack_id for item in call_arg) == ["ack_bar", "ack_foo"] + + def test_send_unary(): manager = make_manager() manager._UNARY_REQUESTS = True @@ -325,7 +431,7 @@ def test_heartbeat_inactive(): def test_open(heartbeater, dispatcher, leaser, background_consumer, resumable_bidi_rpc): manager = make_manager() - manager.open(mock.sentinel.callback, mock.sentinel.on_callback_error) + manager.open(mock.sentinel.callback, False, mock.sentinel.on_callback_error) heartbeater.assert_called_once_with(manager) heartbeater.return_value.start.assert_called_once() @@ -363,7 +469,7 @@ def test_open_already_active(): manager._consumer.is_active = True with pytest.raises(ValueError, match="already open"): - manager.open(mock.sentinel.callback, mock.sentinel.on_callback_error) + manager.open(mock.sentinel.callback, False, mock.sentinel.on_callback_error) def test_open_has_been_closed(): @@ -371,10 +477,10 @@ def test_open_has_been_closed(): manager._closed = True with pytest.raises(ValueError, match="closed"): - manager.open(mock.sentinel.callback, mock.sentinel.on_callback_error) + manager.open(mock.sentinel.callback, False, mock.sentinel.on_callback_error) -def make_running_manager(): +def make_running_manager(batch_callback=False): manager = make_manager() manager._consumer = mock.create_autospec(bidi.BackgroundConsumer, instance=True) manager._consumer.is_active = True @@ -382,6 +488,8 @@ def make_running_manager(): manager._leaser = mock.create_autospec(leaser.Leaser, instance=True) manager._heartbeater = mock.create_autospec(heartbeater.Heartbeater, instance=True) + manager._batch_callback = batch_callback + return ( manager, manager._consumer, @@ -470,9 +578,10 @@ def test__get_initial_request_wo_leaser(): assert initial_request.modify_deadline_seconds == [] -def test_on_response(): - manager, _, dispatcher, _, _, scheduler = make_running_manager() +def test__on_response_no_leaser_overload(): + manager, _, dispatcher, leaser, _, scheduler = make_running_manager() manager._callback = mock.sentinel.callback + manager._schedule_callbacks = mock.Mock() # Set up the messages. response = types.StreamingPullResponse( @@ -486,6 +595,9 @@ def test_on_response(): ] ) + # adjust message bookkeeping in leaser + fake_leaser_add(leaser, init_msg_count=0, init_bytes=0) + # Actually run the method and prove that modack and schedule # are called in the expected way. manager._on_response(response) @@ -493,12 +605,101 @@ def test_on_response(): dispatcher.modify_ack_deadline.assert_called_once_with( [requests.ModAckRequest("fack", 10), requests.ModAckRequest("back", 10)] ) + calls = manager._schedule_callbacks.mock_calls + assert len(calls) == 1 + call_arg = calls[0][1][0] # the first positional argument of the call + assert all(isinstance(item, message.Message) for item in call_arg) + assert sorted(item.ack_id for item in call_arg) == ["back", "fack"] + + # the leaser load limit not hit, no messages had to be put on hold + assert manager._messages_on_hold.qsize() == 0 + + +def test__on_response_with_leaser_overload(): + manager, _, dispatcher, leaser, _, scheduler = make_running_manager() + manager._callback = mock.sentinel.callback + manager._schedule_callbacks = mock.Mock() + + # Set up the messages. + response = types.StreamingPullResponse( + received_messages=[ + types.ReceivedMessage( + ack_id="fack", message=types.PubsubMessage(data=b"foo", message_id="1") + ), + types.ReceivedMessage( + ack_id="back", message=types.PubsubMessage(data=b"bar", message_id="2") + ), + types.ReceivedMessage( + ack_id="zack", message=types.PubsubMessage(data=b"baz", message_id="3") + ), + ] + ) + + # Adjust message bookkeeping in leaser. Pick 99 messages, which is just below + # the default FlowControl.max_messages limit. + fake_leaser_add(leaser, init_msg_count=99, init_bytes=990) + + # Actually run the method and prove that modack and schedule + # are called in the expected way. + manager._on_response(response) + + dispatcher.modify_ack_deadline.assert_called_once_with( + [ + requests.ModAckRequest("fack", 10), + requests.ModAckRequest("back", 10), + requests.ModAckRequest("zack", 10), + ] + ) + + # one message should be scheduled, the leaser capacity allows for it + calls = manager._schedule_callbacks.mock_calls + assert len(calls) == 1 + call_arg = calls[0][1][0] # the first positional argument of the call + assert len(call_arg) == 1 + assert isinstance(call_arg[0], message.Message) + assert call_arg[0].message_id == "1" + + # the rest of the messages should have been put on hold + assert manager._messages_on_hold.qsize() == 2 + while True: + try: + msg = manager._messages_on_hold.get_nowait() + except queue.Empty: + break + else: + assert isinstance(msg, message.Message) + assert msg.message_id in ("2", "3") + + +def test__schedule_callbacks_no_batch(): + manager, _, _, _, _, scheduler = make_running_manager(batch_callback=False) + manager._callback = mock.sentinel.callback + + messages = [ + mock.create_autospec(message.Message, instance=True, ack_id="ack_foo"), + mock.create_autospec(message.Message, instance=True, ack_id="ack_bar"), + ] + manager._schedule_callbacks(messages) + + assert scheduler.schedule.call_count == 2 + expected_calls = [ + mock.call(mock.sentinel.callback, messages[0]), + mock.call(mock.sentinel.callback, messages[1]), + ] + scheduler.schedule.assert_has_calls(expected_calls, any_order=True) + + +def test__schedule_callbacks_with_batch(): + manager, _, _, _, _, scheduler = make_running_manager(batch_callback=True) + manager._callback = mock.sentinel.callback + + messages = [ + mock.create_autospec(message.Message, instance=True, ack_id="ack_foo"), + mock.create_autospec(message.Message, instance=True, ack_id="ack_bar"), + ] + manager._schedule_callbacks(messages) - schedule_calls = scheduler.schedule.mock_calls - assert len(schedule_calls) == 2 - for call in schedule_calls: - assert call[1][0] == mock.sentinel.callback - assert isinstance(call[1][1], message.Message) + scheduler.schedule.assert_called_once_with(mock.sentinel.callback, messages) def test_retryable_stream_errors(): diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_subscriber_client.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_subscriber_client.py index 8bdb414c6280..8419ddf075a8 100644 --- a/pubsub/tests/unit/pubsub_v1/subscriber/test_subscriber_client.py +++ b/pubsub/tests/unit/pubsub_v1/subscriber/test_subscriber_client.py @@ -68,12 +68,14 @@ def test_subscribe(manager_open): creds = mock.Mock(spec=credentials.Credentials) client = subscriber.Client(credentials=creds) - future = client.subscribe("sub_name_a", callback=mock.sentinel.callback) + future = client.subscribe( + "sub_name_a", callback=mock.sentinel.callback, batch=mock.sentinel.batch + ) assert isinstance(future, futures.StreamingPullFuture) assert future._manager._subscription == "sub_name_a" manager_open.assert_called_once_with( - mock.ANY, mock.sentinel.callback, future.set_exception + mock.ANY, mock.sentinel.callback, mock.sentinel.batch, future.set_exception ) @@ -91,6 +93,7 @@ def test_subscribe_options(manager_open): future = client.subscribe( "sub_name_a", callback=mock.sentinel.callback, + batch=mock.sentinel.batch, flow_control=flow_control, scheduler=scheduler, ) @@ -100,5 +103,5 @@ def test_subscribe_options(manager_open): assert future._manager.flow_control == flow_control assert future._manager._scheduler == scheduler manager_open.assert_called_once_with( - mock.ANY, mock.sentinel.callback, future.set_exception + mock.ANY, mock.sentinel.callback, mock.sentinel.batch, future.set_exception )