feat: respect RetryInfo metadata from retryable exceptions

Author: daniel-sanche

google/cloud/bigtable/data/_helpers.py

@@ -25,6 +25,7 @@
 from google.api_core import exceptions as core_exceptions
 from google.api_core.retry import exponential_sleep_generator
 from google.api_core.retry import RetryFailureReason
+from google.rpc.error_details_pb2 import RetryInfo
 from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
 
 if TYPE_CHECKING:
@@ -288,11 +289,10 @@ def set_next(self, next_value: float):
         self._next_override = next_value
 
     def __next__(self) -> float:
+        next_backoff = next(self.subgenerator)
         if self._next_override is not None:
             next_backoff = self._next_override
             self._next_override = None
-        else:
-            next_backoff = next(self.subgenerator)
         self.history.append(next_backoff)
         return next_backoff
 
@@ -308,3 +308,15 @@ def get_attempt_backoff(self, attempt_idx) -> float:
         if attempt_idx < 0:
             raise IndexError("received negative attempt number")
         return self.history[attempt_idx]
+
+    def set_from_exception_info(self, retry_info: RetryInfo):
+        """
+        Use a RetryInfo object to set the next sleep time.
+
+        If a problem is encountered, this method does nothing.
+        """
+        try:
+            retry_seconds = retry_info.retry_delay.ToTimedelta().total_seconds()
+            self.set_next(retry_seconds)
+        except Exception:
+            pass

google/cloud/bigtable/data/_metrics/tracked_retry.py

@@ -26,6 +26,7 @@
 from grpc import StatusCode
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.retry import RetryFailureReason
+from google.rpc.error_details_pb2 import RetryInfo
 from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
 from google.cloud.bigtable.data._metrics import ActiveOperationMetric
@@ -43,10 +44,14 @@
 
 def _track_retryable_error(
     operation: ActiveOperationMetric,
+    backoff_generator: TrackedBackoffGenerator,
 ) -> Callable[[Exception], None]:
     """
     Used as input to api_core.Retry classes, to track when retryable errors are encountered
 
+    If an excemption is encountered with Retryinfo set, it will inform the backoff generator
+    to give it a chance to override the next backoff value
+
     Should be passed as on_error callback
     """
 
@@ -59,6 +64,11 @@ def wrapper(exc: Exception) -> None:
                     rpc_error.initial_metadata()
                 )
                 operation.add_response_metadata({k: v for k, v in metadata})
+                # check for RetryInfo:
+                if exc.details:
+                    info_matches = [field for field in exc.details if isinstance(field, RetryInfo)]
+                    if info_matches:
+                        backoff_generator.set_from_exception_info(info_matches[0])
         except Exception:
             # ignore errors in metadata collection
             pass
@@ -127,7 +137,7 @@ def tracked_retry(
     kwargs.pop("sleep_generator", None)
     return retry_fn(
         sleep_generator=operation.backoff_generator,
-        on_error=_track_retryable_error(operation),
+        on_error=_track_retryable_error(operation, operation.backoff_generator),
         exception_factory=_track_terminal_error(operation, in_exception_factory),
         **kwargs,
     )