turned backoff_generator into a class

Author: daniel-sanche

google/cloud/bigtable/data/_helpers.py

@@ -98,23 +98,32 @@ def _attempt_timeout_generator(
         yield max(0, min(per_request_timeout, deadline - time.monotonic()))
 
 
-def backoff_generator(initial=0.01, multiplier=2, maximum=60):
+class BackoffGenerator:
     """
-    Build a generator for exponential backoff sleep times.
+    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 the send(idx) method. This is
-    used by the Metrics class to track the sleep times used for each attempt.
+    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.
     """
-    history = []
-    subgenerator = exponential_sleep_generator(initial, multiplier, maximum)
-    while True:
-        next_backoff = next(subgenerator)
-        history.append(next_backoff)
-        sent_idx = yield next_backoff
-        while sent_idx is not None:
-            # requesting from history
-            sent_idx = yield history[sent_idx]
+
+    def __init__(self, initial=0.01, multiplier=2, maximum=60):
+        self.history = []
+        self.subgenerator = exponential_sleep_generator(initial, multiplier, maximum)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> float:
+        next_backoff = next(self.subgenerator)
+        self.history.append(next_backoff)
+        return next_backoff
+
+    def get_attempt_backoff(self, attempt_idx) -> float:
+        """
+        returns the backoff time for a specific attempt index, starting at 0.
+        """
+        return self.history[attempt_idx]
 
 
 def _retry_exception_factory(

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

@@ -31,6 +31,7 @@
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
+    from google.cloud.bigtable.data._helpers import BackoffGenerator
 
 
 LOGGER = logging.getLogger(__name__)
@@ -138,7 +139,7 @@ class ActiveOperationMetric:
     """
 
     op_type: OperationType
-    backoff_generator: Generator[float, int, None] | None = None
+    backoff_generator: BackoffGenerator | None = None
     # keep monotonic timestamps for active operations
     start_time_ns: int = field(default_factory=time.monotonic_ns)
     active_attempt: ActiveAttemptMetric | None = None
@@ -192,9 +193,8 @@ def start_attempt(self) -> None:
         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
-            backoff_ns = int(
-                self.backoff_generator.send(len(self.completed_attempts) - 1) * 1e9
-            )
+            backoff = self.backoff_generator.get_attempt_backoff(len(self.completed_attempts) - 1)
+            backoff_ns = int(backoff * 1e9)
         else:
             backoff_ns = 0