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

* feat: make futures compatible with as_completed() The futures implementation is adjusted to work well with the built-in function with the same name in `concurrent.futures` package. * Fix two unit tests in pre-Python 3.8 If setting a result/exception on a concurrent.futures.Future object, an exception is raised only in Python3.8+, thus we conditionally disable two unit tests. This behavior change is fine, though, because users should never use the set_result() and set_exception() methods directly. * Cover missing code line with a test * Use double underscore for internal cancelled flag * Prefix manager reference with double underscore * Remove Future's completed parameter altogether This parameter is unlikely to be used by any 3rd party code, but even if it is, it's better to cause a loud error rather than silently changing its effect to a no-op.

Author: plamut

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

@@ -380,7 +380,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,