diff --git a/google/cloud/pubsub_v1/futures.py b/google/cloud/pubsub_v1/futures.py index ba861e40c..8243bf146 100644 --- 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,6 +39,14 @@ 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 @@ -45,7 +54,7 @@ class Future(google.api_core.future.Future): # 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 = [] @@ -53,36 +62,45 @@ def __init__(self, completed=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,9 +158,12 @@ 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. @@ -150,13 +171,20 @@ def set_result(self, 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. @@ -164,23 +192,17 @@ def set_exception(self, 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 index fa8a79998..cd7812812 100644 --- 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 index 11349d5d4..a6c1f1243 100644 --- 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