feat: retryable resource exhausted handling (#366)

BigQuery Storage Read API will start returning retryable
RESOURCE_EXHAUSTED errors in 2022 when certain concurrency limits are
hit, so this PR adds some code to handle them.

Tested with unit tests and system tests. System tests ran successfully
on a test project that intentionally returns retryable
RESOURCE_EXHAUSTED errors.

Co-authored-by: Tim Swast <swast@google.com>

google/cloud/bigquery_storage_v1/client.py

@@ -47,6 +47,7 @@ def read_rows(
         retry=google.api_core.gapic_v1.method.DEFAULT,
         timeout=google.api_core.gapic_v1.method.DEFAULT,
         metadata=(),
+        retry_delay_callback=None,
     ):
         """
         Reads rows from the table in the format prescribed by the read
@@ -108,6 +109,12 @@ def read_rows(
                 specified, the timeout applies to each individual attempt.
             metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata
                 that is provided to the method.
+            retry_delay_callback (Optional[Callable[[float], None]]):
+                If the client receives a retryable error that asks the client to
+                delay its next attempt and retry_delay_callback is not None,
+                BigQueryReadClient will call retry_delay_callback with the delay
+                duration (in seconds) before it starts sleeping until the next
+                attempt.
 
         Returns:
             ~google.cloud.bigquery_storage_v1.reader.ReadRowsStream:
@@ -122,20 +129,15 @@ def read_rows(
             ValueError: If the parameters are invalid.
         """
         gapic_client = super(BigQueryReadClient, self)
-        stream = gapic_client.read_rows(
-            read_stream=name,
-            offset=offset,
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-        return reader.ReadRowsStream(
-            stream,
+        stream = reader.ReadRowsStream(
             gapic_client,
             name,
             offset,
             {"retry": retry, "timeout": timeout, "metadata": metadata},
+            retry_delay_callback=retry_delay_callback,
         )
+        stream._reconnect()
+        return stream
 
 
 class BigQueryWriteClient(big_query_write.BigQueryWriteClient):

google/cloud/bigquery_storage_v1/reader.py

@@ -17,12 +17,14 @@
 import collections
 import io
 import json
+import time
 
 try:
     import fastavro
 except ImportError:  # pragma: NO COVER
     fastavro = None
 import google.api_core.exceptions
+import google.rpc.error_details_pb2
 
 try:
     import pandas
@@ -79,16 +81,17 @@ class ReadRowsStream(object):
     If the pandas and fastavro libraries are installed, use the
     :func:`~google.cloud.bigquery_storage_v1.reader.ReadRowsStream.to_dataframe()`
     method to parse all messages into a :class:`pandas.DataFrame`.
+
+    This object should not be created directly, but is returned by
+    other methods in this library.
     """
 
-    def __init__(self, wrapped, client, name, offset, read_rows_kwargs):
+    def __init__(
+        self, client, name, offset, read_rows_kwargs, retry_delay_callback=None
+    ):
         """Construct a ReadRowsStream.
 
         Args:
-            wrapped (Iterable[ \
-                ~google.cloud.bigquery_storage.types.ReadRowsResponse \
-            ]):
-                The ReadRows stream to read.
             client ( \
                 ~google.cloud.bigquery_storage_v1.services. \
                     big_query_read.BigQueryReadClient \
@@ -106,6 +109,12 @@ def __init__(self, wrapped, client, name, offset, read_rows_kwargs):
             read_rows_kwargs (dict):
                 Keyword arguments to use when reconnecting to a ReadRows
                 stream.
+            retry_delay_callback (Optional[Callable[[float], None]]):
+                If the client receives a retryable error that asks the client to
+                delay its next attempt and retry_delay_callback is not None,
+                ReadRowsStream will call retry_delay_callback with the delay
+                duration (in seconds) before it starts sleeping until the next
+                attempt.
 
         Returns:
             Iterable[ \
@@ -116,11 +125,12 @@ def __init__(self, wrapped, client, name, offset, read_rows_kwargs):
 
         # Make a copy of the read position so that we can update it without
         # mutating the original input.
-        self._wrapped = wrapped
         self._client = client
         self._name = name
         self._offset = offset
         self._read_rows_kwargs = read_rows_kwargs
+        self._retry_delay_callback = retry_delay_callback
+        self._wrapped = None
 
     def __iter__(self):
         """An iterable of messages.
@@ -131,9 +141,12 @@ def __iter__(self):
             ]:
                 A sequence of row messages.
         """
-
         # Infinite loop to reconnect on reconnectable errors while processing
         # the row stream.
+
+        if self._wrapped is None:
+            self._reconnect()
+
         while True:
             try:
                 for message in self._wrapped:
@@ -152,14 +165,53 @@ def __iter__(self):
             except _STREAM_RESUMPTION_EXCEPTIONS:
                 # Transient error, so reconnect to the stream.
                 pass
+            except Exception as exc:
+                if not self._resource_exhausted_exception_is_retryable(exc):
+                    raise
 
             self._reconnect()
 
     def _reconnect(self):
         """Reconnect to the ReadRows stream using the most recent offset."""
-        self._wrapped = self._client.read_rows(
-            read_stream=self._name, offset=self._offset, **self._read_rows_kwargs
-        )
+        while True:
+            try:
+                self._wrapped = self._client.read_rows(
+                    read_stream=self._name,
+                    offset=self._offset,
+                    **self._read_rows_kwargs
+                )
+                break
+            except Exception as exc:
+                if not self._resource_exhausted_exception_is_retryable(exc):
+                    raise
+
+    def _resource_exhausted_exception_is_retryable(self, exc):
+        if isinstance(exc, google.api_core.exceptions.ResourceExhausted):
+            # ResourceExhausted errors are only retried if a valid
+            # RetryInfo is provided with the error.
+            #
+            # TODO: Remove hasattr logic when we require google-api-core >= 2.2.0.
+            #       ResourceExhausted added details/_details in google-api-core 2.2.0.
+            details = None
+            if hasattr(exc, "details"):
+                details = exc.details
+            elif hasattr(exc, "_details"):
+                details = exc._details
+            if details is not None:
+                for detail in details:
+                    if isinstance(detail, google.rpc.error_details_pb2.RetryInfo):
+                        retry_delay = detail.retry_delay
+                        if retry_delay is not None:
+                            delay = max(
+                                0,
+                                float(retry_delay.seconds)
+                                + (float(retry_delay.nanos) / 1e9),
+                            )
+                            if self._retry_delay_callback:
+                                self._retry_delay_callback(delay)
+                            time.sleep(delay)
+                            return True
+        return False
 
     def rows(self, read_session=None):
         """Iterate over all rows in the stream.

google/cloud/bigquery_storage_v1beta2/client.py

@@ -48,6 +48,7 @@ def read_rows(
         retry=google.api_core.gapic_v1.method.DEFAULT,
         timeout=google.api_core.gapic_v1.method.DEFAULT,
         metadata=(),
+        retry_delay_callback=None,
     ):
         """
         Reads rows from the table in the format prescribed by the read
@@ -109,6 +110,12 @@ def read_rows(
                 specified, the timeout applies to each individual attempt.
             metadata (Optional[Sequence[Tuple[str, str]]]): Additional metadata
                 that is provided to the method.
+            retry_delay_callback (Optional[Callable[[float], None]]):
+                If the client receives a retryable error that asks the client to
+                delay its next attempt and retry_delay_callback is not None,
+                BigQueryReadClient will call retry_delay_callback with the delay
+                duration (in seconds) before it starts sleeping until the next
+                attempt.
 
         Returns:
             ~google.cloud.bigquery_storage_v1.reader.ReadRowsStream:
@@ -123,20 +130,15 @@ def read_rows(
             ValueError: If the parameters are invalid.
         """
         gapic_client = super(BigQueryReadClient, self)
-        stream = gapic_client.read_rows(
-            read_stream=name,
-            offset=offset,
-            retry=retry,
-            timeout=timeout,
-            metadata=metadata,
-        )
-        return reader.ReadRowsStream(
-            stream,
+        stream = reader.ReadRowsStream(
             gapic_client,
             name,
             offset,
             {"retry": retry, "timeout": timeout, "metadata": metadata},
+            retry_delay_callback=retry_delay_callback,
         )
+        stream._reconnect()
+        return stream
 
 
 class BigQueryWriteClient(big_query_write.BigQueryWriteClient):