improved retry instrumentation

Author: daniel-sanche

google/cloud/bigtable/data/_async/_mutate_rows.py

@@ -21,7 +21,7 @@
 import google.cloud.bigtable_v2.types.bigtable as types_pb
 import google.cloud.bigtable.data.exceptions as bt_exceptions
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._helpers import _tracked_exception_factory
+from google.cloud.bigtable.data._helpers import _retry_exception_factory
 
 # mutate_rows requests are limited to this number of mutations
 from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT
@@ -106,7 +106,8 @@ def __init__(
             self.is_retryable,
             metric.backoff_generator,
             operation_timeout,
-            exception_factory=_tracked_exception_factory(self._operation_metric),
+            exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory),
+            on_error=self._operation_metric.track_retryable_error(),
         )
         # initialize state
         self.timeout_generator = _attempt_timeout_generator(

google/cloud/bigtable/data/_async/_read_rows.py

@@ -32,7 +32,7 @@
 from google.cloud.bigtable.data.exceptions import _RowSetComplete
 from google.cloud.bigtable.data.exceptions import _ResetRow
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._helpers import _tracked_exception_factory
+from google.cloud.bigtable.data._helpers import _retry_exception_factory
 
 from google.api_core import retry as retries
 
@@ -123,7 +123,8 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
             self._predicate,
             self._operation_metric.backoff_generator,
             self.operation_timeout,
-            exception_factory=_tracked_exception_factory(self._operation_metric),
+            exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory),
+            on_error=self._operation_metric.track_retryable_error(),
         )
 
     def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:

google/cloud/bigtable/data/_async/client.py

@@ -73,7 +73,6 @@
 from google.cloud.bigtable.data._helpers import _WarmedInstanceKey
 from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
-from google.cloud.bigtable.data._helpers import _tracked_exception_factory
 from google.cloud.bigtable.data._helpers import _validate_timeouts
 from google.cloud.bigtable.data._helpers import _get_error_type
 from google.cloud.bigtable.data._helpers import _get_retryable_errors
@@ -1392,7 +1391,8 @@ async def execute_rpc():
                 predicate,
                 operation_metric.backoff_generator,
                 operation_timeout,
-                exception_factory=_tracked_exception_factory(operation_metric),
+                exception_factory=operation_metric.track_terminal_error(_retry_exception_factory),
+                on_error=operation_metric.track_retryable_error(),
             )
 
     @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"})
@@ -1526,7 +1526,8 @@ async def mutate_row(
                 predicate,
                 operation_metric.backoff_generator,
                 operation_timeout,
-                exception_factory=_retry_exception_factory,
+                exception_factory=operation_metric.track_terminal_error(_retry_exception_factory),
+                on_error=operation_metric.track_retryable_error(),
             )
 
     @CrossSync.convert

google/cloud/bigtable/data/_async/metrics_interceptor.py

@@ -80,16 +80,6 @@ def wrapper(self, continuation, client_call_details, request):
 
     return wrapper
 
-
-def _end_attempt(operation, exc, metadata):
-    """Helper to add metadata and exception to an operation"""
-    if metadata is not None:
-        operation.add_response_metadata(metadata)
-    if exc is not None:
-        # end attempt. If it succeeded, let higher levels decide when to end operation
-        operation.end_attempt_with_status(exc)
-
-
 @CrossSync.convert
 async def _get_metadata(source) -> dict[str, str|bytes] | None:
     """Helper to extract metadata from a call or RpcError"""
@@ -129,7 +119,6 @@ def register_operation(self, operation):
         When registered, the operation will receive metadata updates:
         - start_attempt if attempt not started when rpc is being sent
         - add_response_metadata after call is complete
-        - end_attempt_with_status if attempt receives an error
 
         The interceptor will register itself as a handeler for the operation,
         so it can unregister the operation when it is complete
@@ -149,23 +138,17 @@ def on_operation_cancelled(self, op):
     async def intercept_unary_unary(
         self, operation, continuation, client_call_details, request
     ):
-        encountered_status: Exception | StatusCode | None = None
         metadata = None
         try:
             call = await continuation(client_call_details, request)
             metadata = await _get_metadata(call)
-            if CrossSync.is_async:
-                encountered_status = await call.code()
-            elif isinstance(call, Exception):
-                # sync unary calls return exception objects without raising
-                encountered_status = call
             return call
         except Exception as rpc_error:
             metadata = await _get_metadata(rpc_error)
-            encountered_status = rpc_error
             raise rpc_error
         finally:
-            _end_attempt(operation, encountered_status, metadata)
+            if metadata is not None:
+                operation.add_response_metadata(metadata)
 
     @CrossSync.convert
     @_with_operation_from_metadata
@@ -195,11 +178,13 @@ async def response_wrapper(call):
             finally:
                 if call is not None:
                     metadata = await _get_metadata(encountered_exc or call)
-                    _end_attempt(operation, encountered_exc, metadata)
+                    if metadata is not None:
+                        operation.add_response_metadata(metadata)
 
         try:
             return response_wrapper(await continuation(client_call_details, request))
         except Exception as rpc_error:
-            # handle errors while intializing stream
-            _end_attempt(operation, rpc_error, await _get_metadata(rpc_error))
+            metadata = await _get_metadata(rpc_error)
+            if metadata is not None:
+                operation.add_response_metadata(metadata)
             raise rpc_error

google/cloud/bigtable/data/_helpers.py

@@ -31,7 +31,6 @@
     import grpc
     from google.cloud.bigtable.data._async.client import _DataApiTargetAsync
     from google.cloud.bigtable.data._sync_autogen.client import _DataApiTarget
-    from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
 
 """
 Helper functions used in various places in the library.
@@ -121,56 +120,6 @@ def _retry_exception_factory(
     return source_exc, cause_exc
 
 
-def _tracked_exception_factory(
-    operation: "ActiveOperationMetric",
-) -> Callable[
-    [list[Exception], RetryFailureReason, float | None],
-    tuple[Exception, Exception | None],
-]:
-    """
-    wraps and extends _retry_exception_factory to add client-side metrics tracking.
-
-    When the rpc raises a terminal error, record any discovered metadata and finalize
-    the associated operation
-
-    Used by streaming rpcs, which can't always be perfectly captured by context managers
-    for operation termination.
-
-    Args:
-        exc_list: list of exceptions encountered during operation
-        is_timeout: whether the operation failed due to timeout
-        timeout_val: the operation timeout value in seconds, for constructing
-            the error message
-        operation: the operation to finalize when an exception is built
-    Returns:
-        tuple[Exception, Exception|None]:
-            tuple of the exception to raise, and a cause exception if applicabl
-    """
-
-    def wrapper(
-        exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None
-    ) -> tuple[Exception, Exception | None]:
-        source_exc, cause_exc = _retry_exception_factory(exc_list, reason, timeout_val)
-        try:
-            # record metadata from failed rpc
-            if (
-                isinstance(source_exc, core_exceptions.GoogleAPICallError)
-                and source_exc.errors
-            ):
-                rpc_error = source_exc.errors[-1]
-                metadata = list(rpc_error.trailing_metadata()) + list(
-                    rpc_error.initial_metadata()
-                )
-                operation.add_response_metadata({k: v for k, v in metadata})
-        except Exception:
-            # ignore errors in metadata collection
-            pass
-        operation.end_with_status(source_exc)
-        return source_exc, cause_exc
-
-    return wrapper
-
-
 def _get_timeouts(
     operation: float | TABLE_DEFAULT,
     attempt: float | None | TABLE_DEFAULT,