Repository: googleapis/google-cloud-python

Issue:
Pubsub: How to flush unsent messages on program exit?
I have set caching configuration which may cache some messages and messages are not sent to pubsub yet. Is there a `flush` method when I want to exit the program.


Discussion:
A workaround is to use [Seek](https://cloud.google.com/pubsub/docs/replay-qs#seek_to_a_timestamp) and seek to the current timestamp. This way, messages that were published before this time are marked as acknowledged, and messages retained in the subscription that were published after this time are marked as unacknowledged. 
@plamut do you have any advice about this question?
I'm not 100% sure if I understand the question correctly, but does "caching" refer to the [batch settings](https://google-cloud.readthedocs.io/en/latest/pubsub/types.html#google.cloud.pubsub_v1.types.BatchSettings) used when publishing messages? E.g the `max_latency` setting that sets how long the publisher client should wait before actually publishing a batch of messages to the server (for messages added by the [publisher.publish()](https://github.com/googleapis/google-cloud-python/blob/master/pubsub/google/cloud/pubsub_v1/publisher/client.py#L172) method)?

If so, one way of forcing the publisher to actually send the messages to the server (e.g. on program exit) is to manually invoke the [batch.commit()](https://github.com/googleapis/google-cloud-python/blob/master/pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py#L140) method.

The publisher client internally stores message batches for each topic [here](https://github.com/googleapis/google-cloud-python/blob/master/pubsub/google/cloud/pubsub_v1/publisher/client.py#L168), thus forcing a batch to commit messages early could be done with the following:

```py
topic_batch = publisher._batches.get(TOPIC_PATH)
if topic_batch is not None:
    topic_batch.commit()
```
(invoking `commit()` is thread-safe, and will simply do nothing if batch commit is already in progress)

---

@mckingho Does that answer your question? Or maybe what @anguillanneuf wrote earlier? If not and "caching configuration" refers to something else, we would need a bit more info.

Looking forward to hearing again from you!
@plamut: I think the user wants to have a method to ensure that all messages have successfully been sent... i.e. all futures have completed.
We may need documentation (or a new method) for how to wait for shut down until all futures are complete.

Resolving patch:
diff --git a/pubsub/google/cloud/pubsub_v1/futures.py b/pubsub/google/cloud/pubsub_v1/futures.py
--- a/pubsub/google/cloud/pubsub_v1/futures.py
+++ b/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.
diff --git a/pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py b/pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py
--- a/pubsub/google/cloud/pubsub_v1/publisher/_batch/thread.py
+++ b/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
diff --git a/pubsub/google/cloud/pubsub_v1/publisher/client.py b/pubsub/google/cloud/pubsub_v1/publisher/client.py
--- a/pubsub/google/cloud/pubsub_v1/publisher/client.py
+++ b/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):
@@ -187,20 +188,19 @@ def _batch(self, topic, create=False, autocommit=True):
         """
         # If there is no matching batch yet, then potentially create one
         # and place it on the batches dictionary.
-        with self._batch_lock:
-            if not create:
-                batch = self._batches.get(topic)
-                if batch is None:
-                    create = True
-
-            if create:
-                batch = self._batch_class(
-                    autocommit=autocommit,
-                    client=self,
-                    settings=self.batch_settings,
-                    topic=topic,
-                )
-                self._batches[topic] = batch
+        if not create:
+            batch = self._batches.get(topic)
+            if batch is None:
+                create = True
+
+        if create:
+            batch = self._batch_class(
+                autocommit=autocommit,
+                client=self,
+                settings=self.batch_settings,
+                topic=topic,
+            )
+            self._batches[topic] = batch
 
         return batch
 
@@ -242,12 +242,17 @@ 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:
+            RuntimeError:
+                If called after publisher has been stopped
+                by a `stop()` method call.
         """
         # 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):
             raise TypeError(
-                "Data being published to Pub/Sub must be sent " "as a bytestring."
+                "Data being published to Pub/Sub must be sent as a bytestring."
             )
 
         # Coerce all attributes to text strings.
@@ -266,11 +271,38 @@ def publish(self, topic, data, **attrs):
         message = types.PubsubMessage(data=data, attributes=attrs)
 
         # Delegate the publishing to the batch.
-        batch = self._batch(topic)
-        future = None
-        while future is None:
-            future = batch.publish(message)
-            if future is None:
-                batch = self._batch(topic, create=True)
+        with self._batch_lock:
+            if self._is_stopped:
+                raise RuntimeError("Cannot publish on a stopped publisher.")
+
+            batch = self._batch(topic)
+            future = None
+            while future is None:
+                future = batch.publish(message)
+                if future is None:
+                    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.
+        """
+        with self._batch_lock:
+            if self._is_stopped:
+                raise RuntimeError("Cannot stop a publisher already stopped.")
+
+            self._is_stopped = True
+
+            for batch in self._batches.values():
+                batch.commit()


Test patch:
diff --git a/pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py b/pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py
--- a/pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py
+++ b/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(RuntimeError):
+        client.publish("topic1", b"msg2")
+
+    with pytest.raises(RuntimeError):
+        client.stop()
+
+
 def test_gapic_instance_method():
     creds = mock.Mock(spec=credentials.Credentials)
     client = publisher.Client(credentials=creds)