instrumented read_rows

Author: daniel-sanche

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

@@ -17,6 +17,8 @@
 
 from typing import Sequence, TYPE_CHECKING
 
+import time
+
 from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB
 from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB
 from google.cloud.bigtable_v2.types import RowSet as RowSetPB
@@ -29,13 +31,16 @@
 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 _retry_exception_factory
+from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
 
 from google.api_core import retry as retries
 from google.api_core.retry import exponential_sleep_generator
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
 if TYPE_CHECKING:
+    from google.cloud.bigtable.data._metrics import ActiveOperationMetric
+
     if CrossSync.is_async:
         from google.cloud.bigtable.data._async.client import (
             _DataApiTargetAsync as TargetType,
@@ -64,6 +69,7 @@ class _ReadRowsOperationAsync:
         target: The table or view to send the request to
         operation_timeout: The total time to allow for the operation, in seconds
         attempt_timeout: The time to allow for each individual attempt, in seconds
+        metric: the metric object representing the active operation
         retryable_exceptions: A list of exceptions that should trigger a retry
     """
 
@@ -75,6 +81,7 @@ class _ReadRowsOperationAsync:
         "_predicate",
         "_last_yielded_row_key",
         "_remaining_count",
+        "_operation_metric",
     )
 
     def __init__(
@@ -83,6 +90,7 @@ def __init__(
         target: TargetType,
         operation_timeout: float,
         attempt_timeout: float,
+        metric: ActiveOperationMetric,
         retryable_exceptions: Sequence[type[Exception]] = (),
     ):
         self.attempt_timeout_gen = _attempt_timeout_generator(
@@ -101,6 +109,7 @@ def __init__(
         self._predicate = retries.if_exception_type(*retryable_exceptions)
         self._last_yielded_row_key: bytes | None = None
         self._remaining_count: int | None = self.request.rows_limit or None
+        self._operation_metric = metric
 
     def start_operation(self) -> CrossSync.Iterable[Row]:
         """
@@ -109,10 +118,13 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
         Yields:
             Row: The next row in the stream
         """
+        self._operation_metric.backoff_generator = TrackedBackoffGenerator(
+            0.01, 60, multiplier=2
+        )
         return CrossSync.retry_target_stream(
             self._read_rows_attempt,
             self._predicate,
-            exponential_sleep_generator(0.01, 60, multiplier=2),
+            self._operation_metric.backoff_generator,
             self.operation_timeout,
             exception_factory=_retry_exception_factory,
         )
@@ -127,6 +139,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
         Yields:
             Row: The next row in the stream
         """
+        self._operation_metric.start_attempt()
         # revise request keys and ranges between attempts
         if self._last_yielded_row_key is not None:
             # if this is a retry, try to trim down the request to avoid ones we've already processed
@@ -137,20 +150,20 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
                 )
             except _RowSetComplete:
                 # if we've already seen all the rows, we're done
-                return self.merge_rows(None)
+                return self.merge_rows(None, self._operation_metric)
         # revise the limit based on number of rows already yielded
         if self._remaining_count is not None:
             self.request.rows_limit = self._remaining_count
             if self._remaining_count == 0:
-                return self.merge_rows(None)
+                return self.merge_rows(None, self._operation_metric)
         # create and return a new row merger
         gapic_stream = self.target.client._gapic_client.read_rows(
             self.request,
             timeout=next(self.attempt_timeout_gen),
             retry=None,
         )
         chunked_stream = self.chunk_stream(gapic_stream)
-        return self.merge_rows(chunked_stream)
+        return self.merge_rows(chunked_stream, self._operation_metric)
 
     @CrossSync.convert()
     async def chunk_stream(
@@ -210,6 +223,7 @@ async def chunk_stream(
     )
     async def merge_rows(
         chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None,
+        operation_metric: ActiveOperationMetric,
     ) -> CrossSync.Iterable[Row]:
         """
         Merge chunks into rows
@@ -222,6 +236,7 @@ async def merge_rows(
         if chunks is None:
             return
         it = chunks.__aiter__()
+        is_first_row = True
         # For each row
         while True:
             try:
@@ -304,7 +319,17 @@ async def merge_rows(
                         Cell(value, row_key, family, qualifier, ts, list(labels))
                     )
                     if c.commit_row:
+                        if is_first_row:
+                            # record first row latency in metrics
+                            is_first_row = False
+                            operation_metric.attempt_first_response()
+                        block_time = time.monotonic()
                         yield Row(row_key, cells)
+                        # most metric operations use setters, but this one updates
+                        # the value directly to avoid extra overhead
+                        operation_metric.active_attempt.application_blocking_time_ms += (  # type: ignore
+                            time.monotonic() - block_time
+                        ) * 1000
                         break
                     c = await it.__anext__()
             except _ResetRow as e:

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

@@ -125,6 +125,7 @@
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._helpers import RowKeySamples
     from google.cloud.bigtable.data._helpers import ShardedQuery
+    from google.cloud.bigtable.data._metrics import ActiveOperationMetric
 
     if CrossSync.is_async:
         from google.cloud.bigtable.data._async.mutations_batcher import (
@@ -1020,14 +1021,18 @@ async def read_rows_stream(
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
 
-        row_merger = CrossSync._ReadRowsOperation(
-            query,
-            self,
-            operation_timeout=operation_timeout,
-            attempt_timeout=attempt_timeout,
-            retryable_exceptions=retryable_excs,
-        )
-        return row_merger.start_operation()
+        with self._metrics.create_operation(
+            OperationType.READ_ROWS, streaming=True
+        ) as operation_metric:
+            row_merger = CrossSync._ReadRowsOperation(
+                query,
+                self,
+                operation_timeout=operation_timeout,
+                attempt_timeout=attempt_timeout,
+                metric=operation_metric,
+                retryable_exceptions=retryable_excs,
+            )
+            return row_merger.start_operation()
 
     @CrossSync.convert
     async def read_rows(
@@ -1077,7 +1082,9 @@ async def read_rows(
         )
         return [row async for row in row_generator]
 
-    @CrossSync.convert
+    @CrossSync.convert(
+        replace_symbols={"__anext__": "__next__", "StopAsyncIteration": "StopIteration"}
+    )
     async def read_row(
         self,
         row_key: str | bytes,
@@ -1117,15 +1124,28 @@ async def read_row(
         if row_key is None:
             raise ValueError("row_key must be string or bytes")
         query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1)
-        results = await self.read_rows(
-            query,
-            operation_timeout=operation_timeout,
-            attempt_timeout=attempt_timeout,
-            retryable_errors=retryable_errors,
+
+        operation_timeout, attempt_timeout = _get_timeouts(
+            operation_timeout, attempt_timeout, self
         )
-        if len(results) == 0:
-            return None
-        return results[0]
+        retryable_excs = _get_retryable_errors(retryable_errors, self)
+
+        with self._metrics.create_operation(
+            OperationType.READ_ROWS, streaming=False
+        ) as operation_metric:
+            row_merger = CrossSync._ReadRowsOperation(
+                query,
+                self,
+                operation_timeout=operation_timeout,
+                attempt_timeout=attempt_timeout,
+                metric=operation_metric,
+                retryable_exceptions=retryable_excs,
+            )
+            results_generator = row_merger.start_operation()
+            try:
+                return results_generator.__anext__()
+            except StopAsyncIteration:
+                return None
 
     @CrossSync.convert
     async def read_rows_sharded(
@@ -1264,20 +1284,17 @@ async def row_exists(
                 from any retries that failed
             google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
         """
-        if row_key is None:
-            raise ValueError("row_key must be string or bytes")
-
         strip_filter = StripValueTransformerFilter(flag=True)
         limit_filter = CellsRowLimitFilter(1)
         chain_filter = RowFilterChain(filters=[limit_filter, strip_filter])
-        query = ReadRowsQuery(row_keys=row_key, limit=1, row_filter=chain_filter)
-        results = await self.read_rows(
-            query,
+        result = await self.read_row(
+            row_key=row_key,
+            row_filter=chain_filter,
             operation_timeout=operation_timeout,
             attempt_timeout=attempt_timeout,
             retryable_errors=retryable_errors,
         )
-        return len(results) > 0
+        return result is not None
 
     @CrossSync.convert
     async def sample_row_keys(
@@ -1334,6 +1351,7 @@ async def sample_row_keys(
         async with self._metrics.create_operation(
             OperationType.SAMPLE_ROW_KEYS, backoff_generator=sleep_generator
         ):
+
             @CrossSync.convert
             async def execute_rpc():
                 results = await self.client._gapic_client.sample_row_keys(
@@ -1598,14 +1616,14 @@ async def check_and_mutate_row(
             false_case_mutations = [false_case_mutations]
         false_case_list = [m._to_pb() for m in false_case_mutations or []]
 
-        async with self._metrics.create_operation(
-            OperationType.CHECK_AND_MUTATE
-        ):
+        async with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE):
             result = await self.client._gapic_client.check_and_mutate_row(
                 request=CheckAndMutateRowRequest(
                     true_mutations=true_case_list,
                     false_mutations=false_case_list,
-                    predicate_filter=predicate._to_pb() if predicate is not None else None,
+                    predicate_filter=predicate._to_pb()
+                    if predicate is not None
+                    else None,
                     row_key=row_key.encode("utf-8")
                     if isinstance(row_key, str)
                     else row_key,
@@ -1656,9 +1674,7 @@ async def read_modify_write_row(
         if not rules:
             raise ValueError("rules must contain at least one item")
 
-        async with self._metrics.create_operation(
-            OperationType.READ_MODIFY_WRITE
-        ):
+        async with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE):
             result = await self.client._gapic_client.read_modify_write_row(
                 request=ReadModifyWriteRowRequest(
                     rules=[rule._to_pb() for rule in rules],