googleapis · plamut · Apr 13, 2021 · Apr 13, 2021 · Apr 14, 2021
diff --git a/google/cloud/pubsub_v1/futures.py b/google/cloud/pubsub_v1/futures.py
@@ -14,6 +14,7 @@
 
 from __future__ import absolute_import
 
+import concurrent.futures
 import threading
 import uuid
 
@@ -38,51 +39,68 @@ class Future(google.api_core.future.Future):
             with that model. The ``wait()`` and ``set()`` methods will be
             used. If this argument is not provided, then a new
             :class:`threading.Event` will be created and used.
+
+        condition (Optional[Any]): A condition variable, with the same interface as
+            :class:`threading.Condition`. This is provided so that callers
+            with different concurrency models (e.g. ``threading`` or
+            ``multiprocessing``) can supply a condition that is compatible
+            with that model. The ``acquire()`` and ``release()`` methods will be
+            used. If this argument is not provided, then a new
+            :class:`threading.Condition` will be created and used.
     """
 
     # This could be a sentinel object or None, but the sentinel object's ID
     # can change if the process is forked, and None has the possibility of
     # actually being a result.
     _SENTINEL = uuid.uuid4()
 
-    def __init__(self, completed=None):
+    def __init__(self, completed=None, condition=None):
         self._result = self._SENTINEL
         self._exception = self._SENTINEL
         self._callbacks = []
         if completed is None:
             completed = threading.Event()
         self._completed = completed
 
+        # The following is here for compatibility with concurrent.futures.as_completed().
+        if condition is None:
+            condition = threading.Condition()
+        self._condition = condition
+        self._state = concurrent.futures._base.RUNNING  # Until it goes to FINISHED.
+        self._waiters = []
+
     def cancel(self):
         """Actions in Pub/Sub generally may not be canceled.
 
-        This method always returns False.
+        This method always returns ``False``.
         """
         return False
 
     def cancelled(self):
         """Actions in Pub/Sub generally may not be canceled.
 
-        This method always returns False.
+        This method always returns ``False``.
         """
         return False
 
     def running(self):
-        """Actions in Pub/Sub generally may not be canceled.
+        """A future is always considered running until it goes to the FINISHED state.
 
         Returns:
-            bool: ``True`` if this method has not yet completed, or
-                ``False`` if it has completed.
+            bool: ``True`` if the future has not yet finished, or
+                ``False`` if it has transitioned to the FINISHED state.
         """
         return not self.done()
 
     def done(self):
-        """Return True the future is done, False otherwise.
+        """Return ``True`` if the future is done, ``False`` otherwise.
 
-        This still returns True in failure cases; checking :meth:`result` or
+        This still returns ``True`` in failure cases; checking :meth:`result` or
         :meth:`exception` is the canonical way to assess success or failure.
         """
-        return self._exception != self._SENTINEL or self._result != self._SENTINEL
+        # The internal state is set to FINISHED when either of _exception or _result
+        # is set to a non-sentinel value, thus just checking the state is fine.
+        return self._state == concurrent.futures._base.FINISHED
 
     def result(self, timeout=None):
         """Resolve the future and return a value where appropriate.
@@ -121,11 +139,11 @@ def exception(self, timeout=None):
         if not self._completed.wait(timeout=timeout):
             raise exceptions.TimeoutError("Timed out waiting for result.")
 
-        # If the batch completed successfully, this should return None.
+        # If the corresponding batch completed successfully, this should return None.
         if self._result != self._SENTINEL:
             return None
 
-        # Okay, this batch had an error; this should return it.
+        # Okay, the corresponding batch had an error; this should return it.
         return self._exception
 
     def add_done_callback(self, callback):
@@ -140,47 +158,51 @@ def add_done_callback(self, callback):
         Returns:
             None
         """
-        if self.done():
-            return callback(self)
-        self._callbacks.append(callback)
+        with self._condition:
+            if self._state != concurrent.futures._base.FINISHED:
+                self._callbacks.append(callback)
+                return
+
+        return callback(self)
 
     def set_result(self, result):
         """Set the result of the future to the provided result.
 
         Args:
             result (Any): The result
         """
-        # Sanity check: A future can only complete once.
-        if self.done():
-            raise RuntimeError("set_result can only be called once.")
+        with self._condition:
+            # Sanity check: A future can only complete once.
+            if self._state == concurrent.futures._base.FINISHED:
+                raise RuntimeError("set_result can only be called once.")
+
+            self._result = result
+            self._state = concurrent.futures._base.FINISHED
+            self._completed.set()
 
-        # Set the result and trigger the future.
-        self._result = result
-        self._trigger()
+            for waiter in self._waiters:
+                waiter.add_result(self)
+
+        for callback in self._callbacks:
+            callback(self)
 
     def set_exception(self, exception):
         """Set the result of the future to the given exception.
 
         Args:
             exception (:exc:`Exception`): The exception raised.
         """
-        # Sanity check: A future can only complete once.
-        if self.done():
-            raise RuntimeError("set_exception can only be called once.")
+        with self._condition:
+            # Sanity check: A future can only complete once.
+            if self._state == concurrent.futures._base.FINISHED:
+                raise RuntimeError("set_exception can only be called once.")
 
-        # Set the exception and trigger the future.
-        self._exception = exception
-        self._trigger()
+            self._exception = exception
+            self._state = concurrent.futures._base.FINISHED
+            self._completed.set()
 
-    def _trigger(self):
-        """Trigger all callbacks registered to this Future.
+            for waiter in self._waiters:
+                waiter.add_exception(self)
 
-        This method is called internally by the batch once the batch
-        completes.
-
-        Args:
-            message_id (str): The message ID, as a string.
-        """
-        self._completed.set()
         for callback in self._callbacks:
             callback(self)
diff --git a/google/cloud/pubsub_v1/publisher/futures.py b/google/cloud/pubsub_v1/publisher/futures.py
@@ -43,10 +43,4 @@ def result(self, timeout=None):
             Exception: For undefined exceptions in the underlying
                 call execution.
         """
-        # Attempt to get the exception if there is one.
-        # If there is not one, then we know everything worked, and we can
-        # return an appropriate value.
-        err = self.exception(timeout=timeout)
-        if err is None:
-            return self._result
-        raise err
+        return super().result(timeout=timeout)
diff --git a/tests/unit/pubsub_v1/test_futures.py b/tests/unit/pubsub_v1/test_futures.py
@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import concurrent.futures
 import threading
+import time
 
 import mock
 import pytest
@@ -26,25 +28,36 @@ def _future(*args, **kwargs):
 
 
 def test_constructor_defaults():
-    with mock.patch.object(threading, "Event", autospec=True) as Event:
+    event_patcher = mock.patch.object(threading, "Event", autospec=True)
+    condition_patcher = mock.patch.object(threading, "Condition", autospec=True)
+
+    with event_patcher as Event, condition_patcher as Condition:
         future = _future()
 
     assert future._result == futures.Future._SENTINEL
     assert future._exception == futures.Future._SENTINEL
     assert future._callbacks == []
     assert future._completed is Event.return_value
+    assert future._condition is Condition.return_value
+    assert future._state == concurrent.futures._base.RUNNING
+    assert future._waiters == []
 
     Event.assert_called_once_with()
+    Condition.assert_called_once_with()
 
 
 def test_constructor_explicit_completed():
     completed = mock.sentinel.completed
-    future = _future(completed=completed)
+    condition = mock.sentinel.condition
+    future = _future(completed=completed, condition=condition)
 
     assert future._result == futures.Future._SENTINEL
     assert future._exception == futures.Future._SENTINEL
     assert future._callbacks == []
     assert future._completed is completed
+    assert future._condition is condition
+    assert future._state == concurrent.futures._base.RUNNING
+    assert future._waiters == []
 
 
 def test_cancel():
@@ -146,3 +159,45 @@ def test_set_exception_once_only():
     future.set_exception(ValueError("wah wah"))
     with pytest.raises(RuntimeError):
         future.set_exception(TypeError("other wah wah"))
+
+
+def test_as_completed_compatibility():
+    all_futures = {i: _future() for i in range(6)}
+    done_futures = []
+
+    def resolve_future(future_idx, delay=0):
+        time.sleep(delay)
+        future = all_futures[future_idx]
+        if future_idx % 2 == 0:
+            future.set_result(f"{future_idx}: I'm done!")
+        else:
+            future.set_exception(Exception(f"Future {future_idx} errored"))
+
+    all_futures[2].set_result("2: I'm done!")
+
+    # Start marking the futures as completed (either with success or error) at
+    # different times and check that ther "as completed" order is correct.
+    for future_idx, delay in ((0, 0.8), (3, 0.6), (1, 0.4), (5, 0.2)):
+        threading.Thread(
+            target=resolve_future, args=(future_idx, delay), daemon=True
+        ).start()
+
+    try:
+        # Use a loop instead of a list comprehension to gather futures completed
+        # before the timeout error occurs.
+        for future in concurrent.futures.as_completed(all_futures.values(), timeout=1):
+            done_futures.append(future)
+    except concurrent.futures.TimeoutError:
+        pass
+    else:  # pragma: NO COVER
+        pytest.fail("Not all Futures should have been recognized as completed.")
+
+    # NOTE: Future 4 was never resolved.
+    expected = [
+        all_futures[2],
+        all_futures[5],
+        all_futures[1],
+        all_futures[3],
+        all_futures[0],
+    ]
+    assert done_futures == expected