feat: make publish futures compatible with concurrent.futures.as_completed()

google/cloud/pubsub_v1/futures.py

@@ -14,173 +14,45 @@
 
 from __future__ import absolute_import
 
-import threading
-import uuid
+import concurrent.futures
 
 import google.api_core.future
-from google.cloud.pubsub_v1.publisher import exceptions
 
 
-class Future(google.api_core.future.Future):
+class Future(concurrent.futures.Future, google.api_core.future.Future):
     """Encapsulation of the asynchronous execution of an action.
 
     This object is returned from asychronous Pub/Sub calls, and is the
     interface to determine the status of those calls.
 
     This object should not be created directly, but is returned by other
     methods in this library.
-
-    Args:
-        completed (Optional[Any]): An event, with the same interface as
-            :class:`threading.Event`. This is provided so that callers
-            with different concurrency models (e.g. ``threading`` or
-            ``multiprocessing``) can supply an event that is compatible
-            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.
     """
 
-    # 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):
-        self._result = self._SENTINEL
-        self._exception = self._SENTINEL
-        self._callbacks = []
-        if completed is None:
-            completed = threading.Event()
-        self._completed = completed
-
-    def cancel(self):
-        """Actions in Pub/Sub generally may not be canceled.
-
-        This method always returns False.
-        """
-        return False
-
-    def cancelled(self):
-        """Actions in Pub/Sub generally may not be canceled.
-
-        This method always returns False.
-        """
-        return False
-
     def running(self):
-        """Actions in Pub/Sub generally may not be canceled.
+        """Return ``True`` if the associated Pub/Sub action has not yet completed.
 
-        Returns:
-            bool: ``True`` if this method has not yet completed, or
-                ``False`` if it has completed.
+        Returns: bool:
         """
         return not self.done()
 
-    def done(self):
-        """Return True the future is done, False otherwise.
-
-        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
-
-    def result(self, timeout=None):
-        """Resolve the future and return a value where appropriate.
-
-        Args:
-            timeout (Union[int, float]): The number of seconds before this call
-                times out and raises TimeoutError.
-
-        Raises:
-            concurrent.futures.TimeoutError: If the request times out.
-            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
-
-    def exception(self, timeout=None):
-        """Return the exception raised by the call, if any.
-
-        Args:
-            timeout (Union[int, float]): The number of seconds before this call
-                times out and raises TimeoutError.
-
-        Raises:
-            concurrent.futures.TimeoutError: If the request times out.
-
-        Returns:
-            Exception: The exception raised by the call, if any.
-        """
-        # Wait until the future is done.
-        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 self._result != self._SENTINEL:
-            return None
-
-        # Okay, this batch had an error; this should return it.
-        return self._exception
-
-    def add_done_callback(self, callback):
-        """Attach the provided callable to the future.
-
-        The provided function is called, with this future as its only argument,
-        when the future finishes running.
-
-        Args:
-            callback (Callable): The function to call.
-
-        Returns:
-            None
-        """
-        if self.done():
-            return callback(self)
-        self._callbacks.append(callback)
+    def set_running_or_notify_cancel(self):
+        raise NotImplementedError(
+            "Only used by executors from `concurrent.futures` package."
+        )
 
     def set_result(self, result):
-        """Set the result of the future to the provided result.
+        """Set the return value of work associated with the future.
 
-        Args:
-            result (Any): The result
+        Do not use this method, it should only be used internally by the library and its
+        unit tests.
         """
-        # Sanity check: A future can only complete once.
-        if self.done():
-            raise RuntimeError("set_result can only be called once.")
-
-        # Set the result and trigger the future.
-        self._result = result
-        self._trigger()
+        return super().set_result(result=result)
 
     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.")
-
-        # Set the exception and trigger the future.
-        self._exception = exception
-        self._trigger()
-
-    def _trigger(self):
-        """Trigger all callbacks registered to this Future.
-
-        This method is called internally by the batch once the batch
-        completes.
+        """Set the result of the future as being the given exception.
 
-        Args:
-            message_id (str): The message ID, as a string.
+        Do not use this method, it should only be used internally by the library and its
+        unit tests.
         """
-        self._completed.set()
-        for callback in self._callbacks:
-            callback(self)
+        return super().set_exception(exception=exception)

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

@@ -378,7 +378,7 @@ def publish(self, message):
 
                 # Track the future on this batch (so that the result of the
                 # future can be set).
-                future = futures.Future(completed=threading.Event())
+                future = futures.Future()
                 self._futures.append(future)
 
         # Try to commit, but it must be **without** the lock held, since

google/cloud/pubsub_v1/publisher/futures.py

@@ -25,6 +25,20 @@ class Future(futures.Future):
     ID, unless an error occurs.
     """
 
+    def cancel(self):
+        """Actions in Pub/Sub generally may not be canceled.
+
+        This method always returns ``False``.
+        """
+        return False
+
+    def cancelled(self):
+        """Actions in Pub/Sub generally may not be canceled.
+
+        This method always returns ``False``.
+        """
+        return False
+
     def result(self, timeout=None):
         """Return the message ID or raise an exception.
 
@@ -43,10 +57,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)

google/cloud/pubsub_v1/subscriber/futures.py

@@ -28,9 +28,9 @@ class StreamingPullFuture(futures.Future):
 
     def __init__(self, manager):
         super(StreamingPullFuture, self).__init__()
-        self._manager = manager
-        self._manager.add_close_callback(self._on_close_callback)
-        self._cancelled = False
+        self.__manager = manager
+        self.__manager.add_close_callback(self._on_close_callback)
+        self.__cancelled = False
 
     def _on_close_callback(self, manager, result):
         if self.done():
@@ -47,12 +47,14 @@ def cancel(self):
         """Stops pulling messages and shutdowns the background thread consuming
         messages.
         """
-        self._cancelled = True
-        return self._manager.close()
+        # NOTE: We circumvent the base future's self._state to track the cancellation
+        # state, as this state has different meaning with streaming pull futures.
+        self.__cancelled = True
+        return self.__manager.close()
 
     def cancelled(self):
         """
         returns:
             bool: ``True`` if the subscription has been cancelled.
         """
-        return self._cancelled
+        return self.__cancelled

tests/unit/pubsub_v1/publisher/test_futures_publisher.py

@@ -20,6 +20,14 @@
 
 
 class TestFuture(object):
+    def test_cancel(self):
+        future = futures.Future()
+        assert future.cancel() is False
+
+    def test_cancelled(self):
+        future = futures.Future()
+        assert future.cancelled() is False
+
     def test_result_on_success(self):
         future = futures.Future()
         future.set_result("570307942214048")

tests/unit/pubsub_v1/subscriber/test_futures_subscriber.py

@@ -31,13 +31,12 @@ def make_future(self):
 
     def test_default_state(self):
         future = self.make_future()
+        manager = future._StreamingPullFuture__manager
 
         assert future.running()
         assert not future.done()
         assert not future.cancelled()
-        future._manager.add_close_callback.assert_called_once_with(
-            future._on_close_callback
-        )
+        manager.add_close_callback.assert_called_once_with(future._on_close_callback)
 
     def test__on_close_callback_success(self):
         future = self.make_future()
@@ -71,8 +70,9 @@ def test__on_close_callback_future_already_done(self):
 
     def test_cancel(self):
         future = self.make_future()
+        manager = future._StreamingPullFuture__manager
 
         future.cancel()
 
-        future._manager.close.assert_called_once()
+        manager.close.assert_called_once()
         assert future.cancelled()

tests/unit/pubsub_v1/subscriber/test_subscriber_client.py

@@ -137,7 +137,8 @@ def test_subscribe(manager_open, creds):
     future = client.subscribe("sub_name_a", callback=mock.sentinel.callback)
     assert isinstance(future, futures.StreamingPullFuture)
 
-    assert future._manager._subscription == "sub_name_a"
+    manager = future._StreamingPullFuture__manager
+    assert manager._subscription == "sub_name_a"
     manager_open.assert_called_once_with(
         mock.ANY,
         callback=mock.sentinel.callback,
@@ -164,10 +165,11 @@ def test_subscribe_options(manager_open, creds):
     )
     assert isinstance(future, futures.StreamingPullFuture)
 
-    assert future._manager._subscription == "sub_name_a"
-    assert future._manager.flow_control == flow_control
-    assert future._manager._scheduler == scheduler
-    assert future._manager._await_callbacks_on_shutdown is mock.sentinel.await_callbacks
+    manager = future._StreamingPullFuture__manager
+    assert manager._subscription == "sub_name_a"
+    assert manager.flow_control == flow_control
+    assert manager._scheduler == scheduler
+    assert manager._await_callbacks_on_shutdown is mock.sentinel.await_callbacks
     manager_open.assert_called_once_with(
         mock.ANY,
         callback=mock.sentinel.callback,