feat(pubsub): add stop method

pubsub/google/cloud/pubsub_v1/futures.py

@@ -74,9 +74,7 @@ def running(self):
             bool: ``True`` if this method has not yet completed, or
                 ``False`` if it has completed.
         """
-        if self.done():
-            return False
-        return True
+        return not self.done()
 
     def done(self):
         """Return True the future is done, False otherwise.

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

@@ -74,6 +74,9 @@ def __init__(self, client, topic, settings, autocommit=True):
         self._state_lock = threading.Lock()
         # These members are all communicated between threads; ensure that
         # any writes to them use the "state lock" to remain atomic.
+        # _futures list should remain unchanged after batch
+        # status changed from ACCEPTING_MESSAGES to any other
+        # in order to avoid race conditions
         self._futures = []
         self._messages = []
         self._size = 0

pubsub/google/cloud/pubsub_v1/publisher/client.py

@@ -134,6 +134,7 @@ def __init__(self, batch_settings=(), **kwargs):
         # messages. One batch exists for each topic.
         self._batch_lock = self._batch_class.make_lock()
         self._batches = {}
+        self._is_stopped = False
 
     @classmethod
     def from_service_account_file(cls, filename, batch_settings=(), **kwargs):
@@ -242,7 +243,14 @@ def publish(self, topic, data, **attrs):
             instance that conforms to Python Standard library's
             :class:`~concurrent.futures.Future` interface (but not an
             instance of that class).
+
+        Raises:
+            ValueError:
+                If called after publisher has been stopped
+                by a `stop()` method call.
         """
+        if self._is_stopped:
+            raise ValueError("Cannot publish on a stopped publisher.")
         # Sanity check: Is the data being sent as a bytestring?
         # If it is literally anything else, complain loudly about it.
         if not isinstance(data, six.binary_type):
@@ -274,3 +282,24 @@ def publish(self, topic, data, **attrs):
                 batch = self._batch(topic, create=True)
 
         return future
+
+    def stop(self):
+        """Immediately publish all outstanding messages.
+
+        Asynchronously sends all outstanding messages and
+        prevents future calls to `publish()`. Method should
+        be invoked prior to deleting this `Client()` object
+        in order to ensure that no pending messages are lost.
+
+        .. note::
+
+            This method is non-blocking. Use `Future()` objects
+            returned by `publish()` to make sure all publish
+            requests completed, either in success or error.
+        """
+        if self._is_stopped:
+            raise ValueError("Cannot stop a publisher already stopped.")
+
+        self._is_stopped = True
+        for topic in self._batches:
+            self._batches[topic].commit()

pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py

@@ -201,6 +201,36 @@ def test_publish_attrs_type_error():
         client.publish(topic, b"foo", answer=42)
 
 
+def test_stop():
+    creds = mock.Mock(spec=credentials.Credentials)
+    client = publisher.Client(credentials=creds)
+
+    batch = client._batch("topic1", autocommit=False)
+    batch2 = client._batch("topic2", autocommit=False)
+
+    pubsub_msg = types.PubsubMessage(data=b"msg")
+
+    patch = mock.patch.object(batch, "commit")
+    patch2 = mock.patch.object(batch2, "commit")
+
+    with patch as commit_mock, patch2 as commit_mock2:
+        batch.publish(pubsub_msg)
+        batch2.publish(pubsub_msg)
+
+        client.stop()
+
+        # check if commit() called
+        commit_mock.assert_called()
+        commit_mock2.assert_called()
+
+    # check that closed publisher doesn't accept new messages
+    with pytest.raises(ValueError):
+        client.publish("topic1", pubsub_msg)
+
+    with pytest.raises(ValueError):
+        client.stop()
+
+
 def test_gapic_instance_method():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)