PubSub: Release the state lock before calling the publish api

pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py

@@ -187,56 +187,62 @@ def _commit(self):
                 _LOGGER.debug("Batch is already in progress, exiting commit")
                 return
 
-            # Sanity check: If there are no messages, no-op.
-            if not self._messages:
-                _LOGGER.debug("No messages to publish, exiting commit")
-                self._status = base.BatchStatus.SUCCESS
-                return
-
-            # Begin the request to publish these messages.
-            # Log how long the underlying request takes.
-            start = time.time()
-
-            try:
-                response = self._client.api.publish(self._topic, self._messages)
-            except google.api_core.exceptions.GoogleAPIError as exc:
-                # We failed to publish, set the exception on all futures and
-                # exit.
-                self._status = base.BatchStatus.ERROR
-
-                for future in self._futures:
-                    future.set_exception(exc)
+        # Once in the IN_PROGRESS state, no other thread can publish additional
+        # messages or initiate a commit (those operations becme a no-op), thus
+        # it is safe to release the state lock here. Releasing the lock avoids
+        # blocking other threads in case api.publish() below takes a long time
+        # to complete.
+
+        # Sanity check: If there are no messages, no-op.
+        if not self._messages:
+            _LOGGER.debug("No messages to publish, exiting commit")
+            self._status = base.BatchStatus.SUCCESS
+            return
+
+        # Begin the request to publish these messages.
+        # Log how long the underlying request takes.
+        start = time.time()
+
+        try:
+            response = self._client.api.publish(self._topic, self._messages)
+        except google.api_core.exceptions.GoogleAPIError as exc:
+            # We failed to publish, set the exception on all futures and
+            # exit.
+            self._status = base.BatchStatus.ERROR
+
+            for future in self._futures:
+                future.set_exception(exc)
+
+            _LOGGER.exception("Failed to publish %s messages.", len(self._futures))
+            return
+
+        end = time.time()
+        _LOGGER.debug("gRPC Publish took %s seconds.", end - start)
+
+        if len(response.message_ids) == len(self._futures):
+            # Iterate over the futures on the queue and return the response
+            # IDs. We are trusting that there is a 1:1 mapping, and raise
+            # an exception if not.
+            self._status = base.BatchStatus.SUCCESS
+            zip_iter = six.moves.zip(response.message_ids, self._futures)
+            for message_id, future in zip_iter:
+                future.set_result(message_id)
+        else:
+            # Sanity check: If the number of message IDs is not equal to
+            # the number of futures I have, then something went wrong.
+            self._status = base.BatchStatus.ERROR
+            exception = exceptions.PublishError(
+                "Some messages were not successfully published."
+            )
 
-                _LOGGER.exception("Failed to publish %s messages.", len(self._futures))
-                return
+            for future in self._futures:
+                future.set_exception(exception)
 
-            end = time.time()
-            _LOGGER.debug("gRPC Publish took %s seconds.", end - start)
-
-            if len(response.message_ids) == len(self._futures):
-                # Iterate over the futures on the queue and return the response
-                # IDs. We are trusting that there is a 1:1 mapping, and raise
-                # an exception if not.
-                self._status = base.BatchStatus.SUCCESS
-                zip_iter = six.moves.zip(response.message_ids, self._futures)
-                for message_id, future in zip_iter:
-                    future.set_result(message_id)
-            else:
-                # Sanity check: If the number of message IDs is not equal to
-                # the number of futures I have, then something went wrong.
-                self._status = base.BatchStatus.ERROR
-                exception = exceptions.PublishError(
-                    "Some messages were not successfully published."
-                )
-
-                for future in self._futures:
-                    future.set_exception(exception)
-
-                _LOGGER.error(
-                    "Only %s of %s messages were published.",
-                    len(response.message_ids),
-                    len(self._futures),
-                )
+            _LOGGER.error(
+                "Only %s of %s messages were published.",
+                len(response.message_ids),
+                len(self._futures),
+            )
 
     def monitor(self):
         """Commit this batch after sufficient time has elapsed.

pubsub/tests/unit/pubsub_v1/publisher/batch/test_thread.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import threading
 import time
 
 import mock
+import pytest
 
 import google.api_core.exceptions
 from google.auth import credentials
@@ -39,7 +41,7 @@ def create_batch(autocommit=False, **batch_settings):
         autocommit (bool): Whether the batch should commit after
             ``max_latency`` seconds. By default, this is ``False``
             for unit testing.
-        kwargs (dict): Arguments passed on to the
+        batch_settings (dict): Arguments passed on to the
             :class:``~.pubsub_v1.types.BatchSettings`` constructor.
 
     Returns:
@@ -149,6 +151,35 @@ def test_blocking__commit():
     assert futures[1].result() == "b"
 
 
+def test_client_api_publish_not_blocking_additional_publish_calls():
+    batch = create_batch(max_messages=1)
+    api_publish_called = threading.Event()
+
+    def api_publish_delay(_, messages):
+        api_publish_called.set()
+        time.sleep(1.0)
+        message_ids = [str(i) for i in range(len(messages))]
+        return types.PublishResponse(message_ids=message_ids)
+
+    api_publish_patch = mock.patch.object(
+        type(batch.client.api), "publish", side_effect=api_publish_delay
+    )
+
+    with api_publish_patch:
+        batch.publish({"data": b"first message"})
+
+        start = datetime.datetime.now()
+        event_set = api_publish_called.wait(timeout=1.0)
+        if not event_set:
+            pytest.fail("API publish was not called in time")
+        batch.publish({"data": b"second message"})
+        end = datetime.datetime.now()
+
+    # While a batch commit in progress, waiting for the API publish call to
+    # complete should not unnecessariliy delay other calls to batch.publish().
+    assert (end - start).total_seconds() < 1.0
+
+
 @mock.patch.object(thread, "_LOGGER")
 def test_blocking__commit_starting(_LOGGER):
     batch = create_batch()