diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 13bcfcc29..5755d35a7 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -288,11 +288,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 diff --git a/google/cloud/bigtable/data/_metrics/tracked_retry.py b/google/cloud/bigtable/data/_metrics/tracked_retry.py index 94d2e5dcb..3f8fea5ff 100644 --- a/google/cloud/bigtable/data/_metrics/tracked_retry.py +++ b/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 @@ -47,6 +48,9 @@ def _track_retryable_error( """ Used as input to api_core.Retry classes, to track when retryable errors are encountered + If an exception 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 +63,13 @@ 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 = next((field for field in exc.details if isinstance(field, RetryInfo)), None) + if info: + # override next backoff with server-provided value + retry_seconds = info.retry_delay.ToTimedelta().total_seconds() + operation.backoff_generator.set_next(retry_seconds) except Exception: # ignore errors in metadata collection pass