feat: support for async bidi streaming apis

google/api_core/bidi.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Bi-directional streaming RPC helpers."""
+"""Helpers for synchronous bidirectional streaming RPCs."""
 
 import collections
 import datetime
@@ -22,6 +22,7 @@
 import time
 
 from google.api_core import exceptions
+from google.api_core.bidi_base import BidiRpcBase
 
 _LOGGER = logging.getLogger(__name__)
 _BIDIRECTIONAL_CONSUMER_NAME = "Thread-ConsumeBidirectionalStream"
@@ -36,21 +37,6 @@ class _RequestQueueGenerator(object):
     otherwise open-ended set of requests to send through a request-streaming
     (or bidirectional) RPC.
 
-    The reason this is necessary is because gRPC takes an iterator as the
-    request for request-streaming RPCs. gRPC consumes this iterator in another
-    thread to allow it to block while generating requests for the stream.
-    However, if the generator blocks indefinitely gRPC will not be able to
-    clean up the thread as it'll be blocked on `next(iterator)` and not be able
-    to check the channel status to stop iterating. This helper mitigates that
-    by waiting on the queue with a timeout and checking the RPC state before
-    yielding.
-
-    Finally, it allows for retrying without swapping queues because if it does
-    pull an item off the queue when the RPC is inactive, it'll immediately put
-    it back and then exit. This is necessary because yielding the item in this
-    case will cause gRPC to discard it. In practice, this means that the order
-    of messages is not guaranteed. If such a thing is necessary it would be
-    easy to use a priority queue.
 
     Example::
 
@@ -62,12 +48,6 @@ class _RequestQueueGenerator(object):
             print(response)
             q.put(...)
 
-    Note that it is possible to accomplish this behavior without "spinning"
-    (using a queue timeout). One possible way would be to use more threads to
-    multiplex the grpc end event with the queue, another possible way is to
-    use selectors and a custom event/queue object. Both of these approaches
-    are significant from an engineering perspective for small benefit - the
-    CPU consumed by spinning is pretty minuscule.
 
     Args:
         queue (queue_module.Queue): The request queue.
@@ -96,6 +76,31 @@ def _is_active(self):
         return self.call is None or self.call.is_active()
 
     def __iter__(self):
+        # The reason this is necessary is because gRPC takes an iterator as the
+        # request for request-streaming RPCs. gRPC consumes this iterator in
+        # another thread to allow it to block while generating requests for
+        # the stream. However, if the generator blocks indefinitely gRPC will
+        # not be able to clean up the thread as it'll be blocked on
+        # `next(iterator)` and not be able to check the channel status to stop
+        # iterating. This helper mitigates that by waiting on the queue with
+        # a timeout and checking the RPC state before yielding.
+        #
+        # Finally, it allows for retrying without swapping queues because if
+        # it does pull an item off the queue when the RPC is inactive, it'll
+        # immediately put it back and then exit. This is necessary because
+        # yielding the item in this case will cause gRPC to discard it. In
+        # practice, this means that the order of messages is not guaranteed.
+        # If such a thing is necessary it would be easy to use a priority
+        # queue.
+        #
+        # Note that it is possible to accomplish this behavior without
+        # "spinning" (using a queue timeout). One possible way would be to use
+        # more threads to multiplex the grpc end event with the queue, another
+        # possible way is to use selectors and a custom event/queue object.
+        # Both of these approaches are significant from an engineering
+        # perspective for small benefit - the CPU consumed by spinning is
+        # pretty minuscule.
+
         if self._initial_request is not None:
             if callable(self._initial_request):
                 yield self._initial_request()
@@ -201,7 +206,7 @@ def __repr__(self):
         )
 
 
-class BidiRpc(object):
+class BidiRpc(BidiRpcBase):
     """A helper for consuming a bi-directional streaming RPC.
 
     This maps gRPC's built-in interface which uses a request iterator and a
@@ -227,6 +232,8 @@ class BidiRpc(object):
             rpc.send(example_pb2.StreamingRpcRequest(
                 data='example'))
 
+        rpc.close()
+
     This does *not* retry the stream on errors. See :class:`ResumableBidiRpc`.
 
     Args:
@@ -240,40 +247,14 @@ class BidiRpc(object):
             the request.
     """
 
-    def __init__(self, start_rpc, initial_request=None, metadata=None):
-        self._start_rpc = start_rpc
-        self._initial_request = initial_request
-        self._rpc_metadata = metadata
-        self._request_queue = queue_module.Queue()
-        self._request_generator = None
-        self._is_active = False
-        self._callbacks = []
-        self.call = None
-
-    def add_done_callback(self, callback):
-        """Adds a callback that will be called when the RPC terminates.
-
-        This occurs when the RPC errors or is successfully terminated.
-
-        Args:
-            callback (Callable[[grpc.Future], None]): The callback to execute.
-                It will be provided with the same gRPC future as the underlying
-                stream which will also be a :class:`grpc.Call`.
-        """
-        self._callbacks.append(callback)
-
-    def _on_call_done(self, future):
-        # This occurs when the RPC errors or is successfully terminated.
-        # Note that grpc's "future" here can also be a grpc.RpcError.
-        # See note in https://github.com/grpc/grpc/issues/10885#issuecomment-302651331
-        # that `grpc.RpcError` is also `grpc.call`.
-        for callback in self._callbacks:
-            callback(future)
+    def _create_queue(self):
+        """Create a queue for requests."""
+        return queue_module.Queue()
 
     def open(self):
         """Opens the stream."""
         if self.is_active:
-            raise ValueError("Can not open an already open stream.")
+            raise ValueError("Cannot open an already open stream.")
 
         request_generator = _RequestQueueGenerator(
             self._request_queue, initial_request=self._initial_request
@@ -322,7 +303,7 @@ def send(self, request):
             request (protobuf.Message): The request to send.
         """
         if self.call is None:
-            raise ValueError("Can not send() on an RPC that has never been open()ed.")
+            raise ValueError("Cannot send on an RPC stream that has never been opened.")
 
         # Don't use self.is_active(), as ResumableBidiRpc will overload it
         # to mean something semantically different.
@@ -343,20 +324,15 @@ def recv(self):
             protobuf.Message: The received message.
         """
         if self.call is None:
-            raise ValueError("Can not recv() on an RPC that has never been open()ed.")
+            raise ValueError("Cannot recv on an RPC stream that has never been opened.")
 
         return next(self.call)
 
     @property
     def is_active(self):
-        """bool: True if this stream is currently open and active."""
+        """True if this stream is currently open and active."""
         return self.call is not None and self.call.is_active()
 
-    @property
-    def pending_requests(self):
-        """int: Returns an estimate of the number of queued requests."""
-        return self._request_queue.qsize()
-
 
 def _never_terminate(future_or_error):
     """By default, no errors cause BiDi termination."""
@@ -544,7 +520,7 @@ def _send(self, request):
             call = self.call
 
         if call is None:
-            raise ValueError("Can not send() on an RPC that has never been open()ed.")
+            raise ValueError("Cannot send on an RPC that has never been opened.")
 
         # Don't use self.is_active(), as ResumableBidiRpc will overload it
         # to mean something semantically different.
@@ -563,7 +539,7 @@ def _recv(self):
             call = self.call
 
         if call is None:
-            raise ValueError("Can not recv() on an RPC that has never been open()ed.")
+            raise ValueError("Cannot recv on an RPC that has never been opened.")
 
         return next(call)