Checking Pub / Sub Batch.will_accept in thread-safe way. (googleapis#4616)

dhermes · web-flow · commit 4a9a950e0336 · 2017-12-19T08:16:11.000-08:00
Fixes googleapis#4575.
diff --git a/pubsub/google/cloud/pubsub_v1/publisher/batch/thread.py b/pubsub/google/cloud/pubsub_v1/publisher/batch/thread.py
@@ -37,7 +37,7 @@ class Batch(base.Batch):
     """A batch of messages.
 
     The batch is the internal group of messages which are either awaiting
-    publication or currently in-flight.
+    publication or currently in progress.
 
     A batch is automatically created by the PublisherClient when the first
     message to be published is received; subsequent messages are added to
@@ -72,9 +72,9 @@ def __init__(self, client, topic, settings, autocommit=True):
         self._topic = topic
         self._settings = settings
 
-        self._commit_lock = threading.Lock()
-        # These objects are all communicated between threads; ensure that
-        # any writes to them are atomic.
+        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.
         self._futures = []
         self._messages = []
         self._size = 0
@@ -133,17 +133,24 @@ def status(self):
     def commit(self):
         """Actually publish all of the messages on the active batch.
 
-        This synchronously sets the batch status to in-flight, and then opens
-        a new thread, which handles actually sending the messages to Pub/Sub.
-
         .. note::
 
             This method is non-blocking. It opens a new thread, which calls
             :meth:`_commit`, which does block.
+
+        This synchronously sets the batch status to "starting", and then opens
+        a new thread, which handles actually sending the messages to Pub/Sub.
+
+        If the current batch is **not** accepting messages, this method
+        does nothing.
         """
         # Set the status to "starting" synchronously, to ensure that
         # this batch will necessarily not accept new messages.
-        self._status = base.BatchStatus.STARTING
+        with self._state_lock:
+            if self._status == base.BatchStatus.ACCEPTING_MESSAGES:
+                self._status = base.BatchStatus.STARTING
+            else:
+                return
 
         # Start a new thread to actually handle the commit.
         commit_thread = threading.Thread(
@@ -155,7 +162,7 @@ def commit(self):
     def _commit(self):
         """Actually publish all of the messages on the active batch.
 
-        This moves the batch out from being the active batch to an in-flight
+        This moves the batch out from being the active batch to an in progress
         batch on the publisher, and then the batch is discarded upon
         completion.
 
@@ -164,7 +171,7 @@ def _commit(self):
             This method blocks. The :meth:`commit` method is the non-blocking
             version, which calls this one.
         """
-        with self._commit_lock:
+        with self._state_lock:
             if self._status in _CAN_COMMIT:
                 self._status = base.BatchStatus.IN_PROGRESS
             else:
@@ -213,13 +220,13 @@ def monitor(self):
         This simply sleeps for ``self._settings.max_latency`` seconds,
         and then calls commit unless the batch has already been committed.
         """
-        # Note: This thread blocks; it is up to the calling code to call it
-        # in a separate thread.
-        #
+        # NOTE: This blocks; it is up to the calling code to call it
+        #       in a separate thread.
+
         # Sleep for however long we should be waiting.
         time.sleep(self._settings.max_latency)
 
-        # Commit.
+        _LOGGER.debug('Monitor is waking up')
         return self._commit()
 
     def publish(self, message):
@@ -235,24 +242,34 @@ def publish(self, message):
             message (~.pubsub_v1.types.PubsubMessage): The Pub/Sub message.
 
         Returns:
-            ~google.api_core.future.Future: An object conforming to
-                the :class:`concurrent.futures.Future` interface.
+            Optional[~google.api_core.future.Future]: An object conforming to
+            the :class:`~concurrent.futures.Future` interface or :data:`None`.
+            If :data:`None` is returned, that signals that the batch cannot
+            accept a message.
         """
         # Coerce the type, just in case.
         if not isinstance(message, types.PubsubMessage):
             message = types.PubsubMessage(**message)
 
-        # Add the size to the running total of the size, so we know
-        # if future messages need to be rejected.
-        self._size += message.ByteSize()
-
-        # Store the actual message in the batch's message queue.
-        self._messages.append(message)
-        if len(self._messages) >= self._settings.max_messages:
+        with self._state_lock:
+            if not self.will_accept(message):
+                return None
+
+            # Add the size to the running total of the size, so we know
+            # if future messages need to be rejected.
+            self._size += message.ByteSize()
+            # Store the actual message in the batch's message queue.
+            self._messages.append(message)
+            # Track the future on this batch (so that the result of the
+            # future can be set).
+            future = futures.Future()
+            self._futures.append(future)
+            # Determine the number of messages before releasing the lock.
+            num_messages = len(self._messages)
+
+        # Try to commit, but it must be **without** the lock held, since
+        # ``commit()`` will try to obtain the lock.
+        if num_messages >= self._settings.max_messages:
             self.commit()
 
-        # Return a Future. That future needs to be aware of the status
-        # of this batch.
-        f = futures.Future()
-        self._futures.append(f)
-        return f
+        return future
diff --git a/pubsub/google/cloud/pubsub_v1/publisher/client.py b/pubsub/google/cloud/pubsub_v1/publisher/client.py
@@ -196,9 +196,8 @@ def publish(self, topic, data, **attrs):
         batch = self.batch(topic)
         future = None
         while future is None:
-            if batch.will_accept(message):
-                future = batch.publish(message)
-            else:
+            future = batch.publish(message)
+            if future is None:
                 batch = self.batch(topic, create=True)
 
         return future
diff --git a/pubsub/tests/unit/pubsub_v1/publisher/batch/test_thread.py b/pubsub/tests/unit/pubsub_v1/publisher/batch/test_thread.py
@@ -97,6 +97,19 @@ def test_commit():
     assert batch.status == BatchStatus.STARTING
 
 
+def test_commit_no_op():
+    batch = create_batch()
+    batch._status = BatchStatus.IN_PROGRESS
+    with mock.patch.object(threading, 'Thread', autospec=True) as Thread:
+        batch.commit()
+
+    # Make sure a thread was not created.
+    Thread.assert_not_called()
+
+    # Check that batch status is unchanged.
+    assert batch.status == BatchStatus.IN_PROGRESS
+
+
 def test_blocking__commit():
     batch = create_batch()
     futures = (
@@ -217,21 +230,35 @@ def test_publish():
     )
 
     # Publish each of the messages, which should save them to the batch.
-    for message in messages:
-        batch.publish(message)
+    futures = [batch.publish(message) for message in messages]
 
     # There should be three messages on the batch, and three futures.
     assert len(batch.messages) == 3
-    assert len(batch._futures) == 3
+    assert batch._futures == futures
 
     # The size should have been incremented by the sum of the size of the
     # messages.
-    assert batch.size == sum([m.ByteSize() for m in messages])
+    expected_size = sum([message_pb.ByteSize() for message_pb in messages])
+    assert batch.size == expected_size
     assert batch.size > 0  # I do not always trust protobuf.
 
 
-def test_publish_max_messages():
-    batch = create_batch(max_messages=4)
+def test_publish_not_will_accept():
+    batch = create_batch(max_messages=0)
+
+    # Publish the message.
+    message = types.PubsubMessage(data=b'foobarbaz')
+    future = batch.publish(message)
+
+    assert future is None
+    assert batch.size == 0
+    assert batch.messages == []
+    assert batch._futures == []
+
+
+def test_publish_exceed_max_messages():
+    max_messages = 4
+    batch = create_batch(max_messages=max_messages)
     messages = (
         types.PubsubMessage(data=b'foobarbaz'),
         types.PubsubMessage(data=b'spameggs'),
@@ -240,27 +267,29 @@ def test_publish_max_messages():
 
     # Publish each of the messages, which should save them to the batch.
     with mock.patch.object(batch, 'commit') as commit:
-        for message in messages:
-            batch.publish(message)
+        futures = [batch.publish(message) for message in messages]
+        assert batch._futures == futures
+        assert len(futures) == max_messages - 1
 
         # Commit should not yet have been called.
         assert commit.call_count == 0
 
         # When a fourth message is published, commit should be called.
-        batch.publish(types.PubsubMessage(data=b'last one'))
+        future = batch.publish(types.PubsubMessage(data=b'last one'))
         commit.assert_called_once_with()
 
+        futures.append(future)
+        assert batch._futures == futures
+        assert len(futures) == max_messages
+
 
 def test_publish_dict():
     batch = create_batch()
-    batch.publish({'data': b'foobarbaz', 'attributes': {'spam': 'eggs'}})
+    future = batch.publish(
+        {'data': b'foobarbaz', 'attributes': {'spam': 'eggs'}})
 
     # There should be one message on the batch.
-    assert len(batch.messages) == 1
-
-    # It should be an actual protobuf Message at this point, with the
-    # expected values.
-    message = batch.messages[0]
-    assert isinstance(message, types.PubsubMessage)
-    assert message.data == b'foobarbaz'
-    assert message.attributes == {'spam': 'eggs'}
+    expected_message = types.PubsubMessage(
+        data=b'foobarbaz', attributes={'spam': 'eggs'})
+    assert batch.messages == [expected_message]
+    assert batch._futures == [future]
diff --git a/pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py b/pubsub/tests/unit/pubsub_v1/publisher/test_publisher_client.py
@@ -95,13 +95,20 @@ def test_publish():
     batch = mock.Mock(spec=client._batch_class)
     # Set the mock up to claim indiscriminately that it accepts all messages.
     batch.will_accept.return_value = True
+    batch.publish.side_effect = (
+        mock.sentinel.future1,
+        mock.sentinel.future2,
+    )
 
     topic = 'topic/path'
     client._batches[topic] = batch
 
     # Begin publishing.
-    client.publish(topic, b'spam')
-    client.publish(topic, b'foo', bar='baz')
+    future1 = client.publish(topic, b'spam')
+    future2 = client.publish(topic, b'foo', bar='baz')
+
+    assert future1 is mock.sentinel.future1
+    assert future2 is mock.sentinel.future2
 
     # Check mock.
     batch.publish.assert_has_calls(
@@ -138,7 +145,9 @@ def test_publish_attrs_bytestring():
     client._batches[topic] = batch
 
     # Begin publishing.
-    client.publish(topic, b'foo', bar=b'baz')
+    future = client.publish(topic, b'foo', bar=b'baz')
+
+    assert future is batch.publish.return_value
 
     # The attributes should have been sent as text.
     batch.publish.assert_called_once_with(
@@ -158,8 +167,8 @@ def test_publish_new_batch_needed():
     batch2 = mock.Mock(spec=client._batch_class)
     # Set the first mock up to claim indiscriminately that it rejects all
     # messages and the second accepts all.
-    batch1.will_accept.return_value = False
-    batch2.will_accept.return_value = True
+    batch1.publish.return_value = None
+    batch2.publish.return_value = mock.sentinel.future
 
     topic = 'topic/path'
     client._batches[topic] = batch1
@@ -168,23 +177,23 @@ def test_publish_new_batch_needed():
     batch_class = mock.Mock(spec=(), return_value=batch2)
     client._batch_class = batch_class
 
-    # Begin publishing.
-    client.publish(topic, b'foo', bar=b'baz')
+    # Publish a message.
+    future = client.publish(topic, b'foo', bar=b'baz')
+    assert future is mock.sentinel.future
+
+    # Check the mocks.
     batch_class.assert_called_once_with(
         autocommit=True,
         client=client,
         settings=client.batch_settings,
         topic=topic,
     )
-
-    # The attributes should have been sent as text.
-    batch1.publish.assert_not_called()
-    batch2.publish.assert_called_once_with(
-        types.PubsubMessage(
-            data=b'foo',
-            attributes={'bar': u'baz'},
-        ),
+    message_pb = types.PubsubMessage(
+        data=b'foo',
+        attributes={'bar': u'baz'},
     )
+    batch1.publish.assert_called_once_with(message_pb)
+    batch2.publish.assert_called_once_with(message_pb)
 
 
 def test_publish_attrs_type_error():
