Merge branch 'csm_1_data_model' into csm_2_instrumentation

Author: daniel-sanche

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

@@ -92,7 +92,8 @@ def register_operation(self, operation):
         operation.handlers.append(self)
 
     def on_operation_complete(self, op):
-        del self.operation_map[op.uuid]
+        if op.uuid in self.operation_map:
+            del self.operation_map[op.uuid]
 
     def on_operation_cancelled(self, op):
         self.on_operation_complete(op)

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

@@ -17,9 +17,15 @@
 
 from google.cloud.bigtable.data._metrics.data_model import OperationType
 from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
+from google.cloud.bigtable.data._metrics.data_model import ActiveAttemptMetric
+from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric
+from google.cloud.bigtable.data._metrics.data_model import CompletedAttemptMetric
 
 __all__ = (
     "BigtableClientSideMetricsController",
     "OperationType",
     "ActiveOperationMetric",
+    "ActiveAttemptMetric",
+    "CompletedOperationMetric",
+    "CompletedAttemptMetric",
 )

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

@@ -143,8 +143,11 @@ class ActiveOperationMetric:
     """
 
     op_type: OperationType
-    uuid: str = str(uuid.uuid4())
-    backoff_generator: TrackedBackoffGenerator | None = None
+    uuid: str = field(default_factory=lambda: str(uuid.uuid4()))
+    # create a default backoff generator, initialized with standard default backoff values
+    backoff_generator: TrackedBackoffGenerator = field(
+        default_factory=lambda: TrackedBackoffGenerator(initial=0.01, maximum=60, multiplier=2)
+    )
     # keep monotonic timestamps for active operations
     start_time_ns: int = field(default_factory=time.monotonic_ns)
     active_attempt: ActiveAttemptMetric | None = None
@@ -198,15 +201,16 @@ def start_attempt(self) -> ActiveAttemptMetric:
                 INVALID_STATE_ERROR.format("start_attempt", self.state)
             )
 
-        # find backoff value
-        if self.backoff_generator and len(self.completed_attempts) > 0:
-            # find the attempt's backoff by sending attempt number to generator
-            # generator will return the backoff time in seconds, so convert to nanoseconds
+        try:
+            # find backoff value before this attempt
+            prev_attempt_idx = len(self.completed_attempts) - 1
             backoff = self.backoff_generator.get_attempt_backoff(
-                len(self.completed_attempts) - 1
+                prev_attempt_idx
             )
+            # generator will return the backoff time in seconds, so convert to nanoseconds
             backoff_ns = int(backoff * 1e9)
-        else:
+        except IndexError:
+            # backoff value not found
             backoff_ns = 0
 
         self.active_attempt = ActiveAttemptMetric(backoff_before_attempt_ns=backoff_ns)
@@ -409,19 +413,19 @@ def _handle_error(message: str) -> None:
         full_message = f"Error in Bigtable Metrics: {message}"
         LOGGER.warning(full_message)
 
-    async def __aenter__(self):
+    def __aenter__(self):
         """
-        Implements the async context manager protocol for wrapping unary calls
+        Implements the async manager protocol
 
         Using the operation's context manager provides assurances that the operation
         is always closed when complete, with the proper status code automaticallty
         detected when an exception is raised.
         """
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    def __aexit__(self, exc_type, exc_val, exc_tb):
         """
-        Implements the async context manager protocol for wrapping unary calls
+        Implements the context manager protocol
 
         The operation is automatically ended on exit, with the status determined
         by the exception type and value.