feat: add data model for client side metrics (#1187)

This PR revives googleapis/python-bigtable#923,
which was de-priotirized to work on the sync client. This PR brings it
back, working with both async and sync. It also adds a grpc interceptor,
as an improved way to capture metadata across both clients
---
## Design
The main architecture looks like this:
<img width="651" height="631"
alt="300137129-bebbb05a-20f0-45c2-9d38-e95a314edf64 drawio (1)"
src="https://github.com/user-attachments/assets/c8318ac8-5f18-4027-9f64-4a40a8ab1d79"
/>
Most of the work is done by the ActiveOperationMetric class, which is
instantiated with each rpc call, and updated through the lifecycle of
the call. When the rpc is complete, it will call `on_operation_complete`
and `on_attempt_complete` on the MetricsHandler, which can then log the
completed data into OpenTelemetry (or theoretically, other locations if
needed)
Note that there are separate classes for active vs completed metrics
(`ActiveOperationMetric`, `ActiveAttemptMetric`,
`CompletedOperationMetric`, `CompletedAttemptMetric`). This is so that
we can keep fields mutable and optional while the request is ongoing,
but pass down static immutable copies once the attempt is completed and
no new data is coming
---------
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
Co-authored-by: Mattie Fu <mattiefu@google.com>

Author: 3 people

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

@@ -88,6 +88,7 @@
 fromgoogle.cloud.bigtable.data.row_filtersimportStripValueTransformerFilter
 fromgoogle.cloud.bigtable.data.row_filtersimportCellsRowLimitFilter
 fromgoogle.cloud.bigtable.data.row_filtersimportRowFilterChain
+fromgoogle.cloud.bigtable.data._metricsimportBigtableClientSideMetricsController
 
 fromgoogle.cloud.bigtable.data._cross_syncimportCrossSync
 
@@ -1039,6 +1040,8 @@ def __init__(
 default_retryable_errorsor ()
  )
 
+self._metrics=BigtableClientSideMetricsController()
+
 try:
 self._register_instance_future=CrossSync.create_task(
 self.client._register_instance,
@@ -1753,6 +1756,7 @@ async def close(self):
 """
  Called to close the Table instance and release any resources held by it.
  """
+self._metrics.close()
 ifself._register_instance_future:
 self._register_instance_future.cancel()
 self.client._remove_instance_registration(

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

@@ -13,11 +13,21 @@
 # limitations under the License.
 from __future__ importannotations
 
+fromtypingimportSequence
+
+importtime
+fromfunctoolsimportwraps
+
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportActiveOperationMetric
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportOperationState
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportOperationType
+
 fromgoogle.cloud.bigtable.data._cross_syncimportCrossSync
 
 ifCrossSync.is_async:
 fromgrpc.aioimportUnaryUnaryClientInterceptor
 fromgrpc.aioimportUnaryStreamClientInterceptor
+fromgrpc.aioimportAioRpcError
 else:
 fromgrpcimportUnaryUnaryClientInterceptor
 fromgrpcimportUnaryStreamClientInterceptor
@@ -26,6 +36,57 @@
 __CROSS_SYNC_OUTPUT__="google.cloud.bigtable.data._sync_autogen.metrics_interceptor"
 
 
+def_with_active_operation(func):
+"""
+ Decorator for interceptor methods to extract the active operation associated with the
+ in-scope contextvars, and pass it to the decorated function.
+ """
+
+@wraps(func)
+defwrapper(self, continuation, client_call_details, request):
+operation: ActiveOperationMetric|None=ActiveOperationMetric.from_context()
+
+ifoperation:
+# start a new attempt if not started
+if (
+operation.state==OperationState.CREATED
+oroperation.state==OperationState.BETWEEN_ATTEMPTS
+ ):
+operation.start_attempt()
+# wrap continuation in logic to process the operation
+returnfunc(self, operation, continuation, client_call_details, request)
+else:
+# if operation not found, return unwrapped continuation
+returncontinuation(client_call_details, request)
+
+returnwrapper
+
+
+@CrossSync.convert
+asyncdef_get_metadata(source) ->dict[str, str|bytes] |None:
+"""Helper to extract metadata from a call or RpcError"""
+try:
+metadata: Sequence[tuple[str, str|bytes]]
+ifCrossSync.is_async:
+# grpc.aio returns metadata in Metadata objects
+ifisinstance(source, AioRpcError):
+metadata=list(source.trailing_metadata()) +list(
+source.initial_metadata()
+ )
+else:
+metadata=list(awaitsource.trailing_metadata()) +list(
+awaitsource.initial_metadata()
+ )
+else:
+# sync grpc returns metadata as a sequence of tuples
+metadata=source.trailing_metadata() +source.initial_metadata()
+# convert metadata to dict format
+return {k: vfor (k, v) inmetadata}
+exceptException:
+# ignore errors while fetching metadata
+returnNone
+
+
 @CrossSync.convert_class(sync_name="BigtableMetricsInterceptor")
 classAsyncBigtableMetricsInterceptor(
 UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor
@@ -35,21 +96,33 @@ class AsyncBigtableMetricsInterceptor(
  """
 
 @CrossSync.convert
-asyncdefintercept_unary_unary(self, continuation, client_call_details, request):
+@_with_active_operation
+asyncdefintercept_unary_unary(
+self, operation, continuation, client_call_details, request
+ ):
 """
  Interceptor for unary rpcs:
  - MutateRow
  - CheckAndMutateRow
  - ReadModifyWriteRow
  """
+metadata=None
 try:
 call=awaitcontinuation(client_call_details, request)
+metadata=await_get_metadata(call)
 returncall
 exceptExceptionasrpc_error:
+metadata=await_get_metadata(rpc_error)
 raiserpc_error
+finally:
+ifmetadataisnotNone:
+operation.add_response_metadata(metadata)
 
 @CrossSync.convert
-asyncdefintercept_unary_stream(self, continuation, client_call_details, request):
+@_with_active_operation
+asyncdefintercept_unary_stream(
+self, operation, continuation, client_call_details, request
+ ):
 """
  Interceptor for streaming rpcs:
  - ReadRows
@@ -58,21 +131,42 @@ async def intercept_unary_stream(self, continuation, client_call_details, reques
  """
 try:
 returnself._streaming_generator_wrapper(
-awaitcontinuation(client_call_details, request)
+operation, awaitcontinuation(client_call_details, request)
  )
 exceptExceptionasrpc_error:
 # handle errors while intializing stream
+metadata=await_get_metadata(rpc_error)
+ifmetadataisnotNone:
+operation.add_response_metadata(metadata)
 raiserpc_error
 
 @staticmethod
 @CrossSync.convert
-asyncdef_streaming_generator_wrapper(call):
+asyncdef_streaming_generator_wrapper(operation, call):
 """
  Wrapped generator to be returned by intercept_unary_stream.
  """
+# only track has_first response for READ_ROWS
+has_first_response= (
+operation.first_response_latency_nsisnotNone
+oroperation.op_type!=OperationType.READ_ROWS
+ )
+encountered_exc=None
 try:
 asyncforresponseincall:
+# record time to first response. Currently only used for READ_ROWs
+ifnothas_first_response:
+operation.first_response_latency_ns= (
+time.monotonic_ns() -operation.start_time_ns
+ )
+has_first_response=True
 yieldresponse
 exceptExceptionase:
 # handle errors while processing stream
-raisee
+encountered_exc=e
+raise
+finally:
+ifcallisnotNone:
+metadata=await_get_metadata(encountered_excorcall)
+ifmetadataisnotNone:
+operation.add_response_metadata(metadata)

packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py

@@ -23,6 +23,7 @@
 fromgoogle.cloud.bigtable.data.read_rows_queryimportReadRowsQuery
 
 fromgoogle.api_coreimportexceptionsascore_exceptions
+fromgoogle.api_core.retryimportexponential_sleep_generator
 fromgoogle.api_core.retryimportRetryFailureReason
 fromgoogle.cloud.bigtable.data.exceptionsimportRetryExceptionGroup
 
@@ -248,3 +249,61 @@ def _get_retryable_errors(
 call_codes=table.default_mutate_rows_retryable_errors
 
 return [_get_error_type(e) foreincall_codes]
+
+
+classTrackedBackoffGenerator:
+"""
+ Generator class for exponential backoff sleep times.
+ This implementation builds on top of api_core.retries.exponential_sleep_generator,
+ adding the ability to retrieve previous values using get_attempt_backoff(idx).
+ This is used by the Metrics class to track the sleep times used for each attempt.
+ """
+
+def__init__(self, initial=0.01, maximum=60, multiplier=2):
+self.history= []
+self.subgenerator=exponential_sleep_generator(
+initial=initial, maximum=maximum, multiplier=multiplier
+ )
+self._next_override: float|None=None
+
+def__iter__(self):
+returnself
+
+defset_next(self, next_value: float):
+"""
+ Set the next backoff value, instead of generating one from subgenerator.
+ After the value is yielded, it will go back to using self.subgenerator.
+
+ If set_next is called twice before the next() is called, only the latest
+ value will be used and others discarded
+
+ Args:
+ next_value: the upcomming value to yield when next() is called
+ Raises:
+ ValueError: if next_value is negative
+ """
+ifnext_value<0:
+raiseValueError("backoff value cannot be less than 0")
+self._next_override=next_value
+
+def__next__(self) ->float:
+ifself._next_overrideisnotNone:
+next_backoff=self._next_override
+self._next_override=None
+else:
+next_backoff=next(self.subgenerator)
+self.history.append(next_backoff)
+returnnext_backoff
+
+defget_attempt_backoff(self, attempt_idx) ->float:
+"""
+ returns the backoff time for a specific attempt index, starting at 0.
+
+ Args:
+ attempt_idx: the index of the attempt to return backoff for
+ Raises:
+ IndexError: if attempt_idx is negative, or not in history
+ """
+ifattempt_idx<0:
+raiseIndexError("received negative attempt number")
+returnself.history[attempt_idx]

packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/__init__.py

@@ -0,0 +1,35 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+fromgoogle.cloud.bigtable.data._metrics.metrics_controllerimport (
+BigtableClientSideMetricsController,
+)
+
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportActiveOperationMetric
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportActiveAttemptMetric
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportCompletedOperationMetric
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportCompletedAttemptMetric
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportOperationState
+fromgoogle.cloud.bigtable.data._metrics.data_modelimportOperationType
+fromgoogle.cloud.bigtable.data._metrics.tracked_retryimporttracked_retry
+
+__all__= (
+"BigtableClientSideMetricsController",
+"OperationType",
+"OperationState",
+"ActiveOperationMetric",
+"ActiveAttemptMetric",
+"CompletedOperationMetric",
+"CompletedAttemptMetric",
+"tracked_retry",
+)