From 9c7b04e0c41cae5ddb70a3d4ceaf90641876963e Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Tue, 7 May 2019 18:46:22 +0200
Subject: [PATCH 1/8] Allow skipping Pub/Sub Message autolease

In certain cases automatically leasing Message instances upon creation
might not be desired, thus an optional parameter is added to Message
initializer that allows skipping that.

The default behavior is not changed, new Message instances *are*
automatically leased upon creation.
---
 pubsub/google/cloud/pubsub_v1/subscriber/message.py | 12 ++++++++----
 .../tests/unit/pubsub_v1/subscriber/test_message.py | 13 +++++++++++--
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/message.py b/pubsub/google/cloud/pubsub_v1/subscriber/message.py
index 56dde9a7f6b8..9a2b35416f08 100644
--- a/pubsub/google/cloud/pubsub_v1/subscriber/message.py
+++ b/pubsub/google/cloud/pubsub_v1/subscriber/message.py
@@ -70,7 +70,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 +85,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
@@ -98,7 +101,8 @@ def __init__(self, message, ack_id, request_queue):
 
         # 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.
@@ -208,8 +212,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)
diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py
index 98a946ae75c6..8c22992f7a2b 100644
--- a/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py
+++ b/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py
@@ -33,7 +33,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 +48,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 +83,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.
 

From 097380324058baa1c25ebfe70134b68e2bbfe8f5 Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Tue, 7 May 2019 20:22:08 +0200
Subject: [PATCH 2/8] Directly lease received Messages w/o request queue

Leasing messages through a request queue in dispatcher causes a race
condition with the ConsumeBidirectionalStream thread. A request to
pause the background consumer can arrive when the Bidi consumer is just
about to fetch the the next batch of messages, and thus the latter gets
paused only *after* fetching those messages.

This commit synchronously leases received messages in the streaming
pull manager callback. If that hits the lease management load limit,
the background consumer is paused synchronously, and will correctly
pause *before* pulling another batch of messages.
---
 .../_protocol/streaming_pull_manager.py       | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

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..734be555a48f 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
@@ -443,13 +443,28 @@ def _on_response(self, response):
             for message in response.received_messages
         ]
         self._dispatcher.modify_ack_deadline(items)
+
+        lease_requests = []
+
         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,
+            )
+            lease_requests.append(
+                requests.LeaseRequest(ack_id=message.ack_id, byte_size=message.size)
             )
-            # TODO: Immediately lease instead of using the callback queue.
             self._scheduler.schedule(self._callback, message)
 
+        # TODO: Since the number of received messages can cause an overflow of
+        # the leaser, we need to assure that only a portion of them are actually
+        # leased (and their callbacks scheduled). The rest need to be kept in an
+        # internal buffer until the leaser again has enough room to accept them.
+        self.leaser.add(lease_requests)
+        self.maybe_pause_consumer()
+
     def _should_recover(self, exception):
         """Determine if an error on the RPC stream should be recovered.
 

From f227602586518c45b3ccc334a77d01aa4278b9f8 Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Wed, 8 May 2019 18:10:32 +0200
Subject: [PATCH 3/8] Add streaming pull message holding buffer

If the PubSub backend sends too many messages in a single response that
would cause the leaser overload should all these messeges were added to
it, the StreamingPullManager now puts excessive messages into an
internal holding buffer.

The messages are released from the buffer when the leaser again has
enough capacity (as defined by the FlowControl settings), and the
message received callback is invoked then as well.
---
 .../_protocol/streaming_pull_manager.py       |  85 +++++++++--
 .../subscriber/test_streaming_pull_manager.py | 144 +++++++++++++++++-
 2 files changed, 211 insertions(+), 18 deletions(-)

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 734be555a48f..0a399aabe90e 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
@@ -116,6 +117,11 @@ 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()
+
         # The threads created in ``.open()``.
         self._dispatcher = None
         self._leaser = None
@@ -217,7 +223,11 @@ def maybe_pause_consumer(self):
                 self._consumer.pause()
 
     def maybe_resume_consumer(self):
-        """Check the current load and resume the consumer if needed."""
+        """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.
+        """
         # If we have been paused by flow control, check and see if we are
         # back within our limits.
         #
@@ -227,10 +237,46 @@ def maybe_resume_consumer(self):
         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 %s", self.load)
+            _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 it.
+
+        If there are currently no messageges on hold, or if the leaser is
+        already overloaded, this method is effectively a no-op.
+        """
+        while True:
+            if self.load >= 1.0:
+                break  # already overloaded
+
+            try:
+                msg = self._messages_on_hold.get_nowait()
+            except queue.Empty:
+                break
+
+            self.leaser.add(
+                [requests.LeaseRequest(ack_id=msg.ack_id, byte_size=msg.size)]
+            )
+            _LOGGER.debug(
+                "Released held message to leaser, scheduling callback for it, "
+                "still on hold %s.",
+                self._messages_on_hold.qsize(),
+            )
+            self._scheduler.schedule(self._callback, msg)
 
     def _send_unary_request(self, request):
         """Send a request using a separate unary request instead of over the
@@ -431,9 +477,10 @@ def _on_response(self, response):
         After the messages have all had their ack deadline updated, execute
         the callback for each message using the executor.
         """
-
         _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
@@ -444,7 +491,7 @@ def _on_response(self, response):
         ]
         self._dispatcher.modify_ack_deadline(items)
 
-        lease_requests = []
+        invoke_callbacks_for = []
 
         for received_message in response.received_messages:
             message = google.cloud.pubsub_v1.subscriber.message.Message(
@@ -453,17 +500,23 @@ def _on_response(self, response):
                 self._scheduler.queue,
                 autolease=False,
             )
-            lease_requests.append(
-                requests.LeaseRequest(ack_id=message.ack_id, byte_size=message.size)
-            )
-            self._scheduler.schedule(self._callback, message)
-
-        # TODO: Since the number of received messages can cause an overflow of
-        # the leaser, we need to assure that only a portion of them are actually
-        # leased (and their callbacks scheduled). The rest need to be kept in an
-        # internal buffer until the leaser again has enough room to accept them.
-        self.leaser.add(lease_requests)
-        self.maybe_pause_consumer()
+            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 callbacks for %s new messages, new total on hold %s.",
+            len(invoke_callbacks_for),
+            self._messages_on_hold.qsize(),
+        )
+        for msg in invoke_callbacks_for:
+            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/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
index cbd02e28ac6c..22585675a324 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
@@ -113,6 +115,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 +227,66 @@ 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)
+    )
+    # 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._scheduler.schedule.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
+
+    # 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)
+
+    schedule_calls = manager._scheduler.schedule.mock_calls
+    assert len(schedule_calls) == 2
+    for _, call_args, _ in schedule_calls:
+        assert call_args[0] == mock.sentinel.callback
+        assert isinstance(call_args[1], message.Message)
+        assert call_args[1].ack_id in ("ack_foo", "ack_bar")
+
+
 def test_send_unary():
     manager = make_manager()
     manager._UNARY_REQUESTS = True
@@ -470,8 +549,8 @@ 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
 
     # Set up the messages.
@@ -486,6 +565,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)
@@ -500,6 +582,64 @@ def test_on_response():
         assert call[1][0] == mock.sentinel.callback
         assert isinstance(call[1][1], message.Message)
 
+    # 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
+
+    # 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
+    schedule_calls = scheduler.schedule.mock_calls
+    assert len(schedule_calls) == 1
+    call_args = schedule_calls[0][1]
+    assert call_args[0] == mock.sentinel.callback
+    assert isinstance(call_args[1], message.Message)
+    assert call_args[1].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_retryable_stream_errors():
     # Make sure the config matches our hard-coded tuple of exceptions.

From 9b010605f0796a73f3a0e65bee4cba9c6727238e Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Thu, 9 May 2019 10:52:53 +0200
Subject: [PATCH 4/8] Make a few streaming pull methods thread-safe

With the StreamingPullManager._on_response() callback adding received
messages to the leaser synchronously (in the background consumer
thread), a race condition can happen with the dispatcher thread that
can asynchronously add (remove) messages to (from) lease management,
e.g. on ack() and nack() requests.

The same is the case with related operations of maybe pausing/resuming
the background consumer. This commit thus adds locks in key places,
assuring that these operations are atomic, ant not subject to race
conditions.
---
 .../pubsub_v1/subscriber/_protocol/leaser.py  | 44 +++++++--------
 .../_protocol/streaming_pull_manager.py       | 53 ++++++++++---------
 2 files changed, 52 insertions(+), 45 deletions(-)

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 0a399aabe90e..a7001e5ce85a 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
@@ -121,6 +121,7 @@ def __init__(
         # 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
@@ -217,10 +218,13 @@ 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 load and held messages and resume the consumer if needed.
@@ -228,26 +232,27 @@ def maybe_resume_consumer(self):
         If there are messages held internally, release those messages before
         resuming the consumer. That will avoid leaser overload.
         """
-        # 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)
+        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.

From 7e5b9d778b8d507a634ef6e0b06e924be2d4ddfe Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Fri, 10 May 2019 18:08:07 +0200
Subject: [PATCH 5/8] Add system test for PubSub max_messages setting

---
 pubsub/tests/system.py | 113 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 113 insertions(+)

diff --git a/pubsub/tests/system.py b/pubsub/tests/system.py
index e8921e039164..13e81d281f42 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,6 +25,9 @@
 
 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 test_utils.system import unique_resource_id
@@ -206,6 +210,85 @@ class CallbackError(Exception):
         with pytest.raises(CallbackError):
             future.result(timeout=30)
 
+    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 +319,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)

From 1fe79d97664d14687faffc207d9abd4ae8ff51a1 Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Wed, 15 May 2019 13:28:18 +0200
Subject: [PATCH 6/8] Only allow one ack/nack call on Message instance

---
 .../cloud/pubsub_v1/subscriber/message.py     | 43 +++++++++++++++----
 .../unit/pubsub_v1/subscriber/test_message.py | 23 ++++++++++
 2 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/message.py b/pubsub/google/cloud/pubsub_v1/subscriber/message.py
index 9a2b35416f08..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.
 
@@ -99,6 +104,10 @@ def __init__(self, message, ack_id, request_queue, autolease=True):
         # 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.
         if autolease:
@@ -183,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.
@@ -242,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/unit/pubsub_v1/subscriber/test_message.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_message.py
index 8c22992f7a2b..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
@@ -163,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")

From c38c2d990f4abebe7afc523ed47a0cf4b0a94cbc Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Tue, 14 May 2019 13:56:39 +0200
Subject: [PATCH 7/8] Add batch callback option to streaming pull

Subscriber client's async streaming pull can now be configured to
invoke a callback for batches of received messages, as opposed to
invoking a callback for each received message.
---
 .../_protocol/streaming_pull_manager.py       |  91 ++++++++++++---
 .../cloud/pubsub_v1/subscriber/client.py      |  44 +++++--
 .../subscriber/test_streaming_pull_manager.py | 109 ++++++++++++++----
 .../subscriber/test_subscriber_client.py      |   9 +-
 4 files changed, 196 insertions(+), 57 deletions(-)

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 a7001e5ce85a..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
@@ -53,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)
 
 
@@ -259,11 +270,13 @@ def _maybe_release_messages(self):
 
         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 it.
+        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
@@ -273,15 +286,25 @@ def _maybe_release_messages(self):
             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, scheduling callback for it, "
-                "still on hold %s.",
+                "Released held message to leaser, still on hold %s.",
                 self._messages_on_hold.qsize(),
             )
-            self._scheduler.schedule(self._callback, msg)
+
+        if not released_messages:
+            return
+
+        _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
@@ -351,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`.
@@ -371,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(
@@ -480,7 +515,9 @@ 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(
             "Processing %s received message(s), currenty on hold %s.",
@@ -516,12 +553,30 @@ def _on_response(self, response):
                 self._messages_on_hold.put(message)
 
         _LOGGER.debug(
-            "Scheduling callbacks for %s new messages, new total on hold %s.",
+            "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(),
         )
-        for msg in invoke_callbacks_for:
-            self._scheduler.schedule(self._callback, msg)
+        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/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
index 22585675a324..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
@@ -62,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)
@@ -75,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
@@ -231,6 +253,9 @@ 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
@@ -243,7 +268,7 @@ def test__maybe_release_messages_on_overload():
 
     assert manager._messages_on_hold.qsize() == 1
     manager._leaser.add.assert_not_called()
-    manager._scheduler.schedule.assert_not_called()
+    manager._schedule_callbacks.assert_not_called()
 
 
 def test__maybe_release_messages_below_overload():
@@ -251,6 +276,8 @@ def test__maybe_release_messages_below_overload():
         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)
@@ -279,12 +306,12 @@ def test__maybe_release_messages_below_overload():
     ]
     _leaser.add.assert_has_calls(expected_calls)
 
-    schedule_calls = manager._scheduler.schedule.mock_calls
-    assert len(schedule_calls) == 2
-    for _, call_args, _ in schedule_calls:
-        assert call_args[0] == mock.sentinel.callback
-        assert isinstance(call_args[1], message.Message)
-        assert call_args[1].ack_id in ("ack_foo", "ack_bar")
+    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():
@@ -404,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()
@@ -442,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():
@@ -450,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
@@ -461,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,
@@ -552,6 +581,7 @@ def test__get_initial_request_wo_leaser():
 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(
@@ -575,12 +605,11 @@ def test__on_response_no_leaser_overload():
     dispatcher.modify_ack_deadline.assert_called_once_with(
         [requests.ModAckRequest("fack", 10), requests.ModAckRequest("back", 10)]
     )
-
-    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)
+    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
@@ -589,6 +618,7 @@ def test__on_response_no_leaser_overload():
 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(
@@ -622,12 +652,12 @@ def test__on_response_with_leaser_overload():
     )
 
     # one message should be scheduled, the leaser capacity allows for it
-    schedule_calls = scheduler.schedule.mock_calls
-    assert len(schedule_calls) == 1
-    call_args = schedule_calls[0][1]
-    assert call_args[0] == mock.sentinel.callback
-    assert isinstance(call_args[1], message.Message)
-    assert call_args[1].message_id == "1"
+    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
@@ -641,6 +671,37 @@ def test__on_response_with_leaser_overload():
             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)
+
+    scheduler.schedule.assert_called_once_with(mock.sentinel.callback, messages)
+
+
 def test_retryable_stream_errors():
     # Make sure the config matches our hard-coded tuple of exceptions.
     interfaces = subscriber_client_config.config["interfaces"]
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
     )

From 49fa1b27c4994ab9f346cb40d9c87daa43cc5479 Mon Sep 17 00:00:00 2001
From: Peter Lamut <inbox@peterlamut.com>
Date: Fri, 17 May 2019 21:32:44 +0200
Subject: [PATCH 8/8] Add system test for subscriber batch callbacks

---
 pubsub/tests/system.py | 41 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/pubsub/tests/system.py b/pubsub/tests/system.py
index 13e81d281f42..988b5e94d4cc 100644
--- a/pubsub/tests/system.py
+++ b/pubsub/tests/system.py
@@ -28,7 +28,7 @@
 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
 
@@ -210,6 +210,45 @@ 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
     ):